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

Abstract: Reinforcement learning (RL) is a central approach for improving reasoning capabilities in LLMs, where training efficiency depends critically on how problems are sampled during optimization. Existing adaptive curriculum learning methods typically prioritize prompts of intermediate difficulty, treating problem selection as a standard bandit problem with independent arms and overlooking the structured, heterogeneous nature of the task space. In this work, we frame problem sampling as a manifold-structured bandit problem with endogenous non-stationarity: problems are related through the model's latent representation space, and sampling decisions can steer how learning signals evolve across that space. To operationalize this perspective, we introduce Bayesian Manifold Curriculum (BMC), a structure-aware framework that organizes problems into a hierarchical task tree and applies Bayesian learning to guide sampling. Empirically, we find that different sampling strategies induce non-trivial tradeoffs between productivity (learning signal), diversity (coverage of the task manifold), and utility (evaluation relevance). These results show that prioritizing difficulty alone is insufficient for strong downstream performance, highlighting the importance of incorporating structure and type-awareness into problem sampling.

• The paper presents the Bayesian Manifold Curriculum (BMC) framework, which leverages latent policy geometry to guide adaptive curriculum learning for LLMs.
• It employs hierarchical Thompson sampling and latent task trees to balance productivity, diversity, and utility during prompt selection.
• Empirical evaluations demonstrate BMC's effectiveness in improving out-of-distribution performance and mitigating type imbalance in curriculum learning.

Bayesian Manifold Bandits and the Latent Geometry of LLM Curriculum

Introduction

This paper, "Manifold Bandits: Bayesian Curriculum Learning over the Latent Geometry of LLMs" (2606.19750), advances the methodology of adaptive curriculum learning in RL post-training for LLMs by proposing a structure- and type-aware scheduling method, Bayesian Manifold Curriculum (BMC). BMC leverages the geometry of latent policy representations to organize the problem space and guide prompt selection, moving beyond the prevailing paradigm of scalar, independent-arm bandit algorithms that prioritize prompts solely based on policy-dependent difficulty (e.g., reward variance). The authors demonstrate that neither maximizing learning signal nor brute-force coverage is sufficient for optimal downstream performance—instead, structured, utility-aware sampling that incorporates policy-induced type-awareness is necessary for high-performing and robust learning.

Problem Statement: Beyond Independent-Arms Bandits in Curriculum Learning

Standard RL-based curriculum learning for LLMs, particularly in group-relative policy optimization (GRPO/GSPO/DAPO), treats each prompt as an independent bandit arm. Prompts are sampled to maximize observed within-group reward variance, leading to optimization along a productivity axis (e.g., learning signal). However, the inherent heterogeneity of large mathematical and instruction-following data leads to dataset and frontier imbalance: policy-dependent learning signal becomes concentrated on specific types, which may be common or rare in data (see Figure 1).

Existing methods either overfit to majority problems, neglect minority or underrepresented types, or require heavy wall-clock compute to repeatedly regenerate effective training batches (e.g., dynamic sampling per DAPO). Furthermore, standard curriculum learning does not exploit latent correlations between prompts—crucially, in the context of LLMs, policy (model) representations cluster based on implicit task type and difficulty structure, as revealed by geometric analysis of latent spaces.

Latent Task Trees: Local Manifold Discovery in LLM Embeddings

The authors introduce Latent Task Trees as a hierarchical, model-centric partitioning of the prompt space (Figure 2). The construction method recursively clusters normalized intermediate-layer embeddings via PCA (retaining 95% local variance), UMAP (for neighborhood structure preservation), and HDBSCAN (for cluster assignment). Subdivision halts based on a local "Chart Test" (saturation of intrinsic dimension/connectivity) or cluster size minimums.

This methodology aligns the scheduler's operating structure to the actual latent task manifold as encoded by the policy network, as opposed to external annotation or naïve topic grouping. Importantly, this allows the automatic recovery of fine-grained, policy-aligned, and domain-agnostic prompt regions, supporting diverse LLM architectures (Figures 12–19). Comprehensive analyses show minimal compute overhead (<<2% total training time) and empirical robustness to modest latent drift (Figure 3).

Figure 2: Recursive construction of the Latent Task Tree from policy embeddings, yielding a multi-scale, policy-aligned partition of the task space.

Bayesian Manifold Curriculum (BMC): Hierarchical, Structure-Aware Bandit Learning

BMC utilizes the Latent Task Tree for guided curriculum learning. Each prompt maintains a non-stationary Bayesian belief (logit-normal parameterized) over its expected learning signal (prompt-level reward variance), adapting beliefs as rollouts are updated and incorporating policy drift via staleness and surprise-based forgetting, all within a lightweight empirical Bayes framework.

