Curvature-Aware Gradient Estimation (CAGE)

• Curvature-Aware Gradient Estimation (CAGE) is a family of methods that explicitly integrates second-order information to improve gradient accuracy and optimization convergence.
• CAGE leverages techniques such as Hessian-aligned simplex schemes, inexact Newton methods, and Kalman filtering to mitigate the bias and variance of first-order estimates.
• Empirical results demonstrate that CAGE methods achieve lower mean-squared error, faster convergence, and enhanced performance in bilevel optimization, reinforcement learning, and quantization-aware training.

Curvature-Aware Gradient Estimation (CAGE) refers to a family of methodologies that augment or replace classical gradient estimates in optimization or learning problems with versions that explicitly incorporate curvature—typically via Hessian or higher-order information—to achieve superior accuracy, variance properties, or convergence speed. CAGE methods have been developed independently across diverse subfields: bilevel optimization, numerical differentiation, stochastic optimization, reinforcement learning, quantization-aware training, and Riemannian gradient estimation. While their technical implementations differ substantially, all CAGE approaches leverage either explicit second-order information or data-driven surrogates to mitigate the structural limitations of first-order gradient schemes.

1. Mathematical Principles of Curvature-Aware Estimation

Curvature-aware gradient estimators rely on the observation that first-order Taylor expansions (or stochastic gradient analogues) can be systematically biased or high-variance in complex or high-noise regimes. Explicitly utilizing the Hessian $\nabla^2 f(x)$, or an approximation thereof, enables CAGE methods to:

• Reduce approximation error in finite-difference (or simplex) schemes by aligning sample directions with curvature information (Lengyel et al., 2023).
• Improve the approximation of hypergradients in bilevel problems via inexact Newton methods that exploit the shared Hessian in inner and outer derivatives (Dong et al., 4 May 2025).
• Filter stochastic gradients using predictive models for gradient evolution, governed by dynamics involving per-iteration Hessian-vector products, resulting in Bayesian filtering-based estimates with noise-adapted weighting (Chen et al., 2020).
• Incorporate natural-gradient or Fisher information via blockwise Kronecker-factored updates to minimize gradient estimator variance in RL (Firouzi, 2018).
• Design multi-objective corrected descent directions that penalize constraint violations (such as quantization error) with curvature (or surrogate-curvature) scaled corrections (Tabesh et al., 21 Oct 2025).

