Sample-wise AGOP: Theory & Applications

• Sample-wise AGOP is the empirical second-moment operator computed by averaging outer products of gradients or Jacobians, capturing sensitivity directions in prediction models.
• Estimation procedures for AGOP range from exact gradient methods and kernel approaches to compressive sensing techniques, enabling its use in high-dimensional settings.
• AGOP facilitates feature learning, dimensionality reduction, and optimization by aligning neural representations and enhancing metric learning in diverse applications.

Sample-wise Average Gradient Outer Product (AGOP) is the empirical second-moment operator obtained by forming an outer product from a gradient or Jacobian at each sample and averaging across samples. In scalar regression it appears as

$$\hat G=\frac{1}{n}\sum_{i=1}^n \hat g_i \hat g_i^\top,$$

while in multi-class settings it takes the Jacobian form

$$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$

Across recent work, AGOP is used both in input space—where it encodes directions along which predictions are most sensitive to features—and in parameter space—where it coincides with the empirical Fisher for negative log-likelihood losses and is closely related to Gauss–Newton structure under square loss (Trivedi et al., 2020, Mallinar et al., 2024, DePavia et al., 3 Feb 2025).

1. Definitions and variants

A common parameter-space definition begins with the per-sample gradient

$g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$

the sample-wise gradient outer product

$G_i=g_i g_i^\top,$

and the average

$$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$

A parallel input-space definition replaces loss gradients with Jacobians of the predictor with respect to inputs. For a vector-valued map $f:\mathbb{R}^d\to\mathbb{R}^c$, one writes

$$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$

The literature uses both forms, but distinguishes population operators such as EGOP or EJOP from the empirical sample-wise average that is typically called AGOP (Mallinar et al., 2024, Trivedi et al., 2020).

This distinction is central in regression and classification. For scalar-output regression, EGOP is the expectation $\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$; for multi-class classification, EJOP is $\mathbb{E}[J_f(X)J_f(X)^\top]$. AGOP is the finite-sample estimator obtained by replacing the population expectation with an average over observed samples. In the multi-index setting $f(x)=g(V^\top x)$, the top eigenspace of the population operator identifies the predictive subspace under mild nondegeneracy conditions, which gives AGOP a direct dimensionality-reduction interpretation (Trivedi et al., 2020).

Several papers generalize the basic construction. In structured product spaces, a blockwise AGOP averages $$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$0 over both samples and blocks, so that row- or patch-level feature sharing can be represented explicitly. In convolutional networks, the analogous object is a patch-based AGOP with gradients taken with respect to local patches entering a convolutional layer. In deep-network analyses, layerwise AGOPs are defined on intermediate representations rather than raw inputs, again as sample-wise averages of Jacobian outer products (Radhakrishnan et al., 2024, Beaglehole et al., 2023, Beaglehole et al., 2024).

2. Estimation procedures

The most direct estimators compute exact gradients or Jacobians and average their outer products. In fully connected networks, the mini-batch gradient for layer $$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$1 can be written as

$$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$2

so backpropagation itself already exposes a sample-wise AGOP structure: the layer gradient is an average of rank-1 per-sample outer products. This observation underlies approximate training schemes such as Mem-AOP-GD, which selects only a subset of these outer products and stores the remainder in memory as a bias-correction mechanism (Hernandez et al., 2021).

When gradients are not available in closed form, AGOP can be built from surrogate estimators. In the EJOP estimator for multi-class nonparametric classification, the class-probability function is first smoothed by a kernel, $$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$3-NN, or $$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$4-NN estimator, optionally followed by a softmax normalization. The Jacobian is then approximated with symmetric finite differences,

$$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$5

and unreliable coordinates are gated by empirical-mass indicators $$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$6. The final estimator is the average

$$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$7

followed by spectral decomposition and, if desired, a Mahalanobis-style transformation $$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$8 (Trivedi et al., 2020).

Kernel methods admit exact input-gradient formulas. For kernel ridge regression,

$$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$9

This makes AGOP computation explicit for inner-product, Gaussian, polynomial, Laplace, Matérn, and Rational Quadratic kernels, and supports recursive schemes in which a kernel metric $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$0 is updated from the AGOP of the current predictor (Zhu et al., 14 May 2026, Shen et al., 2024).

