OMN-Mini-1-1.5B-Base is a 1.5B-parameter base language model built on a custom transformer architecture — the Orthogonal Manifold Network (OMN) family. It was trained from scratch (no weight initialization from a prior model) on 193B tokens, and is optimized using a novel dual-manifold strategy: the Muon optimizer handles all 2D weight matrices (82.5% of parameters), while AdamW handles all 1D parameters and the embedding (17.5%).
⚠️ Base Model. This release is pre-trained only and has no instruction tuning. It produces text continuations, not chat answers. Use it as a completion engine, or apply SFT / RFT / DPO to build an assistant.
1. Technical Summary
Property
Value
Parameters
1,503,200,256 (tied embeddings)
Architecture
OMN (custom transformer)
Layers
28
Hidden size
2048
FFN intermediate
5504 (SwiGLU)
Attention
16 heads, 4 KV heads (GQA), head dim 128
Positional encoding
RoPE, θ = 500,000, NTK scaling enabled
Vocabulary
128,256 (Llama-3.1 BPE, BOS=128000, EOS=128001)
Context length
32,768 tokens
Training precision
bfloat16
Token budget
193,000,000,000
Compute
4× NVIDIA H200, ~28 days, 0 crashes
Architectural Highlights
Parallel residual. Each block computes x + attn(RMSNorm(x)) + ffn(RMSNorm(x)) from a single shared RMSNorm — not the pre-norm residual x + ffn(x + attn(x)) used in standard Llama/GPT-NeoX. This halves the normalization cost and is one of the OMN-specific departures from the standard transformer.
QK-Norm. Per-head RMSNorm is applied to query and key projections before RoPE. With Muon's aggressive Newton-Schulz orthogonalization and no spectral cap on attention matrices (removed in v2 — see §4), QK-Norm is the sole attention-stability mechanism.
SwiGLU feed-forward. Standard down(SiLU(gate(x)) * up(x)) with intermediate size 5504.
RoPE. Base frequency θ = 500,000 (5× the Llama 3 default) for better long-context extrapolation at 32K. NTK-aware scaling activates only if sequences exceed max_position_embeddings.
Tied embeddings.lm_head.weight is aliased to embed_tokens.weight (saving ~262M params and tying the input/output manifolds).
Dense attention with causal mask. Eval and short-context inference use an explicit torch.tril causal mask. The training path uses flash-attn-varlen for packing efficiency — these two paths are numerically equivalent for single-document sequences.
2. Dual-Manifold Optimization
The central innovation of the OMN family is optimizing different parameter subsets with different optimizers based on their geometry, not their role.
These matrices live on/near a low-rank Stiefel manifold; Muon keeps updates orthogonal, which stabilizes deep training and improves feature learning per parameter.
1D parameters (RMSNorm γ)
~2.5%
AdamW
Rank-1 — Newton-Schulz undefined.
Embedding / tied head
~15%
AdamW (no weight decay)
1D-token embedding manifold; standard AdamW handles this robustly.
Muon's inner loop applies 5 steps of Newton-Schulz orthogonalization per update:
$$
\mathbf{X} \leftarrow a \mathbf{X} + (b \mathbf{A} + c \mathbf{A}^2) \mathbf{X}, \quad \mathbf{A} = \mathbf{X} \mathbf{X}^\top
$$
with coefficients (a, b, c) = (3.4445, -4.7750, 2.0315) chosen so the iteration contracts on the spectrum [0, 1.28] and pushes the effective singular values toward 1. Learning rate for Muon is scaled by $\sqrt{n/m}$ for each $n \times m$ matrix.
WSD schedule. Learning rate follows a Warmup-Stable-Decay profile: 2B-token linear warmup, 164B-token stable plateau at $\eta_\mu = 0.01$ / $\eta_{\text{adamw}} = 3 \times 10^{-4}$, then a $\sqrt{t}$ cooldown over the final 29B tokens down to $\eta_{\text{min}} = 10^{-5}$.
3. Training Data
193B tokens across 15 domains. Sampling weights set for the desired curriculum:
Packing format: each 32K-token training sequence is a pack of one or more documents delimited by EOS (128001), with BOS (128000) at position 0. Loss is computed only on non-EOS / non-boundary positions (ignore index -100).
4. v2 Architecture Revisions
The v1 prototype was stress-tested by a 22-point technical audit that exposed several real defects. The v2 rewrite (the model released here) addressed every one:
D01 / D02 — sampling-with-replacement & resume replay: replaced with a permutation sampler that consumes sequences without replacement within each epoch, with persistent (rng_state, domain_perms, domain_ptrs) saved every 50 batches and restored on resume.
D04 — mmap-per-sample: replaced with a per-shard _mmap_cache.
D05 / D06 — Frobenius spectral cap (harmful rank collapse):muon_spectral_cap is now null. v2 relies on QK-Norm for attention stability; σ_max telemetry from wandb shows stable plateaus (84–191) rather than the pinned / divergent behavior seen in v1.
D07 — weight decay on norms/embedding: optimizer now splits into decay / no-decay groups; norms and embedding get weight_decay=0.
D08 — embedding init clobbered by weight tying: tying is applied afterself.apply(self._init_weights), preserving N(0, 0.02) embedding initialization.
D09 / D10 — MuonClip: removed entirely.
D13 — logits materialization: replaced with Liger fused linear cross-entropy + z-loss (LSE regularization, weight 1e-4) on the training path.
D15 — 28 independent RoPE caches: replaced with a single shared rotary_emb module.
D19 — save cadence: now expressed in optimizer steps.
D22 — AdamW LR floor: correctly decays to min_lr = 1e-5.
5. Evaluation
All benchmarks run through lm-evaluation-harness 0.4.12 against the averaged checkpoint (omn_final_averaged.pt = uniform mean of optimizer steps 734k/735k/736k in fp32, cast to bf16). Code eval uses a sandboxed subprocess harness (wall-clock + CPU-time rlimits, no network).
Benchmark
Domain
Shot / Metric
Score
HumanEval
Code
pass@1 (n=5)
23.9%
MBPP (sanitized)
Code
pass@1 (n=5)
26.1%
PIQA
Physical intuition
5-shot, acc
74.8%
ARC-Easy
Science reasoning
25-shot, acc_norm
71.9%
BoolQ
Reading comprehension
0-shot, acc
62.4%
HellaSwag
Commonsense completion
10-shot, acc_norm
60.6%
WinoGrande
Commonsense reasoning
5-shot, acc
59.4%
ARC-Challenge
Science reasoning (hard)
25-shot, acc_norm
41.0%
OpenBookQA
Elementary science
acc_norm
36.2%
MMLU
Broad knowledge (57 subjects)
5-shot, acc
27.1%
CommonsenseQA
Commonsense knowledge
0-shot, acc
20.5%
GSM8K
Math (word problems)
8-shot CoT, exact_match
16.5%
5.1 Capability Profile
The evaluation reveals a coherent capability profile:
Strengths — structural / algorithmic tasks. The model performs well on benchmarks that require learning procedural patterns (code, physical reasoning, commonsense completion). HumanEval at 23.9% is notably strong given the 193B-token budget — comparable base models at this size trained on the standard Llama token mix typically score in the single digits to low teens on HumanEval, because their data is overwhelmingly web text (~3–5% code). Our ~20.7% code allocation translated directly into working Python generation.
Weaknesses — broad factual knowledge. MMLU (27.1%, near the 25% random baseline for 4-choice questions) and CommonsenseQA (20.5%, near the 20% baseline for 5-choice) indicate that the model has not yet crystallized a wide factual knowledge base. This is the expected, honest consequence of the 193B-token budget; peer base models at this size (SmolLM2-1.7B, Qwen2.5-1.5B, Llama-3.2-1B) train on 1,000B–18,000B tokens. Knowledge is the last capability small models develop and the one that most scales with tokens.
5.2 Comparison to Peers
A rigorous head-to-head requires same-shot, same-metric comparisons against published base-model numbers (not -Instruct variants). Approximate reference points (verify against each model's official technical report before citing):
Model
Params
Tokens
HumanEval
MMLU
GSM8K
OMN-Mini-1-1.5B-Base
1.5B
193B
23.9%
27.1%
16.5%
SmolLM2-1.7B-Base
1.7B
~1,000B
~28%
~52%
~45%
Qwen2.5-1.5B-Base
1.5B
~18T
—
~52%
—
Phi-3-mini-base
3.8B
~3.8T
—
~69%
—
Llama-3.2-1B-Base
1.2B
~9T
~12%
~32%
—
The headline read: OMN's architecture translates tokens into code/reasoning efficiently (HumanEval at 193B tokens is competitive with much larger token budgets), but it needs more tokens to close the knowledge gap that drives MMLU / GSM8K.
6. Usage
OMN is a custom architecture. Inference requires the src/ directory in the working directory alongside model.pt and tokenizer/.
import torch
from transformers import AutoTokenizer
from src.config import OMNConfig
from src.model import OMNModel
tok = AutoTokenizer.from_pretrained("./tokenizer")
ck = torch.load("model.pt", map_location="cpu", weights_only=False)
cfg = ck["config"]
if isinstance(cfg, dict):
cfg = OMNConfig(**cfg)
m = OMNModel(cfg)
m.load_state_dict(ck["model_state_dict"], strict=True)
m = m.to(torch.bfloat16).to("cuda").eval()
# forward(input_ids) returns logits [B, T, V] when labels is None
ids = torch.tensor([[128000] + tok.encode("The capital of France is")], device="cuda")
logits = m(ids)
For generation, use nucleus sampling (top_p=0.95, temperature=0.8). Greedy decoding on this model can lead to repetition loops on long continuations (a property of the base-model untrained tail, not a defect).
Dependencies:torch, transformers. flash_attn and liger_kernel are only required for the training path, not inference.
7. Limitations
Base model only. No instruction-following, no chat, no tool use until post-trained.
Confabulation under uncertainty. The model invents fluent-sounding facts for low-frequency entities rather than hedging. RLHF / DPO is required to fix this.
Multi-step reasoning gap. Single-step arithmetic is reliable; multi-step / symbolic reasoning degrades noticeably. Consistent with the 1.5B / 193B-token regime.
Repetition under greedy decoding. Observed on long continuations. Mitigate with sampling.
English-centric. Non-English performance is substantially lower.
