OMN-Mini-1-1.5B-Base

OMN-Mini-1-1.5B-Base

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

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.

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 after self.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).

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):

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.

---

8. License & Provenance

• Model weights: © Aethr Intelligence. (Set your license here.)
• Tokenizer: Derived from Meta's Llama 3.1 tokenizer — review the Llama 3.1 Community License before commercial or external use.
• Training data: Mixed public corpora. Data licensing should be reviewed before any commercial deployment.

---

9. Citation

(To be added once the technical report / paper is published.)

10. Contact

Aethr Intelligence (aethrai)