LMD3: Language Model Data Density Dependence

Authors: john  kirchenbauer
garrett  honke
gowthami  somepalli
jonas  geiping
daphne  ippolito
katherine  lee
tom  goldstein
david  andre
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-21 22:30:06

Abstract: We develop a methodology for analyzing language model task performance at the individual example level based on training data density estimation. Experiments with paraphrasing as a controlled intervention on finetuning data demonstrate that increasing the support in the training distribution for specific test queries results in a measurable increase in density, which is also a significant predictor of the performance increase caused by the intervention. Experiments with pretraining data demonstrate that we can explain a significant fraction of the variance in model perplexity via density measurements. We conclude that our framework can provide statistical evidence of the dependence of a target model’s predictions on subsets of its training data, and can more generally be used to characterize the support (or lack thereof) in the training data for a given test task.

Main paper: 
 introduction: 
 "With the right dataset and the right benchmark, you can make a deep learning model fake absolutely any ability. It's easiest if the benchmark and the dataset are the same thing, but otherwise you can also achieve it with a dataset that is a dense sampling of the distribution the benchmark is sampled from." -@fchollet The goal of this study is to perform a careful series of experiments designed to provide concrete evidence for, or against, the hypothesis that the accuracy of a model on a test question is strongly impacted by the density of the training set around that question. To do this, we explicitly estimate a probability density function on the training distribution. We then ask whether it is possible to explain a significant amount of a language model's abilities by directly estimating how dense the training data distribution is at each test point. We quantify the density of training data in embedding space using a classic technique in statistics, Kernel Density Estimation (KDE), which relies on measurements of the similarity between a test point and its neighbors. We present an illustration of the general process in Figure 1 .To determine whether such a simple technique offers any explanatory power in the realm of modern large language models, and to assess the impact of density on test performance, we begin by manipulating density through contamination of a training dataset with copies (or near copies) of test samples. We then observe the correlation between accuracy on test samples and the training density near those samples. We confirm that 1) a properly configured kernel density estimate can detect the increased density near these training points and 2) the elevated performance is reliably predicted by the density estimate value for the question texts.We then study the natural density variation within a pretraining corpus and uncover a more complex and subtle downstream effect. Higher density test points have lower perplexity ("as expected") when density is estimated using only the nearest neighbors in the pretraining corpus for each query, although the effect size is small. However, when estimating density using a random subset of training points, we observe the opposite correlation. KDE q = 1 N ∑ K(q, x)query corpusFigure 1 : A system level view of the LMD3 pipeline. The corpus of data used to train a LLM and a test set of queries are projected into a vector space using a neural embedding model. For each resulting query vector, a density estimate with respect to the training corpus is computed. The resulting density estimates can be used to infer the model's ability to respond to a question-like query or simply reproduce the tokens in the query sequence based on whether the relative density is higher or lower at that point in sample space.We interpret our findings within the context of a wide array of prior work on data attribution, data contamination, and the counter-intuitiveness of high-dimensional statistics, and conclude that simple measures of dataset density are indeed sharp enough tools to be predictive of per-sample test set performance. At the same time, we observe that-at least at the 7B scale where models do not dramatically interpolate their data-the effect of dataset leakage can be surprisingly small. We wrap up with a discussion of applications of our methodology to instance and group-wise error analysis, dataset filtering and supplementing interventions, and the determination of test contamination from the training corpus.2 Related Work 
memorization and contamination: 
 Large language models have been shown to memorize sequences from their training data.The key, replicated finding across this literature is that the more a sample is repeated in the training data the more likely it is to be memorized (Carlini et al., 2022; Hernandez et al., 2022) . 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 expect a density based analysis to highlight the causal relationship between repetition and elevated sampling likelihood. Additional work has shown that a similar relationship holds for fuzzy or near-duplicated samples (Ippolito et al., 2022) implying that a density measure might also capture this weaker dependence relation.Recent work has studied the impact of data contamination during the pretraining and finetuning processes. Despite the fact that training solely and repeatedly on verbatim or paraphrased test samples leads to significantly elevated performance (Yang et al., 2023) , the impact can be much less pronounced in realistic scenarios where leaked samples are mixed into a larger corpus during pretraining (Jiang et al., 2024) . Since leakage of test queries constitutes a practically relevant instantiation of the aforementioned edge case of spikes in the training distribution, we design our initial set of experiments around a leakage intervention.Stated informally, the "data attribution hypothesis" is the belief that specific machine learning model behaviors can be attributed to a specific set of examples from the training data. The critical implicit assumption is that the set of training points the behaviors are attributable to is small with respect to the overall size of the training corpus. Therefore, work on data attribution seeks to design methods for identifying the set of most relevant examples, with notable approaches including the use of optimization related measures (gradients) to develop "influence functions" (Koh & Liang, 2017) , behavioral approximations of model outputs through the lens of kernels (Park et al., 2023) , and even exhaustive models of how a model behaves with respect to leave-one-out resamples of its training dataset (Ilyas et al., 2022) . To benchmark these techniques, curated datasets of queries and corresponding fact sets to which correct responses could/should be attributed have also been developed (Aky ürek et al., 2022) . In this work, we take an orthogonal approach by both relaxing the attribution hypothesis' assertion that the relevant set of training data for each query must be small, and switching away from relying on characteristics of the final model and instead directly analyzing the training corpus itself.The "Similarity Hypothesis"Quite related to the data attribution hypothesis is the "similarity hypothesis"-the near truism that models will make predictions that are similar to their training data. While prior work has shown that structured features of training data such as the elevated occurrence of specific named entities correlates with performance on factual test questions concerning those entities (Kandpal et al., 2023) , similar results involving more general notions of traintest interrelatedness are fewer and far between (at least in the modern LLM literature). That said, recent work such as Tirumala et al. (2023) has shown that density quantification can be used to guide dataset pruning routines for improved sample complexity during model training. Also published contemporaneously to our research activities, Yauney et al. (2023) showed that a wide variety of similarity metrics were "not enough to explain performance across benchmarks". While certain aspects of our results do suggest that the story is subtle, rather than simple, in contrast to their rather definitive conclusions, we posit that significant variance in LLM performance can be explained given the right measure and lens of analysis. 
