On Support Samples of Next Word Prediction
Authors: Yuqian Li, Yupei Du, Yufang Liu, Feifei Feng, Mou Xiao Feng, Yuanbin Wu
Accepted at: acl25

On Support Samples of Next Word Prediction

Yuqian Li1*, Yupei Du2*, Yufang Liu1,
Feifei Feng3, Mou Xiao Feng3, Yuanbin Wu1,
1School of Computer Science and Technology, East China Normal University,
2Utrecht University, the Netherlands,
3Midea Group,
liyuqian44@gmail.com, y.du@uu.nl, ybwu@cs.ecnu.edu.cn

Abstract

Language models excel in various tasks by making complex decisions, yet understanding the rationale behind these decisions remains a challenge. This paper investigates data-centric interpretability in language models, focusing on the next-word prediction task. Using representation theorem, we identify two types of support samples—those that either promote or deter specific predictions. Our findings reveal that being a support sample is an intrinsic property, predictable even before training begins. Additionally, while non-support samples are less influential in direct predictions, they play a critical role in preventing overfitting and shaping generalization and representation learning. Notably, the importance of non-support samples increases in deeper layers, suggesting their significant role in intermediate representation formation. These insights shed light on the interplay between data and model decisions, offering a new dimension to understanding language model behavior and interpretability.1

1 Introduction

Language models make decisions. From selecting an answer from QA benchmarks to generating reasoning steps for grade school math problems in GSM8k, language models are applauded and criticized for their decision-making ability. Like all other AI systems, understanding how decisions are made is an important research topic.

Support of decisions can be explored in different parts of language models. One common attribution is the patterns of neuron activations (mechanistic interpretability): some neurons are more sensitive to a specific decision and some are not. The circuit formed by active neurons mechanistically explains the model’s decision (Elhage et al., 2021). Mechanistic interpretability targets on model parameters (and hidden states), however, it is not the only possible place to explore the rationale behind decisions. From the perspective of compression, parameters are the compressed version of training data. So, a natural question is whether we can trace model decisions back to their remote origin, the data.

We follow this point of view and study data-centric interpretability of language models. Concretely, we will focus on decisions about choosing the next word (the essential decision of language models), and try to answer the question “Which samples contribute the most when the model decides to predict a token \( v \)?”. For general machine learning problems, there are two typical methods: one is based on counterfactual argument (a sample is important if the decision will change if it is removed from the training set (Koh and Liang, 2017)), another is based on representation theorems (a sample is important if it occupies a large part of prediction parameters (Yeh et al., 2018)). The first one requires computing Hessians of model parameters which could be expensive for analyzing large language models (Grosse et al., 2023b). We thus adopt the more efficient representation-based methods. Specifically,

• We illustrate the contribution of different samples to model parameters based on a simple representation theorem for the next word predictor. The results remind us of two types of support samples (those contribute a lot), one attracts the predictor to predict \( v \), while another pushes the predictor away from predicting \( v \).

• We show that, given a dataset and a model configuration, being a support sample is an intrinsic property of that sample, in the sense that 1) non-support samples have limited influence on learning support samples, and 2)

* Equal contribution.
1Our source code is publicly available at https://github.com/liyuqian44/On-Support-Samples-of-Next-Word-Prediction.
it is possible to predict (with 80% accuracy) whether a sample is a support one at the very early phase of language model training, even without any training.

• We find that though non-support samples are less noticeable in the next-word predictor, it may play an important role in controlling generalization and representations learning. Specifically, we observe that 1) without any non-support samples, learning the predictor will suffer from unnecessary overfitting, and a small number of non-support samples can greatly alleviate the problem. 2) the proportion of non-support samples increases as the layer goes deeper, which means some of them are support samples of (thus contributing a lot to) intermediate representations.

2 Next Word Predictor Representation

We aim to find important samples to support a language model’s next-word decisions. To measure importance, we first connect parameters (in particular, prediction heads at the last layer) and training samples with a representation theorem, then support samples are defined to be those samples contributing significantly to parameters.

Denote \((\mathbf{x}, \mathbf{y})\) to be a training sample, where \(x = x_1, x_2, ..., x_{t-1}\) is a sentence prefix, \(y = x_t \in V\) is the target token to predict given the prefix, and \(V\) is the vocabulary. \(D = \{(\mathbf{x}_i, y_i)\}_{i=1}^N\) is a set of samples. We consider a decomposition of language models with a representation function \(\phi(\mathbf{x}) \in \mathbb{R}^d\) and a token prediction function \(f(\mathbf{x}) = \arg\max_v p(v|\mathbf{x})\), where \(p(v|\mathbf{x}) = \frac{1}{Z} \exp \theta_v^T \phi(\mathbf{x})\). Let \(\theta_\phi, \theta_v\) be the parameter of \(\phi(\cdot)\) and \(f(\cdot)\), \(\theta_V = [\theta_{v_1}; \theta_{v_2}; ...; \theta_{v_{|V|}}]\) pack all predictor parameters (also called language model heads), and \(\theta = [\theta_V; \theta_\phi]\) pack all parameters.

Theorem 1 (Crammer and Singer 2001; Yeh et al. 2018 ). Assume \(\theta\) is a stationary point of the loss function \(L(\theta, D) = -\frac{1}{N} \sum_{i=1}^N \log p(y_i|x_i) + \lambda \| \theta \|_2^2\), then \(\theta_v\) equals
\[
\theta_v = \frac{1}{2N\lambda} \sum_{i=1}^N (\mathbb{1}(y_i = v) - p(v|x_i)) \phi(x_i)
\]
where \(\mathbb{1}(\cdot)\) equals 1 when the argument is true, 0 otherwise.

The theorem gives a decomposition of the predictor \(\theta_v\) using training samples \(\phi(\mathbf{x})\). The coefficients of samples \(\alpha_i \triangleq \mathbb{1}(y_i = v) - p(v|x_i)\) describe their importance in \(\theta_v\). A large \(\alpha_i\) implies a bigger influence, and it also implies the sample is difficult to learn: a large \(\alpha_i\) means \(p(v|x_i)\) is small when \(y_i = v\), or \(p(v|x_i)\) is large when \(y_i \neq v\), in both cases the probability of the ground truth token is small.

Definition 2. The support samples of predicting token \(v\) are defined to be \(S_v = \{(\mathbf{x}_i, y_i)|\tau \leq |\alpha_i|\}\). The support samples of the full language model are \(S = \cup_v S_v\). The non-support samples are \(\bar{S}\).\footnote{We set hyperparameter \(\tau = 0.9\) in this paper as justified in Appendix B.}

