Matrix Optimization

Follow this topic Ask about this topic

Momentum

3 papers in the last four weeks, against 1 the four weeks before. 0.0% of all new papers.

Jul 13Week of Sep 28

Latest papers 48

Sep 29, 2026cs.LG

Normalize-Then-Precondition: A Hierarchical Approach to Marginal Scale and Interaction Geometry for LLM Training

Matrix optimizers have emerged as a promising direction, with Muon standing out as a prominent design. Revisiting Muon through its full-Gram representation, we observe that it jointly processes marginal-scale and interaction information. This opens an alternative way to organize geometric information hierarchically, motivating the Normalize-Then-Precondition framework. Specifically, it first uses diagonal-Gram information to construct a marginally normalized update, then applies spectral preconditioning to its directional interaction geometry. Building on this framework, we develop NormPre with NormPre-G and NormPre-L adopting global and localized spectral preconditioning, grounded in spectral-norm steepest descent and a regularized formulation followed by leading mode selection, respectively. To enable large-scale training, NormPre-G uses Newton-Schulz iterations and NormPre-L employs randomized sketching to approximate the leading interaction eigenspace. Theoretically, we establish

\mathcal{O}(T^{-1/2})

convergence guarantees for simplified versions of NormPre. Across extensive pretraining experiments on GPT-2 Small, LLaMA and Qwen3, both variants consistently outperform AdamW, Muon and MANO under matched training budgets. Further efficiency and spectral analyses reveal the complementary strengths of two variants and characterize their performance-efficiency trade-off. We open-source our code through a GitHub repository at https://github.com/zx-gong/NormPre.

Zixuan Gong, Zeyu Gan, Jiaye Teng +1

Gaoling School of Artificial Intelligence Renmin University of China · School of Statistics and Management Shanghai University of Finance and Economics

Language Model Pretraining Matrix Optimization LLM Training

Sep 23, 2026cs.LG

Riemannian Structure and Optimization for a Class of Low-Parametric Orthogonal Matrices

In this paper, we are concerned with matrices formed by block-diagonal factors interleaved with fixed permutations -- a flexible family of structured matrices. This class has recently drawn interest in deep learning architectures for its balanced expressivity-efficiency trade-off, yet efficient computational strategies for working with it remain to be found. We approach this problem through Riemannian geometry and examine under what conditions this class admits a smooth manifold structure. For the practically important case of orthogonal two-factor matrices, we derive the essential Riemannian tools and propose efficient algorithms for their implementation. The algorithms leverage automatic differentiation, support parameter sharing within each factor, and avoid explicit dense matrix construction. We test them within the Riemannian optimization framework on the best matrix approximation problem and for parameter-efficient fine-tuning of large language models. Beyond the two-factor setting, we study the geometric and matrix-theoretic properties of factorizations with a larger number of block-diagonal factors.

Ali Aliev, Maxim Rakhuba

HSE University

Riemannian Optimization Matrix Optimization Matrix Factorization

Sep 15, 2026math.OC

Optimization over covariance matrices with a parameterized metric

The choice of Riemannian metric can strongly influence the convergence of gradient-based optimization over covariance matrices. Euclidean, Bures-Wasserstein and affine-invariant metrics are common choices, but their relative effectiveness depends on the objective. We introduce a two-parameter family defined by

X^{p}LX^{q}+X^{q}LX^{p}=U

, solved for

L

at each tangent vector

U

, that contains all three as exact members, at

(0,0)

,

(1,0)

and

(1,1)

, and extends past them. We treat the choice of member as a particular way of preconditioning for a given problem. To this end, we analyze the conditioning of the Riemannian Hessian at the solution. We show that it obeys a lower bound that depends on

(p,q)

only through the exponent

r=p+q

. When the Euclidean Hessian is a pure power that mixes no eigendirections, the member

p=q=r/2

attains that bound, and a closed-form criterion identifies the other members that do. We discuss ways to tune

r

for a given problem. Experiments on real covariance data confirm the predicted conditioning and the benefit of tuning

r

. A task covariance example shows a further gain from tuning the shape.

Yibang Li, Bamdev Mishra, Pratik Jawanpuria +1

Nanyang Technological University, Singapore · Microsoft India · Indian Institute of Technology Bombay, India

Riemannian Optimization Second-Order Optimization Matrix Optimization Riemannian Geometry

Aug 12, 2026cs.LG

MOON: Multi-Objective OrthoNormalized Updates for Multitask Learning

Multi-objective optimization (MOO) has demonstrated significant success in multi-task learning by mitigating task conflicts through gradient manipulation. However, most existing methods flatten model parameters into vectors and perform gradient manipulation under Euclidean geometry, thereby overlooking the matrix structure prevalent in modern architectures such as Transformers. In this paper, we show that gradient manipulation in Euclidean space does not generally yield the steepest descent direction under matrix geometry, potentially limiting optimization efficiency. Drawing from the theory of steepest descent for matrix-valued parameters, we propose MOON (Multi-Objective OrthoNormalized Updates), which performs gradient manipulation under spectral--nuclear norm geometry and uses the orthonormalized manipulated gradient for parameter updates. Theoretically, for smooth non-convex objectives, we establish convergence of the averaged Pareto-stationarity measure at rates of

\mathcal{O}(T^{-1/2})

in the deterministic setting and

\mathcal{O}(T^{-1/4})

under stochastic gradients. Empirical results across various benchmarks show that MOON consistently improves both optimization efficiency and final multi-task performance. Our code is available at https://github.com/KunlinLyu/MOON.

