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

## Abstract

Abstract Large language models (LLMs) demonstrate remarkable emergent abilities to perform in-context learning across various tasks, including time series forecasting.
This work investigates LLMs’ ability to estimate probability density functions (PDFs) from data observed in-context;
such density estimation (DE) is a fundamental task underlying many probabilistic modeling problems.
We leverage the Intensive Principal Component Analysis (InPCA) to visualize and analyze the in-context learning dynamics of LLaMA-2 models.
Our main finding is that these LLMs all follow similar learning trajectories in a low-dimensional InPCA space, which are distinct from those of traditional density estimation methods like histograms and Gaussian kernel density estimation (KDE).
We interpret the LLaMA in-context DE process as a KDE with an adaptive kernel width and shape.
This custom kernel model captures a significant portion of LLaMA’s behavior despite having only two parameters.
We further speculate on why LLaMA’s kernel width and shape differs from classical algorithms, providing insights into the mechanism of in-context probabilistic reasoning in LLMs.

## 1 Introduction

Modern Large Language Models (LLMs) showcase surprising emergent abilities that they were not explicitly trained for (Brown et al., ; Dong et al., ),
such as learning from demonstrations (Si et al., ) and analogies in natural language (Hu et al., ).
Such capacity to extract patterns directly from input text strings, without relying on additional training data, is generally referred to as in-context learning (ICL).

Recently, LLMs have been shown to achieve competitive performance on various mathematical problems, including time-series forecasting (Gruver et al., ), inferring physical rules from dynamical systems (Kantamneni et al., ; Liu et al., ),
and learning random languages (Bigelow et al., ).
To solve these tasks, an LLM must possess some capacity for probabilistic modeling — the ability to infer conditional or unconditional
probability distribution structures by collecting relevant statistics from in-context examples (Akyürek et al., ).

We investigate LLMs’ ability to perform density estimation (DE), which involves estimating the probability density function (PDF) from data observed in-context.
Our core experiment is remarkably straightforward.
As illustrated in Figure 1, we prompt LLaMA-2 models (Touvron et al., ) with a series of data points $\{X_{i}\}_{i=1}^{n}$ sampled independently and identically from an underlying distribution $p(x)$.
We then observe that the LLaMA’s predicted PDF, $\hat{p}_{n}(x)$, for the next data point
gradually converges to the ground truth as the context length $n$ (the number of in-context data points) increases..(^1^11Our prompts consist purely of comma-delimited multi-digit numbers, without any natural language instructions.)

To interpret the internal mechanisms (Olsson et al., ; Bietti et al., ; Dai et al., ; von Oswald et al., ) underlying an LLM’s in-context DE process, we leverage Intensive Principal Component Analysis (InPCA) (Teoh et al., ; Quinn et al., ; Mao et al., )
to embed the estimated PDF at each context length $\{\hat{P}_{n}\}$ in probability space (Figure 2).
These visualizations reveal that the in-context DE process of these models all follow similar low-dimensional paths,
which are distinct from those of traditional density estimation methods like histograms and Gaussian kernel density estimation (KDE).

By studying the geometric features of these in-context DE trajectories, we identify a strong bias towards Gaussianity, which we argue is a telltale feature of
kernel-based density estimationRosenblat (); Wand & Jones (); Silverman ().
This observation inspires us to model LLaMA in-context DE process as a kernel density estimator with adaptive kernel.
Despite having only two parameters, kernel shape and width, this bespoke KDE model captures LLaMA’s in-context learning trajectories with high precision.

Figure: Figure 1: In-context density estimation experiment. LLaMA-2 is prompted with 200 numbers sampled from $p(x)$ (left in red) (Appendix A.7), and predicts the PDF $\hat{p}(x)$ (right in blue) for the next number.
Refer to caption: x1.png

##### Main contributions.

- 1.
We introduce a framework for probing an LLM’s ability to estimate unconditioned PDFs from in-context data.
- 2.
We apply InPCA to analyze the low-dimensional geometric features of LLaMA-2’s in-context density estimation trajectories.
- 3.
We interpret LLaMA-2’s in-context DE algorithm as a form of adaptive kernel density estimation,
providing evidence for a dispersive induction head mechanism (Olsson et al., ; Akyürek et al., ).

## 2 Background

##### In-context learning of stochastic dynamical systems.

Gruver et al. () pioneered LLM-based time-series prediction by prompting LLMs with time-series data serialized as comma-delimited, multi-digit numbers. They observed competitive predictions for future states encoded as softmax probabilities.
Their data serialization scheme proved simple yet effective, allowing for easy log-likelihood evaluation of predicted next numbers.
It has since become a common prompting method (Jin et al., ; Requeima et al., ; Zhang et al., ; Liu et al., ) for evaluating LLMs’ numerical abilities.

Building on this prompting technique, Liu et al. () tested LLMs’ ability to in-context learn the transition rules of various stochastic and chaotic dynamical systems.
They introduced a recursive algorithm, termed Hierarchy-PDF, to extract an LLM’s predicted PDF for the next data point based on histories observed in-context:
$\hat{p}(x_{t}|X_{0},\dots,X_{t-1})$. This allowed them to rigorously compare the predicted PDF against the ground-truth transition probability $p(x_{t}|x_{t-1})$.
Our work leverages the data serialization technique proposed by Gruver et al. () and extracts LLMs’ probabilistic predictions using the Hierarchy-PDF algorithm introduced by Liu et al. ().

##### Density Estimation.

DE (Izenman, ) is a long-standing statistical problem with both classical (Silverman, ; Rosenblat, ; Scott, ; Lugosi & Nobel, ; Parzen, )
and modern machine-learned based solutions Papamakarios et al. (); Sohl-Dickstein et al. ().
It underlies the learning of more complex stochastic processes.
For example, in the first-order Markov process studied in (Liu et al., ; Bigelow et al., ; Zekri et al., ),
the transition probability $p(x_{t}|x_{t-1})$ has to be estimated in-context.
This can be viewed as a conditional density estimation problem, where for each possible value of $x_{t-1}$, one needs to estimate the density of $x_{t}$. In other words, learning a Markov process involves performing multiple density estimations;
one for each conditioning state.
Our current focus on unconditional density estimation thus serves as a stepping stone towards understanding these more complex learning tasks.

Figure: Figure 2: Visualization pipeline of LLaMA’s in-context density estimation process. (a) Data points are independently sampled from a ground truth distribution (Gaussian in this example), then serialized as comma-delimited, two-digit numbers to prompt LLaMA-2 models. (b) Hierarchy-PDF extracts LLaMA’s estimated density function $\hat{p}_{n}$ at each context length $n$. (c) InPCA reveals low-dimensional structures in density estimation trajectories, capturing 91% of pairwise Hellinger distances in 2 dimensions. Visual guides: gray - uniform PDF representing maximal ignorance, deep blue - ground truth PDF, and pink - 1D submanifold of centered Gaussians with variances ranging from $\infty$ to 0. All three LLaMA-2 models exhibit similar DE trajectories, geometrically bounded between the geodesic and the Gaussian sub-manifold.
Refer to caption: x2.png

## 3 Methodology

Here, we explain how we prompt a LLaMA-2 model to perform density estimation over in-context data, and how we extract and analyze its in-context learning trajectories with InPCA.
Our methodology consists of 5 steps, the first 3 of which are illustrated in Fig. 2.

1. Sampling and prompt generation.
Starting from a ground-truth probability density function $p(x)$, we generate a series of independent samples from it $X_{0},\dots,X_{t-1}$.
Consistent with Gruver et al. (), we then serialize and tokenize this series into a text string consisting of a sequence of comma-delimited, 2-digit numbers, forming a prompt that looks like
“6 1 , 4 2 , 5 9 , … , 8 1 , 3 2 , 5 8 , ” (Figure 2 (a)).

2. Extracting estimated densities with Hierarchy-PDF.
Upon prompting with such a text string, we read out LLaMA’s softmax prediction over the next token,
yielding probabilities for 10 tokens (0-9)(^2^22The logits of all other tokens are ignored.), creating a coarse,
10-bin PDF spanning $x\in(0,100)$. We then read out the next token, which refines one of the tens bins by further dividing it into 10 ones bins.
This process is repeated recursively 10 times, until all bins are refined, yielding a predicted PDF for the next state $p(x_{t})$, which
is a discrete PDF object consisting of $10^{N}$ bins, where $N$ is the number of digits used in representing each number.(^3^33We use the Hierarchy-PDF algorithm (Liu et al., ), which performs this recursive search efficiently for transformers.)
We interpret the extracted PDF $\hat{p}(x_{t})$ as the LLM’s estimation of the ground-truth $p(x)$.

