LLM pretraining optimizers stop orthogonalizing layer by layer: TUON stacks layers into a tensor and beats Muon and TEON on four model scales
Synopsis
Large language model pretraining optimizers are moving from handling each weight matrix separately to stacking the momenta of one module's layers into a third-order tensor. Layer-wise Muon ignores cross-layer correlation, while matricization-based tensor optimizers handle only two-layer stacks and QKV. TUON applies a unitary transform along the layer mode, polarizes each transformed slice, and inverts the transform; it lowers validation loss versus Muon and TEON on GPT-124M/350M and Qwen-0.6B/4B at comparable step time.
Interpretation
LLM pretraining optimizers can now update several layers of the same Transformer module as one object rather than orthogonalizing each layer alone. Matrix-aware Muon treats layers independently and neglects cross-layer correlations, while tensor-aware TEON relies on unfolding the tensor into a matrix. Across GPT-2 124M/350M and Qwen3 0.6B/4B, TUON reaches validation losses of 2.9357, 2.6925, 2.8887 and 2.7598, below Muon's 3.0488, 2.7404, 2.9056 and 2.7636 and TEON's 2.9442, 2.6935, 2.8978 and 2.7604.
The mechanism mixes layers first and separates them afterwards: a unitary transform along the layer mode recombines layer information in a transform domain. Each transformed slice is then polarized, meaning its nonzero singular values are set to one so only directions remain, and an inverse transform returns the update to the original tensor domain. The default DecorQ transform comes from an SVD of the layer Gram matrix of the momentum tensor and is refreshed every 200 iterations; on Qwen3-4B, Newton–Schulz and PolarExpress polarization give 2.7629 and 2.7638 at two-layer stacking and 2.7598 and 2.7569 at full stacking.
Stacking more layers stops hurting performance, which unfolding-based tensor optimizers did not achieve. TEON is best at two-layer stacking and degrades at full stacking, whereas TUON with DecorQ improves from two layers to full stacking on all four models. GPT-2 124M moves from 3.0299 to 3.0233, GPT-2 350M from 2.8183 to 2.8143, Qwen3 0.6B from 2.8938 to 2.8887 and Qwen3 4B from 2.7629 to 2.7598.
Perspective
The results cover four model scales, GPT-2 124M/350M and Qwen3 0.6B/4B, two pretraining corpora, OpenWebText and C4, and two module scopes, attention projections alone and QKVO plus MLP; all cross-method comparisons stay within one model family and one training horizon.
The gain from moving to full-layer stacking is real but bounded, and the authors state it remains below the theoretical maximum, attributing this to shared structure degrading as layers become more diverse. Most ablations train for 20k iterations, unlike the 50k-iteration GPT-2 main runs.