LMD3: Language Model Data Density Dependence

(Responses in comment order, space permitting.)

In Appendix A.1.1 we discuss segmentation in some depth and in Sec. 8, we remark that it is an overlooked detail in similar works. However, in Sec 5.4 in the main body, we now realize that there is no pointer to Appendix A.1.1, a mistake on our part.

The off-the-shelf embedding model we use is detailed in Appendix A.2.1. We believe that this model was trained on a mixture of retrieval specific data including MS-MARCO. While we can not completely rule out the possibility of overlap between MMLU and those data without significant work, the likelihood that the retrieval specific training data for that model contains any significant fraction of MMLU test question texts is exceedingly unlikely in our opinion.

In response to this comment and the suggestions of other reviewers, we are choosing to remove the histograms from Fig. 3 in the main body. Please see the attached set of updated figures for mockup of the updated version. In addition, Algorithm 1 will be promoted to the main body using the extra page of space allowed in the updated draft.

We appreciate the suggestion for improved clarity via a "Setup-Result-Setup-Result" ordering and will consider it in the draft update process.

While we cannot entirely rule overlap out, none of our results or claims hinge on "OoD" separation between MMLU and the Pile being strict. Rather, in the pretraining scale experiments we report relative differences where the training corpus and the model are both constant.

We would be happy to add something like the following in the limitations section:

"Conducting a similar series of controlled experiments where examples of test queries are leaked into a full scale pretraining dataset would be interesting but was computationally out of scope for this work. We expect that results could vary and that the KDE hyperparameters explored in this work might not be optimal for that setting."

Regarding algorithmic details and computational cost, see response to reviewer vEjv.

Agreed, we will tighten the specificity of the claim about benchmark invalidation in the introduction or omit it altogether.

We thank the reviewer for the insightful suggestion and have prepared a version of Fig. 2 where the x-axis represents "counts" i.e. if we train for 2 epochs, and have leaked 1 exact copy and 1 paraphrase, the total count would be $2*(1 + 1) = 4$.

(Responses in comment order, space permitting.)

Beyond questions about an embedding model's capacity to represent certain data and questions about the diversity of its training data, we remark that interpreting distances in high dimensional spaces can be difficult. However, we would like to present a preliminary analysis that didn't make it into the review copy. In early experiments we considered embeddings derived from the hidden states of the LLM you are analyzing as an alternate to the standalone lightweight retrieval encoder used throughout the work. See the attached document for this analysis; we will incorporate it into the updated draft.

Regarding the differences between KDE and top-k distance aggregations please see the response to reviewer vEjv.

First we note that due to computational limits, and the cost of preliminary explorations, we chose not to run all possible permutations of the finetuning experiment focusing on those determined to be most critical to exploring our hypothesis.

We have updated Fig. 2 slightly in the attached document for improved clarity. While it's intuitive that training on an exact copy of a test question causes a larger increase in model performance on the corresponding questions than training on a single paraphrase, the fact that it has a larger effect on accuracy than training on all three paraphrases is slightly surprising. We adopt a purely descriptive stance in interpreting these results.

In the pretraining scale experiments we focused on loss/perplexity as these measures can be calculated on any text sample rather than accuracy measures that are constrained to only task-like data.

This (2404.04125) is a very interesting work, thank you for pointing it out! While it is difficult to identify a precise correspondence between visual concept frequencies and the repetition of test questions or paraphrases thereof, we agree that the there is some relationship between these findings.

While we cannot recover the precise segement texts for each neighbor easily at this time, we can report some anecdotal observations that nearest neighbors for pretraining segments were indeed sometimes duplicates or outher chunks of the same document.

While we did not directly compare the difference between cosine similarity and unnormalized distances, the sentence-transformers embedding model we used emits unit normalized vectors by default.

Thank you for your continued engagement in the review process!

First, we would like to remark on the close relationship between duplication and density, as it seems to be a relevant distinction to discuss. In the review copy (Section 2.1 sentence 2) we note that "Repetition is an important extremal case of elevated density corresponding to a spike in the density function at a specific point in the sample space" and as such we expected duplication to have a strong effect. By design, we included it in our controlled experiments. We also agree that it is relevant to our in the wild experiments. However, we would like to underscore the fact that the presence or lack thereof of duplicates, which are expected to affect both density measures and top-k distance measures, does not invalidate any particular experiment. It also would not invalidate any particular claim about density as duplication is an instance of increased density, albeit an extreme one.

Additionally, regarding para=3, we chose to include paras=1, paras=12, and paras=123 in order to make the results cleaner and easier to understand. The cumulative relationship between the chosen paraphrase settings means that density would be expected to grow across those three levels of leakage in an easy to interpret manner.

Next, in light of recent comments 1. and 2. the authors took a step back and considered an alternate modeling setup, one where we directly ask "can the two features, density and top-k distance, can explain different aspects of variance?" We realized that this could be achieved by running an particular regression model that asks whether the two features explain significant parts of the variance in the performance outcome, when controlling for the variance explainable by the other.

loss_rank_acc ~ gaussian_0.1 + avg_10_neighbor_distance + (1 | query_length) + (1 | task/idx)

Our predictors ("Fixed effects") are gaussian_0.1 and avg_10_neighbor_distance and target is loss_rank_acc. The other features of query_length and the subtask id from the MMLU subject list and the particular question id (i.e. task/idx) are treated as "Random effects" that the model attempts to marginalize out.

