Foundation model training unifies all stages under one method
Training Foundation Models with a Single Algorithm
Machine Learning
Summary
Foundation models usually need different training steps, each using different methods and data. The authors show that all these steps can be combined into a single training method that treats everything as one learning task. This means instead of switching algorithms, training just changes what data to focus on and when. Their tests show this approach keeps the models' broad abilities, improves accuracy for computing effort, and helps them keep learning better over time.
What this means in practice
- •For machine learning engineers: Simplify foundation model workflows by using a single training method mixing diverse data types and learning strategies.
- •For cloud ai service operators: Reduce operational complexity and improve efficiency by unifying different foundation model training stages into one algorithm.
Authors
Yuda Song, Alvin Zhang, Fahim Tajwar, Daman Arora, Jeff Schneider, Daniel Khashabi, Ruslan Salakhutdinov, J. Andrew Bagnell, Aarti Singh, Andrea Zanette
Abstract
Foundation models are typically trained in stages: pretraining, midtraining, supervised fine-tuning (SFT), and reinforcement learning (RL), with changes in both data and learning algorithms. We show that this algorithmic separation is not fundamental. 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. This perspective turns the foundation-model training lifecycle into a problem of choosing which data to train on, in what proportions, and when, with conventional staged training as a special case. We instantiate this framework in a unified training system that mixes static data and RL generations within a single trainer. Across controlled and real-world experiments, unified training improves the accuracy-compute tradeoff and task coverage while preserving broader capabilities, and provides a stronger starting point for continual learning. This work lays the groundwork for a unified science of foundation-model training, connecting mathematical principles, training infrastructure, and the development of model capabilities.