Curvature information is assimilated via explicit Hessian computation, Hessian-vector products, data-driven surrogate modeling, or optimizer statistics (e.g., Adam's second moment $v_t$).

2. Key Algorithmic Frameworks

Multiple CAGE frameworks are distinguished by their operational context and algorithmic design.

Domain	CAGE Methodology	Core Curvature Mechanism
Bilevel Optimization	Inexact Newton-based hypergradient estimation (Dong et al., 4 May 2025)	Iterative Hessian-vector Newton subproblems
Numerical Differentiation	Curvature-Aligned Simplex Gradient (CASG) (Lengyel et al., 2023)	Hessian-aligned simplex construction, blockwise Hadamard
Stochastic Optimization	Curvature-aware gradient filtering (Chen et al., 2020)	Online Kalman filter with Hessian-transport dynamics
Reinforcement Learning	Kronecker-factored curvature (KFAC) augmented control variate (Firouzi, 2018)	KFAC Fisher blocks for natural gradient steps
Quantization-Aware Training (QAT)	Pareto-gradient correction with curvature-weighted quant error (Tabesh et al., 21 Oct 2025)	Local curvature scaling via optimizer second moments (Adam)
Riemannian Gradient Estimation	CAGE estimator on manifolds (Wang et al., 2021)	Curvature-compensated differences via Greene–Wu convolution

In bilevel optimization, CAGE solves coupled quadratic subproblems (lower-level variable update and curvature-corrected hypergradient) via inexact Newton steps, achieving quadratic error reduction per step (Dong et al., 4 May 2025). In numerical differentiation, the curvature-aligned simplex selects evaluation points to minimize mean-squared error under a Hessian-informed model, yielding substantial efficiency gains in noisy, high-dimensional settings (Lengyel et al., 2023). KF-LAX in RL utilizes per-layer KFAC approximations to perform natural-gradient variance reduction (Firouzi, 2018).

3. Theoretical Guarantees and Complexity Analysis

CAGE frameworks improve the dependence of convergence rates and estimator variance on condition number, noise levels, or dimension:

• In deterministic bilevel optimization, NBO-GD (CAGE) achieves $O(\kappa^3/\epsilon)$ outer iterations and $O(\kappa^3/\epsilon)$ gradient-vector and $O(\kappa^4/\epsilon)$ Hessian-vector products, outperforming amortized implicit differentiation methods by an $O(\kappa)$ factor in gradient calls (Dong et al., 4 May 2025).
• In the curvature-aligned simplex (CASG), the estimator attains the optimal possible mean-squared error for any (d+1)-point finite-difference scheme, with theoretical analysis showing that step-sizes and directions adapt optimally to the eigenspectrum of the Hessian (Lengyel et al., 2023).
• For stochastic optimization under quadratic noise, the curvature-aware filter approach yields unbiased estimates whose variance contracts as $O(1/t)$, enabling robust convergence with constant step size (Chen et al., 2020).
• KF-LAX provably matches RELAX in unbiasedness but reduces gradient estimator variance by 30–50%, halving sample complexity on discrete RL tasks (Firouzi, 2018).
• In QAT, CAGE delivers ergodic $O(1/\sqrt{T})$ convergence to Pareto-stationarity in non-convex settings, eliminating the persistent quantization-induced bias of the straight-through estimator (Tabesh et al., 21 Oct 2025).

4. Implementation Details and Practical Overheads

CAGE methods introduce varying computational and memory overheads:

• Inexact Newton for bilevel CAGE: Each outer iteration involves $T=1$ to $10$ inner Newton-like steps (via conjugate gradient or GD on the Hessian), with the cost dominated by Hessian-vector products. Empirically, a single inner step ($v_t$0) achieves most of the possible performance gain (Dong et al., 4 May 2025).
• CASG: Preprocessing costs $v_t$1 for Hessian eigendecomposition, but per-gradient estimation is $v_t$2. The construction is practical for moderate $v_t$3 and justified when function evaluations are expensive (Lengyel et al., 2023).
• CAGE filter in SGD: Per-iteration cost is 1–1.6$v_t$4 vanilla SGD, with additional memory for per-sample gradients and Hessian-vector products (Chen et al., 2020).
• KF-LAX: Requires estimating and inverting per-layer KFAC factors, $v_t$5 per wide layer, with Tikhonov damping and batchwise updating to keep the cost manageable (Firouzi, 2018).
• CAGE for QAT: The decoupled variant adds a single quantization operation and vector subtraction per step. Coupled variants rescale the correction according to Adam's $v_t$6 statistics without extra Hessian computation (Tabesh et al., 21 Oct 2025).
• Riemannian CAGE: Requires two-point evaluations per direction and $v_t$7 vector operations per estimate, with complexity controlled by the number of directions $v_t$8 and dimension $v_t$9 (Wang et al., 2021).

5. Empirical Results and Comparative Performance

CAGE methods generally demonstrate substantial improvements over first-order or curvature-agnostic baselines:

• Bilevel optimization: NBO-GD with $O(\kappa^3/\epsilon)$0 inner step matches or outperforms AmIGO-GD with $O(\kappa^3/\epsilon)$1 in outer-iteration count and wall-clock time. NSBO-SGD converges faster than SOBA, AmIGO, and SABA, finding lower test error in less time (Dong et al., 4 May 2025).
• Numerical differentiation: CASG yields orders-of-magnitude lower MSE than forward differences on ill-conditioned or indefinite Hessians, achieving accuracy close to central differences with half the function calls; global surrogates retaining past evaluations nearly match true-Hessian performance (Lengyel et al., 2023).
• Stochastic optimization: On toy problems, CAGE matches batch GD; in deep networks, CAGE's filtered gradients are $O(\kappa^3/\epsilon)$2 closer (in $O(\kappa^3/\epsilon)$3) to the true full-batch gradient than naive minibatch estimates, though there is no consistent improvement over well-tuned Adam/SGD in test accuracy (Chen et al., 2020).
• Reinforcement learning: KF-LAX reaches high cumulative reward in %%%%24$O(\kappa^3/\epsilon)$25%%%% fewer environment steps than RELAX, with variance reduction clearly observed on both toy and Atari-style tasks (Firouzi, 2018).
• Quantization-aware training: In Llama-style pretraining (W4A4 regime), CAGE recovers over 10% of the quantization-induced loss, systematically exceeding QuEST and outlier-mitigation approaches across model scales; the Pareto-gradient avoids the stationary quantization bias of STE-based training (Tabesh et al., 21 Oct 2025).
• Riemannian setting: The geodesic-sphere CAGE estimator achieves lower MSE than Nesterov–Spokoiny or Gaussian schemes, especially in high-dimensional or curved spaces (Wang et al., 2021).

6. Extensions, Limitations, and Future Directions

CAGE is now established across several fundamental learning, optimization, and control settings, but key limitations and open directions remain:

• Generalization to non-smooth quantization operators and ultra-low bitwidths in QAT (Tabesh et al., 21 Oct 2025).
• Further acceleration in deep network optimization, where curvature-aware filtering steadies gradient estimation but does not outperform adaptive first-order optimizers in nonconvex regimes (Chen et al., 2020).
• Scalability of CASG and Riemannian CAGE to very high-dimensional problems, due to cubic preprocessing cost, though techniques like block-wise partitioning or global-model surrogates ameliorate this (Lengyel et al., 2023).
• Joint optimization of quantization and sparsification, and principled adaptation of weightings in Pareto-type CAGE frameworks.
• Potential for multi-step or meta-learned versions of curvature-aware filtering to address short-horizon biases in stochastic optimization (Chen et al., 2020).
• Empirical tuning of trade-off parameters (e.g., $O(\kappa^3/\epsilon)$6 in QAT CAGE), with a need for automated selection mechanisms (Tabesh et al., 21 Oct 2025).

A plausible implication is that continued integration of low-rank, data-driven, or structure-exploiting curvature approximations will further reduce the practical cost of CAGE methods, enabling widespread adoption as model and dataset scales increase.

7. Summary of Core CAGE Approaches

Approach	Reference	Setting	Mechanism	Empirical Advantage
NBO-GD/NSBO-SGD	(Dong et al., 4 May 2025)	Bilevel optimization	Inexact Newton hypergradient	Improved iteration complexity, fast convergence
CASG	(Lengyel et al., 2023)	Numerical differentiation	Hessian-aligned simplex	Order-of-magnitude MSE reduction vs. FD
KF-LAX	(Firouzi, 2018)	RL gradient estimation	Blockwise KFAC in LAX/RELAX	Halved sample complexity, reduced variance
CAGE Filter	(Chen et al., 2020)	Stochastic opt.	Kalman filter, Hessian transport	5× closer to full-batch gradient, robust in noise
Riemannian CAGE	(Wang et al., 2021)	Manifold grad. est.	Geodesic sphere finite differences	Lower MSE, bias-neutralized on manifolds
QAT CAGE	(Tabesh et al., 21 Oct 2025)	Quantization training	Pareto-corrected gradient, Adam scaling	10–20 % gap closure vs. STE, no quantization bias

CAGE thus provides a general template for curvature-informed gradient estimation, unifying several previously distinct threads across contemporary machine learning, optimization, and computational mathematics.