Efficient Construction of Model Family through Progressive Training Using Model Expansion

1. Problem Statement

The standard practice in Large Language Model (LLM) development involves creating model families (e.g., 1B, 4B, 8B, 70B parameters) to accommodate diverse computational constraints. However, the conventional approach trains each model size independently from scratch.

• Inefficiency: This results in an additive computational cost. Training a family of models $[X_1, X_2, \dots, X_n]$ requires $\sum FLOPs(X_i, T_i)$, which is significantly higher than the cost of training only the largest model.
• Goal: The authors aim to construct a complete model family with reduced total computational cost while maintaining or improving the performance of individual models compared to independent training.

2. Methodology: Progressive Training

The authors propose Progressive Training, a method that constructs a model family by incrementally expanding smaller pre-trained models into larger ones, rather than training them independently.

Core Mechanism

1. Initialization: Start by training the smallest model ($X_1$) from scratch.
2. Model Expansion: For subsequent sizes ($X_{i+1}$), initialize the parameters using a model expansion method ($f$) applied to the previously trained model ($X_i$).
   • The paper utilizes bert2BERT (Chen et al., 2022), which allows for both width expansion (duplicating linear layer weights) and depth expansion (stacking pre-trained layers).
   • This provides a high-quality initialization ($\theta_{init}^{i+1}$) rather than random initialization.
3. Token Allocation Strategy: To ensure the total computational cost matches that of training only the largest model ($X_n$) from scratch, the training tokens ($T_i$) for intermediate models are reduced.
   • The constraint is defined as: $\sum_{i=1}^n FLOPs(X_i, T_i^{prog}) = FLOPs(X_n, T_n^{scratch})$.
   • The "saved" FLOPs from not training smaller models from scratch are effectively reallocated to the training of the largest model within the same total budget.

Learning Rate Adjustment

The authors introduce a strategic adjustment of the maximum learning rate based on model size:

• Smaller models (e.g., 1B) use a higher learning rate ($1.5 \times 10^{-3}$) to accelerate convergence.
• Larger models (e.g., 8B) use a lower learning rate ($3.0 \times 10^{-4}$) to ensure training stability, as larger models are more sensitive to high learning rates.

3. Experimental Setup

• Model Family: Sizes ranging from 1B to 8B parameters (1B, 2B, 4B, 8B), doubling at each stage.
• Architecture: Llama-style architecture with GPT-2 tokenizer.
• Data: FineWeb-Edu and WikiText.
• Baselines:
  • Independent: Each model trained from scratch.
  • Progressive: Models trained via expansion.
  • Progressive+Fixed Data: A control setting where the total unique data is capped at the amount used for the 8B independent model (reusing data for later stages) to rule out data volume as a performance driver.
• Scaling Laws: Experiments conducted under both "Chinchilla law" (20 tokens/param) and "2x Chinchilla law" (40 tokens/param) settings.

4. Key Results

A. Computational Efficiency

• Cost Reduction: Progressive training reduces the total computational cost of building the entire family by approximately 25% compared to independent training.
  • Example: Under the 2x Chinchilla setting, independent training required ~20.4 ZFLOPs, while progressive training required only ~15.4 ZFLOPs (equivalent to training just the 8B model).
• Performance: Despite using fewer total tokens for intermediate models, the progressive models achieved comparable or superior performance to independently trained models across perplexity and downstream benchmarks (LAMBADA, WinoGrande, PIQA, etc.).

B. Impact of Learning Rate Adjustment

• Applying size-dependent learning rates further improved efficiency.
• The combination of progressive training and learning rate adjustment achieved a 31% reduction in FLOPs while maintaining performance, compared to a 26% reduction with fixed learning rates.
• The 8B progressive model with learning rate adjustment outperformed the independent 8B model on almost all metrics.

C. Consistency Across Model Sizes

A unique finding is that progressive training fosters behavioral consistency across the model family:

• KL Divergence: The Kullback-Leibler (KL) divergence between adjacent model sizes (e.g., 1B vs. 2B) was significantly lower in progressive families than in independent families. This indicates that models in the progressive family share more similar internal probability distributions and decision-making processes.
• Speculative Decoding: This consistency translates to practical inference gains. When using a 4B progressive model as a "draft" for an 8B progressive "generator":
  • Acceptance Rate: Increased from 87.2% (Independent draft) to 93.2% (Progressive draft).
  • Inference Speed: Generation time decreased by approximately 14.8%.

5. Key Contributions

1. Efficient Construction Framework: Demonstrated that model expansion can be applied iteratively (progressively) to build a full family of models, reducing total training costs by ~25% without sacrificing quality.
2. Performance Superiority: Showed that progressive models, especially with learning rate tuning, can outperform independently trained models, challenging the notion that "from-scratch" training is necessary for optimal performance.
3. Behavioral Consistency: Identified and quantified that progressive training creates a more coherent model family with aligned output distributions, which is crucial for techniques like speculative decoding and safe deployment scaling.
4. Data Efficiency Analysis: Proved that performance gains are not due to exposure to more unique data (via the "Fixed Data" control experiment) but rather the quality of initialization and the training dynamics.

6. Significance

This work offers a computationally efficient paradigm for the NLP community. As the demand for model families grows to serve edge devices (small models) and complex reasoning tasks (large models), the additive cost of independent training becomes a bottleneck.

• Resource Constraints: The method allows organizations with limited GPU resources to deploy a full spectrum of model sizes.
• Deployment Flexibility: The increased consistency across sizes simplifies the deployment pipeline, allowing seamless switching between model sizes without retraining users or recalibrating safety filters.
• Future-Proofing: The approach is compatible with emerging model expansion techniques, suggesting a scalable path for future LLM development.