preliminaries: kernel density estimation: 
 Density estimation is the general problem of estimating a probability density function from data. In our setting, we aim to estimate a distribution that is intractable, namely the function P : R d → R, where for our domain we have represented all of the points x ∈ X from our training data as d-dimensional vectors, and we would like to approximate the likelihood of drawing any given sample from the distribution over natural language sequences represented by our corpus. The tool we choose to use for this approximation is the Kernel Density Estimate.Definition 3.1. Kernel Density Estimate (KDE) For a training corpus X c = {x 0 , x 1 , ...x n-1 } ∈ R d with a bandwidth parameter h > 0 and a kernel function K h : R d × R d → R, for a query vector x q the KDE at x q over X c denoted KDE X c (x q ) is given as:KDE X c (x q ) = 1 |X c | ∑ x∈X c K h (x, x q )However, because the sum is over all n training samples, the cost of computing the KDE for realistic language model training datasets is intractable. To apply KDE in this setting, we turn to DEANN -Density Estimation from Approximate Nearest Neighbors, a clever optimization proposed and analyzed by Karppa et al. (2022) that enables the scaling of KDE to large datasets X c by decomposing the full KDE into the exact contributions of close neighbors, and approximate contributions from the rest of the corpus. A formal definition of the approximation algorithm that we use when computing KDEs on pretraining-scale datasets is provided in Algorithm 1. We discuss the details of the approximation scheme further in Appendix A.1. 
the lmd3 methodology: 
 For a chosen data point (text sequence) we want to study how the dataset density at that point is related to model performance at that point. We measure performance on "task-like" text sequences using model accuracy, and on unstructured or "webtext-like" data using the loss function. Our method takes as input a model and its training data, as well as one or more test/query sets of interest. The training data are embedded using an embedding model, and the nearest neighbors for each query sequence are retrieved from the vectorized corpus. Then, either, the neighbors are combined with a random subset of corpus vectors to compute a density estimate at each query point, or, when scale permits, the density estimate is computed exactly at each query point.We take for granted access to performant, pretrained neural sequence embedding models as general purpose similarity/distance functions for pairs of text segments. We consider two types of embedding spaces in which to perform density estimates, the features produced by an off-the-shelf retrieval model and the features derived from the hidden states of the model under analysis. While in preliminary investigations we explored working with the latter "self-derived" features, we failed to find a significant difference between the two in our small scale (fine-tuning) experiments, and so we stick with the former for large-scale experiments. The use of LLM hidden states as retrieval features is an active area of study and future versions of our methodology could benefit from those results, so we leave a more comprehensive evaluation of exactly what self-derived LLM embeddings represent that general purpose retrieval models do not to future work. 1We utilize the software package developed by the DEANN authors (Karppa et al., 2022) as an extremely efficient parallelized implementation of the exact KDE calculation for euclidean kernels when the data (X q , X c ) fits in memory. Since bandwidth selection is an empirical process, we perform an initial evaluation across a range of bandwidths in preparation for our paraphrase experiments (Appendix A.10).For the finetuning-scale experiments, where the corpus contains ∼ 100, 000 samples, we are able to directly use the exact version of the KDE. For experiments with a pretraining dataset, we leverage the authors' exact KDE implementation as well as their random sampling based KDE approximation as subroutines. Algorithm 1 describes how we combine this with a scalable neighbor search system to implement our "fully decomposed" version of the nearest neighbor search-based approximate KDE. "Fully decomposed" simply means that before running the KDE calculations themselves, as a preprocessing step for the query set, the neighbors for each example are retrieved from the training corpus via a massively distributed vector search engine (exactly, in our case, as opposed to approximately as suggested by DEANN). 2 Then, the random component for each density calculation is selected in a hierarchical manner via an initial large random sample without replacement, and a second sampling step, also without replacement, to select the final random complement for each query's individual density estimate. In Section 5.3 and Section 6.1, "KDE" refers to the exact KDE computed over the entire finetuning corpus, while in Section 5.4 and Section 6.2 it refers to the final value yielded by Algorithm 1-the weighted average of the local and random components of the KDE approximation over the pretraining corpus, except where denoted specifically as solely the "local" or "random" component. 1 We also remark that running multi-billion parameter embedding models over billions of sequences is costly and therefore such models are not as amenable to the pretraining scale experiments we ultimately perform.2 While the details of the neighbor search system we use in this work for the pretraining scale experiments are proprietary, we include the finetuning-scale implementation in our code. The neighbor search system can be substituted with any vector search engine that scales to the data under analysis, even if that scale necessitates an approximation scheme such those implemented in the FAISS library (Johnson et al., 2019) . Karppa et al. (2022) argue that this additional approximation does not bias the overall KDE approximation results any further.Input: A corpus X c of n text embeddings x c ∈ R d , a k nearest neighbor search subroutine over vectors in X c , NN k (•), a kernel function K and bandwidth parameter h together with the corpus over which it is computed X, defining a KDE X (•), two random sample size parameters m 1 and m 2 (m 2 < m 1 ≪ n), a dataset of query embeddings X q , x q ∈ R d .Output: Z q , ∈ R |X q |>0 , an approximation of the KDE for each x q ∈ X q . Randomly sample w/o replacement X 1 of size m 1 from X c for all x q ∈ X q do X nn ← NN k (x q ) ∈ R k×d X ′ 1 = {x ∈ X 1 |x / ∈ X nn } Randomly sample w/o replacement X 2 of size m 2 from X ′ 1 . z nn ← KDE X nn (x q ) z rand ← KDE X 2 (x q ) z q ← ( k n )z nn + ( n-k n )z 
