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

TL;DR

The paper investigates whether large language models can estimate unconditioned probability density functions from in-context observations by prompting models with in-context data and tracking the evolving PDF via Hierarchy-PDF. Using Intensive PCA (InPCA), the authors reveal common, low-dimensional density-estimation trajectories across LLMs that lie between a Bayesian-histogram path and a KDE-like path, with a strong bias toward Gaussian-like behavior. They introduce a Bespoke KDE with adaptive kernel width $h$ and shape $s$ that closely reproduces the LLM trajectory, revealing rapid bandwidth decay and evolving kernel shape as context grows, suggesting a kernel-based mechanism in in-context probabilistic reasoning. The findings support a kernel-like interpretation of in-context DE and propose dispersive induction heads as a mechanism linking discrete in-context learning with continuous probabilistic estimation, offering a geometric lens to study LLMs and guiding future interpretable analyses.

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

Figures (41)

• Figure 1: In-context density estimation experiment. LLaMA-2 13b is prompted with 200 numbers sampled from $p(x)$ (left in red) (Appendix \ref{['random_pdf_generation']}), and predicts the PDF $\hat{p}(x)$ (right in blue) for the next number.
• Figure 2: Visualization pipeline of LLMs' in-context density estimation process. (a) Data points are independently sampled from a ground truth distribution (Gaussian in this example), then serialized as comma-delimited, two-digit numbers to prompt LLMs, such as LLaMA-2, Mistral-v0.3, and Gemma. (b) Hierarchy-PDF extracts each LLM's estimated density function $\hat{p}_n$ at each context length $n$. (c) InPCA reveals low-dimensional structures in density estimation trajectories, capturing 91% of pairwise Hellinger distances in 2 dimensions. Visual guides: gray - uniform PDF representing maximal ignorance, deep blue - ground truth PDF, and pink - 1D submanifold of centered Gaussians with variances ranging from $\infty$ to 0. All LLMs investigated in this work exhibit similar DE trajectories, which are geometrically bounded between the geodesic and the Gaussian sub-manifold.
• Figure 3: In-context density estimation trajectories for Gaussian targets. Top row: 2D InPCA embeddings of DE trajectories for Gaussian targets of decreasing width (left to right). Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 92% of pairwise Hellinger distances between probability distributions.
• Figure 4: In-context density estimation trajectories for uniform distribution targets. Top row: 2D InPCA embeddings of DE trajectories for uniform targets of decreasing width (left to right). Bottom row: Corresponding ground truth distributions. These 2D embeddings capture 89% of pairwise Hellinger distances between probability distributions.
• Figure 5: Bespoke kernel interpolates various common kernel shapes.