We have seen that support samples are roughly hard-to-predict samples, here we can also inspect the opposite side by considering the relationship between non-support samples and easy-to-predict samples. Specifically, we consider \(M = \{(\mathbf{x}, \mathbf{y})|f(\mathbf{x}) = \mathbf{y}, p(v|\mathbf{x}) > \gamma\}\), which are samples correctly predicted after training, and defined to be “memorized” by language models (Tirumala et al., 2022). The following fact shows that the memorized samples are a kind of non-support samples.\footnote{The original definition of memorized samples in (Tirumala et al., 2022) sets \(\gamma = 0\). We can see a subtle rationale for introducing the confidence threshold \(\gamma\): the claim here says that even a sample’s correct label can be predicted, if the prediction is not confident enough, it could still be a support one.}

Claim 3. If \(\gamma, \tau \geq 0.5\), the non-support samples \(\bar{S}\) contains \(M\).

Proof. When \(\gamma \geq 0.5\), \(p(y|\mathbf{x}) \geq \tau\) implies \(f(\mathbf{x}) = \mathbf{y}\), and \(M = \{(\mathbf{x}, \mathbf{y})|p(v|\mathbf{x}) \geq \gamma\}\).

At the same time, for any \((\mathbf{x}, \mathbf{y}) \in M\), if there exists a \(v\) such that \((\mathbf{x}, \mathbf{y}) \in S_v\),
• if \(v = y\), \(\alpha = 1 - p(y|\mathbf{x}) < 0.5 \leq \tau\),
• if \(v \neq y\), since \(p(y|\mathbf{x}) > 0.5\), \(\alpha = p(v|\mathbf{x}) < 0.5 \leq \tau\).

There are contradictions \(\alpha < \tau\) in both cases, thus for all \(v, (\mathbf{x}, \mathbf{y}) \notin S_v\). \hfill \(\square\)

3 Support Samples at First Glance

We now profile support samples. Our exemplar language model is trained from scratch with GPT-2 architecture (117M parameters, with \(l = 12\) Transformer layers and hidden size \(d = 768\)) and wikitext-2 dataset (2.37M training samples). The vocabulary size \(|V| = 50257\). We will investigate models with more parameters (345M, 774M, 1.5B) and larger training set (wikitext-103 with 117M
Figure 1: Number distribution of support samples for each token \( v \). Blue bars are the top 42% tokens with the most support samples. Orange bars are the remaining 58% tokens which have less than 10 support samples. The table shows the number of the top 5 tokens.

Figure 2: Proportion of support samples in different POS tags. “Random” means that a random sampling on the full dataset will pick 54% (which is the proportion of support samples) support samples. The figure shows that the proportions of support samples in two POS-tag groups significantly differ from the random proportion.

samples) in specific experiments, and report results on larger models in Appendix E.

Number of support samples. Our first observation is that, with sufficient training (loss converges) and standard model selection (best validation set performances), among 2.37M samples, there are 1.29M support samples (54%). The proportion is remarkably high: more than half of the samples are important for the next-word predictor. The larger amount of support samples implies fewer patterns are discovered during training: the predictor fails to make the correct decision by keeping a few representative samples.4

By depicting the number of support samples for individual \( v \in V \) (Figure 1), we further find the distribution is highly screwed: very few \( v \) contribute a large number of support samples. It implies that though the overall support samples are many, prediction patterns of most tokens are clear (i.e., 58% tokens have support samples less than 10). We additionally provide a concrete example of the support and non-support samples for a specific \( v \) in Appendix C.

POS tags of support samples. To inspect support samples more closely, we roughly annotate POS tag of the next token \( y \) given its prefix \( x \).5

4Therefore, we may also have a discuss on the definition of memorized samples: in fact, akin to the memory of a computer, we can imagine the support samples (rather than non-support samples or previously defined memorized samples (Claim 3)) are stored (or memorized) in parameters, and they will be looked up later for token prediction.

5https://spacy.io/

and see whether samples with different target POS tags tokens will have different probabilities to be support samples. Figure 2 shows that 86% of verbs are support samples while the proportion is only 20% for punctuations. In general, support samples are more prevalent in semantically rich POS tags (e.g., ADV, NUM, ADJ, VERB, NOUN), while less common in semantically light tags (e.g., PART, AUX, DET, ADP, PUNCT). It indicates that relatively, the predictor keeps more information about verbs than punctuations in its parameters.

Two types of support samples By revisiting the representation of \( \theta_v \), we can group support samples into two types. Specifically, Equation 1 can be rewrite as (dropping \( (N\lambda)^{-1} \)),

\[
\sum_{y_i = v} (1 - p(v|x_i)) \phi(x_i) + \sum_{y_i \neq v} (-p(v|x_i)) \phi(x_i)
\]

For each target token \( v \), we define support samples from the first summation to be Type-1, and the second summation Type-2. Namely, Type-1 are support samples with the same target token as \( v \), but their prediction confidence is low. Type-2 are samples with different target tokens, but they are confident in predicting the wrong target \( v \).

To understand the roles of Type-1 and Type-2 samples, we conduct the following counterfactual argument: what if we remove (subtract) one support sample type of \( v \) from all prediction parameters \( \theta_V \). Let \( (x', y') \) be a testing sample. Empiri-
<table>
  <tr>
    <th rowspan="2">Removed Set</th>
    <th colspan="2">v="hens"</th>
    <th colspan="2">v="ction"</th>
  </tr>
  <tr>
    <th>full-set</th>
    <th>y = v</th>
    <th>full-set</th>
    <th>y = v</th>
  </tr>
  <tr>
    <td>-∅</td>
    <td>3.28</td>
    <td>0.24</td>
    <td>3.28</td>
    <td>0.56</td>
  </tr>
  <tr>
    <td>-\( S_v \)</td>
    <td>4.45</td>
    <td>2.37</td>
    <td>3.52</td>
    <td>1.50</td>
  </tr>
  <tr>
    <td>Type-1 of \( S_v \)</td>
    <td>3.35</td>
    <td>16.73</td>
    <td>3.28</td>
    <td>33.09</td>
  </tr>
  <tr>
    <td>Type-2 of \( S_v \)</td>
    <td>4.75</td>
    <td>0.00</td>
    <td>3.90</td>
    <td>0.00</td>
  </tr>
  <tr>
    <td>-random</td>
    <td>3.75</td>
    <td>0.27</td>
    <td>3.80</td>
    <td>0.51</td>
  </tr>
</table>

