Density estimation with LLMs: a geometric investigation of in-context learning trajectories

We are grateful to the reviewer for the very detailed and insightful review. We are pleased to note that the reviewer appreciated the “novel approach”, rigor of the research “well-supported by both theoretical and empirical analyses”, clarity of the paper, and that this work is “a significant contribution to both theoretical and practical applications in AI”. We provide a response to the reviewer’s comments and questions below.

Broader datasets

We thank the reviewer for suggesting additional experiments to clarify this point.

Currently, the main body of the paper only discusses experiments with simple symmetric distributions because these are, so far, the only cases where the learning trajectory admits faithful 2D embeddings, where the distances between the embedded points accurately reflect the statistical distance between the corresponding distributions.. In appendix A.7, we include experiments on randomly-generated ground-truth distributions, some of which are multi-modal, for example, the one on page 26. The one on page 25 is approximately a skewed distribution. The LLM is indeed able to estimate all of these distributions, with the DE trajectory well-approximated by the adaptive KDE model. It’s just that the 2D trajectory visualization provides little insights on these more complex cases.

In the updated appendix, we will include additional experiments on explicitly bimodal distributions, such as as Gaussian mixture suggested by reviewer acu3, as well as heavy-tailed distributions, such as Cauchy, as suggested by reviewer mikv. We will also include skewed distributions suggested by the reviewer.

Choice of Gaussian kernel for comparison

We thank the reviewer for the comment. Among our main findings is the fact that LLM’s in-context DE process could be interpreted as an adaptive KDE. This finding is independent of the Gaussian KDE, which only serves as a visual comparison in the InPCA analysis. As explained in the paragraph 331 to 336, many kernel-based density estimators share a similar bias towards Gaussianiety, as further illustrated in figure 13 of the Appendix. In the main paper, therefore, we use Gaussian KDE as a stand-in for all conventional kernel methods.

In our updated manuscript, we will clarify the role played by the Gaussian KDE, to avoid potential miscommunication regarding our main findings.

How to verify dispersive induction heads

Our paper focuses on showing how real foundation models perform probabilistic reasoning, and that a kernel-like algorithm emerges from simple next-token prediction training. Eliciting its mechanistic basis is therefore beyond the scope of our current work.

That being said, we thank the reviewer for raising a highly interesting and non-trivial question for future research. In the mechanistic interpretation community, researchers typically train smaller and more interpretable models, and then use linear probing and ablations to isolate the circuit of interest.

Note that the mechanistic basis of the original induction head was quite complicated, and took the course of several separate works to develop (explained in Appendix A.1.). We therefore think it might be appropriate to focus on introducing the dispersive induction head concept in our current paper, and leave the door open for researchers to better understand the mechanistic basis of it.

Comparison with Other Transformer Models

We agree with the reviewer’s suggestion that experiments on other models could help make our results more comprehensive. We are currently working on repeating our experiments on Gemma and Mistral, and will share the results as soon as we can. However, we would like to emphasize that our work is not to be interpreted as a benchmarking paper. Our goal is to understand how LLMs perform data-driven probabilistic modeling, instead of which LLMs is the most performant, or how they compare to other established methods.

From density estimation to probabilistic reasoning

We regard density estimation as a cornerstone problem that underpins many data-driven probabilistic modeling tasks such as time-series prediction, and conditioned densities estimation. These connections are discussed in section two: background, specifically the paragraph from 101 to 110.

A recent line of work investigates the possibility of using LLMs to make probabilistic predictions conditioned on language [1]. For example, one could provide the LLMs with natural language hints before feeding the serialized numerical data, for example “the following is a series of numbers sampled independently from a distribution: …”

[1]: Requeima, J., Bronskill, J.F., Choi, D., Turner, R.E., & Duvenaud, D.K. (2024). LLM Processes: Numerical Predictive Distributions Conditioned on Natural Language. ArXiv, abs/2405.12856.

A natural future direction, therefore, is to study how the adaptive kernel changes based on natural language instructions.

