Large-Scale Training Data Attribution with Efficient Influence Functions

We appreciate your reviewing effort. Here, we attempt to address your concern as best as we can.

---

Limited Applicability to Specific Layer Types

In L174, we noted that most popular layers in neural networks including convolutional and linear layers can be formulated as 2D matrix multiplication. In Section 4.1, we indeed performed LDS and brittleness experiments with ResNet-9, which include many convolutional layers.

---

Potential Impact of Low-Rank Projections on Attribution Accuracy

In Section 4.1, we quantitatively compare EK-FAC (which does not use low-rank gradient projections) and our proposed method LoGra. We show that Logra obtains comparable performance to EKFAC, despite being order of magnitude more efficient (please see Table 1). In our updated manuscript, we will run additional LDS experiments on the effect of LoGra rank.

---

Lack of Discussion on Dimensionality Reduction Efficiency

Thanks for your comment. To compute the covariance matrix for PCA, we ran one additional training epoch to collect forward/backward activations for all training examples. The eigendecomposition of this covariance matrix is computationally efficient - in our experiments with large-scale models like Llama-3 8B, this step required less than one minute to complete. Furthermore, since the eigenvectors are only used for the initialization of our LoGra encoder/decoder, this eigendecomposition step is a one-time computation that does not need to be repeated. Lastly, while it might be appealing to ensemble the scores when using random projections, we note that our results in Section 4.1 did not use any ensembles and got comparable performance with the PCA.

Thanks again for your review, and we hope our rebuttal cleared your concern and made you evaluate our work more positively!

Thanks for the comment.

---

Limited Applicability of Specific Layer Type

Following your suggestion, we will clarify in our final manuscript that LoGra is mainly designed for layers/modules that involve matrix multiplication (e.g., linear and convolutional layers), given that a vast majority of parameters in most large networks are from them.

---

Trade-Offs Between Projection Ranks, Performance, and Efficiency

We want to note that the result in our previous comment actually shows the effect of varying low-rank dimensions for our LDS experiments in Figure 4.

• • We note that both $x_t$ and $Dx_o$ are vectors, and thus their Kronecekr product is also a vector. More specifically, Kronecker product of m-dimensional vector and n-dimensional vector yields mn-dimensional vector. We believe there is no error in our equations.*

I always assumed that for two vectors $x \otimes y = x y^{\top}$ is a matrix. Suppose you are right, does it mean that Kronecker product of two matrices is a 4 dimensional tensor?

UPDATE: I see where I'm wrong, that checks out.

• • To answer this question, let's first define the "unit" of the example. In our language modeling experiments, we adopted each "sequence" as our unit of the example. *

This appears to be very different from the conventional GNH. Are there any comparison and where in the paper do you define what your GNH is? The traditional GNH comes from local quadratic approximation of the KL divergence, which is also averaged over all tokens. Can your GNH have similar interpretation?

UPDATE: also note the difference between FIM and empirical FIM. You would have to generate the whole sequences \hat{y} if you view them as a unit.

Thanks for your reviewing effort! We try our best to address the raised questions and concerns in our response. We hope this helps you evaluate our work more positively :)

---

Q. While reading the paper, I often found myself having to refer to the previous papers to clarify some of the formulas. I think an average reader would benefit from self-contained notation: what is Kronecker product, what is Kronecker product of vectors, background for K-FAC approximation of the Hessian, and what is Hessian itself (the Gauss-Newton one).

A1. Thanks for the suggestion. We acknowledge that some parts of our manuscript may not be easily digestible, especially for readers who are not familiar with the Hessian and Kronecker product. In our final manuscript, we will include a separate appendix section that introduces notations and basics of the Gauss-Newton Hessian and the K-FAC approximation of it.

---

Q2. The details of Hessian calculation have been omitted. I do not understand how this calculation was performed, in particular, how do the authors manage to calculate the gradients token-wise.

