PaperlayerView exact revision on arXiv

Source-linked AI summary

AC-ODM: Actor--Critic Online Data Mixing for Sample-Efficient LLM Pretraining

Jing Ma, Chenhao Dang, Mingjie Liao

arXiv:2505.23878v2cs.LGcs.AI

TL;DR

LLM pretraining needs data-mixing strategies that adapt to evolving training dynamics while balancing sample efficiency and practical flexibility. AC-ODM uses reinforcement learning to dynamically weight domains by gradient alignment, outperforming prior methods in convergence and downstream accuracy while reaching optimal validation perplexity with 66% fewer training steps than baselines.

  • Problem

    Efficient LLM pretraining requires data-mixing methods that account for domain distribution because it affects sample efficiency, convergence speed, and downstream accuracy.

  • Method

    AC-ODM uses an actor–critic policy to dynamically optimize domain weights from model state, rewarding gradient alignment to maximize constructive interference during pretraining.

  • Results

    Across evaluations, AC-ODM outperforms prior methods in convergence and downstream accuracy, reaching optimal validation perplexity with 66% fewer training steps than baselines.

  • Takeaways & Limitations

    AC-ODM offers a flexible reinforcement-learning framework for improving pretraining efficiency through cross-domain gradient alignment.

  • Takeaways & Limitations

    AC-ODM assumes meaningful domain partitions, and proxy mode relies on transferring a policy from a smaller model to the target model across studied settings.

Abstract

from arXiv · show

Optimizing pretraining data composition is pivotal for LLM generalization. While dynamic mixing outperforms static strategies by capturing evolving training dynamics, current methods fail to reconcile computational efficiency with sample efficiency and structural flexibility for diverse pipelines.We introduce Actor--Critic Online Data Mixing (AC-ODM), which approaches data mixing from a reinforcement learning perspective with a parameterized policy that we theoretically prove to act as a dynamic linear surrogate maximizing the constructive interference of gradients. To enhance practical flexibility, AC-ODM supports two operational modes: (i) a proxy mode for fixed, pre-prepared corpora, where a policy learned on a small model is transferred to a larger target; and (ii) a non-proxy mode for direct end-to-end training from scratch without priors. Empirically, AC-ODM significantly outperforms prior methods in convergence speed and downstream accuracy across various architectures. On Pythia-1B, it reaches optimal validation perplexity using up to 66% fewer training steps than competitive baselines, delivering a 27.5% relative improvement in MMLU accuracy and a 2.23 x higher pass@1 on HumanEval, all while incurring a virtually negligible (0.4%) per-step wall-clock increase and only 2% additional memory overhead. Code is available at https://github.com/DANG-ai/AC-ODM.

1. Introduction

AC-ODM frames online data mixing as reinforcement learning, using a parameterized policy to adapt domain weights during pretraining. It targets the limitations of existing methods by combining dynamic responsiveness with computational efficiency, sample efficiency, and structural flexibility.

  • Motivation: Pretraining-corpus quality, domain coverage, and mixture strongly influence LLM sample efficiency, convergence speed, and downstream accuracy.These factors make data-mixing optimization a critical frontier in efficient LLM pretraining.
  • Prior Work: Dynamic data-mixing strategies generally outperform static baselines because they respond to the model’s evolving capabilities and deficits.Prior static approaches determine domain weights offline, whereas dynamic methods adjust distributions during training.
  • Challenge: Existing dynamic methods often lack a unified balance among computational efficiency, sample efficiency, and structural flexibility across diverse pretraining pipelines.Sophisticated selectors can impose prohibitive runtime overhead, while lightweight heuristics may not support direct pretraining from scratch and pipelines without priors.
  • AC-ODM: AC-ODM treats LLM pretraining as an RL environment in which an Actor dynamically optimizes domain weights from the model’s real-time state.The framework approaches data mixing from an RL perspective rather than relying solely on previous heuristics.
  • AC-ODM: AC-ODM uses an Actor-Critic network and a gradient alignment reward to capture intra-domain interactions and steer pretraining toward faster convergence.Its theoretical analysis presents the reward as a first-order proxy for spectral coherence that maximizes effective gradient-update magnitude.

2. Related Work

