quaternion-llm

Scaling linear language models with SU(2) quaternion recurrence.

The recurrent state lives on S³ (the unit 3-sphere) via quaternion multiplication. Unitary evolution eliminates magnitude decay entirely — information loss occurs only through rotational interference, a geometrically slower process than the exponential eigenvalue decay of standard SSMs.

Key properties

• No magnitude decay. ‖s_t‖ = 1 at every timestep, unconditionally, by construction.
• Native word order. Hamilton product non-commutativity encodes token order algebraically.
• Trig-free critical path. Pure GEMMs + pointwise arithmetic. No sin, cos, or exp.
• O(n) training, O(1) inference. Blelloch parallel scan at train time; fixed-size state at inference.

Recurrence

s_t = R_t ⊗ s_{t-1} ⊗ q_t     where R_t, s_{t-1}, q_t ∈ S³

R_t (from Cayley transform) controls how to integrate. q_t (normalized token projection) controls what to integrate. Both near identity = pass through. Both far from identity = major context update.

Structure

quaternion_llm/
  primitives.py     # hamilton_product, cayley_transform, normalize_quat
  scan.py           # sequential + parallel scan (PyTorch reference)
  layer.py          # QLLMLayer, magnitude gating
  model.py          # full model + configs
kernels/
  scan_kernel.py    # fused Triton selective quaternion scan
  magnitude_kernel.py
train/
  train.py
  data.py
  config.py
eval/
  benchmarks.py
  retention.py      # dilution horizon diagnostics
scripts/
  verify_math.py    # validates unit norm, gradients, associativity
  benchmark_kernel.py

Status

• PyTorch reference implementations complete and validated (make verify, make test)
• Training stack complete — TinyStories, SmolLM2 tokenizer, warmup+cosine LR, gradient accumulation
• Dilution horizon diagnostic in eval/retention.py (make eval-retention)
• Triton kernel not yet implemented (kernels/scan_kernel.py) — GPU required

See nano-quaternion-llm for the educational single-file implementation.

Related

• Inspired by qllm2 (complex-valued S¹ recurrence)
• Moves from U(1)/S¹ to SU(2)/S³ with purely multiplicative recurrence