mimo-7b-mamba2-hybrid-1B-OPD
Stage-3 on-policy distillation, 1B rung of the WSD ladder. The student generates; the teacher scores its own rollouts on reasoning traces. Warmup -> flat plateau -> linear decay, keyed on consumed tokens rather than steps, so each rung extends the previous one's pre-decay weights rather than restarting.
Measured
| metric | score |
|---|---|
| AIME24 pass@1 (think) | 42.08 |
| AIME24 pass@8 (think) | 73.33 |
| AIME25 pass@1 (think) | 31.25 |
| AIME25 pass@8 (think) | 50.00 |
| MATH-500 (think) | 89.90 |
| MATH-500 (no-think) | 62.85 |
| GSM8K (no-think) | 69.30 |
| MMLU (5-shot) | 48.24 |
| PIQA | 71.11 |
| HellaSwag | 56.38 |
| ARC-Easy | 57.95 |
| ARC-Challenge | 39.25 |
| WinoGrande | 58.56 |
| NIAH single @32K | 100.00 |
| NIAH multiquery @32K | 100.00 |
| NIAH multikey @32K | 89.00 |
About this arm
Part of a 2x2 study converting full-attention models into linear-attention
hybrids (uniform 1:4 retention -- every 4th layer keeps softmax attention, the
other 27 are replaced) and recovering capability by staged distillation rather
than pretraining. This repo is the MiMo-7B-RL-0530 x Mamba2 cell; the
matched MiMo x Gated DeltaNet cell is at arianraje/mimo-7b-gdn-hybrid-*
and is directly comparable -- same teacher, same packed data, same token
budgets, same schedules, differing only in the mixer.
Pipeline: surgery/init -> stage-1 hidden-state alignment (100M tok) -> stage-2a forward-KL KD @4k (600M tok) -> stage-2b long-context KD @32k (294M tok) -> stage-3 on-policy distillation, laddered 200M -> 400M -> 600M -> 800M -> 1B generated tokens.
Loading
The architecture is custom (model_type: mimo_mamba2). Remote code is
included, so trust_remote_code=True is all that is needed:
from transformers import AutoTokenizer, AutoModelForCausalLM
repo = "arianraje/mimo-7b-mamba2-hybrid-1B-OPD"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
repo, trust_remote_code=True, dtype="bfloat16", device_map="auto")
For vLLM serving, download the snapshot, put it on sys.path, and call
register() from vllm_mimo_mamba2.py before constructing the engine.
Mamba2's fast path needs mamba-ssm. Without it is_fast_path_available is
False and the mixer falls back to a materializing reference implementation:
correct, but far slower and far more memory-hungry.
Evaluations
Compact raw evaluation JSONs are under evals/. *.generations.json dumps
are deliberately excluded.