models and data: 
 For our experiments we use a transformer-based sequence embedding model from the sentence-transformers library (Reimers & Gurevych, 2019) to generate the feature space in which we compute density estimates. For the controlled finetuning experiments we train the base Llama 2 7B model (Touvron et al., 2023) on doctored versions of the training set for the MMLU benchmark for 2 epochs as this regime produces meaningful differences between experimental settings whilst still being realistic (and resulting in a model that has not completely degenerated in other regards). For the pretraining experiments we analyze the performance of the 6.9B model from the Pythia suite (Biderman et al., 2023) as a function of a version of its training corpus, The Deduplicated Pile (Biderman et al., 2022) . We provide further details on these choices and experimental parameters in Appendix A.2.In order to develop our density methodology, we design a synthetic experimental setting in which we can directly tune the specific amount of support for test queries present in the training data. Roughly defined, an idealized paraphrase of a test example is the transform from natural language sequence x to x ′ such that the semantics of x and x ′ are equivalent. Access to such a transformation function allows us to perform controlled experiments where we take a test sample x and intervene to increase the amount of specifically relevant training data for this sample x by mixing paraphrases x ′ into the training data. This can be done in a number of ways, and we discuss the design choices for this process further in Appendix A.3.In our first series of experiments, we take 1, 000 random questions from the MMLU test set of 14, 042 questions and paraphrase them 3 times. We also consider an exact copy of the selected test questions as a "paraphrase" with perfect similarity, or distance 0 to the original query. For each test question we paraphrase, we compute the cosine similarity between the original query and each sample in the set of 3 paraphrases based on their embeddings. Considering the paraphrases for each query sorted in ascending order by similarity to the original query as Para 1, Para 2, and Para 3 respectively, we then refer to the collection of two paraphrases {Para 1, Para 2} as Paras = 1,2 and the collection of all three paraphrases {Para 1, Para 2, Para 3, } as Paras = 1,2,3.Then we finetune our base language model on datasets constructed by taking the "auxiliary train" split of the MMLU dataset and mixing in various combinations of the paraphrases and/or an exact copy of the original query, which we denote the presence of, or lack thereof, using Exact = 1 and Exact = 0 respectively. We evaluate the performance of the models trained on these different datasets and analyze the impact that the inclusion of paraphrases and exact copies of test queries has on the trained model. "Rank Accuracy" denotes the scoring method used by the HuggingFace Open LLM Leaderboard for evaluating MMLU, first proposed by Brown et al. (2020) .In-the-Wild: In and Out-of-Distribution Queries, Pretraining Scale Finally, we turn our focus to the loftier goal of relating the test time behavior of a language model to its pretraining data. In addition to the increased scale with which we apply the technique, the increased generality of the pretraining distribution requires us to curate sets of queries to represent different test time scenarios we might be concerned with. We consider two classes of query sets: in-distribution (ID) and out-of-distribution (OoD), where for the former we choose the simple, explicit definition of any query x q taken directly from the training corpus X c as ID, and consider anything else to be OoD. The training corpus is The Deduplicated Pile, re-segmented into ∼ 3.5B individual text sequences. For these experiments we compute the approximate KDE using parameters k = 1, 000 nearest neighbors for the local component, m 1 = 1, 000, 000 samples as a base random sample set, with m 2 = 10, 000 samples per query selected as the random complement, exclusive to the k nearest neighbors.We randomly sample, without replacement, 10, 000 segments from the training corpus itself. These are drawn from precisely the same set of segments over which neighbor retrieval and KDE computation are performed. Since these are simply webtext segments, there is no notion of "response" so we just compute perplexity (PPL) of the text segment under the LLM.We employ the same 14, 042 test questions as in the finetuning experiments to act as a set of queries that are not contained in the training data. Since there are ground truth responses for the queries, we can compute perplexities on both the query texts as well as on the correct target response when conditioned on the query text.OpenOrca Random 10k (OoD): As another set of diverse sequences outside of the train distribution, we consider a random selection of 10, 000 questions from the OpenOrca curation project's main dataset. Since these also have reference responses, we can compute perplexities on both the query and responses for these samples. 
results: 
 Considering the relative novelty of our proposed methodology and in order to de-risk the main experiments, we performed a series of validation analyses regarding the embedding model used for the paraphrasing experiments, the bandwidth hyperparameter of the KDE, and basic series of finetuning runs to calibrate our expectations about exactly how much a model should overfit when trained repeatedly on verbatim leaks of test questions. We discuss these preliminary experiments in section Appendix A.10. For the main experiments we consider bandwidths of {0.01, 0.05, 0.1, 1.0} for the exponential kernel, and {0.1, 0.2, 0.5, 1.0} for the gaussian kernel though we only report a small selection across the main body figures, primarily gaussian at 0.1 as a "narrow" setting tuned for discrimination in the leakage experiments and 0.5 as a "wide" setting for capturing a broader range of similarities in the pretraining-scale experiments. the model has trained on the leaked questions for. Right) We plot the same performance measure as a function of KDE, with gaussian kernel and bandwidth 0.5. We see that the trend in accuracy according to our density measure corresponds with the trend in accuracy according to the known degree of leakage, effective epochs.6.1 Controlled Experiments: "Expected" Dependence between Performance and Density.In Figure 2 we demonstrate that training on the various leaky dataset formulations improves performance by visualizing changes in performance across different leakage settings. We enable this by first plotting performance against a handcrafted feature called "effective epochs", where we compute the number of epochs the model has effectively trained on a given test question x t based on the specific set of paraphrases and exact duplicates X p included in the training set. For simplicity, in this figure, we represent exact copies and paraphrases equally such that if we train for 2 epochs, and have leaked 1 exact copy and 1 paraphrase, the total number of times the test question is trained on would be 2 * (1 + 1) = 4.In Figure 2 we observe both a positive trend in performance as a function of effective epochs as well as a positive trend as a function of our KDE measure-the latter of which is computed without access to either ground truths or any outputs of the LLM being analyzed.In Figure 3 , we present accuracies and densities across the full factorial manipulation of exact and paraphrased leaks. The figure shows that the KDE is a discriminative feature between the leak set and non-leak set. To measure the reliability of the correspondence between data density and performance, we perform mixed-effects regressions and discuss these results in detail in Appendix A.5. In summary, we find that the leakage conditions reliably increase accuracy (Exact leak: p < .001, Paraphrased leaks: p < .001) and decrease perplexity (Exact leak: p < .001, Paraphrased leaks: p < .001) on the test questions. Critically, training data density estimates also reliably predict variance in accuracy and perplexity where as density increases, accuracy increases (p < .001) and perplexity decreases (p < .001).6.2 In-the-Wild Experiments: Aggregation-Specific Dependence Relationship.First considering a set of random samples from The Deduplicated Pile itself, ID with respect to the model's training data, in Figure 4 we plot perplexity as a function of KDE by binning the data by KDE values (x-axis) and computing the average accuracy (y-axis) within each bin. In the left chart we see that when measuring the KDE with respect to only the random samples selected in our approximation algorithm we get a noisy, slightly positive trend (which is surprising3 ). However, when we consider the KDE computed only with respect to the nearest neighbors in the corpus, we get a sharp negative trend, in line with results in the We plot the distributions of KDE values (gaussian kernel and bandwidth 0.1) for the test queries that were leaked, exactly and or via paraphrase, and not leaked, for each leakage intervention experiment. Accuracy bar charts) We show the corresponding accuracy breakdown for the leaked and non leaked sets for each experiment. Overall, we find that increasing support for test questions via incorporating paraphrases into the training data increases performance on those test questions, and this increase is magnified by the addition of exact leaks of test questions. The addition of the exact copy of each question also makes the leaked and non-leaked question sets highly separable under our KDE measure as demonstrated by the distinct concentration of "leaked Q" KDE values away from 0.0 in (b). finetuning setting. 4 This is affirmed by the trend relating each query's average distance to its top k neighbors to model perplexity on said query. (Note that x-axes are reversed in the distance charts to make the trend visually congruent with the handedness of the trends for density measurements Figure 4 and Figure 5 .) Next, switching to a set of OoD queries, the MMLU test set, the left chart of Figure 5 shows that the PPL on the query/question texts is not strongly correlated with the local KDE component, but the middle chart shows that it is correlated with the average distance to the top k neighbors in the corpus. Further, we see that the perplexity of the correct response is also correlated with distance to the top k neighbors. While the relationship between query perplexity and local KDE isn't particularly strong, there is a stronger trend as a function of simple neighbor distances. For response perplexity, we see a clear trend as a function of average neighbor distances.As with the finetuning scale experiments, we used mixed-effects regression to measure the reliability of the effect of density on perplexity and report those full results in Appendix A.9. In summary, we find that query perplexity decreases as data density increases for the randomly-sampled ID query set (p < .001) and the OoD OpenOrca dataset (p < .001) but not the OoD MMLU Test query set (p = .95). However, response perplexity does decrease slightly with increased density for the MMLU Test set (p < .001).Since the text length directly effects perplexity values, for Figure 4 and Figure 5 we must manually isolate a subset of the queries where lengths are relatively similar to remove the variance due to length before plotting. We also observe a small number of extremely large outlier perplexity values and drop those rows as well. Details about this process are provided in Appendix A.7. We report only the most insightful combinations here and present a more exhaustive series in Appendix A.8. 
