PaperlayerView exact revision on arXiv

Source-linked AI summary

Mano: Restriking Manifold Optimization for LLM Training

Yufei Gu, Zeke Xie

arXiv:2601.23000v1cs.LGcs.AI

TL;DR

Training LLMs is computationally expensive, while AdamW and Muon capture different aspects of optimization structure. Mano reformulates manifold optimization by projecting momentum onto tangent spaces and constraining it with a rotational Oblique manifold, outperforming AdamW and Muon with lower resource demands.

  • Problem

    Training LLMs requires substantial computational resources, while existing AdamW and Muon optimizers respectively overlook structural information or curvature information.

  • Method

    Mano projects momentum onto tangent spaces and constrains it on a rotational Oblique manifold for LLM training.

  • Results

    Mano outperforms AdamW and Muon in LLM training with lower computational overhead than Muon and a smaller memory footprint than AdamW.

  • Takeaways & Limitations

    The results highlight the potential of geometrically aware manifold techniques combined with modern optimization strategies for efficient LLM training.

  • Takeaways & Limitations

    The empirical scope is constrained by available computational resources, and the convergence analysis covers only a simplified version of Mano.

Abstract

from arXiv · show

While large language models (LLMs) have emerged as a significant advancement in artificial intelligence, the hardware and computational costs for training LLMs are also significantly burdensome. Among the state-of-the-art optimizers, AdamW relies on diagonal curvature estimates and ignores structural properties, while Muon applies global spectral normalization at the expense of losing curvature information. In this study, we restriked manifold optimization methods for training LLMs, which may address both optimizers' limitations, while conventional manifold optimization methods have been largely overlooked due to the poor performance in large-scale model optimization. By innovatively projecting the momentum onto the tangent space of model parameters and constraining it on a rotational Oblique manifold, we propose a novel, powerful, and efficient optimizer **Mano** that is the first to bridge the performance gap between manifold optimization and modern optimizers. Extensive experiments on the LLaMA and Qwen3 models demonstrate that Mano consistently and significantly outperforms AdamW and Muon even with less memory consumption and computational complexity, respectively, suggesting an expanded Pareto frontier in terms of space and time efficiency.

1. Introduction

Training LLMs incurs substantial hardware, energy, and computational costs, motivating more efficient optimizers. Mano revisits manifold optimization for LLMs and reports improved performance with lower resource demands than AdamW and Muon.

  • Training LLMs with billions of parameters requires specialized accelerators and incurs substantial energy and computational costs.
  • AdamW relies on diagonal per-parameter curvature estimates, thereby ignoring spectral information.
  • Traditional manifold optimization has been largely overlooked for LLMs because its assumptions and constraints do not readily fit large-scale model training.
  • Mano restrikes manifold optimization for LLM training through multiple reform strategies.
  • Mano uses less memory and significantly lower computational complexity than modern optimizers such as Adam and Muon.
  • Mano significantly outperforms Adam and Muon in test perplexity across token consumption and wall-clock time, with reduced gradient variance suggesting better convergence.

2. Related Works

Related work spans widely used Adam-based optimizers, matrix-based spectral preconditioning, and geometrically aware optimization. Although manifold methods are established in other deep-learning domains, their use in LLM training remains limited.

  • Adam-based optimizers are widely used because of their simple design, adaptive learning rates, and robust performance across domains.
  • AdamW’s first- and second-moment estimates can impose substantial memory overhead for large language models.
  • Muon applies Newton-Schulz iteration for spectral normalization, producing semi-orthogonal momentum updates across spectral directions.
  • Geometric optimization exploits intrinsic objective-function structure through methods defined on Riemannian manifolds.
  • Manifold optimization has been developed for architectures including CNNs and RNNs but remains largely overlooked in LLM practice.

3. Preliminaries

Riemannian optimization projects gradients onto tangent spaces, updates along the resulting directions, and maps parameters back onto a manifold. Traditional approaches can be computationally expensive and may constrain neural-network optimization when the manifold is poorly aligned with the objective.

  • A Riemannian manifold is a smooth geometric space with a smoothly varying inner product on each tangent space.
  • Riemannian SGD projects the gradient onto the tangent space, takes a manifold-directed step, and maps the parameters back onto the manifold.
  • The exponential map gives a geometrically exact update, while numerical retractions provide cheaper first-order approximations that maintain the manifold constraint.
  • Traditional manifold methods can require expensive decompositions such as SVD or QR, reducing optimization efficiency.
  • Manifold constraints can restrict loss-landscape exploration when the selected geometry does not align with the underlying objective.

4. Methodology

