Dual-Manifold Optimization in OMN

OMN trains different parameter roles through different optimization regimes. Internal Transformer projection matrices — attention Q, K, V, and output projections plus feed-forward gate, up, and down projections, 1,240,465,408 parameters or 82.5216% of the model — follow a Muon path. The shared vocabulary embedding/output matrix and all normalization parameters — 262,734,848 parameters, 17.4784% — follow AdamW. The split is based on parameter role, not merely tensor dimensionality: the vocabulary embedding is itself a large two-dimensional matrix and nevertheless remains on AdamW.

On the Muon path, a Nesterov-style direction is formed from a momentum buffer at coefficient 0.95 and passed through five Newton–Schulz-style polynomial iterations before application, with a rectangular shape scaling of √(max(1, m/n)) adapting the transformed direction to square and tall projections alike. Gradient clipping at 1.0 applies to AdamW parameters only — Muon is never clipped, by explicit recipe comment. On the AdamW path, conventional Euclidean-style optimization runs at betas (0.90, 0.95), epsilon 1e-8, gradient clipping at 1.0 on AdamW parameters, and zero decay on every active group. Both schedules share the token-based phase structure — 2B warmup, stability to 164B, square-root cooldown to a 1e-5 floor — at separate peaks, and the full schedule is drawn in §9.

The implementation and the parameter routing are directly verified against code and checkpoint. What is not verified is comparative advantage: this report does not claim superior loss, sample efficiency, compute efficiency, or downstream capability against a matched optimizer baseline, because no such baseline has been run. The optimization system should therefore be understood as an implemented architectural hypothesis awaiting controlled comparative evaluation.

Not every parameter should have to learn in the same way. A Transformer is commonly represented to an optimizer as one collection of parameters — yet those parameters play very different roles. Large projection matrices transform learned representations from one space into another. Embeddings map discrete symbols into the model space. Normalization gains regulate scale. Each participates in learning differently.

An optimizer generally receives tensors without necessarily encoding why those tensors exist. A conventional adaptive optimizer flattens every matrix into coordinates and rescales each one independently — a perfectly reasonable default that simply discards matrix shape at the door. Uniformity won by simplicity: one rule, one schedule, minimal bookkeeping. OMN asks whether that abstraction should always remain invisible — and treats the question as an experimental hypothesis, not a proven theorem. Rather than assigning one update rule to the entire network, OMN separates its parameters by role: internal Transformer matrices follow a Muon path, while the shared vocabulary embedding/output parameter and normalization parameters follow AdamW. One model, two optimization paths, one training system.

The purpose of this design is not to assume that one optimizer is globally superior. It is to test whether optimizer geometry can be matched more deliberately to parameter role — and to keep that test honest, the report separates what the implementation establishes from what only a matched baseline could establish.

“Dual-manifold” is AETHR project terminology, and it needs a careful definition because the intuitive reading overclaims. Intuitively, the term describes OMN's design philosophy: parameter groups with different structural roles are allowed to follow different assumptions about how useful update directions should be formed — rather than forcing projections, embeddings, and gains through one shared Euclidean update rule. For internal Transformer matrices, Muon transforms the update direction toward an orthogonalized matrix structure. For embeddings and normalization parameters, AdamW is retained.

OMN does not implement classical constrained optimization in which trained weight matrices are forced to remain on the Stiefel manifold. There is no tangent-space gradient projection, no Stiefel retraction of the weights, no persistent orthogonality constraint on W, and no guarantee that WᵀW = I or WWᵀ = I during training. Muon transforms an update tensor before that update is added to the weight — and that distinction, update versus weight, is the entire technical content of the name.