insights and applications: 
 We can summarize our key insights from both groups of experiments succinctly by stating that in extremal cases like test question leakage, one can expect a clear relationship between model performance and apparent training data density in query space. However, at pretraining scale, the specific way in which one aggregates information about relevant support in the dataset matters. Using the abstraction of "effective epochs" one can understand how repeated training on subsets of the data can lead to elevated marginal performance on those subsets but we caution that the number of effective epochs required to achieve drastically inflated test accuracies on benchmarks, such as those reported in Yang et al. (2023) , might be much higher than one would expect. This suggests that any un-evidenced claim that contamination is a reason to wholly discount the benchmark results for models trained on unspecified data distributions (at least at the common 7B parameter scale) should be taken with a grain of salt.Overall we believe that training data density quantification is a grounded methodology for analyzing the failure modes of language models at an instance and group-wise level, since it builds evidence for expected relative performance based directly on aspects of the training data itself. In response to observing relatively high density in certain regions of test query space and lower density in others, we expect that the consistency of model performance could be improved by supplementing data in those regions of weaker support through human data curation as well as automated processes such as machine paraphrasing. 
limitations and conclusion: 
 A few key limitations of our study include the focus on a small set of specific, but very relevant, datasets and models for the field and the fact that the KDE based methodology we propose is not hyperparameter free-the user must choose the kernel parameters. That said, the KDE methodology stands apart from simpler aggregations like the top-k distance measure in its "smoothness". While our experiments show that avg. distance to the top-k neighbors is also a reliable predictor of variance in performance, that metric requires the choice of a specific value of k. Neighbors that happen to lie beyond this boundary will not be accounted for at all in a simple top-k aggregation and thus changing k by just a value of 1, can alter the measurement significantly. In contrast, the KDE measure accounts for the nearest and farthest neighbors simultaneously and varies smoothly under small perturbations of the bandwidth hyperparameter as well as minor changes to which points fall in or outside of the k nearest neighbors set.It is also likely that the preprocessing parameters-length and stride-used to segment the training data before embedding are critical to the outcome of the analysis. We work at a finer segmentation granularity than Tirumala et al. (2023) and Yauney et al. (2023) but we do not ablate our choices enough to prove that this is required to explain meaningful amounts of variance. Finally we also acknowledge that full access to the training dataset of a newly "released" model is increasingly rare. However, given the potential of data-centric analysis techniques such as ours we claim that this is a shortcoming of contemporary norms within the field rather than a true limitation of our methodology. Rather, we implore the community to insist on the release of details sufficient to exactly re-materialize the full training data distribution as a necessary precondition for presenting model releases as bona fide scientific artifacts. 
 
 Appendix: We refer the reader to the excellent work of Karppa et al. (2022) for complete derivations but describe here the critical results that we rely on in this study. Their Lemma 3 states that one can compute an unbiased estimator of the true KDE X c (x q ) by splitting X c of size n into two non-overlapping subsets, X A and X B , computing z A = KDE X A (x q ) and z B = KDE X B (x q ) independently, and then combining them in a weighted sum according to the sizes of X A and X B :z q = |X A | n z A + |X b | n z BLemma 2 in their work, itself a result from prior literature, states that we can also compute an unbiased estimator of KDE X c (x q ) by taking a random sample X m of size m from X c and achieve an approximation KDE X m (x q ) with bounded error that decreases as m increases.Together, the authors use these results to propose Density Estimation from Approximate Nearest Neighbors (DEANN), where they approximate KDE X c (x q ) by computing the contribution of the nearest neighbors (X A ) and the contribution of the rest of the data (X B ). We leverage their results to allow us to perform density estimations at scale across millions or billions of data points. An implementation stylized description is provided as Algorithm 1.While the text documents that comprise a training corpus (which we embed to produce X c ) are in some sense taken as an input to the analysis we propose, the notion of how D c is broken up into the set of text segments D c = {s 0 , s 1 , ...s n-1 } is actually a design choice both during the training of a LLM and during our density analysis. A simple limitation that must be considered is that the embedding models used to transform s i into x i often have a maximum input length they can accept. More critically however, the kernel function that a KDE is based on relies on the ability to compute meaningful distances between points in the embedding space.Choosing the segment size for partitioning the documents in D c defines the sample space over which densities are computed and therefore potentially the outcome of the entire analysis. It is simple to convince yourself of this by considering extreme choices like all of the training tokens being a single segment, or every individual token being its own segment (i.e. devolving to a unigram distribution).Since a comprehensive search over all segmentations is impossible, and we are unaware of conclusive results on optimal segmentation for general retrieval applications let alone our text sequence density estimation problem, for our finetuning experiments we avoid segmenting as the queries are already relatively short, and for the pretraining corpus we analyze, we adopt a simple choice in this work as described in Appendix A.2.5. In Section 8 in the main body we briefly discuss the effect that this choice may have on the results.An important detail when analyzing queries that are suspected or known to be "In Distribution" (ID), i.e. contained in the training data, is to ensure to embed all samples into the space in which the KDE is performed under the "same conditions". This means that if x ∈ X q and x ∈ X c then the embedding vector of the query should be equal to the embedding of the matching doc in corpus, up to some floating point error tolerance. If any segmentation and/or prompt templating is used, it must be applied in the same manner between the train corpus and test query set if any of the queries are intended to be perfect matches with elements in the corpus with a distance ∼ 0.0 between their embeddings. For our leakage experiments and for our ID pretraining experiments, we ensure that this requirement is met.Published as a conference paper at COLM 2024 A.2 Models, Data, HyperparametersFor all the finetuning scale experiments involving the MMLU dataset, we utilize the multi-qa-MiniLM-L6-cos-v1 model from the sentence-transformers library (Reimers & Gurevych, 2019) following the setup of Yang et al. (2023) . We choose this model based on prior experience, literature precedent, generally favorable placement on the sentence-transformers retrieval model leaderboard, and critically it's extreme efficiency.For the pretraining data embeddings, we use a slightly more general purpose variant of the same architecture, all-MiniLM-L6-v2, a similar size model trained on a more general data distribution.We remark that while there are many larger, more competitive retrieval embedding models available, our desired to embed at the billion sample scale limits our choices. We feel that this tradeoff is compensated for by the fact that, at least for the pretraining-scale experiments, we consider very short sequences -just 50 whitespace separated tokens or 3 to 4 English sentences. We rely on the intuition that less representational complexity is required to accurately compare and contrast sequences of this length than when considering much longer documents. We additionally note that most prior work incurs the information loss cost of truncation at lengths of ∼ 512 since fine-grained re-segmentation is seldom performed and most open source retrieval models have context widths shorter than the multiple thousands now commonplace for more modern auto-regressive models.We adopt a focused approach and consider one of the standard state of the art open source large language models, Llama 2 (Touvron et al., 2023) , in its base form as the starting checkpoint for our finetuning scale experiments. We perform full finetuning of all parameters using the axolotl5 finetuning harness with the following features: DeepSpeed ZeRO-2, "target only loss" calculation, sample-packing, warm-up with a linear learning rate schedule, and weight decay. We detail the precise parameters for the different experiments in the Appendix A, as well as the training configuration files provided with the source code release.We also explore computing embeddings based on the last layer hidden states of the Llama 2 model after training on the finetuning data. These features are pooled and normalized in the same manner as the last layer hidden states of the sentence-transformers embedding models. In initial investigations we explored the "last layer, last token" strategy of extracting embeddings, which has been explored in contemporary work including Tirumala et al. (2023) , but did not find compelling evidence that this more expensive embedding process yields features more useful for our density estimation methodology than features from a smaller retriever model.We that said, we admit the choice to restrict ourselves to the general purpose, extremely efficient retrieval models is mostly a practical one enabling pretraining scale experimentation and consider the study of how to optimally leverage the internal representations of autoregressive large language models as important area of research (Wang et al., 2023) .We use the MMLU dataset (Hendrycks et al., 2021b; a) as our test-bed for the density estimation methodology. We focus on this dataset because it features prominently in both the Hugging Face Open LLM Leaderboard task set, as well as recent work competing for state of the art results (Team et al., 2023) in language model training.In order to limit confounding factors, we train the base model on a data distribution similar to that of the test task: the "Auxilliary" training set for MMLU. We format the training samples using the same logic that the EleutherAI LM Evaluation Harness (Gao et al., 2021) implements. We use the precise MMLU test set as filtered and prepared by the harness, noting that this is a slight subset of the original test set released by Hendrycks et al. (2021b) .The Pythia suite (Biderman et al., 2023) is an open science focused language model collection with public code and data. Given that our methodology requires access to the training set, we have limited options for pre-training scale analyses. While none of the Pythia models are directly comparable to Llama 2 due to architectural, training and dataset differences, since it is a relatively modern causal language model it is the most similar option amenable to our analysis at the time of this study.6 A.2.5 The Deduplicated PileThe Pile (Gao et al., 2020 ) is a public multi-domain web corpus for language model pretraining. We use the reworked version that was globally deduplicated (Biderman et al., 2022) before the corresponding set of Pythia models were trained on it. The raw dataset itself contains approximately 134M documents, of widely varying lengths. Due to the input length limitations of embedding models, and more importantly to ensure that distances between relatively short queries and individual samples in the corpus will be meaningful, we re-segment the dataset by first splitting the documents into chunks of 50 whitespace separated tokens (split only on " ") with a stride of 40 tokens, creating an overlap of 10 tokens each with preceding and succeeding segments. The result is an inflation of the original 134M documents into ∼ 3.5B segments.[start]We make this somewhat arbitrary choice for segment length and stride based on the observation that 50 whitespace separated tokens is approximately the length of a few sentences or a short paragraph in English. The current and previous sentences in this paragraph comprise precisely 50 such "words" on their own.[end]We chose our stride in order to limit false negatives during nearest neighbor search without causing the segment count to grow beyond what the implementation can handle7 . While we are unable to ablate this choice due to compute limitations, without obvious prior work to compare to, we accept these uncertainties as part of the cost of exploration.For each approximated KDE computation in the experiments with The Deduplicated Pile, for each query, we use the exact 1, 000 nearest neighbors, and a random complement of 10, 000, drawn for each query from a fixed pre-sample of 1, 000, 000 sequences from the training corpus. These correspond to k, m 2 and m 1 respectively in Algorithm 1.In this work, as in Yang et al. (2023) and other work that utilizes paraphrasing as a black-box transformation in the security and safety domain (Kirchenbauer et al., 2023; Krishna et al., 2023) , we prompt a powerful large language model, GPT4-turbo, to behave as a general purpose paraphrasing model using a template specifically designed to elicit paraphrased samples that are thoroughly transformed with respect to the original queries. Figure 6 shows the the prompt we use to paraphrase the test queries as well as example of an original MMLU test question and its paraphrase. Yang et al. (2023) explores adversarially regenerating paraphrases until a paraphrase with sufficient dissimilarity to the original example is produced (in order to evade leakage Figure 6 : The prompt used to instruct the powerful LLM used as a paraphrasing engine for the MMLU experiments, adopted from Yang et al. (2023) . Similar to their work, we find that GPT-4 Turbo, gpt-4-1106-preview version, is able to reliably paraphrase MMLU questions without losing too much context. We visually point out to the reader the fact that the entire question preamble as well as all answer choices are part of the paraphrasing step.detection measures). In our work, we are instead interested in collecting a spectrum of paraphrases with varying degrees of similarity and dissimilarity to test queries, and therefore, we simply sample k paraphrases for every test query on which we intervene, finding that pairing the prompt with the gpt-4-1106-preview model yields diverse paraphrases that are still faithful to the key details of the original question.In this work, we are interested in measuring the effect of training data density at the example level. We treat example training data density as a fixed effect and other covariates (such as the subject area of an MMLU question) as random effects-noise we wish to remove. Mixed effects modeling works by marginalizing the covariates and fitting a model to predict the dependent variable from the fixed effect(s). The model yields an estimated coefficient for each fixed effect and p-values that, used in conjunction with an α-criterion, determine the statistical significance of the effect. We conclude the fixed effect reliably predicts the variance of the dependent variable when the p-value is below the α-criterion, customarily α = .05. Concretely, if we determined p = .001, we'd expect to mistakenly conclude there was an effect 1 time in 1000 hypothetical experiments.Importantly, the sign of an estimated coefficient indicates the direction of the relationship. So a fixed effect (e.g., density) would have a statistically-reliable positive correlation with a dependent variable (e.g., accuracy) if p < 0.5 and β > 0. Overall, we look for agreement between the observed trends in binned plots of performance marginalized over various ranges of density estimates within a test set, the significance of the fixed effect, and the sign of its regression coefficient.All analyses were performed using R Statistical Software (v4.3.3, R Core Team, 2021) and supporting packages lme4 v1.1-35.2 (Bates et al., 2015) and lmerTest v3.1-3 (Kuznetsova et al., 2017) .We present two classes of analysis for the controlled experiments: manipulation verification regressions and the critical experimental regressions evaluating the effect of training data density. Regressions targeting the DV of rank accuracy are generalized linear mixed models (GLMER) fit by maximum likelihood with the Laplace approximation. Regressions targeting the DV of perplexity are linear mixed models (LMER) fit with restricted maximum likelihood (REML) and we use Satterwaite's method for calculating p-values. The mixed effects structure across all regressions is the same: we treat query length and example nested in MMLU question domain as random intercepts.A.5.1 Leak manipulation verification.rank accuracy ∼ exact leak + paraphrase leaks + (1|len(d q )) + (1|task(d q )/example)Note: Due to the natural ordering of paraphrase exposure (0,1,2,3), we treat the paraphrase variable as an ordered categorical and test its relationship as linear, quadratic, and cubic. <.001it to a supercategory based on the aforementioned mapping. Since this does not yield a balanced subsetting of the training samples (some supercategories are much larger than others) we examine the counts for each supercategory and select 4 with counts between 2, 000 to 4, 000 questions (out of 99, 842 total).Then we take the base language model and train it a collection of datasets where for each split we leave out the group of training examples corresponding to each of the 4 supercategories selected in turn. For each model we measure the impact that the intervention has on the average performance across test questions sharing the supercategory label of the left out samples, as well as the test questions from the other 3 supercategories as reference. We present the results of these experiments in Figure 8 as deltas in density and performance against the control model, which is trained on the full collection of training samples.Figure 8 : (Left) Shows that while the density measurements behave as expected across splits in our leave-one-subject-out experiment, (Right) the performance changes are small and inconsistent. Curiously, leaving out the math subset seems to have an outsize impact on the performance of the model on the biology set, suggesting that determining the influence that specific types of training questions have on the model's performance in this setting could be quite difficult no matter which heuristic was chosen.A.7 In-the-Wild Experiments: Controlling for length While the pretraining dataset analyzed, The Deduplicated Pile, is segemented according to whitespace before embedding, there is still a significant amount of variation in the lengths of the texts after tokenization, or as measured by character length. Since perplexity generally decreases as a function of length under any language model as tokens are easier to predict given more context, we need to control for length in order to highlight any observable differences according to density.We also need to drop outliers with respect to the performance measurement (we drop rows where query PPL > 500 and Response PPL > 60), as these heavily skew the marginalization process. To clean up the data for Figure 4 and Figure 5 in the main body, we first analyze perplexity as a function of query length in characters, and use this view to identify a suitable upper and lower bound for length such that we are able to focus in on any variance that can be explained by changes in density rather than length.A.8 In-the-Wild Experiments: Extended Figure Set While we highlight the most informative figures in the main body (Figure 4 , Figure 5 ) we also present a more complete set of visualizations for the MMLU Test set and Open Orca sampled query set here. (right) For the set of Random 10k (ID) samples from The Pile, we identify this as between 250 and 350 characters, and within this range we see that the query PPL is relatively constant around the mean. Thus, the data is limited to this range for the figures presented in the main body to highlight whatever variance can be attributed to differences in density. (right) For the set of MMLU Test questions, we identify this as between 200 and 400 characters, where the query PPL varies less wildly and where most of the data are concentrated (the majority of the bins are located in the upper left of the left panel). Thus, the data is limited to this range for the figures presented in the main body to highlight whatever variance can be attributed to differences in density. (right) For the set of Open Orca questions, the best we are able to do is limit to between 500 and 2000 characters, where the query PPL varies according to a "U"-shaped and where most of the data are concentrated (the majority of the bins are located in the upper left of the left panel). Thus, the data is limited to this range for the figures presented in the main body to highlight whatever variance can be attributed to differences in density. We see that while there isn't a consistent trend in question or response perplexity according to the local KDE, there is a more clear correlation for the distance to the k nearest neighbors. response perplexity vs distance to k nearest neighbors. Horizontal line denotes the average across all queries. We see that there is a non-monotonic relationship between query perplexity and both of the density heuristics. For response perplexity the trend is noisier. We suspect that the length variation for this query set washes out the effect of relative density differences.Retriever Mistral 7B Recall@5 0.9075 0.8875 Recall@10 0.9375 0.925 NDCG@10 0.939 0.928 Table 7 : During preliminary explorations we set up a retrieval evaluation where we used the original test questions as queries, and tried to retrieve their paraphrases (3 per question plus 1 exact copy) when mixed into the training corpus. The embeddings produced by both the very strong LLM being analyzed for performance effects, and lightweight retrieval model, were equally useful in searching for leaked copies of the test questions when considering the 5, and 10 nearest neighbors (Recall@k) for each test question query. In light of the fact that we wanted to scale our experiments to pretraining corpora, we decided that it was acceptable to use the lightweight retrieval model for the main experiments.Figure 14 Figure 15 : Since the proper KDE bandwidth for a particular problem an empirically derived parameter, we run a small ablation to chose a selection of bandwidths that are likely to produce an informative range of KDE values, especially those useful for identifying the set of queries for which we intervene through the planting of paraphrases amongst the other test questions. We see that as bandwidth is swept from low to higher values, the difference between the KDEs for queries with planted paraphrases and those without becomes smaller and smaller. This suggests that narrow bandwidths close to 0.0 are likely the most useful for developing a reliable estimator of whether a query has relevant paraphrases included in the training dataset. We use a representative selection of bandwidths from across this range in our main experiments so that we are able to examine what effect the bandwidth has on the relationship between KDE and other measures of interest. For the two euclidean kernels we investigate, the bandwidth selections are {0.01, 0.05, 0.1, 1.0} for the exponential kernel, and {0.1, 0.2, 0.5, 1.0} for the gaussian kernel, which we identified as reasonably similar sets for both kernel functions. We discuss the potential limitations of the selection process in the main work. ...but potentially explainable by distances in high dimensions being counterintuitive. Being closer on average to a collection of random other points effectively means being close to nothing in particular.We report this random -local breakdown explicitly here because at pretraining scale, the coefficient used to combine the local and random components according to Algorithm 1, k/N versus (Nk)/N respectively, causes the random component completely wash out the local contribution.github.com/OpenAccess-AI-Collective/axolotlThe Tiny-llama model was trained concurrently to this research, however, the dataset used to train it is still significantly larger than the deduplicated Pile and harder to materialize on disk exactly as it was prepared during the training of that model.A stride greater than or equal to your segment size results in no overlap between corpus segments, or worse, omitted sequences which could incur "false negatives" during neighbor search even for queries with large overlaps with training data. 