Related work frames pretraining data composition as central to LLM generalization and sample efficiency, motivating increasingly sophisticated mixing strategies. Existing approaches span offline static methods and on-the-fly dynamic methods, but dynamic techniques often trade computational overhead against structural flexibility.

  • Data Mixing in LLM Pretraining: Pretraining data composition strongly influences LLM generalization and sample efficiency, often more than pure data volume.Recent technical reports emphasize sophisticated mixing as essential for competitive foundation-model training.
  • Static Data Mixing Strategies: Static methods determine mixture weights offline using proxy-model objectives, scaling laws, regression, clustering, or multidimensional quality assessments.These approaches include minimizing loss gaps, aligning gradients, predicting optimal mixtures, and evaluating corpus quality.
  • Dynamic Data Mixing Strategies: Dynamic methods adjust mixtures on-the-fly to capture feature-learning evolution, extending bandit-based ODM with gradient interactions, bi-level optimization, quality-diversity balancing, or Bayesian optimization.The cited dynamic methods build on ODM while introducing different optimization signals and search strategies.
  • Dynamic Data Mixing Strategies: Dynamic data-mixing methods often face a trade-off between computational overhead and structural flexibility, without a unified framework.This limitation motivates methods that reconcile efficiency with flexible training pipelines.

3. AC-ODM

AC-ODM formulates online data mixing as continuous-control reinforcement learning that adaptively updates domain weights during pretraining. Its alignment-based reward is theoretically a linear surrogate for constructive gradient interference, while proxy and non-proxy modes support different training constraints.

  • Problem formulation: AC-ODM updates the domain-sampling distribution at every iteration to adapt mixing to the model’s evolving training state.Unlike offline mixing, it changes Pαt during training while targeting generalization with negligible computational overhead.
  • RL methodology: The LLM serves as the environment, while a DDPG agent observes training dynamics and outputs actions that update domain weights.States include iteration, per-domain sample counts and losses, loss changes, selected-layer weight norms, and their update magnitudes; actions are mapped to the simplex by softmax.
  • Reward mechanism: AC-ODM rewards domains whose gradients align with the aggregate gradient of the remaining corpus, using an importance-corrected exponential moving average for stability.The importance correction prevents policy collapse toward already-frequent domains.
  • Theoretical analysis: The alignment reward acts as a linear surrogate for cross-term energy, greedily increasing effective gradient magnitude through constructive interference.The policy therefore favors domains whose gradients reinforce one another, rather than directly optimizing the computationally expensive quadratic form.
  • Operational modes: AC-ODM provides proxy and non-proxy operational modes to address different constraints in large-scale training.In proxy mode, the actor and critic are trained with a proxy LLM, the actor is transferred and frozen, and the target LLM is trained without further reward computation or policy updates.

4. Experiment

Across Pythia and LLaMA-style experiments, AC-ODM improves convergence, domain-level generalization, downstream accuracy, and training efficiency over static and dynamic baselines. Its proxy mode is strongest, while non-proxy training retains substantial gains with negligible overhead.

  • Convergence Analysis: AC-ODM-410M reaches the strongest baseline’s optimal validation perplexity with approximately 66% fewer steps on The Pile.On SlimPajama, it requires 65% fewer steps than ODM and 73% fewer than the uniform baseline.
  • Domain-Level Generalization: AC-ODM-410M achieves the best average test perplexity on The Pile and outperforms PiKE in 17 of 22 domains.The gains are most pronounced in small and medium domains while balancing performance across dominant and specialized domains.
  • Downstream Evaluation: +27.5% improvement over ODM on zero-shot MMLU accompanies a 2.23× higher HumanEval pass@1.Against PiKE, AC-ODM-410M gains +5.1% in zero-shot MMLU and +39% relative improvement in HumanEval pass@1.
  • Architecture and Partition Sensitivity: AC-ODM improves training dynamics on a LLaMA-style 0.9B Transformer, with proxy mode reaching target perplexity in fewer steps and achieving lower fixed-budget perplexity.Coarser domain partitions weaken validation perplexity, indicating that meaningful cross-domain structure benefits AC-ODM.
  • Per-step Efficiency: 0.4% per-step overhead distinguishes direct AC-ODM, requiring 2.48 s versus 2.47 s for ODM.PiKE requires 2.53 s per step because of gradient-conflict estimation.
  • End-to-End Speedup: 31.95% fewer total training steps produce a 1.46× end-to-end speedup, surpassing PiKE’s 1.30× speedup.Proxy mode requires only 28.82% of ODM steps for the 1B target and achieves 2.08× overall speedup with a 160M proxy.

5. Limitations

AC-ODM is best suited to corpora with meaningful domain structure, while coarse or overlapping groupings can weaken its reward signal. Its proxy mode depends on transferring a policy across models, and its lightweight reward estimation leaves broader shifts for future study.

  • AC-ODM’s gains are strongest when corpus partitions expose useful cross-domain gradient structure.The method is designed for corpora that can be organized into meaningful domains.
  • Coarse or highly overlapping groupings remain usable but may provide a less discriminative reward signal.This limitation is suggested by the domain-granularity study.
  • The proxy mode relies on transferring a learned policy from a smaller model to the target model.Experiments support this transfer across the studied settings, but broader architecture and corpus shifts remain for future examination.
  • To remain lightweight, AC-ODM estimates rewards from selected parameters and optimizes the mixture of available data.

