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

Arman Adibi; Alireza Jafari; Mohammad Ghavamzadeh; Hadi Daneshmand

Source-linked AI summary

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

Arman Adibi, Alireza Jafari, Mohammad Ghavamzadeh, Hadi Daneshmand

arXiv:2609.08981v1Sep 8, 2026cs.LGcs.AIstat.APstat.COstat.ML

TL;DR

Prior work established in-context learning for inference tasks, but whether frozen transformers can perform unconditional data generation remained open. This paper proves that transformers can simulate diffusion and estimation-free samplers from in-context samples, and reports compatible layerwise distribution and energy dynamics in pretrained models. The authors caution that these observations do not uniquely identify the algorithms executed internally.

Problem
Whether in-context learning extends from supervised prediction to unconditional data generation and generative modeling remains an open question.
Method
The paper constructs transformer parameterizations that use softmax attention for responsibility-weighted empirical means and feedforward layers for sampler updates, then analyzes semantic-topic embeddings across pretrained-model layers.
Results
Transformers can simulate closed-form diffusion and estimation-free samplers from in-context samples, while pretrained models show a two-stage embedding geometry consistent with movement toward and away from a uniform distribution.
Takeaways & Limitations
In-context sampling can be realized as an explicit generative computation, and observed pretrained-model dynamics are compatible with such sampling mechanisms.
Takeaways & Limitations
The theoretical results establish expressivity rather than proving that pretrained language models internally execute these algorithms, and the observed dynamics do not uniquely identify the mechanism.

Abstract

from arXiv · show

A growing body of work establishes that large language models are not mere statistical memorizers, but are capable of in-context learning: performing inference at test time using only examples provided in the prompt, without any parameter updates. Prior theoretical work has shown that this capability extends to supervised learning tasks such as linear regression. We prove that in-context learning extends further to \emph{data generation}: frozen transformers can simulate iterative generative samplers from in-context samples. We first show that transformers can realize closed-form and smoothed closed-form diffusion samplers. The construction identifies a concrete generative role for softmax attention: it computes responsibility weights and weighted empirical averages, while feedforward layers implement Euler updates. To empirically relate these constructions to pretrained language models, we study \emph{semantic-topic sampling}: prompts consisting of words drawn from a common semantic category, such as animals, foods, or cities. Across transformer layers, the normalized hidden states exhibit a two-stage geometry: they move toward a uniform spherical reference in intermediate layers and then return to structured, topic-dependent representations near the output. We further measure an interacting-particle energy on these hidden-state clouds and observe the same U-shape pattern. We then prove that transformers can approximate an energy-based sampler, constructing the same U-shape energy across the layers.

1 Introduction

The paper asks whether transformers can extend in-context learning from supervised prediction to unconditional data generation. It develops theoretical and empirical evidence that frozen transformers can use prompt samples to simulate generative samplers and exhibit related layerwise dynamics.

Research question: The central question is whether a transformer can implement a sampler from i.i.d. in-context samples drawn from an unknown distribution.The sampling task maps x1, . . . , xn ∼ P to a fresh output xn+1 ∼ P.
Motivating experiment: A controlled experiment shows that smile-shaped context points induce samples along an unseen crescent absent from training.This supports the interpretation that the prompt acts as an empirical distribution rather than merely selecting memorized training components.
Theoretical contributions: The paper proves that transformers can simulate closed-form diffusion samplers directly from in-context samples.Softmax attention computes responsibility weights and empirical weighted means, while feedforward layers perform Euler updates.
Empirical analysis: Semantic-topic prompts reveal a two-stage embedding geometry: normalized representations approach a uniform spherical distribution and then return toward structured topic-dependent representations.The analysis uses independently sampled words from common categories such as animals, cities, or foods.
Theoretical contributions: The paper also proves transformer implementations of estimation-free sampling, whose two-step transport mechanism moves empirical data toward a uniform distribution and then toward a target distribution.This construction is intended to match the observed layerwise distribution shift in pretrained transformers.

2 Preliminaries

The paper formulates closed-form diffusion sampling through an explicit Gaussian-mixture score and iterative updates. It then encodes samples and sampler state into transformer tokens so attention and feedforward layers can realize the computation.

