GraSS: Scalable Data Attribution with Gradient Sparsification and Sparse Projection

We sincerely thank Reviewer iWVy for taking the time to review our paper and for their constructive feedback and recognition of our proposed algorithm and its potential impact. Below, we provide point-to-point responses to the questions and concerns raised.

The presentation is difficult to follow

We appreciate the reviewer's feedback regarding the paper's presentation and acknowledge that some paragraphs could be clearer. We are committed to improving readability and flow.

We will thoroughly revise the writing in the next version to enhance clarity and ensure the paper is easier to follow. We welcome any further specific suggestions the reviewer might have to help us achieve this goal.

Can GraSS/FactGraSS also be applied to speed up advanced attribution methods such as SGD-Inf?

Yes, our proposed GraSS and FactGraSS gradient compression algorithms are general and broadly applicable to various data attribution methods that rely on per-sample gradients [1, 2, 3, 4].

For instance, recent advancements like [1,4] (which are conceptually similar to SGD-Inf) utilize techniques akin to those in LoGra [2]. Our methods can be seamlessly integrated into such frameworks to achieve significant acceleration.

For clarity of presentation within this paper, we focused our experimental scope primarily on the Influence Function and its variants TRAK [3]. However, integrating GraSS/FactGraSS with other, more recent data attribution methods like SGD-Inf and works such as [1, 4] is a promising application of our proposed methods.

Does the accuracy of current attribution methods still fall short of what is required in real-world applications?

While we agree that the required accuracy for data attribution methods is application-specific, however, we would like to highlight that existing methods already demonstrate practical utility in many scenarios.

Our work's core contribution is to significantly accelerate gradient-based data attribution, achieving state-of-the-art efficiency for well-known methods like Influence Function and TRAK. By resolving the critical computational bottleneck, we enhance the practical applicability of these valuable tools, making them feasible for larger scales. Our focus is on computational efficiency, an orthogonal problem to inherent algorithmic accuracy limitations.

Error and computational bottleneck remain

We appreciate the reviewer's insightful questions regarding the practical viability of data attribution methods. Regarding the level of error, whether it is acceptable remains application-specific. As discussed above, current data attribution methods have already demonstrated their utility and found practical use in various real-world scenarios. Our work focuses on addressing the computational bottleneck to expand their applicability.

Furthermore, we acknowledge that current gradient-based attribution methods inherently require $N$ per-sample gradients, implying a complexity comparable to at least one training epoch. This is widely considered a hard computational lower bound in the field, with recent advancements [1, 2] that achieve attribution within "one training run" being recognized as significant breakthroughs.

While we agree that for the largest frontier models during their pre-training phase, an attribution cost approaching or exceeding a single training epoch might be prohibitive, this is not universally true. For fine-tuning phases or smaller models, this level of computational expense is often affordable and efficient enough for practical application. Our methods significantly enhance this affordability for the scenarios where it is viable.

References:

[1]: Data Shapley in One Training Run

[2]: What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions

[3]: TRAK: Attributing Model Behavior at Scale

[4]: Capturing the Temporal Dependence of Training Data Influence

We thank the reviewer for acknowledging our rebuttal. We hope our rebuttal has well addressed the reviewer’s concerns about our paper, and we are more than happy to address any follow-up comments if there is any. If not, we hope the reviewer can consider raising the rating. Thank you again for your time.

We sincerely thank Reviewer SuDq for taking the time to review our paper and for their constructive feedback. Below, we provide point-to-point responses to the questions, hopefully addressing the concerns raised.

Better presentation of theoretical results

We appreciate the reviewer's valuable suggestions for enhancing the clarity of our theoretical analysis and computational complexity results.

We will update the paper accordingly, incorporating formal theorems to clearly state our method's computational complexity. Furthermore, we will include a comprehensive table comparing the computational costs of GraSS and FactGraSS, and other popular state-of-the-art methods, to significantly aid exposition. A sample table is provided below:

where $p$ denotes the model size, $k$ denotes the target compression dimension, and $k'$ is a parameter such that for GraSS, $k \leq k' \ll p$, and for FactGraSS, $\sqrt{k} \leq k' \ll \sqrt{p}$. We note that both LoGra and FactGraSS are specific for linear layers, and neither requires materializing the layer-wise gradients.

Regarding the influence functions for GraSS and FactGraSS, we will ensure these are presented as numbered equations for improved clarity.

Discussion on the error between GraSS/FactGraSS compared to the full influence function

We appreciate the reviewer's interest in the statistical properties of the proposed estimators. In the literature, state-of-the-art methods like LoGra [1] and TRAK [2] adopt random projection, often motivated by the approximation error guarantees provided by the Johnson-Lindenstrauss (JL) Lemma. Our approach adopts SJLT, which comes with JL Lemma-type guarantees as well [6]. However, a theoretical error bound w.r.t. the full influence function remains a significant open challenge in the current literature for efficient data attribution methods, including LoGra [1] and TRAK [2]. Arguably, most of the gradient-based efficient data attribution methods are driven by empirical development.

Given this context, we believe our proposed methods are well-motivated. Our extensive empirical validation demonstrates competitive performance, highlighting the practical utility and reliability of GraSS/FactGraSS within the current research landscape. Further theoretical development in this area represents a promising direction for future work.

We will include the above discussion in the next version to better motivate our methods.

Quantitative analysis with LLM is unclear

We acknowledge the reviewer's feedback regarding the clarity of our quantitative analysis with LLMs:

• Our experiments include results for GPT2-small (fine-tuned on Wikitext), presenting both LDS scores (in Table 1) and efficiency metrics. For Llama3-8B (evaluated on a subset of OpenWebText), we report throughput (in Table 2) as an efficiency measure.
• Our experimental setup closely adheres to the methodology established by the recent state-of-the-art [1]. We will provide a more detailed and self-contained description of this setup in the updated version of the paper to enhance clarity.

Scope of the work is limited: only proposes a method for efficiently computing influence functions, without theoretical results

We contend that efficiency is a core and pressing challenge for all recent state-of-the-art gradient-based data attribution methods [1, 3, 4], and our proposed GraSS and FactGraSS methods offer broadly applicable gradient compression techniques, which are crucial components in nearly every modern gradient-based attribution approach [1,2,3,4,5]. Thus, addressing this fundamental efficiency bottleneck significantly advances the practical applicability and scalability of a wide range of gradient-based data attribution algorithms.

Regarding the lack of rigorous statistical theoretical results, we would like to clarify:

1. The key components of our proposed methods, particularly SJLT, come with strong theoretical justifications for their approximation quality [6], as discussed in Lines 138-144.
2. While a full statistical analysis of the end-to-end attribution process is indeed complex, as noted in our response to the second response (i.e., "Discussion on the error between GraSS/FactGraSS compared to the full influence function"), the broader development and evaluation of data attribution methods have been largely driven by extensive empirical research [1, 2, 3].
3. We have conducted comprehensive empirical evaluations on two of the most popular gradient-based data attribution methods, Influence Functions and TRAK [2], demonstrating the effectiveness and practical utility of our approach. Therefore, we believe that the proposed methods offer a substantial improvement in a critical area, supported by both component-wise theoretical guarantees and strong empirical evidence, aligning with the current trajectory of research in data attribution.

References

[1]: What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions

[2]: TRAK: Attributing Model Behavior at Scale

[3]: Capturing the Temporal Dependence of Training Data Influence

[4]: Data Shapley in One Training Run

[5]: Scaling Up Influence Functions

[6]: A Sparse Johnson--Lindenstrauss Transform

We thank the reviewer for acknowledging our rebuttal. We hope our rebuttal has well addressed the reviewer’s concerns about our paper and that the reviewer can consider raising the rating. And we would appreciate further feedback from the reviewer on any remaining concerns. Thank you again for your time.

I thank the author for their thoughtful response. In particular, I liked the table summarizing the complexities of the compared attribution methods, and their commitment to improving the presentation of their theoretical results.