We thank the reviewer for the detailed feedback and noting that the “paper is well written” and that the “approach is novel and the topic is of high interest”. We provide below point-by-point answers to the different comments and suggestions, detailing the additional experiments we chose to perform following the reviewer’s suggestions. Since the experiments are computationally intensive, we will be including them in the paper throughout the discussion period as they come.

Models beyond LLaMA-2

This is an excellent suggestion which will indeed make the paper stronger. Following the reviewer’s comment, we are currently working on repeating our experiments on Gemma and Mistral models (released this year), and will share the results as soon as we can.

The choice of these additional models was made because they share the same tokenizer as Llama-2 (unlike Llama-3), where each multi-digit number is represented uniquely as a succession of 3 tokens (e.g., the number “123” is tokenized as three independent tokens representing “1”, “2”, and “3”). This is not the case for Llama-3 models, which makes the application of the algorithm Hierarchy PDF for extracting the PDF more involved. We will add a comment about this in the revised manuscript.

Choice of distances for InPCA

This is a great comment, Appendix 5.3 provides some explanations regarding the choice of the Hellinger distance (which gave us the most interpretable geometric interpretation in the early experiments that led to this work).

We note that the typical KL divergence is not applicable to inPCA visualization, since it is not symmetric, and hence not a true mathematical distance: the KL divergence from P to Q is not the same as Q to P, rendering a geometric embedding of P and Q impossible. That being said, we do agree that adding some comparison with the KL divergence would be interesting to get a more complete picture of the distributions. Moreover, given the popularity of KL divergence in modern machine learning, it makes sense to use that as the distance measure. To that end, we will repeat some of the visualizations using the symmetrized KL-divergence: $D_{SKL}(P,Q) = D_{KL}(P,Q) + D_{KL}(Q,P)$.

More complex target distributions

Currently, the main body of the paper only discusses these simple cases because these are, so far, the only cases where the learning trajectory admits faithful 2D embeddings, where the distances between the embedded points accurately reflect the statistical distance between the corresponding distributions. In appendix A.7, we include experiments on randomly-generated ground-truth distributions, some of which are multi-modal, for example, the one on page 26. One can appreciate both from the visualization (Figure 30) and the actual distributions (Figure 31), that the estimated density by LLaMA-2 13b is much similar to the bespoke KDE kernel, and not as similar to the classical Gaussian KDE.

In an updated version of the manuscript, we will include additional experiments on explicitly bimodal distributions, such as as Gaussian mixture suggested by reviewer acu3, as well as heavy-tailed distributions, such as Cauchy, as suggested by reviewer mikv.

Uncertainty in fitted KDE parameters

We thank the reviewer for this suggestion. In Figures 6 and all other figures in the manuscript where the adaptive kernel parameters are plotted, we actually visualize the 1-sigma uncertainty of the fitted parameter using shaded regions around the curve. As explained briefly in the paragraph from line 430 to 431, the uncertainty of these fitted parameters are estimated using the inverted Hessian matrix at the point of best fit, which is consistent with the inverse problem community.

We will update the caption of the figure to clarify this point. In our updated appendix, we will also add more theoretical justification for this choice of uncertainty estimation, along the following lines: our uncertainty estimation for the fitted kernel parameters is well-motivated from an information-theoretical perspective. Specifically, when the loss function is the Hellinger distance, or the KL divergence, the Hessian could be interpreted as approximately the fisher information matrix (FIM). The inverse of FIM is a widely accepted estimation of covariance matrix of the fitted parameters.

Skewed ignorance

We thank the reviewer for raising this intriguing hypothesis. Staring at Figure 2 part b, it does seem that the left tail is consistently higher than the right. In our updated appendix, we will visualize the models’ “true” prior at the 0 data limit, by probing the model with only the “start of the sentence” token and nothing else.

Another source of non-uniform bias, in our opinion, is that LLMs are slightly biased towards nice integers such as 0, 1, 2, … ,9, which are the small spikes that can be observed in Figure 2 part b.