In exact totals: Muon carries 1,240,465,408 parameters (82.5216%), AdamW carries 262,734,848 (17.4784%), of 1,503,200,256 unique parameters — the same population the architecture report accounts for block by block, now viewed through the optimizer lens. The routing predicate in code is precise — two-dimensional parameters excluding the embedding and head travel Muon; norms, embeddings, and sub-two-dimensional parameters travel AdamW at zero decay — and applied to the checkpoint's 283 tensors it assigns 196 tensors to Muon and 87 to AdamW's no-decay group, leaving the decay group empty. The counts above are what that predicate produces.

A simplified description of OMN might say Muon is used for 2D tensors and AdamW for 1D tensors. That description is inaccurate, and the inaccuracy is instructive. OMN's shared vocabulary embedding/output parameter is 128,256 × 2,048 — a two-dimensional matrix containing 262,668,288 parameters — yet it remains on AdamW. The routing is therefore not determined by tensor rank. It is determined by parameter role.

The design interpretation is structural: internal Transformer matrices are repeated representation-to-representation transformations, while the vocabulary matrix sits at the discrete token/model interface and is tied to the output projection. Every next-token prediction scores through it, and every input token enters through it — one parameter serving both ends of the sequence model. Because the parameter is shared, gradients from its embedding role and its output-projection role accumulate into a single AdamW update — one more consequence of treating role, rather than shape, as the routing principle. OMN treats it as part of the AdamW path. This is presented as a design decision — the report does not claim Muon was experimentally proven harmful for embeddings, because no such evidence exists.

The Muon path proceeds as a fixed sequence — gradient, momentum state, Nesterov-style direction, normalization, orientation adjustment, five Newton–Schulz-style iterations, rectangular shape scaling, parameter update — implemented exactly as follows.

A momentum-like buffer accumulates gradients at coefficient 0.95, and the Nesterov-style direction adds the scaled buffer back to the current gradient. The buffer is allocated as zeros matching the parameter, so Muon state inherits the model's FP32 representation from the first step; an optional spectral-cap guard with power-iteration estimation exists in the implementation but defaults to inactive in the preserved recipe. G_t is the current gradient, B_t the running buffer, D_t the direction supplied to the matrix transformation — described as momentum/Nesterov-style behavior consistent with the implementation, never as exact second-order optimization:

The direction tensor is converted to FP32, rescaled by its largest absolute element (clamped at 1e-7), normalized by its Frobenius norm with epsilon handling, and transposed when tall so the Gram operation uses the smaller dimension — rows beyond columns trigger the transpose, which is exactly the K, V, gate, and up case. Then five polynomial iterations run, orientation and dtype are restored, and the result is applied. The momentum buffer itself is allocated as zeros matching the parameter, so Muon state inherits the model's FP32 representation from the first step. The exact recurrence:

Applied five times. This is described as a finite Newton–Schulz-style polynomial transformation — an orthogonalizing update transformation — never as exact orthogonalization.

Large Transformer matrices receive gradients containing many correlated directions. Muon transforms the matrix-valued update so its singular directions are treated more uniformly before the update is applied — intuitively, reducing domination by a small number of large singular directions in the raw update. What the transformation does not supply is equally worth stating: Muon carries no per-coordinate adaptivity, no curvature estimate, and no memory across steps beyond its single momentum buffer. It is a geometric reshaping of the present direction — nothing more, and the report claims nothing more. A matched all-AdamW run that equaled or beat the hybrid would falsify the strong form of the hypothesis outright; the report is written so that such an outcome would read as science working, not as a promise broken.

The design hypothesis is that the update geometry of a large learned transformation deserves treatment different from elementwise adaptive rescaling alone. OMN therefore places the matrix transformation itself inside the optimizer. This is the conceptual heart of the system — and it is stated as a hypothesis, with no accompanying claims of guaranteed isotropy, condition-number improvement, rank preservation, or faster convergence unless separately measured.

The implemented update carries an explicit rectangular scaling term:

Here m and n are the matrix dimensions as implemented. The scaling term adjusts the transformed direction for rectangular matrices — and it has real work to do, because the same optimizer path must handle square Q and O projections (2048 × 2048), 4:1 rectangular K and V projections (512 × 2048, four heads of 128 sharing each KV head), and tall FFN gate/up matrices (5504 × 2048) beside the wide down projection (2048 × 5504), all in true PyTorch out-by-in code orientation:

The configuration surface exposes a nominal weight_decay = 0.01, but the audited canonical routing places every active AdamW role — embedding, head, and all normalization gains — into zero-decay groups, while Muon carries no default weight decay at all. Applied to the checkpoint's 283 tensors, the routing predicate assigns 196 tensors to Muon and 87 to AdamW's no-decay group — and zero tensors to the decay group. The decaying AdamW group is empty in the canonical OMN-1 configuration: no parameter in the model trains with nonzero weight decay. The report therefore never writes that OMN trained its AdamW parameters with 0.01 decay. The distinction between configuration surface and effective parameter-group behavior is preserved exactly as implemented.

The preserved OMN-1 recipe peaks Muon at 0.01 and AdamW at 0.0003. The two systems operate on different update constructions and therefore maintain separate learning-rate scales — the raw magnitudes must not be compared as though 0.01 were “33× more aggressive,” because the transformations make such comparisons misleading. Both schedules follow the same token-based phases — 2B-token warmup, stability to 164B tokens, square-root cooldown to a 1e-5 absolute floor across roughly 193B tokens — at their own peaks. The coupling is structural, not coincidental: a single token clock steps the shared scheduler once per accepted accumulation window, and both optimizers step together from that one clock, each at its own scale:

The canonical trainer holds model parameters and optimizer state in FP32 with BF16 autocast for forward computation; Newton–Schulz tensors are represented in FP32, as are RoPE expressions and buffers. The historical averaged checkpoint stores BF16 weights. The policy behind both paths is explicit: a deterministic flag governs TF32 matrix multiplication, cuDNN TF32, and BF16 reduced-precision reduction — permitted by default, disabled under determinism — with the effective policy written into the run record. Under default nondeterministic execution, TF32 matrix multiplication may be permitted — so the report states FP32 representation under the configured hardware execution policy, never strict IEEE FP32 throughout.

Every OMN block contains a shared input RMSNorm, parallel attention, and parallel SwiGLU — and the major internal transformations inside both branches lie on the Muon path while the scale-regulating parameters around them lie on AdamW. Conceptually: normalization on AdamW, transformation on Muon, meeting in the residual stream — a recurring optimizer structure across all 28 blocks. Per block, the account is exact: 10,485,760 attention parameters across Q, K, V, and O plus 33,816,576 feed-forward parameters across gate, up, and down — 44,302,336 Muon-owned transformation parameters per block, with the surrounding gains on AdamW. The residual stream itself carries no parameter — it is the one un-owned channel in the block, accumulating Muon-transformed content under AdamW-regulated scale at every layer. The architecture was designed so that this division can be studied coherently; no causal synergy is claimed without tests.

A simpler experiment could replace AdamW with Muon wherever tensor shape permits — the embedding qualifies by rank, and the gains could be folded into matrix groups. OMN deliberately does not, and the hybrid has a concrete cost worth naming: two learning-rate scales to maintain, two kinds of optimizer state to checkpoint, and separate phase behavior to reason about. In exchange, Muon-carrying parameters hold a single momentum buffer rather than AdamW's two moment estimates — the implementation allocates one zeros-like buffer per Muon parameter against AdamW's standard m-and-v pair. Its optimization thesis is not “Muon is a better optimizer” but “different parameter roles may benefit from different update geometry” — two distinct update regimes preserved inside one training system, so the role hypothesis itself becomes the object of study.

All four are marked as research hypotheses — the programme OMN is built to test, not conclusions it has demonstrated.

“Established” below has a fixed operational meaning: read directly off the implementation, the configuration, or the checkpoint — reproducible by anyone holding the same artifacts. “Not yet established” means requiring a matched experiment that has not been run. The table enforces that boundary without commentary:

