Skip to content
AI.info

The Pulse

CMU Researchers Recast Foundation-Model Training as One Algorithm

A study proposes using one likelihood-based training rule across several foundation-model training stages, with data mixtures changing over time. Experiments found tradeoffs among reasoning performance, retention and data schedules, rather

CMU Researchers Recast Foundation-Model Training as One Algorithm

AI.info Team · Written with AI · fact-checked ·

A 58.47% result after a second training task

A model trained with a 5% data mixture reached 58.47% on MBPP+ pass@16 after coding reinforcement learning, compared with 52.12% for the paper’s staged-training baseline. The result comes from a study of how earlier training affects later adaptation, not from a test on a large commercial model. The researchers first trained a 0.76-billion-parameter model on mathematics, then applied the same coding-RL procedure to the compared models for 4,000 updates.

Pass@16 estimates whether at least one of 16 sampled responses solves a problem; the paper evaluated 378 MBPP+ problems. The 5% model also scored 46.95% on HumanEval+ pass@16, against 45.12% for the staged model. Those figures show an advantage in this experiment, but the paper does not claim that one mixing rate will suit other models or tasks.

One likelihood loss, changing data

Yuda Song, Alvin Zhang, Fahim Tajwar and their coauthors argue that pretraining, midtraining, supervised fine-tuning and certain reinforcement-learning objectives can share a maximum-likelihood training rule. Their proposal, which they call unified training, treats the training process as a changing mixture of fixed examples and model-generated responses rather than a handoff among distinct algorithms.

“By deriving suitable RL objectives as maximum-likelihood fitting to model-generated target distributions, we place all of these stages under a single likelihood-based training algorithm, unified training.”

Yuda Song, Alvin Zhang, Fahim Tajwar and coauthors, researchers and authors of “Training Foundation Models with a Single Algorithm”

In the framework, each data source can receive a different weight as training proceeds. For reinforcement learning with binary rewards, the target distribution can be built from successful generated responses; static text and demonstrations enter the same likelihood objective. The qualification matters: the authors frame their result around suitable RL objectives, not every possible reinforcement-learning method.

A controlled test favors mixing at some budgets

The paper’s controlled study trains a 124-million-parameter model from scratch, using a fixed 20-billion-token pretraining corpus and a 510-million-token midtraining corpus. Staged training uses the corpora in sequence before reinforcement learning; unified training introduces midtraining data during the end of pretraining and continues mixing it during RL. Researchers test the models on logic puzzles, including larger puzzles held out from both midtraining and RL.

When they compare checkpoints at equal total compute, the result varies with the budget and metric. Staged training leads on accuracy at intermediate budgets, while unified training reaches higher accuracy at the largest budgets and performs especially well on the larger-puzzle generalization test. Staged training remains competitive on coverage for puzzles within the tested range, so the findings are not a universal win for a single recipe.

The paper leaves the schedule unresolved

The authors also find that mixing midtraining data during pretraining can help preserve general capabilities in their controlled experiments. In controlled reasoning-puzzle experiments, adding pretraining text during RL lowered held-out pretraining loss in all eight configurations, while reasoning accuracy and coverage declined in most comparisons. The general benchmark average improved in seven of the eight configurations, with changes smaller than 0.7 percentage points. The results point to tradeoffs among accuracy, coverage and retention rather than a mixture that maximizes all three.

Scale is another limit. The experiments center on 124-million- and 0.76-billion-parameter models, and the authors say they do not establish predictive scaling laws or a generally preferred data schedule. Most runs focus on one RL task at a time, and adapting a foundation model to a new downstream task still requires another training stage. The paper’s concrete claim is narrower: one training loss can accommodate several data sources, while deciding what to mix and when remains an open design problem.

Source

Explore

More articles