| | --- |
| | language: |
| | - en |
| | license: apache-2.0 |
| | tags: |
| | - language-model |
| | - sample-efficient |
| | - pretraining |
| | - transformer |
| | library_name: transformers |
| | pipeline_tag: text-generation |
| | arxiv: 2602.02522 |
| | --- |
| | |
| | # IMU-1 Base |
| |
|
| | This repository contains the IMU-1 Base model, a sample-efficient 430M parameter language model introduced in the paper [IMU-1: Sample-Efficient Pre-training of Small Language Models](https://huggingface.co/papers/2602.02522). |
| |
|
| | IMU-1 is trained on 72B tokens and approaches the benchmark performance of models trained on 56× more data. |
| |
|
| |
|
| | ## Model Details |
| |
|
| | | Parameter | Value | |
| | |-----------|-------| |
| | | Parameters | 430M | |
| | | Hidden dim | 1,152 | |
| | | Layers | 30 | |
| | | Attention heads | 18 | |
| | | KV heads (GQA) | 6 | |
| | | Vocab size | 49,152 | |
| | | Max context | 1,152 | |
| | | Training tokens | 72B | |
| |
|
| | ### Architecture |
| |
|
| | IMU-1 uses a validated recipe combining recent advances: |
| | - **QK-norm attention** with learnable scale |
| | - **Per-head gating** (sigmoid-based) |
| | - **Value residual learning** |
| | - **LayerNorm scaling** (depth-dependent) |
| | - **GQA** (grouped query attention) |
| | - **SwiGLU** activation |
| | - **RoPE** positional encoding |
| |
|
| | ### Training |
| |
|
| | - **Optimizer:** NorMuon with cautious weight decay, muP parametrization |
| | - **Schedule:** Three-stage WSD (Warmup-Stable-Decay) |
| | - **Post-processing:** Checkpoint EMA (β=0.8) |
| |
|
| | ## Usage |
| |
|
| | ```python |
| | from transformers import AutoModelForCausalLM, AutoTokenizer |
| | |
| | model = AutoModelForCausalLM.from_pretrained( |
| | "thepowerfuldeez/imu1_base", |
| | trust_remote_code=True |
| | ) |
| | tokenizer = AutoTokenizer.from_pretrained("thepowerfuldeez/imu1_base") |
| | |
| | text = "The quick brown fox" |
| | inputs = tokenizer(text, return_tensors="pt") |
| | outputs = model.generate(**inputs, max_new_tokens=50) |
| | print(tokenizer.decode(outputs[0])) |
| | ``` |
| |
|
| | **Note:** This model uses custom modeling code. You must pass `trust_remote_code=True` when loading. |
| |
|
| | ## Benchmark Results |
| |
|
| | | Benchmark | Score | |
| | |-----------|-------| |
| | | HellaSwag (0-shot) | 51.1 | |
| | | ARC-Easy | 71.4 | |
| | | ARC-Challenge | 41.1 | |
| | | PIQA | 70.2 | |
| | | Lambada (OpenAI) | 51.3 | |
| | | Winograd | 74.7 | |
| | | WinoGrande | 55.2 | |
| | | BoolQ | 59.5 | |
| | | **CORE (centered)** | **30.2** | |
| |
|
| | ## Training Stages |
| |
|
| | | Stage | Iterations | Tokens | Data | |
| | |-------|------------|--------|------| |
| | | 1. Stable | 100k | 29B | DCLM-edu, FineWeb-edu | |
| | | 2. Decay | 100k | 28B | Higher quality filters | |
| | | 3. Midtrain | 65k | 14B | Instruction, reasoning, code | |
| |
|
| | ## Resources |
| |
|
| | - **Training Code:** [sample_efficient_gpt](https://github.com/thepowerfuldeez/sample_efficient_gpt) |
| | - **Stage 1 Data:** [1218_imu1_base_stable_corpus](https://huggingface.co/datasets/thepowerfuldeez/1218_imu1_base_stable_corpus) |
| | - **Stage 2 Data:** [1226_imu1_base_decay_corpus](https://huggingface.co/datasets/thepowerfuldeez/1226_imu1_base_decay_corpus) |
| |
|
| | ## Citation |
| |
|
| | ```bibtex |
| | @misc{grigorev2026imu1sampleefficientpretrainingsmall, |
| | title={IMU-1: Sample-Efficient Pre-training of Small Language Models}, |
| | author={George Grigorev}, |
| | year={2026}, |
| | eprint={2602.02522}, |
| | archivePrefix={arXiv}, |
| | primaryClass={cs.LG}, |
| | url={https://arxiv.org/abs/2602.02522}, |
| | } |
| | ``` |
| |
|
| | ## License |
| |
|
| | Apache 2.0 |
| |
|