Shiji Zhou, Kunlin Lyu, Lei Zhang +2

Institute of Artificial Intelligence, Beihang University · Beijing Advanced Innovation Center for Future Blockchain and Privacy Computing, Beihang University · Beijing Academy of Artificial Intelligence (BAAI) +1

Neural Network Optimization Multi-Task Learning Matrix Optimization Multi-Task Optimization Multi-Objective Optimization

Aug 9, 2026cs.LG

Exact Rank-Space KL Projection for Shared-Marginal Low-Rank Factors: Application to Doubly Stochastic Clustering

We study exact Kullback--Leibler (KL) projection for low-rank factorizations whose two nonnegative factors have prescribed row marginals and a shared, learned column marginal. For arbitrary positive row marginals of equal total mass, the joint KL projection reduces exactly to a strictly convex gauge-fixed dual with only

r-1

effective variables; its Hessian is a sum of categorical covariance terms and admits

O((n+m)r)

matrix-free Hessian--vector products. The projection theorem is objective-independent. We then specialize this geometry to doubly stochastic (DS) graph learning through

W=U\operatorname{Diag}(g)^{-1}V^\top

, where row-simplex factors with a common column mass induce an exactly DS graph without materializing an

n\times n

optimization variable. Combined with observed-edge sparse fitting, a stochastic anchor-reduced manifold regularizer, and Bregman backtracking, the resulting mirror-descent method preserves exact feasibility at every accepted step. Under a nonvanishing latent-mass condition, it satisfies sufficient decrease and an

O(1/N)

mirror-stationarity bound, while strictly positive accumulation points are KKT stationary. Matched clustering experiments show competitive accuracy, feasibility residuals near numerical precision, and favorable anytime behavior without a dense learned graph.

Enliang Hu

Graph Structure Learning Low-Rank Matrix Decomposition Clustering Matrix Optimization Matrix Factorization

Aug 6, 2026math.OC

Muon on the Stiefel Manifold Admits an Exact Closed-Form Update

We study Muon, a recently proposed matrix-aware optimization method, in the context of the Stiefel manifold. This manifold consists of matrices with orthonormal columns and is ubiquitous in machine learning and scientific computing. Existing extensions of Muon to this manifold rely on heuristic, approximate, or iterative updates with varying computational efficiency. We show that the corresponding Stiefel Muon update admits an exact closed-form solution and use this result to develop Skewon, a practical algorithm for orthogonality-constrained optimization with an efficient implementation. We further establish first-order convergence guarantees for Skewon in the smooth non-convex setting.

Mikhail Solonko, Molozhavenko Alexander, Maxim Rakhuba

HSE University, Moscow, Russia.

Stiefel Manifold Optimization Riemannian Optimization Muon Optimizer Nonconvex Optimization Constrained Optimization Matrix Optimization

Aug 5, 2026cs.LG

The Loss Does Not See the Basis, but Adam Does

Gradient descent on a factored model

W = UV^\top

is implicitly biased toward low-rank solutions, while Adam, starting from the same small initialization, is not. We trace the difference to the gauge symmetry of the loss, its invariance under

(U, V) \mapsto (UQ, VQ)

. Gradient flow's low-rank mechanism is available to an optimizer only if that optimizer is gauge-equivariant, a condition necessary for the transfer but not sufficient for low-rank recovery. Gradient descent, momentum, "shared-scalar" Adam, Muon, and Shampoo satisfy it. Adam, RMSProp, and the other coordinate-wise methods do not. A structure theorem characterizes the memoryless equivariant rules as exactly the Gram-determined left preconditioners, and a transfer theorem carries gradient flow's pathwise properties to common-scalar flows. We then sort nine update rules on underdetermined matrix sensing by recovery error against the planted ground truth. A one-parameter family from coordinate-wise to shared-scalar preconditioning restores the bias monotonically, isolating anisotropy as the cause. A "spectral schedule" reconciles two opposing reports about Muon: equal-rate updates recover exactly low-rank targets but lose their edge as the spectral tail grows. In transformers, Adam separates two gauge-equivalent initializations at the first step, where the equivariant optimizers stay at float precision, and ends with the per-head invariants

W_Q^\top W_K

56% apart in relative Frobenius distance, a gap no per-head rotation can close. On two hyperspectral datasets at matched training loss, gradient descent cuts held-out error by 43-44% at the lowest sampling density, and at lower effective rank. Basis choice is therefore not a tuning detail but a decision about which interpolant the optimizer selects.

Devender Singh

Department of Computer Science Memorial University of Newfoundland

Low-Rank Matrix Decomposition Matrix Optimization Adaptive Gradient Methods Matrix Factorization

Aug 5, 2026cs.LG

MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training

Memory-efficient matrix optimizers such as Sinkhorn gradient descent remove most AdamW optimizer state for dense Transformer matrices, but direct application to Mixture-of-Experts (MoE) training is unreliable. We study this failure in a controlled 110M-parameter nanowhale DeepSeek-style MoE pretraining setting. A SAGE/Sinkhorn hybrid reduces optimizer state from 0.883GB to 0.331GB but degrades evaluation loss to 3.8265, far above the AdamW baselines observed in the same setup (3.58--3.64 across the seeds we study). We show that routed MoE expert matrices are the dominant failure point: their gradients are conditional, temporally varying, and poorly served by stateless Sinkhorn normalization. We propose MESH, a hidden-momentum Sinkhorn update for MoE experts. MESH restores a temporal first-moment signal through the gradient-buffer lifecycle, without storing the expert first moment as optimizer state. MESH is an optional block-preconditioned variant that adds a coarse neuron/block inverse-RMS multiplier. Across ablations, temporal smoothing before matrix normalization is the primary causal ingredient; block/neuron preconditioning can improve the memory-quality frontier, but is not established as universally necessary. In two additional seeds, MESH and MESH-B reduce optimizer-state memory by 62.5% and peak PyTorch CUDA allocation by about 12.6% relative to AdamW, with a modest evaluation-loss gap. Full-state diagnostic variants recover AdamW-like performance in ablations, supporting the conclusion that MoE experts need temporal smoothing, but not necessarily full coordinate-wise AdamW state.

