From Softmax to Score: Transformers Can Effectively Implement In-Context Denoising Steps
Authors: Paul Rosu
Lawrence Carin
Xiang Cheng
Venue: NeurIPS 2025 poster ()
Keywords: Transformer
in context learning
diffusion denoising
manifold projection
Primary area: deep_learning

Abstract: Transformers have emerged as powerful meta-learners, with growing evidence that they implement learning algorithms within their forward pass. We study this phenomenon in the context of denoising, presenting a unified framework that shows Transformers can implement (a) manifold denoising via Laplacian flows, (b) score-based denoising from diffusion models, and (c) a generalized form of anisotropic diffusion denoising. Our theory establishes exact equivalence between Transformer attention updates and these algorithms. Empirically, we validate these findings on image denoising tasks, showing that even simple Transformers can perform robust denoising both with and without context. These results illustrate the Transformer’s flexibility as a denoising meta-learner. Code available at https://github.com/paulrosu11/Transformers_are_Diffusion_Denoisers.

From Softmax to Score: Transformers Can Effectively Implement In-Context Denoising Steps
Paul Rosu Duke University paul.rosu@duke. edu
Lawrence Carin Duke University Icarin@duke... .du
Xiang Cheng" Duke University xiang. chenngdukeeeeeee edu
Abstract
Transformers have emerged powerful meta-leamers, with growing evidence that they implement learning algorithms within their forward pass. We study this phenomenon in the context of denoising, presenting unified framework that shows Transformers can implement (a) manifold denoising via Laplacian flows, (b) score-based denoising from diffusion models, and (c) a generalized form of anisotropic diffusion denoising. Our theory establishes exact equiva- lence between Transformer attention updates and these algorithms. Empirically, we validate these findings on image denoising tasks, showing that even sim- ple Transformers can perform robust denoising both with and without context. These results illustrate the Transformer's flexibility denoising meta-learner. Code available Inttps:/// h//ggttuu..... httpauu...ooppaaoossu/////// wttpresuuutttaaaaaoooorrrreee CeeiuusiooDDnniieeee
Introduction
Two dominant paradigms in generative modeling diffusion models and Transformer based se- quence models-harl independently achieved remarkable success. Denoising diffusion probabilis- tic models (DDPMs) generate data by reversing stochastic noise process, reaching state-of-the-an results in image synthesis [4 2 Transformers, trained autoregressively, power large language models (LL.Ms) that excel at text generation [15] This convergence prompts natural ques- tion: can attention mechanisms perform diffusion-like denoising, and what does this imply about the nature of generative modeling?
Recent work shows these paradigms are deeply connected. Diffusion Transformers are competitive with U-Nets as backbones of diffusion models [14] 22. while diffusion- language models like Diffusion-LM [22] offer controllable text generation, suggesting attention can drive iterative de- noising in both vision and language. Beyond empirical advances, Transformers have demonstrated remarkable capabilities meta-learning, particularly through in-context learning, where they emu- late learning algorithms during inference without explicit parameter updates [3] Recent studies show that Transformers can approximate gradient-based updates in their forward pass 7 11 effec- tively bridging meta-learning and classical optimization [6]. From theoretical standpoint, Trans- former dynamics have recently been linked to continuous-time probability flows and Wasserstein gradient flows (8). offering principled lens for understanding attention-driven generation. Yet our theoretical understanding of Transformers as denoisers remains limited. How could attention layers collectively implement stepwise denoising process generative flow through nigh-dimensiooal space? While invertible attention has been proposed in flow-based models [19), a general framework for attention diffusion operator is still lacking.
We seek to develop theoretical insight into how Transformers can perform iterative denoising, by analyzing the capability of Transformers to implement various denoising algorithms, and drawing
"Correspondence: roul.rosu@dukeeco edu, xiang. chengdduke.c edu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

connections to their meta-learning behavior in denoising contexts. Understanding attention as generative flow could unify the strengths of both paradigms, and illuminate why Transformers are such powerful generative learners.
Denoising is central generative modeling, yet theory in-context learning (ICL) has largely focused on supervised settings with clean labels. Real contexts are noisy or partially observed, motivating unsupervised ICL. We ask whether Transformers can perform stepwise denoising during inference, and show that attention naturally implements structured denoising updates and, when trained end-to-end, can learn geometry- anisotropic diffusion. This yields a unified view of score-based diffusion and manifold learning as Transformer-ippleennte denoising and connects to recent in-context diffusion generation (e.g.. (20)).
1.1 Outline and Main Contributions
The main contributions of our work are as follows:
1. Transformers Implementing Manifold Denoising. In Section we consider the problem of in-context manifold denoising [00] One presented a set of points sampled from an arbi- trary. context-dependent manifold perturbed by Gaussian noise, and the goal is to denoise these points. We show in Lemma 3 that the Transformer can implement a Laplacian- -asud manifold- denoising algorithm that iteratively applies a kernel-weighted update [10]. This result provides clear theoretical link between manifold denoising and the attention mechanism of Transformers, highlighting their potential in learning implicit manifold structures in-context.
2. Transformers as Exact Score-Based Diffusion Denoisers. Section 4 examines the role of
Transformers in the context of score-based diffusion generative models 17 We show rigor- ously that the Transformer architecture can implement the exact score-based denoising algo- rithm through suitable prossaattention construction (Lemma 4 Section 4.2 we further ex- plore the Transformer's 's performance in an in-context denoising setup, demonstrating empirically that it effectively generalizes from context-free context-dependent denoising scenarios, thus unifying these two important perspectives through the same attention-based mechanism.
3. Efficient Approximate Score-Based Denoising via Learnable Witness Tokens. Recognizing the impracticality of exact score-based denoising. Section 4.3 introduces approach utilizing small set of learnable tokens ("witnesses") that approximate the exact score computation. This significantly improves both accuracy and computational efficiency, demonstrated in our empir- ical results in Section The use of learnable witnesses bears a connection to several established methods in kernel approximation.
4. Generalization to Anisotropic Diffusion and Learned Attention Parameters. In Section we generalize beyond isotropic kernels and standard diffusion models by studying anisotropic diffusion processes. We theoretically prove (Lemma 5 that Transformers, when their Query, Key, and Value parameters are suitably aligned with the diffusion coefficient matrix, can exactly implement the reverse ODE of anisotropic diffusion. This theoretical insight is substantiated em- pirically in Section 4.5, demonstrating that the Transformer can learn more efficient, geometry- adaptive denoising algorithms beyond the standard score-based denoising framework. We visualize this process in Figure
Collectively, these contributions elucidate how attention and cross-attention can naturally realize a broad spectrum of denoising algorithms. Our results reveal the remarkable flexibility and effective- ness of Transformers for both context-dependent and context-independdnt denoising.
2 Transformer and Kernel Weighted Update
Given n tokens 2(1) _())....,,20 (n) E Rdd, we define the matrix Zo = [,(,)...,, 2()) E Rddx, For L-llaee Transformer, we let WV, W W.K Rdddd denote the value, query and key parameter matrices of layer é, and let W.S E gdxd parameterize linear module. For convenience, let We [WY W WK, W and let W (We)t Attn denotes the standard attention:
Attnaad (2: Wz) = W. Z smax (Z WTK WºZ)).
(1)

---

Clean
Noiov
LG
Clean
0
O
O
4
4
4
7
7 7
between Successive layers
Figure I: Visualization of diffusion steps a denoising algorithm implemented by truined Trans- former with diagonal W W W matrices. The Transformer is modelled by the casee-based anisotropic diffusion reversal algorithm in Lemma [|] Section 4.5 The first columns show the clean and noisy images. Each subsequent column shows the output given layer of 6-layer Transformer. We visualize the difference between consecutive layers' output at the bottom. Notice that each Transformer layer denoises roughly non-overlapping patch.
where smax applies column-wise softmax. We use diagonal mask that each token does not attend to itself. This standard when disallowing self-attention (diagonal mask) and also common in manifold denoising, where the update at each position is informed by its neighbors rather than itself (see [IO]. Section 3.1). For notational simplicity, we omit the explicit mask matrix. Let
rbf(.,) _dxn x -d Rxx" be the matrix valued RBF kernel, defined [rb V)l.
exp (-1/2 2 ( v Let [ogrbf (U, V)l == [log(rb( (U,V))li We define Attnyrt as
Att (2; W/) = w Z (m 2)).
(2)
We highlight that difference between the standard Attnaad and Attnrrr as follows:
|smax Z W. w&Z & exp (WW (O) (WW
[smax (logrbf (W. z, W si) x exp (WW 20) ( (W -0)) 10/2 -I 10
Given Attn E (Attngga, Attnrbe) }, we define the Transformer via the iterative update:
Ze+1 W Attn (e:W W w²).
In many subsequent applications of interest, ||W 2 11 and ||W 200 1/ are constant, or close to constant, due to concentration of norms of high-dimensional Gaussian vectors. this case, we verify that Attnata and Attnye behave identically:
Lemma (Equivalence of Attnead and Attnyee on the sphere (with W w identity-scaled). scaled)). If
(4) 12 for some constant C, and if w and w are scalings of the identity matrix, then ||
At Attmgra(Ze, W) = Attnrag(Ze Zz, W).
We defer proof of Lemma to Appendix A We also empirically verify the similar performance of AttData and Attn,rr in Figure
2.1 Transformer Construction for Weighted Update.
Given tokens (2 )= C Rdd let F {2(1).. ...),........ {m)) denote the subset of frozen tokens.
Much of our discussion in this paper revolves around the following kernel-weightee-uuddee scheme:
() 241 (1-a0)22
ve Σ ()) =²))
for
jØF. (4)
itj
1 1 j} exp(-112/) (0113/2077)
ke(i,)) =
(5)
Exs, exp(epppppp -112 (1)
-00)
13/(0))))
3

---

