Industrializing Orthogonal Optimizers: Overcoming LLM Training Memory Bottlenecks

1. The Memory Consumed by Optimizers: The First Cost of Trillion-Parameter Training

In recent technical reports on next-generation open-source large language models (LLMs), the most notable detail for engineering teams is not the parameter scale itself, but an engineering breakthrough: ultra-large scale training has adopted orthogonal optimizers (a class of matrix-level optimization methods represented by Muon) and successfully scaled them to the trillion-parameter level. This marks the first industrial validation of a technology previously confined to small-scale community experiments. To grasp the significance of this, we must first examine the memory ledger of Adam-class optimizers.

Since the introduction of Adam, there have been virtually no substantial challengers in the deep learning optimizer ecosystem. Its core is adaptive learning rates: maintaining a first-order moment $m_t$ (exponential moving average of gradients) and a second-order moment $v_t$ (exponential moving average of squared gradients) for each parameter to dynamically adjust the step size. This design delivers strong versatility and stability, but at a high memory cost—every model parameter requires two additional fp32 memory spaces.

Taking mixed-precision training as an example, the memory footprint of a single parameter consists of: 2 bytes for bf16 weights, 2 bytes for bf16 gradients, 4 bytes for fp32 master weights, 4 bytes for fp32 first-order moment, and 4 bytes for fp32 second-order moment. The optimizer state ($m_t + v_t$) alone consumes 8 bytes, which is several times larger than the weight itself (2 bytes).

For a 1.6T parameter model, these two state variables alone require approximately 12.8TB of VRAM:

def estimate_optimizer_state(num_params, bytes_per_state=4, num_states=2):
    """Estimate VRAM usage of optimizer states (in bytes)"""
    return num_params * num_states * bytes_per_state
n = 1.6e12  # 1.6T parameters
print(f"mt + vt approx {estimate_optimizer_state(n) / 1e12:.1f} TB")
# Output: mt + vt approx 12.8 TB

The real cost of this overhead is not just expense, but the squeeze on resources. The training memory budget is rigid: whatever the optimizer states occupy is subtracted from the memory available for activations; and activations directly determine the upper limits of batch size and context length. For every 1TB the optimizer state expands, you must either reduce the batch size, truncate the sequence, or add more sharding devices. The cost of all three eventually hits the total training bill. For models where long context is a core capability, this creates a hard ceiling.

2. The Origin of Orthogonal Optimizers: An Unexpected Breakthrough in a Community Contest

In late 2024, a community-driven “minimalist training competition” attracted numerous researchers. The rules were simple: train a mini language model with hundreds of millions of parameters to a target loss using the least compute and time. While nearly all participants were squeezing the last bits of performance out of Adam’s hyperparameter tuning, one researcher submitted a fundamentally different approach: orthogonalizing the parameter gradients of hidden layers before performing the update. The result was a cliff-like drop in training steps, shattering the existing record. This method was subsequently refined into an independent optimizer, now widely known as Muon.

1. Two Worldviews: Element-wise vs. Matrix-level

Adam’s worldview is “element-wise”: a 1000x1000 weight matrix is treated as 1 million independent scalars, each assigned an exclusive step size. The problem with this approach is that it completely ignores the geometric structure within the matrix.

The orthogonal optimizer’s worldview is “matrix-level”: a 1000x1000 matrix is an indivisible high-dimensional geometric object. It projects the “matrix where gradient and momentum are fused” into an orthogonal matrix—where all singular directions are mutually orthogonal and energy is uniformly distributed—before adding this matrix to the weights in a single step at a unified scale.

To use a metaphor: training is like rolling millions of iron balls toward a valley. Adam is a micromanager, measuring the slope and fine-tuning the step length for each ball individually. The orthogonal optimizer, conversely, organizes the iron balls into a formation, letting them advance collectively with uniform, mutually constraining steps. The improvement in command efficiency is obvious.

The mathematical rationale is that the pathological directions of the loss surface are typically concentrated in a few singular directions of the gradient matrix. Element-wise optimization is easily “hijacked” by these strong directions, causing back-and-forth oscillation. Orthogonalization flattens the spectral norm, spreading the update energy evenly across all directions. This is equivalent to executing steepest descent under the spectral norm metric, utilizing gradient information more “fairly” at every step.

2. Reconstructing the Memory Ledger: From 8 Bytes to 2 Bytes

The most direct benefit of the orthogonal optimizer is eliminating the second-order moment. It only needs to maintain a single momentum buffer, which in practice can be stored in bf16. The state footprint per hidden layer parameter drops from Adam’s 8 bytes to approximately 2 bytes. For a 1.6T parameter model, the optimizer state is compressed from ~12.8TB to ~3.2TB. The nearly 10TB of saved VRAM can be directly converted into larger batch sizes, longer contexts, or higher data throughput. On the computation side, orthogonalization consists of only a few small-scale matrix multiplications, taking a negligible amount of time compared to forward and backward propagation.