Fixed effects:
                           Estimate Std. Error  z value Pr(>|z|)    
(Intercept)               1.416e+00  3.514e-01    4.029 5.60e-05 ***
gaussian_0.1              4.625e+04  5.372e+00 8609.363  < 2e-16 ***
avg_10_neighbor_distance -3.595e-01  8.711e-02   -4.127 3.68e-05 ***
---
Signif. codes:  0 ‘***’ 0.001 ‘**’ 0.01 ‘*’ 0.05 ‘.’ 0.1 ‘ ’ 1

We want to look at the Estimate of the coefficient, and the Pr(>|z|) or p-value representing the significance of the predictive relationship for each feature. Interpreting the coefficients, we observe that as density goes up, rank accuracy goes up ($+$ coef) and as neighbor distance goes up rank accuracy goes down ($-$ coef) as expected. Next, moving to the p-values, the model suggests these relationships are accounting for unique variance, i.e., the small p-values indicate that both terms remain reliable predictors of variance when included in the same model (*** level or better).

This evidence clearly shows that despite the intuitive and theoretical ways in which top-k distances and the kernel density measure are related, in practice, they can each be used to explain different parts of model behavior. We agree with the reviewer that studying this relationship further is an important area for future work. Overall, we are very thankful to the reviewer for prompting us to undertake this analysis! We will perform a more complete set of similar regressions for the camera ready copy that helps the work more clearly communincate the difference between our proposed methodology and the top-k distance baseline.

Finally, regarding 3. we agree that the relationship between multiple choice task performance and the likelihoods assigned to individual responses is complicated. That said, the experiment that the reviewer proposes has roughly two outcomes. Either, the same trend in PPL on the correct answer as a function of density or neighbor distance is observed when computing PPL on incorrect answers, or it is not/the trend is different. This experiment would not really build any extra evidence about the overall hypothesis regarding performance and training density. In particular, finding that a trend or significant predictive relationship exists in both evaluation setups (MMLU+correct answer and MMLU+incorrect answer), would actually only suggest to you that PPL is probably a weak proxy for discontinuous evaluations like MCQ accuracy. However, this is a known result in the literature, and the models we had access to are too weak to solve MMLU well in the first place. When it was possible to do so, in our controlled experiments, we focused on the more meaningful task accuracy measure rather than PPL.

(Responses in comment order, space permitting.)

In the attached set of figures we have mocked up some requested changes, a few of which we discuss below.

We agree that the first row of subplots in Fig. 3 as presented in the review copy has relatively low information content. Therefore, we propose removing the histograms entirely from Fig. 3 in the main body, and then in the Appendix, presenting a "zoomed in" version of the histograms for the row of cases where only paraphrases were leaked.

As the reviwer points out, our experiments suggest that the simple average of the distances between a query and its $k$ nearest neighbors is also a discriminative feature. To be incorporated in the updated draft:

"The KDE methodology stands apart from simpler aggregations like the top-$k$ distance measure in its "smoothness". The top-$k$ average requires that you choose a specfic value of $k$ and neighbors that lie beyond this boundary will not be accounted for. Changing $k$ by just a value of $1$, can drastically change the mean. In contrast, the KDE measure accounts for the nearest and farthest neighbors simultaneously and varies more smoothly under perturbations of the $k$ nearest neighbor set. While one does still need to choose a kernel bandwidth for the KDE, informally, there always exists some $\delta$ such that the change in density at query point $q$ for bandwidth $h$ versus bandwidth $h+\delta$ is bounded and small."

Overall, we believe that choosing k, which can have a discontinuous effect on the estimate, is more fraught than tuning the bandwidth and it is a limitation of our evaluation rather than the method that we didn't identify a scenario where the KDE and top-k distance approaches diverge in discriminative power.

We apologize for the lack of detail. For the larger scale nearest neighbor searches over pretraining scale datasets, we used a system called ScaNN (Scalable Nearest Neighbors, (Guo, 2020)). In line with this request and comments made in another review, we will promote Algorithm 1 to the main body in the updated draft. Additionally, we can add a discussion to accompany the algorithm:

"Using a brute force maximum inner product search, the time complexity of the neighbor search step for $M$ queries over a corpus of $N$ segments is $O(M*N)$. However, with more clever tree data structures and quantization schemes, this runtime can be reduced significantly."

Thank you for your response! I disagree with the decision to remove the histograms entirely. Imo the zoomed in version is fine and informative, as long as you are clear about the differing x-scales.

I appreciate the discussion of sensitivity to k. However, it is not immediately obvious to me why it would be preferable to tune the KDE bandwidth than the k. "There always exists some $\delta$ such that the change in density at query point $q$ for bandwidth $h$ versus bandwidth $h+\delta$ is bounded and small" -- is it also true that it is easy to find such delta? Are optimal values for one hyperparameter expected to hold more stable in across settings?

Also, fwiw I agree with reviewer y6ai's assessment that it would be extremely informative to see isolated Paras=3 in Figure 2, given that 1, 2, and 3 are sorted in "ascending order by similarity." I would appreciate some sort of verbal commitment to running this ablation before camera ready.

I am willing to slightly increase my rating, but I believe there is room for improvement before publication (Edit: increased rating from 4->5)

We appreciate the vote of confidence from the reviewer in both the quality of our analysis and the completeness of our presentation.