Table 1: Removing support samples of \( v = "hens" \) and “ction”. The number of support samples \( |S_v| = 10 \), and half of them are Type-1. When removing the samples, we subtract them from all \( \theta_v \). We consider four sets to remove, \( S_v \), Type-1/Type-2 of \( S_v \), and a random subset with the same number. We report the loss of all training samples (full-set) and loss of a subset which contains samples with \( v \) as the target token (\( y = v \)).

cally, we assume \( \phi(x_i)^T \phi(x') \geq 0 \) for all \( i \).

If we subtract Type-1 samples of \( v \), the logit \( q_{\theta_v}^T \phi(x') \) will be smaller since a positive term \( \sum_{y=v}(1-p(y_i|x_i))\phi(x_i)^T\phi(x') \) is dropped. With a softmax activation, the probability of predicting \( v \) could be greatly decreased, and the model resists predicting \( v \). Similarly, when subtracting Type-2 samples, a negative term is dropped, and the model tends to predict \( v \).6 Therefore, we would think that the existence of Type-1 support samples pull the predictor towards predicting \( v \), while Type-2 samples push the predictor away from predicting \( v \).

Table 1 lists a concrete example for which we conduct the above subtracting experiments. The results show that, for these two when removing Type-1, predictions on \( v \) fail, and predictions on other samples can be kept (even improved, compared with random subtracting: 3.35 vs. 3.75, and 3.28 vs. 3.80). When removing Type-2, predictions on \( v \) are perfect, while other samples are negatively affected. It shows that holding Type-2 support samples is crucial for the predictor, which is new to the perplexity-based data selection principle (only Type-1 samples are important).

To further understand Type-2, we draw a network to illustrate connections of next-token targets. In Figure 3, we link tokens according to Type-2 supporting relations: each directed edge from \( v \) to \( u \) indicates that a sample with \( v \) as the target token is the Type-2 support sample of \( u \). We also list hubs with most in-degrees and out-degrees.

6We also simply (and empirically) ignore the operations’ potential complex influences on other predictors (Figure 1).

<table>
  <tr>
    <th>v</th>
    <th>out-degree</th>
    <th>@</th>
    <th>and</th>
    <th>.</th>
    <th>in</th>
  </tr>
  <tr>
    <td></td>
    <td>30</td>
    <td>28</td>
    <td>27</td>
    <td>22</td>
    <td>17</td>
  </tr>
  <tr>
    <th>v</th>
    <th>=</th>
    <th>of</th>
    <th>s</th>
    <th>.</th>
    <th>@</th>
  </tr>
  <tr>
    <td>in-degree</td>
    <td>242</td>
    <td>153</td>
    <td>144</td>
    <td>134</td>
    <td>84</td>
  </tr>
</table>

Figure 3: An illustration of token relations according to Type-2. The number of nodes is 1690 (3% of \( V \), not all tokens have Type-2 support samples). The number of edges is 1993. We lists tokens with top k(=5) out-degree and in-degree in the network.

4 Importance of Support Samples

Support samples contribute the most to \( \theta_v \) (LM heads). This leads us to pose a question: If the support samples are defined to be the most important samples, can we only use them (ignore those non-support samples) during training?

We examine alternative ways of removing non-support samples: removing with sampling 1) hard: removing all non-support samples; 2) random: uniformly removing the same number of samples as the non-support; 3) soft: using the coefficient \( \alpha_i \) as the probability for sampling, meaning that some non-support samples will be retained. We experiment with two training configurations: only training LM heads (other parameters inherited from the original trained LM) and full model training. The result is shown in Table 2.

For training LM heads, the answer is: Yes, we can remove non-support samples (around 40%-50% of the samples) while keeping the same performances, with weighted sampling or even random sampling. However, for training full models, the answer is No.

An interesting phenomenon is that soft sampling outperforms random sampling for training LM heads (fixed representation space), but the opposite occurs when training full model. The former tells non-support samples are not necessary for
learning heads (in the sense of soft removing). The latter implies non-support samples are important for learning representations (parts of the model beyond the heads). Otherwise, the model trained only with support samples will perform like only head tuning with support samples. A deeper exploration of this point is presented in Section 6.

Curiously, for tuning heads, hard sampling performs worse than soft sampling.(more or even exclusively support samples does not lead to better performance, and meanwhile random sampling with fewer support samples performs worse.) Observing the changes in support samples, we find that although hard sampling slightly reduces original support samples (1.29→1.18), it turns more unseen non-support samples into support(1.29→1.90). This makes the predictor’s hyperplane unclearer (more support samples), and leads to a performance drop. It means without any non-support samples, learning the predictor will suffer from unnecessary overfitting to support. A small number of non-support samples can greatly alleviate the problem (soft better than hard), while support samples are still key for tuning heads (soft better than random). Regarding full model training, the lack of non-support samples poses a more significant issue, thus random sampling surpasses soft sampling, because non-support samples are more important for parts of the model beyond the heads as mentioned in the previous paragraph.

See Table 2 again, a pattern seems to emerge: the fewer the support samples, the lower the loss. We already see that the learned next-word predictor has many support samples. The next question is: Is it possible to make the hyperplane clearer by reducing some support samples?

We try more data, various regularization methods (L2 and dropout). Figure 4 first shows that the trends of loss and support proportion are largely consistent. And more data can help reducing support samples to make the hyperplane clearer and the predictor better, while L2 and embedding dropout fail.

In conclusion, we get 1) For tuning heads, support samples are key, but some non-support samples are also necessary. We can remove non-support samples with soft sampling. 2) Non-support samples are important for learning representations. 3) Fewer support samples indicate the clearer hyperplane and better performance of the predictor. More data vitals.

5 Predicting Support Samples

We have another discovery when training with subsets (not shown in Table 2): the loss on original support samples are almost the same for models trained with and without non-support samples (6.17 vs. 6.43). 7

The result implies that non-support samples have limited influence on better fitting support samples, and it roughly suggests that the hardness of support samples might be “intrinsic”. To make the problem clear, we cast an interesting question, can we predict whether a sample is a support sample or not without training the language model?

Surprisingly, we have a positive answer. We will show that given initial parameters, with simple classifiers (linear and MLP) and a limited number of training samples (annotations of support and non-support samples), one could recognize samples with high accuracy. It may confirm our “intrinsic” property intuition. The following are details.

