Darrien McKenzie

Ph.D. Student @UCSanDiego

[Darrien McKenzie @darrienmckenzie]

RL training for LLMs involves exposure to problems in the “Goldilocks zone” of difficulty: not too hard, not too easy.

But difficulty is not the only thing that matters.

Problem type matters, too.

And LLMs do not see or organize problem types the way humans do.

This is the starting point of Manifold Bandits.

The fundamental issue is this: when using RL to train LLMs, there will always be problems that are more productive for the policy at a given training iteration. Across a dataset, there are an astronomical number of possible batch orderings; uniform sampling is unlikely to consistently land on the most useful ones by chance.

This is why many adaptive curriculum learning and difficulty-aware methods focus on the policy’s “optimal learning zone,” avoiding examples that are already trivial or still impossible for the LLM to solve.

But problem type is usually treated differently.

Task decompositions within a training set are often manually defined according to human semantics, or ignored entirely, with an entire dataset/domain treated as “one task,” despite the immense heterogeneity seen even in small training datasets.

It is generally understood that LLMs may not share our concept of difficulty. But it seems less recognized that LLMs may not share our concept of problem type, either. In both cases, training the LLM according to our own semantics can obscure fine-grained learning dynamics that could otherwise be exploited.

In this work, we seek to derive an adaptive curriculum learning method that caters not only to the policy’s ability, but also its perception: its latent organization of tasks.

Driven by the manifold hypothesis and the Bitter Lesson, we leverage computation to derive curriculum structure from the policy’s latent geometry, then use Bayesian inference over that structure to guide search and learning over the training set. The result is Bayesian Manifold Curriculum (BMC): an algorithm that does not just try to find problems of the right difficulty, but instead orchestrates training effort across diverse and interacting problem types.

More technically, we frame problem sampling for LLMs as a manifold-structured bandit problem with endogenous non-stationarity. The tl;dr is that the problems/arms are diverse, structured, and interact through policy updates, so the typical non-stationary bandit framing is not quite the right fit for training LLMs. BMC is derived from this new framing.

🌐Website: https://t.co/iKVjNK92up
📰Paper: https://t.co/j08GktD6Qv
💻Code: https://t.co/ki6O3ojMWO

🧵1/N

[Darrien McKenzie @darrienmckenzie]

There are two central lessons to take from this work.

First: there is far more to curriculum learning than difficulty.

A sampling method can be highly productive, but have low type coverage, limiting diversity and potentially harming generalization.

A sampler can be productive and diverse, but still allocate toward problem types that are less useful for the capabilities we care about, limiting utility and harming benchmark performance.

A sampler can be diverse and relevant to the desired capabilities, but focus on problems that are too easy or too hard, limiting productivity and wasting compute.

And so on.

Curriculum learning is not one-dimensional. It is better described as a multi-objective allocation problem across productivity, diversity, utility, and likely other variables we have not yet isolated.

The second lesson is a familiarly bitter one: LLMs are not humans.

They do not see or organize problems the way we do.
They do not necessarily share our notion of difficulty.
They do not necessarily share our concept of “task,” “domain,” or “problem type.”
And they may not transfer abilities between problems in the way we expect.

Human semantics can be useful. But forcing them too deeply into the training process can obscure the fine-grained learning dynamics that actually govern policy improvement.

The policy has its own rich geometry: the geometry of its latent space, the geometry of its gradients, the geometry of its optimizer trajectories.

Curriculum learning should use it.

🌐Website: https://t.co/iKVjNK92up
📰Paper: https://t.co/j08GktD6Qv
💻Code: https://t.co/ki6O3ojMWO

11/N

[Darrien McKenzie @darrienmckenzie]

RL training for LLMs involves exposure to problems in the “Goldilocks zone” of difficulty: not too hard, not too easy.

But difficulty is not the only thing that matters.

Problem type matters, too.

And LLMs do not see or organize problem types the way humans do.

This is the starting point of Manifold Bandits.

The fundamental issue is this: when using RL to train LLMs, there will always be problems that are more productive for the policy at a given training iteration. Across a dataset, there are an astronomical number of possible batch orderings; uniform sampling is unlikely to consistently land on the most useful ones by chance.