In high dimensions, direct gradient computation may be too expensive. A different estimator combines simultaneous perturbation with compressive sensing: random linear measurements of the gradient are obtained from only two function evaluations per perturbation step, averaged over repeated perturbations, and then inverted by constrained $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$1 recovery. This yields per-sample gradient estimates $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$2, from which the AGOP estimator

$g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$3

is formed (Borkar et al., 2015).

3. Spectral theory and subspace recovery

The principal theoretical role of AGOP is to recover low-dimensional structure from high-dimensional predictors. In multi-index regression, if $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$4 and the gradient outer product of $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$5 is nonsingular, then the top $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$6 eigenvectors of $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$7 span $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$8. The EJOP work extends this idea from scalar regression to multi-class classification and proves operator-norm consistency of the empirical estimator under smoothness, bounded-gradient, and lower-bounded density assumptions. The same analysis yields consistency of eigenvalues and eigenspaces through standard perturbation arguments, with one concrete asymptotic schedule given by $g_i=\nabla_\theta \ell(f_\theta(x_i),y_i),$9 and $G_i=g_i g_i^\top,$0 (Trivedi et al., 2020).

A stronger separation result appears in kernel regression for multi-index models. There, a kernel ridge regressor may still have large prediction error, yet the AGOP of the fitted predictor already recovers the central subspace. Specifically, if the target $G_i=g_i g_i^\top,$1 has degree $G_i=g_i g_i^\top,$2, accurate prediction requires $G_i=g_i g_i^\top,$3, but if a lower-degree component of degree $G_i=g_i g_i^\top,$4 already contains all relevant directions, AGOP-based subspace recovery occurs in the regime $G_i=g_i g_i^\top,$5 for any $G_i=g_i g_i^\top,$6. This establishes a prediction–representation separation: identifying the relevant subspace can be statistically easier than learning the full target function (Zhu et al., 14 May 2026).

In linear Recursive Feature Machines, AGOP becomes especially explicit. For the structured linear model studied in low-rank matrix recovery, the blockwise AGOP of the current predictor collapses to

$G_i=g_i g_i^\top,$7

and the update

$G_i=g_i g_i^\top,$8

induces a spectral reweighting equivalent to IRLS-$G_i=g_i g_i^\top,$9 with $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$0 when $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$1. Fixed points of the iteration are first-order critical points of a singular-value penalty, and for $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$2 the method inherits classical IRLS recovery guarantees (Radhakrishnan et al., 2024).

A modal analogue, LMOPG, replaces the conditional mean $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$3 by the conditional mode $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$4. Its population target is $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$5, and under a rank condition its span equals the central subspace. The local modal estimator is proved consistent and asymptotically normal under assumptions (A1)–(A5), which is significant because mean-based OPG can miss central-subspace directions when error distributions are symmetric about zero (Li et al., 2024).

4. AGOP as a mechanism of feature learning

A recurring theme in recent work is that AGOP is not merely a diagnostic statistic but an explicit feature-learning mechanism. Recursive Feature Machines operationalize this idea by alternating between predictor fitting and an AGOP-based metric update

$$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$6

with modular-arithmetic experiments using $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$7. In that setting, RFM produces a sharp transition from near-random test accuracy to perfect test accuracy while training loss is identically zero, and the transition is tracked by gradual AGOP alignment and the emergence of block-circulant features. The same work argues that these features implement the Fourier Multiplication Algorithm and reports that the square-root of the network input AGOP correlates strongly with the neural feature matrix, with Pearson correlation above $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$8 across modular tasks (Mallinar et al., 2024).

Neural-network studies describe a closely related phenomenon through the Neural Feature Ansatz. In fully connected networks, high correlation between $$G_{\mathrm{AGOP}}=\frac{1}{n}\sum_{i=1}^n g_i g_i^\top.$$9 and an input-space AGOP at layer $f:\mathbb{R}^d\to\mathbb{R}^c$0 can be rewritten as alignment between the left singular structure of $f:\mathbb{R}^d\to\mathbb{R}^c$1 and a pre-activation tangent-feature covariance. Early-time dynamics under gradient flow show that this alignment is driven by the interaction between SGD-induced weight changes and tangent features. A convolutional counterpart, the Convolutional Neural Feature Ansatz, states that $f:\mathbb{R}^d\to\mathbb{R}^c$2 is proportional to a patch-based AGOP taken with respect to layer-$f:\mathbb{R}^d\to\mathbb{R}^c$3 patches, and empirical studies report very high correlations—often exceeding $f:\mathbb{R}^d\to\mathbb{R}^c$4—for pretrained AlexNet, VGG, and ResNet models (Beaglehole et al., 2024, Beaglehole et al., 2023).