Denote \( \zeta(\mathbf{x}, \mathbf{y}, \theta^{(t)}) \) to be a feature map which extracts features of a sample \( (\mathbf{x}, \mathbf{y}) \) according to model parameter \( \theta^{(t)} \) at checkpoint \( t \) (especially, we are interested in the initial state \( \theta^{(0)} \)). \( h_{\eta}(\cdot) \) is a binary classifier (with parameter \( \eta \)) which takes a feature vector \( \mathbf{z} = \zeta(\mathbf{x}, \mathbf{y}, \theta^{(t)}) \) and outputs \( s \in \{0, 1\} \) indicating whether the input is a support sample or not. To learn the classifier, we need a training set \( E = \{(\mathbf{z}_j, s_j)\}_{j=1}^K \). We experiment with the following configurations.

The feature map \( \zeta \) We consider three types of features at checkpoint \( t \),

• the last hidden vectors \( \phi_{\theta^{(t)}}(\mathbf{x}) \in \mathbb{R}^d \), which is \( \phi(\mathbf{x}) \) evaluated with checkpoint \( \theta^{(t)} \). Recalling that support samples are defined according to \( p(v|\mathbf{x}) \) which is a generalized linear model of \( \phi(\mathbf{x}) \), last hidden vectors seem to be a natural and effective choice.

• concatenation of all hidden vectors \( [\phi_{\theta^{(t)}}^1(\mathbf{x}); \phi_{\theta^{(t)}}^2(\mathbf{x}); \ldots; \phi_{\theta^{(t)}}^l(\mathbf{x})] \in \mathbb{R}^{dl} \), where \( \phi_{\theta^{(t)}}^j(\mathbf{x}) \) is the hidden vector of layer \( j \). The feature includes intermediate representations.

• gradient features \( \nabla_{\theta} p(y|\mathbf{x})|_{\theta^{(t)}} \). Instead of directly using the gradients, which are in the

7The loss of hard sampling (training only with support samples) on the original samples is 6.43 and of random sampling is 6.17. In fact, the original model is 5.30, but the comparison may not be fair since the size of the training sets are different. We thus take the random sampling 6.17 as our baseline.
<table>
  <tr>
    <th rowspan="2">Method</th>
    <th rowspan="2">Test Loss</th>
    <th colspan="2"># Support Samples(M)</th>
  </tr>
  <tr>
    <th>new training set (after removal)</th>
    <th>original training set</th>
  </tr>
  <tr>
    <td>-</td>
    <td>5.08</td>
    <td>1.29</td>
    <td>1.29</td>
  </tr>
  <tr>
    <td>hard</td>
    <td>5.64</td>
    <td>1.29→1.18</td>
    <td>1.29→1.90</td>
  </tr>
  <tr>
    <td>soft</td>
    <td>5.13</td>
    <td>0.87→0.87</td>
    <td>1.29→1.33</td>
  </tr>
  <tr>
    <td>random</td>
    <td>5.18</td>
    <td>0.70→0.74</td>
    <td>1.29→1.38</td>
  </tr>
</table>

(a) only training LM heads

<table>
  <tr>
    <th rowspan="2">Method</th>
    <th rowspan="2">Test Loss</th>
    <th colspan="2"># Support Samples(M)</th>
  </tr>
  <tr>
    <th>new training set (after removal)</th>
    <th>original training set</th>
  </tr>
  <tr>
    <td>-</td>
    <td>5.08</td>
    <td>1.29</td>
    <td>1.29</td>
  </tr>
  <tr>
    <td>hard</td>
    <td>6.57</td>
    <td>1.29→1.28</td>
    <td>1.29→2.36</td>
  </tr>
  <tr>
    <td>soft</td>
    <td>5.69</td>
    <td>0.87→1.01</td>
    <td>1.29→1.76</td>
  </tr>
  <tr>
    <td>random</td>
    <td>5.47</td>
    <td>0.70→0.73</td>
    <td>1.29→1.45</td>
  </tr>
</table>

(b) full model training

Table 2: Performance of the predictor and Change of support samples when removing non-support samples from original training set with various sampling methods. -: no removing (original training set). hard: removing all non-support samples; random: randomly removing the same number of samples as the non-support; soft: using the coefficient \( \alpha_i \) as the probability for weighted sampling, meaning that some non-support samples will be retained. We count support samples on two types of dataset: one is new training set (after removal) as the initial number of support samples varies for each method, another is original training set.

![Three line plots showing the change of Loss and Support proportion on different training sets.](page_104_682_627_246.png)

Figure 4: The change of Loss and Support proportion on different training sets. Data Size: we sample subsets of different sizes from wikitext-103.L2: we set the weight decay parameter to adjust L2 regularization during training. Dropout: we try different dropout rates at the embedding layer.

extremely high dimension space, we perform a random projection of the gradient vector to a low dimension (4096). 8

Classifier We consider two classifiers,

• linear classifiers, \( h_{\eta}(\mathbf{z}) = \eta^T \zeta(\mathbf{x}, \mathbf{y}, \theta^{(t)}) \).
The learnable parameter \( \eta \) equals the dimension of input features.

• MLPs, \( h_{\eta}(\mathbf{z}) = w^T W_2 \sigma(W_1 \zeta(\mathbf{x}, \mathbf{y}, \theta^{(t)})) \),
where \( \sigma \) is a non-linear activation function, and \( \eta = [w; W_1; W_2] \). We try different model capacities by varying the size of \( W_1, W_2 \).

Training set We collect training set \( E \) for this classification task from the training samples of the language model \( D \). We first train the language model with \( D \), and get annotations of support/non-support samples according to the last checkpoint. We split the set to training, validation and testing set with 8:1:1 ratio. 9 We also test smaller versions of the training set.

Checkpoints We extract features from different checkpoints of the language model. Two special checkpoints are,

• checkpoint \( \theta^{(0)} \) which is the initial language model without any training.

• the last checkpoint \( \theta^{(T)} = \theta \) which is the parameter defining the gold standard annotation of support/non-support samples. We may think the performances of classifiers at the last checkpoint are performances upperbounds: a classifier can see all necessary information for recognizing support samples.

8The time complexity of gradient features(gradients are only computed offline once) and the space complexity(gradients are projected to low dimension 4096) are acceptable.

9In the experiment, we use wikitext-2 set. The binary classifier has 1.90M training samples, 0.24M validation samples, and 0.24M testing samples.
As \( \theta^{(0)} \) could be hard, we are also interested in the performances at the early checkpoints of training.

The experimental results (Figure 5) show that with MLP and gradient feature, the prediction accuracy can reach 80% without training, just 5% shy of the upper limit (last checkpoint). The classifier of the last checkpoint still cannot achieve 100% accuracy, suggesting that the cut between support and non-support samples is not clear. If we allow extracting features from early training phases, the prediction accuracy can be better.

