🤖 machine learning

Manifold Bandits: Bayesian Curriculum Learning over the Latent Geometry of Large Language Models

This paper introduces Bayesian Manifold Curriculum (BMC), a structure-aware framework that reframes problem sampling in LLM reinforcement learning as a manifold-structured bandit problem to optimize the trade-off between learning productivity, task diversity, and evaluation utility by leveraging the latent geometric relationships between tasks rather than relying solely on difficulty-based selection.

Original authors: Darrien McKenzie, Nicklas Hansen, Xiaolong Wang

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Darrien McKenzie, Nicklas Hansen, Xiaolong Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a brilliant but very young student (the Large Language Model) how to solve complex puzzles. You have a massive library of puzzles ranging from "How do I tie my shoes?" to "How do I build a rocket?"

In the past, researchers tried to teach this student using two main strategies:

Random Guessing: Just picking puzzles out of a hat. This is slow because the student might get bored with easy ones or give up on impossible ones.
The "Goldilocks" Strategy: Only picking puzzles that are "just right"—not too easy, not too hard. This works well, but it has a flaw: it treats every puzzle as a completely separate, isolated item. It ignores the fact that "building a rocket" and "flying a plane" are related. If the student learns about aerodynamics in one, they should be better at the other.

This paper introduces a new, smarter way to teach called Bayesian Manifold Curriculum (BMC). Here is how it works, using simple analogies:

1. The "Mental Map" (The Latent Task Tree)

Instead of treating every puzzle as a separate, unrelated item, BMC looks at the student's internal brain map.

Imagine the student has a mental map where similar ideas are located near each other. "Math puzzles" are in one neighborhood, "Coding puzzles" in another, and "Science puzzles" in a third. Within the "Math" neighborhood, there are smaller streets for "Algebra" and "Calculus."

BMC builds a hierarchical tree (like a family tree) based on this mental map. It doesn't need a human teacher to label these groups; it figures out the structure by looking at how the student's brain represents the problems.

The Analogy: Think of it like organizing a library not just by "Fiction" and "Non-Fiction," but by creating a dynamic map that shows which books are actually similar in storytelling style, even if they are in different sections.

2. The "Smart Tour Guide" (Bayesian Learning)

Once the map is built, BMC acts as a tour guide who decides which puzzles to show the student next. It uses a method called Bayesian Learning, which is basically "learning from experience and updating beliefs."

The "Surprise" Factor: If the student solves a puzzle they were expected to fail, the guide gets excited. If they fail a puzzle they were expected to ace, the guide gets curious.
The Ripple Effect: This is the magic part. If the student learns something new about "Algebra," the guide doesn't just update the "Algebra" section of the map. Because the map knows that "Algebra" is connected to "Calculus," the guide also updates their belief about how well the student will do on Calculus problems, even if they haven't seen one yet.
The Analogy: Imagine you are teaching someone to drive. If they get better at parallel parking, you don't just think, "Good at parking." You also update your belief that they are probably getting better at "maneuvering in tight spaces" in general, which helps you decide what to teach them next.

3. The Three Balancing Acts

The paper argues that a good teacher needs to balance three things, which the authors call Productivity, Diversity, and Utility.

Productivity (The "Learning Signal"): Are we picking puzzles that actually teach the student something new? (i.e., puzzles that are hard enough to be interesting but not impossible).
Diversity (The "Coverage"): Are we only teaching the student about one type of puzzle (like only math), or are we making sure they see a wide variety of problems (math, coding, science)? If you only teach math, the student might fail at science.
Utility (The "Real-World Test"): Are we teaching the specific skills the student will need for the final exam? Sometimes, the most "productive" puzzles (the ones that give the best learning signal) aren't the ones that help the most on the final test.

4. The Result: A Smarter Curriculum

The researchers tested this on a math dataset. They found that:

Old methods (like just picking "medium difficulty" puzzles) often got stuck in a loop, ignoring rare but important types of problems.
BMC managed to find the "sweet spot." It kept the student challenged (Productivity) while ensuring they didn't miss out on rare problem types (Diversity).
The "Utility" Twist: They even created a version called BMC-T that could be told, "Hey, the final exam is going to be mostly about Chinese math problems." The system then subtly shifted its teaching focus to cover those specific areas better, without losing the benefits of the other types of problems.

The Bottom Line

This paper says that to teach an AI effectively, you can't just look at how hard a problem is. You have to understand how the AI sees the relationship between problems. By building a map of these relationships and using a smart, updating guide to pick the next lesson, you can teach the AI faster, more broadly, and better prepare it for the specific tests it will face.

It's the difference between a teacher who just hands you random worksheets and a teacher who understands your brain, knows how your skills connect, and designs a custom curriculum that hits the sweet spot between challenge, variety, and real-world relevance.

Technical Summary: Manifold Bandits – Bayesian Curriculum Learning over the Latent Geometry of Large Language Models

1. Problem Statement

Reinforcement Learning (RL) has become a primary method for enhancing the reasoning capabilities of Large Language Models (LLMs), particularly through group-relative optimization algorithms like GRPO, DAPO, and GSPO. These methods rely on reward variance within a group of rollouts to generate policy gradients. If a model consistently solves or consistently fails a prompt, the reward variance is zero, yielding no learning signal and wasting computational resources.

Existing adaptive curriculum learning methods attempt to address this by prioritizing prompts of intermediate difficulty (those likely to yield non-zero variance). However, these approaches typically frame problem selection as a standard multi-armed bandit problem with independent arms. This framing overlooks two critical dynamics of LLM training:

Endogenous Non-Stationarity: The reward distribution over prompts evolves not just due to external factors, but as a direct consequence of the model's own updates and sampling decisions.
Structured Heterogeneity: Problems are not independent; prompts of similar types share reasoning patterns and latent features. Ignoring these relationships leads to inefficient exploration.

Furthermore, current strategies often focus solely on productivity (maximizing immediate learning signal) while neglecting diversity (coverage of the task space) and utility (relevance to downstream evaluation benchmarks). This can lead to narrow coverage and suboptimal generalization.

2. Methodology

The authors propose Bayesian Manifold Curriculum (BMC), a framework that reframes problem sampling as a manifold-structured bandit problem with endogenous non-stationarity. The methodology consists of three core components:

A. Latent Task Trees

Instead of treating prompts as independent arms, BMC constructs a hierarchical Latent Task Tree directly from the policy model's internal representations.

Construction: The tree is built recursively using embeddings from an intermediate layer of the policy model. At each node, embeddings are reduced via PCA, projected via UMAP, and clustered using HDBSCAN.
Stopping Criterion: Recursion halts when a region passes a "Chart Test," verifying that the local manifold is approximately smooth and coherent (low intrinsic dimensionality and connectivity).
Purpose: This hierarchy approximates the model's latent task manifold, organizing prompts into multi-scale regions that reflect the model's internal perception of task relationships, rather than relying on external annotations or fixed taxonomies.

B. Bayesian Manifold Curriculum (BMC)

BMC performs structured decision-making over the Latent Task Tree using a three-step cycle:

Top-Down Selection (Hierarchical Thompson Sampling): To form a training batch, agents traverse the tree from root to leaf. At each node, they sample reward estimates for child nodes from a Gaussian belief distribution ( $\mathcal{N}(\mu, \sigma^2)$ ) and select the child with the highest sample. This balances exploration (diversity) and exploitation (productivity) across different levels of abstraction.
Non-Stationary Belief Modeling (Bayesian Filtering): Leaf nodes (individual prompts) maintain beliefs about their expected learning signal (reward variance). These beliefs are updated online using a lightweight Bayesian filter inspired by the Normal–Inverse-Gamma model. The update mechanism incorporates surprise (standardized prediction error) and staleness (time since last sampling) to adapt to the model's evolving capabilities.
Bottom-Up Belief Propagation (Empirical Bayes): After updates, information propagates upward. Parent node beliefs are estimated as precision-weighted averages of their children's beliefs, adjusted for subtree heterogeneity. This allows observations from one prompt to inform beliefs about related prompts in the same region, exploiting the latent structure.

C. Utility-Aware Extension (BMC-T)

To address the misalignment between training productivity and downstream evaluation performance, the authors introduce BMC-T. This variant biases sampling toward regions of the training distribution that overlap with a target distribution (e.g., a specific evaluation benchmark) in the latent space. It treats the Latent Task Tree as an approximate intervention model, assuming that training on nearby tasks has stronger transfer effects.

3. Key Contributions

Latent Task Trees: A hierarchical method for approximating the structure of the task manifold directly from model embeddings, providing a multi-scale representation of prompt relationships without external labels.
Bayesian Manifold Curriculum (BMC): A structure-aware adaptive curriculum framework that performs Bayesian decision-making over these trees, enabling efficient exploration while accounting for endogenous, model-induced non-stationarity.
Productivity–Diversity–Utility Trade-offs: A diagnostic analysis demonstrating that adaptive curricula cannot be characterized by downstream performance alone. The authors distinguish three axes:
- Productivity: Learning signal (reward variance).
- Diversity: Coverage of the task manifold.
- Utility: Relevance to specific evaluation benchmarks.
  They show these axes are not interchangeable; optimizing one does not guarantee optimization of the others.

4. Experimental Results

The authors evaluated BMC on mathematical reasoning tasks using the DAPO-Math-17K dataset and Qwen3 models (4B and 8B) with GSPO and GRPO.

Training Efficiency (Productivity): BMC and its "Difficulty Only" ablation achieved learning speeds comparable to Dynamic Sampling (which regenerates batches to ensure variance) but with significantly lower wall-clock time, remaining close to the efficiency of Uniform Sampling.
Coverage and Information Sharing (Diversity): BMC successfully balanced the diversity of Tree Only (which emphasizes coverage) with the productivity of Difficulty Only. Metrics showed that BMC provided better coverage of underrepresented regions (rarity-weighted exposure) and that the Latent Task Tree explained significant variation in learning signals (structure gain), confirming that learning signals are correlated with latent structure.
Evaluation Performance (Utility):
- Dynamic Sampling performed poorly on less-represented benchmarks (e.g., Chinese math, OOD science) due to implicit filtering of rare prompt types.
- BMC achieved strong performance on the out-of-domain GPQA-Diamond benchmark, suggesting that jointly optimizing productivity and diversity promotes capabilities not captured by standard in-distribution math evaluations.
- BMC-T demonstrated that utility can be steered: targeting specific benchmarks (e.g., AIME2024) improved performance on those specific benchmarks while preserving overall productivity, confirming that different problem types carry different evaluation relevance.

5. Significance and Claims

The paper argues that difficulty alone is insufficient for strong downstream performance. Maximizing learning signal (productivity) does not guarantee better outcomes if the model is not exposed to the types of problems that drive generalization to target tasks.

The authors claim that problem selection should be framed as a manifold-structured bandit problem over the policy's latent geometry. By deriving curriculum structure from the policy's own learned representations (adhering to the "Bitter Lesson" of using data and compute to learn structure rather than imposing human definitions), BMC offers a way to:

Navigate the trade-offs between productivity, diversity, and utility.
Exploit endogenous non-stationarity where sampling decisions shape future reward landscapes.
Generalize across domains, languages, and modalities without requiring hand-defined taxonomies or external embedding models.

The work highlights an "evaluation deadzone" phenomenon: a model may be trained productively on a specific set of problems, but if those problems do not align with the evaluation distribution, downstream performance may not improve. BMC-T provides a mechanism to address this by steering training effort toward target-relevant regions of the latent manifold.