Efficient Sketches for Training Data Attribution and Studying the Loss Landscape

Summary

The paper proposes efficient sketching algorithms designed to address memory constraints in large-scale models. The key idea is to eliminate the multiplication of a dense projection matrix multiplication (large matrix materialization) found in the existing sketching algorithms such as in FJL and FastFood. The authors provide theoretical guarantees for their proposed algorithms: AFFD, AFJL, and QK. They showcase the advantages of these methods in three different contexts: training data attribution, intrinsic dimension estimation, and eigenvalue estimation of the Hessian.

Strengths

- The paper addresses an important topic that is likely to be of interest to the NeurIPS community, with potential applications across various domains. - The authors provide a well-motivated problem statement, clearly articulating the need for efficient sketching algorithms. - The results presented are comprehensive and interesting, covering multiple aspects of the proposed algorithms and demonstrating their effectiveness across several use cases. - While the authors have not provided code with the submission, the appendix contains sufficiently detailed information to enable the reproduction of the results.

Weaknesses

- While the paper is generally well-written, there are several areas that require improvement: a) Some notations are not clearly explained. b) Certain statements lack precision, potentially leading to ambiguity (specific examples are provided in the detailed comments section). - The authors do not sufficiently address the limitations of the proposed approach. The brief mention in the conclusion does not adequately address the potential drawbacks or constraints of the method (e.g., the limitations of QK). - The paper would benefit from a more comprehensive comparative analysis, such as the sketching algorithm used in TDA [1, 2]. [1] Park, S. M., Georgiev, K., Ilyas, A., Leclerc, G., & Madry, A. (2023). Trak: Attributing model behavior at scale. arXiv preprint arXiv:2303.14186. [2] Xia, M., Malladi, S., Gururangan, S., Arora, S., & Chen, D. (2024). Less: Selecting influential data for targeted instruction tuning. arXiv preprint arXiv:2402.04333.

Questions

I have also included suggestions in addition to questions. - The authors should be more specific when discussing Training Data Attribution (TDA) methods. Rather than using the broad term "TDA," they should explicitly mention that their work focuses on gradient-based TDA methods, such as TracIn [3] and GradDot [4]. TDA is a wide-ranging concept that encompasses various approaches, some of which do not align with the description provided in line 22. - In the introduction, the authors state that random projections have been implemented using dense matrices for TDA. Several works perform random projections without explicitly materializing dense matrices. The same holds for the experiment section in 5. - The statement “memory constraints limited their investigation of ID in generative tasks to 500k” in line 98 is unclear. What does ID stand for? - In line 103, the authors claim that their work makes influence function computations more efficient. Please clarify this point, considering that influence functions typically require inverse Hessian-vector product (IHVP) estimation, which isn't directly addressed in Section 5.1. Is this claim primarily based on improving methods like the Arnoldi iteration? - In Section 5.1, the conclusion that "layer selection coupled with dense projections faces severe scalability limitations" relative to AFFD, QK, and AFJL is not clear. Consider including tabular results (similar to Table 5) in the main paper to support this claim. - Do you anticipate any challenges in applying these techniques to even larger language models? Minor Comments - Line 129: Replace "D" with "$D$" for consistent mathematical notation. - Line 150: Clarify that $B$ and $H_D$ now have dimensions $D \times D$. - Line 172: Specify the relevant section when mentioning "unacceptable early TPU results" for clearer cross-referencing. - Line 227: Add a period at the end of the sentence. - In the case of TDA, the rankings between the gradient dot products are important. Are there any reasons why authors do not use the Spearman correlation? - In line 311, it would be helpful if the authors provided a direct explanation of why explicit sketching provides substation speed up. - Figure 1 has low resolution, making it difficult to read when printed. [3] Pruthi, G., Liu, F., Kale, S., & Sundararajan, M. (2020). Estimating training data influence by tracing gradient descent. Advances in Neural Information Processing Systems, 33, 19920-19930. [4] Charpiat, G., Girard, N., Felardos, L., & Tarabalka, Y. (2019). Input similarity from the neural network perspective. Advances in Neural Information Processing Systems, 32.

Rating

6

Confidence

4

Soundness

2

Presentation

2

Contribution

3

Limitations

The paper would benefit from a more comprehensive discussion of limitations.

Summary

The authors present a framework for scalable gradient sketching and HVP sketching. They introduce three algorithms AFFD, AFJL, QK and provide guarantees for sketching. The paper focuses on three applications: training data attribution, intrinsic dimension computation, and Hessian spectra analysis. These are all implemented with a focus on pre-trained language models.

Strengths