The same mechanism has been invoked in deep neural collapse. Deep RFM constructs a multilayer model by repeatedly projecting features with $f:\mathbb{R}^d\to\mathbb{R}^c$5, where $f:\mathbb{R}^d\to\mathbb{R}^c$6 is the AGOP at layer $f:\mathbb{R}^d\to\mathbb{R}^c$7, and then applying a fixed random feature map. Empirically, most NC1 reduction occurs at the AGOP projection step rather than at the random feature map, and the asymptotic analysis shows contraction of the Gram matrix toward a collapsed fixed point under suitable kernel-linearization conditions (Beaglehole et al., 2024).

AGOP also appears in optimization narratives that connect training dynamics to feature learning. In the catapult regime of large-learning-rate GD or SGD, spikes in the training loss occur in a low-dimensional NTK eigenspace, and these spikes are associated with improved alignment between the model AGOP and the AGOP of the true predictor. Smaller batch sizes induce more catapults, larger AGOP alignment, and better test performance; in the reported synthetic example, AGOP alignment rises from $f:\mathbb{R}^d\to\mathbb{R}^c$8 at batch size $f:\mathbb{R}^d\to\mathbb{R}^c$9 to $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$0 at batch size $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$1, while test loss drops from $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$2 to $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$3 (Zhu et al., 2023).

5. Applications

One direct application is metric learning for nonparametric classification. The EJOP estimator defines a Mahalanobis-like distance

$$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$4

or equivalently a spectral rescaling by $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$5. On raw-pixel MNIST, the reported classification error drops from $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$6 with Euclidean distance and $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$7 with ReliefF-scaled Euclidean distance to $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$8 for EJOP with $$G(f)=\frac{1}{n}\sum_{j=1}^n J(x^{(j)})J(x^{(j)})^\top.$$9-NN and $\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$0 for EJOP with $\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$1-NN. The same study reports competitive performance relative to LMNN, ITML, and MLR on additional datasets, and also uses EJOP as a cheap initialization for metric learning (Trivedi et al., 2020).

A second application is local explanation in tabular classification. AGOP-IxG computes an empirical training-set AGOP from predicted-class logit gradients, truncates it to the top-$\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$2 eigenspace defined by $\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$3, and filters a sample gradient before applying an Input$\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$4Gradient-style baseline subtraction. On three synthetic multi-class tabular benchmarks, AGOP-IxG leads on Spearman rank correlation and noise feature mass in all cases, and on top-$\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$5 precision for the interaction task. It is also reported to be approximately $\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$6 to $\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$7 faster than SHAP on these settings. Under ROAR on Adult Income and Credit Card Default, all methods cluster within about $\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$8 relative AUC, which is presented as evidence that AGOP-IxG is optimized for local rather than global attribution (Katakam, 15 May 2026).

Continual test-time adaptation provides a third use case. GOLD maintains a feature-space AGOP

$\mathbb{E}[\nabla f(X)\nabla f(X)^\top]$9

as a streaming proxy for the classifier-sensitive subspace, initialized at $\mathbb{E}[J_f(X)J_f(X)^\top]$0. The top eigenspace is used as a low-rank adapter basis, and the reported cumulative spectral energy shows that $\mathbb{E}[J_f(X)J_f(X)^\top]$1–$\mathbb{E}[J_f(X)J_f(X)^\top]$2 eigenvectors capture more than $\mathbb{E}[J_f(X)J_f(X)^\top]$3 of the energy. The paper further reports that alignment between the AGOP-derived subspace and the single-step golden subspace stabilizes above $\mathbb{E}[J_f(X)J_f(X)^\top]$4, with only about $\mathbb{E}[J_f(X)J_f(X)^\top]$5 batch time spent on AGOP calculation, about $\mathbb{E}[J_f(X)J_f(X)^\top]$6 on eigendecomposition, and a trainable-parameter ratio below about $\mathbb{E}[J_f(X)J_f(X)^\top]$7 (Lai et al., 23 Mar 2026).