6. Conclusion

AC-ODM establishes a rigorous reinforcement-learning framework that reformulates pretraining data selection and grounds its reward in optimization geometry. It maximizes update spectral coherence so data mixtures constructively interfere to accelerate convergence.

  • AC-ODM reformulates pretraining data selection from a heuristic process into a principled reinforcement learning problem.
  • Its reward signal is theoretically grounded in optimization geometry and explicitly maximizes the spectral coherence of updates.
  • By promoting constructive interference among data mixtures, AC-ODM accelerates convergence rather than merely reducing conflicts.

Impact Statement … A.2. AC Networks Configuration

AC-ODM is presented as a reinforcement-learning method intended to improve pretraining sample efficiency and reduce associated resource use. The configuration uses a 16-layer Transformer language model and six-layer actor and critic networks with distinct output activations.

  • Impact Statement: AC-ODM targets improved sample efficiency by achieving comparable or better model quality with fewer pretraining steps.The anticipated benefits include reduced compute, energy use, and carbon emissions, alongside a lower barrier to developing capable language models.
  • Impact Statement: More efficient pretraining may reduce compute, energy consumption, and carbon emissions associated with large-scale model training.The impact statement also notes that lowering development barriers could amplify both benefits and risks of capable language models.
  • A. Model Configuration: The language model uses sequence length 1024 and a 16-layer Transformer with hidden size 2048 and 16 attention heads.Rotary positional embeddings are incorporated.
  • A.1. LLM Model Configuration: FlashAttention is used to optimize memory access and reduce computation overhead during training.The model is trained with Adam, and its learning rate uses a linear warm-up for 833 iterations.
  • A.2. AC Networks Configuration: Both actor and critic networks are fully connected six-layer neural networks with 1024 neurons per hidden layer.Every non-output layer uses layer normalization followed by ReLU activation.
  • A.2. AC Networks Configuration: The actor output uses softmax, whereas the critic output uses identity activation.This distinguishes the actor’s output processing from the critic’s output processing.

B. Ablation study of selected layers

The ablation study finds that later Transformer blocks provide the strongest proxy for reward computation, with AC-ODM remaining robust to the precise layer subset. Layers 12,14,16 achieve the lowest perplexity and slightly outperform contiguous later layers 14,15,16.

  • Layer selection: Layers 12,14,16 attain the lowest perplexity among the evaluated layer subsets.This identifies the best-performing proxy for reward computation.
  • Layer selection: Later Transformer blocks slightly outperform contiguous later layers 14,15,16.The result matches or exceeds mid and early layer choices.
  • Layer selection: AC-ODM is robust to the exact layer subset, despite small but consistent absolute differences.The findings suggest that mid-to-late representations provide a more informative signal.

C. Zero shot accuracy on downstream tasks · D. subEffect of proxy model size

AC-ODM-410M delivers the strongest zero-shot downstream accuracy, while proxy-learned policies show that larger and prelearned actors improve training outcomes over joint online learning. Proxy size matters: 160M and 410M policies outperform joint AC-ODM, whereas 70M is insufficient.

  • C. Zero shot accuracy on downstream tasks: 0.62528 is AC-ODM-410M’s best average accuracy, exceeding ODM by +0.0341 absolute and +5.8% relative.AC-ODM-410M achieves the highest accuracy on every task.
  • C. Zero shot accuracy on downstream tasks: AC-ODM-410M’s gains are consistent across commonsense and reasoning benchmarks, with the largest improvement on MMLU.
  • C. Zero shot accuracy on downstream tasks: The non proxy AC-ODM improves over ODM on average but trails proxy mode, supporting policy learning with a proxy model before target guidance.
  • C. Zero shot accuracy on downstream tasks: Removing any AC-ODM state component degrades performance.The ablation defines “Impr.” as the relative change in perplexity compared with using all components.
  • D. subEffect of proxy model size: For a 1B target, policies were learned from 70M, 160M, and 410M proxy LLMs and compared with joint AC-ODM training.
  • D. subEffect of proxy model size: 2.8, 2.65, and 2.48 are the training losses for the 70M, 160M, and 410M proxies, respectively; their validation perplexities are 20.3, 15.5, and 12.1.
  • D. subEffect of proxy model size: Policies from 160M and 410M consistently outperform joint AC-ODM because a prelearned actor adapts from the first step while an online actor converges.
  • D. subEffect of proxy model size: The 70M proxy performs worst, suggesting insufficient capacity.