* The scalable approaches to sketching are well-motivated and interesting. It seems important to figure out useful sketching routines that scale to high-dimensional parameter spaces that we are seeing in pre-trained language models. * The theory appears to be a strength of the paper. It appears that there are two approaches taken in the paper, the direct approach that implements a gradient in the original parameter space and the implicit approach that implements a gradient in the sketched dimension. It is interesting to see the comparison made between these two different approaches. * The introduction of the FFT as a faster pre-conditioner appears novel.

Weaknesses

* Presentation is a weakness of the paper. Since the paper is proposing three different algorithms along with three different applications, it makes it challenging for a reader to follow the contributions and comparisons being made. For example, there seems to be two separate topics being discussed: 1) New theory/algorithms for scalable sketching, and 2) Implications for Pre-trained large language models. Additionally 2) aims to cover three different insights including layer selection, intrinsic dimension, and LLM Spectra. It is a challenge to cover all of the aforementioned points in enough detail in a 9 page paper. * Table 1 is not clear how it is related to the proposed algorithms. In particular, TDA (section 5.1) just writes in bold “Our findings indicate the unreliability of layer selection (Table 1)”. Layer selection does not seem to be well defined, nor related to AFFD, AFJL, QK. * Table 3 is also unclear. Which algorithm should one use in practice? * Minor: Figure 1 is hard to read and is too small. A single sentence in the caption telling the reader why these plots are significant (rather than just a description) would help link it back to the contributions. * Minor: Related to presentation is the use of citations and acronyms. It seems like many acronyms are not defined, e.g. AFFD, AFJL, and QK. Also, citations should not read “While [11] developed…”. It should be “ While X et al. 2024 [11] developed…”. There is a command for this in Latex.

Questions

* Which components of the 6 contributions in the introduction are the main contributions of the paper? * Just above equation (2), “A sketch of the HVP can be obtained as…”: Is this a contribution of the paper, or is there a reference to this available? * In practice, which of the algorithms should one use and why? * Could the authors define what the meaning of an influence score is for this paper?

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors focus on scaling their approach to GPUs and TPUs and adequately talk about the memory constraints and improvements made by each approach. No code was included with the paper.

Summary

This paper introduces new methods for sketching high-dimensional gradients and HVPs. These are important building blocks for tools like training data attribution and Hessian spectrum analysis. Their methods introduces both empirical and theoretical improvements over prior methods, and are used to demonstrate new insights about properties of pre-trained LMs.

Strengths

- Clearly written and organized - Addresses a practically relevant and timely technical problem (projecting high-dim gradients and HVPs) - Provides interesting new insights (e.g., that intrinsic dimension is not that low for LMs)

Weaknesses

Nothing major, but: - Misses some important related work ([1], [2]) - but gives good coverage otherwise - Gradient inner product is not a good TDA estimate (many papers have found this now, e.g, [1][3]) - Estimating inner product is fine, but be clear that's the goal then [1] TRAK: Attributing model behavior at scale https://arxiv.org/abs/2303.14186 [2] Faithful and Efficient Explanations for Neural Networks via Neural Tangent Kernel Surrogate Models https://arxiv.org/abs/2305.14585 [3] Training Data Attribution via Approximate Unrolled Differentiation https://arxiv.org/abs/2405.12186

Questions