Performance improves with more data and more parameters (MLP with increasing intermediate dimension) but eventually levels off. The upper limit of gradient feature is around 80%, of last hidden state is about 65%, and of all hidden states is 66%. Gradient outperforms hidden states, particularly in the early training stages. Even with more parameters, hidden states still cannot catch up.

In brief, we successfully predict whether a sample is a support sample without training (80% accuracy) only using a simple classifier (MLP) with a randomly projected gradient feature.

6 Non-support Samples

In Section 2, from the counterfactual (removing non-support samples), we preliminarily infer that non-support samples are crucial for learning model parameters beyond the heads. Now we further show that non-support samples play a role in learning representations.

In Transformers architectures, higher-layer representations are typically better than lower-layer ones. Therefore, if higher-layer representations “contain” more non-support samples, could this indicate that non-support samples contribute to learning representations? We probe each layer of LM (learning new heads for each hidden states with the original training object of the LM).

Figure 6 describes the number and POS distribution of non-support samples across different layers.\footnote{We additionally present the probing result of memorized samples in Appendix D, which as a subset of non-support samples, show similar outcomes to non-support.} Overall, representations of the LM "contain" more non-support samples at higher layers. In the first six layers, the number of non-support samples remains relatively low and stable, with a sudden surge observed from the fifth to the sixth layer. The growth in higher layers is markedly more pronounced than in lower layers. For POS, the LM shows some recognition of NUM in the first layer. Notably, PUNCT is heavily contained at the sixth layer.

Now we can answer: non-support samples indeed play a role in learning representations. Conversely, they can also be viewed as a visualization of representational capacity: the LM’s representational ability undergoes a qualitative leap starting at the sixth (half) layer.

![Bar chart showing Layer-wise Non-support Sample Distribution](page_1012_1047_393_246.png)

Figure 6: The number and POS distribution of non-support samples across different layers of the LM. Overall, representations of the LM "contain" more non-support samples at higher layers. In the first six layers, the number of non-support samples remains relatively low and stable, with a sudden surge observed from the fifth to the sixth layer. The growth in higher layers is markedly more pronounced than in lower layers. For POS, the LM shows some recognition of NUM in the first layer. Notably, PUNCT is heavily contained at the sixth layer.

7 Related Work

Difficult instances and model training Many studies have investigated the relationship between the difficulty of training instances and their role for model’s learning, and observe that keeping harder instances and a small amount of easy instances can maintain or even improve model’s generalization performance. A wide range of difficulty measures have been adopted, such as “forgetting” behavior, i.e., after which this training example is always correctly predicted (Toneva et al., 2019), training loss magnitude (Jiang et al., 2021; Swayamdipta et al., 2020), gradient norm (Paul et al., 2021), the number of layers needed for prediction (Baldock et al., 2021), perplexity (Kwok et al., 2024), and hidden representations’ distance to other instances (Sorscher et al., 2022). Different from
Figure 5: Performance of the classifier under various variables. Left: Performance changes across different training stages. The accuracy achieves nearly 80% even without training, and only improves to 85% at the last/best step. Middle: Performance improves with more parameters (MLP with increasing intermediate dimension) but eventually levels off. The upper limit of gradient is around 80%, of last hidden state is about 65%, and of all hidden states is 66%. Due to the varying feature dimensions (last hidden state: 768; all hidden states: 9984; gradient: projected to 4096) and significant differences in parameter counts, only the portion within 350k parameters is shown in the figure for clarity. In practice, all hidden states can achieve 66% with more parameters. Right: Performance enhances with the growth of training data. All three figures demonstrate that, in terms of features, gradient outperforms hidden states, particularly in the early training stages. Even with more parameters, hidden states still cannot catch up.

these works that aim to identify small but informative subsets for training, we study the specific roles of support and non-support examples in learning representations and decision boundaries.

Predictive Data Attribution Predictive data attribution seeks to answer the question: what would have happened if we had trained on a different dataset? (Madry et al., 2024). Most existing approaches rely on the gradients of individual instances.\footnote{A notable exception is the work of Ilyas et al. (2022), which uses a linear model of instances to predict test accuracy.} One line of research uses influence functions (Koh and Liang, 2017) and its approximations (Koh et al., 2019; Park et al., 2023; Grosse et al., 2023a). However, other studies have observed them to be fragile (Basu et al., 2021) and do not answer the question of leave-one-out retraining (Bae et al., 2022). The other line of works approximates the training dynamics, by unrolling the updates during each step (Hara et al., 2019; Pruthi et al., 2020; Bae et al., 2024).

8 Conclusion

This work has taken a data-centric approach, to study the training samples that contribute the most to the prediction of each token using a representation theorem. Concretely, we make three contributions,

• We observe that support samples constitute over half of the training set. Interestingly, some predictions are supported by a large number of samples, while others are supported by only a few (\S3). suggesting language models may employ two distinct modes of prediction. Furthermore, we identify two types of support samples: those that help predict the correct tokens and those that prevent incorrect token predictions.

• We further study the roles of support samples in training using removing-retraining experiments. We find that support and non-support samples are respectively more useful for learning prediction heads and backbone representations (\S4 and \S6). Interestingly, we observe that a few non-support samples are always needed for successful training.

• Inspired by the observation that non-support samples do not contribute to learning support samples, we investigate whether the difficulty of support samples is intrinsic by predicting them without training language models. Surprisingly, we observe that using only the gradients of a randomly initialized language model can predict support samples with over 80% accuracy, validating our hypothesis (\S5).

9 Limitation

On the ambiguity between support and non-support samples. In section 6, we have noted that even when using features from the last checkpoint for training, the classifier’s accuracy on the test set is still not perfect. This suggests that the
boundary between support and non-support samples is not well-defined. Directly dividing training samples into support and non-support based on a threshold may have inherent issues.

Beyond argmax: rethinking decoding and sample importance. The argmax decoding enters the representation theorem through the MLE loss function(Theorem 1), which is common in LM training. If the language model is trained in a different object function, both the representation of next-word prediction heads and the definition of important samples could be different.

However, for inference, the analysis is not limited to the argmax decoding. For example, when estimating importance scores, one could use sampling instead of argmax prediction. The sampling-based decoding may better reflect the nature of language data, which often does not align with a deterministic majority-vote interpretation. This may help address the concern that the observed next token \( y \) should not be assumed to be the argmax outcome of the underlying distribution.

