OMN: Architecture and Training Design

OMN was developed as an experimental foundation for investigating the Transformer itself, rather than treating architecture as a fixed implementation detail. The programme starts from a simple observation: a language model contains several structurally different kinds of parameters, several distinct numerical regimes, and a training pipeline full of implicit accounting decisions — and the default research posture is to hold nearly all of that fixed while varying data and scale.

The first OMN implementation is a dense, decoder-only autoregressive Transformer: 28 decoder blocks at hidden width 2,048 with a 5,504-wide SwiGLU feed-forward path, 16 query heads against 4 key/value heads at head dimension 128, a 128,256-token vocabulary, a configured position ceiling of 32,768, and 1,503,200,256 unique parameters. Each block applies one learned RMSNorm to its input and feeds the normalized state to attention and feed-forward branches in parallel, so that x_l+1 = x_l + Attention(u_l) + SwiGLU(u_l).

Optimization is assigned by parameter role rather than by tensor rank alone. Internal transformer projection matrices — 1,240,465,408 parameters, 82.5216% of the model — are routed through Muon, which transforms update directions through Newton–Schulz-style orthogonalization. The shared embedding/output matrix and all normalization parameters — 262,734,848 parameters, 17.4784% — are routed through AdamW with zero decay on the active groups. The preserved training recipe accounts for approximately 193B consumed input tokens: this is a measure of data exposure, not of unique tokens, and the report keeps that distinction explicit throughout.

Several historical benchmark, spectral, truncation, and long-context measurements exist in the project documentation, but their raw provenance is incomplete or internally inconsistent. They are intentionally excluded from the central claims of this report. The purpose of what follows is not to claim architectural superiority, but to document the system precisely enough that its design hypotheses, implementation, and unresolved questions can be evaluated independently.

Large language models often inherit a familiar stack of architectural and optimization decisions and concentrate experimentation on data, scale, post-training, or application behavior. The block topology, the optimizer, the attention scaling behavior, and the training-state bookkeeping arrive as settled defaults.

OMN takes a different starting point. It treats the architecture itself — how information flows through each block, how attention regulates scale, how different parameter roles are optimized, and how training state is accounted for — as part of the research surface.

Scale is an extraordinarily powerful driver of model capability. But scale alone does not answer whether the existing parameter budget is being structured and optimized in the most appropriate way. OMN was designed around a narrower question: can a Transformer be made into a better experimental object by matching more of its training behavior to the structure of the model itself?

This motivation produced four design principles — parameter structure, attention stabilization, block structure, and training correctness — developed below.

The table below records the audited OMN-1 reference implementation — the configuration established directly by the current code. Where the historical checkpoint corroborates a value, the report says so; the two artifacts are never merged into one.

The AdamW population decomposes just as exactly. Beyond the shared embedding, the model carries 66,560 normalization parameters: per layer, one 2,048-wide block gain plus 128-wide Q and K gains, across 28 layers, with a final 2,048-wide output gain — 28 × 2,304 + 2,048. The 128-wide Q/K shape independently confirms that each gain vector is shared across its layer's heads rather than allocated per head.

Each OMN decoder block contains exactly one learned RMSNorm at its input. The normalized hidden state feeds both the self-attention branch and the SwiGLU feed-forward branch, and their outputs are added to the incoming residual stream together:

“Shared normalization” means attention and FFN share the same normalized input within a block. It does not mean one global normalization layer reused across all 28 layers — every layer carries its own learned RMSNorm. And the FFN never consumes the attention-updated residual from the same block; both branches read u_l and write back independently.

OMN uses 16 query heads and 4 key/value heads at head dimension 128, so each KV head serves four query heads. The projection structure, verified in the implementation, is:

Queries and keys are independently RMS-normalized over the head dimension before rotary position encoding. Each normalization carries a learned gain vector of width 128 that is shared across the corresponding heads in that layer — there is no per-head gain vector in the implementation, and the report does not describe one.

The order is load-bearing and must be shown exactly: projection → Q/K RMSNorm → RoPE → attention. Normalization never follows RoPE.

OMN applies rotary position encoding with a theta (base) of 500,000, applied after Q/K normalization. The implemented position ceiling is 32,768, and RoPE buffers and rotation expressions in the current implementation compute in FP32.

The tokenizer metadata contains a larger generic model_max_length value of 131,072. That metadata does not expand the architecture: OMN's implemented position ceiling is 32,768, and this report makes no claim about demonstrated retrieval capability at that span. Configured length and learned capability are different objects of study.

The FFN uses SwiGLU — a standard expressive gated structure, included here as part of the experimental block, not presented as an OMN invention:

