Table of Contents
Fetching ...

Revisiting Orbital Minimization Method for Neural Operator Decomposition

J. Jon Ryu, Samuel Zhou, Gregory W. Wornell

TL;DR

This work revisits the orbital minimization method (OMM) to learn the top-$k$ eigensubspace of positive semidefinite operators, providing a simple linear-algebraic proof of its consistency and extending it from finite matrices to operators on Hilbert spaces. It introduces higher-order OMM variants, nesting schemes for ordered eigenvectors, and connections to Sanger's streaming PCA, enabling neural networks to decompose PSD operators without explicit orthogonalization. The authors validate the approach across diverse tasks, including Laplacian-based representation learning for reinforcement learning, Schrödinger-equation solvers, and self-supervised image representation learning, showing competitive performance and practical insights such as stability issues and spectrum-shifting tricks. The results illustrate that classical numerical methods can yield principled, scalable tools for modern spectral learning in machine learning and scientific computing, with potential extensions to quantum and graph-structure problems.

Abstract

Spectral decomposition of linear operators plays a central role in many areas of machine learning and scientific computing. Recent work has explored training neural networks to approximate eigenfunctions of such operators, enabling scalable approaches to representation learning, dynamical systems, and partial differential equations (PDEs). In this paper, we revisit a classical optimization framework from the computational physics literature known as the \emph{orbital minimization method} (OMM), originally proposed in the 1990s for solving eigenvalue problems in computational chemistry. We provide a simple linear-algebraic proof of the consistency of the OMM objective, and reveal connections between this method and several ideas that have appeared independently across different domains. Our primary goal is to justify its broader applicability in modern learning pipelines. We adapt this framework to train neural networks to decompose positive semidefinite operators, and demonstrate its practical advantages across a range of benchmark tasks. Our results highlight how revisiting classical numerical methods through the lens of modern theory and computation can provide not only a principled approach for deploying neural networks in numerical simulation, but also effective and scalable tools for machine learning.

Revisiting Orbital Minimization Method for Neural Operator Decomposition

TL;DR

This work revisits the orbital minimization method (OMM) to learn the top- eigensubspace of positive semidefinite operators, providing a simple linear-algebraic proof of its consistency and extending it from finite matrices to operators on Hilbert spaces. It introduces higher-order OMM variants, nesting schemes for ordered eigenvectors, and connections to Sanger's streaming PCA, enabling neural networks to decompose PSD operators without explicit orthogonalization. The authors validate the approach across diverse tasks, including Laplacian-based representation learning for reinforcement learning, Schrödinger-equation solvers, and self-supervised image representation learning, showing competitive performance and practical insights such as stability issues and spectrum-shifting tricks. The results illustrate that classical numerical methods can yield principled, scalable tools for modern spectral learning in machine learning and scientific computing, with potential extensions to quantum and graph-structure problems.

Abstract

Spectral decomposition of linear operators plays a central role in many areas of machine learning and scientific computing. Recent work has explored training neural networks to approximate eigenfunctions of such operators, enabling scalable approaches to representation learning, dynamical systems, and partial differential equations (PDEs). In this paper, we revisit a classical optimization framework from the computational physics literature known as the \emph{orbital minimization method} (OMM), originally proposed in the 1990s for solving eigenvalue problems in computational chemistry. We provide a simple linear-algebraic proof of the consistency of the OMM objective, and reveal connections between this method and several ideas that have appeared independently across different domains. Our primary goal is to justify its broader applicability in modern learning pipelines. We adapt this framework to train neural networks to decompose positive semidefinite operators, and demonstrate its practical advantages across a range of benchmark tasks. Our results highlight how revisiting classical numerical methods through the lens of modern theory and computation can provide not only a principled approach for deploying neural networks in numerical simulation, but also effective and scalable tools for machine learning.
Paper Structure (38 sections, 2 theorems, 34 equations, 8 figures, 3 tables)

This paper contains 38 sections, 2 theorems, 34 equations, 8 figures, 3 tables.

Key Result

Theorem 1

Let $r\le d$ denote the rank of a PSD matrix $\mathsf{A}$ with EVD $\mathsf{A}=\sum_{i=1}^r \lambda_i {\bf w}_i{\bf w}_i^\intercal$. Then, In particular, its minimizer $\mathsf{V}^\star$ satisfies $\mathsf{V}_{1:k\wedge r}^\star(\mathsf{V}_{1:k\wedge r}^\star)^\intercal=\mathsf{W}_{1:k \wedge r}\mathsf{W}_{1:k\wedge r}^\intercal$.

Figures (8)

  • Figure 1: Summary of experimental results of the Laplacian representation learning. On the right panel, we plot the top-11 eigenvalues of the Laplacians. It shows that the hard instances (i.e., GridMaze-{26,32} and GridRoom{32,64}) have very small spectral gaps (marked with x). Error bar indicates one standard deviation.
  • Figure 2: Visualization of learned eigenfunctions for the 2D hydrogen atom. OMM$_{\mathsf{seq}}$ performs as well as LoRA$_{\mathsf{seq}}$, which is also known as NeuralSVD Ryu--Xu--Erol--Bu--Zheng--Wornell2024.
  • Figure 3: Grid environments for Laplacian representation learning.
  • Figure 4: Summary of the 2D hydrogen experiment in Figure \ref{['fig:hydrogen']}. The shaded region indicates $\pm$ standard deviations.
  • Figure 5: Visualization of the first 15 of learned eigenfunctions with OMM$_{\mathsf{seq}}$ on the 2D harmonic oscillator. The first row shows the raw learned parametric eigenfunctions. The second row presents the eigenfunctions aligned to the ground-truth degenerate subspaces via the orthogonal Procrustes procedure as instructed in Ryu--Xu--Erol--Bu--Zheng--Wornell2024. The third row shows the ground-truth eigenfunctions.
  • ...and 3 more figures

Theorems & Definitions (5)

  • Theorem 1: Consistency of the OMM-$p$ objective
  • proof : Proof sketch
  • Theorem 2
  • proof : Proof of Theorem \ref{['thm:main']}
  • proof : Proof of Theorem \ref{['thm:jnt_nesting']}