3. Visualizing DE trajectories with InPCA.
When there are very few in-context data, LLaMA’s estimated $\hat{p}_{0}(x)$ is close to a uniform distribution over the domain $(0,100)$, showing a state of neutral ignorance, which is a reasonable Bayesian prior (Xie et al., ).
However, as the number of in-context data ($n$) increases, LLaMA gathers more information about the underlying distribution.
As a result, the estimated $\hat{p}_{n}(x)$ gradually converges to the ground truth $p(x)$ (Figure 2 (b)).
The series of estimated PDFs $\hat{p}_{0}(x),\hat{p}_{1}(x),...,\hat{p}_{n}(x)$ over context length $n$ forms what we term the "in-context DE trajectory".
These trajectories, we argue, offer important clues about how LLaMA-2 performs in-context learning.

However, with our 2-digit representation, each $\hat{p}(x)$ is a vector living in a 99-dimensional space(^4^44$10^{2}-1$, where the $-1$ dimension is from the normalization constraint that $\int p(x)dx=1$.),
making direct analysis of the trajectory challenging.
We therefore use InPCA to embed these PDF objects into a low-dimensional Euclidean space (Figure 2 (c)), and then analyze their geometric features.(^5^55The idea of using InPCA to visualize the training dynamics of machine learning systems has been explored in Quinn et al. (); Mao et al. (), which is further discussed in Appendix A.2)

4. Comparing with classical DE algorithms.
To interpret the low-dimensional geometric features revealed by InPCA, we embed the DE trajectories of well-known algorithms — specifically, kernel density estimator (KDE) and Bayesian histograms — in the same two-dimensional (2D) space.
Surprisingly, we find that the DE trajectories of LLaMA-2, KDE, and Bayesian histograms are simultaneously embeddable in this same low-dimensional subspace.
Moreover, the DE trajectories of LLaMA-2 are geometrically bounded between those of KDE and Bayesian histograms.

5. Explaining LLaMA-2 with bespoke KDE.
We observe that LLaMA’s in-context DE trajectories consistently remain close to those of KDE, suggesting that LLaMA might be employing a kernel-like algorithm, internally.
To investigate this hypothesis, we optimize a KDE algorithm with fitted kernel parameters (using only two parameters).

We now provide a detailed explanation of how to apply InPCA to analyze in-context DE trajectories.
We will also introduce two classical DE algorithms that serve as insightful comparisons.

### 3.1 InPCA visualization

To visualize the density estimation (DE) trajectories of various algorithms in a lower-dimensional space,
we employ Intensive Principal Component Analysis (InPCA) (Quinn et al., ; Mao et al., ; Teoh et al., ).
Given a set of PDFs $\{p_{i}\}_{i=1}^{m}$, InPCA aims to embed them as points $\{P_{i}\}_{i=1}^{m}$ in a lower-dimensional space $\mathbb{R}^{d}$ ($d\ll m$), such that the Euclidean distances between embedded points
closely approximate the statistical distances between the corresponding PDFs:

$$ $\|P_{i}-P_{j}\|_{2}\approx D(p_{i},p_{j}),\quad\forall i,j\in\{1,\ldots,m\},$ (1) $$

where $D(\cdot,\cdot)$ is a chosen statistical distance measure between PDFs.
In this work, we choose $D(\cdot,\cdot)$ to be the *Hellinger distance* (Hellinger, ), defined as:

$$ $D^{2}_{\text{Hel}}(p_{i},p_{j})=\frac{1}{2}\int\left|p_{i}(x)^{1/2}-p_{j}(x)^{ 1/2}\right|^{2}\,\textup{d}x.$ (2) $$

The Hellinger distance is ideal for visualizing our PDF trajectories for the following reasons(^6^66See Section A.5.3 for an in-depth comparison of the Hellinger distance against other divergence measures.):
1) it locally agrees with the KL-divergence (Liese & Vajda, ), a non-negative measure commonly used in training modern machine learning systems, including LLMs (Touvron et al., ), and 2) it is symmetric and satisfies the triangle inequality, making it a proper distance metric suitable for geometric analysis (Quinn et al., ; Teoh et al., ).

Having chosen an appropriate distance measure, we proceed with InPCA involving three steps:

- 1.
Calculate the Hellinger distance between each pair of PDFs, $D^{2}_{\text{Hel}}(p_{i},p_{j})$. This results in a pairwise distance matrix $\boldsymbol{D}\in\mathbb{R}^{m\times m}$, whose $\boldsymbol{D}_{ij}$ entry equals $D^{2}_{\text{Hel}}(p_{i},p_{j})$.
- 2.
Obtain a “centered” matrix $W=-\frac{1}{2}L\boldsymbol{D}L$, where $L_{ij}=\delta_{ij}-1/m$. This step is closely related to multi-dimensional scaling (MDS) (Chen et al., ).
- 3.
Perform the eigenvalue decomposition $W=U\Sigma U^{T}$.
The diagonal entries in $\Sigma$, sorted in descending order, represent the eigenvalues of $W$.
The embedding coordinates for the PDFs are given by $U\Sigma^{1/2}$, where the columns of $U$ are the corresponding eigenvectors.

The eigenvalues in $\Sigma$ represent the amount of variance explained by each dimension in the embedded probability space.
The cumulative fraction of total variance captured with an increasing number of dimensions is shown in Fig. 2 (c).
In our main experiments, approximately 90% of the Hellinger distance variance can be captured in just two dimensions, enabling faithful 2D visualization of the DE trajectories.

### 3.2 Classical DE algorithms

##### Kernel density estimation.

Kernel Density Estimation (KDE) is a non-parametric method for estimating the probability density function of a random variable based on a finite data sample.
Given $n$ samples $X_{1},X_{2},\ldots,X_{n}$ drawn from some distribution with an unknown density function $f$, the kernel density estimator $\hat{p}_{h}(x)$ is a function over the support of $x$ defined as:

$$ $\hat{p}_{h}(x)=\frac{1}{nh}\sum_{i=1}^{n}K\left(\frac{x-X_{i}}{h}\right),$ (3) $$

where $K$ is the kernel function and $h>0$ is the bandwidth (smoothing parameter).
The kernel function $K$ is typically chosen to be a symmetric probability density function, such as the Gaussian kernel $K(u)=\frac{1}{\sqrt{2\pi}}e^{-\frac{1}{2}u^{2}}$.
We discuss other kernel shapes in Appendix A.5.2.

The optimal bandwidth schedule is a central object of study in classical KDE literature (Appendix A.5).
By analyzing the Asymptotic Mean Integrated Squared Error (AMISE) (Equation 12),
researchers have derived a widely accepted scaling for the optimal bandwidth.

$$ $h(n)=Cn^{-\frac{1}{5}}$ (4) $$

where $n$ is the sample size and $C$ is a pre-coefficient (Equation 13).
In the low-data regime, a larger bandwidth provides more smoothing bias, which compensates for data sparsity.
While the $n^{-1/5}$ scaling is widely accepted, determining the pre-coefficient $C$ is more challenging (Wand & Jones, ; Silverman, ).
Unless otherwise noted, we use $C=1$ for classical KDE in this paper(^7^77In practice, various heuristics have been proposed to determine $C$, such as Silverman’s rule of thumb (Silverman, ) (see Appendices A.5.1 and A.8).).

##### Bayesian histogram.

The Bayesian histogram (Lugosi & Nobel, ) is another non-parametric method for density estimation, and can be formulated as follows:

$$ $\hat{p}_{n}(x)=\frac{n_{i}+\alpha}{n+\alpha M},$ (5) $$

where $\hat{p}_{n}(x)$ is the estimated probability density for bin $i$ containing $x$, $n_{i}$ is the number of observed data points in bin $i$, $n$ is the total number of observed data points,
$M$ is the total number of bins.(^8^88$M=100$ since our most refined estimation from Hierarchy-PDF consists of 100 one-digit bins), and $\alpha$ is the prior count for each bin.
We set $\alpha=1$, effectively populating each bin with one "hallucinated" data point prior to observing any data (Jeffreys, ).
This choice ensures that the histogram algorithm starts from a state of maximal ignorance, consistent with LLaMA’s in-context DE in the low data regime(^9^99A curious feature which we foreshadowed in Section 3 and further discuss throughout Section 4).

