Density estimation with LLMs: a geometric investigation of in-context learning trajectories
Authors: Toni J.B. Liu
Nicolas Boulle
Raphaël Sarfati
Christopher Earls
Venue: ICLR 2025 Poster (2025)
Keywords: Large language models
in-context learning
Intensive Principal Component Analysis
density estimation
Mechanistic  interpretations
High-dimensional learning dynamics
dimensionality reduction
Primary area: interpretability and explainable AI
TL;DR: We visualize and analyze LLMs' trajectories of density estimations from data observed in context.

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.
Our codebase, along with a 3D visualization of an LLM's in-context learning trajectory, is publicly available at https://github.com/AntonioLiu97/LLMICL_inPCA.

Published as a conference paper at ICLR 2025
DENSITY ESTIMATION WITH LLMS: A GEOMETRIC IN- VESTIGATION OF IN-CONTEXT LEARNING TRAJECTO- RIES
Toni J.B. Liu Department of Physics Cornell University, USA 1134998coznel1.es
Raphaäl Sarfati
School of Civil and Environmental Engineering
Cornell University, USA
raphael warfatticooneel........
Nicolas Boullé Department of Mathematics Imperial College London, UK
Christopher Earls
Center for Applied Mathematics School of Civil and Environmental Engineering
n. n.boul @o@ imperialli....mmm..mmmmmmmmiimme .c.uk
Cornell University, USA earl worneell.cem
ABSTRACT
Large language models (LLMs) demonstrate remarkable emergent abilities to per- form in-context learning across various tasks, including time-series forecasting. This work investigates LLMs' ability estimate probability density functions (PDFs) from data observed in-context; such density estimation (DE) is a funda- mental task underlying many probabilistic modeling problems. We use Intensive Principal Component Analysis (InPCA) to visualize and analyze the in-context learning dynamics of L.I A-M.2, Gemma, and Mistral. Our main finding is that these LLMs all follow similar learning trajectories in low-dimensional InPCA space, which are distinct from those of traditional density estimation methods such histograms and Gaussian kernel density estimation (KDE). We interpret these LLMs' in-context DE process as a KDE algorithm with adaptive kernel width and shape. This custom kernel model captures significant portion of the LLMs' behavior despite having only two parameters. We further speculate on why the L.Ms' kernel width and shape differ from classical algorithms, providing insights into the mechanism in-context probabilistic reasoning in I.IMs.. Our codebase, along with 3D visualization of LLM's in-context learning trajectory, is publicly available at https: http//gitttuu www.:///.ooo/ttoooooooooooo LLMICL inPCA
INTRODUCTION
Modern Large Language Models (LLMs) showcase surprising emergent abilities that they were not explicitly trained for (Brown al. 2020 Dong et al. (2024), such as learning from demonstrations et al. 2023) and analogies in natural language (Hu et al. 2023). Such capacity to extract patterns directly from input text strings, without relying on additional training data, is generally referred as in-context learning (ICL).
Recently, LLMs have been shown achieve competitive performance various mathematical problems, including time series forecasting (Gruver et al. 2024). inferring physical rules from dynamical systems (Kantamneni et al., 2024, Liu et al., 2024). and learning random languages (Bigelow al.) 2024). To solve these tasks, an LLM must possess modeling - the ability to infer conditional or unconditional probability distribution structures by collecting relevant statistics from in-context examples (Akyürek et al. 2024).
capacity for probabilistic
We investigate LLMs' ability to perform density estimation (DE), which involves estimating the probability density function (PDF) from data observed in in-context. Our core experiment is remarkably straightforward. As illustrated in Figure T we prompt LLMs such as LLaMA-2 (Touvron et al.)
2023), Gemma (Gemma Team al. 2024 and Mistral (Tiang al. 2023) with series of data

---

Published as a conference paper at ICLR 2025
points {i+ sampled independently and identically from underlying distribution p(z). We then observe that the LLMs' predicted PDF, (2), for the next data point gradually converges to the ground truth the context length (the number of in-context data points) increases
To interpret the internal mechanisms (Olsson et al. 2022 Bietti et al,, 2023 Dai et al. 2023
von Oswald et al. 2023) underlying LLM's s in-context DE process,
Intensive Principal
Component Analysis (InPCA) (Teoh et al. 2020; Quinn et al. 2019 Mao et al. 2024 to embed the
estimated PDF at each context length {P,) in the probability space (Figure Our visualizations
reveal that the in-context DE processes the tested LLMs 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 strong bias towards Gaussianity, which we argue a telltale feature of kernel-based density estimation (Rosenblat) 1956; Wand & Jones 1994] Silverman 2018). This observation inspires us model an LLM's 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 the in-context learning trajectories with high precision.
p(x)
P(x)
62
00
0.2
JDx ofu
Secaled
50
100
125
150
175
200
context length
Figure 1: In-context density estimation experiment. LLaMA-2 13b prompted with 200 numbers sampled from p(x) (left in red) (Appendix A.9,, and predicts the PDF p(x:) (right in blue) for the next number.
Main contributions.
We introduce a framework for probing an LLM's ability to estimate unconditioned PDFs from in-context data.
2. We apply InPCA analyze the low-dimensional features in the in-context learning trajecto- ries of various LLMs.
3. We interpret the LLMs' in-context DE algorithm as a form of adaptive kernel density estimation, providing evidence for dispersive induction head mechanism (Olsson al. 2022 Akyürek et al. 2024).
BACKGROUND
In-context learning of stochastic dynamical systems. Gruver et al. (2024) pioneered LLM-base based time-series prediction by prompting an LLM with time-series data serialized comma-delimitee 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 al. 2024 Requeima et 2024 Zhang al. 2024
Liu al. 2024 for evaluating LLMs'
numerical abilitiee.
Building on this prompting technique, Liu (0244) tested LLMs' ability to in-context learn the transition rules of various stochastic and chaotic dynamical systems. They introduced a recursive algorithm, termed Hierarchy-PDr. to extract an I.LM's predicted PDF for the next data point based
Our prompts consist purely of comma-delimited multi-digit numbers, without any natural language instruc-
tions.
2

---

Published as a conference paper at ICLR 2025
on histories observed in-context: p(z; (u. ................................ X1-1). This allowed them to rigorously compare the predicted PDF against the ground -runt transition probability p(ze[#/-1)- Our work leverages the data serialization technique proposed by Gruver et al. (2024) and extracts LLMs' probabilistic predictions using the Hierarchy-PDDF algorithm introduced by Liu et al.
2024
Density Estimation. DE (Izenman 1991) is a Long-standing statistical problem with both classical
(Silverman 2018 Rosenblat 1956 Scott 1979 Lugosi & Nobel, [996] Parzen 1962) and modern
prcchnne-learnet solutions (Papamakarios et al. 2021 Sohl--Dicsttnn et al. 2015). It underlies the
learning of more complex stochastic processes. For example, in the first-order Markov process studied in @Luu et al. 2024, Bigelow et al. 2024 Zekri et al. 2024). ; the transition probability p(zz|_;;1) has to be estimated in context. This can be viewed as conditional density estimation problem, where for each possible value z;-1, 1, the density of zy must be estimated. 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.
METHODOLOGY
Here, explain how we prompt LLM 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 steps, the first of which are illustrated in Fig. 2
1. Sampling and prompt generation. Starting from ground-truth probability density function p(x), we generate series of independent samples from Xo...., X Consistent with Gruver et al. (2024), we then serialize and tokenize this series into a text string consisting of sequence of comma-delimited 2-digit numbers, forming a prompt that looks like "6 1 4,,2,,5 59,... ..1.22... (a)).
8, (Figure (a)).
2. Extracting estimated densities with Hierarchy-PDF, Upon prompting with such a text string, we read out the LLM's softmax prediction over the next token, yielding probabilities for 10 tokens (0-9 creating a coarse, 10-bin PDF spanning x E (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(z:), which is a discrete PDF object consisting of 10º bins. where N is the number of digits used in representing each number We interpret the extracted PDF p(x;) as the LLM's estimation of the ground-trutt p(x).
3. Visualizing DE trajectories with InPCA. When there are very few in-context data, an LLM's estimated (x) is close to uniform distribution over the domain (0, 100), showing a state of neutral ignorance, which is reasonable Bayesian prior (Xie al., 2022). J. However, the number of in- context data (n) increases, the LLM gathers more information about the underlying distribution. As a result, the estimated pr, (x) gradually converges to the ground truth p(x) (Figure 2 (b)). The series of estimated PDFs (x) P1 (I), .... (I) over context length forms what we term the "in-conted DE trajectory". These trajectories, we argue, offer important clues about how the LLM performs in-context learning.
However, with our 2-digit representation, each p(x:) vector living 99-dimensional space 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
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, ker- nel density estimator (KDE) and Bayesian histograms - in the same wwo-dimensional (2D) space. Surprisingly, we find that the DE trajectories of LLMs, KDE, and Bayesian histograms are simultane-
TTh logits of all other tokens ignored. We use the Hicrarchy-PDF algorithm [Liu et al. 2024), which performs this recursive search efficiently for transformers.
$ $102 - 1, where the -1 dimension is from the normalization constraint that p(z)dz= = 1.
'The idea of using InPCA to visualize the training dynamics of machine learning systems has been explored
Quinn al. 2
Mao et al. (2024), which is further discussed in Appendix A.2
3

---

Published as a conference paper at ICLR 2025
a
true density function
307
sample & tokenize
100
10
5 ...5, 6 4
Hierarchy-PDF
LLM
rescaled
estimated density functions
LLaMA 130
Mistral v0.3 7b
Germa 2b
5=5
context length (m)
Pm
30
10-2
n 200
10"2
10
rescaied x
InPCA embeddings of density estimation trajectories
cumulative explained variance
explained Freeciom
0.00
Miciical Component
0.0
Increasing context
-0.
10P number dimensions
201
-0.4
truch
-0.6
LLaMA-2 13b
Mistral-vO.3 7b
Gemma 2b
-1.25
-1.00
-0.75
-0.50
-0.25
0.25 0.50
Principel Component
geodesic
gaussian manifold
Figure 2: Visualization pipeline of LLMs' in-context density estimation process. (a) Data points are independently sampled from ground truth distribution (Gaussian this example), then serialized comma-delin delma-de, two-digit numbers prompt LLMs, such LLaMA-2, Mistral- v0.3, and Gemma. (b) Hierarchy-PDF extracts each L.LM's estimated density function Pr. at each context length (c) InPCA reveals -iweeennniionl structures in density estimation trajectories, capturing 91 % pairwise Hellinger distances 2 dimensions. Visual guides: gray uniform PDF representing maximal ignorance, deep blue- ground truth PDF, and pink- ID submanifold of centered Gaussians with variances ranging from oc to 0. All LLMs investigated in this work exhibit similar DE trajectories, which are geometrically bounded between the geodesic and the Gaussian sub-manifold.
ously embeddable in this same oww-dimensional subspace. Moreover, the DE trajectories of LLMs geometrically bounded between those of KDE and Bayesian histograms.
Explaining LLaMA-2 with bespoke KDE. We observe that LLMs' in-context DE trajectories consistently remain close to those of KDE, suggesting that LLMs might be employing a kernel-iie algorithm, internally. To investigate this hypothesis, we optimize a KDE algorithm with fitted kernel parameters (using only two parameters).
4

---

Published as a conference paper at ICLR 2025
We now provide detailed explanation of how to apply InPCA 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 sower-dimensionll space, we employ Intensive Principal Component Analysis (InPCA) (Quinn et al.) 2019 Mao et al.
2024 Teoh al. 2020). Given of PDFs {pi){=1. InPCA aims embed them as points {P:) :=1
in owerr-immensionl space Hd (d & m), such that the Euclidean distances between embedded
points closely approximate the statistical distances between the corresponding PDFs:
Py-- D(p. Vi,j {L...., m).
(1)
where .). is chosen statistical distance measure between PDFs, In this work, we choose D(:,)) to be the Hellinger distance [Hellinger 1909), defined as:
pc(x))//2 py (2))///3 2
(he p5)
dx.
The Hellinger distance is ideal for visualizing PDF trajectories for the following reasons it
locally agrees with the KI-ddieegeenc (Liese & Vajda 2006), a non-negative measure commonly
used in training modern machine learning systems, including LLMs (Touvron al. (023), and 2) it is symmetric and satisfies the triangle inequality, making it a proper distance metric suitable for
geometric analysis (Quinn al. 2019 Teoh al. 2020).
Having chosen an appropriate distance measure, we proceed with InPCA involving three steps:
1. Calculate the Hellinger distance between each pair of PDFs, (Ps, pý). This results in pairwise distance matrix D E Rmxmm whose Dij entry equals (Pt. Pÿ).
2. Obtain "centered" matrix W LDL, where Lıj- 8.j 1/m. This step is closely related to multi-dimensional scaling (MDS) (Chen et al. 2007).
3. Perform the eigenvalue decomposition W UEUT. The diagonal entries in E, sorted in descending order, represent the eigenvalues of W. The embedding coordinates for the PDFs are given by // where the columns of U are the corresponding cigenvectors.
The eigenvalues in Σ 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 ALGORITTMS
Kernel density estimation. Kernel Density Estimation (KDE) is a non- non-parametric method for estimating the probability density function of random variable based finite data sample. Given n samples X1, X2....,X8 X,, drawn from some distribution with an unknown density function f. the kernel density estimator pa(x)) function over the support of defined as:
Pa. (z) = nh 1=1
Σ
K
T.
X;
where K is the kernel function and > is the bandwidth (smoothing parameter). The kernel function K is typically chosen to be symmetric probability density function, such the Gaussian kernel K(u)= 22T
2u² We discuss other kernel shapes in Appendix A.7.2
The optimal bandwidth schedule is a central object of study in classical KDE literature (Appendix A.7) By analyzing the Asymptotic Mean Integrated Squared Error (AMISE) (Equation 15). researchers have derived a widely accepted scaling for the optimal bandwidth.
h(n)
Cn -
(4)
'See Appendices A.5 and A.6 -i--divergecce.
f.6 for viusalizations using other distance measures such as 1.2 and the symmetrized
5

---

Published as a conference paper at ICLR 2025
where is the sample size and C is pre-coefficient (Equation 16 bandwidth provides more smoothing bias, which compensates for data sparsity. While the n 1/5 scaling widely accepted, determining the pre-coefficient C' is more challenging (Wand Jones 1994: Silverman 2018). Unless otherwise noted, we use C' 1 for classical KDE in this paper
In the low-data regime, larger
Bayesian histogram. The Bayesian histogram (Lugosi & Nobel 1996> is another non-parametric method for density estimation, and can be formulated follows:
71
(x)
(5)
71 + a.
where (x:) is the estimated probability density for bin containing n, is the number of observed data points in bin i, n is the total number observed data points. MM is the total number of bins and is the prior count for each bin. We set "hallucinated" data point prior to observing any data (Jeffreys 1946). This choice ensures that the histogram algorithm starts from state of maximal ignorance, consistent with LLaMA's in-context DE in the low data regime)
1, effectively populating each bin with one
4 EXPERIMENTS AND ANALYSIS
We visualize and analyze the learning trajectories of an LLM 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 the Ignorance and Truth points (Mao al.
2024).
Gaussian submanifold: A ID manifold of centered Gaussians with variances ranging from 30 to O.
For brevity, this section focuses on LLaMA-2 13B, as other LLMs yield qualitatively similar results, detailed in Appendix(A.4 While LLaMA-2 has a context window of 4096 tokens (equivalent to 1365 comma-delimited, 2-digit data points), we limit our analysis to context length of 200. This limitation is based on our observation that LLaMA's DE accuracy typically plateaus beyond this point.
4.1
GAUSSIAN TARGET
We begin our analysis with Gaussian target distributions of varying widths.
Wide Gaussian target. A wide Gaussian distribution serves simplest target, 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 I.LMMA 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 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
In practice, various heuristics have been proposed to determinc C, such Silverman's 's rule of thumb Silverman 2018 (see Appendices A.7.I and A.1
M = 100 since our most refined estimation from Hierarchy-PDe consists of 100 onc-digit bins "A curious feature which we foreshadowed in Section 3 and further discuss throughout Scction

---

Published as a conference paper at ICLR 2025
LL 13b
0.6
hnllggra grodeete
Princtpll Components
04
ganiss an mantrold
0.
-044
-1.0
-0.5
-0 50
-0.5
nón
075
0.50
-1.00
-075
-010
-025
0.0
0.5
0.50
-1 -1.00
-0
-0.50
-025
3.00
025
0.50
Pricipall Comsonert
Figure 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% pairwise Hellinger distances between probability distributions.
not unique to Gaussian kernels; as demonstrated in Figure 16 KDEs with alternative kernel shapes (e.g., exponential and parabolic) also exhibit strong propensity for the Gaussian submanifold. One potential explanation for this phenomenon lies the nature KDE itself. KDE can be expressed as convolution of two distributions: the empirical data distribution and the kernel shape,
n
Xi
(x(
K
(F Kn)(z),
(6)
nh
where Fn is the empirical distribution and K2 is the scaled kernel. This convolution operation tends to produce Gaussian-like distributions due to the central limit theorem (Fischer 201T). Consequently, DE trajectory near the Gaussian manifold may indicate kernel-style density estimation algorithm.
4.2 UNIFORM TARGET
The Gaussian distribution easily arises in data with additive noise (Fischer 2011) and therefore likely dominates the training data (Touvron al., (023).. What's more, the Gaussian distribution is everywhere smooth, which makes it very easy to estimate from a function approximation point of view Vore T993). Uniform distributions, the other hand, feature non-differentiable boundaries;
difficult to represent by both parametric (De Vore 1993) and kernel-based (Wand & Jones) 1994 Chapter 2.9) methods. For these reasons, we now investigate the in-context DE trajectory with uniform targets.
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 sub- manifold and converge to the target distribution.
Narrow uniform target. the uniform target narrows, Gaussian KDE's performance deteriorates, while LLM-based in-context DE successfully reaches the target. Notably, the LLM exhibits less bias towards the Gaussian aub-manifold, as compared with the narrow Gaussian target case.
In the low-data regime, Gaussian KDB and Bayesian histogram follow the same trajectories that they previously followed for the narrow Gaussian target. However, the LLM already appears to differentiate this target by taking path further from the Gaussian sub-maunn manifold and closer to the geodesic. This behavior suggests that LLMS' in-context DE algorithm is more flexible and adaptive than classical KDE with pre-determined width schedule and shape.

---

Published as a conference paper at ICLR 2025
0.6
Liatt22 130 hategrem
Aincipll Compenert
geodeesi geuss manifeld
0.2
0.-
truch
-0.
-9.75
0.50
-0225
0.00
0.25
0.50
0.75
-2.75
-0.50
-0.25
0.00
0.5
0.50
0.75
-0.75
-9.50
-0.25
0.90
025
0.50
075
Princioal component
10
10
rescaled
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 Corresponding ground truth distributions. These 2D embeddings capture 89% of pairwise Hellinger distances between probability distributions.
row:
4.3 THE KERNEL INTERPRETATION OF IN-CONTEXT DENSITY ESTIMATION
Based on our observations in Sections 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 emulate LLaMA's learning trajectory.
4.3.1 BESPOKE KDE MODEL
construct our bespoke KDE model in two steps:
Step 1: Parameterize kernel shape
We introduce a flexible kernel function K2(x) parameterized by a shape parameter st
-(s))=)
K2 (2)x)
Z(s)
where Z(s) 2r 1) normalizes the kernel to integrate 1, and b(s)
+1) 3
scales it to
maintain unit variance. This parameterization allows us to interpolate between common kernel shapes such exponential (s 1), Gaussian
2), and tophat (s - oo), visualized in Figure
Bespoke kernel
0.3
Bespoke kemet at
Beepoke kernel at = 2.0
Bespoce kemei at 5.0
Bespoke kerel at 999
675
Figure 5: Bespoke kernel interpolates various common kernel shapes.
We visualize the DE trajectories with these three common kernels (s 1, and oc) in Figure 16 of Appendix A.7.2 As the shape parameter (s) decreases, the DE trajectories gradually shift from outside the Gaussian submanifold to inside, moving closer LLaMA's trajectory. This trend suggests that our parameterization (Equation 7 may be able to capture LLaMA's behavior by allowing to take values below extending beyond the range of common kernel shapes.
We augment the standard KDE formula with kernel shape, resulting in two hyperparameters: and s.
n
K. x
X;
(1) nh
(8)
Step 2: Optimize kernel bandwidth (h) and shape (s)
8

---

Published as a conference paper at ICLR 2025
For given DE trajectory P1 (z)..... distance at each context length i:
(z)) optimize our bespoke KDE to minimize the Hellinger
min +.
Dia ( (x))
yielding the "bespoke KDE" bandwidth schedule {he J==1 and shape schedule {sy 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 [A88 and A.9 We implement this optimization numerically using SciPy, and estimate param- eter uncertainties from the inverse Hessian of the loss function (Kl-divergence) at the optimum. This provides error bars for our fitted kernel shape and bandwidth parameters (Figure(6). In Appendix we claborate the information ceimattio--thooeticc motivation for our uncertainty quantification method.
4.3.2 BESPOKE KDE ANALYSIS
30
34-
10
24-
100
107 Oonttet LeneeT
Figure 7: Bespoke KDE trajectories (o) fitted LLaMA-2 13b's DE tra- jectory (4) on a narrow Gaussian tar- get.
Figure 6: Optimized kernel shape (s) and bandwidth (h) schedules for the bespoke KDE model, fitted LLaMA-2
13b's DE trajectory on a narrow Gaussian target.
Figure 6 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. Two key observations emerge:
1. The fitted kernel width decays significantly faster than the standard KDE bandwidth schedule (n-3). 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 (c.g., Gaussian KDE with s = 2), LLaMA seems to implicitly employ an adaptive kernel. The shape parameter s evolves from 0.1 to 1, with increasing uncertainty for n Σ 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 bespoke KDE that closely imitates an LLM'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 (.9)). Notably, the I.MSs DE performance on these irregular distributions can still be effectively described using our bespoke KDE process (Appendix(A.3) and A.9). Although the DE trajectories of these random and
PDFs lack low-dimensional embeddings, and thus provide less geometric insight, they nonetheless corroborate findings from the cases with more regular distributions.
This consistency across diverse target PDF types strengthens our core hypothesis about the in-context DE mechanism underlying LLMs. Specifically, we posit that an LLM'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.

---

Published as a conference paper at ICLR 2025
5 CONCLUSIONS
Inspired by emergent in-context abilities of LLMs to continue stochastic time series, the current work explores the efficacy foundation models operating as kernel density estimators; 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 comparison between histogram, Gaussian kernel density estimation, and L.M in-context kernel density estimation. Through the lens InPCA, becomes clear that profitable characterization of this in-context learning can be made in terms of two-parameter, adaptive kernel density estimation framework; hinting at mechanistic basis (Appendix A.I. that points in the direction of future research.
Future direction: towards dispersive induction heads. Recent research has identified induction heads fundamental circuits underlying many in-context learning behaviors in discrete-state stochas- tic systems (Olsson et al. 2022; Bigelow et al. 2024). These emergent circuits increase the predicted probability of token combinations that observed in-context. However, such discrete mechanisms are insufficient to explain the in-context learning of continuous stochastic systems, such as the den- sity estimation tasks we have studied. To address this gap, our Bespoke KDE analysis in Section 4.3.2 reveals that LLMs might possess kernel-like induction mechanism, which we term dispersive induction head. This is an extension the induction head concept, and operates as follows:
Unlike standard induction heads (Akyürek et al. 2024). dispersive induction head increases the predicted probability of not only exact matches, but also 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. 2024) and continuous (Gruver al., 2024 Liu al. 2024) in-context learning mechanisms in transformers (Dong et al. 2024).
ACKNOWLEDGEMENTS
This work was supported by the SciAl Center, and funded by the Office of Naval Research (ONR), under Grant Numbers N00014-23-1-2729 and N00014-23-1-27116
REFERENCES
Mohammad Ahsanullah, BM Golam Kibria, and Mohammad Shakil. Normal and student's 's distributions and their applications, volume 4. Springer, 2014.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models, 2023. URL https://
wwwiii.org/aai abs 2211. 11211.116.. 15661
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Architec- tures and algorithms, 2024. URL/https: http://axvvvvvo wrg. wwtp:///g//aaaaaao./o..o.ss 240.s/2971111
Shun-ichi Amari. Information geometry and its applications. Springer, 2016.
Anil Bhattacharyya. On measure of divergence between two statistical populations defined by their probability distribution. Bulletin of the Calcutta Mathematical Society, 35:99-110 10, 1943.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint, 2023. URL https: wt.tpss//arxiv..r00 www.oorgggbbb.ooo.oooo.... /3060000000200000020 00802
Eric J. Bigelow, Ekdeep Singh Lubana, Robert P. Dick. Hidenori Tanaka, and Tomer D. Ullman. In- context learning dynamics with random binary sequences, 2024. URL https: //trxxx..ggg/////// abs/ 2211 0. 17639
10

---

Published as a conference paper at ICLR 2025
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- -osbe Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Fric Sigler, Mateusz I Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https
7/arxiv. org/ abs 2005. 14165
Chun-houh Chen, Wolfgang Karl Härdle, and Antony Unwin Handbook of data visualization. Springer Science & Business Media, 2007.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers, 2023.
URL httpttt/ss//xiiiv... wtt.:/orrgga.o.oeoooooo.oooose
10559
RA Vevor. Constructive Approximation. Springer-Verrs 1993.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li. Heming Xia, Jingjing Xu, Zhiyong Wu, Tianyu Liu, Baobao Chang, Xu Sun, Lei Li, and Zhifang Sui. survey on in-context learning, 2024. htttp::::///x org. 2bxzg/abs/330.... 00234
Albert Einstein. On the electrodynamics of moving bodies. 1905.
Hans Fischer. history of the central limit theorem: from classical to modern probability theory, volume 4. Springer, 2011.
GemmaTeem, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone. Amélie Héliou, Andrea Tacchetti, Anna Bulanova, Antonia Paterson, Beth Tsai, Bobak Shahriari, Charline Le Lan, Christopher A. Proquette-Chox Clément Crepy, Daniel Cer, Daphne Ippolito. David Reid, Elena Buchatskaya, Eric Ni, Eric Noland, Geng Yan, George Tucker, George-Chrissee Christian Muraru, Grigory Rozhdcstvenskiy. Henryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob Austin, James Keeling, Jane Labanowski, Jean-Baptiste Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu, Justin Mao-Jones, Katherine Lee, Kathy Yu, Katie Millican, Lars Lowe Sjoesund, Lisa Lee, Lucas Dixon, Machel Reid, Maciej Mikula, Mateo Wirth, Michael Sharman, Nikolai Chinaev, Nithum Thain, Olivier Bachem, Oscar Chang, Oscar Wahltinez, Paige Bailey, Paul Michel, Petko Yotov, Rahma Chaabouni, Ramona Comanescu, Reena Jana, Rohan Anil, Ross Mcllroy, Ruibo Liu, Ryan Mullins, Samuel L. Smith, Sebastian Borgeaud, Sertan Girgin. Sholto Douglas, Shree Pandya, Siamak Shakeri, Soham De, Ted Klimenko, Tom Hennigan, Vlad Feinberg, Wojeiech Stokowiec, Yu hui Chen, Zafarali Ahmed, Zhitao Gong, Tris Warkentin, Ludovie Peran, Minh Giang, Clément Farabet, Oriol Vinyals, Jeff Dean, Koray Kavukcuoglu, Demis Hassabis, Zoubin Ghahramani, Douglas Eck, Joelle Barral, Fernando Pereira, Eli Collins, Armand Joulin, Noah Fiedel, Evan Senter, Alck Andreev, and Kathleen Kenealy. Gemma: Open models based gemini research and technology, 2024. URL httts:://azxv.. wtwpg//ab//244
08295
Nate Gruver, Mare Finzi, Shikai Qiu, and Andrew Gordon Wilson. Large language models are zero- shot time series forecasters, 2024. URL httthttts://arattt www.ooor///b/oooooooooo hal.ittppss//axxvu..rrrra////.3. 07820
Ernst Hellinger. Neue begründung der theorie quadratischer formen von unendliehvielen veränder- lichen. Journal für die reine und angewandte Mathematik, 1909(136):21-271, 1909.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https
09685
Xiaoyang Hu, Shane Storks, Richard L. Lewis, and Joyce Chai. In-context analogical reasoning with 17626
pre-trained language models, 2023. URL htt.:/ttps::/aaww..or//////
Alan Julian Izenman. Review papers: Recent developments in nonparametric density estimation. Journal of the american statistical association, 86(413):205-224, 1991.

---

Published as a conference paper at ICLR 2025
Harold Jeffreys. An invariant form for the prior probability in estimation problems. Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences, 186(1007):453 461, 1946.
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux. Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023. URL https: //axivv org/ 22300. 0682.
Ming Jin, Shiyu Wang. Lintao Ma, Zhixuan Chu, James Y. Zhang, Xiaoming Shi, Pin- Yu Chen, Yuxuan Liang, Yuan-Fang Li, Shirui Pan, and Qingsong Wen. Time-llm: Time series forecasting by reprogramming large language models, 2024. URL https://arxiv..eeem wwwpxxi.orrg/aaa.. 2310 01728
Subhash Kantamneni, Ziming Liu, and Max Tegmark. How do transformers "do" physics? investigat- ing the simple harmonic oscillator, 2024. URL https:://axxiiii http://awxii..oooab///////////// 2405 17209
Solomon Kullback. Information theory and statistics. Courier Corporation, 1997.
Friedrich Liese and Igor Vajda. On divergences and informations in statistics and information theory.
IEEE Transactions on Information Theory, 52(10):4394 4412, 2006.
Toni J. B. Liu, Nicolas Boullé, Raphaël Sarfati, and Christopher J. Earls. Lims learn governing
principles of dynamical systems, revealing an in-context neural scaling law, 2024. |
//arxiv.org abs wttp:////aabb//22..0...
Gábor Lugosi and Andrew Nobel. Consistency of data-driven histogram methods for density estimation and classification. The Annals of Statistics, 24(2):687-706, 1996.
Jialin Mao, Itay Griniasty, Han Kheng Teoh, Rahul Ramesh, Rubing Yang, Mark K. Transtrum, James P. Sethna, and Pratik Chaudhari. The training process of many deep networks explores the same low-dimensional manifold. Proceedings of the National Academy of Sciences, 121(12), URLhttp:://da doi org/
March 2024. ISSN -6490. doi: 10. (0073//aas.231002211
10. 10.1073/pnas wttpppas.23100000 231C002121
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022. URL pttttps://axxiv.org
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference, 2021. URL
/ettps:/aaxxv..or http://aaa.vv.oo//a///11.1222 02762
Emanuel Parzen. On estimation of probability density function and mode. The annals of mathemat- ical statistics, 3363):005511076, 1962.
Katherine N Quinn, Colin B Clement, Francesco De Bernardis. Michael D Niemack, and James P Sethna. Visualizing probabilistic models and data with intensive principal component analysis. Proceedings of the National Academy of Sciences, 1 116(28):136221133)
2019.
Alfréd Rényi On measures of entropy and information. In Proceedings of the fourth Berkeley symposium on mathematical statistics and probability, volume 1: contributions to the theory of statistics, volume 4, pp. 547-562. University of California Press, 1961.
James Requeima, John Bronskill, Dami Choi, Richard Turner, and David Duvenaud. Llm pro- cesses: Numerical predictive distributions conditioned on natural language. arXiv preprint arXiv:2405. /2856. 2024.
M Rosenblat. Remarks on some nonparametric estimates of a density function. Ann. Math. Stat, 27: 832- 837, 1956.
12

---

Published as a conference paper at ICLR 2025
David W Scott. On optimal and bata-based histograms. Biometrika, 66(3):605 610, 1979.
Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems, 14(02):69-106; 2004.
Chenglei Si, Dan Friedman, Nitish Joshi, Shi Feng, Danqi Chen, and He He. Mcasuring inductive biases of in-context learning with underspecified demonstrations, 2023. URL https:///rxxii
org/ arg/abs/ 2305. 13299
Bernard W Silverman. Density estimation for statistics and data analysis. Routledge, 2018.
Jascha
Eric A. Weiss, Niru Maheswaranathan, and Ganguli. Deep unsupervised learning using nonequilibriu thermodynamies, 2015. URL httpss:/aaxii..rgg/
1503. 03586
Han Kheng Teoh, Katherine N. Quinn, Jaron Kebt-D, Colin B. Clement, Qingyang Xu, and James P. Sethna. Visualizing probabilistic models in minkowski space with intensive sym- metrized kullback-leibler embedding. Physical Review Research, 2()), August 2020. ISSN 1103//
2643-1564. doi: 10.1 2011100/hhyyeeeeeeerrh2.2222222 URL http: http//xxxcoo doi PhysRevResearch, rch., 0332221
CRL.ttpp:.
MTCAJ Thomas and Thomas Joy. Elements of information theory. Wiley-Interscience,, 2006.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nic, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL (ttt:ttss: httpp:s//arxvv..oo org. aabs 2307 09288
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2023. MtRpssttt:///axttt //arxiv wwtpp//p/ppab...oocp////// 07677
Matt P Wand and M Chris Jones. Kernel smoothing. CRC press, 1994.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in- context learning implicit bayesian inference, 2022. URL httpps:/aaxxiv.rrg///a
2111.02080
Oussama Zekri, Abdelhakim Benechehab, and levgen Redko. Can Ilms predict the convergence of stochastic gradient descent?, 2024. URL httpps:/aa.xii......// htarg//bb/2488... 2408 01
abs
Xiyuan Zhang, Ranak Roy Chowdhury, Rajesh Gupta, and Jingbo Shang. Large language models for time series: A survey, 2024. https:
01801
13

---

Published as a conference paper at ICLR 2025
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, response to prompt information. As an example, Olsson et al. (2022) identified "induction heads" in pre-trained LLMs, which simulate 1-gram models for language learning. Then, [Aküüek et al. (2024) 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 in-context regression via gradient descent (yon Oswald et al. 2023 Dai et al. 2023). These works proposed to understand induction heads as a specific case of in-context gradient descent. Akyürek et al. (2023) showed that trained in-context learners closely match predictions of gradient-based optimization algorithms. More recently, Kantamneni et al. (2024) 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 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.) 2024, Gruver et al. (202)) that focus 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. (202T) 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 al. (2024) showed that, during training, neural networks spanning wide range of architectures and sizes trace out similar low-dimensional trajectories in the space of probability distributions. Their work focuses on learning trajectories Pk=0; Pk=t, where k indexes the training epoch, and Pę is a high- -igh-dimensional PDF describing the model's 's probabilistic classification of input data. Key to their observations is a technique termed Intensive Principal Component Analysis (InPCA), visualization tool (Quinn et al.) 2019) that embeds PDFs as points in low dimensional spaces, such that the geometric (Fuclidean or Minkowski) distances between embedded points reflect the statistical distances c.g. the Hellinger Bhattacharyya distance between the corresponding PDFs.
Inspired by Mao et al. (2024) and Quinn et al. (019), we extend this line of inquiry to investigate whether the in-context learning dynamics of LLMs also follow cow-dimensional trajectories in probability space. and p. is a PDF of dimension 10º² where N is the number of digits used in the multi-digit representation (see 3). Each f.. describes the model's estimation of the underlying data distribution certain context length. Our findings show that in-context density estimation traces dow-dimensional paths.
A.3 BESPOKE KDE CANNOT IMITATE ALL DE PROCESSES
This section examines the ability of Bespoke KDE imitate various density estimation (DE) processes, including LLaMA-2 models (7b, 13b, 70b), Bayesian histogram, and standard Gaussian KDE.
A.3.1 DE TRAJECTORIES OF LLAMA-2 7B, 13B, AND 70BF vs. BESPOKE KDE IMITATIONS
Figure 8 iillustrates 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, expected. 3. However, it struggles to accurately capture the Bayesian histogram's 's trajectory. To further investigate these observations, we analyze the fitted kernel parameters for Gaussian KDE and Bayesian histogram:
14

---

Published as a conference paper at ICLR 2025
ignorance
truth
LLaMA-2 13b
Principal Component
LLaMA-2 70b
0.0
gaussian
histogram
KDE fitted to LLaMA-2 7b
-022
KDE fitted to LLaMA-2 70b KDE fitted to gaussian KDE
fitted histogram
-0.4
geodesic
gaussian manifold
-0.6
-1.00 -0.75
-0.50
-0.25
0.00
0.25
0.50
Principal Component
Figure 8: Comparison of original DE trajectories (A) and their bespoke KDE imitations (o) for LLaMA-2 models (7b, 13b, 70b), Gaussian KDE, and Bayesian histogram. The 2D embeddings capture 97% of pairwise Hellinger distances between probability distributions.
panrias
Figure Bespoke KDE parameters fitted to Bayesian histogram. The constant fitted kernel width likely reflects the fixed bin width, while the near-zero shape parameter indicates highly peaked distribution. The high uncertainties suggest fundamental mismatch between the two models.
Figure 9: Bespoke KDE parameters fitted to standard Gaussian-kernel KDE. The model accurately recovers the Gaussian kernel shape 2) and bandwidth schedule (h = n -
Figure p demonstrates that bespoke KDE accurately recovers the parameters of standard Gaussian KDE. In contrast, Figure reveals significant challenges in fitting Bespoke KDE to the Bayesian histogram: The fitted kernel width remains constant, likely reflecting the fixed bin width of the histogram. 2. The near-zero shape parameter suggests 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.. LLMs and Gaussian KDE), it cannot capture fundamentally different approaches like the Bayesian histogram.
15

---

Published as a conference paper at ICLR 2025
A.3.2
META-INPCA EMBEDDING OF TRAJECTORIES; THE LLAMA-2 SUITE
15
LLaMA-2
principal component
LLaMAA2 13b
LLaMA-2 70b
gaussian
histogram
KDE fitted LLaMA-2
KDE fitted to LLaMA-2 130
-5
KDE fitted LLaMA-2 700
KDF fitted gaussian KDF
KDE fitted to histogram
-80
-60
-40
principal companent
Figure 11: Meta-InPCA embedding of DE trajectories (A) and their bespoke KDE imitations (o).
This 2D embedding captures 94% of pairwise meta-distances between trajectories.
To quantify the similarity between the bespoke KDE model and various DE processes, introduce a meta-distance measure for trajectories and apply InPCA at higher level abstraction.
Trajectory distance metric. We define meta meta-distance between two trajectories {p, Dp> and {4 ) as the sum of Hellinger distances between corresponding points at each context length:
Deros Deros((Pi,, (qi))
(Pal qi)
(10)
where Dua is the Hellinger distance defined in Equation
Meta-InPCA procedure. Given a set of trajectories traj traj,. we: 1. Compute the pairwise
trajectory distance matrix glxt 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 entire DE trajectory.
Observations. Figure II 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) 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 still certain nuances in LLaMA-2's algorithm that doesn't 't fully encapsulate.
16

---

Published as a conference paper at ICLR 2025
A.4 META-INPCA EMBEDDING OF TRAJECTORIES: LLAMA, GEMMA, AND MISTRAL
This section documents additional experiments on three other recently released LLMs: Gemma-2b and -7b (GemmaTTem al., 2024) and Mistral-7b-v0.3 (Jiang al. 2023). These models share similar tokenizers as LLaMA and work well with the Hierarchy- PDF algorithm Liu et al. (2024).
LaMe 7t
0.6
Mistrai-/b-rt
gemma-77
goodasic
04
gaussan manifoid
Prrippa Comporent ent
Ignorance
02
bunh
-02
0.00
Component
900
10
rescaled
Figure 12: Comparison of in-context density estimation trajectories for LLaMA, Gemma, and Mistral.
Top row: 2D InPCA embeddings of DE trajectories for narrow and wide Gaussian targets. Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 91 % of pairwise Hellinger distances between probability distributions.
the in-context DE trajectories of LLaMA-2, Mistral v0.3, and Gemma, are strikingly similar, despite the fact that they were built and trained entirely independently by different teams.
As shown in Figure 12
Trajectory Embedding
20
LLaMA
LLaMA 13b
15
LLaMé--2 70b
prinipal comoonert2
1.0
gemme-76
Mistral-B--0..3
gaussian KDe
Histogrem
KDE ftted to LLaMA To
KDE fited to LLaMA-2 13b fiteed
-5
çemma-20
-10
KOE fitee to cemma 7b
KDE fiteed to Mistral-7-v0..3
comparent
Figure 13: This 2D embedding captures % pairwise meta-distances between trajectories.
embedding of DE trajectories (A) and their bespoke KDE imitations (o).
To quantify the similarity of these models' learning trajectories, we use the meta-InPCA embedding introduced in Appendix A.3.2 to visualize the pairwise trajectory distance from these models. As shown in Figure 13 Gemma and Mistral are opll-approximated by bespoke KDE counterparts. In fact, all kernel-like methods quite similar to each other, forming cluster the right corner of the embedding, with Bayesian histogram isolated on the left, similar to Figure II
Interestingly, Gemma-2b seems similar fistral-7b-v0.3 and LlaMA-2 70b, than Gemma- 7b. Models from the same suites are not always the most alike in terms of in-context DE trajectories.
To summarize, the kernel-like density estimation process observed in this work is not limited to any specific suites of LLMs. We therefore speculate that LLMs might spontancously converge to universal in-context DE algorithm.
17

---

Published as a conference paper at ICLR 2025
A.5
APPENDIX: COMPARISON OF INPCA AND STANDARD PCA
a
InPCA embeddings using Hellinger distance
Ignorance truth
LLaMA-2
0.2
LLaMA-2 13b
LLaMA-2 70b
Principal 2omponent
geodesic
0.0
gaussian manifold
cumulative explained variance
-0.2
2.00
expliained fraction
-1.25
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.5(
Principal Component
dlimensione
b
PCA embeddings using L2 distance
0.4
cumulative explained variance
Principal Component
0.2
expleined frection
2.3
2:50
0.0
0.80
0.2
0.75
number dimensions
Principal Component
Figure 14: Comparison 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.
18

---

Published as a conference paper at ICLR 2025
A.6 INPCA EMBEDDINGS WITH SYMMETRIZED KL KL-DIVERGENCE AND BHATTACHARYYA DISTANCE
a
InPCA embeddings using Bhattacharyya distance
cumui ative explalned variance
1.000
exp explained traction
0.425
0.900
0.875
0.850
10F
number climensions
Agnorance
buth
-2.6
LLaMA-2 76
-1.0
-0.5
-0.5
-0.4
-0.2
LLaHA-2 230
Component
2
geutsien merstord
InPCA embeddings using symmetrized KL-divergence
Inncipal Comperent
0.55
cumulative explained variance
1.00
explained fraction
0.95
Princpal Component
0.90
30
2.0
Corponent
2.0
0.80
13
10a
101
10a
1.-
dimensions
-05
Component
Figure 15; InPCA embeddings of LLaMA-2 in-context DE trajectories using symmetrized KL- divergence and Bhattacharyya distance. The target function narrow Gaussian distribution as shown in Figure 2 Red dots denote negative eigenvalue, and red axis denotes imaginary principal component. (a) 2D InPCA embedding preserves 90% pairwise Bhattacharyya distances. (b) 3D InPCA embedding preserves 87% of pairwise symmetrized KL-diieegeence.
Bhattacharyya distances (Quinn et al. 2019) symmetrized KL-divergence (Teoh et al. 2020) are two other statistical divergences commonly used for InPCA embeddings. They are defined
19

---

Published as a conference paper at ICLR 2025
-In (p.9) p(x) In p(z) dr / ((x)) In g(2) p(x)
- -In( Vo(z)g(o)
(11)
dr
(12)
((x)
Figure shows the InPCA embeddings resulting from these two distances. Both of these distances locally agrees with the Hellinger distance as well as the KL-divergence (Teoh et al. 2020). In fact, the Bhattacharyya distance has deep physical connection to Hellinger distance, which was discussed in Quinn et al. (2019)
However, unlike the Hellinger distance used in this work, both Bhattacharyya distance and sym- metrized KL-divergence violate the triangle in-equality (Quinn et al. 2019), and are therefore not valid distance metrics. Consequently, the eigenvalue decomposition step (Section 3.1 might result in negative eigenvalues, with the corresponding coordinates being imaginary. A linear space with mixture of real and imaginary axes is generalization of the Minkowski space, which is useful in special relativity (Einstein, 1905), where events are represented as points in spacetime with one timelike (imaginary) dimension, and up three spacelike (real) dimensions. The Minkowski metric between two points p = (zr, t) and 9 = (x'; t') is given by:
(p. g) (x (x-x')² (t-
(13)
In our visualization context, the imaginary axis (colored red in Figure 15 plays a role similar to the time dimension Minkowski space. The further separated two points are along the imaginary axis, the closer they are in real distance.
As shown in Panel (a), visualization using the Bhattacharyya distance is quantitatively the same as its Hellinger counterpart, shown in Figure Visualization with symmetrized KL-divergence, shown in the top figure in Panel (b), is also qualitatively similar, except that here LLaMA's estimated density no longer seems to converge to the ground truth (blue dot). However, this is a false impression. The bottom figure Panel (b) reveals a third principal component which is imaginary. The end point of LLaMA'ss estimated density is separated from the ground truth (blue dot) along both the real and imaginary axes, by about the same distance. The positive and negative distances from spacelike and timelike seprations roughly cancel out, which means the statistical distance between LLaMA's estimated density and the ground truth is actually small.
Negative eigenvalues from these metries might be physically meaningful in certain circumstance, such as shown in Teoh et al. (2020): Quinn et al. (2019): Mao et al.) (2024). However, for the purpose of visualizing in-context DE trajectories, they complicate the geometric analysis without adding much insight.
20

---

Published as a conference paper at ICLR 2025
A.7 ADDITIONAL NOTES ON KERNEL DENSITY ESTIMATION
This section provides more background on classical results regarding optimal bandwidth schedules {-r and kernel shapes. These results are derived by minimizing the Mean Integrated Squared
Error (MISE) (Wand & Jones 1994 Chapter 2.3):
I ((in(t) p(x))²² dr dr]
MISE(ylim) Exop
(14)
where Pk is the kernel density estimate with bandwidth and p the true density.
A.7.1 OPTIMAL BANDWIDTH
While MISE (Equation 14 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 1994, Chapter 2.4), an approximation of MISE that becomes increasingly accurate as the sample size grows large:
1 R(K) 11 (2(K))
AMISE(h)
(15)
nh
where R(S) = f fx))2 dr, p2(K) f (xx) dr, and p" is the second derivative of the true density.
AMISE provides more manageable form for mathematical analysis. By minimizing AMISE with respect to ht we can derive the optimal bandwidth:
R(K)
1/5
Ropt
(K)² (M))R(()(("))) = CC
(16)
It's important to note that this expression involves many unknown quantities related to the true density p. such as the average curvature R(y"", which is not known a priori. The key insight from this derivation is therefore the n 1155 scaling the optimal bandwidth. In practice, there are many methods for estimating the pre-coefficient C' from data, such as Silverman's 's rule of thumb (Silverman) 2018):
IQR 1.34 ) m
h = 0.9
(17)
where â 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 15 The n -1/5 scaling reveals a fundamental trade-off in kernel density estimation: as more data becomes available, we can afford to narrower kernel, but the rate at which
narrow the kernel is relatively slow.
A.7.2 OPTIMAL KERNEL SHAPF
The kernel function K can take various forms. Common choices include:
Gaussian kernel: K(u)= = 1 fu²
/2T
Exponential kernel: K(u) = _ee
Epanechnikov kernel: K(u) [1 u²) for |u| < 0 otherwise
for |u| < otherwise
Tophat kernel: K(u)=
The Epanechnikov kernel is key result from optimal kernel theory it minimizes the MISE.
As shown in Figure 16 different kernel shapes can lead to varying density estimation trajectories. This visualization compares the performance of LLaMA-131 13b with various kemel density estimators in 2D InPCA embedding, capturing 87% of the pairwise Hellinger distances between probability distributions.
21

---

Published as a conference paper at ICLR 2025
0.8
Principal Component 2
0.6
truth
Ignorance
LLaMA-2 13b
0.4
exponential KDE
gaussian KDE
epanechnikov KDE
tophat KDE
0.0
geodesic
gaussian manifold
-0.2
-0.4
-1.5
-1.0
-0.5
0.0
0.5
Principal Component
Figure 16: 2D InPCA embedding of density estimation trajectories for LLaMA-134 and KDEs with various kernel shapes: exponential the DE trajectories gradually shift from outside the Gaussian submanifold to inside, moving closer to LLaMA's trajectory. The ground truth is narrow Gaussian distribution. This visualization captures 87% of the pairwise Hellinger distances between probability distributions.
1). Gaussian = 2), and tophat (s co). As decreases,
A.7.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 2006); measure motivated by information theory and commonly used in training modern machine learning systems, including LLMs (Touvron et al. 2023). Mathematically, this agreement means for distributions p and q that close: (2Htr(p.9)
(L((l))) 2. Unlike the glosely-related Bhattacharyya distance (Bhattacharyya 1943). KL-divergence, or cross-entropy (Thomas & Joy 2006), 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 3. Although the L2 distance |pp - e113 = p(x) ((x)) dr is also a proper distance metric, does not locally agree with information woommation-thhorttiio divergence measure (Amari 2016) and is therefore not suited to measure distance between PDFs
In contrast, classical KDE theory focuses on minimizing the Mean Integrated Squared Error (MISE), which is fundamentally an L2-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 varianer decompositions (Wand & Jones 1994 Chapter 2.3).
2. It leads to closed-form solutions for optimal kernels and bandwidths under certain assumptions.
It provides a tractable objective function for theoretical analysis and optimization.
As consequence of its popularity and mathematical tractability, MISE has been used to derive heuristic bandwidth schedules, such as Silverman's rule (Equation 17).
We note that modern machine learning often prefers loss functions from the f-divergence family (Rényi 1961), such as KL-divergence, cross-entropy, and Hellinger distance (Equation 2}. These measures have strong information-theoreel motivations and are often more appropriate for proba- bilistic models.
InPCAA emhedding with Bhattacharyya distance or symmetrized Kl.-divergence typically results in negative distances in Minkowski space, which are harder to interpret (Teoh et al. (000))
"Consequently, PCA embedding with such naive distance measure results erratic trajectories with obscure geometrical features, explored Appendix 14
22

---

Published as a conference paper at ICLR 2025
Despite the prevalence of --ivergeences in machine learning, the best of our knowledge, there rigorous derivations of optimal kernel shapes or bandwidth schedules based on these measures. This gap presents 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, speculate that this gap the reason why the optimal kernel shape employed by LLMs differs significantly from those in classical optimal kernel theory (namely, Gaussian and Epanechnikov).
23

---

Published as a conference paper at ICLR 2025
A.8 ADAPTIVE KERNEL VISUALIZATION; COMPARING LLAMA-2 WITH KDE
This section provides a comprehensive visual analysis of DE processes across various target distribu- tions, 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 emulate LLaMA-2's 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
shaps
Parno wid8)
1.
Figure 18: 2D inPCA cap- turing 94% Hellinger vari- ances
30
152
Contect
Figure 17: Fitted kernel width and shape schedule
context length
LLaMA-2 13b
gaussian KDE
fitted fitted_KDE KDE
fitted keme
164
10
10'
Donety
10-1
10:2
2
55.0
-25
2.5
5.0
context length 20
LLaMA-2 13b
gaussian KDE
fitted KDE
fitted komel
gaussan kome
101
109
Direity
10
10-2
10
20
3.0
2.5
0.0
25
5.0
context length 80
LLaMA-2 13b
gaussian KDE
fitted KDE
10.
100
Donsty
10
-5.0
-2.5
2.5
5.0
context length 140
LLaMA-2
gaussian
fitted KDD
fitted keme!
geu55 an keme
10'
201
10
10
10
true PDF
Figure 19: Fitted kemel visualization with narrow gaussian target
24

---

Published as a conference paper at ICLR 2025
Wide Gaussian distribution target
10
fiieed betrel
beae shaps
grodustc
36-
Nome 00
204
1.0
30
Figure 21: 2D inPCA cap- turing 88% of Hellinger vari- ances
Lungth
Figure 20: Fitted kernel width and shape schedule
context length 5
LLaMA-2
gaussian
fitted KDE
gaussan
10'
109
10
10
10
context length 10
LLaMA 13b
gaussian KDE
fitted KDE
gaussien
1n
100
10-1
10
-5.0
2.5
5.0
context length 20
LLaMA-2 13b
gaussian KDE
fitted KDE
fitted kemel
gauSs aussan kome
Danstt
102
1.09
10-1
10
10
10
2
10
5.0
2.5
25
5.0
context length 140
LLaMA-2
gaussian
fitted, KDE
kemel
gaussian keme
10b
10"
10
10
truc PDF
10
10a
10
10
10
Figure 22: Fitted kernel visualization with wide gaussian target
25

---

Published as a conference paper at ICLR 2025
Narrow uniform distribution target
20
GIASEA
MM
shop2
LUPPIO 13
4.5
ATE
032
Fann dih
Mormll shote 00
029
18
1.2
-03
Figure 24: 2D inPCA cap- turing 92% Hellinger vari- ances
20
15
Figure 23: Fitted kernel width and shape schedule
context length 0
LLaMA 13b
gaussian KDE
fitted_ KDE
fiteed boree
gauss keme
10
Damiity
10-1
1003
10
10a
5.0
255
5.0
context length 5
LLaMA-2
gaussian
fitted_ KDE
fitted keme!
geuss an keme
1.0
10
Doreddy
10
10
10
10
context length 10
LLaMA-2
gauzsian
fitted KDI
fitted keme!
gauss an kame
10'
202
10"
10
10
10
context length 60
LLaMA
gaussian
fitled kemel
geuss en keme
101
30
10"
10
10
10
context length 100
LLaMA-2 13b
gaussian KDE
fitted KDE
komel
10*
100
Dansty
10"1
10
2
-5.0
-2.5
2.5
5.0
true PDF
2
Figure 25: Fitted kernel visualization with narrow uniform target
26

---

Published as a conference paper at ICLR 2025
Wide uniform distribution target
10
Itte ternet sndi
Rildd batnel shaps
grodsscc
Heccpa Compaeei
1.2
N00 00010 w550
Musoo
30-
00
Figure 27: 2D inPCA cap- turing 83% of Hellinger vari- ances
Length
Figure 26: Fitted kernel width and shape schedule
context length 5
LLaMA.2
gaussian
fitted KDE
gaussan
10'
109
10
10
10
context length 10
LLaMA 13b
gaussian KDE
fitted KDE
gaussien
1n
180
Demiiyy
10-1
10
context length 50
LLaMA-2 13b
gaussian KDE
fitted KDE
gauzs
10'
100
10
context length 80
LLaMA 13b
gaussian KDE
fitted KDE
gauss
10
100
Deriity
10-1
context length 140
LLaMA-2 13b
gaussian KDE
fitted KDE
geuzs
10'
100
10
-2.5
2.5
5.0
10
21th
20
Figure 28: Fitted kernel visualization with wide uniform target
27

---

Published as a conference paper at ICLR 2025
A.9 RANDOMLY GENERATED PDFs
For completion, we also test LLaMA-2's in-context DE ability on randomly generated PDFs. As noted Appendix a.7 Chapter 2.9 of Wand & Jones (1994), PDFs with high average curvature (defined as R(p") = [T"(")
are more diffiult to estimate from data. Therefore, we want to be able to control the average curvature R(p") of the generated random PDFs. We employ 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 Gaussian process with predefined covariance matrix. The process is defined over the interval [0, 1], discretized into N₂ 102 points, where is the precision parameter. The covariance matrix is constructed using squared exponential kernel (Seeger 20044:
(x: y)²
K(z, v) = exp
(18)
2/2
where x and y are points in the domain, and I is the correlation length, crucial hyperparameter in our generation process.
Correlation Length (I): The parameter 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 I result more irregular distributions with higher average curvature. These PDFs can have sharper features with higher curvatures.
As 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 corre- lation length 7, we can systematically explore how different estimation methods perform on targets of varying regularity and curvature. Given randomly generated PDF, we can calculate its average cur- vature using two methods: numerical differentiation and analytical derivation. Numerically, we approximate the second derivative using finite differences and then compute the average curvature as:
p(x;-1) 2p(xx) + p ((x-1))
Raumeric(p')
(19)
N
(ax)²
where is the number of discretization points and Ar the spacing between points. Analytically, we can leverage the fact that the derivative of Gaussian process is itself Gaussian process Seeger (2004). For a Gaussian process with squared exponential kernel k(x,x') = exp(
2
expected average curvature can be derived as:
3
Ponalyticat(p") E1 (z" = 4/3 /
(20)
In numerical experiments, 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.
28

---

Published as a conference paper at ICLR 2025
Randomly generated distribution at low curvature (1 0.5)
Biteel lernel
est beree stape
20-
Innizall Conppretn
Nome showe (g)
Name midi
224
Figure 30: 2D inPCA cap- turing 70% of Hellinger vari- ances
ConGONt Langoh
Figure 29: Fitted kernel width and shape schedule
context length 10
LLaMA-2 13b
gaussian KDE
fitted_Ki KDE
fitled kerrel
gaussian kerdd
0.3
0.4
0.0
10
10
-3.0
-25
C.0
25
300
context length 20
LLaMA-2 13b
gaussian KDE
fittd_KDE
fitted kerrel
gaussian kerd
0.75
Densty
0.2
0.25
0.0
10
0
10
-30
-75
0.0
7.5
50
context length
LLaMA-2 13b
gaussian KDE
fitted KDE KDE
fitled kerrel
keend
10
0.0
6.0
50
context length
LLaMA- 13b
gaussian KDE
fitted
fittedKKDE KDE
gaussian
6.5-
6.0
context length 140
LLaMA-2 13b
gaussian KDE
fitted KDE
Ritedd Qaussian
0.0
true PU!
0.
01
0.0
Figure 31: Fitted kernel visualization with randomly generated target
29

---

Published as a conference paper at ICLR 2025
Randomly generated distribution at medium curvature (1 0.1)
STNTRIOR
Lanioz L1B
GUOCREC
Gmdd Smadc GP
Patr Mdd
Shape
Figure 33: 2D inPCA cap- turing 73% Hellinger vari- ances
100
10
10
Quitrot Leagth
Figure 32: Fitted kernel width and shape schedule
context length 10
LLaMA-2
gaussian
fitted_KD KDE
fitted kernel
çaussian kerd
0.4
Danety
0.2
0.2
0.0
-7.
1.0
context length 20
LLaMA-2 13b
gaussian KDE
fitted KDE
fittod karrel
gausstan kand
0.3
Danatt
0.0
10
-5.0
-2.
0.0
5.0
context length 50
LLaMA-2
gaussian
fitted
kerrel
gaussian kornd
Deneby
10-
0.2
0.5
0.0
context length
LLaMA-2 13b
gaussian KDE
fitted_KDE KDE
gaussien
0.1
context length 140
LLaMA-2 13b
gaussian KDE
fitted KDE
fitted kerrel
0.4
Decalty
2-
0.7
30
-2.5
50
true PDF
0.1
30
Figure 34: Fitted kernel visualization with randomly generated target
30

---

Published as a conference paper at ICLR 2025
Randomly generated distribution at high curvature (1 = 0.02)
fetd terel with
fited lard Shopo
produsi
Conted Length
Figure 36: 2D inPCA cap- turing 55% of Hellinger vari- ances
Figure 35: Fitted kernel width and shape schedule
context length 10 gaussian KDE
LLaMA-2 13b
fitted_ KDE
0.2
context length 40
LLaMA-2 13b
gaussian KDE
fitted KDE
0.3
Densty
0.0
context length 90
LLaMA-2 13b
gaussian KDE
fitted KDE
0.2
0.0
context length 130
LLaMA-2 13b
gaussian KDE
fitted KDE
2
true POF
10
10
0.3
0.
Figure 37: Fitted kernel visualization with randomly generated target
31

---

Published as a conference paper at ICLR 2025
A.10 HEAVY-TAILED DISTRIBUTIONS
In this section, we visualize the in-context DE trajectories of Student's t-distribution (Ahsanullah al. 2014). The t-distribution is parametrized
I 1( Un
(1
(x =
(21)
where v controls the amount of probability mass in the tails. Smaller v leads to heavier tails. For 1, the t-distribution reduces to the Cauchy distribution:
p
(22)
π + x2)
which has undefined mean and variance due to its extremely heavy tails.
Figure 38: In-context density estimation trajectories for t-distribution targets. Top row: 2D InPCA embeddings of DE trajectories for Mistral-7b-v0.3, Middle row: 2D InPCA embeddings of DE trajectories for LLaMA-2 13b. Bottom row: Corresponding ground truth distributions, with v = 0.1. and These 2D embeddings capture 84% of pairwise Hellinger distances between probability distributions.
LLaMA and Mistral trace out similar in-context DE trajectories, as shown in Figure 38 eventually converge the ground truth distribution. However, unlike the cases with Gaussian and uniform targets, the trajectories for learning t-distributions are more erratic. This is likely due to the fact that t-distribution with small v has large variance (diverging if S 1), where the thick tails have non-trivial probability mass. This in turn makes extreme samples occur more frequently, causing instability in any kernel-like density estimators.
Both models
A.11 ADDITIONAL NOTES ON BAYESIAN HISTOGRAM
The Bayesian histogram with uniform prior starts near ignorance by construction, and gradually converges to ground truth distribution with increasing data. Its DE trajectories serve as informative visual comparisons against kernel-like methods. This work does not focus on Bayesian histogram per se. However, for completion, we briefly visualize the learning trajectories of Bayesian histograms at different levels of prior bias, a, as defined Equation 5 This results are shown in Figure 39
Unburdened by strong prior, Bayesian histograms with smaller a converge faster to the ground truth. However, its learning trajectories are also more erratic. This because in the absence of strong regularization, any new data point would cause significant update in the estimated distribution (height the corresponding bin). This effect especially pronounced for high-variance distribution, where extreme data occur often.
Bayesian histograms with larger a, on the other hand, follow the geodesic more conservatively in smooth trajectories. However, they also take more data converge to the true distribution. Note that unlike the other visualizations in this work, each trajectory in Figure 39 consists n = 500 data points, instead of 200.
32

---

Published as a conference paper at ICLR 2025
Iedtguum with
hittogenn with phann grodecx
22-
0.1
-0.93
3.5
675
1.0
-5.0
-0.5
500
032
0.79
120
-2.0
25
510
0.21
332
0.75
1.00
15
Figure 39: Density estimation trajectories for Bayesian histogram at different uniform bias a. Top row: 2D InPCA embeddings of DE trajectories for Mistral-7b-vO.3. Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 87% of pairwise Hellinger distances between probability distributions.
A.12 UNCERTAINTY QUANTIFICATION FOR BESPOKE KDE PARAMETERS
In this section, provide information-thherreiee justifications for estimated uncertainty of the fitted KDE parameters, outlined in Section 4.3.1 On high level, our method is based on the observation that the Hessian matrix at optimal fit approximately equal to the Fisher information matrix (Kullback 1997).
Let P(r) denote an LLM's estimated PDF from data, which treat the ground truth fit Let Qo(x) denote the PDF estimated by parametrized KDE with parameter vector 0 = (h, Technically, both P(x)) and (2(x) are conditioned on the of data observed in-context, {X, ji For simplicity, we drop this conditioning in our notation.
During the fitting process, the loss function is defined as the Kl--divereence between P(x) and Qo
= P(x) log
P(x) dr
DxL()))20) =
(23)
Qo(z)
The second derivative (Hessian) of the loss function with respect to is:
Q+(2) dr (rP(QQ)))) = =- P(x)V (n(())
The Fisher Information Matrix (FIM) defined as:
(0) = -Eq V In [o()))
(m dr
Therefore, at the optimal 6*, if Q+(x) x P(x). then the Hessian the KL divergence provides a good approximation the FIM:
(k*) = Dri (PIIQo) = I(0")
The covariance matrix of 9 can then be estimated as the inverse of the FIM, a relationship known as the Cramér-Rao bound (Thomas & Joy 2006). Under regularity conditions and asymptotic normality, this bound becomes an equality:
Cov(A) I(0)) (PI
Said another way, the uncertainty in the fitted kernel parameters, such bandwidth (h) and shape (s), can be estimated from the inverted Hessian matrix of KL--divergence loss function.
33

---

Published as a conference paper at ICLR 2025
A.13 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
with C' Howeerrr,, noted in Appendix A.7.1 in practice one often estimates the pre-coefficient from data, such as Silverman's rule thumb (Equation 17). In this section, we replicate Figures
and from Section 4 and Silverman's rule-of-thumb bandwidth instead of C = 1. As shown in Figures 40 and [1 even with Silverman's pre-coefficient, Gaussian KDE still shows clear bias
towards the Gaussian submanifold.
Gaussian target
LLMMA-2 1.3
gausstan KUE
Pinceal Componeit?
histegram
geodesic gaussian
0.0
-0.2
-0.75
-0.25 8.00
-0.75
-0500
-0.25
0.00
n25
-0.15 -0.50
-0.25
c.00
335
a5o 0.75
Component
10
ieccaled X
Figure 40: In-context density estimation trajectories for Gaussian targets. Top row: 2D InPCA embeddings of DE trajectories for Gaussian targets of decreasing width (left right). Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 94% of pairwise Hellinger distances between probability distributions.
Uniform target
LLaMte2 130
hettgremm
Princpel Component.
0.4
geodesic
geuss maniffil
butth
0.0
-0.4
0.00 0.75
0.50
-0.50 0.00
0.75
-0.50
0.25
0.50
0.75
Figure 41: 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.
With Silverman's heuristic bandwidth, the Gaussian KDE algorithm can now efficiently converge to the ground truth target distribution, performing on par with even surpassing LLaMA-13b's in- context DE in terms of speed. This outcome is not surprising for two reasons: LLaMA 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, Gaussian KDE is specifically designed for DE tasks.
34

---

Published as a conference paper at ICLR 2025
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 's bandwidth, features much shorter trajectory lengths, and therefore provides fewer geometric insights compared to our previous visualizations with fixed bandwidth.
35


Reviews:
[{"review_id": "HhqlTw8fel", "reviewer": "Reviewer_mikv", "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).\n\nThis, 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.", "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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oo0uQIe2ex", "reviewer": "Reviewer_KbiX", "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.", "questions": "1.\tIt 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?\n\nBy 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.\n\n2.\tWhen 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)?\n\nTypos:\nIt seems like there is a missing citation (squared exponential kernel (?)) in line 1253.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.\tProbability 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.\n\n2.\tThe 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.\n\n3.\tThe paper has a clear structure and is easy to follow.", "weaknesses": "1.\tAlthough 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.\n\n2.\tWhen 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Uv7cwQJT2H", "reviewer": "Reviewer_acu3", "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.", "questions": "While my main concerns relate to the empirical results, as outlined in the 'weaknesses' section, I have some smaller questions below:\n\n> a uniform distribution over the domain (0, 100), showing a state of neutral ignorance, which is a reasonable Bayesian prior\nI'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. \nAs 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'?\n\n> While LLaMA-2 has a context window of 4096 tokens (equivalent to ∼1365 comma-delimited, 2-digit data points), we limit our\nanalysis to a context length of n = 200.\nCould 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)\n\nThe 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)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and clearly presented.\n\nThe approach is novel, and the topic is of high interest.\n\nThere is a large number of supportive experiments and these are presented with well chosen visualisations.", "weaknesses": "I have a few concerns regarding the scope of the experiments:\n\n- Models\nThe 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.\n\n- Distances\nWhen 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.\n\n- Distributions\nThe 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?\n\n- Uncertainties\nIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hxwr9WPfDs", "reviewer": "Reviewer_2CMq", "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.", "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?\n\n2. 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?\n\n3. 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?\n\n4. 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?\n\n5. 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?\n\n6. 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?\n\n7. 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 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.\n\nQuality: 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.\n\nClarity: 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.\n\nSignificance: 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.", "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. \n\n1. 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.\n\n2. 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.\n\n3. 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.\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)