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

TL;DR

IMU-1 is a small language model trained efficiently on a relatively small dataset, achieving near state-of-the-art performance through innovative architectural and optimization techniques.

Contribution

The paper introduces a novel training recipe and architectural modifications that enable small models to perform comparably to much larger models.

Findings

Achieves benchmark performance with 430M parameters on 72B tokens

Demonstrates effectiveness of architectural interventions and optimization strategies

Provides open-source code and trained weights for reproducibility

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