AGOP has also been used to improve optimization and computational efficiency. EGOP-based orthonormal reparameterization aligns adaptive optimizers such as Adagrad and Adam with dominant gradient directions, and the accompanying theory ties possible gains to the stable-rank decay of the EGOP spectrum (DePavia et al., 3 Feb 2025). In DNN training, Mem-AOP-GD approximates backpropagation by selecting only $\mathbb{E}[J_f(X)J_f(X)^\top]$8 sample-wise outer products per layer and retaining the remainder in memory; experiments on energy-efficiency regression and MNIST report significant reductions in the number of outer products with competitive or improved learning curves at moderate $\mathbb{E}[J_f(X)J_f(X)^\top]$9 (Hernandez et al., 2021).

In scientific machine learning and chemoinformatics, AGOP-driven RFM has been used for interpretable QSPR modeling. With multi-scale hybrid fingerprints, RFM is reported to achieve state-of-the-art solubility prediction across nine benchmark datasets, with Matérn and Laplace kernels performing best among the studied kernels. Local and global explanations are obtained from gradients, sample-wise AGOPs, and the learned metric $f(x)=g(V^\top x)$0, and the reported performance on ESOL and FreeSolv is approximately $f(x)=g(V^\top x)$1 RMSE and $f(x)=g(V^\top x)$2 RMSE, respectively (Shen et al., 2024).

6. Limitations, ambiguities, and open directions

The reliability of AGOP depends strongly on the fidelity of the underlying gradients or Jacobians. The nonparametric EJOP estimator assumes a continuously differentiable target on a $f(x)=g(V^\top x)$3-envelope, bounded and uniformly continuous gradients, and a lower-bounded input density on compact support. Practical guidance in that setting includes smoothing, symmetric finite differences, and boundary gating through the events $f(x)=g(V^\top x)$4, but the same work notes that sharp decision boundaries, insufficient local mass, and severe boundary effects can degrade the estimate (Trivedi et al., 2020).

Computational cost is a second recurring limitation. Exact AGOP formation is often $f(x)=g(V^\top x)$5 or worse, with an additional $f(x)=g(V^\top x)$6 eigendecomposition in dense settings. Kernel versions incur $f(x)=g(V^\top x)$7 storage for Gram matrices and potentially $f(x)=g(V^\top x)$8 solves; patch-based and layerwise AGOPs require expensive Jacobian extraction; streaming CTTA variants require maintaining a dense $f(x)=g(V^\top x)$9 matrix; and compressive-sensing estimators trade gradient access for repeated randomized function evaluations and an $$\hat G=\frac{1}{n}\sum_{i=1}^n \hat J_i \hat J_i^\top.$$00 recovery stage (Zhu et al., 14 May 2026, Beaglehole et al., 2023, Lai et al., 23 Mar 2026, Borkar et al., 2015).

There are also substantive modeling caveats. The block-circulant structures observed in modular arithmetic are explicitly described as task-specific and tied to cyclic-group structure, so similar geometry need not appear in tasks without comparable symmetry. In that same setting, the learned AGOP features are reported not to be low-rank, which limits methods that rely only on low-rank assumptions. In attribution, AGOP-IxG improves local fidelity on controlled synthetic data, but ROAR results on real data show clustering rather than dominance, indicating that local attribution quality and global feature ranking remain distinct objectives (Mallinar et al., 2024, Katakam, 15 May 2026).

Open problems in the cited work are largely estimator- and representation-oriented. Proposed directions include better adaptive Jacobian estimators such as local polynomial or graph Laplacian methods, tighter convergence rates and finite-sample bounds, extensions to structured outputs and deep feature spaces, class-conditional or mean-subtracted AGOP variants for attribution, and scalable low-rank or incremental eigensolvers for online use (Trivedi et al., 2020, Katakam, 15 May 2026, Lai et al., 23 Mar 2026). Taken together, these directions suggest that AGOP is best understood not as a single fixed statistic but as a family of sample-wise gradient-covariance operators whose empirical usefulness depends on how faithfully local sensitivities are estimated and how effectively their spectrum is exploited.