Diffusion preliminaries: Score-based diffusion transports a simple reference distribution toward a target distribution using the score of a Gaussian-smoothed data density.The score is the gradient of the smoothed log-density and is typically approximated by a denoising-trained neural network.
Closed-form diffusion: For an empirical distribution, Gaussian smoothing yields an explicit mixture whose score points from the current state toward a responsibility-weighted sample mean.The weighted mean is the main data-dependent computation; the update direction is a fixed linear combination of the current state and this mean.
Closed-form diffusion: The sampler advances iteratively by the Euler recurrence zs+1 = zs + hsvts(zs).The construction computes the velocity from in-context samples without a trained score network.
Smoothed sampling: Smoothing replaces the weighted mean at the current state with an average over perturbed states, recovering the closed-form method when σ = 0.Attention evaluates the weighted means for fixed perturbations, averages them, and the feedforward layer applies the resulting affine update.
Transformer construction: The transformer input stores in-context samples and an initial sampler state, while separate token blocks provide sample data, generated-state memory, and intermediate-computation scratch space.A state readout extracts the z-block of the final row as the generated sample.
Transformer construction: A standard softmax self-attention transformer is parameterized so its final state readout matches a specified iterative generative sampler.The architecture maps an embedding matrix to an output matrix of the same size, with frozen parameters chosen for the sampling computation.

3 Simulation of Closed-form Diffusion with Transformers

The paper proves that standard transformers can simulate closed-form and smoothed closed-form diffusion samplers from in-context samples. The construction uses softmax attention for normalized responsibility weights and feedforward layers for iterative updates, with parameters fixed independently of the realized dataset.

The proof shows that transformers can implement iterative closed-form diffusion updates from prompts containing empirical samples and arbitrary initial states.This establishes the simulation result for the prompt encoding and state initialization used in the theorem.
Theorem 1 establishes that a standard transformer can simulate closed-form diffusion for every empirical dataset and initial state.The transformer uses the same prompt encoding and state readout as the construction described in the section.
The same architecture can simulate a smoothed closed-form diffusion recursion, with parameter choices depending on smoothing settings but not on the realized dataset.The stated dependencies include N, d, τ, and, for smoothing, σ and E.
Softmax attention computes the normalized responsibility weights required by the closed-form diffusion update, unlike linear attention constructions that omit this normalization.The construction uses standard softmax self-attention, while feedforward layers support the iterative computation.
Closed-form diffusion explains expressivity but does not reproduce the observed U-shaped layer dynamics in pretrained language models.The paper therefore distinguishes the ability to implement a sampler from identifying the mechanism learned by a trained transformer.

4 Mechanistic Analysis of Pretrained Transformers

The analysis treats pretrained transformers as shaping in-context word embeddings through the layers. Across semantic-topic and natural-text prompts, embeddings move toward a uniform spherical distribution in intermediate layers before becoming structured again.

Scope: The paper explicitly studies embedding-distribution shaping rather than claiming pretrained language models exactly implement a sampling algorithm.This scope statement limits the mechanistic interpretation of the observed layerwise dynamics.
Measurement: The study measures normalized word-embedding distributions against a uniform spherical reference using Gaussian-kernel MMD2.Prompts contain i.i.d. words from a shared semantic category, such as animals, cities, or foods.
Observed dynamics: Pretrained models exhibit a two-stage pattern: embeddings approach a uniform spherical distribution in middle layers and then return toward non-uniform structure.This pattern appears across three pretrained foundation models, including two LLaMA variants and Cerebras-GPT.
Natural-text prompts: The same qualitative behavior appears for natural CBT sentences, despite computing metrics only once for repeated words.The CBT analysis measures MMD distance between word embeddings and uniform unit-sphere samples.
Figure 2: Figure 2 plots layer index against MMD2 for animal, city, and food prompts, with several models showing a U-shaped profile.Intermediate layers move closer to the uniform reference distribution.

5 Physics of In-context Sampling

The paper connects the observed U-shaped embedding dynamics to interacting-particle energy and estimation-free sampling. It proves that a transformer can simulate the sampler, first uniformizing the empirical distribution and then reversing that transport toward the target.

Energy interpretation: In the logarithmic limit, the interacting-particle energy has minimizers that converge toward a uniform spherical distribution as the number of particles grows.This theoretical property motivates using the energy to explain layerwise embedding evolution.
Energy interpretation: Pretrained-transformer embeddings show a U-shaped energy profile across layers when prompts contain i.i.d. words from one topic.The energy pattern resembles the MMD profile, and optimizing the energy yields a proven sampling method.
Estimation-Free Sampling: Estimation-Free Sampling alternates energy minimization by gradient descent with energy maximization by inverse gradient descent.The forward phase moves the empirical distribution toward uniformity, while the backward phase transports a reference sample toward the empirical distribution.
Transformer simulation: The EFS mechanism first makes the empirical distribution uniform and then non-uniform, matching the qualitative U-shaped dynamics observed in pretrained transformers.The theory connects intermediate uniformization with later recovery of target structure.
Transformer simulation: Under the assumption that EFS iterates remain in a compact set, transformer parameters can be chosen to implement EFS for in-context sampling.The construction depends on a particular input encoding and fixed EFS parameters.