However, I'm still not convinced by the discussion regarding the theoretical analysis as presented. While I accept that other papers in this space do not provide rigorous justification, I would like to see, at the very least, some formal analysis of the impact of using the sparse JL transform vs the canonical JL transform. In particular, the role of sparsity should be formally discussed. For instance, $nnz(A)$, there is a maximum number of nonzero entries in the projection matrix for the JL bounds to hold for the SJLT (I believe that it is at least of the order $1/\varepsilon$). So even if the gradients are naturally sparse, it seems odd to me that you can choose $s=1$ (as noted on line 144).

We thank the reviewer for the insightful follow-up, and we would love to discuss how SJLT affects the performance compared to canonical JLT. The short answer is that, SJLT enjoys asymptotically the same JL-type guarantees compared to, e.g., dense Gaussian projection. Furthermore, by letting $s$ as a tunable parameter, our proposed methods still outperform the baselines significantly, both in theory and in practice. Please see below for a detailed discussion.

Recall that the canonical JL Lemma guarantees:

(Informal) For any $d$ and $0 < \epsilon < 1 / 2$, there exists a probability distribution on $k \times d$ real matrices $S \in \mathbb{R} ^{k \times d}$ with $k = \Theta (\epsilon ^{-2})$ such that "for any" $x \in \mathbb{R} ^d$, $\mathbb{P} _S ((1 - \epsilon ) \lVert x \rVert _2 \leq \lVert Sx \rVert _2 \leq (1 + \epsilon ) \lVert x \rVert _2) > 1 - \delta$. Note that we hide the dependency of $\delta$ in all bounds for simplicity.

First, we want to clarify a potential confusion about the roles of “per-sample gradient sparsity” and “projection matrix sparsity” here: observe that the role of per-sample gradient sparsity (i.e., $\lVert x \rVert _0$) is irrelevant in terms of the JL-type guarantees, as the statement is for any vectors $x \in \mathbb{R} ^d$. Instead, the importance of the per-sample gradient sparsity only lies in the computational aspect of $Sx$: the sparser the per-sample gradient is, the faster we can obtain the projected per-sample gradient $Sx$.

With this in mind, we next discuss $s$ for SJLT. [1] shows that there exists such a distribution over $S$ with $s = \Omega (\epsilon ^{-1})$, where $s$ is the number of non-zero entries of each column of $S$. We see that:

1. Choosing $s$ is independent of input sparsity.
2. As these bounds are asymptotic, they inform very little about how to choose $s$ in practice. With extensive empirical testing across different orders of $d$ and $k$, we find that in practice, $s = 1$ usually suffices and achieves relative errors similar to the fully dense Gaussian projection matrix (e.g., Figure 2 in the paper).

On the other hand, keeping $s$ as a tunable parameter will not change the competitiveness of our proposed method much:

1. Complexity of GraSS goes from $O(k’)$ to $O(sk’)$ for some $k \leq k’ \ll \sqrt{p}$ and $1 \leq s \ll k$.
2. Complexity of FactGraSS goes from $O({k’}^2)$ to $O((sk’)^2)$ for some $\sqrt{k} \leq k’ \ll \sqrt{p}$ and $1 \leq s \ll \sqrt{k}$.

In both cases, our proposed methods outperform the respective SOTA baselines easily under any reasonable choice of $s$ in practice. We previously omitted this in our discussion due to the space limit. We will incorporate the above discussion into the next version.

References:

[1]: Sparser Johnson-Lindenstrauss Transforms

We sincerely thank Reviewer 49cd for taking the time to review our paper and for their constructive feedback. Below, we provide point-to-point responses to the questions, hopefully addressing the concerns raised.

Comparison with RapidIn.

We appreciate the reviewer’s suggestion to include a more direct comparison with RapidIn. Here, we provide a detailed discussion and empirical comparison that we will also integrate into the next version.

Theoretical Complexity: In terms of computational complexity, the core RapidIn’s gradient compression method requires $O(\lambda p)$, where $\lambda$ is a hyperparameter that is of order $20$ to $100$, and $p$ is the model size. Compared to our proposed methods, which require only $O(k’)$ computational time for some $k < k’ \ll p$, we argue that our methods outperform RapidIn significantly.

