BOTS: A Unified Framework for Bayesian Online Task Selection in LLM Reinforcement Finetuning

AI summary

Overview

• Research area: Reinforcement finetuning (RFT) of large language models, specifically dynamic/online task selection and curriculum design for RL training.
• Technical level: Advanced. The paper assumes familiarity with Bayesian posterior updating, Beta–Bernoulli conjugacy, multi-armed bandits, Thompson sampling, and GRPO-style RL training.
• Scope: The paper proposes BOTS, a Bayesian framework that fuses explicit and implicit evidence about task difficulty to select "just right" tasks online during LLM reinforcement finetuning, and validates it on the GURU dataset across math, code, and logic domains with Qwen2.5-1.5B-Instruct and Qwen2.5-7B.

What This Paper Is About

Reinforcement finetuning is highly sensitive to which tasks the model trains on: uniform sampling wastes computation on tasks the model has already mastered or cannot yet solve, and tasks of "just right" difficulty (success probability near 0.5) are the most informative. Existing online selection methods are limited — oversampling-based approaches pay extra rollout costs, while non-oversampling approaches rely on only one evidence source, either historical evaluations (explicit) or inter-task correlations (implicit), leaving information underused. BOTS recasts online task selection as Bayesian inference over the model's evolving capabilities and fuses both evidence types to select tasks more efficiently.

Key Contributions

1. A unified Bayesian framework for online task selection. BOTS is described as the first unified and extensible framework for Bayesian Online Task Selection in LLM reinforcement finetuning, recasting task selection as a Bayesian inference problem that handles non-stationarity (the model keeps changing) and partial observability (only evaluated tasks are directly observed).

2. A generalized update rule that fuses two complementary evidence sources. The framework jointly incorporates explicit evidence from direct evaluations of selected tasks and implicit evidence inferred for unselected tasks, controlled by tunable coefficients λ (prior mixing / discounting) and ρ (explicit vs. implicit balance). Proposition 1 shows this generalized-Bayes update preserves the Beta posterior family.

3. An ultra-light interpolation-based plug-in for implicit evidence. To make implicit evidence practical, the authors estimate unselected-task difficulty by linearly interpolating between two reference models of distinct capability (weak vs. strong), using a momentum-smoothed relative capability coefficient. It requires no extra rollouts and adds negligible overhead.

4. Thompson sampling for principled exploration–exploitation balance. Task selection draws samples from each task's Beta posterior and picks the tasks whose sampled success probability is closest to a target difficulty p* = 0.5.

Main Findings

• BOTS improves data efficiency and performance over baselines and ablations. Across diverse domains and LLM scales, the recommended configuration (λ = 0.1, ρ = 0.1) outperforms both out-of-framework baselines (Random, Offline) and within-framework baselines (BOTS-MoPPS, BOTS-DOTS), achieving 10 first-place and 6 second-place finishes out of 18 reported metrics in the Qwen2.5-1.5B-Instruct comparison across Math, Code, and Logic.

• Implicit evidence is critical for cold-start; explicit evidence is vital for long-term accuracy. With ρ = 0 (implicit evidence ignored), training with sparse explicit feedback behaves almost identically to random sampling. With ρ > 0, the Effective Task Ratio rises sharply early on, primarily by filtering out unsolvable (p = 0) tasks.

• Over-reliance on implicit evidence degrades long-term performance. Large ρ values (e.g., 0.5, 1.0) show a declining Effective Task Ratio as training progresses, driven by an increased tendency to select tasks whose success probability is 1, and they underperform on aggregated TTB and BSF.

• A small positive ρ works best. ρ ∈ {0.05, 0.1, 0.2} balances the two evidence sources, using implicit evidence for early acceleration while letting accumulating explicit evidence dominate later.

• Both too-small and too-large λ hurt. Small λ (e.g., 0.0, 0.05) fails to track capability improvement, producing a declining "< 1" ratio in mid-to-late training as stale difficulty estimates cause ineffective tasks to be selected. Large λ (e.g., 0.5, 1.0) aggressively discounts history, yielding small effective sample sizes and near-random, overly exploratory selection.

• Moderate λ achieves the best trade-off. λ ∈ {0.05, 0.1, 0.2} keeps difficulty estimates adaptive while controlling uncertainty, giving the best TTB and BSF results; λ = 0.1 is recommended as default.

• Theoretical characterization of posterior uncertainty. Proposition 2 shows that, under the update rules, the effective sample size satisfies lim inf n_t = n_0 + (ρ/λ)·n and lim sup n_t = n_0 + (1/λ)·n, linking λ and ρ directly to how confident the posterior is and therefore how exploratory Thompson sampling becomes.

• Negligible train-time overhead. BOTS introduces measured overhead of ≤ 0.2% of total training time, since no extra rollouts are required online and all updates reduce to simple vector operations. Because overhead is so small, the main results report TTB and BSF in training steps rather than wall-clock time.

• Reference-model tags are often already available. The authors note that datasets such as GURU supply Qwen2.5-7B-Instruct and Qwen3-30B-A3B scores as difficulty tags for offline filtering, which BOTS reuses for online selection. Where such tags do not exist, enabling implicit evidence costs a one-time rollout that does not recur during training.

Methodology in Plain English

Each task is modeled as a Bernoulli trial: for a given model and task, the reward is 1 (correct) or 0 (incorrect) with some unknown success probability. BOTS keeps a Beta distribution over that success probability for every task, and that Beta distribution is parameterized by accumulated success and failure counts.