6 Experiments

Experiments test whether transformers generate samples from prompt-defined distributions and whether hidden states undergo an intermediate transport-like phase. Controlled two-moons experiments and pretrained-language-model analyses support both patterns, while energy measurements reinforce the interpretation.

Experimental questions: The experiments ask whether transformers generate samples from prompt-defined distributions and whether hidden states show an intermediate transport-like phase.These questions are evaluated in both controlled and pretrained-model settings.
Controlled sampling: A small GPT-2-style model generates new smile-shaped points when supplied with smile examples, although that component was absent from training.The controlled setup uses two-dimensional point-cloud sequences and a causal next-token objective.
Controlled sampling: In a learned two-moons sampler, early layers concentrate points, middle layers spread them toward uniformity, and later layers recover the context-specified geometry.Intermediate hidden states are projected back into data space to inspect this evolution.
Pretrained models: Pretrained models show U-shaped MMD profiles: intermediate layers approach a uniform spherical reference before later layers return toward topic-dependent structure.The analysis uses semantic-category prompts and also observes the pattern with natural CBT text.
Energy evidence: The EFS-style interaction energy follows the same middle-layer regularization pattern as the MMD curves.This supports an interacting-particle interpretation in which later layers recover structured topic-dependent geometry.

7 Discussion and Limitations

The discussion identifies two limitations: the U-shaped layerwise pattern weakens in smaller Qwen2.5 models, and the theory establishes expressivity without proving that pretrained models execute the proposed algorithms internally.

Failure mode across model scales: Smaller Qwen2.5 models show only a short or weak movement toward the uniform reference, indicating an incomplete intermediate transport phase.The authors relate this failure mode to insufficient effective depth or capacity, since transformer depth controls available iterative computation steps.
Gap between expressivity and mechanism: The theoretical results show that some transformer parameters can implement closed-form diffusion and EFS, but do not prove pretrained language models execute these algorithms internally.The experiments provide compatible mechanistic evidence, but the observations do not uniquely identify the internal algorithm.
Failure mode across model scales: The Qwen2.5 U-shaped pattern is weaker, shorter, and less robust across model scales than the clearer profiles observed in larger pretrained models.The curves still suggest movement toward a uniform spherical reference followed by movement away from it.
Generative AI statement: The authors used generative AI tools for code implementation, proof development and verification, and language editing, with the work independently reviewed and validated.The authors state that they take full responsibility for the paper's content.

A Related Work

Related work studies transformers as context-based predictors and algorithm executors, while diffusion and particle methods provide generative procedures. This paper connects these threads by asking whether frozen transformers can execute samplers specified by in-context data.

In-context learning: Prior in-context learning research mainly studies prediction from prompt examples, including learned function classes and optimization-like computations.Examples include linear functions, decision trees, shallow neural networks, ridge regression, gradient descent, and algorithm selection.
Transformers as algorithm executors: Transformers can serve as computational resources: attention aggregates prompt information, while depth supports iterative optimization and structured updates.Prior work includes temporal-difference procedures and discrete optimal transport or sorting-type computations.
Diffusion and score-based generative models: Diffusion models generate through iterative reverse processes, whereas this paper isolates their computational structure rather than proposing a new diffusion model or training a score network.The closed-form setting makes the score and sampling dynamics explicit functions of empirical samples and the current state.
In-context generation with diffusion models: Unlike Prompt Diffusion, which trains a diffusion model to perform in-context visual tasks, this work asks whether a frozen transformer can implement diffusion-style and particle-based samplers.The central object is a frozen transformer executing a sampling algorithm specified by the prompt.
Estimation-Free Sampling: Estimation-Free Sampling avoids explicit score estimation through deterministic particle optimization, transporting particles toward a reference distribution and then mapping a reference point back toward data.The paper uses EFS to test whether in-context sampling extends beyond diffusion-style score dynamics.
Positioning: The paper introduces in-context sampling as a common computational view in which prompts specify instances, attention computes interactions, and depth performs iterative updates.This framework connects prediction-oriented in-context learning with closed-form diffusion and Estimation-Free Sampling.