Masato Fujitake

FA Research, Fast Accounting Co., Ltd.

Neural Network Optimization Memory-Efficient Optimization Matrix Optimization Mixture-of-Experts Models Memory-Efficient Training

Aug 4, 2026cs.LG

Joint Affine Spectral Shaping: Coupling Weight and Bias Updates Beyond Weight-Only Muon

Matrix spectral optimizers reshape weight-update spectra but usually delegate vector-valued biases to a separate optimizer. We study whether this separation is neutral. We formulate each affine layer as a joint momentum matrix

A=[M_W,αm_b]

and apply a capped regularized-inverse spectral map to the complete matrix, producing both the weight and physical bias updates. A strict five-seed ablation on a four-layer BERT-mini trained from scratch on IMDb compares exact-SVD Muon, weight-only inverse shaping, affine-probe inverse shaping, and the proposed joint regularized inverse (JRI). Weight-only inverse shaping raises validation-loss-selected test accuracy from

84.903\pm0.242\%

to

85.562\pm0.308\%

and lowers selected test loss from

0.3479

to

0.3345

. Allowing bias to alter the joint SVD while retaining an independent Adam bias update does not improve over weight-only inverse shaping. Using the transformed bias jointly raises selected test accuracy to

85.738\pm0.180\%

and lowers test loss to

0.3291

, with all five seeds improving relative to the probe baseline. During the peak-performance window, JRI preserves the eligible weight-update norm while reducing the bias-update norm from

0.02095

to

0.00301

, lowers boundary-function share from

86.58\%

to

78.97\%

, and changes the cosine between weight-induced boundary motion and explicit bias from

+0.030

to

-0.137

. An independent 22-seed replication yields

85.743\pm0.203\%

selected test accuracy. These results identify joint affine spectral allocation as a small but consistent extension to weight-only spectral optimization.

Gongyue Zhang, Honghai Liu

State Key Laboratory of Robotics and Systems, Harbin Institute of Technology, Shenzhen, Shenzhen, 518055, China.

Neural Network Optimization Muon Optimizer Matrix Optimization Spectral Methods

Jul 27, 2026cs.LG

Low-Rank Dependence Decomposition via Accelerated Symmetric Non-negative Matrix Factorization

Symmetric non-negative matrix factorization (SymNMF) recovers latent group structure from a dependence matrix, but its dense, quadratic-memory objective has confined prior work to moderate sizes. We present a large-scale GPU study of seven algorithm families (over 30 configurations) on absolute Pearson correlation and tail pairwise dependence matrices from Extreme Value Theory, two proxies for empirical risk-factor estimation on large portfolios. A trace-identity reformulation eliminates all

n \times n

intermediates, so a single GPU reaches

n \approx 10^5

and multi-node distribution scales to

n = 10^6

and beyond. Under a two-phase protocol, eleven methods converge at moderate scale; six remain efficient enough at

n = 10^5

(five AdaGrad-family plus ADMM), and five AdaGrad-family methods still converge at

n = 10^6

: AdaGrad, RMSprop, and three we introduce (Piecewise AdaGrad, Row-Stochastic SVRG, Block-SVRG AdaptGrow). At

n = 10^6

the fastest solver tracks the matrix spectrum: Block-SVRG AdaptGrow wins on the flat, ill-conditioned tail-dependence spectrum, where its lower per-iteration cost decides a long factorization, and full-batch AdaGrad wins on the dominant-low-rank correlation spectrum, where the run is short. We also benchmark spherical K-means as a hard-label baseline: cheaper when angular cluster structure is present, yet provably degenerate once the matrix collapses toward a single common factor, where the soft factorization remains necessary.

Lavinia Ghita, Dhruv Desai, Jake Goldberg +1

NVIDIA

Non-Negative Matrix Factorization Low-Rank Matrix Decomposition Matrix Optimization Adaptive Gradient Methods Stochastic Gradient Descent

Jul 24, 2026stat.ML

Graph-Based Correlation Matrix Generation: A Convex Optimization Approach

This work addresses the generation of theoretical correlation matrices with prescribed sparsity patterns associated to graph structures. We propose a novel convex optimization framework in which an initial matrix is projected onto an elliptope under a positive semidefiniteness constraint. Several numerical schemes are implemented and compared. The problem falls within the broader class of matrix completion, where off-diagonal entries corresponding to absent edges are fixed to zero and diagonal entries are fixed to one. Beyond this structural constraint, the approach offers greater flexibility than existing methods by allowing control over the mean of the off-diagonal entry distribution, enabling the generation of correlation matrices that better reflect realistic data. This procedure is not designed to yield a uniform distribution over the feasible set; rather, it provides a principled and tunable way to construct correlation matrices suitable for benchmarking statistical methods for graphical model inference. Theoretical guarantees on the existence of solutions are established, both in the general setting and under the additional mean constraint. Simulation studies illustrate the properties of the generated matrices with respect to graph structure. The methodology is applied to two real-world datasets from neuroscience and finance, and a comparison with GAN-based correlation matrix generation is provided.

