Calibrated Target Noise Recovers Curvature from the Gradients
Authors: Arash Jamshidi, Katsiaryna Haitsiukevich, Aristides Gionis, Kai Puolamäki
Organizations: Department of Computer Science, University of Helsinki, Helsinki, Finland · Division of Theoretical Computer Science, KTH Royal Institute of Technology, Digital Futures, Stockholm, Sweden · Department of Computer Science, Institute for Atmospheric and Earth System Research, University of Helsinki, Helsinki, Finland
Estimating curvature information from gradients alone is a fundamental problem in machine learning. The challenge becomes particularly acute in settings where only aggregate mini-batch gradients are available, yet curvature information is still required for applications such as preconditioning in optimisation, sampling, and early stopping. Perhaps surprisingly, we show that a simple target-perturbation scheme suffices to recover such curvature information from gradients alone. Our method injects zero-mean noise into the targets, with variance calibrated to the local output curvature. The covariance of the resulting noisy batch gradients then recovers practical curvature matrices, without requiring per-sample gradients, second-order derivatives, or other internal network quantities. For a broad class of non-linear models, our scheme recovers the Generalised Gauss--Newton (GGN) matrix; for feed-forward ReLU networks, it additionally recovers the diagonal blocks of the population Hessian. We demonstrate that the resulting estimator is simple, lightweight, provably optimal, and easy to integrate into existing training pipelines. Experiments demonstrate accurate curvature recovery, faster optimisation when the estimator is used for preconditioning, and the utility of the estimated GGN for early stopping.
Figures & tables
Centered estimator
Swamping estimator
1: Input: parameter w , noise multiplier β .
1: Input: parameter w , noise multiplier β , center μt .
2: Draw a mini-batch {(xi,yi)}i=1n .
2: Draw a mini-batch {(xi,yi)}i=1n .
3: Compute ηi=fw(xi) , ∀i∈[n] .
3: Compute ηi=fw(xi) , ∀i∈[n] .
4: Set vi=nA′′(ηi) , ∀i∈[n] .
4: Set vi=nA′′(ηi) , ∀i∈[n] .
5: Sample ξi∼Unif{+vi,−vi} , ∀i∈[n] .
5: Sample ξi∼Unif{+vi,−vi} , ∀i∈[n] .
Algorithm 1 Centred and swamping noisy-gradient curvature (CNGC) estimators.
Figure 1: Relative operator norm between GGN and its estimates showing the effects of centring ( left ), swamping ( middle ) and the added noise distribution ( right ) on curvature recovery. Both centring and swamping with a large noise multiplier make the estimates robust. Rademacher noise distribution yields a more accurate estimator with lower variance.
Figure 2: Comparison of GGN estimators with an MLP for a classification task ( left ) and a CNN for a regression task ( middle ) by the relative operator norm of the difference of the full GGN matrix and its estimators. The wall-clock time for calculation of the estimators for the CNN model ( right ). The proposed CNGC estimator performs consistently across classification and regression tasks and is computationally cheaper than the full GGN calculation.
Figure 3: The effect of calibrated noisy gradient covariance as a preconditioner (blue lines) for a classification task (telescope dataset) with an MLP, a time series classification (occupancy detection dataset) and regression (air quality dataset) with a 1D CNN and a transformer encoder.
Optimiser
MLP Telescope
CNN Occupancy
CNN Air
Transformer Occupancy
Transformer Air
GD
0.22±0.00
0.42±0.01
0.18±0.00
3.85±0.03
1.81±0.02
Adam
0.25±0.00
0.46±0.01
0.20±0.00
4.05±0.06
1.89±0.02
Sophia
0.26±0.00
0.49±0.01
0.22±0.00
4.42±0.06
2.05±0.02
SophiaC (ours)
0.26±0.00
0.50±0.00
0.22±0.00
4.41±0.04
2.06±0.02
GD, PC (ours)
0.35±0.01
0.66±0.01
0.29±0.00
4.90±0.05
2.29±0.02
Table 1: Wall-clock time per training epoch (seconds). The results are mean ± one standard deviation over 10 random seeds.
Optimiser
MLP Telescope
CNN Occupancy
CNN Air
Transformer Occupancy
Transformer Air
Test-loss threshold
0.3402
0.1701
0.4804
0.1767
0.2526
GD
22.29±0.23†
8.36±0.11†
18.22±0.35†
77.08±0.66†
181.03±2.00†
Adam
16.02±0.17
3.81±0.05
20.34±0.22†
44.71±0.66
169.94±2.04
Sophia
7.62±0.09
4.10±0.04
16.78±0.28
44.37±0.56
96.71±0.76
SophiaC (ours)
7.42±0.11
4.12±0.05
12.00±0.14
39.90±0.38
86.46±0.92
GD, PC (ours)
1.65±0.02
0.96±0.02
2.60±0.02
10.01±0.11
16.32±0.16
Table 2: Wall-clock time (seconds) to reach within 10% of the best test loss achieved in each experiment. The corresponding thresholds are listed below. † indicates that the threshold was not reached within the training budget; the full elapsed training time is reported instead. Bold and italic indicate the fastest and second-fastest methods, respectively. The results are mean ± one standard deviation over 10 random seeds.
Figure 4: The effect of calibrated noisy gradient covariance for early stopping (red vertical lines) for the same tasks as in the optimisation experiment compared to the baseline.
Figure 5: Illustrative example comparing CNGC and GGN with the true Hessian for a small convolutional network with 12 parameters. Top: regression, bottom: classification.
Figure 6: Relative operator norm between the spectrum of the true Hessian and the curvature estimator with centring and swamping with noise multiplier equal to 1 and to 10. Both centring and swamping with noise, with a multiplier equal to 10, capture the eigenspectrum of the true Hessian.
Figure 7: Top: Comparison of GGN estimators with an MLP for a classification task ( left ) and a CNN for a regression task ( middle ) by the relative operator norm of the difference of the eigenvalue spectrum of the full GGN matrix and its estimates. The wall-clock time for the calculations of GGN estimates for the MLP model ( right ). Bottom: Comparison of GGN estimators with the true Hessian for a CNN on a regression task. The relative operator norm of the difference between (positive) Hessian eigenvalues and the eigenvalues of the estimates ( left ), the relative operator norm of the difference between full matrices of the estimates and the Hessian ( right ).
Figure 8: The effect of calibrated noisy gradient covariance as a preconditioner for logistic regression and MLPs on classification tasks for telescope, bank marketing, and adult income datasets.
Optimiser
MLP Telescope
CNN Occupancy
CNN Air
Transformer Occupancy
Transformer Air
GD, PC (ours)
0.3005±0.0032
0.1080±0.0065
0.4053±0.0200
0.1136±0.0573
0.1456±0.0096
Exact GGN blocks (oracle)
0.2981±0.0040
0.1082±0.0130
0.3887±0.0214
0.0918±0.0058
0.1384±0.0118
Appendix
Table 3: Best test loss achieved using CNGC and exact GGN blocks. The results are mean ± one standard deviation over 10 random seeds.
Optimiser
MLP Telescope
CNN Occupancy
CNN Air
Transformer Occupancy
Transformer Air
GD, PC (ours)
4.50±0.50
1.10±0.30
11.00±3.82
3.70±5.46†
7.50±2.29
Exact GGN blocks (oracle)
3.40±0.80
1.10±0.30
3.70±1.10
2.40±1.62
7.20±3.16
Appendix
Table 4: Number of training epochs required to reach within 10% of the best test loss for each optimiser and model–dataset combination. † marks an entry for which the threshold was not reached within the training budget for some of the seeds. The results are mean ± one standard deviation over 10 random seeds.
Stopping criterion
MLP Telescope
CNN Occupancy
CNN Air
Transformer Occupancy
Transformer Air
CNGC (ours)
0.3075±0.0034
0.1279±0.0127
0.4647±0.0234
0.1223±0.0580
0.1843±0.0271
GradStop
0.3056±0.0039
0.1350±0.0196
0.4799±0.0419
0.1227±0.0610
0.2411±0.1246
Validation loss
0.3061±0.0039
0.1273±0.0197
0.5223±0.2019
0.1078±0.0137
0.1646±0.0096
Appendix
Table 5: Test loss obtained with each early-stopping criterion using patience 3. The results are mean ± one standard deviation over 10 random seeds.
Figure 9: The calculated stopping points on the occupancy detection dataset with GradStop and the proposed criterion with CNGC.
The standard convergence analysis of mini-batch stochastic gradient descent (SGD) models gradient noise using a single variance term that treats all parameter directions equally, ignoring the fact that noise in high-curvature directions has less impact because learning rates are already constrained there. We introduce Curvature-Weighted Gradient Diversity (CWGD), a geometry-aware measure that weights per-sample gradient diversity by the inverse square root of the Hessian, providing a tighter proxy for the effective optimization noise. For strongly convex quadratic objectives with diagonal Hessians and isotropic noise, we prove that a CWGD-modulated cosine learning-rate schedule can reduce the asymptotic optimization error floor by up to a factor of two compared with standard cosine annealing. We implement this idea as CWGD-Cosine using a Hutchinson-based diagonal Hessian estimator that is exact for quadratic objectives. Across a range of condition numbers, batch sizes, and noise structures, CWGD-Cosine consistently achieves approximately 20% lower final optimization error than standard cosine annealing while incurring negligible overhead in the quadratic setting. We also identify and correct a degenerate curvature estimator, analyze the robustness of the proposed estimator, and explicitly discuss the limitations of the method, including Hessian staleness in non-convex optimization. These results establish CWGD as a principled geometry-aware measure of optimization noise and motivate future extensions to more general learning problems.
Muhammad Hamza, Ayush Goel
Indian Institute of Technology Kharagpur · Indian Institute of Technology Kharagpur (IIT KGP)
Many machine learning techniques rely on approximating a loss function's curvature, but this is notoriously hard to do at the scale of modern deep networks. Surprisingly, no previous work has exploited the curvature constraints that arise from well known weight-space symmetries in loss landscapes. By analytically averaging over group actions that leave the loss invariant, we construct structured Hessian approximations from single gradients that can be tractably estimated, stored, and inverted. The choice of user-specified symmetry group directly governs the trade-off between approximation accuracy and computational cost. Moreover, our framework provides a unifying theoretical lens for viewing existing methods; in particular, a specific choice of symmetry group recovers Shampoo/Muon-like curvature estimates. We validate our method on a range of network architectures, and deploy it to second-order optimization benchmarks, including a small language model. Our curvature estimation framework might find applications in other machine learning problems such as uncertainty estimation, continual learning, compression/pruning, training data attribution, and more.
Artem Artemev, Rui Xia, Benjamin M. Boyd +4
MediaTek, Cambridge, UK · University of Cambridge, Cambridge, UK · Vector Institute, Toronto, Canada
As an alternative to the standard geometric analyses, we give an exact, information-theoretic analysis of stochastic gradient descent (SGD) and its variants. We show that a preconditioned SGD step is the posterior-mean update of a Gaussian Bayes model, and that its one-step regret splits into an intrinsic-time cost and a change in comparator information. The split extends to an identity for the objective itself. Convex convergence, strict-saddle-point escape, the link between flatness and generalization, the standard learning-rate schedules, adaptive optimizers, and the noisy, momentum, heavy-tailed, and gradient-free variants of SGD each correspond to a term or a special case of this identity. We measure its terms on synthetic and real training runs. On real networks it attributes the slack of classical convergence bounds to the terms their derivations drop and separates optimizers that reach the same training loss. That separation follows the number and consistency of their steps. Its relation to which of them generalizes better differs between networks. For gradient-free SGD the identity determines how a curvature preconditioner should enter the update. The sharpness-based generalization certificate it yields, with a data-independent isotropic prior, is vacuous at network scale unless the curvature spectrum is nearly flat across all parameters.