Empirical Comparison: As the actual runtime of each method highly depends on the specific implementation, it is tricky to compare different methods. Nevertheless, we run the official implementation of RapidIn to compare. Below are the results using RapidIn’s public codebase and default hyperparameters ($\lambda = 20$):

MLP+MNIST (RapidIn):

• $k$=2048: LDS=0.1529, Time=4.5928
• $k$=4096: LDS=0.1547, Time=4.6150
• $k$=8192: LDS=0.1554, Time=4.7306

ResNet9+CIFAR2 (RapidIn):

• k=2048: LDS=0.0234, Time=16.6967
• k=4096: LDS=0.0235, Time=16.7008
• k=8192: LDS=0.0223, Time=16.7235

MusicTransformer+MAESTRO (RapidIn):

• $k$=2048: LDS=0.1268, Time=43.1692
• $k$=4096: LDS=0.1338, Time=43.2024
• $k$=8192: LDS=0.1432, Time=43.1993

Compared to Table 1 in our paper, our methods consistently achieve significantly lower projection times, often by several orders of magnitude, while yielding higher LDS values. We again note that the comparison has some caveats: first, our compression algorithms are implemented with a CUDA kernel, while theirs use only high-level PyTorch APIs; second, RapidIn is different from TRAK used in the paper: the former only requires one model checkpoint, while TRAK requires several (5~10) models ensemble, and the reported wall time is the total projection time for each attribution method. Besides all these caveats, we highlight that the CUDA kernel implementation is also a key contribution of this work, and it requires non-trivial efforts to implement a CUDA-kernel-equivalent algorithm for RapidIn.

Experiments focus on only a few models and only one billion-scale model.

We respectfully clarify that our experimental setup includes five different models, each evaluated on a distinct dataset, covering a diverse range of architectures and domains. For the billion-scale setting, we adopt the largest experimental configuration used in LoGra [1], which serves as a strong benchmark for evaluating attribution methods at scale. We believe this setup sufficiently demonstrates the efficiency and scalability of our proposed methods relative to the current state-of-the-art.

Stability and implementation of Selective Mask.

We address the two questions below, respectively:

Choosing Selective Mask: The procedure for learning the Selective Mask is described in Section 3.2.2. Here, we provide additional implementation details and clarifications that were omitted from the main text due to space constraints:

• Forcing exact $k$: A naive optimization for Selective Mask can not guarantee that the final $S^{\ast}$ contains exactly $k$ active indices. A simple workaround is to select the top-$k$ indices based on their sigmoid values, which corresponds to adaptively setting the activation threshold to ensure exactly $k$ active indices.
• Linear Layer: For linear layers, we can derive a factorized Selective Mask by decomposing the mask w.r.t. the gradients factorization, according to Eq. 2. Specifically, one can reformulate Eq. 1 and leverage the Kronecker product structure to simplify inner product calculations. As a result, training Selective Mask does not require computing full layer-wise gradients, providing similar computational and memory efficiency as in FactGraSS. This additional discussion will be included in the next version.

Stability: We refer to the ablation studies in Table 1 (specifically the standalone Selective Mask (SM) columns), where we evaluate SM across multiple settings and varying target projection dimensions. The results demonstrate that SM is stable across runs and consistently outperforms the baseline Random Mask in terms of LDS, highlighting its effectiveness and robustness. Furthermore, we conducted additional ablation studies to assess the sensitivity and impact of using the Selective Mask alone as a gradient compression method on attribution results. Below, we report the mean LDS and the standard error of the mean across five different random seeds:

MLP+MNIST (TRAK with SM):

• $k$=2048: LDS=0.3924 (±0.0030)
• $k$=4096: LDS=0.4192 (±0.0034)
• $k$=8192: LDS=0.4274 (±0.0032)

ResNet9+CIFAR2 (TRAK with SM):