Ali Fakhar, K{é}vin Polisano, Ir{è}ne Gannaz +1

UGA · Univ. Grenoble Alpes, CNRS, Grenoble INP, Inria, LJK, F-38000 Grenoble, France · G-SCOP\_GROG, G-SCOP, Grenoble INP, UGA +2

Matrix Optimization Matrix Completion

Jul 16, 2026cs.LG

Muse: Representation Geometry of Muon Beyond Normalized Momentum

Muon-style optimizers apply a polar map to matrix momentum, but their updates also depend on the representation of each parameter block before orthogonalization. We study this representation choice as a form of optimizer geometry and introduce {\method}, a family of Muon-style optimizers that shares the same momentum rule and Newton--Schulz backend across native, nearest-square, skinny, and vector representations. Each Frobenius-isometric representation induces a distinct polar steepest-descent geometry, in which the shorter matrix dimension determines the number of supported singular channels, the pullback scaling, and the constants in stochastic nonconvex convergence bounds. In a teacher--student model, curvature collapse and an isotropic Marchenko--Pastur spectral profile connect early-stage dissipation to the represented nuclear-to-squared-Frobenius norm ratio. Pretraining experiments on LLaMA2-130M and LLaMA2-600M, together with fixed-momentum diagnostics, show that balanced non-native representations can match the performance of the native representation, whereas reducing the shorter dimension weakens the scaling and singular-channel support, leading to behavior that increasingly resembles normalized momentum.

Da Chang, Qiankun Shi, Lvgang Zhang +4

1Pengcheng Laboratory · 2Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences +3

Representation Geometry Momentum Methods Muon Optimizer Matrix Optimization

Jul 10, 2026cs.LG

Graph-Regularized Low-Rank Matrix Completion by Variable Projection

We address the low-rank matrix completion problem by incorporating graph regularization into the existing Riemannian Trust-Region Matrix Completion (RTRMC) framework. The latter uses the geometry of the low-rank constraint to remodel the problem as an unconstrained optimization problem on a single Grassmann manifold. Our approach, named Graph-Regularized RTRMC (GR-RTRMC), exploits the inherent relationships between rows and columns of the matrix. By using these relationships, we aim to improve the accuracy and robustness of matrix completion, particularly in scenarios where the underlying data exhibits strong correlations between rows or columns.

Benoît Loucheur, P. -A. Absil, Michel Journée

ICTEAM Institute, UCLouvain, 1348 Louvain-la-Neuve, Belgium. · Department of Climatology, Royal Meteorological Institute of Belgium, 1180 Uccle, Belgium.

Low-Rank Approximation Matrix Optimization Matrix Completion

Jul 2, 2026cs.LG

Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials

Machine learning interatomic potentials (MLIPs) have become a hallmark of AI for scientific simulation. While efforts on new architectures and datasets have led to increasingly accurate and general models, the choice of optimizer for training has largely remained unexplored, defaulting to Adam and its variants in the community. Here, we implement and systematically compare a class of recently proposed matrix-structured optimizers, including Muon, SOAP, and the hybrid SOAP-Muon, for training NequIP and Allegro MLIP models. We find that these optimizers can substantially outperform Adam in both convergence speed and final accuracy. SOAP and SOAP-Muon emerge as robust and consistently strong methods, while Muon only provides partial gains relative to Adam. The improvements are particularly pronounced under partial force supervision. Our results indicate that optimizer choice is an overlooked yet impactful design axis for MLIPs.

Gil Harari, Yoel Zimmermann, Ola Tangen Kulseng +4

John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA, USA · Robert Bosch LLC Research and Technology Center, Watertown, MA, USA

Machine Learning Interatomic Potentials Neural Network Optimization Muon Optimizer Matrix Optimization

Jun 30, 2026math.OC

Direction-Magnitude Decomposition for Low-Rank Matrix Optimization: Faster Convergence and Saddle-to-saddle Dynamics

Low-rank matrix optimization is often carried out via the Burer-Monteiro (BM) formulation, but choosing the factorization rank

r

is delicate and can substantially slow optimization. We propose a unified framework, termed direction-magnitude decomposition (DMD), that decomposes the optimization variable to improve optimization efficiency even when the target rank is unknown. We develop two DMD-based approaches and establish their theoretical advantages on the canonical problem of matrix factorization. The first, overparameterized DMD, uses a rank

r

larger than necessary and enjoys faster convergence as

r

increases. The second, recursive DMD, is motivated by the incremental eigenpair learning, or saddle-to-saddle, behavior of overparameterized DMD. It achieves lower memory and computational costs, complementing overparameterized DMD. Both approaches are exponentially faster than gradient descent applied to the BM formulation. Numerical experiments on matrix factorization, sensing, and completion corroborate our theoretical findings and demonstrate the practical effectiveness of DMD.

Yudong Wei, Liang Zhang, Bingcong Li +1

Department of Mathematics, ETH Zurich, Zurich, Switzerland · Department of Computer Science, ETH Zurich, Zurich, Switzerland

Low-Rank Matrix Decomposition Nonconvex Optimization Matrix Optimization Matrix Factorization Gradient Descent Dynamics

Jun 29, 2026cs.LG

ITSPACE: Monotone Gaussian Optimal Transport Updates