A genuine Stiefel constraint would hold a weight on the manifold throughout optimization — for appropriate orientation, WᵀW = I or WWᵀ = I at every step — maintained by tangent-space gradient projection, retraction (Cayley, exponential-map, or QR-based), explicit constrained parameterization, or related Riemannian machinery. OMN implements none of those operations on W. It is appropriate to discuss orthogonalized update geometry, Stiefel-like motivation, matrix geometry, and Newton–Schulz orthogonalization of updates. It is not appropriate to say OMN weights remain on the Stiefel manifold, that OMN is a fully Riemannian neural network, or that every Muon-trained matrix is orthogonal.

“Dual-manifold” is intentionally broader than a description of one optimizer primitive. It names the architectural thesis that the model should not be viewed as one homogeneous Euclidean parameter collection when designing its training system — where “homogeneous” means flattened coordinate by coordinate, with matrix shape forgotten at the optimizer boundary. OMN uses two practical optimization regimes to express that idea: matrix-aware transformed updates for internal Transformer transformations, and AdamW for vocabulary and scale-regulating roles. The name therefore describes the research programme and parameter-geometry thesis — not a claim of strict mathematical manifold constraints on every weight. That lets the strong intellectual identity stand while remaining technically exact.

The historical OMN-1 configuration — 1,503,200,256 unique parameters across 28 blocks, 82.5216% Muon against 17.4784% AdamW, approximately 193B consumed input tokens at peak learning rates 0.01 and 0.0003 — warmed over 2B tokens, held stable to 164B, cooled by square-root decay to a 1e-5 floor — demonstrates that the optimization architecture can be implemented and carried through a large training run. It does not, by itself, establish the causal advantage of the optimizer split: the run contains no matched all-AdamW counterfactual, and no benchmark results are placed here. The run's value to this report is architectural rather than comparative: it proves the two-regime system trains end to end at scale.

The correct scientific path is a matched comparison, proposed here — not presented as completed work. Same architecture, initialization policy, tokenizer, data, token order, batch, schedule budget, precision, hardware, and evaluation panel; Arm A runs dual-manifold OMN, Arm B runs an all-AdamW baseline, with an alternative Muon routing as a potential Arm C. Measure training and validation cross-entropy against consumed tokens, stability, gradient and update norms, wall-clock throughput, optimizer-state memory, downstream capability, and reproducibility variance — across multiple seeds where feasible.

These are research questions, not promises.

The long-term importance of this work is not whether one particular 1.5B checkpoint used Muon. The larger question is whether parameter-role-aware optimization remains useful as AETHR develops Darwin across scales — compact Type-S regimes, larger-capability Type-M regimes, and specialized derivatives. Whether role-aware optimization matters more, less, or differently at 5504-wide feed-forward layers versus the wider transformations of a future Type-M model is itself an open scaling question — one the Darwin lineage is designed to answer empirically. Future models are not claimed to use the identical optimizer recipe. OMN provides the experimental foundation from which those decisions can be tested.

Primary sources, repo-relative: src/optimizer.py · src/model.py · src/config.py · src/training.py · src/scheduler.py · configs/phase1_193b.yaml. Audited source fingerprint 8d2014ae9b3a17e472b8e63092dcce70319f2f18f14b829316b6948af8fa02df; checkpoint SHA-256 fd8c59f3b7452e85d583b4b3c5d49ff04ae196027c9b8944760ad574979cc728; implementation Git HEAD 6d2eb04, with the standing qualifier that the audited working tree contained additional modifications, so HEAD alone does not uniquely identify the audited implementation.

The result is not a model whose weights are permanently constrained to two mathematical manifolds. It is a Transformer whose optimization system explicitly acknowledges that different parameter roles may warrant different update geometry. That distinction is the central idea behind OMN's Dual-Manifold Optimization programme — stated once more because everything else on this page depends on it.