Top-down, hierarchical Thompson sampling is employed: for each batch element, the sampler descends the tree, recursively selecting the most promising sub-region according to uncertain reward beliefs—thus, batch elements are dispersed across productive, yet diverse, regions (Figure 4).

Figure 4: BMC's batched hierarchical Thompson sampling and bottom-up empirical Bayes belief propagation across the task tree.

After batch rollouts, prompt-level posteriors are updated, and empirical Bayes belief propagation aggregates this information up the tree, incorporating a random-effects term that accounts for intra-subtree heterogeneity.

This approach enables coordinated, structure-aware credit assignment, and crucially, information sharing across related prompts, mitigating under-sampling of minority types and supporting adaptive, multi-scale curriculum learning.

Empirical Evaluation: Productivity, Diversity, and Utility

The authors design a thorough diagnostic framework to analyze curriculum learning strategies along three axes:

• Productivity: The effective ratio (fraction of sampled prompts with non-zero learning signal) and learning speed (training-set accuracy over time) quantify the sample efficiency.
• Diversity: Rarity-weighted exposure and structure gain (variance explained in learning signal by tree clusters over random partitions) capture coverage of underrepresented regions and evidence that learning signal naturally aligns with latent structure.
• Utility: Downstream evaluation performance on both in-distribution and OOD benchmarks assesses alignment between areas of high learning signal and actual target utility.

Experimental results (Figures 4–7, 11) show:

• Dynamic Sampling achieves maximal effective ratio by construction but is compute-prohibitive and suffers from minority prompt skipping (Figure 5, 23).
• Thompson-sampling baselines (Difficulty Only) increase learning signal but do not guarantee type diversity.
• Simple diversity-first policies (Tree Only) underperform due to lack of productivity focus.
• BMC interpolates between productivity and diversity, achieving strong OOD transfer (e.g., on GPQA-Diamond; Figure 6), and improves rarity exposure and structure gain (Figure 7).
• Utility-aware extensions (BMC-T) steer evaluation outcomes towards specific target benchmarks/distributions, decoupling improvements along the productivity/diversity axes from true utility (Figure 8).

  Figure 9: Training efficiency for several sampling strategies, showing trade-offs between productivity and wall-clock cost.

  

  Figure 7: BMC achieves balanced improvement in both rarity-weighted exposure and structure gain, indicating coverage of the latent task manifold and alignment with policy-induced structure.

  

  Figure 6: Downstream benchmark results; BMC demonstrates strong generalization to OOD benchmarks, reflecting the importance of structured coverage.

  

Figure 8: Utility-aware extensions alter downstream evaluation profiles despite similar productivity, empirically validating utility as an independent axis.

Extensions, Analysis, and Limitations

BMC outperforms non-structure-aware bandit methods in settings with high conceptual heterogeneity and type imbalance (see DAPO-Math-17K, AlphaMed19K, GURU-92k, and BarExamQA), confirming the necessity of aligning curriculum with latent policy geometry in RLVR for LLMs. The authors examine drift in latent geometry and demonstrate the stability of fixed trees for the training horizons studied.

Medical-domain experiments (Figure 10) emphasize domain-general applicability, though also expose the fact that structure-aware diversity may not always lead to the best utility unless aligned with evaluation. Discussion of staleness/uncertainty growth, alternative prioritization objectives (e.g., frontier expansion), and connections to structured, causal, and multi-task bandit literature are provided.

Implications and Future Directions

This work problematizes the notion that productivity (difficulty) maximization suffices for curriculum design in LLM RL. Empirical evidence shows type-mismatch (train-test) can occur even with high learning signal; diversity and utility require explicit modeling. The latent geometry-based approach decouples problem selection from either heuristic or externally sourced taxonomies and enables scalable, domain-agnostic, and policy-aligned curriculum adaptation.

Points for further development include revisiting curriculum as a causal intervention problem (to model how sampled prompts affect neighboring regions in the latent manifold), combining with gradient-alignment criteria for explicit utility steering, integrating more granular rollout allocation, and improving computational aspects of periodic recursive tree reconstruction.

Conclusion

The Bayesian Manifold Curriculum framework substantially generalizes independent-arm bandit scheduling for RLVR-based LLM post-training by leveraging the policy’s own representation geometry to construct a dynamic, multi-scale structure over the prompt space. This enables sampling strategies that balance signal, coverage, and relevance, elevating curriculum learning from scalar difficulty to structured, policy-driven orchestration over the LLM's latent task manifold. The framework is broadly extensible and provides a foundation for further advances in efficient, effective RL-based LLM training.

---