B Proof of Theorem 1

The proof constructs transformer blocks that simulate closed-form diffusion by having attention compute sample-dependent interactions and feedforward layers perform Euler updates. Induction across layers yields both unsmoothed and smoothed sampler recursions.

Prompt encoding and layer schedule: Each token stores data, state, and scratch blocks, with the state token carrying the current sampler state and the scratch block used for intermediate computations.The construction uses one transformer layer per Euler step and resets scratch space after each update.
One closed-form diffusion step: Softmax attention is designed to compute the responsibility-weighted empirical interaction k_t(z) from the state token and data tokens.Query, key, and value projections make the relevant logits and weighted data values available at the state token; terms independent of the data index cancel in softmax.
One closed-form diffusion step: The feedforward and residual updates transform the attention result into z + h v_t(z), implementing one closed-form diffusion Euler step while leaving data tokens unchanged.The block also resets the scratch block.
Closed-form diffusion construction: Induction over layers shows that stacking the blocks produces the closed-form diffusion recursion, with the final state token storing z_L.The invariant preserves the data tokens and the state-token representation at every layer.
Smoothed closed-form diffusion: The smoothed construction replaces k_t(z) with an average of perturbed interactions computed across multiple attention heads.The same induction then yields the smoothed closed-form diffusion recursion, and σ = 0 recovers the unsmoothed method.

C Proof of Theorem 2

The proof encodes EFS particles, reference points, and sampler states into token blocks, then uses attention and feedforward maps to approximate forward and backward interaction fields. Finite-stage stability converts these approximations into an approximation of the EFS output.

Prompt and token representation: The EFS prompt stores initial particles and instance-specific quantities, while transformer weights compile the fixed parameters τ = (K, γ, s, ϵ).Particle, reference, state, scratch, and norm blocks support the forward and backward phases.
Forward and backward phases: During the forward phase, particle tokens store successive EFS iterates, and the reference row is copied into the state token before the backward phase.The final readout extracts the state token's z-block as the approximate EFS output.
Forward interaction construction: Attention heads and tokenwise feedforward maps approximate Gaussian-kernel interaction sums uniformly on a compact domain.The compact-domain assumption bounds logits and permits approximation of the required normalization and kernel factors.
Backward interaction construction: Backward-stage attention applies the same interaction construction to the moving state token, yielding a uniform approximation of the particle interaction field.Multiple kernel heads combine through output coefficients, while a uniform-attention head supplies the affine attractive term.
EFS gradient field: Combining attractive and repulsive components yields a uniform approximation of the EFS gradient field.The feedforward update then advances the backward state through the finite sequence of stages.

D.1 Smile Experiment

The experiments test whether transformers can generate samples from in-context data and whether pretrained models reshape representations through a depth-dependent geometric process. Results include held-out smile generation, layerwise point-cloud evolution, and U-shaped uniformization patterns across semantic prompts and model families.

Synthetic sampling: Held-out smile-shaped samples enable the transformer to generate new IID samples along the same geometry, despite smiles being absent from training.The controlled experiment uses a small GPT-2-style decoder-only transformer trained with causal next-token prediction.
Synthetic layerwise geometry: The model’s hidden-state point clouds are projected back to two dimensions after selected transformer layers to inspect their geometric transformation across depth.The two-moons experiment uses 5,000 points, 1,000 training sequences, and 64 points per input sequence.
Semantic-topic geometry: Layerwise MMD2 measures distance from the normalized hidden-state cloud to a uniform spherical reference, with smaller values indicating greater closeness.The curves report trial means with shaded 95% confidence intervals.
Semantic-topic geometry: Across several large language models, hidden-state geometry follows a U-shaped profile: topic structure weakens toward a uniform sphere in middle layers before returning near the output.The corresponding logarithmic interaction energy shows the same qualitative pattern, supporting the interpretation of intermediate regularization followed by semantic structure recovery.
Model-scale comparison: The uniformization effect varies across model families and scales, with smaller models showing shorter or weaker movement toward the uniform spherical phase.The pattern appears in both IID semantic-topic prompts and CBT story prompts, including GPT-2 and Qwen2.5 families.
Failure mode: Insufficient depth can leave the transport process incomplete before a stable uniform-like phase and later semantic reconstruction form.The paper interprets shallow or smaller models as a failure mode of in-context sampling because depth controls the number of sampler-like computation steps.