## 4 Experiments and analysis

We visualize and analyze the learning trajectories of LLaMA on two types of target (ground truth) distributions: uniform and Gaussian. To provide context and facilitate interpretation of the 2D space, we embed the following additional reference points and trajectories:

- •
Ignorance: A point representing maximum entropy (uniform distribution over the entire support).
- •
Truth: A point representing the ground truth distribution.
- •
Geodesic: The shortest trajectory connecting Ignorance and Truth points (Mao et al., ).
- •
Gaussian submanifold: A 1D manifold of centered Gaussians with variances ranging from $\infty$ to 0.

Our experiments focus on LLaMA-2 13b, unless otherwise noted.
While LLaMA-2 has a context window of 4096 tokens (equivalent to $\sim$1365 comma-delimited, 2-digit data points),
we limit our analysis to a context length of $n=200$.
This limitation is based on our observation that LLaMA’s DE accuracy typically plateaus beyond this point.

### 4.1 Gaussian target

Figure: Figure 3: In-context density estimation trajectories for Gaussian targets. Top row: 2D InPCA embeddings of DE trajectories for Gaussian targets of decreasing width (left to right). Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 92% of pairwise Hellinger distances between probability distributions.
Refer to caption: extracted/5914712/figures/gaussian_0.3_0.1_7b_13b_70b_seed3_Hellinger_2D_grouped.png

We begin our analysis with Gaussian target distributions of varying widths.

Wide Gaussian target.
A wide Gaussian distribution serves as our simplest target, as it is close in Hellinger distance to the uniform distribution (total ignorance).
In this scenario, both LLaMA and Gaussian KDE successfully approximate the target PDF within 200 data points.

Narrow Gaussian target.
However, as the Gaussian target narrows, Gaussian KDE lags behind, while LLaMA maintains its ability to closely approximate the target distribution.

In all three cases, the Bayesian histogram, by design, begins exactly at maximal ignorance, and then follows the geodesic trajectory.
Despite following this geometrically shortest path, the Bayesian histogram is the slowest to converge to the target distribution; likely due to the strong influence of its uniform prior.

Gaussian KDE, on the other hand, starts closer to the target, thanks to its unfair advantage of having a Gaussian-shape kernel.
Interestingly, Gaussian KDE consistently lingers on the Gaussian sub-manifold throughout the DE process.
This behavior of lingering on the Gaussian sub-manifold is not unique to Gaussian kernels; as demonstrated in Figure 13,
KDEs with alternative kernel shapes (e.g., exponential and parabolic) also exhibit a strong propensity for the Gaussian submanifold.
One potential explanation for this phenomenon lies in the nature of KDE itself.
KDE can be expressed as a convolution of two distributions: the empirical data distribution and the kernel shape, as

$$ $\hat{p}_{h}(x)=\frac{1}{nh}\sum_{i=1}^{n}K\left(\frac{x-X_{i}}{h}\right)=(F_{n }*K_{h})(x),$ (6) $$

where $F_{n}$ is the empirical distribution and $K_{h}$ is the scaled kernel.
This convolution operation tends to produce Gaussian-like distributions due to the central limit theorem (Fischer, ).
Consequently, a DE trajectory near the Gaussian manifold may indicate a kernel-style density estimation algorithm.

### 4.2 Uniform target

The Gaussian distribution easily arises in data with additive noise (Fischer, ) and therefore likely dominates the training data (Touvron et al., ).
What’s more, the Gaussian distribution is everywhere smooth, which makes it very easy to estimate from a function approximation point of view (DeVore, ).
Uniform distributions, on the other hand, feature non-differentiable boundaries; difficult to represent by both parametric (DeVore, ) and kernel-based (Wand & Jones, , Chapter 2.9) methods.
For these reasons, we now investigate the in-context DE trajectory with uniform targets.

Figure: Figure 4: In-context density estimation trajectories for uniform distribution targets. Top row: 2D InPCA embeddings of DE trajectories for uniform targets of decreasing width (left to right). Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 89% of pairwise Hellinger distances between probability distributions.
Refer to caption: extracted/5914712/figures/uniform_0.5_0.3_0.1_7b_13b_KDE_hist_Hellinger_2D_grouped.png

Wide uniform target.
For a wide uniform target, both LLaMA and Gaussian KDE initially move rapidly towards the target distribution.
However, they then linger at the point on the Gaussian sub-manifold nearest to the uniform target.
As more data streams in, they slowly depart from the sub-manifold and converge to the target distribution.

Narrow uniform target.
As the uniform target narrows, Gaussian KDE’s performance deteriorates,
while LLaMA’s in-context DE successfully reaches the target.
Notably, LLaMA exhibits less bias towards the Gaussian sub-manifold, as compared with the narrow Gaussian target case.

In the low data regime, Gaussian KDE and Bayesian histogram follow the same trajectories that they previously followed for the narrow Gaussian target.
However, LLaMA already appears to differentiate this target by taking a path further from the Gaussian sub-manifold and closer to the geodesic.
This behavior suggests that LLaMA’s in-context DE algorithm is more flexible and adaptive than classical KDE with pre-determined width schedule and shape.

### 4.3 The Kernel Interpretation of In-Context Density Estimation

Based on our observations in Sections 4.1 and 4.2, we propose a kernel-based interpretation of LLaMA’s in-context density estimation algorithm.
The bias towards Gaussian submanifolds in LLaMA’s trajectories suggests a KDE-like approach.
To test this hypothesis, we develop a bespoke KDE model with adaptive kernel shape and bandwidth, optimized to emulate LLaMA’s learning trajectory.

#### 4.3.1 Bespoke KDE Model

We construct our bespoke KDE model in two steps:

Step 1: Parameterize kernel shape

We introduce a flexible kernel function $K_{s}(x)$ parameterized by a shape parameter $s$:

$$ $K_{s}(x)=\frac{b(s)e^{-|b(s)x|^{s}}}{Z(s)}$ (7) $$

where $Z(s)=2\Gamma(\frac{1}{s}+1)$ normalizes the kernel to integrate to 1, and $b(s)=\sqrt{\frac{\Gamma(\frac{3}{s}+1)}{3\Gamma(\frac{1}{s}+1)}}$ scales it to maintain unit variance.
This parameterization allows us to interpolate between common kernel shapes such as exponential ($s=1$), Gaussian ($s=2$), and tophat ($s\to\infty$), as visualized in
Figure 5.

Figure: Figure 5: Bespoke kernel interpolates various common kernel shapes.
Refer to caption: extracted/5914712/figures/bespoke_kernel_interpolation.png

We visualize the DE trajectories with these three common kernels ($s=1,2$, and $\infty$) in Figure 13 of Appendix A.5.2.
As the shape parameter ($s$) decreases, the DE trajectories gradually shift from outside the Gaussian submanifold to inside, moving closer to LLaMA’s trajectory.
This trend suggests that our parameterization (Equation 7) may be able to capture LLaMA’s behavior by allowing $s$ to take values below 1, extending beyond the range of common kernel shapes.

We augment the standard KDE formula with kernel shape, resulting in two hyperparameters: $h$ and $s$.

$$ $\hat{p}_{h,s}(x)=\frac{1}{nh}\sum_{i=1}^{n}K_{s}\left(\frac{x-X_{i}}{h}\right)$ (8) $$

Step 2: Optimize kernel bandwidth (h) and shape (s)

For a given DE trajectory $\hat{p}_{1}(x),\ldots,\hat{p}_{n}(x)$, we optimize our bespoke KDE to minimize the Hellinger distance at each context length $i$:

$$ $\min_{s_{i}\in(0,\infty),h_{i}\in(0,\infty)}D_{\text{Hel}}(\hat{p}_{i}(x)\| \hat{p}_{h_{i},s_{i}}(x))$ (9) $$

yielding the "bespoke KDE" bandwidth schedule $\{h_{i}\}_{i=1}^{n}$ and shape schedule $\{s_{i}\}_{i=1}^{n}$, which together prescribe a sequence of fitted kernels of changing widths and shapes.
We visualize such sequences of fitted kernels, and compare them against the Gaussian kernel with standard $n^{-1/5}$ width schedule in Appendix A.6 and A.7.
We implement this optimization numerically using SciPy, and estimate parameter uncertainties from the inverse Hessian of the loss function (Hellinger distance) at the optimum.
This provides error bars for our fitted kernel shape and bandwidth parameters (Figure 7).

#### 4.3.2 Bespoke KDE Analysis

