Natural W-TRAK Certification Framework

• Natural W-TRAK Certification Framework is a robust method that guarantees data attribution by applying a natural Wasserstein metric derived from the model’s feature covariance.
• It overcomes Euclidean robustness issues by neutralizing spectral amplification, yielding non-vacuous certified attribution intervals for deep neural networks.
• Empirical evaluations on CIFAR-10 demonstrate that the framework certifies 68.7% of ranking pairs and achieves AUROC of 0.970 for label-noise detection.

The Natural W-TRAK Certification Framework provides certified robustness guarantees for data attribution in machine learning models, ranging from convex estimators to deep neural networks. Attribution methods such as TRAK quantify how individual training examples influence test predictions, but conventional certification approaches based on Euclidean geometry become vacuous when applied to modern neural networks. The Natural W-TRAK framework resolves this by introducing a geometry derived from the model’s feature covariance, yielding the first non-vacuous certified attribution intervals for large-scale neural models (Li et al., 9 Dec 2025).

1. Mathematical Foundations

The framework’s central concept is the Natural Wasserstein metric, adapted to the geometry of the model’s feature space. For each input $z$, let $\phi(z) \in \mathbb{R}^d$ denote its feature embedding. Define the regularized sample feature covariance as $$Q = \mathbb{E}_{P_n}[\phi(z)\phi(z)^\top] + \lambda I$$.

Natural distance: The distance between two points $z$ and $z'$ in this geometry is

