Research map / Learning theory and algorithms / Matrix and tensor methods
The Expected Loss of Preconditioned Langevin Dynamics Reveals the Hessian Rank
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives a closed-form expected loss for preconditioned Langevin dynamics near stationary points (via an Ornstein–Uhlenbeck approximation) and shows that under a specific preconditioner–noise relation it equals the Hessian rank, yielding an iterative Hessian rank estimator for neural networks.
Problem
Estimating the rank of the loss Hessian of neural networks is computationally hard because the Hessian is too large to store or decompose, and the relation between preconditioned Langevin dynamics and loss near stationary points is not well characterized.
Approach
Near a stationary point, preconditioned LD reduces to a multivariate Ornstein–Uhlenbeck process; its steady-state expected loss is Tr(Sigma G P J P^-1)/4, which, when Sigma G = sigma^2 I, reduces to sigma^2 rank(H)/4, so the Hessian rank can be estimated by simply running LD and averaging the loss.
Key results
- For linear networks, the proposed estimator's normalized RMSE was 0.3% to 6.9% across ranks 8^2 to 256^2, versus 12.2% to 19.2% for the U&S method.
- On DnCNN (753 parameters), the proposed estimate of rank 133 captured more than 98.5% of the Hessian spectral energy, while U&S's rank 34 captured only 88%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace U&S with the proposed Langevin-dynamics rank estimator | Linear NN Hessian rank estimation · Linear NN Hessian ranks 8², 16², 32², 64², 128², and 256² | normalized RMSE: 12.2% -> 6.9%; 18.5% -> 2.8%; 19% -> 1.2%; 19.2% -> 0.3%; 17.9% -> 2%; 18.7% -> 4.6% | Helped |
| Replace U&S with the proposed Langevin-dynamics rank estimator | DnCNN Hessian rank estimation · MNIST DnCNN | eigenvalue energy captured: 88% -> more than 98.5% | Helped |
| Use the identity matrix instead of the Adam preconditioner (ablation) | Proposed Hessian rank estimation · MNIST DnCNN | stability | Hurt |
Surprising findings
- Plain gradient descent as the preconditioner (G=I) was unstable for deep networks, so an Adam-based preconditioner was needed in practice.
- The U&S rank estimator systematically underestimates rank, with roughly constant ~19% normalized error across network sizes.
Limitations
- Theory relies on a quadratic (second-order) approximation near a stationary point and an ergodicity/large-time assumption.
- Hessian-rank estimation for nonlinear networks is validated only on a small 753-parameter DnCNN where the Hessian can be computed exactly.
- Plain GD preconditioning (G=I) is unstable for deep networks, so the method requires an adaptive preconditioner such as Adam.
Experimental setup
Datasets: MNIST
Benchmarks: linear networks with d=8 to 256 (analytic Hessian rank d_x d_y) DnCNN on MNIST (753 parameters, numerically computed Hessian spectrum)
Metrics: normalized RMSE of Hessian rank estimate cumulative normalized Hessian eigenvalue energy captured by expected loss over time vs theory
Baselines: U&S matrix rank estimation method
Builds on: stochastic gradient Riemannian Langevin dynamics Adam Hessian-trace/rank estimation of Ullrich and Sagun-style spectral methods
Closest papers by meaning
- Preconditioned Langevin Dynamics with Score-based Generative Models for Infinite-Dimensional Linear Bayesian Inverse Problems (NeurIPS 2025 · similarity 0.90)
- Sample as you Infer: Predictive Coding with Langevin Dynamics (ICML 2024 · similarity 0.88)
- Dimension-Uniform Discretization Analysis of Preconditioned Annealed Langevin Dynamics for Multimodal Gaussian Mixtures (NeurIPS 2026 · similarity 0.87)
- Universal Sequence Preconditioning (NeurIPS 2025 · similarity 0.87)
- Generalization of noisy SGD in unbounded non-convex settings (ICML 2025 · similarity 0.87)
- High-dimensional Limit of SGD for Diagonal Linear Networks (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.