A2. To answer this question, let's first define the "unit" of the example. In our language modeling experiments, we adopted each "sequence" as our unit of the example. This choice was motivated by the fact that we want to identify the most influential sequences instead of tokens in our experiments. Therefore, the gradient is also computed at the sequence level instead of the token level (this can also be inferred from Equation (5)). For implementation details, we computed per-sequence gradient using PyTorch hooks.

Next, we adopt the Fisher information matrix approximation of the Hessian. For convex loss functions, such as cross entropy and mean-squared error, the Hessian equals the Gauss-Newton Hessian. Furthermore, for exponential family negative log-likelihood functions, such as cross entropy, the Gauss-Newton Hessian is equivalent to the Fisher Information Matrix (FIM). We refer the reviewer to [1] for a more technical justification. Hence, we compute the block-diagonal FIM by simply averaging the outer product of each training sequence gradient following the definition of FIM. We would like to highlight that this Hessian approximation with Fisher is common in the literature [2, 3, 4].

---

Q3. In equation (5), should $x_t \otimes Dx_0$ not be a matrix? However, the left-hand side (LHS) is a vector. The same applies to Equation (6). I believe everything would be correct if the vec(⋅) operation were removed—am I right?

A3. We note that both $x_t$ and $Dx_o$ are vectors, and thus their Kronecekr product is also a vector. More specifically, Kronecker product of m-dimensional vector and n-dimensional vector yields mn-dimensional vector. We believe there is no error in our equations.

---

Q4.1. I believe there are some limitations in the interpretation of Lemma 1 that could be mentioned. First, … The assumption that they follow the same distribution is impractical, as training and generated texts are typically very different. Additionally, the variance is calculated per token, while g_{tr} and g_{te} are both averages over multiple dependent tokens.

A4.1. We first remind the reviewer that gradient for our language model experiments are computed at the sequence level. The model output is sampled from $p(x|\theta^*)$, and $\theta^*$ is obtained by solving $\arg\min_\theta L(D_{tr}|\theta)$. Combined with the fact that pre-trained language models have seen a variety of "prompts" from web-scale data, we believe that there is a meaningful overlap between training data distribution and the model output distribution. However, we also agree with the reviewer that their distributions may not perfectly overlap.

Q4.2. coefficient distribution and normalization. We agree with the reviewer that while $E[c_i^2]\approx 1$ holds in our derivation, its practical implication can be obfuscated, especially when $c_i$ follows the heavy-tailed distribution. In fact, we suspect that the heavy-tail noise in coefficients can be one potential root cause of failure cases in Section 4. In general, we believe studying the gradient projection in influence functions under heavy-tailed noise would be an interesting future research direction.

---

If there is anything that can help you evaluate our work more positively, please let us know.

[1] Martens, James. New insights and perspectives on the natural gradient method. JMLR, 2020. [2] Bae et al. If influence functions are the answer, then what is the question?." NeurIPS, 2022. [3] Kwon et al. Datainf: Efficiently estimating data influence in lora-tuned llms and diffusion models. ICLR, 2024. [4] Park et al. Trak: Attributing model behavior at scale. ICML, 2023.

Thanks a lot for the clarifications!

My question about GNH in the case where your treat each sequence as a single "unit" still stands.

Let's say we only consider the case of classification, hence GNH and FIM are equivalent, and both approximate the KL (in the sense they define the quadratic form that approximate the KL). If you treat a whole sequence as a unit, this will be a different KL than what usually used in proximal methods. In fact it will be rather intractable version of KL, since it would require calculation of all $p(y)$ for all y which grows exponentially with sequence length. It seems to me that in that case the approximation by PBRF will break, if the divergence in PBRF is calculated per token [1, 2]. Which by the way I was wondering why the comparison to PBRF is not included in your evaluations.

Q1: If you treat the whole sequence as a "unit", then to calculate FIM we need to take the expectation of $\nabla \ell(\hat{y} | x)$. What is $x$ in that case?

