Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

IMU-1: Sample-Efficient Pre-training of Small Language Models

George Grigorev

TL;DR

IMU-1 demonstrates that a 430M-parameter decoder-only LM trained on 72B tokens can approach benchmark performance of models trained hundreds to thousands of times more data by integrating a validated recipe of architectural innovations (QK-Norm attention, per-head gating, normalized value residuals, LayerNorm scaling), optimization advances (NorMuon, cautious weight decay, Z-loss, muP parametrization), and flexible learning rate schedules (WSD with post-hoc EMA). The work provides extensive ablations and a complete reproducible pipeline, including data mixtures and implementation details, supporting strong data efficiency and practical transfer to small-LM pretraining. Key findings show significant relative improvements when combining interventions, with EMA contributing additional gains, and WSD enabling checkpoint reuse across stages. The results suggest that careful combination of architecture, optimization, and data curation can yield competitive small-model performance with far fewer tokens, impacting how practitioners approach cost-effective LLM pretraining at sub-billion scales.

Abstract

We present IMU-1, a 430M-parameter language model trained on 72B tokens that approaches the benchmark performance of models trained on 56x more data. We describe a validated training recipe combining recent architectural interventions (QK-norm attention, per-head gating, value residuals, LayerNorm scaling) with optimization advances (NorMuon with cautious weight decay, muP parametrization) and a three-stage training schedule with post-hoc checkpoint EMA. We provide ablations for each component and release code, weights and data to enable reproduction: https://huggingface.co/thepowerfuldeez/imu1_base

IMU-1: Sample-Efficient Pre-training of Small Language Models

TL;DR

IMU-1 demonstrates that a 430M-parameter decoder-only LM trained on 72B tokens can approach benchmark performance of models trained hundreds to thousands of times more data by integrating a validated recipe of architectural innovations (QK-Norm attention, per-head gating, normalized value residuals, LayerNorm scaling), optimization advances (NorMuon, cautious weight decay, Z-loss, muP parametrization), and flexible learning rate schedules (WSD with post-hoc EMA). The work provides extensive ablations and a complete reproducible pipeline, including data mixtures and implementation details, supporting strong data efficiency and practical transfer to small-LM pretraining. Key findings show significant relative improvements when combining interventions, with EMA contributing additional gains, and WSD enabling checkpoint reuse across stages. The results suggest that careful combination of architecture, optimization, and data curation can yield competitive small-model performance with far fewer tokens, impacting how practitioners approach cost-effective LLM pretraining at sub-billion scales.

Abstract

We present IMU-1, a 430M-parameter language model trained on 72B tokens that approaches the benchmark performance of models trained on 56x more data. We describe a validated training recipe combining recent architectural interventions (QK-norm attention, per-head gating, value residuals, LayerNorm scaling) with optimization advances (NorMuon with cautious weight decay, muP parametrization) and a three-stage training schedule with post-hoc checkpoint EMA. We provide ablations for each component and release code, weights and data to enable reproduction: https://huggingface.co/thepowerfuldeez/imu1_base
Paper Structure (60 sections, 8 equations, 2 figures, 11 tables)

This paper contains 60 sections, 8 equations, 2 figures, 11 tables.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Related Work
  4. Sample-efficient pretraining.
  5. Learning rate schedules.
  6. Small language model architectures.
  7. Optimization advances.
  8. Method
  9. Model Architecture
  10. Ablation model.
  11. Architectural Interventions
  12. QK-Norm Attention
  13. Per-Head Gated Attention
  14. Normalized Value Residual Learning
  15. LayerNorm Scaling
  16. ...and 45 more sections

Figures (2)

  • Figure 1: Training loss vs. tokens. Vertical lines mark stage transitions (29B, 57B).
  • Figure 2: Ablation results. (a) Architectural interventions versus baseline. (b) Optimization progression from baseline through NorMuon with CWD. (c) Training curves comparing baseline versus all architectural interventions combined. (d) QK-Norm effect on activation kurtosis across layers---lower kurtosis indicates more stable activations.