$$d_\mathrm{Nat}(z, z') = \sqrt{(\phi(z) - \phi(z'))^\top Q^{-1} (\phi(z) - \phi(z'))}$$

This is a Mahalanobis distance in the “whitened” feature space, aligning perturbations with the model’s learned representation.

Wasserstein-Robust Influence Functions (W-RIF):

In convex settings, influence is classically given by

$$I(z_i, z_\mathrm{test}) = -g_\mathrm{test}^\top H^{-1} g_i$$

where $g_i = \nabla_\theta \ell(\hat{\theta}; z_i)$ and $$H = \mathbb{E}_{P_n}[\nabla^2\ell(\hat{\theta}; z)]$$. The robust certified interval at radius $\varepsilon$ (in the chosen metric) is

$$I^{\mathrm{range}}_\varepsilon(z_i, z_\mathrm{test}) = \left[ \inf_{Q: W_1(Q, P_n) \leq \varepsilon} I_Q(z_i, z_\mathrm{test}), \ \sup_{Q: W_1(Q, P_n) \leq \varepsilon} I_Q(z_i, z_\mathrm{test}) \right]$$

where $I_Q$ denotes influence after retraining on distribution $Q$.

A functional Taylor expansion yields:

$$I_Q(z_i, z_\mathrm{test}) = I(z_i, z_\mathrm{test}) + \int S(z) \, d(Q - P_n)(z) + O(\|Q - P_n\|^2)$$

For sensitivity kernel $S$ that is $L_S$-Lipschitz in the ground metric, Kantorovich duality gives

$$I^{\mathrm{range}}_\varepsilon(z_i, z_\mathrm{test}) = I(z_i, z_\mathrm{test}) \pm \varepsilon L_S + O(\varepsilon^2)$$

Choosing $\varepsilon \geq \operatorname{diam}(Z) / n$ guarantees the coverage of leave-one-out influences, constituting a certified robust coverage result.

2. Limitations of Euclidean Robustness and Spectral Amplification

In deep neural networks, direct application of Euclidean-metric certification to attribution methods such as TRAK is ineffective due to spectral amplification. For a test point and candidate training point, linearized TRAK forms:

$$\mathrm{TRAK}(z_\mathrm{test}, z_i) = \phi(z_\mathrm{test})^\top Q^{-1} \phi(z_i)$$

The corresponding Euclidean Lipschitz bound on $\mathrm{TRAK}$ with respect to $\phi(z_i)$ is dominated by the spectral condition number of $Q$:

$$L_\mathrm{Euc} \lesssim \frac{1}{\lambda_{\min}(Q)} \|\phi(z_\mathrm{test})\|_2$$

For instance, on CIFAR-10 with a ResNet-18 last layer, $\kappa(Q) \approx 2.7 \times 10^5$ and $L_\mathrm{Euc}$ can exceed $7.7 \times 10^7$, rendering certification intervals vacuous (0% of actual attribution rankings are certifiable).

Spectral amplification arises because ill-conditioned feature covariances inflate Euclidean distances, undermining robustness certification in deep feature spaces.

3. Natural W-TRAK Certification for Deep Networks

Adopting the Natural metric, the sensitivity of the TRAK score to perturbations is controlled by the model’s feature geometry, eliminating the spectral amplification. The Natural W-TRAK interval for robust attribution is given by:

$$W\mathrm{TRAK}^{\mathrm{Nat}}_\varepsilon(z_\mathrm{test}, z_i) = \mathrm{TRAK}(z_\mathrm{test}, z_i) \pm \varepsilon L_\mathrm{Nat}(z_\mathrm{test}, z_i)$$

Where

$$L_\mathrm{Nat}(z_\mathrm{test}, z_i) \leq 2 \sqrt{SI(z_\mathrm{test})} \sqrt{SI(z_i)} R_\mathrm{whit}$$

Here, $SI(z) = \phi(z)^\top Q^{-1} \phi(z)$ is the Self-Influence, and $R_\mathrm{whit} = \max_j \sqrt{SI(z_j)}$ is the training manifold radius in whitened space.

Because both the metric and the attribution functional share $Q^{-1}$ structure, the worst-case $1/\lambda_{\min}(Q)$ amplification cancels, yielding non-vacuous intervals.

4. Self-Influence and Attribution Instability

Self-Influence $SI(z)$ quantifies the per-point geometric instability of attribution:

$$SI(z) = \phi(z)^\top Q^{-1} \phi(z) = \|Q^{-1/2} \phi(z)\|_2^2$$

The per-point Lipschitz constant of the TRAK score map $z_i \mapsto \mathrm{TRAK}(z_\mathrm{test}, z_i)$ under the Natural metric is $\sqrt{SI(z_i)}$ up to constant factors. High $SI(z)$ identifies training points whose attribution scores are susceptible to perturbations, providing both theoretical guarantees on robust attribution and a foundation for leverage-based anomaly detection.

Empirically, using $SI(z)$ for label-noise detection on CIFAR-10 with 10% label corruption, the method achieves AUROC $= 0.970$ and top-20% SI scores capture $94.1\%$ of noisy labels.

5. Algorithmic Implementation

The practical computation of Natural W-TRAK intervals involves the following steps:

1. Compute the regularized feature covariance $Q$ and its inverse.
2. Compute all Self-Influence scores $SI_i = \phi_i^\top Q^{-1} \phi_i$ and for the test point.
3. Cap $SI_\mathrm{test}$ at twice the maximum $SI_j$ for numerical stability.
4. Determine $R_\mathrm{whit} = \max_j \sqrt{SI_j}$.
5. For each candidate point $i$, compute $$L_{\mathrm{Nat},i} = 2 \sqrt{SI_\mathrm{test}} \sqrt{SI_i} R_\mathrm{whit}$$ and $\tau_i = \phi_\mathrm{test}^\top Q^{-1} \phi_i$.
6. Return the certified interval $$[\tau_i - \varepsilon L_{\mathrm{Nat},i}, \tau_i + \varepsilon L_{\mathrm{Nat},i}]$$ for each $i$.

Computational complexity is $O((n+1) \cdot d)$ for feature computation, $O(n d^2 + d^3)$ for covariance construction and inversion, and $O(n)$ for scalar computations, where $d$ is the feature dimension (typically $\sim5000$ for deep networks).

6. Empirical Performance and Significance

On CIFAR-10 (50,000 train, 10,000 test) with last-layer ResNet-18 gradients:

• Euclidean W-TRAK certifies 0% of all test–train ranking pairs at a fixed radius $\varepsilon$.
• Natural W-TRAK certifies 68.7% of ranking pairs at the same $\varepsilon$.
• The ratio $L_\mathrm{Euc} / L_\mathrm{Nat} \approx 76 \times$, consistent with $\sqrt{\kappa(Q)}$.

Additionally, in the context of label-noise detection:

• Utilizing $SI(z)$ as an anomaly score yields AUROC $=0.970$, AP $=0.796$.
• The top 20% of points by $SI$ capture 94.1% of corrupted labels.

These results demonstrate a substantial improvement in provable attribution robustness relative to prior methods based on the Euclidean ground metric.

7. Broader Implications and Extensions

Measuring distributional perturbations in the feature-induced Mahalanobis geometry directly aligns the certification process with the form of the attribution functional, neutralizing spectral ill-conditioning. The reduction in worst-case sensitivity by a factor of $O(\sqrt{\kappa(Q)})$ enables meaningful robust attribution analysis at scale. While the framework is derived for TRAK, this principle generalizes to any attribution method representable as a quadratic form $f(\phi) = \phi^\top A \phi$, provided perturbations are measured in the metric induced by $A$.

Self-Influence unifies certified robust attribution with leverage-based anomaly and outlier detection, offering theoretical support for established data cleaning heuristics. By overcoming the spectral amplification barrier, Natural W-TRAK sets precedent for non-vacuous certified influence in deep neural settings (Li et al., 9 Dec 2025).