The arithmetic is worth stating because it closes the parameter accounts exactly. Per block, the three FFN matrices hold 2048 × 5504 × 3 = 33,816,576 parameters; across 28 blocks that is 946,864,128. Attention contributes 10,485,760 per block — Q and O at 2048 × 2048, K and V at 512 × 2048 — for 293,601,280 across the model. Together those reach 1,240,465,408: precisely the Muon-routed population in Figure 2. The remainder is the 262,668,288-parameter shared embedding plus 66,560 normalization gains on the AdamW path. Nothing is unaccounted for.

OMN shares a single parameter between token embeddings and the output vocabulary projection, avoiding a second independent vocabulary matrix. The verified independent-parameter saving is:

Direct inspection of the historical checkpoint confirms the mechanism: the file contains separate serialized entries for the embedding and the head whose bytes are identical, while unique-parameter accounting deduplicates the tied role — 1,765,868,544 serialized entries against 1,503,200,256 unique parameters.

Internal linear matrices are orthogonally initialized with gain 0.85; embeddings are drawn from a Gaussian with σ = 0.02. Both values are established directly by the configuration.

OMN does not route parameters solely by tensor rank. Its optimizer assignment is role-specific, implemented as a three-way partition of the named parameters:

In parameter counts, the Muon path carries 1,240,465,408 parameters (82.5216%) and the AdamW path 262,734,848 (17.4784%). Note what this implies about the AdamW population: the shared embedding accounts for 262,668,288 of those parameters, leaving 66,560 normalization gains — and every active AdamW group runs at zero decay, as developed below.

This hybrid design is substantial enough to receive its own technical report — Matrix-Aware Hybrid Optimization in OMN, currently in preparation and marked forthcoming at the end of this page. What follows here is a summary, not a duplication.

Muon maintains a Nesterov-style momentum buffer at 0.95 and applies five Newton–Schulz-style iterations per step, with peak learning rate 0.01 in the preserved recipe:

At a high level, each update tensor is converted to FP32, magnitude-normalized, oriented so the Gram matrix uses the smaller dimension, transformed through five Newton–Schulz-style polynomial iterations, and restored to its original orientation and dtype. The exact polynomial, for the record:

One nuance matters more than the rest. Although the preserved YAML specifies weight_decay = 0.01, the implementation partitions AdamW parameters into decay and no-decay groups — and the routing places the embedding, the head, and every normalization gain in the no-decay group at weight decay 0.0. This report therefore does not claim the normalization vectors or the embedding were actually decayed: the implemented routing says otherwise, and the routing is what ran.

The current canonical trainer holds model parameters and optimizer state in FP32 with BF16 autocast for forward computation. RoPE buffers and rotation expressions compute in FP32, and Newton–Schulz tensors are represented in FP32. The historical averaged checkpoint, by contrast, stores BF16 weights.

The execution policy itself is explicit rather than ambient. A deterministic flag governs the hardware allowances: TF32 matrix multiplication, cuDNN TF32, and BF16 reduced-precision reduction are each set to the negation of that flag — permitted under default nondeterministic execution, disabled under determinism — and the effective policy is written into the run record alongside the configuration.

One qualification is required: default nondeterministic execution may permit TF32 matrix multiplication. The report therefore states that Newton–Schulz tensors are represented in FP32 subject to the configured hardware execution policy — never that all Muon arithmetic is full IEEE FP32.

The training objective is cross-entropy plus a Z-loss term that penalizes excessive log-partition magnitude, with preserved weight 10⁻⁴:

The mean runs over valid targets only — boundary-masked positions, developed next, are excluded from both terms. And a standing rule for any downstream evaluation: perplexity must be computed from pure cross-entropy, never from the combined training objective. The two are different quantities.

The current reference implementation packs multiple documents into each 32,768-token row with strict segment discipline, anchored to vocabulary ID 128001 as EOS under base seed 42 — reproducible from the configuration rather than emergent from data order. An EOS token closes an attention segment, and row boundaries terminate context likewise: packed attention cannot cross document or segment boundaries, and the token immediately after a boundary is excluded as a next-token target (masked to −100 in the labels).

Two details are easy to get wrong and are stated exactly. First, prediction of EOS itself remains trainable — only the cross-boundary transition is masked. Second, positions normally continue through the packed row unless explicit position IDs are provided; the report does not claim positions reset at every document boundary by default.

The current reference path runs: approved documents → tokenizer → corpus builder (document-end EOS, optional BOS, exact deduplication, protected n-gram screening) → uint32 shards with manifests and tail accounting → global domain/sequence planner → packed rows with boundaries and labels → OMN forward → valid-target cross-entropy plus Z-loss → globally weighted accumulated gradients → Muon and AdamW → checkpoint → resume into the next data state.