Q2: In your response you mention that empirical FIM and FIM behave similar. Do you mean that you actually implemented the case where you generate the whole sequences for calculation of FIM? How did you do it? Usually FIM is approximated by empirical FIM in the realizible case. I don't think that on sequence level you can treat it as realizable.

Q3. I also looked more closely at evaluations in Section 4. In the GPT2 case, the random LoGRA generally performs not worse or even better than PCA LoGRA. Does it mean that some assumptions of Lemma 1 can be broken? Lemma 1 seems to suggest that PCA is the right thing to do.

[1] Bae et al. If influence functions are the answer then what is the question

[2] Grosse et al. Studying large language model generalization with influence functions

Thanks for the response! In our comment below, we attempt to address Q1 & Q2 together, and Q3 separately.

---

Q1: If you treat the whole sequence as a "unit", then to calculate FIM we need to take the expectation of $\nabla \ell (\hat{y}|x)$. What is $x$ in that case?

Q2: In your response you mention that empirical FIM and FIM behave similar. Do you mean that you actually implemented the case where you generate the whole sequences for calculation of FIM? How did you do it? Usually FIM is approximated by empirical FIM in the realizible case. I don't think that on sequence level you can treat it as realizable.

A. We again note that the FIM is defined as covariance of the score function, and it coincides with the KL divergence under a specific condition. In our work, we computed the FIM following its definition (i.e., covariance of the score function), instead of the KL interpretation.

More formally, the loss function (or negative log likelihood) of each sequence is defined as follows:

$L(x|\theta) = -\log p(x|\theta) = - \sum_{t=1}^T log_{\hat{y}|x_{1:t-1}}(x_t|x_{1:t-1};\theta)$

Therefore, when we take each sequence as a unit, the gradient for each unit corresponds to the sum of token-wise gradients. For the FIM computation, we directly average the outer product of each sequence gradient. To compute the true FIM, we sample the next-token at each time step from the model output distribution, and performed the Monte-Carlo approximation over many training sequences.

---

Q3. I also looked more closely at evaluations in Section 4. In the GPT2 case, the random LoGRA generally performs not worse or even better than PCA LoGRA. Does it mean that some assumptions of Lemma 1 can be broken? Lemma 1 seems to suggest that PCA is the right thing to do.

A. We discussed this issue in L338-347. The implication of Lemma 1 is that we can expect better influence estimations when we successfully keep larger components of the Hessian/FIM. To make PCA compatible with LoGra, we adopt the KFAC approximation of the FIM (L238-239). However, unlike naive-MLP (without weight sharing) or convolutional networks where there exist specialized KFAC techniques, we note that the Transformer architecture lacks the specialized KFAC approximation. Therefore, there is an increased chance that our Kronecker product estimation of the FIM may be inaccurate in the first place, and thus larger components we computed with our inaccurate Kronecker product estimation may not be actually larger components. Exploring the specialized KFAC for the Transformer architecture would be an interesting future research problem.

---

We hope our comment addressed some of your concerns. Let me know if you have any other questions!

To compute the true FIM, we sample the next-token at each time step from the model output distribution, and performed the Monte-Carlo approximation over many training sequences.

If you sample next-token, then the gradients have to be calculated per-token. If you calculate the whole gradient, then the whole sequence needs to be sampled. What you are describing is just some conditioner, I don't see how you can interpret it as FIM. If you claim that two things are the same empirically, I would love to see that comparison. The paper also needs to describe precisely what is your Hessian in the main part and how you motivate this choice.

It looks there is too many issues in this paper. I'll lower the score for now, maybe it will change later after discussion with reviewers and AC.

We appreciate your valuable and positive feedback. Here, we attempt to address your questions.

---

Q1. Pythia-1.4B Failure Cases

