Imperfect Influence, Preserved Rankings: A Theory of TRAK for Data Attribution

TL;DR

This paper provides a theoretical analysis of the TRAK algorithm, characterizing its performance and quantifying the errors introduced by the approximations on which the method relies, and shows that although the approximations incur significant errors, TRAK's estimated influence remains highly correlated with the original influence and therefore largely preserves the relative ranking of data points.

Abstract

Data attribution, tracing a model's prediction back to specific training data, is an important tool for interpreting sophisticated AI models. The widely used TRAK algorithm addresses this challenge by first approximating the underlying model with a kernel machine and then leveraging techniques developed for approximating the leave-one-out (ALO) risk. Despite its strong empirical performance, the theoretical conditions under which the TRAK approximations are accurate as well as the regimes in which they break down remain largely unexplored. In this paper, we provide a theoretical analysis of the TRAK algorithm, characterizing its performance and quantifying the errors introduced by the approximations on which the method relies. We show that although the approximations incur significant errors, TRAK's estimated influence remains highly correlated with the original influence and therefore largely preserves the relative ranking of data points. We corroborate our theoretical results through extensive simulations and empirical studies.

Figures (17)

• Figure 1: Experimental results for $3$-class classification with $p=100$. Left two panels: Results for the dependent case, where $\mathbf{z}_i=\mathbf{z}_{\mathrm{new}}$. Right two panels: Results for the independent case, where $\mathbf{z}_i$ and $\mathbf{z}_{\mathrm{new}}$ are independent. Top row: Displays $\mathcal{I}^{\rm Linear}$ versus $\mathcal{I}^{\rm True}$. Although the deviation from the red $y=x$ line suggests large errors between $\mathcal{I}^{\rm Linear}$ and $\mathcal{I}^{\rm True}$, the two quantities still exhibit strong correlation. Bottom row: Displays $\mathcal{I}^{\rm ALO}$ versus $\mathcal{I}^{\rm Linear}$.
• Figure 2: Experimental results for a three-class classification problem with $p = 100$ and $d=200$. The x-axis shows $\mathcal{I}^{\rm True}$, while the y-axis shows $\mathcal{I}^{\rm TRAK}$ after projection. Left two panels: Results for the dependent case, where $\mathbf{z}_i=\mathbf{z}_{\mathrm{new}}$. Right two panels: Results for the independent case, where $\mathbf{z}_i$ and $\mathbf{z}_{\mathrm{new}}$ are independent. From top to bottom, the projection dimension is $k = 150$, $k = 100$, and $k = 50$, respectively. As is clear, as the number of projections decreases the correlation becomes smaller.
• Figure 3: CIFAR-10: Correlation between $\mathcal{I}^{\rm True}$ and $\mathcal{I}^{\rm Linear}$, and between $\mathcal{I}^{\rm Linear}$ and $\mathcal{I}^{\rm ALO}$. Results are aggregated over $100$ held-out test points and $100$ training points ($10{,}000$ pairs).
• Figure 4: Left two panels: Binary logistic regression with $p=100$, $n=1024$ and $n=2048$. The x-axis shows $\mathcal{I}^{\rm True}$; the y-axis shows $\mathcal{I}^{\rm ALO}$. Right two panels: Poisson regression under the same $p$ and $n$ settings.
• Figure 5: Binary logistic regression with projection ($p=100$, $n=1024$ and $n=2048$). The x-axis is $\mathcal{I}^{\rm True}$; the y-axis is $\mathcal{I}^{\rm TRAK}$. From left to right, projection dimensions are $k = 75$, $50$, and $25$.

Theorems & Definitions (22)

• Example 2.2: Linear model
• Example 2.3: Neural network
• Proposition 2.4
• Proposition 2.5
• Theorem 2.6
• Remark 2.7
• Remark 2.8
• Theorem 2.9
• Remark 2.10
• Remark 2.11