and _0) = for j E F. Motivated by 4 we implement this update using single attention module together with fixed mask that freezes the columns in F.
Let Ze = 1zp [m)] E Rddn, and let Mf E R denote the diagonal freeze mask with We define the Transformer with frozen-token mask by
(Mr)aj (
(wf Ze + Attnrer (Z W (...W) Wx ) + Ze(Luxn - MF))
Ze+1
(6)
We verify below that (6 can exactly implement the werghl-weighted update 44 (proof deferred to Appendix[A] A}:
Lemma 2 (Transformer implements kernel-weighted update.). By choosing W 1/0elaxds W.K 1/0eLaxd, W. Telaxd: W (1 ae ldxd the masked Trans- former in (6) exactly implements the kernel-weighted update in (4)
When a subset of tokens is permanently frozen (e.g.. a large train- ing set), it is convenient store them separately and update only the remaining tokens. Let zF (2(1, 2(2) .m)) E Rdxm denote the matrix whose columns the frozen tokens. Let Z = (2(m+1), 2(m+2), ,...,. [")] denote the matrix of not-frozen tokens. Let cross attention modules CrAttnata (U, We) == W. smax (UTW? w z) and CrAttn (U, Z; We) ==
Cross-attention form.
WY U smax (W&U, U, W Z)) be defined analogously to Attneta and Attnebs- We consider the Transformer with cross attention:
Ze+1 = W Ze + Attnybt
Zg; WV W w + CCrAttnorr (z² WW W W (7)
We will use (Z, F-; W) to denote the value of Ze, when evolved according to 77. initialized at Zo ZZ and "²; i.e. the output of the é-layer Transformer in ( In the single-query setting used for score-based denoising (Section 4 coincides with (6) (the self-attention term vanishes), so we will use the cross-attention notation for convenience.
Manifold Denoising via Laplacian-based ODE
In this section, we are motivated by the following in-context manifold denoising problem, originally proposed in [0] Let M be a m-dimensional manifold embedded via an embedding T .M- -> Rd. We identify a point a E .M with its d-dimensional embedding T (a).. Let x(1) ..... ((1) be sampled from M, and we observe 2( = (x) + c(f), where e(i) ~ N(00, o² (aIx))). The goal is to recover the (f.)s. This problem is particularly challenging because (1) the manifold M is unobserved. and (2) the manifold M context-dependes. Therefore, solve this problem in-context, the Transformer needs be able learn (implicitly or explicitly) the manifold structure, from the observations (
[00 propose the following manifold-denoising algorithm based on the Laplace-Beltram operator: Let the adjacency matrix W E g='= be defined as Wiy := exp( (1 - -2-0) 201) (20²)) and set W. 0 to remove self- Let D be the diagonal matrix with D;; Σド= Wyy. The RBF
Laplacian is defined as = Inxn D-1W. The denoising algorithm based on a Laplacian- flow (,2(0 (1) = ? Σラ= Lig 205 (t), whose Fuler discretization gives the discrete-time algorithm:
20 ( (0)
( 2841 24
Se
Σ Cij el))
(8)
In Lemma 3] we verify that K is an instance of the Bernel-weightee update (4) (with F= therefore implemented by the atteetitttnntion Transformer construction in Lemma
2), and
Lemma 3 (Reformulation of Manifold Denoising). The manifold denoising algorithm in ex- actly equivalent to (4), with F 2, ag = O, a; a, and 7 òg- Consequently, the algorithm implemented by the Transformer construction in Lemma
The proof of Lemma 3 is just by algebra; we defer it to Appendix B

---

3.1 Experiments for Manifold Denoising
We present experimental evidence demonstrating the Transformer construction in Lemma can indeed act as an effective manifold denoiser. The problem setup is very similar to the above: Let
x(1), ...... x(n) be clean images sampled from training set. We perturb each image via 2(1)
a00)) + e(), where e(i) ~ (O, I²I), for some noise level a. The input matrix Zu= = [x(1)..... 2 [n)
and the Transformer (Z, We) are defined in Section Motivated by Lemma 2 we impose the constraint that WY, W W W
are each (trainable) scalar multiple of identity. The training loss is given by Ez. [ Ei= (Zu; W)] - x Throughout, we report noise levels via
a signal-to-noise ratio (SNR) defined as SNR := moise ddams i.e., the standard deviation of the additive Gaussian noise divided by that of the clean data distribution.
Figure 2a shows that the test loss decreases with the context length; this agrees with the fact that the discrete RBF Laplacian L becomes a better approximation to the Laplace-Beltrami operator with increasing number of samples. Figure 2b shows the test loss of the intermediate layer outputs of 6-layer Transformer; the test loss decreases with number of layers, which is consistent with (8. as each layer implements one more discretization step of
Loss Vs Context Size
Loss of Output of Intermediate Layers
3.57
moon A.1 ne]
moan n=]
3.300
5.225
Test Locs
2.500
Test
0.00
6.0
3225
5.5
100
200
2010
400 500
Context Size
Layer Index
(a) Loss Against Context Length
(b) Loss Against Layer f. Context size =100.
Figure 2: (a) Test-set RMSE on MNIST versus context size n for the Laplacian based denoising Transformer under noise level of SNR 1. Each point shows the converged loss of separately trained 6-layer model; (b): Test RMSE of intermediate output at layer l with fixed n = 100.
4 Score-Based Diffusion Denoising
As in standard based diffusion models [1 18 17), our training objective is defined over the empirical distribution á (equivalently denoted po below), the uniform distribution over training sam- ples (z(()) } We assume these samples are drawn i.i.d. from an unknown population distribution p. While optimization is with respect to P. the modeling goal is generalize to p. Unless otherwise stated, all evaluations reported in this section are computed on unseen test examples.
Score-based diffusion generative models 18 motivate the following problem: Let Po denote the uniform distribution over a (large) training set of samples (1) (n)) }. Let Po * (O, tf),
where * denotes convolution. As we review in Section 4.1 getting access to V log (xx) allows one to turn (almost pure) Gaussian noise x into sample from po by reversing diffusion process. In practice, one trains a neural network (0 l). parameterized by weights 0, to approximate the
(scaled) score, by minimizing the denoising objective E
[|so t) - tV log (x))
Conveniently, V log Pr (x) is exactly equal to a brighted-averrgee between x and each (a ) Consequently, we verify in Lemma that the Transformer (7) can implement score-based denoising.
4.1 Standard Score-based Diffusion Denoising
Consider the variance-exploding SDE, defined as
Zo ~ Po:
dr. = d.B.,
(9)
5

---

where B, denotes the standard Brownian motion. Let denote the distribution of 24. We verify that Pt = Po (0, By the Fokker Planck equation, pt evolves as
Ur (x) == 1 (o)
कe pi(s) - div(p
(10)
Consequently, if 20 ~ Po and dz; - logpr(z))dt, then 2ς has the same distribution as dr in (9) The above follows directly from integration-by-parts. The diffusion process of (9) can be
reversed by reversing the velocity field (2)). Concretely, for any fixed time T, let yo pr, and dys- = log o then us= pr-s. The Euler discretization of the ODE for ys is given by
Ue+1 Ie + de
2/2 V Voggoo
0+ = o dr,
(11)
where of is the time T - at step is (equivalently, the Gaussian variance of ÿt).
From the definition of pa and pt, it follows that
E ee =|==~<(000 11
(2(00 (z(0) -x)
(x(x))
1 nZ
Σ
+ #1|a-a001
V Voogp((1)
(12)
E 1-1-0 1
where Ze is the normalization constant for t1). Observe that the RHS of 12 is exactly the keight-weighted average from 4 Combining (12) and (TT) then implies that the score-based denoising algorithm (11 is equivalent to (4) which we formalize in the following lemma:
Lemma 4 (Score-based Denoising as Kernel-Weighted Update). The score-based denoising algo- and "e de/ (20
rithm 11 is equivalent to 4 with z =(==1) Dt. F -1,...,m,, ar
Consequently, the Transformer construction in Lemma 2 can implement score-based denoising 11
We defer the proof to Appendix We highlight an important difference from the manifold denoising setup in Section here, the set of training samples is typically large, and does not vary with context. Instead, the training samples encode global context-free knowledge. On the other hand, the presence of similar tokens in-context can indeed improve the denoising accuracy; show this empirically in Figure 5 in the next section. Thus the same attention architecture can be well-suited for both in-context, and context-free learning.
In Section 4.3 we discuss a generalization of the construction of Lemma where we train the cross- attention tokens in 7 to learn an approximate score that significantly improves generalization error and computational cost. In Section 4.4 we discuss a generalization of the score-based denoising sequence (1T) itself, which also leads to significant accuracy improvements.
4.2 Experiments for Score-Based Denoising
We empirically validate Lemma 4 Our experiment verifies that the Transformer construction of Lemma 4 does indeed implement exact score-based denoising. The setup is as follows: for each input context, we sample zo!1) .((),,,(")) clean images from training set. We let 2(() x(1) The
query image is (x++1), We perturb the query image by Gaussian noise: 2n++1) = ((++1) + c(++))
where c(++1) ~ ~²),,²T)) where is chosen noise level. The Transformer [F We) is as
defined in (7) Motivated by Lemma 2 impose the constraint that WY. W.O, W.K W
each a (trainable) scalar multiple of identity. The training loss is given by the reconstruction error
Ez (Z W)](++) (n++1))) [[]].
Single-query score-based denoising. Figure shows two architectures: {RBBF" "STD") refer to Transformers that use Attn (2) and Attneed. (I) respectively. For each architecture, "trained" means the Transformer is trained on the above loss; "theory" means the Transformer parameters fixed at the construction in Lemma 4 (for a geometrically decreasing σы schedule). Figure plots the test loss of the per-layer oulput of 6-layer Transformers. We observe the following:
1. The reconstruction error decreases monotonically with layers, suggesting that the Transformer is implementing an iterative denoising algorithm at training convergence. "RBF-trained" appears to denoise more aggressively in the first few layers compared to the theory constructions.

---

2. RBF-theory and STD-theory have almost identical losses, supporting our claim in Lemma the RBF and STD attentions behave similarly under our constructions.
that
In Figure we visualize the clean and noisy query images, as well as the per-layer intermediate
outputs of the RBF-trained and RBF-theory Transformers from Figure 3 In Figure 4a| and 4b) we see that the trained Transformer can exactly recover the clean training set image, as does the exact score-denoising algorithm. In Figure 4c] and [d] we see that test images are not being recovered; instead the denoising algorithm produces similar image from the training set. Though undesired, this is the expected behavior for exact -eorsdden. In Figure 44 and 4f we see the "generation" capabilities of each Transformer from pure noise; again, the "generated" image is in fact sample from the training set (as expected of exact score denoising). The generated outputs of RBF-trained and RBe-th are almost identical in all cases.
MNIST Denoising Progression Layer (L=6)
CIFAR-10 Denoising Progression per Layer (L-6)
16
saT theary
10
14
FDT tou nod incon ztn, nwd)
bntkod 222 ne): mained (neas .15 31
1.2
20
Layer index E
Layer index E
(a) MNIST
(b) CIFARIO
Figure 3: (a) Test loss of intermediate output at layer E versus é on MNIST, (b) CIFAR-10, at an injected noise level corresponding to SNR ss 3 for (a) and for (b). All Transformers have 6- layers total. For CIFAR-100 results and FID on CIFAR-10/100, see App. E.2.
In-context score-based denoising. In Figures and 6 we present different experiment which highlights the in-context aspect of score-based denoising. In this setup, the input consists of the noiseless training samples {x(11, ,..,),,...,.,... x((),, }, as well as m contextual images (o(=+t)...... (n+m)) }. The query image is 1)*+m+1) Crucially, all m contextual images belong to the same class as the query image g(n+m++1), The Transformer input is 20)) = ,.(() for = =1,..-. I..... .n + m, and the query is perturbed with Gaussian noise: 2(++++1) x (n+ +N I² We call m the "context length". The query ((m++++1) attends to the training set (²) ...., (*) via the CrAttn module in 7.. while the context samples query (((+1), ...... .z. (n attend to each other via the Attn module in (with appropriate masking). Experiment details are in Appendix
F.4 E.4
The training loss is again the reconstruction loss. In Figure 5 we plot the test reconstruction loss against the context length. The presence of a few contextual examples leads to a significant im- provement in reconstruction error. This demonstrates that the Transformer is indeed implementing a denoising algorithm that can adapt to its context (at least partially). Figure 6 visualizes this phe- nomenon. The injected noise has high SNR=5, so that the original image is difficult to recover. In the leftmost block, we see that providing same-class contextual images guided the model to correctly generate "6" (even though the intermediate output appears tend towards "5"). In the middle block, the contextual set consisted of images from random classes, and the Transformer had difficulty re- covering the image. In the right block, Transformer trained without context generated different digit from the same starting noisy image.
4.3 Approximate Score-Based Denoising With Learnable Witnesses
In the preceding section, we show that exact score-based denoising can be implemented via cross- attention between the noisy sample, and every other sample in the training dataset. This is in gen- erally practically infeasible. On the flip-side, the training set is rypically not context-dependent Therefore, one can hope to approximate the exact score cross-attending to small set of represen- tative samples. Concretely, let V (w(11) ,.... (())) C. RRd denote set of witnesses g(t), for some
7

---

3
3
3
3
3
3
Taned
3
3
3
3
3
3
Thooy
Theor
(a) MNIST (query from train set)
(b) CIFARIO (query from train set)
Ove
L
2
?
2
2
2
2
2
Thaiwed
Thatned
2
2
2
2
2
2
Theoyy
c MNIST (query from test set)
(d) CIFARIO 10 (query from test set)
L4
Hobbe
12
7
7
7
Tained
Trained
theory
7
7
7
7
Theary
(e) MNIST (random noise query)
(I) CIFARIO (random noise query)
Figure 4: Visualization of denoised images. Query image from {train, test, random noise). In each sub-figure, the 1st row shows the trained Transformer, and the 2nd row shows the Transformer with theoretical weights. First two columns in each figure shows the {clean, noisy) query image respectively. Columns 2-7 show the per-layer intermediate outputs of 6-layer Transformer.
12.0
Test RMSE
Test Loss
0.5
20 Context Length
Context Longth
(a) MNIST: Test Loss Context Length
(b) CIFAR-10: Test Loss Context Length
Figure 5: Test-set RMSE for a 6-layer RBF-attention denoiser function of the number of clean, same-class context images 7a.x- The query token is severely noised (SNR 15 for MNIST, SNR= 10 for CIFAR-10), while the full training set is available through cross-attention. Adding just a few class-matched examples dramatically lowers the loss from the nearly unrecoverable zero-context case. For CIFAR-100 results and FID on CIFAR-10/100, see App. F.3
<< n. The goal is to choose a which well-approximates the kernel- update from
e 10 (2)))
=210
-1
(2(8) -2)
(13)
=12/
n kml e -11))
f=1
Observe that the 2 is equivalent the k-averaged term in In the setting of score-based de- noising, z(1)..... ,..., (n) is drawn from po independently. An intuitively simple choice of V/ is then to
sub-sample 7 points from (2(1) ...., (()]; under this choice, the LHS of (L3) is simply Monte-
Carlo estimate of the RHS of 13). Assume that Po is compactly supported, we can show via basic
algebra that E 10 2 O(- We provide short proof in Lemma 7 in Appendix C
Much more sophisticated approaches exist, and the problem of picking good witness set V has deep connections to number of areas such as Nystrom's method for kernel approximation kernel mean embeddings
[3] and kernel herding [5].
8

---

Clean Yoisy
Danasigg
6
6
2 2 2
nest Sanples
6
6
6
6
2 8
6
9
7
5
5
4
Figure 6: Two leftmost images are the clean and noisy (SNR & 15) query images respectively. We show three sets of six intermediate layer outputs. Left: a Transformer trained using with 100- image contextual set is presented contextual images from the correct class at test time, and generates the correct image. Middle: the same Transformer, when provided contextual images from random classes, fails to generate the right image. Right: a different Transformer, trained without context, generates the digit from the same input. All three scenarios start from the identical noisy input.
Motivated by the above, we treat V learnable parameter, alongside the other Transformer parameters. Lemma|4shone showed that the score-based denoising algorithm (IT) can be implemented by the Transformer defined in (6) (equivalently (7). The witness- basedd Transformer is then defined
Ze+1 W Ze Attnrur
(Ze WY, W w() + CrAttnrrr (Uer Ze: WY W wW*) )
(14)
where Up E Rdxr is the only difference from The full set of parameters of 14 is (Ur, wY, W.O W. Je=t....L In words, column of Ug represents wi(0)) in the above discussion, and Ue can differ from layer layer (analogous & changing with each denoising step.) We ini-
tialize U2 with randomly sampled vectors from (=(())=i..... We show in Figure 7 that [14 has
significant advantages in computation and generalization error (see Section 4.5 for discussion).
4.4 Anisotropic Diffusion and its Reverse ODE
Consider a generalization of the standard diffusion process in (9) instead of dx, dBr, let A. : R+ + Rdxd denote a time-parameterized family of matrices. For simplicity, we assume that A, is symmetric. The anisotropic diffusion is described by
Fo Po: dx, A,d.,
(15)
Throughout Section 4.4 we restrict diffusion schedules satisfying the normalization f A ds TT. Let Pt denote the distribution of T4- Let T be some final time, and let pr = Pn (0, TT). In Appendix C.I we show that (15) admits simple continuous-time reverse ODE ye ~ pr-t- Its Fuler discretization gives denoising algorithm that be implemented by Transformer:
Lemma 5 (Anisotropic Denoising with Transformer). Let (:(1) .i..,,,,
(n) denote a training set. Let 200) x( for L..... .. n. Let (m+1) (=+ + (O, TI) denote the query token. Consider the anisotropic diffusion in (15). where A, is a family of arbitrary symmetric matrices satisfying
TT Adds TI. Let to T, 1= = de, Mo TI, M M, deA? where we identify As with At,- Then there exists a discrete-time denoising algorithum, based on the time-reversal ODE of the anisotropic diffusion in [55) This denoising algorithm can be implemented by the Transformer in 7 with F {I,...,,., n. and parameters
SST
W I 8,A,M, 12, WY = 6.A /2, WP = W M 1/2
(16)
The proof of Lemma uses the same tools from Section 4.1 and we defer Appendix C The importance of Lemma lies in generalizing beyond the isotropic RBF kernel used in based denoising, and justifies the choice of non-identity W W W W matrices, which enables the Transformer to implement an anisotropic denoising algorithm that adapts to the distribution at step& In Section 4.5 we evaluate a witness-variant of the Transformer from Lemma
4.5 Experiments for Witness and Anisotropic Denoising
In the following, we present set of experiments that serves two purposes: First, it demon- strates the advantage of the Transformer with learnable witnesses from Section 4.3 Second, studies the anisotropic diffusion Transformer construction from Lemma 5 with general learnable

---

We, W WY matrices. The setup is identical to the single-query score- basee denoising experi- ment in in Section The one difference the architecture: let the Transformer (= Ue, We) be as defined in (14), where Uj denotes the set of learnable witnesses and We denotes the set of pa- rameters layer e. The training loss given by Ezz [ U, W) - ==|=]-
We initialize the witnesses (U) by sampling a random subset of training examples to set the initial columns of (U). During training, we treat (U) as a continuous, fully learnable parameter matrix and optimize it jointl with (W) via gradient descent. Consequently, at convergence the columns of (U) need not coincide with any training example (e.g.. column of (U) may represent an average or a denoising- end-to-end and connects to classical kernel approximation methods (e.g., Nystrom and related con-
summary of multiple samples). This makes witness selection differentiable
structions) that summarize large kernels with a small set of inducing points 21,9 9 [3]
In Figure 7 we compare the test loss of three models: the Baseline model is the Transformer from Lemma 4 that exactly implements the ssore-hased denoising algorithm in (TT). The "Witness+RBI" model implements the witness- Transsss-Trn from (14), and its weights W are (trainable) scalar mul- tiples of identity. The "Witness+anisotrroicc" model is motivated by Lemma 5 we do not assign the
Transformer parameters as done in (16); instead, we simply make {WS. W. W&. W&} trainable diagonal matrices. We provide details of the implementation in Appendix E.4
With a moderate number of witnesses, both titmess-basele Transformers significantly outperform the baseline Transformer that implements exact score matching. Furthermore, the anisotropic denoising Transformer shows a significant advantage over the RBF (isotropic) denoising Transformer.
In Figure 1 we visualize the per-layer denoising progress of a 6-layer Witness+anisotropic Trans- former (1000 witnesses). In contrast to the standard denoising sequence in Figure 4 where noise is removed uniformly in space. the denoising sequence in Figure [I proceeds patch-wise. Such a patch-wise denoising procedure coincides with the anisotropic denoising algorithm when Ae for each step É is sparse diagonal matrix, whose non-zero entries coincide with the pixel locations of a localized image patch.
Across datasets and witness counts, the witness-based Transformer achieves lower test reconstruc- tion error than the Baseline model that implements exact score-based denoising. Although the Base- line can perfectly fit the empirical train distribution (ÿ), its unconstrained attention tends to general- ize poorly. contrast, the Witness model's learnable summary witnesses provide an inductive bias that improves generalization to the population distribution (p).
MNIST
Beseline (Eaatt Score)
85
Witness-
A
3.0
Teec Frec (R305))
Withess Arisotropic (moan 100 n-
4.4
Witness RBF imean #1o, 3-3)
0.0
7,5
20
Number Witeeses Blg))
Number Vitsesses (loq)
(a) MNIST
(b) CIFARIO
Figure 7: Test loss at SNR RS 3 for (a) and RJ 5 for (b) with a 6-layer model as function of the number of trainable witness tokens т. The horizontal line shows the performance of the exact-score denoiser that uses the full training set, while the two curves plot witness-based approximations (14): one based on anisotropic diffusion (Section[4.4), and one based on standard RBF kernel. All reported metrics are computed on unseen test data to assess generalization. For CIFAR-100 results and FIID on CIFAR-10/100, see App. E.3.T
All experiments are done on a single A5000 GPU, with each experiment taking at most few hours.
10

---

References
[1] E. Akyürek, D. Schuurmans, and J. Andreas. Transformers learn in-context by gradient de- scent. In Proceedings of the 40th International Conference Machine Learning. PMLR, 2023.
[2] J. Bao, C. Dong, D. Zhang, L. Zhang, J. Zhang, H. Liu, Y. Yang, I. Yuan, D. Chen, and F. Wen. All are worth words: A vit backbone for diffusion models. arXiv preprint 222.iv:2009121122 2022.
[3] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, et al. Language models few-shot learners. arXiv preprini 2005:20005.4165,, 2020.
[4] S. C. Y. Chan, A. Santoro, A. K. Lampinen, J. X. Wang, A. Singh, P. H. Richemond, J. Mc- Clelland, and F. Hill. Data distributional properties drive emergent in-context learning in trans- formers. arXiv preprint arXiv:2205.0m 2200::2005.0005,, 2022.
[5] Y. Chen, M. Welling, and A. Smola. super-samples from kernel herding. arXiv preprint arXiv:1203.388
2012.
[6] D. Dai, Y. Sun, L. Dong, Y. Hao, S. Ma, Z. Sui, and F. Wei. Why can gpt learn in-context? language models implicitly perform gradient descent meta-optimizers, arXiv preprint arXiv:2212. 10559, 2022.
[7] S. Garg. D. Tsipras, P. Liang and G. Valiant. What can transformers learn in-context? a case study of simple function classes. arXiv preprint arXiv:2208. 01066, 2022.
[8] B. Geshkovski, C. Letrouit, Y. Polyanskiy, and P. Rigollet. A mathematical perspective transformers. arXiv preprint arXiv:2312.10794, 2023.
[9] A. Gittens and M. W. Mahoney. Revisiting the nyström method for improved large-scale machine learning. The Journal of Machine Learning Research, 17(1)397774444, 2016.
[10] M. Hein and M. Maier. Manifold denoising. Advances in neural information processing systems, 19, 2006.
[11] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840- 6851, 2020.
[12] X. L. Li, J. Thickstun, I. Gulrajani, Liang, and T. B. Hashimoto. Diffusion-Im improves controllable text generation. arXiv preprint arXiv:2205.. 14217, 2022.
[13] K. Muandet, K. Fukumizu, B. Sriperumbudur, B. Schölkopf, et al. Kernel mean embedding of distributions: A review and beyond. Foundations and Trends@ in Machine Learning, 10(1- 2):1-141,2011
[14] W. Peebles and S. Xic. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
[15] A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[16] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[17] J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. arXiv preprint
arXiv:2010002502, 2020.
[18] Y. Song, J. Sohl-Dickstein, D. P. Kingma, Kumar, S. Ermon, and B. Poole. Score-based gen- erative modeling through stochastic differential equations. arXiv preprint arXiv:201 13456, 2020.
[19] R. Sukthanker, Q. Huang, H. Wang, 7. S. Liu, and P. X. Wang. Generative flows with invertible attentions. arXiv preprint wxX.:2200.00399, 03959, 2021.

---

[20] Z. Tang, Z. Yang, M. Khademi, Y. Liu, C. Zhu, and M. Bansal. Codi-2: In-context inter- leaved and interactive any-to-any generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 77425-27434, June 2024.
[21] C. Williams and M. Seeger. Using the Nyström method to speed up kernel machines. Advances in neural information processing systems, 13, 2000.
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have a subsection dedicated discussing limitations. See Section G
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We have several lemmas, and all details and assumptions connected to them are discussed. All proofs are provided and to the best our knowledge are correct.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data provided not)?
Answer: [Ys]]
Justification: All details of our experiments are provided. All code will be released upon publication.
5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes] X
Justification: All data we have considered here is publicly available, and the Appendix we provide details for reproduction of synthetic data.
6. Experimental setting/details
Question: Does paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes] X
Justification: All experimental details are provided and code will be released upon publi- cation.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer:
[Yes]
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [es]]
Justification: We explain what type of GPU was used for our computations, and needed compute time above the Limitations Section.
9. Code of ethics
13

---

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://pwwpppppaap........pp/aaa
Answer: [Yes]
Justification: We have read the code of conduct, and have complied with it.
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA].
Justification: We discuss some broader impact of the work in introduction. The work is mainly foundational and has no direct social impact.
Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: There are no safety issues with the data that we have considered.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: We are using no code or data that would require this.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [Ye]]
Justification: The main new asset is the related code we have developed. It will be freely released and is well documented with its license.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given participants and screenshots, applicable, as well as details about compensation (if any)?
[NA]
Justification: No experiments of this type were performed.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: None of our work requires an IRB.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: This work does not involve any non-standard usages of LLMss
14

---

A Theory for Transformer Construction
Proof of Lemmal
? Recall from Section 2 the following expression for smax:
2 W T W 2) x exp (W) (W (W(220))
|max
(W Z, w 2) exp (W (0) (w, -0) 1//222 1 11
In particular, the x symbol is applied for fixed column j. across rows i.
Under our assumptions that W - allxa and W = blaxd. we can simplify
(W/* -01 abz(1) T 200
IW 13/2 d2l1200) d²C² b" 113/2 = 6²c²
Carefully writing the expression for the smax term in Attnrb gives
exp
Enety (W& z(1)) - (00) 18/2 W 200 1B3/
exp (bb(()) 200 (a² +2)
Σ Enetyy exp (ab(()) (a² + b2)02)
exp (ab2(1) z66)
Σnyj exp (bb(()) 200
-[smax(2 Z W W 2)
This concludes our proof.
Proof of Lemma 2 Recall (6). where (M=);s = 1(i Ifj F, then (M=),s = 0 and the
update gives 2 z( matching the frozen-token rule
Now fix j & F and let 2/4 denote the jth column of Zz. Since (MF),s 1, 66 gives
. [w 2z] (f) + |tttmer (Z W. Wo w ) (1) 2+1
With W = WK = 1 Taxd, the RBF similarity between key and query j is exp( (xx(-11z]" (6)
||
oe
(00 (1/(2/))) and the column-wise softmax therefore equals ke(i,) in (5). Thus, using WV 28 7eLaxx.
Attnybe (Ze: W (.?.W.) (1)
r Σ (2
ig..
Using W = (1 ae - -76) 710 Taxds we obtain
-6) 2141
= - ar - -76)=()) + Σ
Ke ()(
(1 - -0))))=()) 76 Σ Ke((.)) (z) =()))
which is exactly 4
15

---

B Theory for Manifold Denoising
Proof of Lemma 3 Recall L = Lnn D, where Wy == exp(-11z(0) - 21011//220 and
W₂ 0, and Di プド=1 Way. Then for each
Lij 2t
=26
Σ
Wisze
Dis
1=1
s
Plugging this into (8) gives
2141 =/ -
de (e(
2t
-
-5))a()
+ 0:
Σ
Wy (1)
Dí
Di "
jæi
jf1
Since Ejji Wtj/Di; 1, this can be rewritten as
2p+1
28
+ 88
Wij
Dís
(3) (8)
Noting that Wty/D ke(j. with 0g = o, we obtain
with JF 23, ag 0, and 7/ 8z.
C Theory for Score-Based Diffusion (Section
Proof of Lemma f Let (11), _'.)..... (n) denote the training set of unperturbed samples. Let (x+1)
denote the query sample. Let =0 = ((()) for i = I,..., n1. Let =0 (n+1) = ((++1) + NNO0, og I), where og == T the initial noise level. From
2t+1
(8)
(i)
for i 1 ....
E+1
(n+1) (n+1)
81
v
log log
(0 ((ut1))
Poz
Σ Ei=1
(n+1)
-600
2(03
(1()) - 2(
exp
20
(n+1)
=2
+
20
1
(n+1)
x11) 1
Li=1 exp
26 20
(x+1)
x0(+)
1 x((1)
-(1
28
(n+1)
Li=1 exp
2
28
(17)
20²
20
(n+1) x1)) 1?2)
Σ
exp
207
22
Where the second equality is by 11 and the third equality is by the expression the score in 12
Comparing 17 to
see that it is equivalent to letting F ()..)..... (...,,(()),, and og dr. and Consequently, let
6e
T/ 207
Zo [z(a+1)] E Rdx1
2 := ((.,)...., Rdxn
From
Ze+1 W CrAttn (z" Z2,2e;W WY W w
(n+1)
where we leave out the Attnya because Ze contains the single token query 27 assumed in Section 2 that Attn does not have tokens attending to themselves.
and
By the lemma statement, choose
8₁ laxds
1
8,
WV = 20
W = w.k = Laxds
S
1
20
16

---

Recall from Section 2.1
that Cr (U, := W (WW W&Z)). 2)). Here,
U z {[200) 2 } and Z Zz. Thus
Ze + CrAttn Ze we', w.ki)
Z++1
20
dg 7t smax 1ogrbb | 7/ )
Zz +
20
2
Under the definition of CrÅttn, z and Ze = |=g 1 we verify that above identical to
Lemma 6 (Continuity Equation). Let Pt Po =N t1). Its continuity equation is satisfied by
12 V (t)), as defined in 10
Proof of Lemma 6 Because the forward SDE is standard Brownian motion, its density satisfies the
heat equation" d.p. 1/2 Ape- Define the velocity field (x) - Vlog (xx). Then, using the identity Pt Vlog = Vpt:
(p.(u))) -V- Pc (-vo-o2ooo))) V. Ap. 0.pt.
which is exactly 10)
Lemma 7. Consider the setting in Section 4.3
Assume that || S R for all i. Then for any
||*|| S R. we have
27
11 5) =018/(200)
e (|((0) -
(5100
(2 -2)
-(
[10((h) -0100-=11///200
11:00)=
k=1
2
Proof. For convenience,
-
(366)>=11//100 -011/(2007)
(1660) -
Ek=1 k= -
j=
-|
-110=/ (200)
1=1 Σk=1 k= -| =113/(202)) (26)-=)
Let Ar = & E 11(00)) - - (200) (w:(0) =) and B, 1 - e- (لأااا) =113/(200?) Analo- gously let A, == der our assumptions, there exists constants C, o, a' such that =|000 (z0) 2) s C, and - -13=00113/000' a'
e - - (20) - and B, = 1 Σn k= - -) -210/ Un-
We can decompose
A-
An
A,
An
1
Ar
(18)
B
B,
B
B,
B
The first term can be bound as
E
Ar
Anl
C
E
|A
Anl13]
<
O
Bn
B,
TB
For unit-variane Brownian motion the infinitesimal generator J.A. hence Oep: Ap..
17

---

The second term can be bound using a Taylor expansion of f(r) = 1/r:
+ + +c)= - 1
c
+
dtds
,2
28
()++c)-J0
+
e
,2
min
+c3
Thus we can bound the second term of (88 as
1
1
1 1
E
A
SC
B+
Bn
Br
Bn
C2
E
(B, B,)
+
E
(B, - (B,)
Bn
min
Bn, Br
Cº
(Bi B₂)²
C
(B) Bn)
-2a6
C²aa
Caaad
+
= 0(T
aT
72
C.1 Anisotropic Diffusion
Proof of Lemma 5 Let To Po. and dx, A,dB, for some matrix-valued function & A,. For simplicity, assume that At is symmetric, so that A.A! = A7. By the Fokker Planck equation,
1 (z) - Lr (A3V²m) div (2
a (c())
From the RHS, we verify that the continuity equation is satisfied by dz; = - log) The forward ODE is thus given by
dx, ------- 2
1 A - log pc()..
Let pt denote the distribution of æt. Notice that the total Gaussian covariance under dx, = Add is M, = fo Asds. Thus Pt Po (N, M.) The explicit form of V (o)) is
Vlogp(()) M, exp (-(x (x 2o) ( (2M,) -x mo)) (2o - x (po((r)
exp (-(x (x 20) (2.))
(20))dp((zz)
Therefore, the forward ODE is also more explicitly written as
dr, 1 A f exp (-(x xo) (7oo)((MM)))
(2.,) (xx x0)) (xo r)dpo(20))
S Sexp (-(x 2o) (2M,)- - xo vo))dpu(Zo)
F(M,,))
dM, = Aldt.
Let yt FT-t denote the time-reversal of the forward ODE. Then the Euler-discretizatio of y is given by
le+1 = 11 6e
W1+1 + 8:
A? Mt: F FMM,,
(19)
Mres = MM. 5r A7
In the above, AG, is sequence that we assume is known priori. Mr, can be computed from A via the last line. Mo denotes the initial Gaussian noise covariance, thus Mo TT. Subsequently, simplify the notation Mr, MM and At; + Ac.
18

---

In the score-based denoising setup, we have Po be discrete distribution over the training set
x(1), ((n) Thus
F(M,z) DE exp (-(x x ()) (2M) (x
(x(1)))))))))))))))))))))))))) x) (2M) Ei=1 exp (- (2M)) x(1))
Let Z== = [2(1.
[m)] and Ze [2g 12 Observe that we are trying to denoise the query
(x+1) 2p. 28
(m+1)
is ye in 19 Therefore, 19 is equivalent
8z A,A2
(n+1) (x+1)
F(Me, 22 ((ttl)
2e+1 =2
+
2
6, 6,A7M, 7J smax (logrbe (M. Mr ))
6r.A3M M, (n+1)
(x+1)
=2 +
2
2
where the last line can be verified by definition of smax and rbf. Observe that (n+1) above equal Ze.
By pattern matching, we see that the above coincides with the choices of
M,A we = W = M
W = (1 5eA7Mr
WY
2
2
which is exactly the parameter setting in the paper.
D Miscellaneous Theory
The following standard result guarantees that Gaussian distributions, and the sum between a Gaus- sian distribution and an arbitrary unit-length vector, lie on sphere.
Lemma (Radius concentration under Gaussian perturbation.). Let E gd-1, Let NO claxd). Then & (0,1) P(I[W]]B ed ted) 2e 2e-
P(Ile u112 (1 c(d 1))1 2 led) S 2e -d2 /16
Proof of Lemma 8 The first equality is simply follows by (2vd, dus-suu--xpoommtia concentration of d-dimensional Gaussians.
The second equality follows by decomposing into ii = (u, TolTaa Tollz and -g By the same sub-exponential bound applied to un, we have
P ((Il 1122 c(d 1)1 ted/2) < 2c dt2/16
On the other hand, we have
P(II + 11 - 1| ted/2) -P -P (1|4)11 ted/2) s 2c 22- The second inequality then follows by union bound.
E
Experiments
E.1
General Experiment Details
In our experiments, we use both the MNIST (60,000 samples, 10 classes) and CIFARIO (50,000, 10 train/test split of 9:1. All experiments are run a single
classes) datasets. In each case, A5000 GPU.
In all our experiments, we do not use any positional encoding or tokenization of the image. Instead, we simply represent each image as vector of all pixel values. For MNIST, the vector dimension is 784 28 28. For CIFAR10 and CIFAR100, the vector dimension is 3072 3 x 32 32.
19

---

E.2 Additional CIFAR-10/100 Metries and FID
FID protocol (minimal). We report Fréchet Inception Distance (FID) using the Inception pool 3 features (2048-D), with feature normalization enabled. Real images are the held-out test split; fake images are generated by the same pipeline used for the corresponding figure (e.g.. per- layer denoising outputs when the figure plots layerwise losses). We compute FID per configuration and report the mean standard deviation over independent runs (same seeds used for CIFAR-I0 and CIFAR-100 100 for parity). Unless otherwise noted, we use batch sizes matching the main experi- ments and evaluate on the full test split.
E.2.1 Addendum to Fig. score test loss CIFAR-100 and FID
What we add. (i) CIFAR-100 analog of the loss curves (see Fig. 8 (ii) FID for CIFAR-10 and
CIFAR-100 at each layer for the two trained variants (RBF, STD). We follow the same evaluation protocol in App. E.2
CIFAR-100 Denoising Progression per Layer (L=6)
- RBF theory
theory
RBF treined (mean 21c, n-3)
STD trained (mean +10, n=
25
Test Loss
20
10
Layer Index
Figure 8: CIFAR-100 100 RMSE per layer for the curves in Fig.
20

---

Table 1: Per-layer FID and RMSE for CIFAR 10 and CIFAR-100,
Mean + std over seeds.
Dataset
Metric
Layer
RBF trained
STD trained
CIFARIO
FID
347.2712 2.4329
403.9372 + 0.9406
2
198.1683 + 1.4204
321.1340 + 1.1317
3
19.7453 + 1.6279
233.5431 = 1.5499
4
43.6949 + 1.2320
163.7732 = 0.9192
5
21.8445 + 0.4051
114.8899 = 0.5935
6
19.7110 + 0.0940
78.7409 0.1956
RMSE
1
20.4389 = 0.3717
30.8847 = 0.0038
2
11.0475 0.0544
18.2189 = 0.0016
3
9.9866 - 0.0112
12.4526 = 0.0037
4
10.0475 0.0148
10.3822 0.0064
5
10.1168 0.0144
9.8443 + 0.0080
6
10. 0.0141
9.7725 0.0087
CIFARIO0 FID
1
327.7921 + 5.2556
390.9332 = 1.5852
2
300.4900 2.2285
3
95.6904 = 2.5033
212.3552 2.2666
4
47.0779 + 1.8721
150.9074 + 1.3448
5
25.1974 + 1.0298
107.3123 + 1.1695
6
22.5363 0.4453
77.9261
0.9152
RMSE
20.4177 0.3744
30.8679
0.0044
2
10.9805 + 0.0717
18.1712 = 0.0147
3
9.8963 0.0532
12.3639 = 0.0331
4
9.9517 + 0.0570
10.2620 = 0.0485
5
10.0194 + 0.0577
9.7092 0.0565
6
10.0477 = 0.0578
9.6317 + 0.0600
Takeaways FID decreases monotonically with layer index both datasets, indicating progressive denoising. Across layers, RBF trained attains consistently lower (better) FID than STD trained. For RMSE, both datasets converge near 10 at the final layer; STD trained ends slightly lower than RBF trained at layer 6 (CIFAR-10: 9.77 vs. 10.15; CIFAR-100: 9.63 vs. 10.05).
E.3
Addendum to Fig. 5: Context vs. Loss on CIFAR-100 and FID on CIFAR-104 10/100
We report test-set FID and RMSE across context lengths. Values mean + standard deviation over three seeds. Results mirror the trend in Fig. 5: modest context substantially lowers error, with diminishing returns larger context.
Table 2: Addendum to Fig. 5: FID and RMSE vs, context length (Nex). Mean + la over seeds.
CIFAR-10
Bex
RMSE
368.79 i 22.66
133.9387 6.7723
2
102.12 +: 130.07
10.3142 1 0.7683
14.06 1: 0.71
11.1293 0.0403
30.14 14.65
10.3618 0.4681
32
31.80 + 11.76
10,2821 0.4452
60
47.28 + 0.81
9.8351 + 0.0090
C1F
Dax
FID
RMSE
333,26 1 77.20
150.0745 1 35.1599
2
16.55 + 1.23
11.0319 +.111900.111 0.1125
8
21.1644.72 4.74
10.7499 + 0.2683
21.94 4.06
10.6655 0.2589
42.78 4: 1.26
9.9464 + 0.0324
48
44.02 + 0.40
9.9238 + 0.0206
Note. FID decreases sharply with small context; RMSE shows smaller but consistent gains with additional context. CIFAR-100 exhibits the same qualitative pattern CIFAR-10.
21

---

E.3.1 Addendum to Fig. witness/anisotropicc CIFAR-100 and FID
What we add. (i) CIFAR-100 loss vs. #witmess tokens; (ii) FID and RMSE on CIFAR-10 and CIFAR-100 vs. 7 for the Witness and aiinesi-anisotrepic models. Baseline exact-score model metrics from Fig. 7 are also provided for comparison.
Table 3: Fig. addendum: FID and RMME at the final layer vs. #witness tokens (т) for CIFAR-10.
Mean + std over & seeds.
Witness+RBi
Witness+
# Witnesses (7)
FID
RMSE
FID
RMSE
20
339.00 + 0.70
7,8919 : 0.0020
306.09 = 0.34
7.4240 = 0.0016
40
360.82 0.13
7
240.09 0.60
7.2523 + 0.0017
100
312.56 0.02
7.3838 +.000
193.31
7.1351 0.0007
200
267.91 1.01
7.2784 0.0019
173.66 0.24
7.0598 + 0.0005
400
219.15 0.86
7.2307 + 0.0016
161.80 0.70
7.0427 + 0.0012
800
239.34 +: 0.18
7.2414 + 0.0032
168.74 0.64
7.0427 + 0.0019
1000
271.73 + 0.51
7.2924
+ 0.0014
168.89 0.37
7.0462 + 0.0003
Baseline (Exact)
21.57
10.0805
23.57
9.5456
Table 4: Fig. addendumm: FID and RMSE at the final layer vs. #witness tokens (7) for CIFAR-100. Mean + std over & seeds.
Witness+RBI
Witness+Anisotropit
# Witnesses (T)
FID
RMSE
FID
RMSE
20
313.06 0.16
7.8954
0.0020
284.66
6.37
7.4390 = 0.0055
40
337.84 + 0.84
7.7204 + 0.0027
221.24 + 0.46
7.2714 + 0.0035
100
306.16 0.14
7.4170
+ 0.0027
181.34
7.1272 + 0.0032
200
262.20 0.22
7.2780
+ 0.0022
168.53 0.06
7.0650 + 0.0019
400
222.98 0.98
7.2226 0.0038
164.91
0.22
7.0323 + 0.0036
800
236.45 0.14
7.2424
+ 0.0038
163.78
0.34
7.0244 + 0.0048
1000
258.02 0.27
7.2733 0.0034
0.26 167.02
7.0415 + 0.0036
Baseline (Exact)
26.64
9.5048
26.53
9.5242
We discuss a few key observations from Tables and
Exact Baseline: low FID, high RMSE. RMSE measures per-image reconstruction, whereas FID measures distribution-level realism. The Baseline (Exact) model implements the exact score-based denoising via cross-attention to the full training (Section 4.2); ) at test time, this dynamics pulls the noisy query toward an actual training example. This behavior is visible in Figure 4 where exact score denoising reproduces training images rather than the held-out target. Thus the model essentially "memorizes" the training set of real images, yielding excellent FID. However, this same behavior results in a worse RMSE when the noisy test image's ground truth is not in the training set, the model selects different point from the data manifold (i.e., different training image), which often far from the clean test image.
Witness models: high FID, low RMSE. In contrast, the Witness+RBi and Wit- nesotAni models replace full-set cross-attention with a compact, learned witness set (Sec- score approximation that avoids instance-level copying, so the denoising trajectory tracks the actual test instance rather than a nearest training neighbor, i.e. better generalization, and thus lower test RMSE. The cost is higher FID: outputs are less tightly matched to the empirical sample distribution measured by Inception features. Wit- nesstAnisotropic consistently outperforms Witness+RBF, reflecting the benefit of geometry-awart updates.
tion 4.3 E.4.2). This induces geometry-
These results highlight gemeralization-realisn trade-off: the Exact Baseline optimizes distribu- tional realism (FID) by memorizing the training set, whereas Witness models trade some FID for stronger generalization to unseen test inputs.
22

---

E,4 Implementation Details for Specific Experiments
In this section, we describe details of several experiment setups and Transformer implementations:
1. The in-context score-based denoising experiment from Section 4.2 is explained in Ap- pendix E.4.1
2. The experiments comparing Witness-based denoising and Anisotropic Denoising from Sec- tion 4.5 is explained in Appendix E.4.2
E.4.1 Implementation Details for In-Context Score-based Denoising
In this section, we provide further details on the in-context soore-baeed denoising experiment out- lined in Section 4.2 For ease of reference, we repeat below the setup described in the main paper:
Input consists of the noiseless training samples (r(()),..,,.""") J, well contextual images The query image is 2(++m+1), Importantly, all contextual images belong to the same class as the query image x(*+=+1), The Transformer input is 2 x(1) for i 1,..., + m, and the query point is perturbed with Gaussian noise: ((+m++1) = x((+m+1) + NN00 T² We call m the "context length". In terms of implementation, the query x((+++11) attends to the training set ((()), ,((),..., x²()) via the CrAttn module in (7), while the context samples + query 1(++1),.... x....,,(mm.)) attend to each other via the Attn module in (with appropriate masking).
The Transformer architecture is as defined in
Ze+1 W Ze + Attn
Ze; W. Wº. w + CrAttn
27,7::W W.', wk²).
We let
7o (x(+)) p....,,ttttn (atm) , p((tm+))]]] Rdx(m+1)
denote the context plus queries, and let
z = (². E Rdxn
denote the frozen training set. Since the context tokens in the first m columns of Zo are noiseless, enforce that they not updated using additional mask. In summary.
( Ze + Attn,be (Z6 (Zo, WY, W wx) + CrAttnrb (z² WE W. Wo' w.k) (m-1) for i -,...,m.. m..
1e+1
(11 E+1
2(())
Note that 2 denotes the 1th column of Ze and corresponds to token 2 (m+)) Each parameter is constrained to be a scalar multiple of identity, i.e.
W Taxds w = laxd;
W. w Tdxds W ww Taxd,
w. w. Taxd WK' w Taxas WY' = =wy I''aa
so the total trainable parameters are { (WS , wo wh wy wp' wK w/ t=l.... C IR.
E.4.2
Implementation Details for Witness and AnisotropictWitnes Experiments
In this section, provide further details on the in-context score-based denoising experiment out- lined in Section 4.5 For ease of reference, we repeat below the setup described
main paper: Input consists of the noiseless training samples ".("),..... (zx(1),...,.((x) The query image is (x+1)) The Transformer input is 2(0) zx() for i 1,...,7, n, and the query point is perturbed with Gaussian noise: z(n+1) (x+1)) + (A o²)). In terms of implementation, the query (x+1)) attends to the training set z! ,..., (..,,(()) (n) via the CrAttn module in 14
23

---

The Transformer architecture is as defined in (14)
Ze+1 W Zz + CrAttn (Ue, ZUe,Ze Ze; WY", w W
in the above, we leave out the Attn module since not used. For all subsequent discussion, Ue E R is trainable parameter, where S is the number of witnesses in one layer. We initialize each Ug by setting its Tt column to be randomly drawn training sample from x(1), ,... Note that each layer has a separate Ut.
x(0)
For Witness+RBE, we have
2 (n+1) E Rdxx
241 W Zr CrAttnrer (Ue, Zei WE', w W )
where ( corresponds to the query token 2(++1), Each parameter constrained to be a sealar multiple of identity, i.e.
w = wi Laxd, We =w?
O' laxds WK = Laxddd W we Laxd:
so the total trainable parameters the learned witnesses (Uete= RRxS and the (scalar)
parameters {ws. wi w .. wpy
Semr... C.
For AiisoossAiisstrpopccc we have
Zo - (m++1)
E Rdxl
(04 w 1e CrAttnee (Ue, Zes WY", w W wx )
where 2 corresponds to the query token (x+11, Each parameter is constrained be a d agonal matrix, so the total trainable parameters are the learned witnesses {Ue)t= C RdxS and the diagonal parameter matrices (W&, WOO w.'', wX' t=l....b <Rd (each parameterized by the d diagonal scalars).
F Future Directions
Parameter-eflicient denoising, Our characterization of anisotropic denoising with diagonal atten- tion weights suggests that even highly constrained parameterizations can learn powerful denoising algorithms. This opens up opportunities for efficient Transformer variants in generative models- for example, using low-rank or diagonal attention diffusion Transformers, or fine-tuning these components in parameter-efficien learning settings.
Localized, interpretable attention. We observe that trained Transformers often specialize differ- ent layers (or heads) to denoise approximately non-overlapping semantic patches (Figure D This suggests a promising direction for designing sparse, locality-aware attention mechanisms in struc- tured domains like vision. Such mechanisms may be more interpretable and robust to noisy or incomplete context, and they align with the geometry-aware updates formalized in our analysis.
Multi-modal Contexts
Incorporating conditioning signals (e.g., text) cross-attention into our kernel-weighted update view provides principled path to conditional denoising and controllable generation, and can be combined with parameter-efficien anisotropic modules.
G Limitations:
Our analysis and experiments deliberately focus small-scale image datasets and small-depth, single-head Transformers. While this choice isolates the denoising mechanisms we study, leaves open how the same constructions behave on larger, more varied data (e.g., high-resolution images, audio, or text) and deeper, multi-stage architectures. Second, we evaluate only pixel-space RMSE; assessing perceptual quality, class fidelity, downstream generation tasks remains future work.


Reviews:
[{"review_id": "SpqTTJOzEd", "reviewer": "Reviewer_YpJ6", "summary": "This paper studies whether Transformer architectures can inherently perform diffusion-based denoising and draws connections between Transformer attention mechanisms and classical denoising algorithms. It presents a unified theoretical framework showing that Transformers can exactly implement various denoising processes: (a) manifold denoising via Laplacian smoothing flows, (b) score-based denoising as used in diffusion probabilistic models, and (c) anisotropic diffusion processes. The authors prove an equivalence between Transformer attention updates and these denoising algorithms under certain conditions, and introduce a method using learnable “witness” tokens to approximate the diffusion score more efficiently. They validate these insights on image denoising tasks, demonstrating that simple Transformers can perform robust iterative denoising both with additional context (in-context learning scenarios) and without context (standard generative diffusion). These results highlight the Transformer's flexibility as a denoising meta-learner, capable of bridging context-dependent and context-independent generative modeling.", "questions": "See weaknesses", "limitations": "See weaknesses", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n1. Provides proofs linking Transformer attention to manifold, score-based, and anisotropic diffusion denoising, offering a unified mathematical perspective that is novel and well-substantiated.\n\n2. Controlled experiments on MNIST/CIFAR-10 show monotonic error reduction across layers, effective use of context, and clear benefits from learnable tokens, confirming the theoretical claims in practice.\n\n3. Introduces learnable witness tokens and extends the framework to anisotropic diffusion, yielding significant efficiency gains and opening new avenues for Transformer-based generative modeling.\n\nWeaknesses\n\n1. Results are confined to small-scale vision datasets and shallow, single-head Transformers; there is no evidence the approach scales to high-resolution images, deeper models, or other modalities. Even under constrained compute budgets, it would strengthen the paper to report results on more additional small- or medium-scale benchmarks like CIFAR-100, Tiny ImageNet for images, or SpeechCommands and UrbanSound8K for audio, to demonstrate broader applicability.\n\n2. The study reports only pixel-level RMSE, omitting perceptual metrics (e.g., SSIM, FID) and comparisons to state-of-the-art denoisers, so practical denoising quality and generative relevance remain uncertain.", "was_revised": true, "final_justification": "Thanks for the rebuttal that addresses my concerns.", "rebuttal": "We thank the reviewer for their thoughtful comments and constructive suggestions for improving our experiments. We address the main concerns below.\n\n## 1. Generalization beyond CIFAR-10 and Shallow Transformers\n\n> Results are confined to small-scale vision datasets and shallow, single-head Transformers; there is no evidence the approach scales to high-resolution images, deeper models, or other modalities. Even under constrained compute budgets, it would strengthen the paper to report results on more additional small- or medium-scale benchmarks like CIFAR-100, Tiny ImageNet for images, or SpeechCommands and UrbanSound8K for audio, to demonstrate broader applicability.\n\nWe appreciate the reviewer’s suggestion and agree that demonstrating broader applicability is important. To that end, we include **new results on CIFAR-100** using the same denoising implementation described in Section 4. We report RMSE values for CIFAR-100 alongside the existing CIFAR-10 results below. As shown, **RMSE trends are highly consistent across both datasets**, suggesting that the observed denoising behavior is not specific to a particular dataset.\n\nWe view this as a promising indication that our findings extend beyond CIFAR-10. We aim to include more complex/diverse datasets like Tiny Imagenet and UrbanSound8K in the future.\n\n### CIFAR10/CIFAR100 RMSE-vs-Layer\n| Layer | CIFAR-10 (Theory) | CIFAR-10 (Trained) | CIFAR-100 (Theory) | CIFAR-100 (Trained) |\n| ----- | ----------------- | ------------------ | ------------------ | ------------------- |\n| 1     | 30.9              | 20.03 ± 0.10       | 30.9               | 20.05 ± 0.08        |\n| 2     | 18.3              | 10.98 ± 0.02       | 18.3               | 10.93 ± 0.04        |\n| 3     | 12.6              | 9.98 ± 0.02        | 12.6               | 9.91 ± 0.04         |\n| 4     | 10.5              | 10.05 ± 0.03       | 10.6               | 9.97 ± 0.04         |\n| 5     | 10.0              | 10.12 ± 0.03       | 10.1               | 10.04 ± 0.04        |\n| 6     | 9.98              | 10.15 ± 0.03       | 10.0               | 10.06 ± 0.04        |\n\n\n## 2. Lack of Perceptual Metrics and Comparison to SOTA\n> The study reports only pixel-level RMSE, omitting perceptual metrics (e.g., SSIM, FID) and comparisons to state-of-the-art denoisers, so practical denoising quality and generative relevance remain uncertain.\n\nWe agree with the reviewer that perceptual metrics offer valuable insights. In response, we have computed Fréchet Inception Distance (FID) for CIFAR-10 and CIFAR-100 as a function of Transformer depth. The results reported in the table below.\n\nAs expected, the FID improves substantially with network depth and is consistent across datasets. We note that the absolute FID values are significantly higher than those of state-of-the-art generative models. This is expected due to the **highly minimalistic** nature of our architecture:\n\n1. **No positional encoding**: our model does not possess spatial inductive bias.\n2. **Flattened input**: no patching, tokenization, or convolution; attention is computed over raw pixel vectors.\n3. **No multi-head attention, MLP, or full-rank projection layers**: the attention matrices $W^Q, W^K, W^V$ are diagonal, and the Transformer uses only a single head and 6 layers.\n\nBy contrast, state-of-the-art models employ deep, multi-head Transformers or UNet-style CNNs with convolutional inductive bias, positional encodings, and patch-based input representations that are specifically designed to optimize perceptual metrics such as FID.\n\nOur focus in this paper is on **mechanistically understanding the denoising capabilities of the attention mechanism in isolation**. RMSE serves as the most direct measure of denoising accuracy in this setting. Nevertheless, we agree that FID is a useful complementary signal, and we view incorporating more perceptual metrics into future, more expressive variants of our architecture as an exciting future direction.\n\n### CIFAR10/CIFAR100 FID-vs-Layer\n| Layer | CIFAR-10 (Theory) | CIFAR-10 (Trained) | CIFAR-100 (Theory) | CIFAR-100 (Trained) |\n| ----- | ----------------- | ------------------ | ------------------ | ------------------- |\n| 1     | 405               | 343.97 ± 1.36      | 392                | 324.51 ± 2.15       |\n| 2     | 325               | 195.97 ± 1.86      | 301                | 179.15 ± 1.44       |\n| 3     | 238               | 98.03 ± 1.00       | 217                | 95.07 ± 0.91        |\n| 4     | 167               | 43.26 ± 0.77       | 155                | 46.24 ± 1.06        |\n| 5     | 113               | 21.87 ± 0.54       | 108                | 25.25 ± 0.69        |\n| 6     | 78.3              | 19.31 ± 0.15       | 76.6               | 22.88 ± 0.38        |"}, {"review_id": "7IEDTYneLj", "reviewer": "Reviewer_ukLB", "summary": "This paper shows that transformers can perform denoising using in-context examples. Some theory is developed around this idea and experiments confirm the theory.", "questions": "- “uniform distribution over a (large) training set of samples” I am not sure I understand this one. Why do we care more about the uniform distribution over a set of examples as opposed to the true underlying signal distribution? In particular if we want to model the uniform distribution over a set of examples aren’t better equipped with simply sampling from it rather than using a diffusion model?\n- What does “We verify that” mean  (line 148)?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n- The idea of using transformers for in-context learning to do denoising is really original and should be pushed further.\n\n**Weaknesses**:\n- *Code*: code is absent from the submission without justification. This greatly hinders reproducibility and reviewer’s ability to dive into the experimental setup.\n- *Experimental setup and claims*: For each experiment it’s not clear what its purpose is in relationship to the theory laid out just before. For example it’s written:\n> Figure 2a shows that the test loss decreases with the context length; this agrees with the fact that the discrete RBF Laplacian L becomes a better approximation to the Laplace Beltrami operator with increasing number of samples.\n\nBut I don’t understand how that relates to Lemma 3. nor why the the loss is the right thing to look at.\n- *Assumption not discussed*: “We assume that smax masks out the diagonal entries of the similarity matrix” why does this make sense in this context? How does it deviate from typical practice? Normalization operations not presented in section 2, what difference does it make?\n- *Phrasing/typos*: “ We highlight that the difference between the standard Attnstd and Attnrbf as follows”\n- *Equivalence between Attnrbf and Attnrbf on sphere*: the lemma title is a bit misleading as it hides the fact that Q and K weight matrices also have to identity times a scalar.\n- *Notations*: section 2 is extremely notation-heavy which makes it hard to follow.\n- *Transformer and Kernel Weighted Update practical relevance*: it is not mentioned in this section what is the practical relevance of the kernel weighted update. Indeed there is no connection to prior literature using these kinds of transformers or illustration that these are useful in practice or are good approximations. On the other hand it seems exactly designed to fulfill the manifold denoising. A similar note can be said about the unconventional diffusion score modeling setup where the distribution we learn is a convoluted set of diracs.\n- *Checklist*: the checklist has items marked as Yes or No without explanation, for example code or statistical significance.", "was_revised": true, "final_justification": "Theoretical and experimental setup issues were clarified in the discussion. \n\nScore is still pending on seeing the code as well as a plan for improving the clarity of the paper.", "rebuttal": "We thank the reviewer for their thoughtful comments and suggestions. Below, we address the main concerns, organized thematically for clarity.\n\n## 1. Unconventional Diffusion Score Modelling Setup\n\n> “uniform distribution over a (large) training set of samples” I am not sure I understand this one. Why do we care more about the uniform distribution over a set of examples as opposed to the true underlying signal distribution? In particular if we want to model the uniform distribution over a set of examples aren’t better equipped with simply sampling from it rather than using a diffusion model?\n\n> ... A similar note can be said about the unconventional diffusion score modeling setup where the distribution we learn is a convoluted set of diracs\n\nWe thank the reviewer for raising this point. We would like to clarify that our score modeling setup is **not unconventional**—it follows the standard practice in diffusion-based generative modeling.\n\nAs described in lines 137–145 of our paper, our setup corresponds exactly to the **standard formulation of denoising score matching used in nearly all modern diffusion models.** Specifically, the empirical training distribution is **a uniform mixture of Dirac measures at the training samples**, convolved with a Gaussian kernel to define the perturbed data distribution \\$p_t\\$. This “convoluted set of Diracs” is exactly what appears in standard implementations of score-based models such as DDPM [1] and EDM [3].\n\nBecause the true data distribution is unknown in practice, the training procedure *necessarily relies on the empirical distribution*. This is **consistent across virtually all generative modeling pipelines**. As noted by the reviewer, if we had access to the population distribution, *there would be no need for generative modeling at all.*\n\nThe theoretical justification for learning scores on this empirical mixture is well-established. For instance, Appendix B.3 of [3] shows that in the infinite-capacity limit, a model trained on this empirical score-matching objective will converge to the true score of the empirical distribution \\$p_t\\$. The fact that practical models generalize beyond this set of Diracs is due to **architectural inductive biases, such as convolutional structure in UNets or tokenization in DiTs**. Analyzing the generalization properties of these inductive biases is an important research question, but one that is beyond the scope of our paper.\n\n\n## 2. Relevance of Kernel Weighted Update\n> - Transformer and Kernel Weighted Update practical relevance: it is not mentioned in this section what is the practical relevance of the kernel weighted update. Indeed there is no connection to prior literature using these kinds of transformers or illustration that these are useful in practice or are good approximations. On the other hand it seems exactly designed to fulfill the manifold denoising. \n\nWe thank the reviewer for this question. We would like to clarify both the theoretical and practical relevance of the Kernel-Weighted Update (KWU) framework.\n\nFirst, the KWU abstraction captures **both manifold-denoising [10], and score-based denoising [1]**. These two algorithms are very commonly used in practice. An important contribution of our paper is to show that KWU **unifies both paradigms under a single update rule**, thereby providing a common lens through which to interpret diverse denoising behaviors.\n\nImportantly, we prove in Lemmas 2–4 that **Transformers can implement KWU updates**, using standard attention layers. This immediately implies that Transformers have the capacity to simulate both manifold denoising and diffusion denoising dynamics.\n\nAs for the practical relevance: the Transformer architecture we study **is conventional**. The only deviation from standard practice is the use of RBF attention (${Attn}_{\\text{rbf}}$) instead of softmax attention. However, this difference is minor for the following reasons:\n- **we do show in Lemma 1 that these \\$Attn_{rbf}\\$ and \\${Attn}_{smax}$ are exactly identical for inputs of constant norm**, which is almost always the case in practice due to the use LayerNorm in nearly all modern Transformers. \n- We also show in Lemma 8 that the the vectors in the denoising setup **strongly concentrate to a sphere of constant radius**.\n- Empirically, our experiments find that models trained with either attention **behave similarly**.\n\n\n## 3. Relevance of Experiments\n\n> - Experimental setup and claims: For each experiment it’s not clear what its purpose is in relationship to the theory laid out just before. For example it’s written:\n> \n>>    Figure 2a shows that the test loss decreases with the context length; this agrees with the fact that the discrete RBF Laplacian L becomes a better approximation to the Laplace Beltrami operator with increasing number of samples.\n> But I don’t understand how that relates to Lemma 3. \n\nLemma 3 shows that the Transformer is capable of implementing the denoising algorithm of [10], which is based on the Laplacian-based flow. Figure 2a is relevant for two reasons\n1. As the context length increases, the discrete Laplacian approaches the Laplace Beltrami Operator (see e.g. Theorem 1 of [10]). Lemma 3 shows that the Transformer can simulate the discrete manifold flow; **as context size increases, this discrete flow converges to the continuous Laplace–Beltrami dynamics**. We explain this on lines 109-112,and will add a more detailed explanation in the next draft.\n2. An important question in the in-context-learning literature is how the ICL loss scales with context length (e.g. [4])\n\n> nor why the the loss is the right thing to look at.\n\nThe goal of in-context denoising is to recover \\$x^{(i)}\\$ given a noisy observation \\$z^{(i)} = x^{(i)}+\\epsilon\\$. The L2 Euclidean distance between the denoised output and \\$x^{(i)}\\$ is a natural measure of recovery error. One could of course also consider alternatives like L1 loss.\n\n**Please let us know if you would like an explanation of the purpose of the other experiments.**\n\n## 4. Transformer Architecture Questions\n\n> - Assumption not discussed: “We assume that smax masks out the diagonal entries of the similarity matrix” why does this make sense in this context? How does it deviate from typical practice? \n\n\"Masking out diagonal entries of the similarity matrix\" means that each token does not attend to itself. It is a very common setup, and is implied, for instance, by lower-triangular mask matrices which are used in causal self-attention. It is also common to enforce this in manifold denoising, see e.g. Section 3.1 of [10].\n\n\n> - Equivalence between Attnrbf and Attnrbf on sphere: the lemma title is a bit misleading as it hides the fact that Q and K weight matrices also have to identity times a scalar.\n\nWe do state that \\$W^Q\\$ and \\$W^K\\$ are scalings of identity in the sentence on the second line of Lemma 1 (line 86). Additionally, \\$Attn_{rbf}\\$ and \\$Attn_{std}\\$ are functions of vector embeddings. The title highlights the fact that input vector embeddings must lie on the sphere.\n\n> Normalization operations not presented in section 2, what difference does it make?\n- If you mean \"normalization to constant norm for tokens\", that would be exactly LayerNorm, which is commonly used in Transformers. In addition, Lemma 8 provides a concentration bound for the radius \\$x_t \\sim p_t\\$ around \\$\\sqrt{td}\\$, assuming \\$p_0\\$ is normalized.\n- If you mean \"normalization of probability distribution\", that is a natural part of softmax, see e.g. proof of Lemma 1 in Appendix A.\n- If you are referring to some other normalization, please clarify and we would be happy to follow up.\n\n## 5. Miscellaneous Issues\n\n> - Notations: section 2 is extremely notation-heavy which makes it hard to follow.\n\nWe will add more exposition to make the section easier to follow.\n\n> - What does “We verify that” mean (line 148)?\n\nIt means that we can verify that the continuity equation (9) describes the density evolution in (8). There are several ways to verify this, such as via the Fokker Planck equation. We can add a more detailed proof in the appendix in the next draft. The proof can also be found in many existing diffusion model papers such as [2].\n\n\n> - Code: code is absent from the submission without justification. This greatly hinders reproducibility and reviewer’s ability to dive into the experimental setup.\n\nWe have uploaded the code as an anonymized link to the AC in an Official Comment (per the author FAQ). We believe that the AC provide the link to you soon.\n\n> - Checklist: the checklist has items marked as Yes or No without explanation, for example code or statistical significance.\n\nWe will include std bounds for all data in the next draft of our paper. The standard deviation is generally significantly less than \\$1\\%\\$ of the mean value. We provide below the std bounds for Figure 3 as an example. Please let us know if you would like the STD bounds for data in the other figures.\n\n| Layer | RBF Trained           | STD Trained           |\n|-------|------------------------|------------------------|\n| 1     | 20.0300 ± 0.0997       | 30.8837 ± 0.0031       |\n| 2     | 10.9777 ± 0.0206       | 18.2155 ± 0.0091       |\n| 3     | 9.9822 ± 0.0233        | 12.4484 ± 0.0185       |\n| 4     | 10.0486 ± 0.0253       | 10.3784 ± 0.0260       |\n| 5     | 10.1176 ± 0.0257       | 9.8413 ± 0.0297        |\n| 6     | 10.1461 ± 0.0257       | 9.7701 ± 0.0311        |\n\n*\\*RBF Theory and STD Theory have 0 variance.*\n\n## References\n\n[1] Denoising Diffusion Probabilistic Models, *Jonathan Ho, Ajay Jain, Pieter Abbeel*\n\n[2] Score-Based Generative Modeling through Stochastic Differential Equations, *Song et el 2020*.\n\n[3] Elucidating the Design Space of Diffusion-Based Generative Models, *Karras et al 2022*\n\n[4] What learning algorithm is in-context learning? Investigations with linear models, *Akyurek et el 2023*\n\n[10] Manifold Denoising, *Matthias Hein, Markus Maier 2006*."}, {"review_id": "RzCcmuOD20", "reviewer": "Reviewer_eksT", "summary": "This paper proposes to theoretically understand Transformers for in-context and out-of-context diffusion denoising. The tasks include in-context denoising, score-based diffusion denoising, denoising with learnable witness tokens, and anisotropic diffusion.  Through theoretical proof and experiments on MNIST and CIFAR-10, the paper proves that Transformers can perform iterative denoising for various denoising algorithms.", "questions": "Please refer to the Weaknesses.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n- The paper is well-motivated, as performing iterative denoising across layers in a Transformer remains an underexplored area. This work provides valuable insights that could inspire future research.\n- The paper provides solid theoretical derivations for various diffusion denoising tasks, which are supported by experimental results to justify the conclusions.\n\n### Weaknesses\n- The main concern is the practical applicability of the paper. While the paper offers theoretical analysis supporting the use of Transformers for diffusion denoising, **1)** it lacks a quantitative comparison with mainstream methods such as DiT [1] on more practical image generation tasks (such as benchmarks like COCO, ImageNet) and **2)** does not sufficiently discuss the conditions under which Transformers are more suitable than mainstream methods.\n- An important aspect of modern denoising-based image generation methods is the flexibility to **1)** adjust the number of denoising steps to trade off between efficiency and quality, and **2)** to incorporate conditioning information such as text prompts. It would be valuable for the authors to discuss how their proposed method could support or be extended to enable these capabilities.\n- The authors only conduct experiments using a six-layer Transformer and relatively simple benchmarks such as MNIST and CIFAR-10, which raises concerns about the scalability and generalizability of the proposed method to larger models and more complex datasets.\n- The related work could be better organized into a section to provide a clearer understanding of the related tasks and contextualize the contributions of this paper more effectively.\n\n[1] Scalable Diffusion Models with Transformers. ICCV 2023.", "was_revised": true, "final_justification": "Most of the concerns are addressed. I hope the author can organize the content in the rebuttal into the paper. I will keep my rating.", "rebuttal": "We thank the reviewer for their thoughtful and comments and insightful questions. Below, we address each of the key points in turn.\n\n## 1. Practical Applicability\n> The main concern is the practical applicability of the paper. While the paper offers theoretical analysis supporting the use of Transformers for diffusion denoising, 1) it lacks a quantitative comparison with mainstream methods such as DiT [1] on more practical image generation tasks (such as benchmarks like COCO, ImageNet) and 2) does not sufficiently discuss the conditions under which Transformers are more suitable than mainstream methods.\n> [1] Scalable Diffusion Models with Transformers. ICCV 2023.\n\nWe have extended our experiments to the CIFAR100 dataset (requested by reviewer YpJ6), which we provide below. The loss values are very similar to CIFAR10. We aim to include more diverse and complex experiments (such as ImageNet, as well as other modalities) in the future draft.\n\n### CIFAR10/CIFAR100 RMSE-vs-Layer\n| Layer | CIFAR-10 (Theory) | CIFAR-10 (Trained) | CIFAR-100 (Theory) | CIFAR-100 (Trained) |\n| ----- | ----------------- | ------------------ | ------------------ | ------------------- |\n| 1     | 30.9              | 20.03 ± 0.10       | 30.9               | 20.05 ± 0.08        |\n| 2     | 18.3              | 10.98 ± 0.02       | 18.3               | 10.93 ± 0.04        |\n| 3     | 12.6              | 9.98 ± 0.02        | 12.6               | 9.91 ± 0.04         |\n| 4     | 10.5              | 10.05 ± 0.03       | 10.6               | 9.97 ± 0.04         |\n| 5     | 10.0              | 10.12 ± 0.03       | 10.1               | 10.04 ± 0.04        |\n| 6     | 9.98              | 10.15 ± 0.03       | 10.0               | 10.06 ± 0.04        |\n\n## 2. Flexible Denoising\n> An important aspect of modern denoising-based image generation methods is the flexibility to ... adjust the number of denoising steps to trade off between efficiency and quality\n\nWe agree this is an important aspect of modern denoising-based generative models. Our approach naturally lends itself to this flexibility in two ways:\n\n- **Layerwise early stopping**: In Lemma 4, we show that each attention layer implements one denoising step. Figures 2 and 3 empirically demonstrate a monotonic decrease in loss with increasing layers. Thus, the model's output at an intermediate layer could be used as a lower-cost, lower-fidelity sample, this provides a direct mechanism to trade-off generation quality for speed.\n- A multi-step denoising process could be implemented by a *looped transformer* [1], in which the same attention parameters are applied recurrently across layers. Our Transformer construction in Lemma 4 uses the same weight configuration (up to scaling) for all attention layers, and thus extends naturally to this setting without modification. This approach **allows for scaling to many denoising steps with a small number of recurrent attention layers**. Furthermore, as recently explored in [2], test-time recurrence depth can be adapted based on input complexity—suggesting that **adaptive denoising steps** is feasible in our framework as well.\n\n[1] Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? *Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, Sanjiv Kumar 2024*\n\n[2] Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach *Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Tom Goldstein 2025*\n\n## 3. Incorporating conditional information\n> An important aspect of modern denoising-based image generation methods is the flexibility to  ... incorporate conditioning information such as text prompts. It would be valuable for the authors to discuss how their proposed method could support or be extended to enable these capabilities.\n\nThis is a very valuable suggestion. In text-to-image diffusion models, conditioning is typically achieved through cross-attention over textual embeddings. Our framework could naturally extend to this setting:\n1. Both Lemma 4 and Lemma 5 can be modified to incorporate conditioning via cross-attention.\n2. We point out a connection of the above to the \"in-context score-based denoising\" problem on line 200 of the paper. In this setup, the \"conditioning information\" is in the form of \"contextual demonstration images\". The Transformer denoises the query noise, conditioned on contextual demonstrations. This \"conditioning on context\" is achieved by having the query token attend to the demonstration tokens in the self-attention module.\n\n## 4. Scalability and Generalizability\n> - The authors only conduct experiments using a six-layer Transformer and relatively simple benchmarks such as MNIST and CIFAR-10, which raises concerns about the scalability and generalizability of the proposed method to larger models and more complex datasets.\n\nWe acknowledge that the architecture and datasets used in our experiments are relatively simple. However, this choice is deliberate: our primary goal is to provide a *mechanistic understanding* of how attention layers alone can implement in-context denoising.\n\nTo isolate the role of attention, our architecture processes images as a single flattened pixel vector, *without any positional encodings or convolutional components*. Consequently, **the observed denoising behavior can be attributed entirely to the attention mechanism, independent of spatial inductive biases.** To our knowledge, this has not been systematically investigated in prior work.\n\nWe agree that evaluating scalability to more complex datasets and architectures is an important future direction. In particular, adapting our constructions to settings involving positional encodings or patch-based tokenization (as in diffusion Transformers) will help illuminate how architectural components interact with the attention-driven mechanisms we uncover.\n\n> - The related work could be better organized into a section to provide a clearer understanding of the related tasks and contextualize the contributions of this paper more effectively.\n\nWe thank the reviewer for their suggestion. We will include a detailed section on relation between our work, and prior work on 1) in-context learning, 2) score-based diffusion, and 3) manifold denoising."}, {"review_id": "JKfEx7zeLq", "reviewer": "Reviewer_Xi3K", "summary": "The authors of this paper present a theoretical interpretation of transformers from the perspective of denoising processes in diffusion-based generation. They propose a specifically defined forward diffusion process that enables the reverse process to be formulated in a manner equivalent to the kernel-weighted update scheme of transformers, which can be implemented through the attention mechanism inherent to transformer architectures.\n\nBuilding on this theoretical insight, the authors conduct experiments to empirically support their findings. Their approach is validated across various settings, including single-image generation, in-context class-conditioned generation, scenarios with a limited set of learnable witnesses, and anisotropic diffusion settings.", "questions": "1. In Figure 7(a), it appears counterintuitive that the Witness + RBF configuration outperforms the baseline as the number of witnesses increases. From my understanding, the baseline represents a model without the constraints of a limited witness set and, thus, should theoretically serve as the upper bound for the performance achievable by the witness-based model. Could the authors clarify how the Witness + RBF setting surpasses the baseline in this case? Is there an underlying factor, such as regularization effects or inductive bias introduced by the witness set, that explains this observation?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### **Strength**\n\n**1. Novel interpretation of transformer computation as a denoiser**\n\nThe authors provide a theoretical analysis demonstrating that transformer computations can inherently act as a denoiser within the diffusion generation framework. This theoretical claim is further supported by empirical results on MNIST and CIFAR-10 generation tasks, confirming its practical feasibility.\n\n\n### **Weaknesses**\n\n**1. Insufficient alignment between theory and practical implementation**\n\nThe discretized sampling process described in Equation (10) requires a significant number of iterative steps to properly implement the denoising process. However, the practical transformer-based implementation utilizes only 6 layers, which may be insufficient to faithfully approximate the theoretical framework even compared with previous diffusion based methods. This discrepancy raises concerns about the strict validity of the theoretical-to-practical alignment.\n\n**2. Non-differentiability of the witness set selection process**\n\nIt remains unclear how the selection of elements from the set $U_{I}$ (the small set of learnable witnesses) is performed in a differentiable manner. If the witness set selection occurs within the training process, a clear explanation of how this selection can be integrated with gradient-based optimization is required to ensure the method's consistency and practical applicability.\n\n**3. Limited evaluation of generation performance**\n\nWhile test loss provides a useful proxy for evaluating model performance on the test set, it does not fully reflect the effectiveness of the model as a diffusion denoiser. Standard generation quality metrics, such as Fréchet Inception Distance (FID) or Inception Score (IS), should be reported to quantitatively assess the quality and diversity of the generated images. The absence of these metrics weakens the empirical validation of the proposed method.", "was_revised": true, "final_justification": "Initially, I had some concerns, but they were mostly resolved through the authors' rebuttal. After reviewing other reviewers' comments and participating in the discussion, I came to better appreciate the value of the paper and have decided to raise my score accordingly.", "rebuttal": "We thank the reviewer for their thoughtful comments and suggestions. Below, we address the main points of concern, organized thematically for clarity.\n\n## 1. Insufficient alignment between theory and practical implementation\n> The discretized sampling process described in Equation (10) requires a significant number of iterative steps to properly implement the denoising process. However, the practical transformer-based implementation utilizes only 6 layers, which may be insufficient to faithfully approximate the theoretical framework even compared with previous diffusion based methods. This discrepancy raises concerns about the strict validity of the theoretical-to-practical alignment.\n\nWe thank the reviewer for this thoughtful observation. We agree that, in general, fewer denoising steps may increase discretization error in reverse-time diffusion processes. \n\n\n- We note that in Equation (10), **the discretization step size can be made larger to reduce the required number of steps**. In our implementation, we choose a step size that enables the process to be approximated within 6 layers. As shown in Figures 2 and 3, the test loss decreases monotonically with depth, indicating that even with this coarse discretization, the Transformer continues to refine its denoising behavior.\n- Our primary goal is to **understand the mechanistic capability of attention layers to implement denoising dynamics**—not necessarily to match state-of-the-art generative quality. From this perspective, using a small number of layers allows us to clearly attribute performance improvements to specific architectural choices. \n- **Our theoretical result in Lemma 4 holds for arbitrarily deep Transformers**, and our design uses the same parameter configuration across all layers. This naturally opens the door to recurrent implementations (e.g., looped Transformers), where the same layer is reused multiple times to achieve arbitrarily fine discretization. Exploring this direction is an exciting avenue for future work.\n\n\n\n## 2. Non-differentiability of the witness set selection process\n> It remains unclear how the selection of elements from the set $U_I$ (the small set of learnable witnesses) is performed in a differentiable manner. If the witness set selection occurs within the training process, a clear explanation of how this selection can be integrated with gradient-based optimization is required to ensure the method's consistency and practical applicability.\n\nThank you for bringing up this point. Prior to training, we sample a random subset $\\psi^{(1)}...\\psi^{(\\tau)}$ of size $\\tau$ from the training set (of size $n$), and use $\\psi^{(1)}...\\psi^{(\\tau)}$ to initialize $U_\\ell \\in \\mathbb{R}^{d\\times \\tau}$. **Subsequently, we treat the *entire $U_\\ell \\in \\mathbb{R}^{d\\times \\tau}$* as a learnable parameter in continuous space, and optimize it using gradient-descent**. At convergence of training, the columns of $U_\\ell$ do not need to be elements from the training set (e.g. each column of $U_\\ell$ could be an average over multiple training samples.) \n\nThe above process is defined on lines 234-241, and we will add further clarifications in the next draft. We also remark that the optimal choice of $U_\\ell$ has connections to several kernel-related algorithms such as the Nystrom method, as detailed in [5,9,14,21] referenced in the paper. \n\n## 3. Limited evaluation of generation performance\n> While test loss provides a useful proxy for evaluating model performance on the test set, it does not fully reflect the effectiveness of the model as a diffusion denoiser. Standard generation quality metrics, such as Fréchet Inception Distance (FID) or Inception Score (IS), should be reported to quantitatively assess the quality and diversity of the generated images. The absence of these metrics weakens the empirical validation of the proposed method.\n\nWe thank the reviewer for this valuable suggestion. In response, we have computed the Fréchet Inception Distance (FID) across varying Transformer depths. The results for CIFAR-10 and CIFAR-100 are reported in the table below. As expected, the FID improves substantially with network depth and is consistent across datasets.\n\nWe note that the absolute FID values are significantly higher than those of state-of-the-art generative models. This discrepancy is expected due to the minimalistic nature of our architecture:\n\n1. **No positional encoding**: our model does not possess spatial inductive bias and cannot exploit pixel locality.\n2. **No image tokenization**: images are input as flattened vectors, so attention must operate over raw pixel vectors.\n3. **No multi-head attention, MLP, or full-rank projection layers**: the attention matrices $W^Q, W^K, W^V$ are diagonal, and the Transformer uses only a single head and 6 layers.\n\nBy contrast, state-of-the-art models employ deep, multi-head Transformers or UNet-style CNNs with convolutional inductive bias, positional encodings, and patch-based input representations that are specifically designed to optimize perceptual metrics such as FID.\n\nOur focus in this paper is on **mechanistically understanding the denoising capabilities of the attention mechanism in isolation**. This motivates our use of RMSE as a primary metric: it directly evaluates denoising fidelity without entangling architectural heuristics. Nevertheless, we agree that FID is a useful complementary signal, and we view incorporating more perceptual metrics into future, more expressive variants of our architecture as an exciting future direction.\n\n### CIFAR10/CIFAR100 FID-vs-Layer\n| Layer | CIFAR-10 (Theory) | CIFAR-10 (Trained) | CIFAR-100 (Theory) | CIFAR-100 (Trained) |\n| ----- | ----------------- | ------------------ | ------------------ | ------------------- |\n| 1     | 405               | 343.97 ± 1.36      | 392                | 324.51 ± 2.15       |\n| 2     | 325               | 195.97 ± 1.86      | 301                | 179.15 ± 1.44       |\n| 3     | 238               | 98.03 ± 1.00       | 217                | 95.07 ± 0.91        |\n| 4     | 167               | 43.26 ± 0.77       | 155                | 46.24 ± 1.06        |\n| 5     | 113               | 21.87 ± 0.54       | 108                | 25.25 ± 0.69        |\n| 6     | 78.3              | 19.31 ± 0.15       | 76.6               | 22.88 ± 0.38        |\n\n\n## 4. Question: why does witness+RBF surpass baseline?\n> In Figure 7(a), it appears counterintuitive that the Witness + RBF configuration outperforms the baseline as the number of witnesses increases. From my understanding, the baseline represents a model without the constraints of a limited witness set and, thus, should theoretically serve as the upper bound for the performance achievable by the witness-based model. Could the authors clarify how the Witness + RBF setting surpasses the baseline in this case? Is there an underlying factor, such as regularization effects or inductive bias introduced by the witness set, that explains this observation?\n\nThis is a sharp observation. This is related to your question above about how $U_\\ell$ is optimized. We conjecture that the witness-based model achieves better performance than the baseline because **it has the freedom to select (and attend to) witness-vectors outside the training set,** whereas the baseline can only attend to samples from the training set. For instance, a learned witnesses may capture certain higher-level features of the training data. This related to \"inductive bias introduced by the witness set\" as suggested by the reviewer.\n\nIt would certainly be interesting to investigate this phenomenon in more detail in future work."}]

Author rebuttal:

Meta-review: This paper develops a unified theoretical framework showing that Transformers can implement various diffusion-based denoising algorithms, supported by the novel idea of witness tokens for efficient score approximation and empirical validation on MNIST, CIFAR-10, and CIFAR-100. Reviewers found the work well-motivated, rigorous, and conceptually novel, highlighting its mechanistic interpretation of attention as a denoiser and its potential implications for generative modeling.

The main concerns were the limited experimental scope, absence of large-scale evaluations, and lack of comparisons with stronger baselines. The rebuttal partially addressed these by adding CIFAR-100 results, FID evaluations, and clarifying the setup, while acknowledging the need for broader benchmarks in future work.

Overall, the contribution is seen as original and impactful, providing valuable theoretical insights despite modest empirical scale.
Decision: Accept (poster)