E. Ablation study of state components

The ablation shows that all six state features improve policy quality, with per-domain losses and weight norm being most critical. The full state achieves the best perplexity because its components provide complementary information.

  • Feature importance: ≈6.4% degradation follows removing per-domain losses ℓ(θM, B) or weight norm ∥ω∥2, the largest ablation drop.These features indicate that absolute training signal and model-scale dynamics are critical for the actor.
  • Feature importance: −3.69% degradation follows removing the change-of-loss term ∆ℓ(θM, B), confirming its importance.The change-of-loss signal contributes alongside the absolute loss and weight norm.
  • Feature importance: The count of seen samples n and step index t provide smaller but nontrivial gains, while the full state offers the best perplexity.All six features contribute to policy quality, with complementary effects across the state representation.

F. Evolution of Domain Weights During Training

AC-ODM dynamically adjusts domain weights during warmup, stabilizing after approximately 15,000 steps while retaining minor fluctuations that track the evolving LLM state. Compared with ODM, it makes larger early adjustments and favors several high-quality, general-purpose domains across token proportions.

  • Weight evolution: AC-ODM dynamically updates The Pile’s original domain weights and stabilizes after approximately 15,000 training steps, with minor fluctuations afterward.These fluctuations correspond to the evolving state of the LLM.
  • Comparison with ODM: AC-ODM makes more substantial domain-weight adjustments during the first third of training, whereas ODM stabilizes after the first fifth.Both algorithms eventually converge to stable domain weights.
  • Comparison with ODM: Unlike ODM’s nearly constant later-stage weights, AC-ODM continues slight fluctuations, enabling adaptation to evolving LLM states and later parameter updates.ODM’s later-stage constancy indicates limited flexibility in response to parameter updates.
  • Domain trends: High-quality, general-purpose domains tend to gain weight during training regardless of token proportion.Examples include HackerNews, Gutenberg (PG-19), BookCorpus2, StackExchange, USPTO Backgrounds, and Book3.

G. Analysis of Results of MMLU Tasks … H.5. Comparison with RegMix

Across MMLU and additional experiments, AC-ODM benefits from proxy transfer, compact policies, and online adaptation, while remaining effective across larger targets and architectures. Its learned policies stabilize and outperform static mixing baselines.

  • G. Analysis of Results of MMLU Tasks: AC-ODM achieves approximately 21% higher-than-average accuracy in MMLU Social Sciences, indicating strong adaptation to that domain shift.The analysis attributes this result likely to alignment between Social Sciences content and The Pile training distribution.
  • G. Analysis of Results of MMLU Tasks: 19% overall average-accuracy improvement distinguishes proxy-based AC-ODM-410M from direct AC-ODM, including 26% and 17% gains in Social Sciences and Other.AC-ODM-410M consistently improves performance across all MMLU groups.
  • G. Analysis of Results of MMLU Tasks: AC-ODM-410M achieves non-zero accuracy on High School Statistics, Elementary Mathematics, and Management where direct AC-ODM records 0% accuracy.The result supports using a well-trained proxy to capture transferable domain relationships.
  • H.1. Policy Model Size Sensitivity: Policy sizes of roughly 0.25%–0.5% of the target model provide most benefits, whereas very small policies underfit and larger ones add little.This keeps the policy orders of magnitude smaller than the target LLM while retaining capacity for cross-domain interactions.
  • H.2. Reward Stabilization: Mean reward over 22 Pile domains rises rapidly and stabilizes at a high level, indicating convergence toward consistently constructive gradient interactions.The reward need not be monotonic because it is computed from stochastic gradients while the model state changes.
  • H.3. Proxy-Target Scale-Up: AC-ODM maintains a clear validation-perplexity advantage over ODM for a Pythia-12B target trained with a Pythia-1B proxy.The advantage persists throughout the first 25B tokens, supporting transfer to substantially larger targets.
  • H.4. Larger LLaMA-Style Models: Non-proxy AC-ODM remains effective for larger LLaMA-style decoders, including 3B and 7B models beyond the 0.9B setting.Validation perplexity improves consistently as model scale increases.
  • H.5. Comparison with RegMix: AC-ODM improves validation perplexity at every measured checkpoint against RegMix in a half-budget Pythia-1B comparison.The result reinforces the advantage of adapting mixtures online over fixing a globally optimized static mixture.
Loading 2505.23878v2…