Lemma-2.2M

Lemma-2.2M is a 2.2-million-parameter autoregressive language model.

---

Benchmarks

Evaluations are reported in normalized accuracy.

HellaSwag (normalized)	ARC-Easy (normalized)	ARC-Challenge (normalized)	PIQA (normalized)	ArithMark-3 (normalized)	Intelligence Index*
28.29%	37.42%	21.93%	56.42%	37.50%	9.39

*Intelligence Index is calculated according to the Open SLM's Leaderboard's formula.

---

Architecture Overview

Rather than using a standard vanilla Transformer, Lemma-2.2M uses a custom hybrid block structure engineered for high parameter efficiency:

Hyperparameter	Value
Total Parameters	~2.2M
Layers (n_layers)	3
Hidden Dimension (d_model)	192
Attention Heads (n_heads)	4 (head_dim = 48)
FFN Hidden Dimension	384 (2 * d_model)
Context Length (seq_len)	512
Vocabulary Size	2,048 (Tied Embeddings)
Logit Softcapping	30.0

Core Block Components

Each block combines three residual stages with learnable per-channel scaling (gamma) and RMSNorm:

1. Multi-Scale Causal State Mixer (MSCSM): Depthwise dilated 1D causal convolutions ($k \in {3, 5}$, dilations ${1, 2, 5}$) with learnable softmax branch weighting and SiLU gating for local token mixing.
2. Residual Pyramid Attention (RPA): A 3-path hybrid attention mechanism combining:
   • Local Window Attention ($W=8$) with Multi-Octave RoPE ($\theta \in [64, 262144]$ across heads).
   • Retrieved Sparse Attention ($B_{\max}=16$) indexed via a custom Triton causal inverted-index kernel over Spherical Residual Vector Quantization (SRVQ) buckets ($L=1, M=16$).
   • Global Register Compression ($G=8$) via cumulative prefix state pooling.
3. Harmonic Differential StarCarrier FFN: A gated differential feed-forward network utilizing conjugate feature differences, squared ReLU + SiLU ("Star") activation carriers, and internal RMS normalization.

---

Math-Isolated Tokenizer

To prevent subword fragmentation of numbers and equations in a tiny parameter budget, Lemma-2.2M uses a custom 2,048-token Byte-Level BPE tokenizer trained directly on the 8-dataset pre-training mixture.

• Strict Character Isolation: Pre-tokenization enforces regex-level isolation ([0-9+\-*/=()]) so that all digits (0–9) and core arithmetic symbols (+, -, *, /, =, (, )) are guaranteed to remain single, unmerged tokens.
• Zero Digit Merging: Multi-digit numbers are never merged into arbitrary subwords, enabling consistent positional arithmetic learning (reflected in its 37.5% ArithMark-3 score).

---

Training Data Blend

Lemma-2.2M was streamed and dynamically interleaved across 8 datasets using a parallel Deficit Weighted Round-Robin prefetcher (82% General Language, Knowledge & QA / 18% Mathematical Reasoning):

Dataset	Subset / Config	Domain	Mixture Weight
openbmb/Ultra-FineWeb-L3	en-Multi-Style-Synthetic	High-Quality Web & Synthetic	24%
HuggingFaceTB/cosmopedia-v2	cosmopedia-v2	Synthetic Textbooks	20%
roneneldan/TinyStories	default	Narrative & Commonsense	14%
wikimedia/wikipedia	20231101.en	Encyclopedic Knowledge	12%
IFM/TxT360-v2	txt360-qa	Question Answering	12%
openbmb/UltraData-Math	L3-QA-Synthetic	Synthetic Math QA	7%
microsoft/orca-math-word-problems-200k	default	Math Word Problems	6%
nvidia/OpenMathInstruct-2	train_5M / train_1M	Step-by-Step Math Solutions	5%

---

Training & Post-Training Details

Pre-Training

• Training Horizon: 30,000 steps at 65,536 tokens/step ($\approx 1.97\text{B}$ tokens) with micro-batch size $32$ and gradient accumulation $4$.
• Hybrid Muon + AdamW Optimization:
  • Muon Optimizer: Applied to 2D internal projection matrices (lr = 0.030, Nesterov momentum 0.95, 5 Newton-Schulz orthogonalization steps).
  • Fused AdamW: Applied to embeddings, norms, biases, and attention temperature/scale scalars (lr = 1.5e-3, $\beta_1 = 0.9, \beta_2 = 0.95$, weight decay 0.01).
• Curriculum-Decoupled Controller & EMA:
  • Auxiliary SRVQ commitment and load-balancing losses are quadratically warmed up over the first 1,000 steps so early representations form without quantization collapse.
  • Exponential Moving Average (EMA): Shadow weights maintained throughout training with decay 0.995 and bias correction.
• Precision & Acceleration: Trained in float16 mixed precision with torch.compile, TF32 TensorCore matmuls, and custom Triton indexing kernels.

Post-Training

• Following base pre-training, the model underwent ~3 hours of post-training.

---

Intended Use & Limitations

• Intended Use: Research into extreme small language models (SLMs), micro-scale architectural priors, hybrid optimizers (Muon + AdamW), and arithmetic tokenization strategies.
• Limitations: With only 2.2M parameters and a 512-token context window, Lemma-2.2M has limited factual recall and world knowledge compared to larger models and may hallucinate on open-ended facts or complex multi-hop tasks.

License

This model is released under the MIT License.