Research map / Optimization / Optimizers
Revisiting Scalable Hessian Diagonal Approximations for Applications in Reinforcement Learning
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.
Revisits the Becker and LeCun (1989) deterministic Hessian diagonal approximation, adds exact last-layer Hessian diagonals (HesScale), and uses it for second-order optimizers and trust-region step-size scaling in RL.
Problem
Hessian diagonal approximations that are cheap (gradient-like cost) are needed for second-order optimization and step-size scaling, but existing scalable methods such as BL89, AdaHessian and GGN-MC have poor approximation quality or high cost.
Approach
Refine the BL89 layer-wise backpropagated Hessian diagonal approximation by using the exact output-layer Hessian diagonal (e.g., q - q∘q for softmax cross-entropy) instead of BL89's approximation, keeping linear cost; use the resulting diagonal for AdaHesScale optimizers and K-FAC-style trust-region step-size scaling.
Key results
- On a small tanh MLP, HesScale's normalized L1 approximation error is lower than BL89, AdaHessian, GGN-MC and diag(GGN) approximations (HesScale used as the reference at 1.0).
- AdaHesScaleGN consistently outperformed all methods on CIFAR-100 3C3D and AdaHesScale outperformed most methods on ALL-CNN, with per-step cost about 2x Adam for AdaHesScale and 1.25x for AdaHesScaleGN (AdaHessian 3x).
- HesScale-based step-size scaling with PPO made Adam-style optimizers insensitive to step size in all MuJoCo environments except Humanoid, and kept all five UR-Reacher runs stable where Adam deteriorated.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use exact Hessian diagonals at the last layer instead of the BL89 approximation (ablation) | HesScale · MNIST classification, Hessian-diagonal approximation | L1 approximation error | Helped |
| drop the final term in the backpropagation rule (Gauss-Newton variation) (ablation) | HesScaleGN · MNIST classification, Hessian-diagonal approximation | L1 approximation error | Mixed |
| use HesScale rather than alternative stochastic or deterministic diagonal approximations | Hessian-diagonal approximations · MNIST classification | L1 approximation error | Helped |
| use HesScaleGN rather than other optimizers | CIFAR-100 3C3D · CIFAR-100 3C3D | classification performance | Helped |
| use AdaHesScaleGN rather than other optimizers except AdaHesScale | CIFAR-100 ALL-CNN · CIFAR-100 ALL-CNN | classification performance | Helped |
| use AdaHessian rather than the other optimizers | CIFAR-100 3C3D and ALL-CNN · CIFAR-100 | classification performance | Hurt |
| use AdaHesScale rather than Adam and AdaHessian | A2C · MuJoCo: Ant, Walker2d, HalfCheetah, and Hopper | undiscounted return | Helped |
| use AdaHesScale rather than Adam | PPO · MuJoCo environments | undiscounted return | No clear effect |
| use AdaHesScale rather than AdaHessian | PPO · Most evaluated MuJoCo environments | undiscounted return | Helped |
| add HesScale-based step-size scaling (ablation) | AdaHesScale and Adam with PPO · MuJoCo environments except Humanoid | robustness to step-size choice | Helped |
| add HesScale-based step-size scaling (ablation) | Adam · UR-Reacher | undiscounted return stability | Helped |
Surprising findings
- Simply replacing BL89's last-layer Hessian approximation with the exact output-layer diagonal substantially improves approximation quality, despite BL89 being long overlooked.
- AdaHessian performed worse than all methods except BL89 in supervised experiments, consistent with its poor diagonal approximation quality.
Limitations
- Approximation is exact only in special cases; the diagonal approximation ignores off-diagonal Hessian terms and is shown to be accurate mainly on small networks.
- Scalability to larger models is not demonstrated; the paper states this as a future promise.
- RL experiments use small MLPs (two hidden layers of 64 units) and the Humanoid environment did not benefit from step-size scaling.
Experimental setup
Datasets: MNIST EMNIST CIFAR-100 MuJoCo environments UR-Reacher
Benchmarks: CIFAR-100 3C3D (DeepOBS) CIFAR-100 ALL-CNN (DeepOBS) MuJoCo (5 environments) UR-Reacher
Metrics: L1 distance to exact Hessian diagonals (normalized) validation accuracy vs time undiscounted return robustness across step sizes stability across runs
Baselines: BL89 AdaHessian GGN-MC diag(GGN) Adam SGD A2C PPO
Builds on: Becker and LeCun 1989 diagonal Hessian approximation Adam K-FAC BackPACK
Closest papers by meaning
- Bigger, Regularized, Optimistic: scaling for compute and sample efficient continuous control (NeurIPS 2024 · similarity 0.87)
- A Unifying View of Linear Function Approximation in Off-Policy RL Through Matrix Splitting and Preconditioning (NeurIPS 2025 · similarity 0.87)
- Oracle-Efficient Reinforcement Learning for Max Value Ensembles (NeurIPS 2024 · similarity 0.86)
- On the Second-Order Convergence of Biased Policy Gradient Algorithms (ICML 2024 · similarity 0.86)
- On Double Descent in Reinforcement Learning with LSTD and Random Features (ICLR 2024 · similarity 0.86)
- Breaking the Bias Barrier in Concave Multi-Objective Reinforcement Learning (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.