Figure: Figure 6: Optimized kernel shape ($s$) and bandwidth ($h$) schedules for the bespoke KDE model, fitted to LLaMA-2 13b’s DE trajectory on a narrow Gaussian target.
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_gaussian_0.1_hel_dist_False.png

Figure 7 presents the optimized kernel shape and bandwidth schedules for our bespoke KDE model, fitted to LLaMA-2 13b’s in-context DE trajectory for the narrow Gaussian target (Section 4.1).
Two key observations emerge:

1. The fitted kernel width decays significantly faster than the standard KDE bandwidth schedule ($n^{-\frac{1}{5}}$).
This rapid decay explains the much longer trajectory of LLaMA’s in-context DE within the inPCA embedding, as compared to standard Gaussian KDE (Figure 3).

2. Unlike fixed-shape KDE methods (e.g., Gaussian KDE with $s=2$), LLaMA seems to implicitly employ an adaptive kernel.
The shape parameter $s$ evolves from $\sim 0.1$ to $\sim 1$, with increasing uncertainty for $n\gtrsim 20$.

Figure 7 compares LLaMA-2 13b’s in-context DE trajectory with its bespoke KDE counterpart,
demonstrating the close fit achieved by our model.

Significance of Bespoke KDE
The existence of a bespoke KDE that closely imitates LLaMA’s in-context DE processs is significant and non-trivial.
Not all DE processes can be faithfully represented by KDE methods;
for instance, the Bayesian histogram, despite its superficial similarity to a fixed-width tophat kernel KDE,
resists accurate modeling as a bespoke KDE process (Appendix A.3).

While our initial experiments focused on regular distributions like Gaussian and uniform, for better generality we extended our experiments to randomly generated PDFs (Appendix A.7).
Notably, LLaMA-13b’s DE performance on these irregular distributions can still be effectively described using our bespoke KDE process (Appendix A.6 and A.7).
Although the DE trajectories of these random PDFs lack low-dimensional embeddings, and thus provide less geometric insight, they nonetheless corroborate our findings from the cases with more regular distributions.

This consistency across diverse target PDF types strengthens our core hypothesis about LLaMA’s underlying DE mechanism.
Specifically, we posit that LLaMA’s in-context DE algorithm shares fundamental characteristics with kernel density estimation,
albeit with adaptive kernel shape and bandwidth schedules that distinguish it from classical KDE approaches.

## 5 Conclusions

Inspired by emergent, in-context abilities of LLMs to continue stochastic time series, the current work explores the efficacy of foundation models operating as kernel density estimators; a crucial element within such time series forecasting. Through an application of InPCA, we determine that such in-context kernel density estimation proceeds within a common, low dimensional probability space; along meaningful trajectories that allow for a comparison between histogram, Gaussian kernel density estimation, and LLM in-context kernel density estimation. Through the lens of InPCA, it becomes clear that a profitable characterization of this in-context learning can be made in terms of a two-parameter, adaptive kernel density estimation framework; hinting at mechanistic basis (Appendix A.1) that points in the direction of future research.

Future direction: towards dispersive induction heads.
Recent research has identified induction heads as fundamental circuits underlying many in-context learning behaviors in stochastic systems with discrete states (Olsson et al., ; Bigelow et al., ).
These emergent circuits increase the predicted probability of token combinations that are observed in-context.
However, such discrete mechanisms are insufficient to explain the in-context learning of continuous stochastic systems,
such as the density estimation task we’ve studied.
To address this gap, our Bespoke KDE analysis in Section 4.3.2 reveals that LLaMA might possess a kernel-like induction mechanism, which we term dispersive induction head.
This is an extension to the induction head concept, and operates as follows:

- •
Unlike standard induction heads (Akyürek et al., ), a dispersive induction head increases the predicted probability of not just exact matches, but also of similar tokens or words.
- •
The "similarity" is determined by an adaptive kernel, analogous to our bespoke KDE model.
- •
The influence of each observation on dissimilar tokens decays over context length, mirroring the decreasing bandwidth in our KDE model.

This concept of dispersive induction heads could potentially bridge the gap between discrete (Akyürek et al., ) and continuous (Gruver et al., ; Liu et al., ) in-context learning mechanisms in transformers (Dong et al., ).

## Acknowledgements

This work was supported by the SciAI Center, and funded by the Office of Naval Research (ONR), under Grant Numbers N00014-23-1-2729 and N00014-23-1-2716.

## Appendix A Appendix

### A.1 Mechanistic interpretations of in-context learning.

In recent years a paradigm has emerged for LLMs’ in-context learning abilities: the simulation and application of smaller, classical models, with well-understood algorithmic features, in response to prompt information. As an example, Olsson et al. () identified “induction heads" in pre-trained LLMs, which simulate 1-gram models for language learning. Then, Akyürek et al. () extended this analysis to “n-gram heads" for computing distributions of next tokens conditioned on n previous tokens. It has also been observed in the literature that transformers can perform in-context regression via gradient descent (von Oswald et al., ; Dai et al., ). These works proposed to understand induction heads as a specific case of in-context gradient descent. Akyürek et al. () showed that trained in-context learners closely match predictions of gradient-based optimization algorithms. More recently, Kantamneni et al. () identified mechanisms in transformers that implicitly implement matrix exponential methods for solving linear ordinary differential equations (ODEs).

Our approach differs from these aforementioned studies for the following reason:
we do not train LLMs on synthetic data designed to induce specific in-context learning abilities. Instead, we investigate the emergent density estimation capabilities of pre-trained foundation models (the LLaMA-2 suite) without any fine-tuning. This approach aligns more closely with recent works (Liu et al., ; Gruver et al., ) that focus on the inherent mathematical abilities of foundations models, rather than LLMs that are trained to induce certain behaviors.

### A.2 Low-dimensional structures in learning dynamics.

Despite the success of modern neural networks at learning patterns in high-dimensional data, the learning dynamics of these complex neural networks are often shown to be constrained to low-dimensional, potentially non-linear subspaces. Hu et al. () demonstrated that the fine-tuning dynamics of LLMs such as GPT and RoBERTa can be well-captured within extremely low-rank, weighted spaces. More recently, Mao et al. () showed that, during training, neural networks spanning a wide range of architectures and sizes trace out
similar low-dimensional trajectories in the space of probability distributions.
Their work focuses on learning trajectories $p_{k=0},...,p_{k=t}$, where $k$ indexes the training epoch, and $p_{k}$ is a high-dimensional PDF describing the model’s probabilistic classification of input data.
Key to their observations is a technique termed *Intensive Principal Component Analysis (InPCA)*,
a visualization tool (Quinn et al., ) that embeds PDFs as points in low dimensional spaces, such that the geometric (Euclidean or Minkowski)
distances between embedded points reflect the statistical distances - e.g. the Hellinger or Bhattacharyya distance - between the corresponding PDFs.

Inspired by Mao et al. () and Quinn et al. (), we extend this line of inquiry to investigate whether the in-context learning dynamics of
LLMs also follow low-dimensional trajectories in probability space.
and $\hat{p}_{k}$ is a PDF of dimension $10^{N}$, where N is the number of digits used in the multi-digit representation (see 3).
Each $\hat{p}_{k}$ describes the model’s estimation of the underlying data distribution at a certain context length.
Our findings show that in-context density estimation traces low-dimensional paths.

### A.3 Bespoke KDE cannot imitate all DE processes

This section examines the ability of Bespoke KDE to imitate various density estimation (DE) processes, including LLaMA-2 models (7b, 13b, 70b),
Bayesian histogram, and standard Gaussian KDE.

#### A.3.1 InPCA visualization of original DE trajectories vs. Bespoke KDE imitations

Figure: Figure 8: Comparison of original DE trajectories ($\triangle$) and their bespoke KDE imitations ($\circ$) for LLaMA-2 models (7b, 13b, 70b), Gaussian KDE, and Bayesian histogram. The 2D embeddings capture 97% of pairwise Hellinger distances between probability distributions.
Refer to caption: extracted/5914712/figures/gaussian_0.1_7_13_70_KDE_hist_traj_embed_Hellinger_2D_97.png

Figure 8 illustrates the DE trajectories of various models alongside their bespoke KDE counterparts. We observe that:
1. Bespoke KDE successfully imitates the LLaMA suite with high precision.
2. It trivially replicates the Gaussian KDE trajectory, as expected.
3. However, it struggles to accurately capture the Bayesian histogram’s trajectory.
To further investigate these observations, we analyze the fitted kernel parameters for Gaussian KDE and Bayesian histogram:

Figure: Figure 9: Bespoke KDE parameters fitted to standard Gaussian-kernel KDE. The model accurately recovers the Gaussian kernel shape ($s=2$) and bandwidth schedule ($h=n^{-\frac{1}{5}}$).
Refer to caption: extracted/5914712/figures/uniform_0.1_7_13_70_KDE_hist_traj_embed_fitted_to_gaussian_KDE_hel_dist_False.png

Figure 10 demonstrates that bespoke KDE accurately recovers the parameters of standard Gaussian KDE. In contrast, Figure 10 reveals significant challenges in fitting Bespoke KDE to the Bayesian histogram:
1. The fitted kernel width remains constant, likely reflecting the fixed bin width of the histogram.
2. The near-zero shape parameter suggests a highly peaked distribution.
3. High uncertainties in both parameters indicate a fundamental mismatch between Histogram and KDE.

These results highlight that while Bespoke KDE can effectively model certain DE processes (e.g., LLaMA models and Gaussian KDE), it cannot capture fundamentally different approaches like the Bayesian histogram.

#### A.3.2 Meta-InPCA Embedding of Trajectories

Figure: Figure 11: Meta-InPCA embedding of DE trajectories ($\triangle$) and their bespoke KDE imitations ($\circ$). This 2D embedding captures 94% of pairwise meta-distances between trajectories.
Refer to caption: extracted/5914712/figures/uniform_0.1_7_13_70_KDE_hist_traj_embed_real_silverman_trajectory_embedding_dim1_2.png

To quantify the similarity between the bespoke KDE model and various DE processes, we introduce a meta-distance measure for trajectories and apply InPCA at a higher level of abstraction.

Trajectory distance metric.
We define a meta-distance between two trajectories $\{p_{i}\}_{i=1}^{n}$ and $\{q_{i}\}_{i=1}^{n}$ as the sum of Hellinger distances between corresponding points at each context length:

$$ $D_{traj}(\{p_{i}\},\{q_{i}\})=\sum_{i=1}^{n}D_{\text{Hel}}(p_{i},q_{i})$ (10) $$

where $D_{\text{Hel}}$ is the Hellinger distance defined in Equation 2.

Meta-InPCA procedure.
Given a set of trajectories $\text{traj}_{1},...,\text{traj}_{l}$, we:
1. Compute the pairwise trajectory distance matrix $D\in\mathbb{R}^{l\times l}$ using Equation 10.
2. Apply the InPCA procedure described in Section 3.1 to embed these trajectories in a lower-dimensional space.

Unlike previous InPCA visualizations where each point represented a single PDF, in this meta-embedding, each point represents an entire DE trajectory.

Observations.
Figure 11 reveals several key insights:
1. The Bayesian histogram and its bespoke KDE imitation are far apart, confirming the model’s inability to capture this approach.
2. Gaussian KDE almost exactly overlaps with its bespoke KDE imitation, as expected.
3. LLaMA-2 models (7b, 13b, 70b) are relatively close to their bespoke KDE counterparts but do not overlap exactly.
These observations suggest that while the 2-parameter bespoke KDE model captures much of LLaMA-2’s in-context DE behavior,
there are still certain nuances in LLaMA-2’s algorithm that it doesn’t fully encapsulate.

### A.4 Comparison of InPCA and standard PCA

Figure: Figure 12: Comparison of InPCA and PCA visualizations for density estimation trajectories of LLaMA-2 models. (a) 2D InPCA embedding preserves 91% of pairwise Hellinger distances, revealing LLaMA’s DE trajectories geometrically bounded between the geodesic and Gaussian submanifold. (b) In contrast, 2D PCA preserves only 83% of pairwise L2 distances, displaying erratic oscillations in LLaMA’s DE trajectories without clear geometric relationships.
Refer to caption: x3.png

### A.5 Additional notes on Kernel Density Estimation

This section provides more background on classical results regarding optimal bandwidth schedules $\{h_{i}\}_{i=1}^{n}$ and kernel shapes.
These results are derived by minimizing the Mean Integrated Squared Error (MISE) (Wand & Jones, , Chapter 2.3):

$$ $\text{MISE}(p|\hat{p}_{h})=\mathbb{E}_{x\sim p}\left[\int(\hat{p}_{h}(x)-p(x)) ^{2}\,\textup{d}x\right],$ (11) $$

where $\hat{p}_{h}$ is the kernel density estimate with bandwidth $h$, and $p$ is the true density.

#### A.5.1 Optimal Bandwidth

While MISE (Equation 11) provides a reasonable measure of estimation error, it is often analytically intractable.
To overcome this limitation, researchers developed the Asymptotic Mean Integrated Squared Error (AMISE) (Wand & Jones, , Chapter 2.4), an approximation of MISE that becomes increasingly accurate as the sample size grows large:

$$ $\text{AMISE}(h)=\frac{1}{nh}R(K)+\frac{1}{4}h^{4}\mu_{2}(K)^{2}R(p^{\prime \prime}),$ (12) $$

where $R(f)=\int f(x)^{2}\,\textup{d}x$, $\mu_{2}(K)=\int x^{2}K(x)\,\textup{d}x$, and $p^{\prime\prime}$ is the second derivative of the true density.

AMISE provides a more manageable form for mathematical analysis.
By minimizing AMISE with respect to $h$, we can derive the optimal bandwidth:

$$ $h_{opt}=\left(\frac{R(K)}{n\mu_{2}(K)^{2}R(p^{\prime\prime})}\right)^{1/5}=Cn^ {-\frac{1}{5}}.$ (13) $$

It’s important to note that this expression involves many unknown quantities related to the true density $p$, such as the average curvature $R(p^{\prime\prime})$,
which is not known a priori. The key insight from this derivation is therefore the $n^{-1/5}$ scaling of the optimal bandwidth.
In practice, there are many methods for estimating the pre-coefficient $C$ from data, such as Silverman’s rule of thumb (Silverman, ):

$$ $h=0.9\min(\hat{\sigma},\frac{\text{IQR}}{1.34})n^{-\frac{1}{5}},$ (14) $$

where $\hat{\sigma}$ is the sample standard deviation, IQR is the interquartile range, and $n$ is the sample size.
This rule is derived heuristically based on minimizing the AMISE (Equation 12).
The $n^{-1/5}$ scaling reveals a fundamental trade-off in kernel density estimation: as more data becomes available, we can afford to use a narrower kernel, but the rate at which we can narrow the kernel is relatively slow.

#### A.5.2 Optimal kernel shape

The kernel function $K$ can take various forms. Common choices include:

- •
Gaussian kernel: $K(u)=\frac{1}{\sqrt{2\pi}}e^{-\frac{1}{2}u^{2}}$
- •
Exponential kernel: $K(u)=\frac{1}{2}e^{-|u|}$
- •
Epanechnikov kernel: $K(u)=\frac{3}{4}(1-u^{2})$ for $|u|\leq 1$, 0 otherwise
- •
Tophat kernel: $K(u)=\begin{cases}\frac{1}{2}&\text{for }|u|\leq 1\\
0&\text{otherwise}\end{cases}$

The Epanechnikov kernel is a key result from optimal kernel theory as it minimizes the MISE.

As shown in Figure 13, different kernel shapes can lead to varying density estimation trajectories.
This visualization compares the performance of LLaMA-13b with various kernel density estimators in a 2D InPCA embedding, capturing 87% of the pairwise Hellinger distances between probability distributions.

Figure: Figure 13: 2D InPCA embedding of density estimation trajectories for LLaMA-13b and KDEs with various kernel shapes: exponential ($s=1$), Gaussian ($s=2$), and tophat ($s=\infty$). As $s$ decreases, the DE trajectories gradually shift from outside the Gaussian submanifold to inside, moving closer to LLaMA’s trajectory. The ground truth is a narrow Gaussian distribution. This visualization captures 87% of the pairwise Hellinger distances between probability distributions.
Refer to caption: extracted/5914712/figures/gaussian_multi_kernels_Hellinger_2D_87.png

#### A.5.3 MISE vs. valid statistical distances

