Unified training: one algorithm replaces staged model training

A team of ten researchers led by Yuda Song argues that one training algorithm can replace the separate stages of today's foundation models. In "Training Foundation Models with a Single Algorithm", posted to arXiv on October 29, they treat pretraining, midtraining, supervised fine-tuning and reinforcement learning (RL) as the same likelihood-fitting job, which they call unified training. The work has 11 upvotes on Hugging Face daily papers, and the code and starting checkpoints are public.

The results come from small models, 124 million and 0.76 billion parameters, so this is a method to watch, not a recipe for a frontier lab yet.

The problem: every stage forgets a little of the last one

A modern model is trained in a chain. Pretraining on a huge pile of text comes first. Midtraining, a smaller pass on curated data such as math and reasoning traces, follows. Then supervised fine-tuning (SFT) on demonstrations, and RL on tasks with checkable answers. Each stage uses different data and a different objective.

The cost of the chain is that later stages overwrite earlier ones. The paper measures it. In its 124M setting, a model that keeps training on the pretraining text has a held-out loss of about 2.78 nats per token on ClimbMix, a web-text dataset. After the best staged recipe, the same loss is 12.4. The average score on five general benchmarks falls from 42.4% to 35.0%.

The method: one loss, a changing mix of data

The authors build on MaxRL, an RL method that, for tasks scored right or wrong, is equivalent to fitting the model by maximum likelihood to its own correct answers. They extend it so the correct answers can come from any generator, not only the current model.

Once RL looks like likelihood fitting, it has the same form as pretraining or SFT: raise the probability of some target text. The paper puts all sources in one buffer, applies one weighted cross-entropy loss, and gives each source a mixing weight that can change during training. Corpus tokens get unit weight. SFT examples weigh only their demonstration tokens. RL samples weigh only the correct completions, each scaled by one over the pass rate, so rare successes count more.

Staged training is then a special case: all the weight sits on one source at a time. The real questions become which data, in what proportion and when. Generation and training run asynchronously, and the trainer sends new weights to the generators every few updates.

What they measured

Controlled reasoning puzzles, 124M model. The team pretrained on 20B tokens and midtrained on a 510M-token corpus of 796,175 reasoning traces. They then ran RL on puzzles of three kinds: base, extended, and larger puzzles than anything seen in training. According to the paper, the best unified configurations beat the best staged one on all three, with the widest gap on the largest puzzles. Those results appear only in figures, so exact scores are not in the text. The best staged run, with two midtraining epochs and a peak learning rate of 2e-4, reached 24.5% pass@1 on the base puzzles and 5.6% on the largest. In staged training, going from two to four midtraining epochs cut base accuracy to 15.2%.

Retention. Mixing midtraining data into pretraining left the base model almost unchanged: held-out loss between 2.784 and 2.788 on ClimbMix, and a five-benchmark average of 41.9% to 42.4%, against 42.4% for pretraining alone.

Real-world mix, 0.76B model. Here the team mixed 0% to 5% midtraining data into math RL on Math12K. Math accuracy barely moved: MATH-500 pass@1 stayed between 16.1% and 17.0% for every mix and for staged training. Retention changed a lot. Held-out loss on the midtraining corpus rose by 0.047 nats under staged training and fell by 0.147 with a 5% mix. On SciQ, a science question set, a 0.5% mix scored 62.0% against 53.0% for staged. MMLU, a general knowledge test, rose from 37.5% to 40.8% at the 5% mix.

Continual learning. After math RL, the models were trained with RL on coding. All four unified starting points (0.5%, 1%, 2% and 5%) ended with higher HumanEval+ and MBPP+ accuracy than the staged model, even though they started lower. The 5% model was best on all four final metrics.

Where the result is weaker

The authors list the limits themselves.

• The unified mix is not free. After math RL, the staged model scored 15.5% on HumanEval+ and 25.7% on MBPP+ coding tests. With a 0.5% mix the scores were 12.0% and 21.1%. Coding accuracy generally fell as the mix grew, though pass@32 coverage stayed above staged.
• There is no rule for choosing the mixture or schedule, and no mix is best on every metric. Ending the midtraining mix halfway through RL helped coding but gave back part of the retention gain.
• The compute comparisons describe the configurations tried. They are not scaling laws, and staged training stayed competitive at some budgets.
• Most runs train on one RL task at a time. Multi-task training and long sequences of tasks are untested.
• Learning rates and batch sizes were not retuned for each schedule, and a separate downstream adaptation stage is still needed.
• REINFORCE, a simpler RL method, gave higher math accuracy than MaxRL in the 0.76B test (24.8% against 24.0% pass@1), while MaxRL gave higher coverage.

Why it matters

Most of the cost and fragility in model training sits at the handoffs between stages. If a lab can keep older data in the mix while learning new skills, a model can keep learning after release without losing what it had. That is the case the continual-learning result makes, at small scale. Whether it holds at billions of parameters is the open question.

Code and models

The authors released training and generation code under Apache-2.0, including a pinned NeMo-RL copy, and a Hugging Face repository of starting checkpoints for both experiments, as well as the midtraining data. The Hugging Face page lists the trained RL outputs for the 0.76B math runs, with float32 weights. There is also a project page. The paper is 67 pages. A follow-up at larger scale, or by another group, would be the real test.