After each training step, the counts for the selected tasks are updated with real evaluation results (explicit evidence). For tasks that were not selected, counts are updated with predicted results (implicit evidence) supplied by the plug-in. A single update equation mixes: (a) the previous counts, (b) a base prior, (c) the explicit counts, and (d) the pseudo counts, weighted by λ and ρ. Proposition 1 proves that this mixture is still exactly a Beta distribution, so nothing breaks mathematically.

The plug-in's prediction is deliberately cheap. Two reference models of different strength have known empirical success rates on each task. From the current batch, BOTS estimates where the training model sits between the weak and strong reference models on that batch (a coefficient), smooths that coefficient over time with momentum to reduce rollout noise, then linearly interpolates each unselected task's difficulty between its weak and strong reference rates, clipping the result to [0, 1].

Selection then follows Thompson sampling: draw one sample of the success probability from each task's Beta posterior, score each task by how close that sample is to the target p* = 0.5, and train on the top-scoring tasks. The randomness in the posterior draws naturally balances exploiting tasks believed to be near the target difficulty and exploring tasks whose difficulty is still uncertain.

The experiments use GRPO with Qwen2.5-1.5B-Instruct and Qwen2.5-7B, a learning rate of 1e-6, 16 rollouts per task, temperature 1.0, and the GURU math, code, and logic subsets (the Zebra Puzzle is excluded because its reward is non-binary). Results are averaged over 3 random seeds with 95% confidence intervals, and reported with Effective Task Ratio, Time-to-Baseline (at 50%, 75%, 100% target fractions), and Best-so-far (at 25%, 50%, 100% of total steps), evaluated on MATH500, AMC23, and AIME24 for math.

Why This Matters

• Impact on research: The paper provides a principled formulation — not a heuristic — for the observation that explicit and implicit evidence about task difficulty are complementary. By proving the update preserves the Beta family (Proposition 1) and characterizing posterior uncertainty in closed form (Proposition 2), it gives a reusable theoretical scaffold that other online selection methods can be plugged into, since the framework places no restrictions on the form of the implicit-evidence estimator.

• Real-world applications:

  • Designing training curricula for reasoning models in math, competitive-programming, and logic domains, where the GURU subsets are used here.
  • Reducing compute spend in large-scale RL post-training pipelines by avoiding rollouts on tasks that are already solved or currently unsolvable.
  • Reusing existing difficulty tags in public RL datasets (like GURU's reference-model scores) for online selection rather than only offline filtering.
  • Serving as a pluggable selection policy inside existing RFT codebases, given the released code in the Trinity-RFT repository.

• Industry relevance: The measured train-time overhead of ≤ 0.2% of total training time and the lack of extra online rollouts make the method attractive for production RL finetuning, where per-step rollout cost dominates. The one-time cost of reference-model evaluation amortizes when multiple models or checkpoints are trained on the same task pool, which is typical of industrial pipelines.

Future Directions

• Handling extrapolation from reference models. The authors report that when both reference models are stronger or weaker than the training model, implicit evidence becomes less accurate, though BOTS remains robust. Improving implicit evidence under extrapolation is left open.
• Extending the implicit-evidence plug-in. Because the framework places no restriction on the estimator p̃(k, B_t), richer estimators — beyond ultra-light interpolation — could be substituted and tested against this baseline.
• Broader domains and scales. The paper reports experiments on Qwen2.5-1.5B-Instruct and Qwen2.5-7B across math, code, and logic; the provided content does not report results for other model families, sizes, or domains beyond these.
• Untangling hyperparameter sensitivity. The default λ = 0.1 and ρ = 0.1 are reported to work across the evaluated scales and domains, but the paper notes practitioners may need to adjust them for new applications using the guidance in Section 3.5, implying a need for more automatic or adaptive tuning.

Target Audience

Researchers and engineers working on reinforcement learning from verifiable rewards and LLM post-training, particularly those focused on data curation, curriculum learning, or compute-efficient RL pipelines. It is also relevant to applied scientists implementing online task selection within existing RFT frameworks (such as GRPO-based ones), and to readers interested in Bayesian or bandit-based formulations of training-data selection. Beginners will find the algorithmic ideas accessible but will need background in Bayesian inference and RL to follow the derivations.

Authors’ abstract

Reinforcement finetuning (RFT) is a key technique for aligning Large Language Models (LLMs) with human preferences and enhancing reasoning, yet its effectiveness is highly sensitive to which tasks are explored during training. Uniform task sampling is inefficient, wasting computation on tasks that are either trivial or unsolvable, while existing task selection methods often suffer from high rollout costs, poor adaptivity, or incomplete evidence. We introduce BOTS, a unified framework for Bayesian Online Task Selection in LLM reinforcement finetuning. Grounded in Bayesian inference, BOTS adaptively maintains posterior estimates of task difficulty as the model evolves. It jointly incorporates explicit evidence from direct evaluations of selected tasks and implicit evidence inferred from these evaluations for unselected tasks, with Thompson sampling ensuring a principled balance between exploration and exploitation for task selection. To make implicit evidence practical, we instantiate it with an ultra-light interpolation-based plug-in that estimates difficulties of tasks without extra rollouts, adding negligible overhead. Empirically, across diverse domains and LLM scales, BOTS consistently improves data efficiency and performance over baselines and ablations, providing a practical and extensible solution for dynamic task selection in RFT. Code is available at https://github.com/agentscope-ai/Trinity-RFT/tree/main/examples/bots.