Lemma-2.2M-init2 (Preview)

oscar128372/Lemma-2.2M-init2 is an experimental 2.2M-parameter micro language model designed to test token-level numerical reasoning and hybrid sparse-recurrent attention under extreme parameter and compute constraints.

It was pre-trained from scratch on a single NVIDIA T4 GPU for 6 hours (30,000 steps / 1.97B tokens), followed by **1 hour of post-training** on science, QA, and synthetic math corpora.

---

Model Highlights & Subtleties

1. Math-Isolated 2K BPE Tokenizer

To prevent subword fragmentation of numbers and equations in a tiny 2,048-token vocabulary, the byte-level BPE tokenizer enforces a regex pre-tokenization split ([0-9+\-*/=()] with behavior="isolated"). Every digit (0–9) and core arithmetic operator (+, -, *, /, =, (, )) is guaranteed to be isolated as a single, unmerged token—never fused with adjacent whitespace, letters, or multi-digit chunks.

2. UnifiedLM Architecture (2.2M Params)

Rather than a standard vanilla Transformer, Lemma-2.2M-init2 stacks 3 UnifiedTransformerBlock layers (d_model=192, n_heads=4, head_dim=48, tied embeddings, logit softcapping at 30.0), each containing three residual-scaled sublayers:

• Multi-Scale Causal State Mixer (MSCSM): A gated depthwise causal conv mixer combining 3 dilated branches ($k=3, d=1$; $k=5, d=2$; $k=5, d=5$) via learned softmax routing weights for local n-gram induction.
• Residual Pyramid Attention (RPA): A 3-path hybrid attention mechanism with QKV causal conv context injection and L2-normalized QK states:
  • Local Window Path ($W=8$): Exact sliding-window attention with Multi-Octave RoPE (head-specific rotary bases log-spaced from $64$ to $262,144$).
  • Vector-Quantized Retrieval Path ($B_{\max}=16$): Queries and keys are quantized via Spherical Residual Vector Quantization (SRVQ) ($L=1, M=16$ EMA-updated codebooks with load-balancing and margin losses). A custom Triton causal inverted-index kernel retrieves up to 16 window-disjoint historical tokens per head from the top-1 (10 slots) and top-2 (6 slots) quantization buckets.
  • Global Register Path ($G=8$): Learned register keys aggregate historical context prior to the local window via causal cumulative exponential pooling.
• Harmonic Differential StarCarrier FFN: An expanded (hidden_dim=384) feed-forward network utilizing differential key cancellation (including a circular-shifted harmonic difference term) and a dual-stage squared-ReLU + SiLU star gate before internal RMSNorm.

---

Training Details

Compute & Optimization

• Hardware & Precision: 1× NVIDIA T4 GPU using FP16 mixed precision, TF32 matmuls, custom Triton kernels, and torch.compile.
• Pre-training Schedule: 30,000 steps at 65,536 tokens/step (batch_size=32, accum_steps=4, seq_len=512), totaling ~1.97B tokens (~6 hours).
• Post-training: ~1 hour continuation on targeted QA, science, and synthetic math data with warm EMA initialization and linear LR re-warmup.
• Hybrid Optimizer:
  • Muon (lr=0.03, momentum 0.95, 5-step Newton-Schulz orthogonalization) for 2D internal projection matrices.
  • Fused AdamW (lr=1.5e-3, $\beta=(0.9, 0.95)$, weight decay 0.01) for embeddings, norms, and attention temperature/bias scalars.
• Training Controller: Maintains bias-corrected EMA shadow weights (decay=0.995) for evaluation/export and applies a quadratic warmup curriculum to decouple SRVQ commitment/balance auxiliary losses during early steps.

Pre-training Data Mixture (82% General & QA / 18% Math)

Streamed concurrently via an 8-thread deficit weighted round-robin prefetcher:

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

---

Evaluation & Benchmarks

Evaluated zero-shot using the Open SLM Leaderboard's Intelligence Index formula ($P(\text{score}, \text{chance}) = \frac{\text{score} - \text{chance}}{100 - \text{chance}} \times 100$):

Benchmark Breakdown

Benchmark	Primary Metric	Score (acc_norm)	Raw Acc / Aux
HellaSwag (Zero-Shot)	acc_norm (25% ch.)	27.91% ± 0.45	26.92% (raw)
AI2 ARC-Easy	acc_norm (25% ch.)	34.34% ± 0.97	36.32% (raw)
AI2 ARC-Challenge	acc_norm (25% ch.)	20.99% ± 1.19	17.15% (raw)
AI2 ARC Combined	acc_norm (25% ch.)	27.67%	26.74% (mean(E, C))
PIQA (Validation)	acc_norm (50% ch.)	54.95% ± 1.16	57.02% (raw)
ArithMark-3.0	acc_norm (25% ch.)	39.50% ± 1.55	39.70% (raw)

Open SLM Leaderboard Intelligence Index

Pillar Benchmark	Score (%)	Chance (%)	$P(\text{score}, \text{chance})$	Weight
HellaSwag	27.91%	25.00%	+3.88	1.00
ARC Combined	27.67%	25.00%	+3.56	1.00
PIQA	54.95%	50.00%	+9.90	1.00
ArithMark-3.0	39.50%	25.00%	+19.33	0.65
Intelligence Index	—	—	+8.19	$\Sigma w = 3.65$

---

Usage & Limitations

• Custom Modeling Code Required: Because Lemma-2.2M-init2 uses custom modules (ResidualPyramidAttention, SRVQ, MultiScaleCausalStateMixer, HarmonicDifferentialStarCarrierFFN) and a Triton kernel (build_causal_index_fast_kernel), loading the weights requires the UnifiedLM class definition and a CUDA-capable GPU with triton installed. The modeling code is available under the repository. (modeling.py)
• Preview Scope: At 2.2M parameters and a 512-token context window, this checkpoint is a research preview exploring micro-scale architectural efficiency and arithmetic tokenization; it is not intended for production deployment or factual reliability.