In our InPCA analysis, we choose the the Hellinger distance because:
1. It locally agrees with the KL-divergence (Liese & Vajda, ); a measure motivated by information theory and commonly used in training modern machine learning systems, including LLMs (Touvron et al., ).
Mathematically, this agreement means for distributions $p$ and $q$ that are close: $D_{Hel}^{2}(p,q)\approx\frac{1}{2}\mathrm{KL}(p||q)$
2. Unlike the closely-related Bhattacharyya distance (Bhattacharyya, ), KL-divergence, or cross-entropy (Thomas & Joy, ), Hellinger distance is both symmetric and satisfies the triangle inequality, making it a proper metric.
This property makes it more suitable for visualization purposes.(^10^1010InPCA embedding with Bhattacharyya distance or symmetrized KL-divergence typically results in negative distances in Minkowski space, which are harder to interpret (Teoh et al., ))
3. Although the $L_{2}$ distance $||p-q||_{2}^{2}=\int|p(x)-q(x)|^{2}\textup{d}x$ is also a proper distance metric, it does not locally agree with any information-theoretic divergence measure (Amari, ) and is therefore not suited to measure distance between PDFs.(^11^1111Consequently, PCA embedding with such naive distance measure results in erratic trajectories with obscure geometrical features, as explored in Appendix 12).

In contrast, classical KDE theory focuses on minimizing the Mean Integrated Squared Error (MISE), which is fundamentally an $L_{2}$-type distance that disagrees (globally and locally) with information-theoretic divergence measure.
Nevertheless MISE is widely used in classical KDE literature for several reasons:

1. It can be easily analyzed mathematically using bias-variance decompositions (Wand & Jones, , Chapter 2.3).

2. It leads to closed-form solutions for optimal kernels and bandwidths under certain assumptions.

3. It provides a tractable objective function for theoretical analysis and optimization.

As a consequence of its popularity and mathematical tractability, MISE has been used to derive heuristic bandwidth schedules, such as Silverman’s rule (Equation 14).

We note that modern machine learning often prefers loss functions from the $f$-divergence family (Rényi, ),
such as KL-divergence, cross-entropy, and Hellinger distance (Equation 2).
These measures have strong information-theoretic motivations and are often more appropriate for probabilistic models.

Despite the prevalence of $f$-divergences in machine learning, to the best of our knowledge, there are no rigorous derivations of optimal kernel shapes or bandwidth schedules based on these measures.
This gap presents an interesting avenue for future research, potentially bridging classical statistical theory with modern machine learning practices.
Bridging this gap is beyond the scope of this paper.
However, we speculate that this gap is the reason why the optimal kernel shape employed by LLMs differs significantly from those in classical optimal kernel theory (namely, Gaussian and Epanechnikov).

### A.6 Adaptive Kernel Visualization: Comparing LLaMA-2 with KDE

This section provides a comprehensive visual analysis of DE processes across various target distributions,
including Gaussian, uniform, and randomly generated probability density functions.

We present side-by-side comparisons of the DE trajectories for LLaMA-2, Gaussian KDE, and KDE with fitted kernel designed to emulate LLaMA-2’s behavior.
The fitted KDE process closely mirrors LLaMA-2’s estimation patterns,
while both diverge significantly from traditional Gaussian KDE approaches.

Narrow Gaussian distribution target

Figure: Figure 14: Fitted kernel width and shape schedule
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_gaussian_0.1_hel_dist_False.png

Figure: Figure 16: Fitted kernel visualization with narrow gaussian target
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_gaussian_0.1_fitted_kernel_visualization_5_Log_True.png

Wide Gaussian distribution target

Figure: Figure 17: Fitted kernel width and shape schedule
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_gaussian_0.5_hel_dist_False.png

Figure: Figure 19: Fitted kernel visualization with wide gaussian target
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_gaussian_0.5_fitted_kernel_visualization_5_Log_True.png

Narrow uniform distribution target

Figure: Figure 20: Fitted kernel width and shape schedule
Refer to caption: extracted/5914712/figures/fitted_kernel_13_uniform_0.1_hel_dist_False.png

Figure: Figure 22: Fitted kernel visualization with narrow uniform target
Refer to caption: extracted/5914712/figures/fitted_kernel_13_uniform_0.1_fitted_kernel_visualization_0_Log_True.png

Wide uniform distribution target

Figure: Figure 23: Fitted kernel width and shape schedule
Refer to caption: extracted/5914712/figures/fitted_kernel_13_uniform_0.3_hel_dist_False.png

Figure: Figure 25: Fitted kernel visualization with wide uniform target
Refer to caption: extracted/5914712/figures/fitted_kernel_13_uniform_0.3_fitted_kernel_visualization_5_Log_True.png

### A.7 Randomly generated PDFs

For completion, we also test LLaMA-2’s in-context DE ability on randomly generated PDFs.
As noted in Appendix A.5 and Chapter 2.9 of Wand & Jones (), PDFs with high average curvature (defined as $R(p^{\prime\prime})=\int p^{\prime\prime}(x)^{2}dx$) are more difficult
to estimate from data. Therefore, we want to be able to control the average curvature $R(p^{\prime\prime})$ of the generated random PDFs.
We employ a technique for generating random PDF using Gaussian processes.
This allows us to create a wide variety of smooth, continuous distributions with controllable average curvatures.

##### Gaussian Process Generation:

We generate random PDFs using a Gaussian process with a predefined covariance matrix. The process is defined over the interval [0, 1], discretized into $N_{x}=10^{p}$ points, where $p$ is the precision parameter.
The covariance matrix is constructed using a squared exponential kernel (Seeger, ):

$$ $K(x,y)=\exp\left(-\frac{(x-y)^{2}}{2l^{2}}\right)$ (15) $$

where $x$ and $y$ are points in the domain, and $l$ is the correlation length, a crucial hyperparameter in our generation process.

##### Correlation Length ( l 𝑙 l italic_l ):

The parameter $l$ controls the smoothness and regularity of the generated PDFs. Specifically:

- •
Large values of $l$ produce more regular distributions with lower average curvature. These PDFs tend to have smoother features with lower curvatures.
- •
Small values of $l$ result in more irregular distributions with higher average curvature. These PDFs can have sharper features with higher curvatures.

As $l$ increases, the average curvature decreases, indicating a smoother, less curved function.

##### Generation Process:

The random PDF generation involves the following steps:

- 1.
Generate the covariance matrix using the squared exponential kernel.
- 2.
Perform Cholesky decomposition on the covariance matrix.
- 3.
Sample from the Gaussian process using the decomposed matrix.
- 4.
Apply boundary conditions to ensure the PDF goes to zero at the domain edges.
- 5.
Normalize the function to ensure it integrates to 1, making it a valid PDF.

This method allows us to generate a wide range of PDFs with varying degrees of complexity and smoothness, providing a robust test set for our density estimation algorithms. By adjusting the correlation length $l$, we can systematically explore how different estimation methods perform on targets of varying regularity and curvature.
Given a randomly generated PDF, we can calculate its average curvature using two methods: numerical differentiation and analytical derivation.
Numerically, we can approximate the second derivative using finite differences and then compute the average curvature as:

$$ $R_{numeric}(p^{\prime\prime})\approx\frac{1}{N}\sum_{i=1}^{N}\left(\frac{p(x_{ i+1})-2p(x_{i})+p(x_{i-1})}{(\Delta x)^{2}}\right)^{2}$ (16) $$

where $N$ is the number of discretization points and $\Delta x$ is the spacing between points.
Analytically, we can leverage the fact that the derivative of a Gaussian process is itself a Gaussian process Seeger ().
For a Gaussian process with squared exponential kernel $k(x,x^{\prime})=\exp(-\frac{(x-x^{\prime})^{2}}{2l^{2}})$, the expected average curvature can be derived as:

$$ $R_{analytical}(p^{\prime\prime})=\mathbb{E}[\int(p^{\prime\prime}(x))^{2}dx]= \frac{3}{4l^{3}\sqrt{\pi}}$ (17) $$

In our numerical experiments, we find that these two methods agree to high precision, typically within 10% relative error, validating our choice to use Gaussian processes to generate PDFs with controllable curvatures.

Randomly generated distribution at low curvature ($l=0.5$)

Figure: Figure 26: Fitted kernel width and shape schedule
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_random_0.5_seed_15_hel_dist_False.png

Figure: Figure 28: Fitted kernel visualization with randomly generated target
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_random_0.5_seed_15_fitted_kernel_visualization_10_Log_False.png

Randomly generated distribution at medium curvature ($l=0.1$)

Figure: Figure 29: Fitted kernel width and shape schedule
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_random_0.1_seed_15_hel_dist_False.png

Figure: Figure 31: Fitted kernel visualization with randomly generated target
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_random_0.1_seed_15_fitted_kernel_visualization_10_Log_False.png

Randomly generated distribution at high curvature ($l=0.02$)