Decision: Although the reviewers make arguments for and against this paper, including some unresolved questions (possibly due to time as they were posted at the end of the rebuttal period). This is an important area and everyone agree the experiments are well designed and study an important class of problems. The relationships between data and end task performance is a hard one to figure out, and this papers add quite a bit to the discussion even if it doesn't have perfect answers on all fronts.

Reviews: review: paper_summary: paper_summary: This paper proposes a methodology for establishing a relationship between LM task performance and training data density estimation. Some controlled experiments are performed by introducing paraphrased data points. Although the approach is interesting, the results are not always easy to interpret. Nevertheless, I think the paper provides a useful tool for analysis. Many details are provided in the Appendices.
strength_weakness: strength_weakness: As the authors also point out there are some limitations which might restrict the wide application of the methodology.
paper_summary: paper_summary: The authors present a series of studies relating aspects of training data density at the example level to perplexity performance and task performance in language models. The authors perform both 1) a controlled experiment in which they vary the number and type of examples in training data that are "leaked" from the test set, and 2) an analysis of the relationship between perplexity measures and two different estimates of data density. Their claims are that increasing the training data density around certain examples in the test set can boost task performance, and that measures of training data density around test set examples can be predictive of perplexity performance.
strength_weakness: strength_weakness: My main concern is that at least one of the main claims about the usefulness of kernel density estimation (KDE) seems only weakly supported; it is unclear to me whether my perception is due to the evidence gathered itself or simply the authors' presentation of the results.
1. The authors claim that Figure 3 "shows that KDE is a discriminate feature between the leak set and non-leak set." However, the top row of Figure 3 clearly shows overlapping bars in what appears to be a histogram of KDE values. If there is any distinction, it is not visible from that plot with the x-axis scale presented. (Please do let me know if there is some mistake in Figure 3 or if a different scale would show a different pattern.) There appears to be a clear distinction only when the "leaked" queries are exact duplicates of test queries, which significantly dampens any claims of KDE's usefulness as a "discriminative feature" between leaked and non-leaked data points (if it cannot detect "leaked" paraphrases + many other simple methods could detect exact match leakage)
2. Relatedly, Figure 5's caption and 6.2 show that KDE is often but *not consistently* indicative of predictive of perplexity (visible also in Appendix figures). I would like to see the authors discuss more specifically the usefulness and limitations of KDE and top-k neighbor distance -- especially since it seems like the average Top-k neighbor distance may be at least as reliable and KDE seems to feature more heavily in the manuscript