In terms of the sparsity of next-word distribution and the role of argmax decoding. We would think that the learning of language model tries to establish an approximate mapping between context-dependent semantics of a token \( v \) (i.e., presentations of various prefixes ended with \( v \)) and context-independent semantics of the token (i.e., presentations of \( v \) in vocabulary). We would also think that the next-word prediction task is more like a nearest-neighbor search in the learned vocabulary space.

10 Acknowledgement

The authors wish to thank all reviewers for their helpful comments and suggestions. The corresponding authors are Yuanbin Wu and Yufang Liu.

References

Juhan Bae, Wu Lin, Jonathan Lorraine, and Roger Baker Grosse. 2024. Training data attribution via approximate unrolling. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.

Juhan Bae, Nathan Hoyen Ng, Alston Lo, Marzyeh Ghassemi, and Roger Baker Grosse. 2022. If influence functions are the answer, then what is the question? In Advances in Neural Information Processing Systems.

Robert John Nicholas Baldock, Hartmut Maennel, and Behnam Neyshabur. 2021. Deep learning through the lens of example difficulty. In Advances in Neural Information Processing Systems.

Samyadeep Basu, Phil Pope, and Soheil Feizi. 2021. Influence functions in deep learning are fragile. In International Conference on Learning Representations.

Koby Crammer and Yoram Singer. 2001. On the algorithmic implementation of multiclass kernel-based vector machines. J. Mach. Learn. Res., 2:265–292.

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. 2021. A mathematical framework for transformer circuits. Transformer Circuits Thread. https://transformer-circuits.pub/2021/framework/index.html.

Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamile Lukosiuite, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, and Samuel R. Bowman. 2023a. Studying large language model generalization with influence functions. Preprint, arXiv:2308.03296.

Roger B. Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamile Lukosiuite, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, and Samuel R. Bowman. 2023b. Studying large language model generalization with influence functions. CoRR, abs/2308.03296.

Satoshi Hara, Atsushi Nitanda, and Takanori Maehara. 2019. Data cleansing for models trained with sgd. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.

Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, and Aleksander Madry. 2022. Data-models: Predicting predictions from training data. In ICML.

Ziheng Jiang, Chiyuan Zhang, Kunal Talwar, and Michael C Mozer. 2021. Characterizing structural regularities of labeled data in overparameterized models. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 5034–5044. PMLR.

Pang Wei Koh and Percy Liang. 2017. Understanding black-box predictions via influence functions. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1885–1894. PMLR.
Pang Wei W Koh, Kai-Siang Ang, Hubert Teo, and Percy S Liang. 2019. On the accuracy of influence functions for measuring group effects. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.

Devin Kwok, Nikhil Anand, Jonathan Frankle, Gintare Karolina Dziugaite, and David Rolnick. 2024. Dataset difficulty and the role of inductive bias. arXiv preprint arXiv:2401.01867.

Aleksander Madry, Andrew Ilyas, Logan Engstrom, Sung Min (Sam) Park, and Kristian Georgiev. 2024. Data attribution at scale. Tutorial presented at the 41st International Conference on Machine Learning (ICML).

Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. 2023. Trak: Attributing model behavior at scale. In International Conference on Machine Learning (ICML).

Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. 2021. Deep learning on a data diet: Finding important examples early in training. In Advances in Neural Information Processing Systems.

Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan. 2020. Estimating training data influence by tracing gradient descent. In Advances in Neural Information Processing Systems, volume 33, pages 19920–19930. Curran Associates, Inc.

Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari S. Morcos. 2022. Beyond neural scaling laws: beating power law scaling via data pruning. In Advances in Neural Information Processing Systems.

Swabha Swayamdipta, Roy Schwartz, Nicholas Lourie, Yizhong Wang, Hannaneh Hajishirzi, Noah A. Smith, and Yejin Choi. 2020. Dataset cartography: Mapping and diagnosing datasets with training dynamics. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 9275–9293, Online. Association for Computational Linguistics.

Kushal Tirumala, Aram H. Markosyan, Luke Zettlemoyer, and Armen Aghajanyan. 2022. Memorization without overfitting: Analyzing the training dynamics of large language models. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022.

Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, and Geoffrey J. Gordon. 2019. An empirical study of example forgetting during deep neural network learning. In International Conference on Learning Representations.

Chih-Kuan Yeh, Joon Sik Kim, Ian En-Hsu Yen, and Pradeep Ravikumar. 2018. Representer point selection for explaining deep neural networks. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montreal, Canada, pages 9311–9321.

A Proof of Theorem 1

Proof. We first compute derivatives of \( \log p(\mathbf{y}|\mathbf{x}) \),

\[
\frac{\mathrm{d} \log p(\mathbf{y}|\mathbf{x})}{\mathrm{d}\theta_v} = \frac{\mathrm{d}}{\mathrm{d}\theta_v} \left( \theta_v^T \phi(\mathbf{x}) - \log Z \right).
\]

For the normalizer \( Z = \sum_{v'} \exp \left( \theta_{v'}^T \phi(\mathbf{x}) \right) \), we have,

\[
\frac{\mathrm{d} \log Z}{\mathrm{d}\theta_v} = \frac{1}{Z} \exp \left( \theta_v^T \phi(\mathbf{x}) \right) \cdot \phi(\mathbf{x}) = p(v|\mathbf{x}) \cdot \phi(\mathbf{x}).
\]

Therefore,

\[
\frac{\mathrm{d} \log p(\mathbf{y}|\mathbf{x})}{\mathrm{d}\theta_v} = \begin{cases}
(1 - p(\mathbf{y}|\mathbf{x})) \phi(\mathbf{x}), & \mathbf{y} = v \\
-p(\mathbf{y}|\mathbf{x}) \phi(\mathbf{x}), & \mathbf{y} \neq v
\end{cases}
= (\mathbb{I}(\mathbf{y} = v) - p(\mathbf{y}|\mathbf{x})) \phi(\mathbf{x}).
\]

Next, since \( \theta \) is a stationary point, \( \frac{\partial L(\theta, D)}{\partial \theta_v} = 0 \), we have

\[
-\frac{1}{N} \sum_{i=1}^N (\mathbb{I}(y_i = v) - p(v|x_i)) \phi(x_i) + 2\lambda \theta_v = 0
\]
\[
\Rightarrow \theta_v = \frac{1}{2N\lambda} \sum_{i=1}^N (\mathbb{I}(y_i = v) - p(v|x_i)) \phi(x_i)
\]

B Rationale Behind the Choice of Hyperparameter \( \tau \)

We set \( \tau = 0.9 \) empirically, based on the following rationale:

Distribution of sample importance scores \( \alpha \). As shown in Figure 7, we observe a drastic increase of \( \alpha \) from 9.27% to 54.16% from \( \tau = 0.8 \) to \( \tau = 0.9 \), suggesting a natural threshold for different regions.

Ablation studies for threshold parameter \( \tau \) sensitivity. We report the number of support samples with different \( \tau \) as illustrated in Figure 8. The sharp decrease of the number of support samples from 0.9 to 1.0 further justifies our choice.
Figure 7: Distribution of sample importance scores \( \alpha \).

Figure 8: The number of support samples with different \( \tau \).

Softmax sharpen probabilities. Since probabilities are processed through softmax, they tend to be quite sharp. Actually, using a temperature parameter can make the model less sensitive to threshold choices, resulting in a smoother distribution.

Actually, the choice of threshold does not affect the analysis framework in the work: all methods can be generalized to any threshold.

On the other hand, threshold selection is practically significant. From a practitioner’s viewpoint, with a fixed budget, determining an appropriately sized support set and exploring how to adjust thresholds for different support set sizes are valuable directions for future work.

C Concrete Example

We present some support and non-support samples of *etts*, where the target token is *etts* itself, in Table 3. In our framework, *etts* is typically a suffix observed in proper nouns such as *Burnetts*, *Corbetts*, and *Plunketts*. Our statistics reveal that *Burnetts* and *Corbetts*(of support samples) appear once each in the training data, while *Plunketts*(of non-support samples) appears 68 times. This example suggests that less frequent patterns in the training data are more likely to become support to the LM heads’ parameters.

<table>
  <tr>
    <th>Type</th>
    <th>Sample</th>
  </tr>
  <tr>
    <td>Support</td>
    <td>... *The northern edge of the Plunketts Creek drainage basin is formed by Burnetts*</td>
  </tr>
  <tr>
    <td></td>
    <td>... *They are mainly composed of granite that has weathered into more rounded hills with many long scree slopes on their flanks. The highest point of these hills is Glamaig, one of only two Corbetts*</td>
  </tr>
  <tr>
    <td>Non-support</td>
    <td>... *Much of the Plunketts Creek valley is composed of various glacial deposits, chiefly alluvium. Although the Plunketts*</td>
  </tr>
</table>

Table 3: Partial support and non-support samples of the token *etts*.

Figure 9: The number and POS distribution of memorized samples across different layers.

![Bar chart showing distribution of sample importance scores α](page_120_120_480_480.png)

![Bar chart showing layer-wise memorized sample distribution](page_600_120_480_480.png)

D Layer-wise Non-support/Memorized Samples Distribution

Figure 9 presents the probing result of memorized samples, which reveals a trend consistent with non-support samples as shown in Figure 6.

E Scaling Up

We scale our experiments to 345M(GPT-2-Medium architecture), 774M(GPT-2-Large architecture), and 1.5B(GPT-2-XL architecture) models. Our findings remain consistent across these scales as follows.
<table>
  <tr>
    <th rowspan="2">Method</th>
    <th rowspan="2">Test Loss</th>
    <th colspan="2"># Support Samples(M)</th>
  </tr>
  <tr>
    <th>new training set (after removal)</th>
    <th>original training set</th>
  </tr>
  <tr>
    <td>-</td>
    <td>5.02</td>
    <td>1.13</td>
    <td>1.13</td>
  </tr>
  <tr>
    <td>hard</td>
    <td>5.62</td>
    <td>1.13→0.98</td>
    <td>1.13→1.77</td>
  </tr>
  <tr>
    <td>soft</td>
    <td>5.07</td>
    <td>0.73→0.70</td>
    <td>1.13→1.25</td>
  </tr>
  <tr>
    <td>random</td>
    <td>5.12</td>
    <td>0.57→0.53</td>
    <td>1.13→1.23</td>
  </tr>
</table>

(a) only training LM heads

<table>
  <tr>
    <th rowspan="2">Method</th>
    <th rowspan="2">Test Loss</th>
    <th colspan="2"># Support Samples(M)</th>
  </tr>
  <tr>
    <th>new training set (after removal)</th>
    <th>original training set</th>
  </tr>
  <tr>
    <td>-</td>
    <td>5.02</td>
    <td>1.13</td>
    <td>1.13</td>
  </tr>
  <tr>
    <td>hard</td>
    <td>6.67</td>
    <td>1.13→1.12</td>
    <td>1.13→2.36</td>
  </tr>
  <tr>
    <td>soft</td>
    <td>5.77</td>
    <td>0.87→0.70</td>
    <td>1.13→1.78</td>
  </tr>
  <tr>
    <td>random</td>
    <td>5.51</td>
    <td>0.65→0.53</td>
    <td>1.13→1.51</td>
  </tr>
</table>

(b) full model training

Table 4: Performance of the predictor and Change of support samples of the 774M model when removing non-support samples from original training set with various sampling methods.

![Distribution of samples importance scores α across models of various sizes trained on the wikitext-2 dataset.](page_184_670_563_377.png)

Figure 10: Distribution of samples importance scores \( \alpha \) across models of various sizes trained on the wikitext-2 dataset.

![Proportion of support samples in different POS tags of 774M model.](page_820_670_563_377.png)

Figure 11: Proportion of support samples in different POS tags of 774M model.

Distribution of sample importance scores \( \alpha \).
We compare the distribution across models of various sizes trained on the wikitext-2 dataset and find them similar, highlighting the scalability of our method.

Number of support samples. Table 5 lists the number of support samples for models of various sizes trained on the wikitext-2 dataset, showing consistency across scales.

We further replicate most of results using 774M model with the wikitext-2 dataset. Following are some key observations.

POS tags of support samples. Figure 11 shows that for 774M model, 82% of verbs are support samples while the proportion is only 16% for punctuation. This feature is consistent with the 117M model.

Removing non-support samples. We conduct the experiment of removing non-support samples on the 774M model. As shown in Table 4, the conclusions(\S4) still hold at this larger scale. For only training LM heads, we can remove non-support samples(around 50% of the samples) while keeping the same performances, with weighted sampling. Following the same reasoning(\S4), we get for tuning heads, support samples are key, but some non-support samples are also necessary. For full model training, we can’t simply remove non-support samples. Soft sampling outperforms random sampling when training LM heads, but the opposite is true for full model training, suggesting non-support sam-
<table>
  <tr>
    <th>Model Size</th>
    <th>Eval Loss</th>
    <th>#Support Samples(Prop.)</th>
  </tr>
  <tr>
    <td>117M</td>
    <td>5.05</td>
    <td>1.29M (54%)</td>
  </tr>
  <tr>
    <td>345M</td>
    <td>5.01</td>
    <td>1.28M (54%)</td>
  </tr>
  <tr>
    <td>774M</td>
    <td>5.02</td>
    <td>1.13M (47%)</td>
  </tr>
  <tr>
    <td>1.5B</td>
    <td>5.13</td>
    <td>1.26M (53%)</td>
  </tr>