3. Newton-Schulz Iteration: Orthogonalization Without SVD

Performing direct SVD decomposition on GPUs is prohibitively expensive. The industry commonly uses Newton-Schulz iteration to approximate polar decomposition using pure matrix multiplication. A conceptual implementation is as follows (coefficients are from common public implementations):

def orthogonalize(G, steps=5):
    """Newton-Schulz iteration orthogonalization (conceptual code)"""
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.bfloat16()
    X = X / (X.norm() + 1e-7)
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * (A @ A)
        X = a * X + B @ X
    return X

The entire process uses only matrix multiplication, naturally fitting the GPU’s compute architecture and allowing efficient execution in bf16. Another critical engineering detail: because the spectral norm of an orthogonal matrix is 1, a scaling coefficient must be applied based on the shape of the weight matrix to keep the magnitude of each update comparable to AdamW. This is a standard practice in public implementations.

3. From 120M to Trillion Parameters: Three Questions That Must Be Answered

“Formations look good, but no one actually marches in them on the battlefield.” Since its inception, this new method has faced a sharp question: it was initially validated only on mini-models with around 120M parameters. Just because it works at a small scale doesn’t mean it works at a large scale—the history of optimizers is full of examples that shone on small models but failed on large ones. Industrial deployment must answer three questions.

First, does convergence stability hold at scale? In large model training, loss spikes and gradient anomalies are far more frequent than in small models. Orthogonalization inherently smooths out abnormally large singular values in gradients, potentially offering superior robustness against outlier gradients. However, rounding errors in bf16 iterations, momentum coefficients, and their interaction with gradient accumulation still need to be validated in multi-month, multi-node long runs. Recent public industrial practices have provided a preliminary positive answer: this method can stably scale to the trillion-parameter level.

Second, is it compatible with distributed training stacks? The complexity of industrial training lies in parallel strategies, which must be verified item by item. For data parallelism and optimizer state sharding, smaller states mean lighter shards and less communication—a pure win. For tensor and pipeline parallelism, updates happen locally on weight shards, requiring verification of the equivalence of orthogonalization at chunk boundaries. For the learning rate system, orthogonal updates must be scaled by matrix shape and aligned with the unit RMS update of AdamW to avoid introducing extra tuning costs.

Third, does it overthrow the existing hyperparameter ecosystem? The mainstream answer in practice is a hybrid configuration rather than a wholesale replacement. Two-dimensional weight matrices (hidden layers) use the orthogonal optimizer, while one-dimensional parameters (embedding layers, output heads, normalization scaling parameters) continue to use AdamW—one-dimensional parameters lack matrix structure, rendering orthogonalization ineffective. This approach captures both memory and convergence benefits while keeping the risk profile within acceptable limits.

4. A Pilot Roadmap for Engineering Teams

If training memory is currently a bottleneck for your team, you can perform a low-cost validation in four steps:

1. Calculate the Ledger: Use the estimation formula above to quantify the current optimizer state footprint and define the amount of VRAM that can be released.
2. Small-scale A/B Testing: Run two training groups, one with AdamW and one with the hybrid orthogonal optimizer, on the same data and architecture. Control variables to compare loss curves and downstream evaluation performance.
3. Stability Check: Continuously monitor the gradient norm, the spectral norm distribution of the update matrix, and the frequency of loss spikes to ensure no drift over long training cycles.
4. Gradual Scale-up: After successful validation, expand the parameter scale and parallelism. Utilize logging services and monitoring alerts to automate anomaly detection.

During the pilot phase, leveraging a mature container platform can significantly reduce experimental costs. For example, on cloud container platforms like Baidu AI Cloud, GPU resources can be elastically scheduled on demand, and training tasks can be launched or terminated flexibly. Coupled with built-in logging and monitoring alert capabilities to quickly locate training anomalies, the cycle for an “optimizer A/B comparison experiment” can be compressed from months to weeks.

It is also crucial to define the selection boundaries: orthogonal optimizers are not a silver bullet. Their benefits are concentrated in scenarios with massive parameter counts, long contexts, and memory constraints. For small-to-medium parameter scales or training scenarios dominated by embedding layers, AdamW’s mature ecosystem and tuning experience still hold the advantage.

5. Conclusion

Optimizers are one of the slowest-evolving areas in deep learning: the transition from SGD to Adam took nearly two decades, and there has been no true industrial challenger since Adam. If orthogonal optimizers successfully complete the leap from community contests to trillion-parameter training, the significance is more than just a faster optimizer—it represents a restructuring of the entire training cost profile. The same memory budget can be converted into longer contexts, larger batches, or even larger models themselves. For every engineering team throttled by memory constraints, now is the right time to calculate the ledger and launch a small-scale comparative experiment.