Other concerns:
3. No useful information about the proprietary algorithm for example neighbor search. The efficiency of this implementation seems to be an important aspect of the authors' methodology, so it could be very helpful to mention its runtime complexity or even the scale of resources (compute/time) required to run the analyses presented
4. Baseline rank accuracy (0 paraphrases, Exact=0) missing depiction of variance. I would have liked to see it plotted like the others with error bars (or bands if the horizontal line style is kept)


I would not be comfortable accepting this paper in its current state, but I would be sincerely glad to adjust my score if the authors address my concerns, as I find the topic and experimental setting interesting, and my concerns are more about the discussion of the results and the conclusions drawn than the fundamental questions that the authors set out to answer.
questions: questions: See above. Is KDE truly not affected by the existence of "leaked" paraphrases? Additionally, some minor comments/suggestions/questions:
1. Last sentence in section 4.1 does not parse well to me. I would recommend not only splitting it up into multiple sentence but also rearranging the post-comma clause (e.g. --> "Future work can explore ...")
2. The term "doctored" in 5.1 makes that sentence read weirdly to me. Is there a more specific way that can the modifications to the training set can be succinctly described there?
3. It would be helpful to point to 6.1 in Figure 2 for a definition of the "effective epochs" that the accuracy values are plotted against, since 6.1 falls on a later page.
4. Figure 3 caption: "exactly and **/** or via paraphrase"
5. Figure 5 caption: would like to see a pointer to the missing plots in the Appendix
6. Appendix 10.2 and 10.3 seem to contain only figures; unclear which specific figures are meant to belong where
paper_summary: paper_summary: This work studies the data density dependence of language models on the perplexities/accuracy of queries / downstream tasks. In particular, this work provides the framework to assess the dependence of performance on a query given support of the training distribution The work suggests that model performance can be reliably predicted based on deter density dependence in many of the cases.
strength_weakness: strength_weakness: 1. Embedding model bias. I would like to hear more from authors on this beyond the small subsection of the paper. For example, embedding models not trained on code will actually cluster all code data very close by. I think the embedding model bias is a big problem here.
2. I do not understand why should we use KDE, when avg nearest neighbour distance performs better in the last few graphs in the paper? I did not find this discussion in the paper.