We thank the Reviewer for the comments and suggestions, which have helped up to substantially strengthen our paper, as well as expand its scope. We hope that our added experiments and visualizations have adequately addressed your questions and concerns.

As the author-reviewer discussion period is coming to an end, we would be happy to read their reply to answer any additional questions. After that, we may not have the opportunity to respond to your comments.

If you believe we have addressed your comments, we would be grateful if you could acknowledge this and update your score. Otherwise, we are happy to provide further clarification.

Thank you again for your time and valuable review.

Our apologies, due to a glitch in our plotting pipeline, the first plot in Figure 25 was actually at (n=5). It has been corrected now. As shown in the updated Figure 25, the LLM's prior at (n=0) is indeed a uniform distribution with peaks at "nice" integers.

Thanks for catching this!

It seems reasonable that the LLM's prior corresponds to a superposition of a continuous part and those spikes, to represent integers and real numbers.

Thanks again for the insightful discussion. While we are mostly changing the appendices during the reviewer-author discussion period, if accepted, we will update the manuscript to highlight this observation, as well as the lopsidedness towards lower numbers, potentially due to Benford's Law.

We are glad that you found our update manuscript significantly improved. We'd be grateful if you would consider increasing your score if our additional experiments and explanations have addressed your questions.

We thank the reviewer for the insightful review and the positive feedback of the novel approach of using LLMs for PDF estimation explored in this work. We provide below a detailed point-by-point answer to the reviewer’s comments.

More complex settings

While we agree with the reviewer that the experiments focus on relatively simple tasks, this is mainly because we are interested in understanding “how” LLMs grapple with probabilistic reasoning tasks, and these mechanistic interpretations are better elucidated with distributions that admit intuitive visualizations.

As also remarked by reviewer mikv, the idea to generalize the observation to the multivariate case is certainly interesting. However, as we explained in the response to the comment, generalizing the PDF estimation algorithm to multivariate cases has technical challenges since LLMs are not “permutation invariant” with respect to the ordering of multivariate variables. In addition, we are currently repeating some of the experiments using 3-digit representations, to investigate whether the enhanced precision leads to consistent results.

Estimating high curvature PDF with longer context length

We thank the reviewer for pointing out the possibility of non-convergence when learning complex distributions. We are currently repeating the experiments of high curvature random distribution using maximum context length (~1000 states). In our updated appendix, we will plot the loss function (Hellinger distance or KL divergence) as a function of context length. In our earlier experiments, these loss curves typically plateau after at around 400 states.

Bias towards the edges

This is indeed an interesting artifact that deserves a note in the paper, and we appreciate the reviewer for pointing it out.

It turns out that in Figure 2 part b, there are spikes in the predicated PDF not only at 0.0 and 10.0, but also at other “nice” numbers such as 1.0, 2.0, … 9.0. A possible cause for this, in our opinion, is that these rounded numbers appear more often in colloquial, non-scientific speech in the training text, hence biasing the LLMs. As for the bias towards 0.0 and 10.0, we speculate that this is because “0” and “1” are especially privileged numbers in the training data since they may also occur as boolean or binary symbols.

As also observed by reviewer acu3, in the low-data limit, the estimated distribution is slightly skewed towards the lower end. This could be attributed to Benford's Law, which is the empirical observation that lower digits appear more frequently as leading digits in real-life data.

It is quite striking, in our opinion, that LLMs could both inherit such biases induced by training data statistics, and also at the same time perform the unbiased KDE algorithm in the large data limit. Although such superposition likely makes a clean mechanistic interpretation very challenging.

Additional hints for adaptive kernel shape

We appreciate the insightful remark! In the low-data limit, one could indeed eyeball the kernel shape by studying the predicted density distribution. This was actually one of our motivations to parametrize the bespoke kernel in the way we did - so that at some value of s, the kernel shape could imitate the thick-tailed distribution as the reviewer observed in Figure 2 b.

Impact of $\alpha$ in Bayesian histogram