Figure: Figure 32: Fitted kernel width and shape schedule
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_random_0.02_seed_7_hel_dist_False.png

Figure: Figure 34: Fitted kernel visualization with randomly generated target
Refer to caption: extracted/5914712/figures/fitted_kernel_13b_random_0.02_seed_7_fitted_kernel_visualization_10_Log_False.png

### A.8 KDE with Silverman bandwidth schedule

For the majority of our investigations, we have been comparing LLaMA’s in-context DE process with Gaussian KDE with bandwidth schedule

$$ $h(n)=Cn^{-\frac{1}{5}}$ $$

with $C=1$.However, as noted in Appendix A.5.1, in practice one often estimates the pre-coefficient from data, such as Silverman’s rule of thumb (Equation 14).
In this section, we replicate Figures 3 and 4 from Section 4,
and use Silverman’s rule-of-thumb bandwidth instead of $C=1$.
As shown in Figures 35 and 36, even with Silverman’s pre-coefficient,
Gaussian KDE still shows a clear bias towards the Gaussian submanifold.

Gaussian target

Figure: Figure 35: In-context density estimation trajectories for Gaussian targets. Top row: 2D InPCA embeddings of DE trajectories for Gaussian targets of decreasing width (left to right). Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 94% of pairwise Hellinger distances between probability distributions.
Refer to caption: extracted/5914712/figures/gaussian_0.5_0.3_0.1_13b_real_silverman_Hellinger_2D_grouped.png

Uniform target

Figure: Figure 36: In-context density estimation trajectories for Gaussian targets. Top row: 2D InPCA embeddings of DE trajectories for Gaussian targets of decreasing width (left to right). Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 91% of pairwise Hellinger distances between probability distributions.
Refer to caption: extracted/5914712/figures/uniform_0.5_0.3_0.1_7b_13b_KDE_hist_real_silverman_Hellinger_2D_grouped.png

With Silverman’s heuristic bandwidth, the Gaussian KDE algorithm can now efficiently converge to the ground truth target distribution, performing on par with or even surpassing LLaMA-13b’s in-context DE in terms of speed.
This outcome is not surprising for two reasons:
1. LLaMA is a general-purpose LLM that is not optimized for density estimation tasks.
2. We do not specifically prompt LLaMA to perform DE, so it must infer the task from the raw data sequence.
This comparison is inherently unfair to LLaMA, as Gaussian KDE is specifically designed for DE tasks.

We would like to reiterate that the purpose of these visualizations is not to benchmark LLaMA against existing algorithms, but to distill geometric insights.
The Gaussian KDE, with Silverman’s bandwidth, features much shorter trajectory lengths, and therefore provides fewer geometric insights compared to our previous visualizations with fixed bandwidth.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: What about even the bivariate normal case? The biggest issue with kernel density estimation methods is scaling with dimensionality. The refinement scheme seems sensible and easily adapted to the bivariate case, so would it be possible to run this scheme on bivariate normal data to see how well learning happens there. How are the correlations learned, for example?
soundness: 3
strengths: I think the question is a nice one and this is certainly (at least to me) an interesting paper on density estimation. It is not obvious a priori that the LLM will be able to do density estimation as it does. The trajectory analysis is certainly interesting. I would think this paper can allow us to do density estimation in automated ways. The experiments are well-designed for the most part.
summary: This paper is looking at the classical problem of density estimation with the use of large language models. The question of interest here is how the large language model estimates a density by relying on the data accessible to it "in context". A simple experiment is used to validate the proposal, namely, checking whether presenting an LLM (in this case, the LLaMA-2) with a sequence of points sampled iid from an underlying distribution results in the *next* sampled point eventually converge to come from the true distribution (with the increase of the number of observed data points).

This, of course, brings the question of how exactly does convergence of this process occur. The authors propose looking at low dimensional characteristics of the density estimate trajectories. Heuristically, at least, it appears that the LLM performs a kind of adaptive Gaussian estimation.
weaknesses: I wonder what properties of the density are important here. I don't believe checking on a Uniform is sufficient - uniform densities are easy to estimate with a flexible kernel and indeed that is what the experiments show, this is not surprising. What is more of interest is to try a heavy-tailed example like a Cauchy or a t with a small number of degrees of freedom.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1.	It is observed in part b of Figure 2 that, for both n = 10 and n = 50, the estimated density functions appear to have an upward bias at the edges (when x is close to 0 and 10). Could the authors provide some explanation for this?

By the way, it can be seen that when the context length is limited in Figure 2 b, the LLMs' estimation of the Gaussian distribution seems to resemble the scenario where the proposed bespoke KDE method has a smaller s (depicted in Figure 5, s = 1 for example). Therefore, as seen in Figure 13, the estimation using the exponential KDE appears to be closer to LLaMA's prediction pattern than that using the Gaussian KDE. This might also provide us with some insights.

2.	When using Bayesian histogram to estimate a Gaussian distribution, the authors propose that its slow convergence is 'likely due to the strong influence of its uniform prior'. However, when estimating a uniform distribution, it does not seem to converge faster either. Therefore, this explanation is somewhat contradictory. Furthermore, what would happen if \alpha is not set to 1 (for example, if \alpha = 0)?

Typos:
It seems like there is a missing citation (squared exponential kernel (?)) in line 1253.
soundness: 3
strengths: 1.	Probability density estimation is a classic and important problem in statistics and this paper explores the underlying mechanism of utilizing LLMs' in-context capabilities for probability density estimation, which is novel to me.

2.	The authors designed various experiments that are easy to understand, providing support for their viewpoints and interesting findings regarding the potential patterns of density estimation in LLMs.

3.	The paper has a clear structure and is easy to follow.
summary: This paper investigates the behavioral characteristics of LLMs (more specifically, LLaMA-2) in performing probability density estimation from in-context by using InPCA. It finds that compared to traditional KDE and Bayesian histogram estimation, the proposed KDE with adaptive kernel width and shape resembles their behavior more closely.
weaknesses: 1.	Although using LLMs for probability density estimation is relatively novel compared to traditional methods, the scope of the discussion in the paper is somewhat limited (for example, it focuses on the one-dimensional case where continuous probability densities are discretized into 2-digit numbers). It would be interesting to explore whether LLMs can still make effective predictions and the analysis results remain consistent in more complex scenarios.

2.	When estimating more complex distributions, such as those with higher curvature (l = 0.02) as shown in Figures 32 - 34, the context length may limit the model's estimation performance. It would be worth investigating whether increasing the context length would lead to better estimation results.

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 3
questions: While my main concerns relate to the empirical results, as outlined in the 'weaknesses' section, I have some smaller questions below:

> a uniform distribution over the domain (0, 100), showing a state of neutral ignorance, which is a reasonable Bayesian prior
I'm not so sure that is a reasonable Bayesian prior. To quote David Mackay's book, "ignorance does not correspond to a uniform probability distribution". I would be surprised if Llama has a priori a uniform prior over the digits, since 1 is much more often the leading digit than others. 
As David MacKay argues in his book on Information Theory, "ignorance does not correspond to a uniform probability distribution". Given that lower digits appear more frequently as leading digits (in line with Benford's Law), it would be surprising if LLMs such as LLaMA  adopted a uniform prior over digits without explicit conditioning. Could this be the reason we can see in Figures 2 and 3 that the Llama trajectory starts some distance from the grey point marked as 'ignorance'?

> While LLaMA-2 has a context window of 4096 tokens (equivalent to ∼1365 comma-delimited, 2-digit data points), we limit our
analysis to a context length of n = 200.
Could you clarify here whether n=200 relates to max number of tokens or data points? (I would interpret this sentence to imply n = 200 numbers, except I notice Figure 1 seems to illustrate 200 datapoints which would correspond to many more than 200 tokens)

The paper frequently refers to the target function as a probability density function (PDF). However, since the outputs are discrete values, perhaps it would be more accurate to describe this function as a probability mass function (PMF)?
soundness: 3
strengths: The paper is well written and clearly presented.

The approach is novel, and the topic is of high interest.

There is a large number of supportive experiments and these are presented with well chosen visualisations.
summary: This work explores how LLMs infer probability density functions, derived from in-context observations. It uses a variant of Principal Component Analysis to help visualise the learning dynamics. Llama-2 models are chosen for the experiments, and these are fed data from target distributions which are either uniform or Gaussian. The results are found to be comparable to certain forms of kernel density estimation.
weaknesses: I have a few concerns regarding the scope of the experiments:

- Models
The empirical results would be significantly more compelling if they weren't all using the Llama-2 family of models, as then the results could potentially be claimed to apply more generally to LLMs. As it stands, the title of the paper may be considered slightly misleading since conclusions can only be drawn on this specific family of models, and not LLMs in general.

- Distances
When visualising the results, it would be beneficial not to solely rely upon the Hellinger distance but to also show results derived from other distances such as the KL divergence or Wasserstein distance. Even if this does not lead to such compelling visualisations, this would provide a much more complete picture of how well the predicted distributions align with the targets.

- Distributions
The chosen target probability distributions are very simple (Gaussian & uniform) so this also limits the conclusions the reader can draw. How would the LLM respond to data from a bimodal distribution, for example?

- Uncertainties
In order to establish a robust, reproducible set of empirical results, it would be strongly advisable to assign uncertainty estimates to the key quantities involved, such as the bespoke KDE parameters.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. Clarification on Kernel Selection and Adaptation: How were the kernel shapes and bandwidth schedules chosen, and why were they deemed the most appropriate for LLaMA’s in-context learning trajectories? Did you experiment with non-Gaussian kernel shapes beyond those briefly mentioned?

2. Scope of Data Distributions: Why did you primarily focus on Gaussian and uniform target distributions? How might LLaMA’s in-context DE capabilities generalize to more complex, real-world distributions, such as multimodal or skewed distributions?

3. Empirical Validation of Theoretical Concepts:How can you empirically validate the proposed concept of “dispersive induction heads”? What experiments would directly test this idea beyond the adaptive KDE model?

4. Comparison with Other Transformer Models: How does LLaMA-2’s performance in in-context DE compare to other large language models  or other architectures known for probabilistic reasoning?

5. Generalization Beyond DE Tasks: How do you see the techniques and findings in this paper extending to tasks beyond density estimation, such as other forms of probabilistic reasoning or inference? Could these methods be generalized to conditional density estimation or more complex probabilistic models?

6. Impact of Tokenization on DE Process: How does the tokenization process impact the DE process in LLaMA? Could different tokenization schemes affect the observed DE trajectories or results?

7. Role of Model Size: How does the model size (7B, 13B, 70B) impact the in-context density estimation process? Are there any trends or differences in the DE trajectories of these models, and do larger models perform better in this context?
soundness: 3
strengths: Originality: The paper contributes a novel approach to understanding the in-context learning capabilities of large language models (LLMs) by framing their density estimation (DE) processes as an adaptive form of kernel density estimation (KDE). This is an innovative application of KDE in a context that diverges from traditional statistical methods. The use of Intensive Principal Component Analysis (InPCA) to visualize and interpret learning trajectories in probability space adds further originality, combining ideas from both geometric analysis and machine learning in a fresh way. The concept of interpreting LLaMA’s learning trajectory as a kernel-like algorithm, with an adaptive kernel, is particularly creative and represents an advancement in understanding how LLMs process probabilistic information.

Quality: The research is rigorous and well-supported by both theoretical and empirical analyses. The application of InPCA, along with detailed comparisons between LLaMA’s trajectories and traditional DE methods like histograms and Gaussian KDE, demonstrates the authors’ thoroughness in exploring the model’s behavior. The paper is grounded in solid mathematical foundations, using well-established statistical measures such as the Hellinger distance to quantify and compare learning trajectories. The methodology is well-executed, and the visualizations (e.g., 2D embeddings of DE trajectories) are informative and enhance the quality of the research.

Clarity: The paper is generally clear in its explanations of complex concepts, especially given the technical nature of the topic. The authors effectively explain the connection between KDE and LLMs’ in-context learning, providing sufficient background for the reader to follow the argument. The figures and visualizations, such as those illustrating LLaMA’s DE process and its comparison with classical methods, are well-designed and help clarify the main findings. However, some sections, particularly around the technical details of InPCA and bespoke KDE modeling, may require more careful reading by less specialized audiences.

Significance: The significance of the work lies in its contribution to the growing field of in-context learning for LLMs. By providing a new geometric and probabilistic interpretation of LLaMA’s behavior, the paper has the potential to impact future research on both LLMs and kernel-based methods. The idea of modeling LLMs’ probabilistic reasoning with an adaptive kernel framework offers a new lens through which to study these models, which could inform the development of more efficient or interpretable LLMs in various domains. Moreover, the paper opens avenues for future exploration, such as investigating the role of dispersive induction heads in continuous probabilistic learning, making it a significant contribution to both theoretical and practical applications in AI.
summary: The paper titled “Density Estimation with LLMs: A Geometric Investigation of In-Context Learning Trajectories” investigates the ability of large language models (LLMs), specifically LLaMA-2 models, to perform density estimation (DE) from in-context data. Density estimation is crucial for probabilistic modeling, where the goal is to estimate a probability density function (PDF) based on observed data. The authors leverage Intensive Principal Component Analysis (InPCA) to visualize and analyze the in-context learning dynamics of these models, showing that LLaMA-2 follows distinct low-dimensional trajectories in probability space when compared to traditional methods such as histograms and Gaussian kernel density estimation (KDE). A key insight from the paper is that LLaMA-2’s in-context DE process can be interpreted as an adaptive KDE with a dynamic kernel shape and width, which captures the model’s behavior with high precision. This adaptive KDE model, with only two parameters, successfully explains the geometric features of the model’s DE trajectories, offering a deeper understanding of how LLMs perform probabilistic reasoning in-context. The paper also speculates that LLaMA-2 may use a mechanism similar to the dispersive induction head to adjust predictions based on observed data.
weaknesses: While the paper “Density Estimation with LLMs: A Geometric Investigation of In-Context Learning Trajectories” offers a strong contribution, there are several areas where it could be improved. 

1. Insufficient exploration of broader datasets: The paper focuses primarily on Gaussian and uniform distributions for testing LLaMA’s in-context density estimation (DE) capabilities. While these distributions are useful for illustrating the model’s ability to perform DE, they are relatively simple and may not fully capture the complexity of real-world probabilistic learning tasks. Incorporating more complex and diverse target distributions, such as multimodal or skewed distributions, would significantly strengthen the generalizability of the findings. By broadening the scope of the experiments, the authors could better demonstrate the robustness of LLaMA’s DE across a wider variety of data structures.

2. Incomplete explanation of adaptive kernel choices: The core contribution of the paper is the interpretation of LLaMA’s in-context DE process as an adaptive KDE, but the reasoning behind the specific choices of kernel shapes and widths, and their relationship to LLaMA’s behavior, could be made clearer. For instance, while the authors introduce a flexible kernel with a shape parameter, the empirical results seem to focus primarily on the Gaussian kernel and only touch briefly on other possibilities. Expanding this section to include more detailed justifications for the selection of kernel parameters and providing more comparative results with different kernel shapes would offer deeper insights into why LLaMA behaves as it does.

3. Theoretical depth versus empirical validation: The paper proposes several intriguing theoretical concepts, such as the idea of dispersive induction heads, but these remain speculative without much empirical validation. The authors propose that LLaMA employs a dispersive induction mechanism to handle in-context density estimation, yet this idea is not thoroughly tested or backed up with concrete evidence. To improve the work, the authors could design specific experiments or ablation studies to directly probe the existence of this mechanism, or alternatively, provide more substantial theoretical evidence. This would move the idea from speculation to a more solid theoretical contribution.

4. Comparative analysis with more models and baselines: The paper primarily compares LLaMA’s performance against classical KDE and Bayesian histograms. However, it would benefit from including a broader set of baseline models, including other transformer-based LLMs or models specifically designed for probabilistic reasoning. For instance, comparing LLaMA’s performance against other architectures known for in-context learning abilities could provide more insight into how LLaMA’s density estimation compares to other state-of-the-art models. This would make the experimental results more compelling and allow for a better understanding of LLaMA’s unique contributions.

Meta-review: This is an interesting paper that studies the capability of LLMs to perform density estimation when prompted with iid samples from a unknown distribution interpreting next token probabilities as an estimated PDF that can be interpreted as a kind of adaptive Gaussian estimation.   The papers strengths are that it attempts to shed insight into intriguing properties of in-context learning for a classical statistical task. Weaknesses that were raised were: the distributions studied are possibly too simple e.g.1D problems; and the insights derived may be overfitting only one family/instance of models.  That said, there was unanimous agreement that the paper is well written, well motivated and provides some interpretability through clear experiments on how LLMs can solve a classical task given in-context examples.

Decision: Accept (Poster)