On top of that, two weaknesses that question the validity of the hypothesis of this paper:
1. The paper does not discuss the result of adding only Paras=3 in Figure 2. This experiment should not increase the density more than Paras=1 2 3 because the point in yellow should have almost 3x more density as compared to Paras=3. However, if it results in close to equivalent improvement in model performance, it would suggest density dependence may not actually be the true hypothesis here. Some other tests would be, does repeating exact 3x help more than paras=1,2,3 + exact repetition? In general, the "exact" inclusion convulates the hypothesis quite a bit. One example, should not increase the density much, but it does have a very large impact on the model performance (as compared to Paras=1,2,3)
2. In Figure 5, we never discuss accuracy on MMLU and only the response and query PPL. This is not really descriptive enough because as models improve or go to more favorable regions, the perplexity typically improves on all correct and wrong answers. I do not think the graph is suggestive of model performance as of now.
questions: questions: 1. here was a recent paper listed as no zero shot performance without exponential data (https://arxiv.org/abs/2404.04125). I wonder if the authors would like to contrast the results with the same. of course this paper came out during the review. But I am curious if you think exponential amount of data should be added rather than the linear amounts that tried in your work. 
2. secondly I don't understand the goal of figure 3.a, the count figure.
3. When the neighbor distance is low, do you check if there is a paraphrase or exact duplicate? Since you're only aggregating with 10 nearest Neighbors, we can actually check if the samples that were in the train set were in some way duplicates of the query
4. How does the average cosine similarity fare in comparison to the distance between embeddings
paper_summary: paper_summary: The paper considers whether kernel density estimation can be used as a measure to quantify the familiarity (my term, not the paper's) of a test-time example to a language model.  The paper includes a "controlled" experiment where MMLU test instances and/or paraphrases of them are leaked into finetuning data (and effect on their accuracy is measured), and "in the wild" experiments where test instances are leaked into pretraining data (and effect on their perplexity is measured).
strength_weakness: strength_weakness: One matter that's not at all clear from the paper is the matter of text granularity.  In the KDE setup, the paper doesn't make clear what an item in the corpus is.  Presumably it's not a word/token; is it a document?  Something smaller?  It would seem to make a big difference.  I'm fairly certain that this is a distinct decision from "bandwidth" and other kernel hyperparameters.

There is no discussion of the data that the off-the-shelf embedding model was trained on.  It seems worth some discussion as contamination of the embedding model's training data (with test query content) could be a problem.

It's a weakness of the paper that the Algorithm 1 and the details of paraphrase generation are deferred to the appendix.  I suspect other reviewers might prefer some other appendix materials to have been in the main paper, and I empathize with the challenge of keeping to the page limit.  But these two, at least, are essential for understanding the paper.  The details of the mixed effects models and controls for length would also be better in the main paper.  (Some space could be saved by compressing the related work section and figure 3's "count" histograms, could maybe be cut or moved to appendix since they're very low-information.)

I'd strongly recommend reorganizing the paper so that results of an experiment come directly after the explanation of the setup.  The current organization (all experimental setups laid out before discussing any results) creates an unnecessary mental burden on the reader.

Section 5.4:  MMLU test instances are taken as OOD, but the paper doesn't say whether these were explicitly checked against the pretraining corpus.  Perhaps some of them are actually in the pretraining data.

My understanding of the methodology is that "leaked instances" (in the pretraining case at least) are revealed to the model only after the normal pretraining is complete (the authors did not re-pretrain the whole model with the leaks randomly mixed in).  This has implications for the role of pretraining hyperparameters and the generalizability of the conclusions.  The paper should discuss this limitation. 

There is no attempt to quantify the computational cost of the experiments, leading to uncertainty about how practical the method is to apply in practice or to experiment with further.

This claim in the introduction -- "the effect of dataset leakage is small enough that it likely does not invalidate benchmark numbers" -- is not really supported by the quantitative arguments of the paper.  The paper doesn't really inspect any benchmark numbers at all.
questions: questions: Suggestion:  figure 2 is really nice.  I would suggest using the number of instances as the x-axis on the left, it's just easier to interpret and doesn't rely on weird assumptions (i.e., your fractional weighting based on cosine similarity).  The wide error bars on the right side make me want to see plots that show individual points (for a single choice of what is "leaked" at a time), for individual data points' KDE values and accuracy values.  Do we see correlation across instances as well?

Comments: title: Discussion period closing soon
Discussion period closing soon
comment: Thank you for the response and the increased rating!

Regarding the histograms in Figure 3, we're glad that you think the zoomed in version is informative. If there is no conflicting request from the other reviewers, then we can certainly include a final version of Figure 3 that swaps in the zoomed in histograms for the paraphrase only cases, and add a clear description of the axes scaling for that row in the caption.

Regarding the KDE bandwidth $h$ versus the value of $k$, we'd like to clarify that the rebuttal comment about the "$h+\delta$" is not about actually finding such a $\delta$, but rather a comment about the _smoothness_ of density wrt the users choice of $h$, versus the lack of smoothness of top-k wrt the users choice of $k$ (imagine one chose $h$ to be slightly off from "optimal" versus $k$ that is off by 1 or more). 

That said, we'd like to step back and remark that this is actually a classic choice in general for nearest neighbor data structures. Thinking simply of the nearest neighbor search process abstractly, one can either retrieve the "nearest k points" in the space, or can retrieve the "points that are closer than some radius d". A KDE measure with some bandwidth $h$ is actually a relaxation of the latter, as it is effectively building evidence at a query point based on a _soft_ analogue to a radius. 

Relevant to this conversation, we also think that the analysis presented in the recent response to y6ai is quite informative here. We find that since the KDE and the top-k measure seem to account for _different_ aspects of the variance in our setting, that there exists complementarity between the KDE and the top-k measures. Using the extra page of space, we think a concise presentation of this analysis and a discussion of complementarity of the two apporaches would make a good addition towards the end of the camera ready main body.

Finally, in light of the discussion with y6ai and your most recent comments, between now and camera ready, we are happy to agree to run the requested training experiment with just the singular nearest paraphrase (para 3), and analyze it in the same manner as the other settings. This would mean visualizing it with the other settings within Fig. 2, and in Fig. 3.
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)
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.
Thanks for your efforts with the rebuttal response. I do not think my questions are satisfactorily answered, and my main qualms still remain. In order of priority, they are

1. When nearest neighbors work just fine (except for the value of k being brittle?), I feel this paper is trying to make a claim of data dependence that is not really the causal reason for improved performance. Two things that could help in the future: (i) Do an actual ablation showing how duplicates change the density versus how they impact the nearest neighbours score. (ii) Actually show how brittle NN is. In fact, answer to Q3 above further suggests the existence of duplicates.

2. The reason why I suggested the experiment Paras=3 experiment is not just to pad the paper with more experiments, but because I believe this is a critical experiment in understanding the cause behind increase in model performance.

3. The MMLU comment has reason: Often when models become better, their PPL on wrong and right answers go low together. To understand this better, you can do the following : evaluate MMLU performance, but this time by checking PPL on second most likely (but wrong answer). You will likely see the exact same trend.
Hello again and thank you for the time already invested in your review of our work!

As the discussion period is coming to a close today, the authors are wondering whether the reviewer has had a chance to look through our rebuttal and linked figure updates? We'd be happy to try and address any specific lingering concerns from the reviewer that we weren't quite able to treat properly under the tight rebuttal length constraint.
Hello again and thank you for the time already invested in your review of our work!

As the discussion period is coming to a close today, the authors are wondering whether the reviewer has had a chance to look through our rebuttal and linked figure updates? We'd be happy to try and address any specific lingering concerns from the reviewer that we weren't quite able to treat properly under the tight rebuttal length constraint.
Thanks.
Thanks for the response; acknowledged.