One Learning Rule Could Simplify AI Training, Study Finds

Insider Brief

• Researchers developed a unified AI training approach that improved reasoning and preserved earlier skills in small-model experiments by combining existing data and rewarded model-generated answers under one learning rule.
• Mixing data during mathematical training improved retention and subsequent coding learning, but sometimes reduced immediate coding accuracy.
• Results depended on data proportions, timing and computing budgets, with larger models and longer task sequences identified as areas for further study.

Researchers have shown that a single learning rule can combine major stages of artificial intelligence training, improving reasoning and preserving earlier skills in small-model experiments.

The approach, called unified training, mixes existing text and examples with answers generated by a model and evaluated for success. Rather than changing learning algorithms between stages, developers can keep the underlying rule fixed and adjust which data the model encounters, in what proportions and when, according to the study, which was posted on arXiv.

The findings suggest that the organization of AI training — including when models revisit earlier material — can affect both their immediate performance and their ability to learn later tasks. The benefits varied with the data mixture, training schedule and computing budget.

The researchers present the work as a mathematical framework and a working training system. Their experiments also leave several questions open, including how to choose the best data mixture and whether the results extend across larger models and longer sequences of tasks.

Rethinking the Training Pipeline

Foundation models, the broadly trained systems underlying many AI applications, typically develop through several stages. Pretraining exposes them to large collections of text. Midtraining adds more targeted material. Supervised fine-tuning teaches them through examples of desired responses, while reinforcement learning uses feedback on answers the model generates.

The first three stages already share a basic learning principle to adjust the model to make training examples more likely. Reinforcement learning is usually described differently, as a process of improving behavior to earn rewards.

The researchers show how suitable reinforcement-learning objectives can also fit within that likelihood-based approach. Building on a method called Maximum Likelihood Reinforcement Learning, or MaxRL, they express learning from successful generated answers in a form compatible with learning from existing text and demonstrations.

In their implementation, the model generates several answers to a question, and a checking process determines which succeed. Correct answers receive training weight based on the estimated success rate for that question. Incorrect answers receive zero weight. If every answer fails, that question produces no learning update through this mechanism.

Those generated examples join existing material in a shared data buffer. One trainer samples from the combined collection and updates the model, while generators periodically receive newer versions of its parameters.

Under this framework, conventional staged training becomes one possible schedule: give each data source its turn. Unified training also permits sources to overlap and their proportions to change gradually.

Testing Reasoning and Retention

For controlled experiments, the team trained a model with 124 million parameters, the adjustable values that encode what it learns. The runs used a fixed pretraining collection of 20 billion tokens — units of text — and a fixed midtraining collection of 510 million tokens.

The tests involved language-based reasoning puzzles with verifiable answers and adjustable difficulty. Midtraining examples demonstrated reasoning operations, while reinforcement learning introduced new combinations of those operations. Some evaluations used puzzles larger than anything encountered during training.

The researchers measured both the probability that a single sampled answer would succeed and the likelihood of finding a correct answer within 16 attempts. They call the latter measure coverage, distinguishing the range of problems a model can solve across attempts from its reliability on one attempt.

The best unified configurations outperformed staged training on the main reasoning tasks and on larger, unfamiliar puzzles, according to the paper. The controlled reinforcement-learning results were averaged over two training runs with different random seeds.

When the researchers included computing costs across pretraining, midtraining and reinforcement learning, the comparison depended on the available budget. Staged training delivered the best single-answer accuracy at some intermediate budgets. Unified training performed better at the smallest and larger budgets tested in the main settings, and achieved the best accuracy and coverage at almost all budgets when generalizing to larger puzzles.

The team also examined forgetting. Before reinforcement learning, models trained with mixed general and specialized material retained performance closer to the pretraining-only model on general benchmarks. Even the staged configuration selected for its strongest retention lost more ground, the researchers reported.

Trade-Offs Across Skills

A second set of experiments used a larger model with about 760 million parameters and a more varied collection of web text, mathematics, code and reasoning examples.

During mathematical reinforcement learning, the researchers mixed in material from a 240 million-token collection. They varied its configured token share from zero to 5%, comparing that approach with using the additional collection before reinforcement learning.

Mathematical accuracy changed relatively little across the mixtures. Retention benefits were clearer: every nonzero mixture improved the model’s prediction of held-out material relative to its starting point, while staged training worsened that measure.

On SciQ, a science question-answering benchmark, the model using a 2% mixture scored 67%, compared with 53% for the staged model. But adding more existing material generally reduced single-answer coding accuracy, even as coding coverage improved.

Timing changed that balance. Ending the extra data exposure halfway through training produced higher final coding and SciQ accuracy than constant mixing, while preserving less of the improvement in predicting held-out text.

The researchers then gave the mathematically trained models the same additional coding reinforcement learning. All tested unified models finished with higher single-answer accuracy on two coding benchmarks than the staged model, despite starting with weaker coding accuracy.

The model previously trained with a 5% mixture scored 24.96% on HumanEval+ after coding training, compared with 22.94% for the staged model. On MBPP+, it scored 32.69%, compared with 29.81%.

Choosing What Comes Next

The researchers say a central unresolved problem is selecting mixtures systematically. Their experiments identify trade-offs among accuracy, coverage and retention, but no universally preferred schedule.

The computing comparisons describe the configurations tested rather than establish rules that predict performance across model sizes and data budgets. The researchers identify broader scaling studies and closer examination of optimization settings as future directions.

Most experiments also optimized one reinforcement-learning task at a time. Training several tasks together could introduce competing learning demands and different rates of progress, according to the paper.

The researchers propose studying longer task sequences to determine how earlier data choices affect both later learning and the retention of existing skills.

The research team included Yuda Song, Fahim Tajwar, Daman Arora, Jeff Schneider, Ruslan Salakhutdinov, J. Andrew Bagnell, Aarti Singh and Andrea Zanette of Carnegie Mellon University; Alvin Zhang and Daniel Khashabi of Johns Hopkins University; and Bagnell, who is also affiliated with Aurora Innovation.

For a deeper, more technical dive, please review the paper on arXiv. It’s important to note that arXiv is a pre-print server, which allows researchers to receive quick feedback on their work. However, it is not — nor is this article, itself — official peer-review publications. Peer-review is an important step in the scientific process to verify results.