This is why many adaptive curriculum learning and difficulty-aware methods focus on the policy’s “optimal learning zone,” avoiding examples that are already trivial or still impossible for the LLM to solve.

But problem type is usually treated differently.

Task decompositions within a training set are often manually defined according to human semantics, or ignored entirely, with an entire dataset/domain treated as “one task,” despite the immense heterogeneity seen even in small training datasets.

It is generally understood that LLMs may not share our concept of difficulty. But it seems less recognized that LLMs may not share our concept of problem type, either. In both cases, training the LLM according to our own semantics can obscure fine-grained learning dynamics that could otherwise be exploited.

In this work, we seek to derive an adaptive curriculum learning method that caters not only to the policy’s ability, but also its perception: its latent organization of tasks.

Driven by the manifold hypothesis and the Bitter Lesson, we leverage computation to derive curriculum structure from the policy’s latent geometry, then use Bayesian inference over that structure to guide search and learning over the training set. The result is Bayesian Manifold Curriculum (BMC): an algorithm that does not just try to find problems of the right difficulty, but instead orchestrates training effort across diverse and interacting problem types.

More technically, we frame problem sampling for LLMs as a manifold-structured bandit problem with endogenous non-stationarity. The tl;dr is that the problems/arms are diverse, structured, and interact through policy updates, so the typical non-stationary bandit framing is not quite the right fit for training LLMs. BMC is derived from this new framing.

🌐Website: https://t.co/iKVjNK92up
📰Paper: https://t.co/j08GktD6Qv
💻Code: https://t.co/ki6O3ojMWO

🧵1/N

[Darrien McKenzie @darrienmckenzie]

Suppose benchmark performance was the end goal, and not just a passive measurement of generalization. Could we ensure that sampled problems are productive, diverse, and relevant to the specific capabilities we care about?

Sure. We can use the tree for this, too.

To do this, we make an assumption: if problems are closer together in the policy’s latent geometry, they are more likely to affect each other during training.

Put another way: learnability propagates locally. Training on problems involving polynomials will probably help more with other polynomial problems than with arbitrary math problems, for instance.

Why does that matter? Well, all prompts are embedded in the policy’s latent space. That means:

The evaluations themselves lie on the same manifold as the training set.

So, if we target training problems that are geometrically closer to evaluation problems, we may be able to improve those evaluations without training on them directly.

For this, we introduce BMC-T, a targeting variant of BMC. BMC-T builds a tree over both the training set and a chosen target set: evals, a subset of the evals, another dataset, etc. The target examples cannot be sampled (we are not training on the test set). They only provide a utility bias: tree regions with more train/target overlap receive a larger bonus, steering BMC toward training problems near the targets.

Still, if the target set is the evals themselves, they are no longer “held out” in the usual sense. That’s fine for our purpose. Our goal is to show that some training problems are more relevant to some evaluations than others, and that utility is independent from productivity and diversity.

When we target all evals, or AIME2024 specifically, productivity remains roughly constant. But the evaluations change meaningfully, and they change predictably: targeting AIME2024 improves not only AIME2024, but also other competition-math benchmarks.

Thus, utility (evaluation relevance) is different from productivity.

Even if two samplers expose the policy to problems of similar “optimal difficulty,” they may not improve the same evaluations, because not all problem types affect each benchmark equally. Problem type matters.

Moreover, while coarse human-defined categories like “Math” or “Code” are better than nothing, they are not expressive enough to capture the true scale of type heterogeneity. It is clear that this dataset is neither “one task” nor an unstructured bag of prompts.

These results also suggest that the tree is not just “for show.” It partially encodes relationships between problems, and hints at how transfer may occur over the policy’s latent geometry, whether those problems are sampled or not.

Do we know exactly which training problems improve which evals? No. We don’t need to.

The policy knows better.

10/N

[Darrien McKenzie @darrienmckenzie]

Hello world!

Starting this account for research updates. I’m a CS Ph.D. student at UC San Diego working on… stuff. Still figuring out how much I want to put myself in a specific bucket.

Currently wrapping up my first (long overdue) major project. More soon.