In Appendix A.3 of the original manuscript, we presented several hypotheses to explain the suboptimal qualitative results, particularly those obtained with the Pythia-1.4B model. To facilitate the review process, we would like to reiterate the relevant excerpt from the appendix:

Influence functions tend to give a high score for the example that contributes most to decreasing (test) loss at the current weight [1]. At the same time, it is also hypothesized that different layers learn different concepts at different stages of training [2]. Combining these two facts, when interpreting influence analysis results, we need to think about which features the model is most likely learning at the current weight. Here, we specifically discuss two factors: training data quality and training steps. First, if the training data quality is low, then there would be a lot of features (e.g., random email address) that are frequent enough in the training dataset to be considered as learnable patterns. In other words, even though these features look redundant to humans, they may still be useful for decreasing loss from the model perspective. Second, many LLMs are only pretrained for a single epoch, or under-trained to their pretraining dataset. That being said, redundant features from the first point would likely still remain as learning-worthy features at the end of training and are captured by influence functions. In sum, we hypothesize that as the model is well-trained on a high-quality dataset, influence functions would capture more similar data to the query LLM output.

---

Q2. Storage Cost

In L391-400 (and footnote 2), we discussed the trade-off between the storage cost and memory/compute efficiency. In particular, a simple cost analysis based on hourly AWS rates for GPUs and storage in footnote 2 shows that trading-off the storage cost for compute/memory efficiency is largely favorable for practitioners in most cases. If you have any further suggestions on making this discussion clearer, we would greatly appreciate it!

---

Q3. Theoretical Interpretation of Damping

Indeed, our theoretical analysis shows that (overly) aggressive gradient projection can result in significant information loss and an increased likelihood of failure cases. In addition, our theoretical analysis shows that, for a fixed projection dimension, retaining larger components of the gradient is generally preferable. This insight motivated us to design the PCA initialization scheme for LoGra, which leverages the Kronecker-factored approximate curvature (KFAC) approximation of the Hessian, as described in lines 234-246 of the manuscript. Lastly, we investigated the effectiveness and limitations of the PCA initialization in LoGra in the small-scale quantitative experiments from Section 4.1. In L338-347, we discussed that Transformer architecture lacks the specialized KFAC approximation for the Hessian, which led us to use random initialization in our large-scale qualitative experiments.

---

We again greatly appreciate your reviewing effort. If you have any suggestions that could make our paper stronger, do not hesitate to let us know!

[1] Bae et al., If influence functions are the answer, then what is the questions? NeurIPS, 2022 [2] Chen et al., Which layer is learning faster? A systematic exploration of layer-wise convergence rate for deep neural networks. ICLR, 2023.

Thanks for your positive and valuable feedback. Here, we attempt to address your comments:

---

The order complexities listed in sections 2 and 3 can be hard to process since the constants are hidden. It is not clear if I can compare two order complexities directly. Is there any way to make this more concrete and easily comparable?

To address this exact issue, we provided a concrete comparison using one of the models in our experiment (i.e., Llama3-8B) in lines 199-201: "To clearly see this benefit, given the model/projection sizes of 8B/4k, we note that projection matrix sizes are about 1GB and 128TB respectively for LoGra and naive projection."

---

I would recommend adding compute and memory complexity for the methods compared in Figure 4 to get a more complete picture.

We appreciate your thoughtful feedback. In the final manuscript, we will include a detailed analysis in the Appendix covering both the computational and memory complexity of influence calculations, as well as the gradient projection requirements for each method.

---

Please note why quantitative evaluations are not possible for accuracy measures in section 4.2.

The quantitative evaluations proposed in Section 4.1 would require retraining each model at least 500 times, which presents two significant challenges. First, we lack access to the (pre-)training datasets for two of the three models examined in Section 4.2 (GPT2-XL and Llama3-8B). Second, even if we had access to these datasets, the computational resources required to retrain such large models 500 times would be prohibitively expensive and impractical.

---

If you have any further concerns, feel free to let us know. We appreciate your reviewing effort again!