Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

PRISM: Structured Optimization via Anisotropic Spectral Shaping

Yujie Yang

TL;DR

PRISM introduces innovation-augmented spectral shaping to incorporate partial second-order information into the Muon-style spectral descent framework. By forming an innovation-augmented momentum and applying a polar decomposition, PRISM yields a preconditioner $P_t^{PRISM} = (M_t^T M_t + \gamma^2 D_t^T D_t)^{-1/2}$, enabling anisotropic spectral gains that damp high-variance directions while preserving signal directions. This results in a quasi-second-order method with negligible overhead and zero extra memory, improving convergence and stability on large-scale non-convex losses. Empirical results on causal language model pre-training show PRISM outperforms Muon and AdamW, with robustness across damping parameters and wider stable learning-rate ranges.

Abstract

We propose PRISM, an optimizer that enhances first-order spectral descent methods like Muon with partial second-order information. It constructs an efficient, low-rank quasi-second-order preconditioner via innovation-augmented polar decomposition. This mechanism enables PRISM to perform anisotropic spectral shaping, which adaptively suppresses updates in high-variance subspaces while preserving update strength in signal-dominated directions. Crucially, this is achieved with minimal computational overhead and zero additional memory compared to first-order baselines. PRISM demonstrates a practical strategy for integrating curvature-adaptive properties into the spectral optimization paradigm.

PRISM: Structured Optimization via Anisotropic Spectral Shaping

TL;DR

PRISM introduces innovation-augmented spectral shaping to incorporate partial second-order information into the Muon-style spectral descent framework. By forming an innovation-augmented momentum and applying a polar decomposition, PRISM yields a preconditioner , enabling anisotropic spectral gains that damp high-variance directions while preserving signal directions. This results in a quasi-second-order method with negligible overhead and zero extra memory, improving convergence and stability on large-scale non-convex losses. Empirical results on causal language model pre-training show PRISM outperforms Muon and AdamW, with robustness across damping parameters and wider stable learning-rate ranges.

Abstract

We propose PRISM, an optimizer that enhances first-order spectral descent methods like Muon with partial second-order information. It constructs an efficient, low-rank quasi-second-order preconditioner via innovation-augmented polar decomposition. This mechanism enables PRISM to perform anisotropic spectral shaping, which adaptively suppresses updates in high-variance subspaces while preserving update strength in signal-dominated directions. Crucially, this is achieved with minimal computational overhead and zero additional memory compared to first-order baselines. PRISM demonstrates a practical strategy for integrating curvature-adaptive properties into the spectral optimization paradigm.
Paper Structure (36 sections, 16 equations, 5 figures, 2 tables, 1 algorithm)

This paper contains 36 sections, 16 equations, 5 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Adaptive Optimization
  4. Structured Optimization
  5. Methodology
  6. A Preconditioning Perspective for Muon Optimizer
  7. Innovation-Augmented Polar Decomposition
  8. Anisotropic Spectral Shaping
  9. Eigenvalue Decomposition: Signal vs. Noise Energy
  10. Signal Energy ($\|M_t v_k\|^2$)
  11. Noise Energy ($\|D_t v_k\|^2$)
  12. Anisotropic Spectral Gain
  13. Analysis of the Spectral Gain
  14. Case I: High-SNR Subspace ($\text{SNR}_k \gg 1$)
  15. Case II: Low-SNR Subspace ($\text{SNR}_k \ll 1$)
  16. ...and 21 more sections

Figures (5)

  • Figure 1: Comparison of training loss curves for AdamW, Muon, and PRISM.
  • Figure 2: Training loss comparison between PRISM and Muon with Tikhonov damping.
  • Figure 3: Comparison of training loss trajectories for Muon and PRISM under aggressive learning rates.
  • Figure 4: Frobenius norm of the parameters $\|O_t\|_F$ over training steps.
  • Figure 5: The relationship between the empirical SNR and the spectral gain ($\rho_k$) for subspaces.

Theorems & Definitions (1)

  • Definition 3.1