</table>

Table 5: Evaluation loss and number of support samples for models of various sizes trained on the wikitext-2 dataset.

ples are crucial for learning representations beyond the heads.

Reviewer paper summaries: This work examines how training samples influence next word prediction models through the lens of the representer theorem. The authors develop an efficient methodology for analyzing the impact of training samples on large language models' token predictions, making a valuable contribution to our understanding of these systems.

The paper characterises sets of support and non-support samples of different words in next word prediction and studies these sets in different conditions as well as the effect of manipulations to these sets on, for example, the models's loss. 

With this tool the paper carefully tracks and analyses a number of hypotheses about the importance and/or role of support and non-support samples to the formation of decision boundaries and to the quality of learned representations.

The paper is rather clear, well written, and to my assessment rather sound. I particularly appreciate the pace and the attempt to form sketches and predictions prior to designing controlled experiments and gathering observations. I also particularly like how the paper presents the investigation as a collection of controlled observations (gathered with purpose) that help form reasonable explanations (without naively going on to declare grand generalisable conclusions built on impossible to justify inductive reasoning).
Summary of strengths: - The research addresses a fundamental question in LLM interpretability by investigating how support samples influence next word prediction, providing another practical tool in facilitating our understanding.
- The proposed methodology stands out for its elegance and efficiency, offering a practical way to measure and analyze how support samples affect token predictions. This measurement framework provides valuable insights for the field.
- The paper presents a clear, logical structure and supports its claims with comprehensive experimental results. The authors thoroughly demonstrate the significance of both types of support samples through well-designed experiments.

Clear, well-written and well-executed analysis of support and non-support samples (obtained by a form of data attribution technique), how many they are, how they distribute as a function of words and their POS tags, how they contribute to decision boundaries and/or to learned representations, etc.
Summary of weaknesses: - There is a conceptual disconnect between the definition of support samples as sentence prefixes in Section 2 and the POS tag analysis in Section 3. The paper needs to clarify how POS tags are determined for entire sentences and whether analyzing POS tags of individual tokens provides meaningful semantic insights about the complete sentences.
- The mathematical representation of language model heads $\theta_v$ in Section 4 could be clearer, particularly in relation to the definitions established in Section 2. The current presentation risks misleading readers about the relationship between these concepts.

I have no major weaknesses to highlight. 

I have to, in my view, minor weaknesses, which I believe the authors can address for a camera-ready and I'd be satisfied with a good response of how they intend to address it.

1. While I do find the observations and the narrative that joins them of nearly self-evident interest, I think most readers (myself included) would enjoy the authors' views on (potential) implications. There's some of this in the paper, but I think there could be more and more obviously marked as such (eg, gathered in a section). Section 8 gathers the observations and reiterates them, but leaves the question of implications mostly unaddressed. Some implications can be inferred from some passages in related work. However speculative such a section may turn out to be, I believe it helps to demonstrate why the observations contributed here are important on their own.

2. I think the first two paragraphs of the intro could have been clearer that the questions of whether we can trace the model decision back to their origin (line 044) is by no means original to this work and has been asked throughout the literature. In fact, this work is not trying to hide that (after all, the work builds on data attribution techniques from roughly 10-20 years ago), but it seems to me fair to get to that earlier in the introduction.
Comments, suggestions and typos: - Line 181-182, "very few $v$ contribute a large number of support samples" is confusing.

1. You may dispute the following as a matter of style, but I'll make the remark in any case.

I find the opening of the paper a little unpleasant to read. I think you are just trying to remark  that these models are being applied to a variable (but sometimes great) degree of success, and that understanding how models form or contribute to predictions is hence a worthy use of time. As currently worded though (lines 24 to 30), the introduction gives LLMs a rather disputable level of agency (not mentioning it also paints a caricature of researchers, practitioners and users as some sort of spectators). 

In my view (and I think this is not a rare view), these models _don't decide_. We plug them to decision rules (decision making pipelines, whatnot), we take these models to contexts where the decisions they power are regarded (by people or other systems) as automated decisions (or decision aids), typically not before having gathered some confidence in this setup (eg, by means of some amount of testing and development). Even if one considers the amount of technological adaptation required to port models across conditions/applications to be relatively little (arguably true compared to other tools), the mere fact that we test them for a purpose, adjust  hyper parameters on how we interact/prompt or decide, document our successes and failures in chains of related attempts, etc. do come across as an explicit form of continued design (ie, by _people_). It seems of little relevance, to the argument of agency, that this procedure seems simple and/or relatively uniformly applicable across a large collection of problems.

2. Line 57: a sample is important if it occupies a large part of prediction parameters

I do not understand 'occupies' in this sentence.

3. Perhaps the likelihood function in theorem 1 should be denoted L(theta)? (Perhaps L_D(theta) or L(theta; D) if you see the need to emphasise the dataset)

4. Number of support samples

This number must be a function of tau. Is this remark of any importance for the experiments and results as you presented them?

5. In line 210: but they are confident in predicting the wrong target.

By construction of the corresponding term in the sum, you can claim that the "wrong target" isn't the mode of those distributions, not sure you can claim anything about the model being confident (ie, about that probability being 'much higher' than 1/V or some other (larger) threshold for 'confidence'). Though you might have inspect a histogram of those probabilities and perhaps that's where the claim comes from.

6. The analysis assumes argmax decoding from the conditional next-token distributions. Is this in any way important? Couldn't this be generalised to any function (esp a stochastic one)? To me, there's something slightly odd in talking about _the next token_. The fact that, for most x, the dataset contains a single sample of a next token is all but data sparsity (mostly for we have fortuitously mined observations to train next-word predictors out of text collected for some other end). 

If embrace argmax decoding and force myself to think of the data generating process that accommodates it, then I have to imagine the observed y as the "majority vote" sample, which it most certainly isn't (contrast mining <context, next token> pairs from a complete sentence vs. asking K humans to perform next word prediction for that given context and then settle for the most frequent answer -- very different processes). Hence, I'd rather attempt to generalise argmax decoding to categorical sampling in my analysis tool.
Overall assessment scores: 3.5

4