We thank the reviewer for the suggestion. Please note that the Bayesian histogram is not biased towards the general uniform distribution family, but a particular uniform distribution over the range 0 to 10. The target uniform distribution, however, is much narrower, roughly over the range 4 to 6, as shown in Figure 4. We think this discrepancy, along with the high alpha value (1) explains the slow convergence.

According to our earlier experiments, setting alpha to 0 leads to erratic trajectories in the low data limit. This is because any new data point would update the histogram dramatically in the presence of prior. A strong uniform prior therefore not only makes Bayesian histogram start at similar points (ignorance) as the LLMs, but also makes the resulting trajectories smoother and easier to interpret, that is, following the Geodesics.

In our updated appendix, we will visualize the histogram algorithm at different levels of alpha, in order to help the reader better understand the role played by uniform prior.

Missing citation

We thank the reviewer for bringing this to our attention. We will update our manuscript to correct this typo. The intended citation here is a textbook on Gaussian Processes: Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems, 14(02):69–106, 2004.

LLM’s in-context capabilities for PDF estimation. We provide below detailed poiny-by-point answers to the reviewer’s questions.

Heavy-tailed target distributions

We agree with the reviewer that heavy-tailed distributions are more challenging with classical Kernel methods. Following the reviewer’s suggestions, we are currently running additional experiments with Cauchy and t-distributions, and will share the results later during the rebuttal period.

In Appendix A.7 of the current manuscript, we included experiments on randomly generated target PDFs. The in-context DE process of these are also well-described by the adaptive kernel method.

Generalization to multivariate density estimation

It would indeed be interesting to see how much the insights from this paper could generalize to multivariate time series, and indeed, we had thought about it too. However, using pre-trained LLMs to predict multivariate time-series systems is a non-trivial task, as explored in [1]. In particular, current state-of-the-art models struggle with order invariance. For example, the statistical prediction for the 2D tuple (A, B) might differ from that of (B, A). In other words, the formatting of multivariate data thus impacts the probabilistic prediction from LLMs.

Generalizing the Hierarchy-PDF algorithm to the multivariate case could be useful for many future investigations. However, due to the difficulties explained above, we suspect such extension may not be as straightforward as at first sight. We therefore prefer to leave multivariate generalization to future work.

[1]: Requeima, J., Bronskill, J.F., Choi, D., Turner, R.E., & Duvenaud, D.K. (2024). LLM Processes: Numerical Predictive Distributions Conditioned on Natural Language. ArXiv, abs/2405.12856.

The manuscript is now updated with new sections in the appendix that address some of the suggestions raised by reviewers:

A.4 In-context DE trajectories of LLaMA, Gemma, and Mistral Documents additional experiments on three other recently released LLMs: Gemma-2b and -7b and Mistral-7b-v0.3. LLaMA, Gemma, and Mistral are three different suites of models built independently by different teams. Yet, they all demonstrate similar in-context DE trajectories, all of which could be well-approximated by the bespoke KDE method.

A.6 InPCA embeddings with symmetrized KL-divergence and Bhattacharyya distance For completeness, we repeat InPCA embedding with other divergence measures. Unlike L2 embedding (vanilla PCA), the resulting trajectories appear smooth. However, they are difficult to interpret due to the presence of imaginary distances.

We are now working on other additional experiments suggested by reviewers, and will update as the results come.

The manuscript is updated with a new section in the appendix:

A.10 HEAVY-TAILED DISTRIBUTIONS which visualizes the learning trajectories with heavy-tailed distributions. These experiments are performed with LLaMA as well as Mistral-v0.3 for generality.

Dear all reviewers,

The manuscript is updated with new sections in the appendix:

A.11 ADDITIONAL NOTES ON BAYESIAN HISTOGRAM

Visualizes the Bayesian histogram algorithm at different levels of prior bias ($\alpha$).

A.12 UNCERTAINTY QUANTIFICATION FOR BESPOKE KDE PARAMETERS

Explains in detail how the uncertainty in fitted KDE parameters are estimated based on information theory.