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

2026-06-17

AI Summary

Traditional curriculum learning in reinforcement learning with large language models relies solely on task difficulty for sampling, overlooking the structured and heterogeneous nature of the task space, thereby limiting training efficiency. This work formulates problem sampling as a manifold-structured bandit problem with endogenous non-stationarity and introduces a structure-aware sampling framework that uniquely integrates task manifold geometry with Bayesian curriculum learning. The approach constructs a hierarchical task tree from latent representations of the large language model and employs Bayesian inference to guide a sampling strategy that jointly optimizes productivity, diversity, and utility. Experimental results demonstrate that the proposed method achieves a superior trade-off among learning signal quality, task coverage breadth, and evaluation relevance, leading to significant improvements in downstream task performance.

Abstract

Reinforcement learning (RL) is a central approach for improving reasoning capabilities in large language models (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.

Problem

• manifold bandits
• curriculum learning
• large language models
• problem sampling
• latent geometry

Innovation

• Manifold Bandits
• Bayesian Curriculum Learning
• Latent Geometry
• Task Structure
• Non-stationary Sampling