Transformers as In-Context Samplers: From Closed-Form Diffusion to Estimation-Free Sampling

Summary

The paper asks whether frozen transformers can perform in-context data generation: producing a fresh sample from an unknown distribution specified only by i.i.d. examples in the prompt. Theoretically, the authors prove (Theorem 1) that standard softmax self-attention transformers exactly simulate closed-form and smoothed closed-form diffusion samplers: attention heads compute Gaussian-mixture responsibility weights and weighted empirical means from in-context samples, while feedforward layers apply Euler updates; unlike prior work, no modified (RBF) attention is needed because norms are precomputed in the encoding. They further prove (Theorem 2) that transformers uniformly approximate a finite forward-backward particle discretization inspired by Estimation-Free Sampling, using Laplace-quadrature kernel heads with an explicit error budget. Empirically, nine pretrained LLMs (GPT-2 through OPT-66B) are prompted with semantic-category words and natural CBT sentences; across layers, MMD-squared to a uniform spherical reference and an interacting-particle (logarithmic) energy both follow a U-shape: middle layers move toward uniformity, later layers recover topic structure. Smaller Qwen2.5/GPT-2 models show weaker effects, framed as a depth-limited failure mode. A trained GPT-2-style model generates held-out 'smile' geometry absent from training, supporting genuine in-context sampling. The authors are explicit that the theorems are expressivity statements and do not establish that pretrained LLMs implement these algorithms or a U-shaped energy trajectory. Limitations: purely descriptive statistics (5 prompts per topic, no control conditions), engineered encodings far from practical LLMs, and no code release stated.