Covariance matrices serve as compact descriptors of feature distributions in many machine-learning pipelines, including domain adaptation and Gaussian embeddings. Under a centered Gaussian approximation, the unregularized Wasserstein-2 optimal-transport (OT) discrepancy admits a closed form on covariances given by the Bures-Wasserstein (BW) objective on the symmetric positive definite (SPD) cone. We propose ITSPACE (Iterative Transport for Stable Proximal Alignment of Covariance Embeddings), a proximal majorization-minimization method that directly optimizes this exact BW objective through closed-form updates in a square-root factorization. In exact arithmetic, each iteration satisfies a sufficient-decrease inequality for the BW objective; under inexact polar computations, we provide an explicit certificate-gap bound controlling deviations from exact descent. The resulting iterations preserve PSD structure by construction and naturally support rank-restricted factors, making ITSPACE well-suited as a lightweight inner-loop primitive in settings where adaptation must be performed from unlabeled target batches under strict step and compute budgets. Across real-world covariance-alignment benchmarks, ITSPACE reaches low-BW-gap solutions substantially faster than BW-gradient descent, methods based on other covariance geometries, and entropically regularized sample-OT baselines.

Woojoo Na, Jennifer Dy

Department of Computer Engineering, Northeastern University, Boston, USA.

Matrix Optimization Optimal Transport

Jun 29, 2026cs.LG

Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics

Matrix factorization (i.e., problems of the form

\min_{\mathbf{P},\mathbf{Q}} \|\mathbf{M}^\star - \mathbf{P}^\top\mathbf{Q}\|_\mathrm{F}^2

) is a minimal learning problem that exhibits both nonlinear parameter dynamics and representation learning. In this setting, we study how parameter trajectories under the Muon optimizer differ from those of gradient descent. We identify three main dynamical differences: 1) Muon avoids the slow saddle-to-saddle dynamics from small initialization. Muon instead learns all the top modes of

\mathbf{M}^\star

at the same rate, with the smaller modes converging first. 2) Muon remains stable even when the learning rate exceeds the critical threshold set by the local loss sharpness. This frees the learning rate from the condition number of the problem, enabling rapid convergence via exponential learning rate annealing. 3) Once the weights are aligned with each other and the target, Muon flow conserves the matrix quantity

\sqrt{\mathbf{P}^\top \mathbf{P}}-\sqrt{\mathbf{Q}^\top \mathbf{Q}}

, while gradient flow is known to conserve the matrix

\mathbf{P}^\top\mathbf{P} - \mathbf{Q}^\top\mathbf{Q}

. Despite having distinct conserved quantities, both optimizers find the so-called \textit{balanced} solution from vanishing initialization. When training from small random initialization, the weights spontaneously align early in training. We derive the alignment rates in simple settings and show that they predict the empirical alignment rates in general. Finally, we exploit structural properties of Muon to construct a learning rate schedule that achieves near-perfect alignment in only two optimization steps.

Mark Rhee, Jamie Simon, Dhruva Karkada

UC Berkeley · Imbue

Muon Optimizer Matrix Optimization Matrix Factorization Gradient Descent Dynamics

Jun 26, 2026cs.LG

Aurora: A Leverage-Aware Spectral Optimizer

We show that for tall matrix parameters, like projection matrices in the MLP layers, the Muon update can have row norms that are arbitrarily non-uniform. This can lead to a self-reinforcing feedback loop whereby neurons receive persistently small updates and eventually do not contribute meaningfully to network outputs. This problem is effectively mitigated by an additional row normalization step, but current methods do this in a way that moves the Muon update geometry away from the polar factor of the momentum matrix, which we find is undesirable. We propose Aurora, an optimizer that enforces row-uniformity of matrix parameter updates while respecting Muon's polar factor geometry. Aurora outperforms Muon in our pre-training experiments and, when combined with existing methods, achieves state-of-the-art performance among spectral optimizers on the optimizer track of the modded-nanoGPT speedrun. Additionally, we find that Aurora's empirical gains over Muon scale with the MLP expansion factor, suggesting that Aurora may allow for effective training of very wide MLP layers.

Alec Dewulf, Dhruv Pai, Li Yang +2

Tilde Research

Neural Network Optimization Muon Optimizer Matrix Optimization Spectral Methods Neural Network Training Dynamics

Jun 17, 2026cs.LG

Spectral DPPs via NEPv: A Scalable Continuous Relaxation of Determinantal MAP for Diversity-Aware Data Selection

Selecting a small, diverse, high-quality subset from a massive pool of candidates is a recurring primitive in modern machine learning -- data curation and coreset selection for training and fine-tuning large models, active-learning batch acquisition, prompt and exemplar selection for in-context learning, retrieval diversification, and experimental design. Determinantal Point Processes (\DPP s) give a principled, well-calibrated notion of diversity for this task, but their \emph{MAP} objective -- pick a size-

k

subset

S

maximizing

\logdet(L_S)

-- is NP-hard, and the standard greedy and sampling algorithms scale superlinearly in the ground-set size

n

. This cost is prohibitive precisely in the data-centric regime where diversity matters most, where

n

ranges over millions to billions of candidate examples, features, or embeddings. We recast \DPP-MAP as a continuous optimization problem over the Stiefel manifold, and show that its first-order optimality conditions form a \emph{Nonlinear Eigenvalue Problem with eigenvector dependency} (\NEPv) of a previously unstudied form. This \NEPv\ admits a self-consistent field (\SCF) iteration with a spectral-gap-based local contraction guarantee, giving a principled iterative solver where the diversity objective drives an eigenvector-dependent operator. The resulting algorithm, \OurMethod, requires only matrix-vector products with the kernel and runs in time

O\!\big((ndk+nk^2)\,t\big)

for a small number of iterations

t

, scaling near-linearly in

n

and integrating directly with low-rank and feature-map kernels common in ML. This paper focuses on the relaxation, solver, and scaling analysis; full real-data benchmarking is left to a planned empirical study.

Richard Yi Da Xu

Hong Kong Baptist University · TadReamk Limited

Determinantal Point Process Training Data Selection Matrix Optimization Spectral Methods

Jun 14, 2026math.OC

Schattor: Schatten-family methods for deep learning optimization

Modern deep learning optimization features heterogeneous parameter structures, noisy gradients, and highly nonconvex landscapes, posing significant challenges for both algorithm design and theoretical analysis. Motivated by the limitations of SGD and the success of adaptive optimizers, we propose {\it Schattor}, a family of adaptive first-order methods based on Schatten norms. Schattor unifies SGD and the recently proposed matrix-variate adaptive optimizer Muon within a single Schatten-norm-based framework. We establish dimension-free stationarity guarantees for methods in the Schattor family for stochastic matrix optimization problems via a novel matrix martingale moment bound. We also develop multi-block extensions that adaptively balance block-wise optimization progress and prove dimension-free stationarity guarantees in this more general setting.

Bohao Ma, Junyu Zhang, Chuan He

School of Data Science, The Chinese University of Hong Kong (Shenzhen), Shenzhen, Guangdong, China · Department of Industrial Systems Engineering and Management, National University of Singapore, Singapore · Department of Mathematics, Linköping University, Sweden

Deep Learning Optimization Stochastic Optimization Convergence Matrix Optimization Adaptive Gradient Methods

Jun 12, 2026math.OC

Free Heavy-Tailed Lunch for Muon: A Theoretical Justification of Empirical Success

Non-Euclidean optimisation methods with matrix-valued updates, such as Muon and Scion, have recently shown strong empirical performance for training Transformer models, yet their theoretical advantages over Euclidean methods remain poorly understood. We address this gap in the heavy-tailed non-convex regime, where stochastic gradients have bounded

p

-th central moments,

p \in (1,2]

. We show that certain non-Euclidean methods achieve optimal sample complexity under stronger stationarity measures, while Euclidean methods incur additional dimension-dependent costs. As a consequence, for

m \times n

matrices, Muon finds an

\varepsilon

-stationary point in nuclear norm within

\mathcal{O}\left(\min\{m, n\} \frac{Δ_1 L}{\varepsilon^2} \left(\frac σ\varepsilon \right)^{\frac p {p-1}}\right)

samples, absorbing heavy-tailed noise without extra dimension dependence, unlike Euclidean methods. We further prove this sample complexity, including its dimension dependence, is optimal for all first-order methods under nuclear-norm stationarity. Experiments on large language models support our theory. Surprisingly, our results suggest that other Schatten geometries beyond the spectral geometry of Muon can perform competitively in certain settings.

Florian Hübler, Thomas Pethick, Suvrit Sra

Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, Technical University of Munich, Germany · Munich Center for Machine Learning (MCML)

Stochastic Optimization Nonconvex Stochastic Optimization Heavy-Tailed Noise Muon Optimizer Stochastic Optimization Convergence Matrix Optimization+1

Jun 12, 2026cs.LG

Zeta: Dual Whitening for Matrix Optimization via Coordinate-Adaptive Preconditioning

Large-scale neural network training increasingly relies on matrix-aware optimizers that exploit the structure of weight parameters beyond element-wise adaptation. However, existing matrix-aware methods such as Muon have an underappreciated vulnerability: their core operation, Newton-Schulz iteration, depends critically on input conditioning, yet the raw momentum matrices exhibit severe coordinate-wise scale heterogeneity. In this paper, we first verify this scale heterogeneity through a chi-square uniformity test, showing that intra-matrix scale imbalance is prevalent across Transformer layers and that coordinate whitening effectively corrects it. Motivated by this finding, we propose Zeta, a dual whitening optimizer that applies coordinate whitening and spectral whitening in a strictly ordered pipeline. The ordering is not a tunable choice but follows from a mathematical dependency: coordinate whitening establishes the statistical isotropy that spectral whitening requires to function reliably. We further prove that this dual pipeline strictly reduces orthogonalization error relative to pure spectral methods by improving the condition number of the input. Empirically, Zeta matches or surpasses strong baselines across language modeling (0.6B to 8B parameters), mixture-of-experts architectures, and vision tasks, demonstrating that resolving scale imbalance before orthogonalization leads to faster convergence and better generalization. Code is available at https://github.com/AIGCodeOS/aigcode_zeta_optimizer.

Kaiwen Chen, Shuhai Zhang, Zimo Liu +7

South China University of Technology · AIGCode · Hong Kong Baptist University

Neural Network Optimization Matrix Optimization

Jun 11, 2026math.OC

Scalable Deep Unfolding of Conic Optimizers

Deep unfolding (DU) accelerates iterative optimizers by introducing learnable components and training them through unrolled iterations, but extending DU to the large-scale semidefinite programs (SDPs) common in robotics has remained limited. Unrolling a full-update conic solver such as COSMO exposes two obstacles that prior work on learned conic solvers has not: backpropagating through the per-iteration linear-system solve incurs memory quadratic in the problem size once the coefficient matrix is formed explicitly, and backpropagating through the positive semidefinite (PSD) cone projection becomes numerically unstable when eigenvalues coincide. We address the first obstacle with a matrix-free implicit differentiation rule that operates entirely through matrix-vector products, reducing memory from

O(n^2)

to

O(n)

and enabling backpropagation at scales where direct factorization runs out of memory. We address the second with a backward rule based on the Dalečkii--Krein representation of the Fréchet derivative, which remains well-defined under repeated eigenvalues. Together these make it possible to learn lightweight hyperparameter policies and warm-starts for a full-update conic solver. We evaluate on nonlinear covariance steering problems solved via sequential convex programming (SCP), as well as standalone SDPs and second-order cone programs ranging from max-cut and Lovász

\vartheta

SDPs to robust estimation and control problems. The learned policies outperform state-of-the-art solvers across all problems, and can provide up to a 50

\times

speedup depending on the class. When used as a subroutine in SCP, the learned approach delivers over a 30

\times

speedup compared to COSMO.

Alex Oshin, Rahul Vodeb Ghosh, Evangelos A. Theodorou

Georgia Institute of Technology

Semidefinite Programming Algorithm Unrolling Constrained Optimization Matrix Optimization Differentiable Optimization Deep Unfolding

Jun 7, 2026cs.LG

The Spectral Dynamics and Noise Geometry of Muon

Muon replaces a matrix gradient

G=UΣV^\top

by its polar factor

UV^\top

. This keeps the singular directions selected by the gradient, but makes the update spectrum flat. We study the optimization bias created by this operation. Under explicit alignment assumptions, we prove that the polar update is the one-step entropy-maximizing choice among bounded updates that use the gradient singular directions and do not adapt to the current weight spectrum. In an underdetermined regression model, we derive exact singular-value dynamics for continuous-time Muon and identify a measurement-dependent condition under which the normalized spectrum moves toward equal nonzero singular values. This geometry also rules out a common low-rank interpretation: at fixed Frobenius norm, Muon's distinguished state has a flat spectrum, whereas nuclear-norm minimization favors spectral concentration. Controlled matrix-sensing experiments separate the effect from simple gradient rescaling, show that norm-matched gradient descent does not reproduce Muon, and recover the predicted flattening trend across broad ablations. In small NanoGPT pretraining, Muon preserves stable rank, has a broad learning-rate plateau, and improves validation loss relative to AdamW; in a matched small-ViT control, the ranking reverses. The resulting picture is regime-dependent: Muon is not universally superior, but its flat-spectrum bias can help when many spectral directions need to remain active.

Pierfrancesco Beneventano, Mahmoud Abdelmoneum, Tomaso Poggio

Massachusetts Institute of Technology

Muon Optimizer Matrix Optimization Spectral Methods Implicit Regularization Gradient Descent Dynamics

Jun 6, 2026cs.LG

On solving symmetric multi-type orthogonal non-negative matrix tri-factorization problem

We study the symmetric multi-type orthogonal non-negative matrix tri-factorization problem, where several symmetric non-negative matrices are simultaneously approximated by factors of the form

GS_{i}G^{\top}

, with a shared non-negative and orthogonal factor

G

. This model is motivated by clustering and network analysis, where non-negativity improves interpretability and orthogonality gives a natural assignment-type structure to the latent factor. Since the resulting optimization problem is highly non-convex, we develop two heuristic algorithms for computing high-quality local solutions. The first one is a fixed point method derived from the Karush-Kuhn-Tucker conditions after adding a penalty term for the orthogonality constraint. The second one is a three-stage ADAM-based method that combines non-negativity-preserving optimization, orthogonalization, and restricted ADAM refinement on the feasible set. We evaluate both methods on synthetic data, including noisy instances, and on citation network benchmarks. The synthetic experiments show that both algorithms recover factorizations close to the optimum and remain stable under noise. On real networks, the learned embeddings are competitive with or better than standard baselines such as SVD, node2vec, and classical link prediction heuristics in link prediction, node clustering, and node classification tasks.

Rok Hribar, Gregor Papa, Janez Povh +1

Computer Systems Department, Jožef Stefan Institute, Jamova cesta 39, Ljubljana, 1000, Slovenia. · Laboratory for Engineering Design, Faculty of Mechanical Engineering, University of Ljubljana, Aškerčeva cesta 6, Ljubljana, 1000, Slovenia. · Rudolfovo – Science and Technology Centre Novo mesto, Podbreznik 15, Novo mesto, 8000, Slovenia. +1

Non-Negative Matrix Factorization Graph Clustering Matrix Optimization Graph Embedding

Jun 1, 2026cs.LG

A Nonmonotone Gradient-Based Algorithm for Symmetric Nonnegative Matrix Factorization and Graph Clustering

Symmetric nonnegative matrix factorization (Symmetric NMF) approximates a matrix as

WW^T

with nonnegative rectangular factor

W

. It has broad applications in graph clustering and machine learning. In contrast to the NMF, projected gradient methods for the symmetric problem had been associated with slow convergence. To address this, we introduce SNMPBB, the first adaptation of nonmonotone projected Barzilai-Borwein methods to Symmetric NMF, demonstrating that gradient algorithms are significantly more effective than previously understood. We further extend SNMPBB to graph clustering using the graph Laplacian regularization (Graph-SNMPBB) and to large problems with low-rank approximations (LAI-SNMPBB). For all variants we prove global convergence to first-order stationary points and also that Barzilai-Borwein curvature information is preserved with randomized approximations. On synthetic data, SNMPBB achieves 6 times speedup over the alternative SymANLS for similar residuals, with advantages growing at higher ranks. Across six real-world clustering benchmarks, Graph-SNMPBB matches or exceeds SymANLS accuracy. Lastly, LAI-SNMPBB outperforms state-of-the-art LAI-SymPGNCG on 34 SuiteSparse matrices in both runtime and residual quality.

Ryan Swart, Johannes Brust

School of Mathematical and Statistical Sciences, Arizona State University, Tempe, AZ

Non-Negative Matrix Factorization Graph Clustering Matrix Optimization

Jun 1, 2026cs.LG

Riemannian Gradient Descent for Low-Rank Architectures

We explore Riemannian optimization techniques for rank-factored matrix parameters, targeting contemporary deep learning applications. We examine ten points in the algorithm design space: two geometries for rank-

r

matrices, three geometries for rank-

r

partial isometries, and block-matrix variants of these five, where factors are shared across block-rows and block-columns. We apply our methods to the multihead attention parameters in small language models. After tuning learning rates, our methods do not conclusively outperform an AdamW baseline. Our implementations are available online.

Nicholas Knight

NVIDIA

Riemannian Optimization Low-Rank Approximation Matrix Optimization Riemannian Geometry Multi-Head Attention

May 31, 2026cs.LG

Riemannian Optimization for Hadamard Products of Low-Rank Matrices

The elementwise Hadamard product of two low-rank matrices provides a parameter-efficient model for data with multiplicative structure, but its modeling is challenging due to the presence of additional symmetries under coupled row/column scalings between the two factors. In order to leverage the geometry of the space, we formulate the learning of such matrices as optimization on a Riemannian quotient manifold. We propose a novel block-diagonal Riemannian metric derived from the pullback of the Frobenius inner product. The metric is shown to be invariant under these symmetries. We develop a Riemannian gradient descent algorithm that uses a tuning-free Gauss--Newton step size and scales linearly in the number of observed entries per iteration. The versatile framework of Riemannian quotient optimization enables both first-order and second-order Riemannian methods, the latter through a closed-form connection and the Riemannian Hessian. Experiments on real and synthetic datasets illustrate the efficacy of our proposed Riemannian approach.

Pratik Jawanpuria, Ankish Chandresh, Bamdev Mishra

Centre for Machine Intelligence and Data Science, Indian Institute of Technology Bombay, India · Microsoft India

Riemannian Optimization Low-Rank Approximation Matrix Optimization Riemannian Geometry Matrix Factorization

May 27, 2026math.OC

Manifold-based Algorithms for the Hadamard Decomposition

Given a matrix

X

, and two ranks

r_1

and

r_2

, the Hadamard decomposition (HD) looks for two low-rank matrices,

X_1

of rank

r_1

and

X_2

of rank

r_2

, both of the same size as

X

, such that

X\approx X_1\circ X_2

, where

\circ

is the Hadamard (element-wise) product. In most cases, HD is more expressive than standard low-rank approximations such as the truncated singular value decomposition (TSVD), as it can represent higher-rank matrices with the same number of parameters; this is because the rank of

X_1 \circ X_2

is generically equal to

r_1 r_2

. In this paper, we first present some theoretical insights for HD, in particular a useful reformulation

X\approx WH^\top

where

W

and

H

have

r_1 r_2

columns and belong to certain manifolds. These allow us to develop three new algorithms for computing HD. The first one uses the representation

X\approx X_1\circ X_2

and relies on the Manopt toolbox. The other two rely on the reformulation

X\approx WH^\top

: one is a block projected gradient method, and the other is a manifold-based gradient descent algorithm that does not require projection onto the feasible set. The last two algorithms are particularly effective for handling large sparse data. We also propose new initializations that allow us to improve the accuracy of the HD. We compare our algorithms and initialization strategies with the TSVD and with the state of the art. Numerical results show that the new methods are efficient and competitive on both synthetic and real data.

Nicolas Gillis, Subhayan Saha, Stefano Sicilia +1

Department of Mathematics and Operational Research, University of Mons, Mons, Belgium. · SS is a member of the Gruppo Nazionale Calcolo Scientifico-Istituto Nazionale di Alta Matematica (GNCS-INdAM).

Low-Rank Matrix Decomposition Low-Rank Approximation Matrix Optimization Matrix Factorization

May 27, 2026math.OC

Implicit Regularization in Perturbed Deep Matrix Factorization: Spectral Conditions and Stability

This paper studies the stability of low-rank implicit regularization in perturbed deep matrix factorization, where the target matrix is corrupted by a noise matrix. We first derive sufficient spectral conditions under which gradient descent exhibits a low-rank phase in the noiseless setting. These conditions show how the target spectrum, initialization, and step size jointly determine the existence of a nonempty low-rank interval. We then analyze the perturbed gradient descent dynamics, proving convergence guarantees and quantifying how the perturbation affects iteration complexity and eigenvalue recovery. Finally, we show that the low-rank phase persists under perturbation, with explicit dependence on the perturbation size. Numerical experiments support the theoretical findings.

Jingzhe Wang, Hung-Hsu Chou

Department of Informatics and Networked Systems, University of Pittsburgh · Department of Mathematics, University of Pittsburgh

Algorithmic Stability Low-Rank Matrix Decomposition Low-Rank Approximation Matrix Optimization Implicit Regularization Matrix Factorization+1