Mano reformulates manifold optimization as a momentum-driven optimizer using tangent-space projection, rotating Oblique normalization, and Euclidean descent. The design targets efficient LLM training while preserving practical implementation and convergence properties.

  • Reformed Manifold Optimization: Mano projects momentum onto the tangent space and constrains the resulting update on a manifold before applying the parameter update.The method defines tangent-space projection and manifold normalization, while retaining Euclidean objectives and solutions.
  • Manifold Selection and Design: Oblique normalization is selected for computational efficiency, while Table 1 compares average geodesic distances across Oblique, Sphere, and Stiefel manifolds.The table reports distances over 1000 consecutive Qwen3-0.6B AdamW update steps separately for attention projections and MLP layers.
  • Rotational Manifold Scheme: The optimizer alternates column-wise and row-wise normalization across iterations, forming a rotating Oblique manifold rather than a static geometric structure.This rotation addresses the unvalidated assumption that column directions dominate row directions in LLM parameter matrices.
  • The Mano Optimizer: The Mano update uses momentum, alternating manifold normalization, tangent momentum, momentum normalization, rescaling, learning rate, and weight decay.The implementation rotates dimensions by timestep parity and applies a 0.2√n_k rescaling factor for the selected dimension.
  • Efficiency: Mano adds two normalization steps and one tangent-space projection without additional hyperparameters or problem-specific geometric assumptions.Its memory overhead is comparable to SGD-momentum or Muon and approximately halves the footprint of Adam-based optimizers.
  • Efficiency and Convergence: Mano’s FLOP overhead is at most 11/6B relative to the layer baseline, compared with Muon’s 5m/B overhead, while convergence is guaranteed under simplified assumptions.The theorem assumes no momentum and a static Oblique manifold; the proof is provided in Appendix E.

5. Experiments

Across LLaMA and Qwen3 experiments, Mano generally improves later-stage convergence and sample efficiency over AdamW and Muon, while exhibiting lower gradient variance and favorable spectral behavior. Its computational advantage persists in wall-clock comparisons, but ablations expose scaling limitations for static normalization and traditional Riemannian SGD.

  • Experiment Results: Mano ultimately surpasses AdamW and Muon in LLaMA pretraining, despite sometimes converging more slowly initially.Baseline losses plateau while Mano continues reducing loss at a nearly constant rate; larger models reach the crossover later.
  • Experiment Results: Mano transfers its performance advantage from LLaMA to Qwen3 architectures on the Pile dataset.The replicated Qwen3 experiments demonstrate cross-architecture transferability of the advantage.
  • Experiment Results: With 10B training tokens, Mano ultimately achieves the best performance across the three optimizers despite underperforming AdamW during part of LLaMA-130M training.The authors identify larger over-training experiments as future work because of computational constraints.
  • Learning Dynamics: Mano maintains lower gradient variance and higher SNR than Muon under the same momentum coefficient, indicating more stable training dynamics.The comparison uses µ = 0.95 and similar update RMS.
  • Learning Dynamics: Mano’s manifold normalization lifts update spectra while preserving singular-value ordering, unlike Muon’s whitening, which flattens the spectrum.The spectral comparison covers attention and MLP layers in LLaMA-350M.
  • Learning Dynamics: Mano achieves 1.75× and 1.38× faster convergence than Muon on LLaMA-350M and LLaMA-1B, respectively, in one-day experiments.Its normalization time grows linearly with model dimension, contrasting with the reported exponential growth for Muon.
  • Ablation Studies: Standard Riemannian SGD with momentum fails to reach the optimal loss range on LLaMA-350M, whereas Mano reduces loss beyond it.The authors attribute the difference to traditional retractions constraining expressivity and exploration.
  • Ablation Studies: Static column-wise normalization performs comparably on LLaMA-350M but significantly worse on LLaMA-1B, indicating poor model-wise scaling.This ablation motivates Mano’s rotational Oblique manifold scheme.

6. Conclusion

The paper reports that Mano improves LLM training efficiency and performance, while identifying empirical and theoretical limits that remain for future work.

  • 6. Conclusion: The study’s empirical scope is constrained by available computational resources, limiting additional hyperparameter and larger-model over-training experiments.These experiments are left as future work.
  • 6. Conclusion: The convergence analysis covers a simplified Mano variant, so extending theory to momentum dynamics and broader optimization regimes remains future work.The current analysis excludes aspects of the full optimizer.
  • 6. Conclusion: Mano outperforms AdamW and Muon in LLM training, with lower computational overhead than Muon and reduced memory usage compared to AdamW.The conclusion frames Mano as a geometrically aware optimizer for efficient LLM training.

B.1. Hyperparameters

The experiments use standardized training configurations across LLaMA and Qwen model scales, with specified schedules, optimization settings, and extended runs for selected models.

  • B.1. Hyperparameters: Experiments use a cosine-decay learning-rate scheduler with a minimum learning-rate ratio of 0.1, weight decay of 0.1, and gradient clipping at 1.0.LLaMA models use the T5 tokenizer, while Qwen3 models use the generati… tokenizer.
  • B.1. Hyperparameters: Training configurations vary across LLaMA and Qwen model scales, including architecture, sequence length, learning rate, batch size, accumulation steps, and schedule.The configuration table covers the principal scale-dependent experimental settings.
  • B.1. Hyperparameters: Over-trained LLaMA-130M and LLaMA-350M models use a 10B-token corpus and 36000 training iterations, with other hyperparameters unchanged.This extends the training schedule for those two models.

B.2. Additional Empirical Designs for Mano

Additional designs examine optimizer choices, parameter subsets, learning-rate comparability, tensor-general Mano updates, and computational or geometric alternatives.

  • B.2. Additional Empirical Designs for Mano: Nesterov momentum may improve Mano for large models but can occasionally degrade performance for smaller models, with a minor overall effect on ultimate training outcomes.The paper therefore evaluates Nesterov-style momentum as a model-scale-sensitive design choice.
  • B.2. Additional Empirical Designs for Mano: Mano uses AdamW for input and output parameters and one-dimensional biases, following Muon’s implementation rather than applying manifold updates to every parameter.The rotational manifold scheme can otherwise be extended to arbitrary-dimensional parameters.
  • B.2. Additional Empirical Designs for Mano: The experiments constrain parameter-update RMS magnitudes to AdamW’s range of 0.2 to 0.4 to compare optimizers at similar effective step sizes.All models use a uniform learning-rate schedule while this normalization controls learning-rate differences.
  • B.2. Additional Empirical Designs for Mano: For an order-d tensor, Mano applies tangent projection and manifold normalization to the t mod d-th dimension of parameters and updates at each step.The general algorithm combines momentum, dimension-wise normalization, and the resulting parameter update.
  • B.2. Additional Empirical Designs for Mano: Mano regulates update magnitudes geometrically on an Oblique manifold, unlike Adafactor’s EMA-based second-moment normalization and spectral optimizers’ matrix-wide preconditioning.Its implementation uses vector-based operations rather than extensive spectral preconditioning.
  • B.2. Additional Empirical Designs for Mano: Hyperball and SSO also integrate manifold constraints into LLM pretraining, respectively regulating effective step sizes and weight norms or constraining weights and updates to a spectral sphere.These methods provide nearby manifold-based comparison points.

E. Proofs

The proof analyzes a simplified Mano update by projecting gradients onto a tangent space, normalizing the tangent direction, and studying its alignment with the gradient.

  • E. Proofs: The convergence proof considers a simplified Mano rule without momentum and with the Oblique manifold fixed at dimension 0.The fixed dimension has size m.
  • E. Proofs: The simplified update normalizes parameters, removes the gradient component parallel to the normalized parameters, normalizes the tangent vector, and updates weights along it.The displayed rule uses η√m as the tangent-update scale.
  • E. Proofs: The proof defines dimension-wise inner products and norms to express normalization and tangent projection along a selected tensor dimension.These operations provide the geometric quantities used by the update rule.
  • E. Proofs: The gradient–normalized-tangent inner product approaches zero when the full gradient vanishes or the gradient is perfectly parallel to the weight vector.The analysis assumes gradients are not perfectly aligned with weights when a tangential component exists.

E.1. Deterministic Setting

The analysis derives Mano’s convergence behavior under deterministic gradients and then extends the argument to stochastic gradients using an unbiased-noise assumption.

  • Deterministic Setting: Under L-smoothness, the deterministic analysis applies the update rule through the smoothness inequality and substitutes the momentum-gradient inner product.The derivation uses the tangential component bound to relate the inner product to the gradient norm.
  • Deterministic Setting: The deterministic result bounds the minimum gradient norm over T iterations by O(Lm/T).
  • Stochastic Setting: The stochastic extension assumes zero-mean subsampling noise, making the stochastic gradient an unbiased estimator of the true gradient.This is expressed as E[ξk] = 0 and equivalently Eξk[˜gk] = gk.
  • Stochastic Setting: Lemma 3 extends the tangential-component argument to the stochastic setting under the unbiased-gradient condition.
  • Stochastic Setting: The proof of Theorem 1 combines the smoothness descent argument with the stochastic inner product and concludes after applying Lemma 3.
Loading 2601.23000v1…