The statement above describes the current reference implementation. It does not imply the historical 193B corpus was necessarily built with the current repaired builder — that would conflate two artifacts the report keeps separate.

The preserved phase1_193b recipe schedules by consumed input tokens — not by steps or epochs. Both optimizers share the same phase structure at different peaks: Muon at 0.01, AdamW at 0.0003. Warmup spans the first 2B tokens, the stable phase runs to 164B tokens, and cooldown interpolates toward an absolute floor of 1e-5 as 1 − √(cooldown progress), for approximately 193B tokens total.

This is a WSD-style schedule — warmup, stable, decay — and the report does not call it a cosine schedule. The figure below draws the configured curves, normalized per optimizer; it is recipe, not measured telemetry, and no noisy training curve is shown because none is evidenced.

The preserved recipe allocates approximately 64.23% to general, web, and academic sources, 15.02% to math, and 20.74% to code:

Two independent local lines of evidence support approximately 193B consumed input tokens for the historical OMN-1 run. The sampler-state prefixes record 192,990,412,800 tokens; independent step/recipe arithmetic — 736,000 × 32,768 × 1 × 2 × 4 — gives 192,937,984,000 tokens. In per-step terms that is 736,000 steps of 262,144 tokens, with the factor of 8 spanning batch, accumulation, and world parallelism. Both converge on the same neighborhood, and the report claims no more precision than that convergence supports.

Three negative statements are load-bearing. This does not mean 193B unique tokens. It does not mean 193B unique training examples. And it does not imply every sequence was seen exactly once — historical sampler snapshots show duplicate sequence-slot exposure, examined next.

Across 32 preserved snapshots, the recorded consumed-prefix draws total 5,889,600 against 3,935,427 distinct (domain, sequence-slot) identities — 33.1801% duplicate draws — with Domain 0 covered at 1,142,350 of 1,792,838 slots (63.7174%). These measurements demonstrate overlapping historical sequence-slot exposure.

They do not establish a semantic duplicate rate, privacy leakage, a memorization rate, or complete final-corpus coverage, and the report does not present them as any of those. They are included for a narrower reason: overlapping exposure is part of why training-state accounting became an explicit research concern in the current framework.

Current checkpoints preserve model weights, optimizer states, scheduler state, RNG state, the consumed stream boundary, and data-planner state with recovery lineage. The central distinction is between prefetched work and committed, consumed training state: only the boundary between them lets resumed training reconstruct the correct logical stream position. The mechanics are concrete: sampler state persists every 50 batches, and each stream carries a versioned digest binding sequence length, EOS identity, and base seed, so a resumed run can verify it is continuing the same logical stream rather than a lookalike.

The historical averaged checkpoint was inspected directly for this report — tensor by tensor, not via summary documents. The serialized file holds 1,765,868,544 entries; subtracting the 262,668,288-parameter tied duplicate reproduces the unique count of 1,503,200,256 exactly. Every value below was read off the artifact:

One limit is stated plainly: the three source checkpoints used for the historical average were not available locally, so the averaging operation itself could not be recomputed. The report verifies the product, not the operation. Notably, the checkpoint envelope carries its own averaging metadata — source steps and uniform method stored inside the file rather than in a sidecar — which is how the declaration above was read.

This report does not claim that OMN has demonstrated any of the following:

Several historical benchmark, spectral, truncation, and long-context measurements exist in the project documentation, but their raw provenance is incomplete or internally inconsistent — including per-layer tables with template-pattern digits that cannot be regenerated from shipped analysis outputs. They are intentionally excluded from the central claims of this report. That exclusion is a strength of the report, and it is stated openly rather than buried.

OMN is AETHR's research architecture and training foundation. Darwin is the foundation-model family intended to evolve from the lessons of this research. OMN experiments can inform compact Type-S models, larger Type-M models, and specialized Darwin derivatives — but the audited OMN-1 checkpoint itself is not presented as Darwin Type-S1. That mapping has not been decided, and the report does not imply it.

Audited implementation Git HEAD is 6d2eb04, with one qualifier: the working tree contained substantial modifications and new files, so HEAD alone does not uniquely identify the audited source. The current core-source fingerprint is 8d2014ae9b3a17e472b8e63092dcce70319f2f18f14b829316b6948af8fa02df. Historical checkpoint SHA-256: fd8c59f3b7452e85d583b4b3c5d49ff04ae196027c9b8944760ad574979cc728.