- Would benefit from better exposition of the different methods in Section 3. Some type of figure would be more useful to get intuition for differences between different methods (it's a bit hard to keep track with so many different symbols). Happy to revisit if there's a newer version. - Some missing text, e.g., there seems to be no transition at the start of S4

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

none

Summary

Gradient information is useful for various tasks such as training data attribution and intrinsic dimension analysis, but often suffers from huge compute/memory costs, limiting its practical utility. This paper focuses on several popular sketching approaches (e.g., FJL, FFD), points out lookup-based memory as the major bottleneck in enabling efficient computations on modern accelerators, and proposes the improvement using the Kronecker products with theoretical guarantees. Their scalable sketching algorithms are applied to downstream applications including training data attribution, intrinsic dimension estimation, and Hessian analysis, and provide various useful insights, some of which are contradictory to existing common beliefs.

Strengths

I agree that gradient information is useful in understanding important characteristics of neural networks, and sketching is a promising approach to reduce the cost of storing/computing gradient information. Indeed, (random) low-rank gradient projection has been studied in several prior research (e.g., Arnoldi IF, TRAK), with similar downstream applications in mind. That being said, I believe one of the main contributions of this work is in improving efficiency of existing sketching algorithms including FJL and FFD. Their solution of using the Kronecker product is, in my opinion, a smart strategy. Furthermore, their theoretical analyses are very interesting and make their arguments more rigorous. With scalable sketching algorithms, they were able to scale various analyses that were only applicable to small-scale networks to much larger networks, and provided many useful insights. For instance, many people have previously observed that full fine-tuning often achieves better performance than LoRA fine-tuning in LLMs, and the intrinsic dimension analysis in this paper can experimentally corroborates such observation. Overall, I believe the proposed algorithms would be valuable tools for understanding various phenomena in neural networks at scale.

Weaknesses

In overall, I believe this paper is a strong paper *if* I limit the scope to gradient sketching. However, I am confused by several arguments. 1. TDA (Sec 5.1): The authors claimed that layer-selection for TDA is inadvisable by pointing out the low (Pearson) correlation with the full gradient dot product between (x, z) pairs. While a high correlation with full gradient dot products could indicate accurate sketching, it doesn't necessarily imply better TDA performance. There are better options for TDA evaluation such as linear data modeling score from TRAK, mislabel detection, data subset selection, brittleness tests, etc. 2. Sketching for HVP: While the authors cited matrix sketching for dealing with higher-order gradient sketching, I am unsure if theoretical guarantees apply in this case. For instance, if we adopt the Fisher information approximation of the Hessian, we can understand the Hessian as the gradient covariance matrix. If the vector in HVP is some sort of gradients, we can understand HVP, especially in influence functions, as a whitening operation that makes all components equally important. In this case, information loss from sketching is not negligible, and the general sketching argument (ie Eq(1)) doesn't really hold? 3. Clarity: I find notations in the paper to be a bit confusing. For instance, in Eq. (2, 3, 6), they used respectively P, F, \Phi to indicate the sketching process.

Questions

1. I wonder how easy it is to combine proposed sketching algorithms with (data-parallel) distributed computing. While I appreciate general efficiency improvements in AFJL and AFFD, I believe distributed computing becomes necessary at some point for scaling these analyses to large-scale networks and datasets. Can you enable distributed computing simply by wrapping the model or do you need special implementation tricks? 2. How are ground-truth eigenvalues are obtained in Table 1? If those eigenvalues are not exact (e.g. approximated using some iterative methods), then being close to approximate eigenvalues do not necessarily imply better accuracy?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have appropriately discussed limitations.

Reply

I thank the authors for their reply and acknowledge reading their response. As the authors mentioned, it would be helpful to be more explicit about the TDA in the motivation of the paper (and give a more detailed description of why this also applies to quantities derived by preconditioning the gradient), add explicit comparisons to Trak's random projection (as the author did for the general rebuttal and the response to my review), and address other details the authors mentioned they will fix/add in the next revision of the manuscript. In general, it seems that most reviewers, including myself, had trouble following the notations used in the paper, which made the presentation of the work on the weaker side. Relatedly, while minor, I agree with reviewer 9kYE that it was difficult to follow all the acronyms in the paper (e.g. AFFD, AFJL, and QK), and it would be helpful to fix issues with \citep vs \citet. It would make the paper more clear if the authors could address these issues. Given these issues will be resolved, as the authors promised, I have increased my score (but it was difficult for me to give a higher score since the current process does not allow the reviewers to see the final manuscript).

Response

> We focused on models that don't require partitioning weights across devices. For larger models, sketching can be applied to individual partitions. We commit to including an example of lifting single-device code to multi-device code in JAX in Appendix B. Thanks. This would be helpful. > We follow [11, Krishnan et al.], as exact eigenvalues are infeasible due to model size. This estimation is accurate, with error [11, eq(7)] decreasing exponentially in the number of iterations. Thanks for the answer. > Theoretical guarantees for matrix sketching apply here. The Hessian has limited bulk [11, Krishnan et al.], exploited by Arnoldi-based influence functions [21, Schioppa et al. ]. Sketching provides guarantees for approximating matrix spectrums [22, Swartworth et al.] and solving linear systems [20, Sarlos, Thm 12]. I agree with the authors that the theoretical guarantees for sketching and apply, **if** we separately consider the Hessian and the gradient. However, when sketching is applied to both the Hessian and train/test gradients simultaneously in influence computations, there would be an inherent information loss due to the whitening effect of the inverse Hessian---all components in gradients become equally important from the whitening effect of the inverse Hessian. > TRAK authors make the following point: as long as we preserve the inner products to sufficient accuracy, the resulting system has approximately the same evolution as the original one [they refer to tracing the training dynamics, our note]. This justifies replacing the gradient features with their random projections This is true if we only consider the short horizon. In the long horizon, however, small components that contribute minimally to naive dot product may still have a non-negligible influence on the overall training dynamics. > Would replacing P, F, and with throughout the paper address this concern? I believe this would improve the readability. Considering all these, I am willing to increase my score to 6. Thanks again for your comment.

Thanks for the Rebuttal

In light of the additional results and the comments above, I will increase my score. I am still concerned with the overall presentation, but I am hopeful the authors will be able to adjust this.