Bayesian Manifold Curriculum

• Bayesian Manifold Curriculum is a structure-aware adaptive curriculum that uses latent task geometry to guide reinforcement learning in LLMs.
• It constructs a Latent Task Tree using unsupervised clustering and hierarchical Thompson sampling to balance productivity, diversity, and utility.
• The framework models expected learning signals with Bayesian beliefs and online updates to efficiently direct training and overcome non-stationarity.

Searching arXiv for the named papers and related work on Bayesian Manifold Curriculum. arXiv search query: "Bayesian Manifold Curriculum latent task tree manifold bandits" Bayesian Manifold Curriculum (BMC) is a structure-aware adaptive curriculum for reinforcement learning training of LLMs that treats problem sampling as a manifold-structured bandit problem with endogenous non-stationarity. In this formulation, problems are related through the model’s latent representation space, sampling decisions can steer how learning signals evolve across that space, and curriculum design is therefore not reducible to difficulty-only selection over independent prompts. BMC organizes problems into a Latent Task Tree built from the policy’s own hidden representations, models expected learning signal with Bayesian beliefs, selects prompts by hierarchical Thompson sampling, and propagates information bottom-up through empirical Bayes aggregation. Its stated objective is to balance productivity, diversity, and utility, and it has also motivated comparison with earlier Bayesian uncertainty-driven curricula in other domains, including medical image segmentation, where uncertainty is fed back into sampling without an explicit Latent Task Tree (McKenzie et al., 18 Jun 2026, Iwamoto et al., 2021).

1. Problem formulation and conceptual reframing

BMC is introduced for problem sampling and curriculum learning in RL training of LLMs, especially in group-relative policy optimization methods such as GRPO, DAPO, and GSPO. In RL with verifiable rewards, a prompt $x$ is sampled, the policy $\pi_\theta$ generates multiple responses, and a verifier assigns rewards. In group-relative methods, the learning signal comes from within-prompt reward variation across rollouts. If all rollouts for a prompt receive the same reward, the advantage collapses and the prompt contributes little or no gradient signal. The curriculum problem is therefore to choose prompts that are productive for learning, avoid wasting compute on prompts that are all-success or all-failure, preserve coverage over the problem space, and ideally remain relevant to downstream evaluation utility (McKenzie et al., 18 Jun 2026).

The prompt-level learning signal is defined from reward variability across multiple rollouts:

$$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$

where $i$ indexes prompts, $k$ is the number of rollouts per prompt, $r_{ij}$ is the reward of rollout $j$ for prompt $i$, and $\phi:\mathbb{R}\to[0,1]$ is an optional normalization function. For binary rewards, variance is highest near a 50% success rate. BMC therefore prioritizes prompts near the model’s current capability frontier: not too easy, not too impossible, but informative for credit assignment. The paper emphasizes that this is a proxy for productivity, not the policy’s optimization objective itself (McKenzie et al., 18 Jun 2026).

The central conceptual claim of BMC is that prompt sampling for LLM RL is not an independent-arm bandit problem; it is a manifold-structured bandit problem over the latent task geometry of the policy. Two structural properties follow from this reframing. First, there is structured correlation: nearby prompts in latent space may share difficulty, solution patterns, or learning dynamics. Second, there is endogenous non-stationarity: the reward distribution over prompts changes because the model itself changes during training. In this setting, curriculum selection is not merely a matter of identifying static high-value arms. It is also a matter of choosing interventions that alter future learning dynamics across related regions of the task space (McKenzie et al., 18 Jun 2026).

A common misconception addressed by this formulation is that adaptive curricula need only estimate difficulty accurately. BMC argues instead that difficulty alone is insufficient, because prompt selection in RLVR interacts with structure, coverage, and evaluation relevance. This is why the framework explicitly separates productivity, diversity, and utility rather than assuming they coincide (McKenzie et al., 18 Jun 2026).

2. Latent Task Trees and the approximation of task geometry

The hierarchical substrate of BMC is the Latent Task Tree (LTT), which approximates the latent task manifold with a recursive tree structure. The motivation is that task structure is not flat: examples such as math $\rightarrow$ geometry $\pi_\theta$0 Euclidean geometry, coding $\pi_\theta$1 algorithms $\pi_\theta$2 dynamic programming, and science $\pi_\theta$3 physics $\pi_\theta$4 mechanics illustrate why a single flat partition can miss informative substructure. Rather than assigning prompts to one global cluster, BMC recursively discovers finer organization from the policy model’s own hidden representations (McKenzie et al., 18 Jun 2026).

For each prompt, the tree construction begins with a normalized mean embedding from a selected intermediate layer of the policy model. At each node, embeddings are standardized, PCA is run, the minimum number of components explaining 95% variance is retained, the reduced representation is projected with UMAP, and HDBSCAN is used for clustering. Clusters are recursively subdivided if meaningful substructure remains. Recursion stops when the region is deemed chart-like, when HDBSCAN finds at most one cluster, or when clusters are too small (McKenzie et al., 18 Jun 2026).

The chart test is intended to determine whether a region locally behaves like a single smooth patch. It checks intrinsic dimension via Two-NN and local connectivity via a $\pi_\theta$5-nearest-neighbor graph. If the region is sufficiently low-dimensional and connected, it is treated as locally simple and is not further subdivided. The resulting hierarchy is described as multi-scale, local rather than globally forced, unsupervised, model-specific, and cheap relative to RL training (McKenzie et al., 18 Jun 2026).

This architecture matters because it converts latent geometry into a sampling structure. Observations collected at leaves need not remain isolated at the prompt level; they can inform ancestor regions and, indirectly, sibling branches. A plausible implication is that BMC uses the tree not merely as a clustering device, but as an inductive bias about how productivity varies across task space.

3. Bayesian belief modeling and hierarchical Thompson sampling

Once the Latent Task Tree is built, BMC performs sampling with Bayesian beliefs over expected learning signal. The prompt-level belief model is presented as

$\pi_\theta$6

where $\pi_\theta$7 is the current estimate of expected learning signal and $\pi_\theta$8 is uncertainty. Because signals are bounded in $\pi_\theta$9, the implementation uses a Logit-Normal parameterization: the bounded signal is transformed with logit, Gaussian beliefs are maintained in logit space, and the result is mapped back via sigmoid. The stated reason is to avoid assigning probability mass to impossible values outside $$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$0 (McKenzie et al., 18 Jun 2026).

Batch construction proceeds by top-down selection through the tree. At a node with children $$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$1, BMC samples

$$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$2

and chooses

$$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$3

This is hierarchical Thompson sampling. The paper notes that sampling is done without replacement in batch-aware form. If one region seems promising, multiple batch elements can descend there; if several regions appear similarly promising, the batch spreads out. The intended balance is among exploitation of productive regions, exploration of alternative regions, and diversity within the batch (McKenzie et al., 18 Jun 2026).

Because productivity is non-stationary, BMC updates beliefs online with a lightweight Bayesian filter inspired by the Normal-Inverse-Gamma model. The maintained quantities are the mean $$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$4, variance $$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$5, and effective prior weight $$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$6. For an observation $$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$7 at time $$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$8, surprise is defined as

$$y_i = \phi\big(\mathrm{Var}([r_{i1}, r_{i2}, \dots, r_{ik}])\big),$$9

The effective prior weight is updated by

$i$0

the posterior mean by

$i$1

and the variance by

$i$2

Large surprise reduces reliance on prior evidence, while staleness increases uncertainty when a prompt has not been sampled for some time. This is the mechanism by which BMC handles endogenous non-stationarity: old beliefs can be forgotten, surprising evidence can rapidly reshape estimates, and unsampled prompts can become uncertain again (McKenzie et al., 18 Jun 2026).

After leaf-level updates, information is propagated upward through empirical Bayes aggregation. For each internal node, BMC introduces a heterogeneity term

$i$3

assigns weights

$i$4

and computes the parent belief as

$i$5

If children agree, aggregation is tight; if children disagree, heterogeneity inflates uncertainty. This is intended to avoid overconfident averaging in heterogeneous regions. Algorithmically, BMC repeats three steps: top-down selection via hierarchical Thompson sampling, update of prompt-level beliefs from rollout rewards, and bottom-up propagation through the tree (McKenzie et al., 18 Jun 2026).

4. Productivity, diversity, and utility as distinct axes

A defining empirical claim of BMC is that productivity, diversity, and utility are distinct axes and are not automatically aligned. Productivity is measured by reward variance, effective ratio, and training-set pass@1 over time. Diversity is measured by rarity-weighted exposure and structure gain. Utility is assessed on held-out benchmarks such as AIME, AMC, MATH500, CNMO, CCEE, and GPQA-Diamond (McKenzie et al., 18 Jun 2026).

Axis	Measurement	Reported pattern
Productivity	reward variance, effective ratio, training-set pass@1	Dynamic Sampling achieves the highest effective ratio and learning signal
Diversity	rarity-weighted exposure, structure gain	Tree Only emphasizes diversity; BMC interpolates between diversity and productivity
Utility	held-out benchmark performance	No single sampler dominates all benchmarks

For productivity, Dynamic Sampling achieves the highest effective ratio and learning signal, but is expensive in wall-clock time due to repeated regeneration. BMC and Difficulty Only recover much of the training-speed benefit of dynamic sampling while staying much closer to uniform sampling in wall-clock cost. Uniform sampling has weaker learning signal and slower learning progress. The paper’s interpretation is that higher learning signal generally correlates with faster training-set improvement, but maximizing it alone is costly or incomplete (McKenzie et al., 18 Jun 2026).

For diversity, Tree Only emphasizes diversity but lacks productivity, so it does not perform best. Difficulty Only emphasizes productivity but is less diverse. BMC interpolates between them: it increases exposure to rarer regions while still prioritizing productive prompts. The tree also explains more variance in learning signal than random partitions of the same size, which is offered as evidence that learning signal aligns with latent task structure (McKenzie et al., 18 Jun 2026).

For utility, no single sampler dominates all benchmarks. BMC performs especially strongly on GPQA-Diamond, suggesting that broad, structured coverage can help out-of-distribution generalization. Dynamic Sampling can hurt less-represented benchmarks because of problem skipping and imbalance effects. Difficulty Only improves some benchmarks but still tends to prioritize majority, common, or immediately productive types. Tree Only is not best anywhere because it ignores learning signal. The main empirical conclusion is that difficulty alone is insufficient, productivity does not guarantee utility, and diversity helps but not if it ignores learning signal (McKenzie et al., 18 Jun 2026).

5. BMC-T, target awareness, and structural stability

BMC-T is the utility-aware extension of BMC. Its motivation is the observation that productivity and diversity do not by themselves determine downstream performance. Given a training set $i$6 and a target distribution $i$7, the method builds a tree over both training and target examples and computes target overlap within subtrees. For each child node $i$8, a utility score is defined from target mass overlap, train-support gating, and smoothing via a Dirichlet-like prior. The hierarchical Thompson score is then modified as

$i$9

where $k$0 controls utility strength, $k$1 puts the utility term on the same scale as productivity, and $k$2 is normalized child utility (McKenzie et al., 18 Jun 2026).

Empirically, BMC-T preserves productivity roughly as well as BMC, but shifts downstream performance depending on the target distribution used. Targeting AIME-style benchmarks improves English competition math performance, whereas targeting the full evaluation mixture improves broader benchmark coverage and Chinese math performance. This is presented as evidence that utility is a separate axis, not reducible to productivity (McKenzie et al., 18 Jun 2026).

The paper also discusses structure drift. The latent task tree is fixed during main training, but the authors check whether tree structure changes after RL. The reported findings are that tree topology can change somewhat with training, sometimes becoming wider, smaller, or more hierarchical, while periodic reconstruction changes results only modestly. The stated conclusion is that the static tree is a reasonable approximation over the training horizon studied (McKenzie et al., 18 Jun 2026).

Two additional misconceptions are addressed here. The first is that a high-learning-signal sampler should automatically dominate downstream evaluation; the reported benchmark tradeoffs contradict this. The second is that a static latent structure must be exact to be useful; the modest effect of periodic reconstruction suggests that exact real-time remanifoldization is not required in the reported regime.

6. Relation to Bayesian uncertainty-driven curricula outside LLM RL

Although BMC is introduced as a framework for RL-trained LLMs, the broader idea of Bayesian uncertainty guiding curriculum construction appears in earlier work on semantic segmentation of medical images. In the paper “Improving the Reliability of Semantic Segmentation of Medical Images by Uncertainty Modeling with Bayesian Deep Networks and Curriculum Learning,” a Bayesian U-Net implemented with MC Dropout estimates predictive entropy, converts that entropy into an Uncertainty Map, computes sub-image uncertainty as the average uncertainty over pixels in a sub-image, and resamples training patches more densely in uncertain regions through the rule

$k$3

with the stated interpretation that large $k$4 yields denser sampling in high-uncertainty regions. The workflow is iterative: train a model, infer uncertainty, alter the training sampling distribution, retrain on the new resampled dataset, and repeat until no further reliability improvement occurs (Iwamoto et al., 2021).

That segmentation system is applied to iPS cell colony segmentation from phase-contrast microscopy images with three classes—Good, Bad, and BGD—and evaluates reliability by Negative Predictive Value, True Positive Rate, and Uncertainty Accuracy, alongside mean IoU. On a dataset of 59 images of size $k$5, with Bayesian U-Net, 50% dropout in encoder and decoder, $k$6 MC samples, sub-image size $k$7, step $k$8, uncertainty threshold $k$9, and Gaussian width $r_{ij}$0, the baseline single-stage CE model gives NPV: $r_{ij}$1, TPR: $r_{ij}$2, UA: $r_{ij}$3, and IoU: $r_{ij}$4. The best curriculum result gives NPV: $r_{ij}$5, TPR: $r_{ij}$6, UA: $r_{ij}$7, and IoU: $r_{ij}$8, with reported gains of TPR improved by 12%, NPV improved by 9%, UA improved by 6%, and IoU improved by about 1% (Iwamoto et al., 2021).

This earlier method is not explicitly presented as BMC. The differences are explicit: the segmentation paper is in semantic segmentation with spatial resampling of image patches; it uses MC Dropout entropy and a Gaussian step rule for patch resampling; and its curriculum is local and pixel/patch-based. By contrast, BMC in the LLM setting uses a Latent Task Tree, Bayesian belief modeling, hierarchical Thompson sampling, empirical Bayes aggregation, and explicit analysis of productivity, diversity, and utility. Even so, the medical imaging paper is described as strongly aligned with the BMC philosophy because Bayesian uncertainty guides where training should focus next. This suggests a broader family resemblance among Bayesian curriculum methods in which uncertainty is used to reshape sampling over a data manifold, whether that manifold is represented as image patches or as a latent hierarchy of prompts (Iwamoto et al., 2021, McKenzie et al., 18 Jun 2026).