• $k$=2048: LDS=0.3733 (±0.0071)
• $k$=4096: LDS=0.4072 (±0.0028)
• $k$=8192: LDS=0.4326 (±0.0037)

MusicTransformer+MAESTRO (TRAK with SM):

• $k$=2048: LDS=0.1470 (±0.0030)
• $k$=4096: LDS=0.1628 (±0.0039)
• $k$=8192: LDS=0.1651 (±0.0033)

This consistent LDS performance demonstrates the stability of the Selective Mask.

Add quantitative accuracy evaluation for billion-scale models beyond throughput.

We agree that it would be valuable to include additional quantitative analysis beyond throughput. However, computing common metrics such as LDS at the 8B scale is prohibitively expensive in terms of time and memory. This limitation is well recognized in prior work [1], which similarly did not include such measurements at this scale due to the prohibitively high computational cost. Our experimental setup follows closely with [1], aligning with the literature.

Sensitivity of GraSS/FactGraSS to key hyperparameters like $k$ and $k’$.

We have addressed the sensitivity of GraSS/FactGraSS to the hyperparameters $k$ and $k’$, by the experiments in Section 4 of our paper. In these experiments, we consider three different target compression dimensions across various settings, including billion-scale models and datasets. The consistent performance trend observed across these values demonstrates the robustness of our approach to reasonable variations of compression dimensions. Additionally, as discussed in Lines 271-272, we don't include results for larger values of $k$ and $k′$ because increasing $k$ leads to a quadratic complexity blow-up. This would require storing and inverting a matrix that far exceeds our available GPU memory resources. This practical limitation prevents further exploration of such extreme ranges.

In terms of other hyperparameters, we follow the same settings as in [3] to align with the literature.

Why does Random Mask lead to lower accuracy, and when is it effective?

Random Masking, along with Selective Masking and SJLT, serves as a component within our proposed algorithm's pipeline. The experiments involving Random Masking are primarily designed as ablation studies to facilitate comparisons with other algorithmic components and our final proposed methods.

On the other hand, while Random Masking may lead to lower accuracy compared to more sophisticated compression techniques, its key advantage lies in its exceptional efficiency. This makes it the most efficient compression method within the pipeline. Consequently, its inclusion provides a valuable option for scenarios where a trade-off favoring computational efficiency over maximal accuracy is acceptable or necessary.

References

[1]: What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions

[2]: TRAK: Attributing Model Behavior at Scale

[3]: 𝚍𝚊𝚝𝚝𝚛𝚒 : A Library for Efficient Data Attribution

We thank the reviewer for acknowledging our rebuttal. We hope our rebuttal has well addressed the reviewer’s concerns about our paper and that the reviewer can consider raising the rating. And we would appreciate further feedback from the reviewer on any remaining concerns. Thank you again for your time.

We sincerely thank Reviewer HgYa for taking the time to review our paper and for their constructive feedback and recognition of our proposed algorithm and its potential impact. Below, we provide point-to-point responses to the questions and concerns raised.

Only 1 billion-scale model in empirical analysis.

Due to computational constraints, we focused our quantitative analysis on smaller models and efficiency analysis on an 8B-parameter model. It is worth noting that our experimental setup follows closely to the previous state-of-the-art [1]. Specifically, for efficiency analysis, we highlight that:

1. The efficiency analysis is conducted on the 8B model, the largest model [1] considered.
2. The observed speedup is robust and model-agnostic, as supported both by our computational complexity analysis (independent of model) as well as the empirical evidence (up to overall 165% throughput speed up).

Hence, we believe the quantitative analysis results, as well as efficiency results for 8B model, provide a representative and reliable demonstration of the efficacy.

Possibility to perform quantitative analysis for 8B model beyond throughput.

We agree that it would be valuable to include additional quantitative analysis beyond throughput. However, computing common metrics such as LDS at the 8B scale is prohibitively expensive in terms of time and memory. This limitation is well recognized in prior work [1], which similarly did not include such measurements at this scale due to the prohibitively high computational cost.

Reference

[1]: What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions