In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Authors: Matthew Smart
Alberto Bietti
Anirvan M. Sengupta
Venue: ICML 2025 oral (2025)
Keywords: attention
in-context learning
denoising
associative memory
Hopfield network
transformers
Primary area: deep_learning->attention_mechanisms
TL;DR: We show that one-layer transformers perform optimal in-context denoising through a single step of context-dependent associative memory inference.

Abstract: We introduce in-context denoising, a task that refines the connection between attention-based architectures and dense associative memory (DAM) networks, also known as modern Hopfield networks. Using a Bayesian framework, we show theoretically and empirically that certain restricted denoising problems can be solved optimally even by a single-layer transformer. We demonstrate that a trained attention layer processes each denoising prompt by performing a single gradient descent update on a context-aware DAM energy landscape, where context tokens serve as associative memories and the query token acts as an initial state. This one-step update yields better solutions than exact retrieval of either a context token or a spurious local minimum, providing a concrete example of DAM networks extending beyond the standard retrieval paradigm. Overall, this work solidifies the link between associative memory and attention mechanisms first identified by Ramsauer et al., and demonstrates the relevance of associative memory models in the study of in-context learning.

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Matthew Smart Alberto Bietti
Anirvan M. Sengupta 234
Abstract
derstanding the mechanisms behind transformer-based net- works remains an open challenge. This challenge is exacer- bated by the growing scale and complexity of modern large networks. Toward addressing this, researchers studying sim- plified architectures have identified connections between the attention operation that is central transformers and asso- ciative memory models (Ramsauer et al., 2021), providing not only avenue for understanding how such architectures encode and retrieve information but also potentially ways to improve them further.
We introduce in-context denoising, a task that refines the connection between attention-based ar- chitectures and dense associative memory (DAM) networks, also known modern Hopfield net- works. Using a Bayesian framework, we show the- orctically and empirically that certain restricted denoising problems can be solved optimally even by a single-layer transformer. We demonstrate that trained attention layer processes each de- noising prompt by performing single gradient descent update context-aware DAM energy landscape, where context tokens serve as asso- ciative memories and the query token acts as an initial state. This one-step update yields better solutions than exact retrieval of either a context token or spurious local minimum, providing a concrete example of DAM networks extending beyond the standard retrieval paradigm. Overall, this work solidifies the link between associative memory and attention mechanisms first identified by Ramsauer et al., and demonstrates the rele- vance of associative memory models in the study of in-context learning.
The most celebrated model for associative memories in sys- tems neuroscience the so-called Hopfield model (Amari, 1972: Nakano, 1972; Little, 1974; Hopfield, 1982). This model has capacity store "memories" (stable fixed points of recurrent update rule) proportional to the number of nodes (Hopfield, 1982; Amit et al., 1985). In the last decade, new energy functions (Krotov & Hopfield, 2016; Demircigil et al., 2017) were proposed for dense associative memories with much higher capacities. These energy func- tions are often referred as modern Hopfield models. Ram- sauer et al. (2021) pointed out the similarity between the one-step update rule of a certain modern Hopfield network (Demircigil et al., 2017) and the softmax attention layer of transformers, generating interest in the statistical physics and systems neuroscience communities (Krotov & Hopfield, 2021; Krotov, 2023; Lucibello & Mézard, 2024; Millidge et al., 2022). Recent work has extended this concept to improve retrieval by incorporating sparsity (Hu et al., 2023; Wu et al., 2024b; Santos et al., 2024; Wu et al., 2024a), while others have leveraged associative memory principles to de- sign new energy-based transformer architectures (Hoover et al., 2023). However, these extensions and the founda- tional construction in Ramsauer et al. (2021) primarily focus on the specific task of exact retrieval (converging to fixed point), while in practice transformers may tackle many other tasks.
1. Introduction
The transformer architecture (Vaswani et al., 2017) has achieved remarkable success across diverse domains, from natural language processing (Devlin et al., 2019; Brown et al., 2020; Touvron et al., 2023) to computer vision (Doso- vitskiy et al., 2021). Despite their practical success, un-
CCenter for Computational Biology. Flatiron Institute, New York, NY, USA? Center for Computational Mathematics, Flat- iron Institute, New York, NY, USA Center for Computa- tional Quantum Physics, Flatiron Institute, New York. NY, USA Department of Physics and Astronomy, Rutgers Univer- sity, Piscataway. NJ, USA. Correspondence to: Matthew Smart <msmart @ iutironinstitute.rrg.. Anirvan M. Sengupta <anir- vans.physicse gmail.com>.
To explore this connection beyond retrieval, we intro- duce in-context denoising, a task that bridges the behav- ior of trained transformers and associative memory net- works through the lens of in-context learning (ICL). In standard ICL, a sequence model is trained to infer an un- known function g from contextual examples, predicting
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
((X++1) given a sequence of input-output pairs E (Xi,g((i))) (())) (XL+1, -)). Crucially, g is implied solely through the context and differs across prompts performant models are therefore said to "lear g(x) in context". While ICI. has been extensively stud- ied in supervised settings (Garg et al., 2022; Zhang et al., 2024; Akyürek et al., 2023; Reddy, 2024), recent work suggests that transformers may internally emulate gradient descent over context-specific loss function during infer- ence (Von Oswald et al., 2023; Dai et al., 2023; Ahn et al., 2023). This general perspective aligns with our findings.
ruption be defined by X Pnoise (- (-1X The random sequence F X2, .... XL, X) are given as "context" (input) to sequence model F(; 0) which outputs an esti- mate XL+1 of the original (L + 1)-th token The task is to minimize the expected loss [E[[ (X11)1 XL+1)) )] for some loss function I(.-)) Namely, our problem is to find
min E Px= -..IIuu 1t X-Ponsa (2 1 (F((F) X X }]
In practice, we choose X = X + Z, pure token cor- rupted by the addition of isotropic Gaussian noise Z ~ N o2 m)), and our objective function to minimize is the mean squared error (MSE) E[[] [[||8++1 L+1 - X1+111%)
In this work, we generalize ICL to an unsupervised setting where the prompt consists of L samples from a random dis- tribution and the query is noise-corrupted sample from the same distribution. This shift allows us to probe how trained transformers internally approximate Bayes optimal infer- ence, while deepening the connection to associative mem- ory models which are prototypical denoisers. By setting up this problem in this way, we also attempt to answer a few questions. One concerns the memorization generalization remmorizaioo-geneeriizatiiio dilemma in denoising: a Hopfield model's success is usually measured by successful memory recovery, while in-context learning may have to solve completely new problem. An- other question has to do with the number of iterations of the corresponding Hopfield model: why does the Ramsauer et al. (2021) correspondence involve only one iteration of Hopfield energy minimization and not many?
In the following subsection, we explain the pure token dis- tributions for three specific tasks. These tasks are of course structured so that one-layer transformer has the expres- sivity to capture a solution, which, as L + co, provides optimal solution, in some sense. To that end, we derive Bayes optimal estimators for each of the three tasks, under the assumption that we know the original distribution px of pure tokens. In Section 3, we use these estimators as base- lines to evaluate the performance of the denoiser f(E, 0) based on a one-layer transformer trained on finite datasets.
2.2. Task-specific token distributions
We consider three elementary in-context denoising tasks, where the data (vectors in IR²) comes from:
In summary, our contributions are as follows: In Sec- tion 2, we introduce in-context denoising as a framework for understanding how transformers perform implicit inference beyond memory retrieval. In Section 3, we establish that single-layer transformers with one attention head are expres- sive enough to optimally solve certain denoising problems. We then empirically demonstrate that standard training from random weights can recover the Bayes optimal predictors. The trained attention layers are mapped back to dense asso- ciative memory networks in Section 4. Our results refine the general connection pointed out in previous work, offer new mechanistic insights into attention, and provide concrete example of dense associative memory networks extending beyond the standard memory retrieval paradigm to solve a novel in-context learning task.
Linear manifolds (d-dimensional subspaces)
2. Nonlinear manifolds (d-spheres)
3. Small noise Gaussian mixtures (clusters) where the component means have fixed norm
Below we describe the task-specific distributions px and the process for sampling tokens {. The same corruption pro-
cess applics to all cases: X- = XL+1 + Z, Z ~ In).
1.2.1. CASE 1 LINEAR MANIFOLDS
A given training prompt consists of pure tokens sampled from random d-dimensional subspace S of R
2. Problem formulation: In-context denoising
Let P be the orthogonal projection operator to ran- dom d-dim subspace S of R!,, sampled according to the uniform measure, induced by the Haar measure on the coset space (n - d) x O(d), on the Grass- manian G(d,, n). the manifold of all d-dimensional sub- spaces of IR"..
In this section, we describe our general setup. Recurring common notation is described in Appendix A.
2.1. Setup
Each task corresponds to a distribution D over the proba- iid bility distribution of data: px ~ D. Let Xi,..., X141 ~
Let Y ~ AN IT.I)) and define X = PY: we use this procedure to construct the starting sequences (X1,..,, X1+1) of 1 independent tokens.
Px, define the sampling of the tokens. Let the noise cor-
2

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
(a)
Problem formulation
(b) Prompt: Pure tokens from a data distribution and a single corrupted example
(prompts randomily constructed from pre-specilee lask dist lbttio)
sample a task from a task distribution
Query ===+=
sample context
Target
tokens
(a ~ PX
Prediction
== Fs(E)
Geotg
o2 corruptional finel token
query target
(1) Ps 8) 2)x(2)))
Case Linear manifolds
Case 2:
Case 3:
Nonlinear manifolds
Gaussian mixtures
Figure 1. (a) Problem formulation for general in-context denoising task. (b) The three denoising tasks considered here include instances linear and non-linear manifolds as well as Gaussian mixtures. In each case, task embedding E7()) consists sequence pure tokens from the data distribution pX D where D denotes the task distribution, along with single query token that has been corrupted by Gaussian noise. The objective predict the target (i.e. denoise the query) given information contained only in the prompt.
We thus have px (0, P) with the Haar distribution of P characterizing the task ensemble associated with D.
zero. The distribution of tasks D, is decided by the choice of (1))
For our ideal case, we will consider the limit that the vari- ances go to zero. In that case, the density is simply
2.2.2. CASE 2 NONLINEAR MANIFOLDS
We focus on the case of d-dimensional spheres of fixed radius centered at the origin I
K
pxo (x)
Σ
waō(x pa)- a=l
Choose random d+ 1-dimensional subspace V of R", sampled according to the uniform measure, as before, on the Grassmanian G(d + n). The choice of this random subspace generates the distribution of tasks D.
2.3. Bayes optimal denoising baselines for each case
The first tokens in E are "pure samples" from that should provide information about the distribution for our denoising task. Our performance expected be no better than that of the best method, in the case that the token dis- tribution and also the corrupting process are exactly known. This is where the Bayesian optimal baseline comes in. As is well-known, the Bayes optimal predictor of quantity is given by the posterior mean. We use that fact to compute the Bayes optimal loss.
Inside V, sample uniformly from the radius R sphere (once more, a Haar induced measure on a coset space Old + ()/(d)). We use this procedure to construct input sequences X1 = (Iı, IL+1) of L independent tokens.
In practice, we uniformly sample points with fixed norm in Rd and embed them in IR" by concatenating zeros. We then rotate the points by selecting a random orthogonal matrix Q E Raxx
In particular, we seek a function f F. - R such that E xx.x ||IX- - [(X) is minimized. Since the perturba-
tion Z is Gaussian, the posterior distribution of X, given X is
2.2.3. CASE 3- GAUSSIAN MIXTURES (CLUSTERING)
Pure tokens are sampled from weighted mixture of isotropic Gaussians in n-dimensions, (Wa; (Ma, o²) ] The density is
Pxix(" (x (8) (o))
where C(ū) is a normalizing factor (see Appendix A.2 for more explanation). The following proposition sets up baseline which expect to compare results as L 00. The proof is in Appendix B.I.
(x Σ w. C.c -le-Hall /207
Proposition 1. For each task, specified by the input distri- bution px, and the noise model PXIX-
where C, = (2 /2 are normalizing constants. The Ha are independently chosen from a uniform distribution on the radius R sphere of dimension 1, centered around
E Ex,8 [IX s(8)11 2 Eg (TCCov(X1X)) [x)] (2)
3

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
This lower bound is met when [(X)
[[ X
where X x P& and I, is the modified Bessel function of the first kind.
Thus, the Bayes optimal denoiser is the posterior expectation for X given X. The expected loss is found by computing the posterior sum of variances.
Clustering case. For clustering with isotropic Gaussian mixtures (Wa, (Ma, o²))a==> the Bayes optimal predictors for some important special cases are as follows. See Ap- pendix C.3 for the general case.
These optimal denoisers can be computed analytically for both the linear and nonlinear manifold cases (given the variances and dimensionalities). In the Gaussian mixture (clustering) case, it depends on the choice of the centroids which then needs to be averaged over.
Proposition 4. For general isotropic Gaussian model with Ta = =0. Pall = for all a I....,K. K.
Sap (x) = X
Linear ease. For the linear denoising task, pure samples X are drawn from an isotropic Gaussian in a restricted subspace. The following result provides the Bayes optimal predictor in this case, the proof of which is in Appendix C.I.
of
0%
Ea X1108tt0 pa
X
o + +
0
+
0z
Ea WieeltX.)/loftt
If + 0.
Proposition For px corresponding to Subsection 2.2.1, the Bayes optimal answer
Eo w e(Ha pio
(np E[X X]
(8)
00
Eo Wee
(X) [[X|X]
PX,
(3)
0 07
and the expected loss is
In all three cases, we notice similarities between the form of the Bayes optimal predictor, and attention operations in transformers, connection which explore below.
E
1|P-X- - -X++1l
dojo (0) +z2)
(4)
3. In-context denoising with one-layer transformers - Empirical results
Projection
a p(i)
22
In this section, we provide simple constructions of one-layer transformers that approximate (and under certain conditions, exactly match) the Bayes optimal predictors above.
Projection (shrunk)
a p(0
a
(1) Input: Let Px .P...PX (N) iid D, be distributions sampled for one of the tasks. For each distribution p& sample E(i) := (x.,)...,, x.,,
o + oz
x() taking value in (=x(++1 be an input to a sequence model. We also retain the true
Figure 2.
Baseline estimators for the case of random linear mani- folds with projection operator p().
(L + 1)-th token (6)
L+1 for each
Manifold case. In the nonlinear manifold denoising prob- lem, we focus on the case of lower dimensional spheres S (e.g. the circle S R²). For such manifolds, the Bayes optimal answer is given by the following proposition.
Objective: Given an input sequence E(9, return the uncorrupted final token x We consider the mean- squared error loss over a collection of training pairs, (E x (1) 1.11j=11
Proposition 3. For px defined as in Subsection 2.2.2, with P being the orthogonal projection operator to V, the dl + dimensional linear subspace, with R being the radius of sphere S, the Bayes optimal answer is
N C(I) IIFEE -)
(9)
fap = E[X X
where F(E(),, denotes the parametrized function predict- ing the target final token based on input sequence E(i
S 811081 x dS,
(5)
1 (clx ds.
3.1.
One-layer transformer and the attention between the query and pure tokens
1841
R o2
X R
(6) To motivate our choice of architecture, let us start by dis- cussing the linear case.
R X1
||
oz

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
There we have (X) = of PX. Note that, by the %
The proof of the theorem is in Appendix D.3. See Ap- pendix E, particularly Proposition 6, for consideration of convergence rates, Note that the condition of px being supported on sphere is not artificial as, in many practical transformers, pre-norm with RMSNorm gives you inputs on the sphere, up to learned diagonal multipliers.
strong law of large numbers, P
==
1 El= X,XX is oiL
random matrix that almost surely converges component-by component to the orthogonal projection P as -> oo, since, for each 6, X has the expectation of and that X is Gaussian random variable with zero mean and finite covariance matrix. So we could propose
Note that the natural form of attention that is suggested by our formulation of in-context denoising would involve Gaussian kernels:
f(X)
of og + oz
PX (o) + +z))L
Σ
X).
X = Fa(E, Σ. Wr X,e 3IWxxx,-wox> Σ.
X|I
(10)
(14)
The relation between softmax attention and the Gaussian kernel has been noted in (Choromanski al., 2021; Am- brogioni, 2024) and a Gaussian kernel-based attention is
We now consider a simplified one-layer linear transformer (see Appendices D.1 and D.2 for more detailed discussions) which still has sufficient expressive power to capture our finite sample approximation to the Bayes optimal answer. We define
implemented in (Chen et al., 2021). A related Hopfield energy, with WK, Wq. and Wpy proportional to identity matrices, is proposed in (Hoover et al., 2024a).
1 WWV.X1...4 - W.....Kg WxQX
x= = Fi (E,0)
(11)
For the linear case, we use linear attention, but that may not be essential. Informally speaking, the softmax atten- tion model has the capacity to subsume the linear attention model.
taking values in IRR, where XI.I. = [XI,..., Xy] taking values in gTxL, with learnable weights WkQ. WPv E R=x= abbreviated by 0. Note that, when WPv aln, WKQ = BIn, and aB o. F(E,0) should ap- proximate the Bayes optimal answer (x) as L - co. For detailed discussion of the convergence rate, see Ap- pendix E, in general, and Proposition 5, in particular.
Proposition 3.2. As + 0, F(E, (Wv,,eWko) -
4 + I Wrv Σ Ex((x, - XxXXX-XX\Wxo + O(c),
(15)
Similarly, we could argue that the second two problems, the d-dimesional spheres and the o -+ zero limit of the Gaussian mixtures could be addressed by softmax attention
t=1
where X = 1, Et X, is the empirical mean.
X = FFEFFEOOFF 0) = Wp So.t: W x) (12)
See Appendix F for the details of small W WKQ expansion and Appendix F.1 for the proof of Proposition 3.2.
taking values in IR". The function softmax(z) == 1 Ee ed (e
y E R" is applied column- wise.
For case I, note that [[X, = 0 and covariance of X, is finite, E[X] = O, and |||| = O(t), allowing us to drop as L + co. If, in addition, is small, only the second term survives. Thus, F(E, (E to approximate Fin (E, (Wpv, Wk))) when L is large and is small, with V, large. We therefore could use the softmax model for all three cases.
For both problems, namely the spheres and the do 0 Gaussian mixtures, we could have WPV = aln, WKQ 31, with = 1 = 1/02 providing Bayes optimal an- swers as I > xc.
eWkq)) (WV.)) starts
In fact, we could make a more general statement about distributions px where the norm of X is fixed.
3.2. Case 1 - Linear manifolds
Theorem 3.1. If we have task distribution D so that the support of each px is the subset of some sphere, centered around the origin, with a px -eeeeneennt radius R, then the function
The Bayes optimal predictor for the linear denoising task from Section 2.3 suggests that the linear attention weights should be scaled identity matrices with their product sat- isfying a; ootoz Fig. 3 shows that a one-layer net- work of size n 16 trained on tasks with 02 = 1,
Dt, Xxe
F(((X))E.))))0))) 0),00
(13)
El_ (013))00
2. d 8, L 500 indeed achieves this bound, train- ing nearly diagonal weights with the appropriate scale (w (wxk)) (w (6)) = 0.327 x 1/3 (similar weights are learned for each seed, up to sign flip).
converges almost surely to the Bayes optimal answer (ā) for all ā R", as L + co. The optimal parameter 0" refers to Wpy In, WKQ 0z
1
5

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
(a)
Case 1: Linear manifolds
Case 2: manifoos
manifolds
Case Gaussian mixtures
0
1.0
0.08
0.06
0.8
1.0, 06
2.0,
2 0.1,
8,
o
MSE/2n
linear test train
softmax test train
0.05
0.06
0.6
p(O E
0.04
0.04
0.4
fopo (2)
0.03
0.02
0
20
40
25
50
75
100
25
50 Epoch
75
100
Epoch
Epoch
(b)
weighte linear
softmax
weights
weights diagonall
weights
Final weights
0.563
0.501
0.196
1.607
3,043
5.127
0.998
Wso
Wwy
Wro
Wiy
WxO
Wwv
Wxo
Wrv
-35
-1.
3.0
0.0
4.0
Figure 3. (a) Training dynamics for the studied cases using one-layer softmax attention (circles) as well as linear attention (triangles). Solid lines represent the average loss over six seeds, with the shaded area indicating the range for cases 2 and 3. For cach case, the grey dashed baseline indicates the O-predictor, and the pink line indicates the Bayes optimal predictor. All cases use context length of I. = 500, ambient dimension n 16, and are trained with Adam on a dataset of size 800 with batch size 80 and standard weight initialization Wiy Ui,1/vn, 1/vni. (b) Final attention weights WxQ and Wev are shown. For each, we indicate the mean of the diagonal elements. Representative initial weights are displayed for the second and third case.
Fig. 4(a) displays how this bound is approached as the context length L of training samples is increased. In Fig. 4(b) we study how the performance of a model trained to denoise random subspaces of dimension d = 8 is affected by shifts in the subspace dimension at inference time. We find that when provided sufficient context, such models can adapt with mild performance loss to solve more challenging tasks not present in the training set.
3.3. Case 2 Nonlinear manifolds
Fig. 3 (case 2) shows networks of size = 16 trained
to denoise subspheres of dimension d = 8 and radius
R = 1, with corruption o2 = 0.1 and context length L = 500. Once again, the network trains to have scaled identity weights.
We note that although the network nearly achieves the opti- mal MSE on the test set, the weights appear at first glance to deviate slightly from the Bayes optimal predictor of Sec- tion 2.3, which indicated WPV = al, WxQ = 331 with a = 1, 3 = 1/02. 2. To better understand this, we consider a coarse- grained MSE loss landscape by scanning over a and B. See Fig. 6(a) Appendix G. We find that the 2D loss landscape has roughly hyperbolic level sets which is sug- gestive of the linear attention limit, where the weight scales become constrained by their product a²3. Reflecting the symmetry of the problem, we also note mirrored negative solutions (i.e. one could also identify a = -1, B ====,,=-1//- =1/02 from the analysis in Section 2.3). Importantly, the plot shows that the trained network lies in the same valley of the loss landscape as the optimal predictor, in agreement with Fig. 3. Moreover, the shape of the loss landscape suggested that linear attention might also be applicable to this case, which we demonstrate and discuss further in Appendix G.
It is evident from Fig. 3(a) that the softmax network per- forms similarly to the linear one for this task. We can un- derstand this through the small argument expansion of the softmax function mentioned above. The learned weights dis- played in Fig. 3(b) indicate that gsoftmax 0.194 becomes small (note it decreases by factor s 0.344 relative to glinear), while the value scale a sofimax ~ 1.607 becomes larger by a similar factor ~ 1/c to compensate. Thus, al- though the optimal denoiser for this case is intuitively ex- pressed through linear self-attention, can also be achieved with softmax self-attention in the appropriate limit.
Moreover, we find that when the entire prompt undergoes global invertible transformation A / I, the optimal at- tention weights are no longer scaled identity matrices but acquire a structured form determined by A. Both linear and softmax attention layers are able to recover this structure through training; see Appendix H for details and empirical verification.
3.4. Case 3 - Gaussian mixtures
Figure 3 (case 3) shows networks of size 16 trained to denoise balanced Gaussian mixtures with 8 compo-
6

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
(a)
Effect of context length on training
(b)
Shifting the subspace dimension at inference time
weight scaling
Predict
subspace dimension can vary
0.30
Meen diagonal weights o/3
Performance
0.25
0
maintained away
L-30
(+a)
from d=8
MSE
L-50
0.5
L-500
MSE/ h
Ineer projection
train Loss trained
test
$1-1
((0))
0.4
1(6")
Train n-16 model: d-8, L-500
in-context learning
subspace provided only via context
100
200
300
400
500
5
10
15
Figure (a) Trained linear attention network converges Bayes optimal estimator context length increases (n 16, d 8, oñ = o = 1). (b) A network trained to denoise subspaces of dimension d can accurately denoise subspaces of different
dimensions presented at inference time, given sufficient context.
nents that have isotropic variance of 0.02 and centers randomly placed on the unit sphere in R The corruption magnitude is 2 0.1 and context length is I. 500. The baselines show the zero predictor (dashed grey line) well as the optimum from Proposition (4) (pink) and its of - 0 approximation Eq. (8) (grey).
Consider the energy function:
1 2a
Σ expx
E(XI:L,
log
(16)
B
which mirrors the Ramsauer et al. (2021) construction but with a Lagrange multiplier added the first term. Figure illustrates this energy landscape for the spherical manifold case.
The trained weights qualitatively approach the optimal esti- mator for the zero-vare variane limit but with slightly different scaling: while the scale of Wpy is s 1, the W KQ scale is 3 ~ 5.127 1/03. To study this, we provide correspond- ing plot of the 2D loss landscape in Fig. 6(a) in Appendix G. While the symmetry of the previous case has been broken (the context cluster centers { will not satisfy (14) = O), we again find that the trained network lies the anticipated global valley of the MSE loss landscape.
Num. steps:
Num. steps: 50
00o context tokens
1.-
nou
82
4. Connection to dense associative memory networks
@ query
target
* prediction
trajectories
In each of the denoising problems studied above, we have shown analytically and empirically that the optimal weights of the one-layer transformer W x aT, Wx WKQ Rs BT. In the softmax case, the trained denoiser can be concisely expressed as
-1.0
1.0
-1.0
1.0
81
{(X
scaled identity matrices
Figure 5. Gradient descent denoising for the nonlinear manifold case (spheres) in n with d context- aware dense asso- ciative memory network E(XI.L, x) is constructed whose gradient corresponds to Bayes optimal update (trained attention layer). Note that the density of sampled context tokens sculpts the valleys of the energy landscape. Left: the attention step of a one-layer transformer trained the denoising task corresponds gradient descent step. Right: Iterating the denoising process as is conventional for Hopfield networks can potentially degrade the estimate by causing it to become query-independent (e.g. con- verging distant minimum). Here It = 1, 10, L 20 and = 1, 1/02.
4 -=x(X1,, = sof.ma (3XT
re- re-written such that X E R stores pure context tokens.
single
We now demonstrate that such denoising corresponds to one-step gradient descent (with specific step sizes) of energy models related to dense associative memory networks, also known as modern Hopfield networks (Ramsauer et al., 2021; Demircigil al., 2017; Krotov Hopfield, 2016).

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
An operation inherent to the associative memory perspective is the recurrent application of denoising update. Gradient descent iteration s(t + 1) = s(t) y V,E (..,,.,,s(t)) yields
More broadly, our study suggests that trained attention lay- ers can readily adopt structures that facilitate context-aware associative retrieval. We have also noted preliminary con- nections between our work and other architectural features of modern transformers, namely layer normalization and residual streams, which warrant further study.
2) s(t) + + soxtmaxaxaxx 1.2 (t)). stt + 1) = 21
In-context denoising and generative modeling both involve learning about an underlying distribution, suggesting poten- tial relationships between these two tasks. Recently, Pham et al. (2024) invoked spurious states of the Hopfield model as a way of understanding how one can move away from retrieving individual memorized patterns towards general- ization via appropriate mixtures of multiple similar "mem- ories". In our work, one-step updates do not have to land in spurious minimum, but we often operate under cir- cumstances where there are such states (see, for example, the energy landscape in Fig. 5). More generally, analo- gies between energy-basee baaee associative memory and diffu- sion models have recently been noted (Ambrogioni, 2024; Hoover et al., 2024b). Lastly, Bayes optimal denoisers play an important role in the analysis (Ghio et al., 2024) of a very related generative model that is based on stochastic interpolants (Albergo & Vanden- 2023). Although this work focuses the case where it is possible to sample enough tokens from the relevant distributions for certain functions to converge, generative models become important when the distribution is in a prohibitively high- digh-dimensional space making direct sampling difficult. Nonetheless, in- vestigating the precise relationship between our work and different generative modeling approaches would be an inter- esting direction to pursue.
(17)
is now clear that initializing the state to the query s(0) = fb
and taking a single step with size Y = recovers the be- havior of the trained attention model (Fig. 5). The attention mechanism here is thus mechanistically interpretable: the context tokens XI: induce a context-dependen associa- tive memory landscape, while the query acts as an initial condition for inference-time gradient descent. One could naturally consider alternative step sizes and recurrent itera- tion. However, Fig. 5 demonstrates that naive iteration of Eq. (17) has the potential to degrade performance.
Additional details are provided in Appendix I. In particu- lar, the energy model for linear attention is discussed in Appendix I.1.
5. Discussion
Motivated by the connection between attention mechanisms and dense associative memories, here we have introduced in- context denoising, task that distills their relationship. We first analyze the general problem, deriving Bayes optimal predictors for certain restricted tasks. We identify that one- layer transformers using either softmax or lincarized self- attention are expressive enough to describe these predictors. We then empirically demonstrate that standard training of attention layers from random initial weights will readily converge to scaled identity weights with scales that approach the derived optima given sufficient context. Accordingly, the rather minimal transformers studied here can perform optimal denoising of novel tasks provided at inference time via self contaiedd prompts. This work therefore sheds light on other in-context learning phenomena, point we return to below.
Overall, this work refines the connection between dense associative memories and attention layers first identified in (Ramsauer et al., 2021). While we show that one en- ergy minimization step of a particular DAM (associated with a trained attention layer) optimal for the denoising tasks studied here, it remains open question whether multilayer architectures with varying or tied weights could extend these results to more complex tasks by effectively performing multiple iterative steps. This aligns with re- cent studies on in-context learning, which have considered whether transformers with multiple layers emulate gradient descent updates on a context-specific objective (Von Oswald et al., 2023; Shen et al., 2024; Dai et al., 2023; Ahn et al., 2023), and may provide bridge to work on emerging archi- tectures guided by associative memory principles (Hoover et al., 2023). Investigating when and how multilayer atten- tion architectures perform such gradient descent iterations in a manner that is both -eeeext-dependent and informed by large training set represents an exciting direction for fu- ture research at the intersection of transformer mechanisms, associative memory retrieval, and in-context learning.
While practical transformers differ in various ways from the minimal models studied here, we note several key con- nections. Intriguingly, the self-attention heads of trained transformers sometimes exhibit weights WKQ. WPV that resemble scaled identity matrices, i.c. cl +e with small fluc- tuations Eij ~ AO, )²), an observation noted in Trockman & Kolter (2023). This phenomenon motivated their pro- posal of "mimetic" weight initialization schemes mirroring this learned structure. Relatedly, connections to associative memory concepts have been explored in other architectures (Smart& & Zilman, 2021),, which enabled data data-dependent weight initialization strategies to be identified and leveraged.
8

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Software and Data
12):1077-1018, 1985. ISSN 10502947. doi: 10.1103/ PhysRevA..32.007.
Python code underlying this work is available at conea///ooommmtttttttaaoooccncee
Bollé, D., Nieuwenhuizen, M., Castillo, P., and Ver- beiren, T. A spherical hopfield model. Journal of Physics A: Mathematical and General, 36(41): 2003.
Acknowledgements
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
MS acknowledges M. Mézard for very useful feedback on an earlier version of this work. AS thanks D. Krotov and P. Mehta for enlightening discussions on related matters. Our early work also benefited from AS's participation in the deeplearning23 workshop at the Kavli Institute for Theoreti- cal Physics (KITP), which was supported in part by grants NSF PHY- 748958 and PHY-2309133 to KITP. AS thanks Y. Bahri and C. Pehlevan for their patience and willingness to listen to our early ideas at KITP.
Chen, Y., Zeng, Q., Ji, H., and Yang. Y. Skyformer: Re- model self-attention with gaussian kemnel and nyström method. Advances in Neural Information Processing Systems, 34:2122-2133 2021.
Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q., Mo- hiuddin, A., Kaiser, L., Belanger, D. B., Colwell, L. J., and Weller, A. Rethinking attention with performers. In International Conference on Learning Representations, net//orum? forum?
Impact Statement
This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
2021. URL / http:///eeeeeeeeeeeeeeeeeeeeeeea netp:/wwwwwweeeeeeeee.eeeeeeeeee id=Ja 6ZukORRI..
References
Dai, D., Sun, Y., Dong, L., Hao, Y., Ma, S., Sui, Z., and Wei, F. Why can gpt learn in-context? language models implic- itly perform gradient descent as meta-optimizers, 2023. URL https: arxiv. wrg/oorg////...oooooo.o 2212. 10559.
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-context learning. Advances in Neural Infor- mation Processing Systems, 36:45614 2023.
Demireigil, M., Heusel, J., Löwe, M., Upgang, S., and Vermet, F. On a model associative memory with huge storage capacity. Journal of Statistical Physics, 168:288- 299, 2017.
Akyürek, E., Schuurmans, D., Andreas, J., Ma. T., and Zhou, D. What learning algorithm is in-context learn- ing? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. URL https:////..........tt / /pwp//peeeeeeeeeeo/eeeeeee htt.eeeeeeeeeeeeeeeeeeeeeeeeeeee 1d=0g0X4HByN41.
Devlin, J., Chang. M.-W., Lee, K., and Toutanova, K. Bert: Pre-trad Pre-training of deep bidirectional transformers for lan- guage understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, Volume 1 (Long and Short Papers), pp. 4171-4186, 2019.
Albergo, M. S. and Vanden- Fijnde, E. Building nor- malizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL 15571.
org/ abs/
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations, 2021. URL https:// openreview net forumoiddeeeoe (errumiiddiicbaaNTT..
Amari, S.-I. Learning patterns and pattern sequences by self- organizing nets of threshold elements. IEEE Transactions on computers, 100(1 197- 1972.
Ambrogioni, L. In search of dispersed memories: Gen- erative diffusion models are associative memory net- works. Entropy, 26(5), 2024. ISSN 1099-4300. doi: 10. 3390/e26050381. URL https: /wwp///wwwww..///////////// htt///wwmmm..... 0000444040////33////
Fischer, K. H. and Hertz, J. A. Spin Glasses. Cambridge University Press, 1993.
wwwwwwpp...c....wwwwwwwwwwwwwwww wwwpmmppp.oo.. com/
Garg, S., Tsipras, D.,, Liang, P. S., and Valiant, G. What can transformers learn in-context? a case study of simple function classes. In Advances in Neural Information
Amit, D. J., Gutfreund, H., and Sompolinsky, H. Spin- glass models of neural networks. Physical Review A,

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Processing Systems, volume 35, pp. 30583-30598, 2022. URL htttpss//axxiv..orrgggoooo
Krotov, D. and Hopfield, J. Large associative mem- ory problem in neurobiology and machine learning. In International Conference on Learning Representations, 2021. URL https: opppneeeeeeeeeeeeeeeeee net forum?
01066.
Ghio, D., Dandi, Y., Krzakala, F., and Zdehorová, L. Sam- pling with flows, diffusion, and autoregressive neural networks from a spin-glass perspective. Proceedings of the National Academy of Sciences, 121(27))e2311810121, 2024.
id=x4y_ 10x0__0000h..-
Little, W. A. The existence of persistent states in the brain.
Mathematical biosciences, 191-2):101-120, 1974.
Loève, M. Probability theory i. Graduate Texts in Mathe- matics, 1977.
Gradshteyn, S. and Ryzhik, M. Table of Integrals, Series, and Products. Elsevier/Academic Press, Amsterdam, seventh edition, 2007.
Lucibello, C. and Mézard, M. Exponential capacity of dense associative memories. Phys. Rev. Lett., 132: 077301, Feb 2024. doi: 10. 10.1103/Phsseevl.et 132. 077301. URL peeetps//linng aps wtt.:ea..s.............o..o wttpppooogooooooo...eoo.oooooooe org/doi/11/ 1103 / wtwe//eysseeeetteeeeeeeeeeeeeeee 132. 077301.
Hoeffding, W. Probability inequalities for sums of bounded random variables. The collected works of Wassily Hoeffd- ing, pp. 409-426, 1994.
Hoover, B., Liang, Y., Pham, B., Panda, R., Strobelt, H., Chau, D. H., Zaki, M. J., and Krotov, D. Energy transformer. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: !/openreview net /fr//forumm d=MMWWNEXX9S.
Millidge, B., Salvatori, T., Song, Y., Lukasiewicz, T,, and Bogacz, R. Universal hopfield networks: A general frame- work for single-shot associative memory models. In In- ternational Conference on Machine Learning, pp. 15561-
15583. PMLR, 2022.
Hoover, B., Chau, D. H., Strobelt, H., Ram, P., and Krotov, D. Dense associative memory through the lens of random features. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024a.
Nakano, K. Associatron-a model of associative memory. IEEE Transactions on Systems, Man, and Cybernetics, 2: 380-388, 1972.
Pham, B., Raya, G., Negri, M., Zaki, M. J., Ambrogioni, L., and Krotov, D. Memorization to generalization: The emergence of diffusion models from associative memory. In NeurIPS 2024 Workshop on Scientific Methods for Understanding Deep Learning. 2024.
Hoover, B., Strobelt, H., Krotov, D., Hoffman, J., Kira, Z.,, and Chau, D. H. Memory in plain sight: Surveying the uncanny resemblances of associative memories and dif- fusion models, 2024b. URL httpt::////xxitt org/ abs/2309, 16750.
Ramsauer, H., Schäfl. B., Lehner, J., Scidl. P., Widrich, M., Gruber, L.,, Holzleitner, M., Adler, T., Kreil, D. P., Kopp, M. K., Klambauer, G., Brandstetter, J., and Hochreiter, S. Hopfield networks is all you need. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: http:::s//ppeeeevieeee nwt. forum? id-tI8 1i-t188RniIic..
Hopfield, J. J. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the National Academy of Sciences of the United States of America, 2554)2255-- 1982 ISSN 00278424. doi: 10.10733pnas.79..o 79.8.2554.
Hu, J. Y.-C., Yang, D., Wu, D., Xu, C., Chen, B.-Y., and Liu, H. On sparse modern hopfield model. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, 2023.
Reddy, G. The mechanistic basis of data dependence and abrupt learning in in-context classification task. In The Twelfth International Conference on Learning Represen- tations, 2024. URL https: /ppprevieeeeeeeeeeoe/eeeee net/ forum ? i=diiddf6
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: fast autoregressive transformers with linear attention. In Proceedings of the 37th Inter- national Conference on Machine Learning, ICML'20. JML..org, 2020.
Rigollet, and Hütter, J.-C. High-dimension statistics. arXiv preprint arXiv:2310. 19244, 2023.
Santos, S. J. D.,, Niculae, V., Menamee, D. C., and Mar- tins, A. Sparse and structured hopfield networks. In Proceedings of the 41st International Conference on Mu- chine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 43368 43388. PMLR, 21-27 Jul 2024. URL httss:/ http:////pcceeeeeeeeee mlr press/ v235/ sasstssns2244 4a html.
Krotov, D. A new frontier for hopfield networks. Nature Reviews Physics, 5(7):36 367, 2023.
Krotov, D. and Hopfield, J. Dense associative memory for pattern recognition. In Advances in Neural Information Processing Systems, volume 29, 2016.
10

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Shen, L., Mishra, A., and Khashabi, D. Position: Do pre- trained transformers learn in-context by gradient descent? In Proceedings of the s1 International Conference Machine Learning. volume 235 of Proceedings of Ma- chine Learning Research, pp. 44712-44740. PMLR, 21- 27 Jul 2024. URL https: press/v235 4d html.
/ /ttp://dooc.e.e...ooeo/o mlr..
Smart, M. and Zilman, A. On the mapping between hop- field networks and restricted boltzmann machines. In- ternational Conference on Learning Representations, 2021. URL htttp:://peenteiieeet id=RGJbergVIo0.
forum?
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint 2rXiv:2302.13977, 2023.
Trockman, A. and Kolter, J. Z. Mimetic initialization of self- attention layers. Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org. 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Łukasz Kaiser, and Polosukhin, I. At- tention is all you need. In Advances in Neural Information Processing Systems, volume 201e-December, 2017.
Von Oswald, J., Niklasson, E., Randazzo, E.. Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp.
35151-35174. PMLR, 2023.
Wu, D., Hu, J. Y.-C., Hsiao, T.-Y,, and Liu, H. Uniform memory retrieval with larger capacity for modern hop- field models. In Proceedings of the 41st International Conference on Machine Learning, ICML'24. JMLR.org. 2024a.
Wu, D., Hu, J. Y.C., Li, W., Chen, B.-Y., and Liu, H. Stanhop: Sparse tandem hopfield model for memory- enhanced time series prediction. In The Twelfth Interna- tional Conference Learning Representations, 2024b. URL httph:tth///cmmm / //aaiiiiiii////i////aaa///aaaeva or..
Zhang, R., Frei, S., and Bartlett, P. L. Trained transform- ers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024. URL http: h/tmmllorgg/paaeeeeeeeeee wwtpoopppp./////////22. 1042. html.

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
A. Notation
A.I. Recurring notation
n ambient dimension of input tokens.
Tt E IR" the value of the 1-th random input token.
E = (X, in some way been corrupted. takes values (#1...., IL, #) E R=x(L+1), Note: while capital X or X; here denotes a random variable, in Section use X1.L or simply X to refer to the realized matrix input tokens.
the random variable input to the sequence model. The "tilde" indicates that the final token has
I context length = number of uncorrupted tokens.
d dimensionality of manifold S that are sampled from
N number of training pairs
A.2. Bayes posterior notation
(x) is task-dependent (the three scenarios considered here are introduced above).
(T) where ī = x + z. For sum of independent random variables, = X1 + X2, their pdf a convolution (y()) = Px, (x)px, (y z)dr. Thus:
Px(0) = / p2(x)px(x -2)dz
Cz / 6-1111//203 px (3 2)
where Cz = (2ma2))~1) -n/2 a constant.
(x r): This is simply
Pz(# x) Oze 118-~11 1202
Pxxx((1 [): By Bayes' theorem, this is
(xix(=1 [b) Pxxx (5 x)px
P& (ō)
|| x11 /20 pxx (x) ʃe 1/-=z' 13/20% x (xx ddx
Posterior mean:
Expx(x [X] x1 = -f 1 xx (x &)dr
[ x -18 10100/1///// /202 px (x dx
ʃe- ||x x11/211/203 px (x)
12

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
B. Bayes optimal predictors for square loss
B.1. Proof Proposition 1
Proof. Observe that
|| - 1(8)11)
= =Ex
E ExI || f(x)11? [X]
=Bx Exix [1|X - E[X x111 [X]
+ ||[[X X] - -1- 1(8)11)
2 Bx E ExI [|XX - -EXX I 1Xl [x]
= Ex Crro xx]]
Note the final line independent of f. This inequality becomes an equality when f(x)= = [E[X X].
C. Details of Bayes optimal denoising baselines for each case
C.1. Proof of Proposition 2
Proof. The linear denoising task is special case of the result in Proposition 1. Here, X is an isotropic Gaussian in restricted subspace,
21
-x) C(T)xx(Te
203
where C()) is a normalizing factor. The noise can be decomposed into parallel and perpendicular parts using the projection P onto S, i.e.
8 8 + 8 = PX + ( P)X.
so that
|=-8|
|1-- 12
11
20
20
Only the first factor matters for P (x it) since it depends on a. Then, for a E S, the linear subspace supporting px. dropping the independent 1 contribution,
1x 200
12
Px (x)
202
OC
20
03
||
o t Tn
X exp
2 3303
o6 -0
Thus, f(8)
X11
d 9. t
PX.
o7. 76-0z
Using X = + Z, +
P.,, and the independence X and Z
0
2 0
dg + oá
dofo3
E
1X = E
2X
+
27
0 oz
0
+
(o) 0z)
01 +
13

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
C.2. Proof of Proposition 3
Proof. In the nonlinear manifold denoising problem, focus on the case of lower dimensional spheres S (e.g. the circle S' C R²). For such manifolds, have
1e-8,1
202
x Px (z) dr
E[X X = d
12-2,
f 20
px (x)dd
J |e ctt.21)/02 x dS.
ʃ e(2,81)/08 dS,
We have used the fact that || - ænl12 = ||=1|² + ||
2(x, 2n) and that ||>|| is fixed on the sphere.
The integrals can be evaluated directly once the parameters are specified. S is d s-sphr of radius R, then the optimal predictor is again shrunk projection of onto S.
cos sin (d-1) d0 Sd e cos 0/03 c//
In 1
R
So Bllaa cos 0/02 sin (d 1) A d0
|+ R
18 2
02
I
R
I d=
where we used identities involving (y), modified Bessel function the first kind of order v (Gradshteyn & Ryzhik, 2007). The vector R #1 the point on S in the direction of æn. Fill
C.3. Proof of Proposition 4
Proof. For the clustering case involving isotropic Gaussian mixtures with parameters {Way (Ma: 3=)
11
11
2x
Ea
Wi Wal
x dx
B X =
T]
|x-21|
11s- 1? 20
e
202
Σ.
Wa Col
dx
where C₂ = (2ma²)) 2.
We can simplify this expression by completing the square in the exponent and using the fact that the integral of a Gaussian about its mean is zero. This yields
Eo Camama f exp(p ga) dx
E X 5)
E. wi Co. expl-9g. dx
where we have introduced
1
102 +
1 2(02 o2)
9a
2
2 2
|| ma 1 +
16 Mall?
with
02 f. + 0z Ha
ma
02 + +
14

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Doing the integrals and using the expressions for Ca Ml
n1/2 ) exp ( (- (2 0
|| ||ā-Mall (
CetBto3 th 03+02
Ea Wa
2(a2+02))
B X X = []]
Ea Wa 06
0z 202 n/2 exp
||3- 21 +
In the case that the center norms || are independent of a and variances o.a Oo. we have
03 Ea Walta exp
(B,
o
0 +00
E
X
0
00 + +0z E. Wa exp
(i.Ma 02
Note that the limit that do - 0, this becomes expressible by one-layer self-atiention, since one can simply replace the matrix of cluster centers M [M1... (.l.....p P4p implicit in the expression with the context X1.I itself,
Eo Wa eaelmu X1/02 Ho
E X
Ea Wi eae
D. Additional details on attention layers and softmax expansion
D.1. Standard self-attention
Given sequence of Luqq imput tokens z; E R" represented matrix E RaxLo, standard self-attention defines query, key, and value matrices
K = WKX, WoX, V = Wy.X
(A.1)
where WK, WQ E Ruam x and Wy E R Xn The softmax self-attention map (Vaswani al., 2017) then
Attn(x, X, Wv. W Vsoftmax(K3 (FQ)= R
(A.2)
On merging WK. Wi into WKQ = Wo: The simplification WKQ = W (made here and elsewhere) is general only when aa 2 n; in that case, the product WxQ can have rank n and thus it reasonable to work with the combined matrix. On the other hand, if nann < n, then the rank of their product is at most Namn and thus there are matrices in *xx that cannot be expressed W k W Q. A similar point can be made about W Pv. We note that while aatt may be used in practical settings, one often also multiple heads which when concatenated could be (roughly) viewed as single higher-rank head.
We will also use the simplest version of linear attention (Katharopoulos et al., 2020),
Atto X, Wv, WFWO)=
V(KTQ) E R Rua
(A.3)
Leey
D.2. Minimal transformer architecture for denoising
We now consider simplified one-layer linear transformer in term of our variable E = (X1.,, taking values in (1*((++)) and start with the linear transformer which still has sufficient expressive power to capture our finite sample approximation to the Bayes optimal answer in the linear case. Inspired by Zhang et al. (2024), we define
1 Wry EM F.T W
Attm E,, Wrv, WKQ) =
(A.4)
L
taking values in (L+1),)), The additional aspect compared to the last subsection is the masking matrix Mun ((++))(L++)) which is of the form
1t
01x1
Mun
(A.5)
01xL
0
15

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval preventing Wpy X from being added the output.
Note that this more detailed expression is equivalent to the form used in the main text.
X Fism (E,0) ww XI.L W KKQ
With learnable weights WKQ. Wrv E [Rxx abbreviated by we define
F(E,0) 0) = [Attmiin (LE, Wpv, WKQ)).L+1.
(A.6)
Note that, when Wey aIn, WxQ BI,, and a. oā+o3* F(E,0) should approximate the Bayes optimal answer (X) L oc.
Similarly, could argue that the second problems, the d-dimesional spheres and the on -> 0 zero limit of the Gaussian mixtures could be addressed by the full softmax attention
Aut(E, Wpv, WKQ) = Wev Esoftmax (E" W + M)
(A.7)
taking values in [2x((++) where M E ((++))((++)) is masking matrix of the form
0Lx(L+1) [-x0)
M
(A.8)
once more, preventing the contribution of X value to the output. The function softmax(2) = i= हक (e po.r e E E.
applied column-wiss We then define
F(E,0) [Atn(E, Wpv, W WKO)la44
(A.9)
which is equivalent the simplified form used in the main text:
x = F(E,0) = WW X 11 WKQ
D.3. Proof of Theorem 3.1
Proof. Let the support of px be subset of a sphere, centered around the origin, of radius R. Then the function
Ele () Σt e(x.8)/02
+ X,e
(A.10)
1L cxX..8
Both the numerator 1 EE= (Xe and the denominator 1. Et (x are averages of independent and identically distributed bounded random variables. By the strong law of large numbers, as L - oc, the average vector in the numerator converges to almost surely to f eln,51)/o2 dpx (z) for each component, while the average in the denominator almost surely converges f cln,an)/o2 dpx(x), which positive. So, as + oc, the ratio Eq. A. converges almost surely to
f c(mmmi)) x dpx(x)
/ellaail/oz dpx (x)
which is the Bayes optimal answer (T) for all 5 ER.
E. Further discussion of convergence rates as L + oo and the dependence on dimensions
Our analysis primarily focused on the asymptotic behavior as - oo using the strong law of large numbers, which just requires the mean to exist (Loève, 1977). However, in the linear example, our tokens are Gaussian, and in the two nonlinear cases they are bounded. Intuitively, expect error O( } In fact, we can give precise results of the form that the probability of the difference between the empirical sum for the ideal weights departing from the expectation by less
16

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
than V (d,1n Interestingly, these bounds depend on d, the dimension spanned by the tokens, not the ambient dimension n.
I. is greater than 1 - 6. The function C' of the query vector and the function f depend on the problem.
As mentioned before, the results of the previous paragraph refer to the convergence of the finite sample attention expressions for ideal weights, namely those corresponding to Bayes optimal answer. There is second source of enror associated with finite sample estimation of weights, which should also get small L becomes large. Once more the expectation is that the weights are known to eтror O ) for well-converged training procedures, although this is more difficult to guarantee or quantify analytically. Overall we expect the loss (MSE) to go down inversely with some power of L. Fig. 4(a) provides some empirical evidence for this relationship, showing how performance improves with increasing context length.
Notice that the one-layer transformer output is a linear combination of the uncorrupted samples. Hence, if the distribution Px is supported by d-dimensional linear subspace, the estimate X is also that subspace. We can therefore look at convergence restricted to the supporting subspace. Therefore, is the dimensionality the supporting subspace that matters.
Let d-dimensional vector space V be a linear subspace of IR.. We define the maximum norm for V with respect some orthonormal basis (u:) in V as ||=||>,, max4e(1.. ((t,,)) a)) for any æ E V. The conventional maximum norm for R",, of course, is defined with respect to the standard orthonormal basis {e,}i=ı- Since (1,, S || for all i,
d
|
Σ (1),))))
dI ||0||| S Vdll=elloo..
Then, for any x E V C R", ||<|b> S Vdl V=llooo, since (er, ej)l & || < Vdll=llooov for all j E {l...,n). Thus, controlling component-wisg wise error in any orthonormal basis V controls component- wise error in R", in an r-indeppndeen but d-dependent manner, In the following, we give flavor of how we can analyze finite sample estimate errors in V. The maximum norm = is to be understood as Vl for some orthonormal basis choice. Here is the result relevant the linear case described Subsubsection 2.2.1.
Proposition 5. Let X, sid N(0, of Ta), 1.... and let ह
1 Ete ogL
X,XT. Then, for any E (0, 1),
d In d In(
}
Pr | øllo0 < 0115112 max
> >1-8
for some C >>0 O.
Proof. We start by bounding the maximum norm of the difference,
|
alloc S -|لt all2 s 1I Tallop|<=11e.
where || llop is the operator norm.
can be shown that, for any 5 E (0,
{ ( d In(
d In(3)
Pr | lallop < C max
> =1-
for some C > 0 (Rigollet & Hütter, 2023). Combining with the first bound, we get our result.
As to the nonlinear cases, the key result of Theorem 3.1 is the convergence of the numerator Et. X (xXe(X
f ( dx)(x) and the denominator E (X to (x2(02))
I (m..i1))/02 2pxx()))))
In the following, we assume that the support of px is inside a vector space V whose dimension we denote by d (instead of +++ 1, as in the sphere problem). In addition, we refer to the projection of the query on V by ā E V, instead of žn. As usual, the maximum norm in V with respect to some orthonormal basis choice
17

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Proposition 6. Let X std Then, for any 8 E (0, 1),
~ px and X & R = 1.....,. L.
Pr
ex, 2//02 - -
sinh
R||
0z
)
2
In
8
2
2 1 8
t=1
and
Pr
Xze (x.,8)/02 - B[xelx.a/e
< Re
Blll2
2d In
2 - 8.
Proof. We provide the sketch of our proof here, the key ingredient of which is the Hoeffding inequality (Hoeffding, 1994).
For the average + El= e(Xe, each term in the sum is bounded above and below by e inequality leads to
So, the Hoeffding
2Le
Lca
Pr
e (X, 10x -Bee
2 c
2cxp
exxp
= exx
R||2|12) 02
RI
2 sinh?
2
exp
exp
02
Setting d = 2 exp
we get sinh ||| In (2), which gives our first probabilistic inequality.
2innh
R
For each component of the vector average + ELe X5e () the terms in the sum are bounded above and below Rl
by + 2
We use similar arguments involving the Hoeffding inequality, combined with the union bound over all d
coordinates
Le
Pr Σ
- E[Xe(Xx8/003) 2 S 2d
2R222 exp
22112112
Once more, setting the RHS to and solving for get our second probabilistic inequality.
F. Limiting behaviors of the softmax function and softmax attention For small argument
Taylor expansion of the softmax function at zero gives
1 = (LL B0 + 0(33))
softmax
where -= -, (z, +0())) = L(1 + 88 + +(83)) is normalizing factor, with 7 + E, Us The notation 1,
stands for an L-dimensional vector of ones.
Thus, we have
Lemma F.1 (Small argument expansion of softmax), As 3 + 0,
softmax(BB)) = L1 + 36 + (L B0 0(B()) 1 1 (1r + B(v -((v-8)++((33))
1
F.1. Proof of Proposition 3.2
Proof.
F E, 1 (WPP)) = 1 WPvXXy W X).
18

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Using Lemma F.I, O.
1 (I
F F(E (=Der,ewko)) -WPVXXL I 11 e(XT 1 [ X(WKoX))LL) ( +0(23)
1
4 WPvX - Wev Xxxxxxx)))x +
1
(A.11)
where X = Et X₂ is the empirical mean and the notation 1, emphasizes that it is a column vector of ones with
dimension L.
For large argument
As B + oc, the softmax function simply selects the maximum over its inputs (as long as the the maximum is unique):
softmax(3u) if arg max, Uy>
otherwise.
In this case, all attention weight is given single element, and the others effectively ignored.
G. MSE Loss landscape for scaled identity weights
(a)
MSE loss landscape for Fig. (Case
(b)
MSE landaccccpeee for Fig. 3 (Case 3)
1.28
hour id0, Lroj
hour betc llhwaryf
0.48
Scalez We
0.32
LOGG
0.64
0.24
0.48
-1
0.32
0.08
-10.0
-75
-5.0
-25
0.0
2.5
5.0
7.5
10.0
-100
-75
-5.0
-2.5
0.0
2.5
5.0
7.5
10.0
Stalel: Mao
Scalel: Wao
Figure 6. Loss landscape corresponding to Case 2 and Case of Fig. 3. The MSE is numerically evaluated by assuming scaled identity weights WKQ = BT,, (x-axis) and Wpv oT, (y-axis) and scanning over a 50 50 grid. The green point corresponds the heuristic minimizer identified from the posterior mean. In Case it is exact, while in case 3 is an approximation that neglects the residual term (see Proposition 4). The orange point corresponds to learned weights displayed Fig. 3(b), while the white point corresponds to the numerically identified minimum from this 2D scan. These can fluctuate due to the finite context (L = 500) and sampling (N 800 here). both panels, is apparent that the trained weights the heuristic estimator co-occur broad valley (contour) the loss landscape.
The loss landscapes Fig. exhibit large, low-cost valleys with roughly hyperbolic structure that is especially apparent in Case 2. This indicates multiplicative tradeoff the scales of VQ and Wpv. which suggests that linear attention might be applicable here well. For completeness, Figure 7 shows linear attention performance for both cases, demonstrating that it performs quite similarly to softmax for sub-sphere denoising, but less well in the Gaussian mixtures case.
H. Structured optimal weights under prompt transformation
We find that one-layer transformers can learn to undo arbitrary invertible coordinate transformations that warp the denoising tasks. Focusing on the subspace denoising case, suppose each prompt is transformed by fixed invertible square matrix
19

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
(a)
Case 2: Nonlinear manifolds
Case 3: Gaussian mixtures
=31, 8= ==1
p=
=000
0.06
softmax (train/test)
0.08
linear (train/test) A
MSB/nn
0.05
MSB/n
fas R
0.04
0.04
0.02
200
0
50
100
150
Epoch
Epoch
(b)
Finel weights: Inear
Final weights softmax
Final weights linear
Final Final weights softmax
Wioe
Wrv
Wxo
Wko
Wrv
3.0
4.0
Figure 7. Linear attention performance for Cases 2 and Additional empirical results for nonlinear manifolds case (left) and the Gaussian mixtures case (right). (a) Loss dynamies for randomly initialized softmax and linear attention layers. Solid lines represent the average loss over six seeds, with shaded area indicating the range. Training details and parameters follow Fig. 3(a). (b) Representative final attention weights for each layer.
A, i.e. E = (X + +' (AXI.,, AT). If the target remains Ft+11 in the untransformed space, then the optimal attention weights are no longer diagonal, but instead take structured form determined by the transformation matrix:
Wpy= = a WKQ = B(AA")
(A.12)
where a& oz+03 before.
(b)
(c)
Linear attention
(d)
Softmax attention
Theory Wao B(AA7
Learned
Learned Weo
8 0.421
z 0.193
3-
0.50
(a)
0.25
6-
6-
(transform)
1.00
9-
9-
-025
0.50
12-
-0.50
0.25
15
15
12
12
0.500
Theory Way
Learned Way
Learned Way
0075
is 0,792
1.61
-1.00
0.5
3-
0.0
0-
12
-0. 12
15
9
12
15
12
Figure (a) Example transformation used to globally alter the in-context denoising prompts. (b) Structure of the optimal attention weights for this transformed subspace-denoising task. (c.d) Empirically, we find that both lincar attention and softmax attention layers able to leam these structured targets, but with distinct scalings a, 3. Final weights after 500 epochs using Adam, random initializations, and context length L 500; other parameters follow Fig. 3(a).
Notably, we find that both the linear and softmax attention layers are able to learn these structures; see Fig. 8 for an example.
20

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
We the same basic training procedure as the limiting case of A / (no additional coordinate transformation) assumed throughout the main text.
Suppose we still work with transformed coordinates Y, = AX, and Y = AX, but now intend to retrieve YI.+1 = 1 in the new coordinate space (rather than XL+1 as above). In this case, we would be dealing with variables with covariance matrices Σ x AA".. We would need weight matrices that are not simply proportional to identity to deal with the covariance structure, This is also the case for in-context learning linear functions when the input has an anisotropic covariance matrix (Zhang al., 2024; Ahn et al., 2023). Recall in the original setting, we had the sample covariance E[XX"] Σx of and noise Σg = T²T, leading to the estimator, Eq. (10): X ( D X). Here, the sample covariance Σy = oAAPAT. and the noise V = AZ has covariance Σv = o2 AAT. One can show the generalized solution is y = Σy (Ex + Ev) Thus, in the transformed coordinates, the denoising estimate is
1 y = (o +02)L
Σ Y(A Y. t=1
(A.13)
For the relationship of this denoising result in Y to energy models, as discussed in Section 4 and Subsection I.1, need modified energy E(YI,L, ||||| 21 E (A -Y, A s)² and preconditioner proportional to AAT..
I. Additional comments on the mapping from attention to associative memory models
L.1. Linear attention and traditional Hopfield model
We have considered a trained network with linear attention, relating the query X and the estimate of the target X, of the form
x= 18} = 7 Σ X) I.
(A.14)
with 7 = oz40z
With
1 S
E(X1,,, 8) 27
||/1/
XxX?))s
(A.15)
2T.
gradient descent iteration s(t+ stt 7 V V.E( (XE,X., s(t) gives
s( + 1) 7, Xx(x,, x())
making the one-step iteration our denoising operation.
We will call this energy function the Naive Spherical Hopfield model for the following reason. For random memory patterns XI.L, and the query denoting Ising spins 1 the so-called Hopfield energy is
Eloppieaaa XI:L, 8) = 2L
X,X? ss.
(A.16)
We could relax the Ising nature of the spins by letting s E R", with a constraint ||s||2 = n. This is the spherical model (Fischer & Hertz, 1993) since the spin vector & lives on a sphere. If we minimize this energy the optimal s would be aligned with the dominant eigenvector of the matrix _ X, (Fischer Hertz, 1993), and the model will not have retrieval phase (see Bollé al. (2003) for similar model that does). A soft-constrained variant can also be found in Section 3.3, Model C of Krotov & Hopfield (2021).
We could reformulate the optimization problem of minimizing the Hopfield energy, subject to ||s|| R²,
arg min max { SCR"
1
X.X )s 1(52 - -n-)]
2T.
21

---

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
The s-dependent part of the Lagrangian, with À replaced by by gives us the energy function in Eq. A.15 which have called the Naive Spherical Hopfield model.
1
1 s" X, s= x
2 (o0 [08 oZ) 1 2 xxxX) L
1
8(LILL, 8) 27
11s11² 11s11²
(A.17)
2L
For L much larger than n, 1, E X P.P, so its eigenvalues are either or are very close to o.. Hence, for large and øg > O, this quadratic function is very likely to be positive definite. One-step gradient descent brings down the d-dimensional linear subspace spanned by the patterns, but repeated gradient descent steps would take towards zero.
1.2. Remarks on the softmax attention case (mapping to dense associative memory networks)
Regarding the mapping discussed in the main text, we note that there is symmetry condition on the weights WK Wrv that is necessary for the softmax update be interpreted
gradient descent (i.e. conservative flow). general, flow ds/ di. (s) is conservative if it can be written the gradient of potential, i.e. f(s) (sV(s) for some V. For this to hold, the Jacobian of the dynamies (s) v.f must be symmetric.
The softmax layer studied in the main text is fs) Wpv X softmax(X" (KK)s.. We will denote z(s) X WxQ and g(s) = s=f)m both in B._. Then the Jacobian
J(s) Wrv X Og
Wrv (diag(g) - gg"") W
(A.18)
Ös
Observe that Y = X (diag(g) ") XT is symmetric (keeping in mind that g(s) depends on WkQ). The Jacobian symmetry requirement = J therefore places the following constraint on feasible W WxQ, Wpv:
Wpy Y w = WKQ Y WTy.
(A.19)
It is clear that this condition holds for the scaled identity attention weights discussed in the main text. Potentially, it could allow for more general weights that might arise from non-isotropic denoising tasks to be cast as gradient descent updates.
The mapping discussed in the main text involves discrete gradient descent steps, Eq. (17). In general, this update rule retains "residual" term in (t) if we choose a different descent step size 7 L a. Thus, taking K recurrent updates could be viewed as the depthwise propagation of query updates through a K-layer architecture if one were to tied weights. Analogous residual streams commonly utilized in more elaborate transformer architectures to help propagate information to downstream attention heads.
22


Reviews:
[{"review_id": "dFdtjpMI4w", "reviewer": "Reviewer_PjG9", "summary": "The paper considers in-context denoising as a fundamental task of attention (when applied in a prompt-conditioned, auto-regressive manner). When interpreted this way, there are clear connections to Dense Associative Memories (DAMs), and a model trained on a final-token denoising task can approach the optimal Bayesian prediction when knowledge of the true manifold is known in advance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer's thoughtful review describing our paper as 'a pleasure to read' with 'clear exposition and motivations.' We're encouraged by the recognition of our 'strongly defended connection of ICL to Dense AM' and we are particularly thankful that the reviewer found the paper to be of 'sufficient completeness, correctness, novelty, and quality to be accepted to ICML.'\n\n**Re: Strengths and Weaknesses**  \nWe thank the reviewer for their thoughtful feedback on the paper's limitations. While we acknowledge the controlled nature of our experimental setting, this design was deliberate to establish clear theoretical connections and optimal bounds. \n\n- **Synthetic data:** This choice was necessary to derive Bayes-optimal predictors as theoretical baselines. Our findings suggest that with sufficient context, transformers can learn to denoise even when the underlying distribution is not fully known a priori (Fig. 4).\n\n- **Single-layer architecture:** We focused on minimal architectures to isolate fundamental connections between attention and DAM in a novel ICL setting. Our results suggest deeper models might implement multiple gradient steps on a complex energy landscape, and it will be very interesting to explore how this interacts with related but fundamentally distinct efforts including the Energy Transformer (mentioned below by the reviewer.)\n \n- **Identity weights:** While our tasks led to scaled identity weights, we've begun exploring settings with non-isotropic covariance structures that require more complex parameterizations for $W_{KQ}$.\n\n- **Single token denoising:** We deliberately focused on denoising a single token to establish the clearest connection to DAM. Extending to multiple tokens and incorporating sequential information represents important future work.\n\nWe view these limitations not as weaknesses of our approach but as exciting avenues for future research that can build upon the theoretical framework we've established.\n\n**Re: References**  \nWe regret missing the Energy Transformer (ET) citation and are grateful to the reviewer for highlighting it. We will certainly discuss it.\n\nBriefly, while both works explore connections between transformers and DAM networks, their approaches differ fundamentally. Our work examines in-context denoising and demonstrates why a single-step update can be optimal from a Bayesian perspective. In contrast, the ET tackles dataset-specific tasks via multiple iterations using a specialized but elegant design based on Hopfield models. The approaches are complementary but reversed: ET begins with an energy function to construct its architecture, whereas our work shows that standard attention mechanisms naturally learn to perform a gradient step on a context-aware DAM energy landscape. By starting from a vanilla attention layer in a minimal setting -- where true distributions are known and theoretical insights can be gleaned -- our contribution concretely explains why the Ramsauer correspondence involves only one iteration of energy minimization rather than many.\n\nWe likewise will cite Ambrogioni alongside Hoover et al. in the revision (thank you!). \n\n**Re: Typo in Appendix E**  \nWe reviewed and did not find a typo there, but we will change $\\mathbb{1}$ to $\\mathbb{1}_L$ for notational consistency. \n\n**Re: Questions for Authors**\n- **(Q1)** The blue tokens in Fig. 1 represent the $L$ \"pure\" tokens sampled from $p_{X}(x)$. These form the context portion of the prompt, with the query being the corrupted $(L+1)^\\textrm{th}$ token. Our formulation deliberately focuses on tasks where token order isn't relevant to maximize clarity in the transformer-DAM connection without positional embedding complexities. Extending this to sequence-dependent tasks (e.g. trajectory inference in dynamical systems) represents an exciting direction for future work.  \n- **(Q2a)** The apparent outperformance in Fig. 3a (Case 2) represents mild training set overfitting. Test set performance remains bounded by the theoretical optimum.  \n- **(Q2b)** While theory indicated softmax attention was appropriate for nonlinear manifold and mixture cases, linear attention also performs well. We can add a supplementary figure showing these additional results.  \n- **(Q3)** This interpretation is insightful. The scaled identity weights suggest optimal operation attracts the noisy token toward context tokens themselves rather than transformed versions. This aligns with our finding that the network performs gradient descent on a DAM energy landscape where context tokens serve as stored patterns. We've also explored minimal generalizations where $W_{KQ}$ that are not scaled identity arise from more elaborate covariance structures.\n\n--\n\nWe thank the reviewer again for their valuable insights and constructive feedback towards improving the work. We hope our responses have addressed their questions and concerns while further clarifying the contributions of our work.", "claims_and_evidence": "Claims are supported by simple experiments and theoretical analyses.", "methods_and_evaluation_criteria": "The methods and evaluations are sound. It would be nice to see experiments on real data, not only toy data for which we know the underlying distribution $p_\\mathcal{X}$, see weaknesses.", "theoretical_claims": "I did not review the proofs in the appendix for correctness, but I saw no glaring issues in the propositions of the the main paper.", "experimental_designs_or_analyses": "The experimental designs are sufficient to test the claims in this work.", "relation_to_broader_scientific_literature": "This paper contributes to the growing body of evidence that treats Transformers (specifically, the attention operation of Transformers) as a form of Associative Memory. It also builds on a growing interest in formalizing ICL as a form of memory lookup given the previous tokens in the sequence.", "essential_references_not_discussed": "A very related but undiscussed work is the [Energy Transformer](https://arxiv.org/abs/2302.07253), which formalizes the entire Transformer block as a kind of DAM. Specifically, the paper contains several statements that are explicitly addressed by the Energy Transformer:\n\n> [L068-070] why does the Ramsauer et al. (2021) correspondence involve only one iteration of Hopfield energy minimization and not many?\n\n> [L397-399] We have also noted preliminary connections between our work and other architectural features of modern transformers, namely layer normalization and residual streams, which warrant further study.\n\n> [L949-952] taking K recurrent updates could be viewed as the depthwise propagation of query updates through a K-layer architecture if one were to use tied weights. Analogous residual streams are commonly utilized in more elaborate transformer architectures to help propagate information to downstream attention heads.\n\nAdditionally, the work of [Ambrogioni](https://arxiv.org/abs/2309.17290) should be mentioned alongside [Hoover et al.](https://arxiv.org/abs/2309.16750) on [L412]", "other_strengths_and_weaknesses": "### Strengths\n- **Well written**. Clear exposition and motivations makes this paper a pleasure to read. The figures are generally complete and self-descriptive (though a few details can be clarified, see questions).\n- **Strongly defended connection of ICL to Dense AM**. The idea that ICL (and indeed, much of the operation of transformers) can be studied through the perspective of associative memory and gradient descent on the energy landscape has been growing lately, and this paper continues that line of works by thoroughly describing the \"next token denoising\" task as an energy minimization problem that approaches a Bayes optimal answer.\n\n\n### Weaknesses\n- **Limited scalability and connection to real transformers**. The ideas in the paper are good, but the sandbox in which they test the ideas is small. The experiments are on toy data where the underlying true distribution is known (*how can this work be extended when the true underlying distribution is not known?*), and they only use a single attention head for a single transformer layer. They also do not consider the case where \"interesting parameterizations\" of $W_{KQ}$ and $W_{PV}$ may be necessary to solve the task at hand (in their experiments, the training tasks cause these parameters to converge almost to identity). Additionally, the update rule they describe discusses only the denoising of the $L+1$ token, when the power of transformers today comes from their parallelizability, allowing them to evolve all tokens in the input simultaneously. Also, as far as I can conclude, all experiments and propositions treat the input tokens as a set of tokens instead of a sequence of tokens with positional information, as is needed for real tasks.\n\n### Summary\nI like the paper and I think it is well written. The setting in which they choose to study the model, though quite limited, is appropriate for evaluating the claims in the paper. It would be nice for the authors to include some real-data experiments to make it easier for others to build off the work, but I believe the content in this paper is of sufficient completeness, correctness, novelty, and quality to be accepted to ICML. As always, I am happy to increase my score during the review process if the authors can clarify my questions and concerns", "other_comments_or_suggestions": "**Typos**\n\n- Is there an error in Appendix E's taylor expansion of the softmax [L814] and [L822]? It seems like a math symbol did not render correctly.", "questions_for_authors": "Q1: Fig 1 is very helpful, thank you. But it is still unclear what is included in the prompt? Are all blue dots included in the prompt? It seems that the goal of ICL is to project the noisy \"test token\" ($\\tilde{x}$}) back onto the manifold encoded by the context prompt $E$, but I am still not clear if there is any \"sequential information\" (e.g., token positions $[1,\\ldots,N]$) in the prompt?\n\nQ2: Fig3a Case 2 -- why is the bayes optimal predictor have higher loss than the softmax train? Is this \"overfitting\" to the training data? Also, is there a reason you only test linear attention on the linear manifold case instead of including it in the non-linear and GMM results?\n\nQ3: I find it very interesting that the optimal weights for a single transformer layer/head can be expressed as scaled identity matrices. Could one interpretation of this idea be that the optimal predictor for this task is trying to \"attract\" the noisy token to the prompt tokens themselves (and not some transformed version of the prompt tokens)? In which case, it is of no surprise that the optimal parameterization leads to identity matrices.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D8R6eXhCSo", "reviewer": "Reviewer_RPZT", "summary": "The paper explores a link between one-layer transformer attention and associative memory retrieval.\nIt frames an in-context denoising setting with three synthetic data scenarios: linear subspaces, points on a sphere, and Gaussian mixture clusters.\nIt derives Bayes-optimal predictors in each setting and shows that a single attention update approximates these optimal solutions under small-scale assumptions. \nIt interprets the one-layer transformer as one-step gradient descent on a Hopfield-like energy. Iterative updates converge to stored attractors and degrade denoising performance. \nThe paper supports these claims with synthetic experiments, where trained transformers match the predicted optimal performance.\n It concludes that a one-step attention update implements Bayes-optimal denoising in specific synthetic contexts, connecting transformer attention with associative memory beyond exact pattern retrieval.\n\n## update after rebuttal", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's thoughtful engagement with our work, particularly their recognition of our work's **originality** in studying in-context denoising, its **significance** in demonstrating transformers performing optimal denoising updates beyond exact pattern retrieval, and its **clarity** in explaining the proposed approach. We believe these findings establish fundamental connections between attention architectures and Dense Associative Memory networks through Bayes-optimal denoising -- providing insights that extend beyond incremental advances while also establishing a bridge to the study of in-context learning.\n\n**Re: Positioning and Contribution**  \nOur paper makes a specific contribution beyond prior work connecting transformers and Hopfield networks: establishing that for certain denoising tasks, a single attention step can be optimal from a Bayesian perspective. This provides theoretical justification for why the \"one-step correspondence\" noted by Ramsauer et al. (2020) works effectively. Our key insight is that in-context denoising connects associative memory retrieval, attention, and ICL in a common Bayesian framework. We demonstrate that:\n- For certain denoising tasks, one-layer transformers can represent Bayes-optimal predictors  \n- Trained attention layers correspond to single gradient descent steps on a context-dependent energy landscape  \n- This single step can outperform multiple iterations, challenging conventional associative memory paradigms. This extends beyond the common view that associative memories must converge (often through multiple iterations) to be effective.\n\n**Re: Relation To Broader Scientific Literature**  \nWe thank the reviewer for highlighting important related works. We will incorporate these references, particularly Santos et al. (ICML 2024), Wu et al. (ICML 2024), Wu et al. (ICLR 2024), and Hu et al. (NeurIPS 2023) on sparse modern Hopfield networks. We'll also discuss connections to implicit gradient descent literature mentioned by the reviewer. These works complement our findings while our contribution on optimal one-step denoising adds a new perspective.\n\n**Re: Equation referencing**  \nWe are open to doing so (especially if it is an ICML guideline), but regard it as a style choice, as readers may mention specific equations even if they aren't referenced within the text. \n\n**Re: Strengths and Weaknesses**  \nRegarding our energy function differing from standard DAM: This modification was deliberate and essential for our denoising tasks. The Lagrange multiplier term handles continuous state spaces more naturally while maintaining core associative memory dynamics, similar to regularization approaches in other energy-based models.\n\nAs for additional iterations degrading performance: This is a key finding, not a limitation. Traditional DAMs aim to retrieve patterns exactly, but in-context denoising requires blending information from context tokens with the corrupted query. Figure 5 demonstrates why a single step is optimal for certain denoising tasks, providing theoretical insight into when one-step updates outperform iterative approaches.\n\n**Re: Why not higher score?**  \nWe respectfully suggest reconsidering our contribution:\n\n- Our paper provides fundamental insights into why single-step attention can be optimal for denoising tasks, moving beyond merely refining known connections and establishing the ICL connection.\n- Our focus on simplified synthetic scenarios establishes rigorous theoretical foundations, a common approach in theoretical ML papers.\n- The impact extends beyond specific scenarios studied. By establishing when single-step updates can be optimal, we provide insights informing design and analysis of more complex transformer architectures.\n\n**Re: Questions for Authors**  \n- **(1) & (3):** We've shown that for elementary denoising tasks, the optimal solution is a one-layer, single-head attention architecture viewable through the DAM lens. For more complex tasks, multi-layer/multi-head architectures likely become necessary for Bayes-optimality -- an exciting direction for future research.\n- **(2):** This important question motivated Section 4. Traditional views suggest convergence via iterative attention updates are beneficial. Fig. 5 shows how one-step outperforms multiple iterations in our denoising task. Multiple steps push the query toward fixed points that depend on random sampling, causing the system to 'forget' query information - sub-optimal for denoising. This insight reinforces the utility of the single attention step: it strikes an optimal balance between query and context information. We will discuss this in the revised text.\n\n--\n\nThe reviewer's thoughtful feedback and insightful questions are much appreciated. We believe that addressing these points will strengthen the paper while maintaining its core contribution: solidifying the connection between attention mechanisms and associative memory through a novel in-context denoising setting.", "claims_and_evidence": "* **Claim: Single-step attention approximates Bayes-optimal denoisers in specific synthetic settings**\n    * Evidence: The authors derive exact optimal predictors under Gaussian mixtures, linear subspaces, and spherical assumptions. They show near-optimal mean squared error in experiments.\n* **Claim: One-layer transformers match a one-step gradient update on a Hopfield-like energy**\n    * Evidence: The paper presents an energy formulation and demonstrates that learned attention weights align with the gradient of that energy. Experimental training runs confirm that a single attention update outperforms multiple updates in denoising.", "methods_and_evaluation_criteria": "-", "theoretical_claims": "All key claims are backed by formal derivations (with detailed proofs in the appendices). The math seems sound. For example, Proposition 2–4 give closed-form optimal denoisers for each case, and Proposition 3.2 analyzes the softmax attention in a small-argument regime. These ensure the theoretical claims hold under certain limits.\n\nThroughout, assumptions (e.g. isotropic distributions, large ambient dimension, small noise limits) are clearly stated, and the derivations use standard Bayesian inference and asymptotic analysis, which appear logically valid.", "experimental_designs_or_analyses": "**Conclusion:** The experimental design is well aligned with the theory. The analyses are correct for the stated goals. The limited scope leaves open questions about broader generalization.\n\n\n---\n\n**Synthetic Task Design:** The authors focus on three artificial data setups: linear subspaces, points on a sphere, and Gaussian mixture clusters. They construct these tasks to mirror the paper’s theoretical assumptions (e.g., isotropy, symmetry). This design is consistent with the paper’s goal of checking Bayes-optimality in simplified contexts.\n\n**Model Configuration:** They use one-layer transformers with either linear or softmax attention. They match each architecture to a theoretical predictor. The design supports the core claim that a single-step attention update aligns with the derived solutions.\n\n**Training Procedure:** They train these models on the synthetic tasks and measure mean squared error, checking how close each model’s performance is to the analytical Bayes-optimal denoiser. This procedure is valid. They compare different initializations and sometimes vary noise scales to confirm robustness.\n\n**Comparisons:** They contrast linear attention versus softmax attention. They show that both approaches achieve near-Bayes-optimal performance. They analyze how the learned attention weights resemble the predicted identity-like matrix in each scenario.\n\n**Limitations:** They do not test real-world data or more complex tasks. The synthetic setups are narrow in scope. I guess this is fine for a theoretical paper.", "relation_to_broader_scientific_literature": "### **Connection to Modern Hopfield Networks:** \n\nThe paper builds on the established link between transformer attention and modern Hopfield retrieval (Ramsauer et al., 2021). It interprets a single-layer transformer as a one-step gradient update on a Hopfield-like energy. Similar efforts in recent literature but not mentioned in this paper include\n-  [Santos et al. ICML 2024] Sparse and Structured Hopfield Networks https://arxiv.org/abs/2402.13725\n-  [Wu et al., ICML 2024] Uniform Memory Retrieval with Larger Capacity for Modern Hopfield Models https://arxiv.org/abs/2404.03827\n-  [Wu et al., ICLR 2024] STanHop: Sparse Tandem Hopfield Model for Memory-Enhanced Time Series Prediction https://arxiv.org/abs/2312.17346\n-  [Hu et al., NeurIPS 2023] On Sparse Modern Hopfield Model https://arxiv.org/abs/2309.12673\n\nThe paper’s one-step approach complements these ideas but does not discuss them.\n\n---\n\n### **Extension from Exact Retrieval to Denoising:**\n\nPrior Hopfield models often retrieve exact stored patterns, e.g., [Santos et al. ICML 2024]. The paper shifts toward Bayes-optimal denoising, which involves partial or mixed retrieval rather than convergence to a single pattern.\n\n---\n\n### **Broader Transformers and In-Context Learning:**\n\nThe paper adds to literature suggesting transformers perform implicit gradient updates in-context (for examples, see below). Prior works demonstrate in-context meta-learning [Dai et al., ACL 2023] or approximate Bayesian inference [Xie et al., ICLR 2022] , and this paper contributes a denoising viewpoint.\n\n- [Oswald et al., ICML 2023] Transformers learn in-context by gradient descent\n https://arxiv.org/abs/2212.07677\n- [Dai et al., ACL 2023] Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers https://arxiv.org/abs/2212.10559\n- [Bai et al., NeurIPS 2023] Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection https://arxiv.org/abs/2306.04637\n- [Ahn et al., NeurIPS 2023] Transformers learn to implement preconditioned gradient descent for in-context learning https://arxiv.org/abs/2306.00297\n- [Xie et al., ICLR 2022] An Explanation of In-context Learning as Implicit Bayesian Inference\n https://arxiv.org/abs/2111.02080", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "## Strengths \n\n- **Originality:**  \n  Introduces in-context denoising as a new setting, explicitly linking transformer attention to Hopfield memory. Extends previous views that only consider retrieval tasks.\n\n- **Significance:**  \n  Clearly demonstrates the significance of transformers performing optimal denoising updates beyond single-pattern retrieval, providing a formal connection between in-context attention and optimal memory retrieval.\n\n- **Clarity:**  \n  Clearly explains the proposed approach and its advantages, with logical connections between theoretical analysis and experiments.\n\n## Weaknesses  \n\n- **Positioning and Contribution:**  \n  The main concern is the unclear positioning of the paper’s contribution. Transformers implicitly simulate gradient-based updates in-context, as established by existing literature. Prior studies already demonstrate this perspective, though typically limited to single-token retrieval.\n\n    After some reflections, I suggest positioning this paper as extending the known \"transformers-as-implicit-gradients\" concept. Through 3 very specialized examples, this paper analyzes how attention interpolates among multiple context tokens to achieve a Bayes-optimal denoiser, rather than merely retrieving or copying one pattern. Specifically, this paper goes beyond prior studies by:\n\n  - Performs optimization blending multiple context tokens, achieving performance superior to single-token retrieval.\n  - Demonstrates Bayes-optimality explicitly in simplified synthetic tasks.\n  - Clarifies how mixed retrieval outperforms exact retrieval, providing rigor to the implicit-gradient viewpoint in multi-token scenarios.\n\n    This adds a conceptual and theoretical layer to the implicit GD story: it clarifies why partial or “mixed” retrieval can be optimal, whereas previous examples focused on simpler single-token or exact retrieval scenarios. \n\n- **Literature Completeness:**  \n  Several relevant references about the association between attention mechanisms and one-step associative memory retrieval are missing. Including these references improves the contextual completeness of the paper.\n\n- **Equation referencing:**  \n  Some equations have numbers without explicit references in the main text. Equations that are not explicitly cited should not have numbers.\n\n## Why not higher score? \n\n- The paper primarily refines known connections between transformers and Hopfield networks. The core novelty is modest, building on established frameworks rather than providing fundamental insights.\n\n- The theoretical analysis and experiments focus exclusively on simplified synthetic scenarios, limiting generalization to more realistic or practical applications.\n\n- The limited experimental and theoretical scope reduces the broader impact. Expanding experiments or theoretical analysis beyond synthetic settings significantly strengthens this work.", "other_comments_or_suggestions": "see above", "questions_for_authors": "1. whether the one-step Bayes-optimal property holds when data are more complex, structured, or high-dimensional (e.g., text or images)? for example, how read the corresponding attention be like? is it possible to be sth like really attention heads in practice?\n\n2. What happens if we are allowed to perform multiple iterative updates? Do additional steps help or hinder performance? \n\n3. Does the Bayes-optimality of one-step updates hold in 1-layer multi-head attention?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Xg6l48aU41", "reviewer": "Reviewer_kZ23", "summary": "This paper studies the in-context unsupervised denoising of data points in transformers. They show that single layer transformers with a single attention head are sufficient to learn this task and that standard training procedures from random initialization can converge to Bayes optimal solutions. Lastly, they provide connections between the learned transformer models on this task and Dense Associative Memory models, a modern variant of Hopfield networks which have high storage capacity.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's thoughtful review and their high evaluation of our paper (4/5). We agree that this is an important problem and are thus grateful for their positive assessment of our theoretical findings and empirical validations.\n\n**Re: Weaknesses**  \nThe reviewer notes that our paper primarily focuses on synthetic data. This was a deliberate design choice to establish clear theoretical connections between transformer attention and Dense Associative Memory (DAM) networks. By working with well-defined synthetic distributions, we could derive Bayes-optimal predictors as theoretical baselines against which to evaluate our transformer models.\n\nWe agree that extending these insights to more complex, real-world data distributions represents an important direction for future work. The current paper focuses on establishing foundational theoretical results and validating them empirically in controlled settings where optimality can be rigorously established. This approach allows us to identify the key mechanisms by which transformer attention implements optimal denoising, providing insights that can inform future work on more complex data distributions.\n\n**Re: References**  \nWe thank the reviewer for pointing us to the work of Boix-Adsera et al. This paper's finding that adding an identity matrix to the product $W_K^T W_Q$ can improve performance on template reasoning tasks (Observation 1.2) resonates with our work. We will include this reference in our revision (perhaps alongside Trockman & Kolter (2023) in the discussion) to highlight this potential connection and suggest avenues for future investigation.\n\n**Re: Specific Questions**  \n- **(Q1)** Regarding the convergence rate in $L$ as a function of dimension: we appreciate this question and will consider including a more detailed analysis of finite-sample effects in the Appendix. Briefly, our analysis primarily focused on the asymptotic behavior as $L \\rightarrow \\infty$ (using the strong law of large numbers, which just requires the mean to exist). However, in the linear example, our tokens are Gaussian and in the two non-linear cases they are bounded. Intuitively, we expect error $O(\\sqrt\\frac{1}{L})$. In fact, we can give precise probabilistic bounds of the form that the difference between the empirical sum for the ideal weights depart from the expectation by less than $C(\\tilde x)\\sqrt{\\frac{f(d,\\ln\\frac{1}{\\delta})}{L}}$ is greater than $1-\\delta$. The function $C$ of the query vector and the function $f$ depends on the problem. Interestingly, this indicates $d$ depends on the dimension spanned by the tokens (not the ambient dimension $n$). Figure 4(a) provides some empirical evidence for this relationship, showing how performance improves with increasing context length. \n- **(Q2)** Regarding the learned implementation matching theoretical expectations: Yes, we have verified this carefully. Figure 3(b) shows the final learned attention weights, which closely resemble scaled identity matrices as predicted by our theory. We also analyzed the loss landscape with respect to scaling factors $\\alpha$ and $\\beta$, confirming that the learned weights indeed lie in the predicted valley of the loss landscape. The minor deviations from the exact theoretical values can be attributed to finite context effects and the stochastic nature of the training process, but do not contradict our theoretical findings.\n\n--\n\nWe thank the reviewer again for their thoughtful review and insightful questions. Their suggestions have helped us identify important points to clarify and highlight in the final version of our paper.", "claims_and_evidence": "The authors provide ample theoretical and empirical support of their claims. They provide derivations of all of the Bayes optimal predictions $\\mathbb{E}[X|\\tilde{X}]$ for the various settings they consider, they provide constructions which show that transformers can implement these solutions, and they also provide many experiments where they train transformers on these tasks and show the ability of the model to achieve Bayes optimal error.", "methods_and_evaluation_criteria": "Yes, the paper primarily focuses on synthetic data but this is the setting of interest for their theory.", "theoretical_claims": "Yes, I checked the theoretical claims, most of which are easy to verify by following the derivations in the Appendix.", "experimental_designs_or_analyses": "I did not check the experimental design in great detail, but I trust the simulations.", "relation_to_broader_scientific_literature": "This paper studies an important problem of in-context capabilities of transformers, specifically in the context of denoising. They show that transformers can theoretically and empirically learn to denoise several distributions. In addition, they also provide interesting connections to the associative memory literature.", "essential_references_not_discussed": "Possibly of interest is this work https://arxiv.org/abs/2310.09753 which finds that adding additional identity matrix to the product $W_k^\\top W_Q$ can lead to improved performance on template \"reasoning\" tasks. This reminds me somewhat of the fact that you argue that $W_k^\\top W_Q$ should be close to a scaled identity in this work.", "other_strengths_and_weaknesses": "This paper studies an interesting problem, gives precise and testable predictions and verifies them. One weakness at the present moment is the limited scope of data distributions that are considered, but I think that the present contribution is already good.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. The authors provide mainly asymptotic results as context length diverges $L \\to \\infty$. Do they have any sense of the convergence rate in $L$ as a function of the dimension of the space? Specifically how large does $L$ need to be for the finite sample sums in equation 11 or 13 need to be to converge to the correct estimate?\n2. The construction concluded that $W_k^\\top W_q$ should be a scaled version of the identity, but in experiments it doesn't have to be exactly this. The authors then provide a short discussion on why this could be, specifically that there are different values of $W_k^\\top W_q$ and $W_v$ that implement equivalent attention operations. Have the authors checked if in fact the learned implementation matches their theoretical expectations?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lCCKxkeyhp", "reviewer": "Reviewer_BQoS", "summary": "## Update After Rebuttal\n## I maintain my score. Please see below for my reasonings in my response. Overall, I think it is a great paper despite some of my comments.\n-----\nThis work introduces a concept, in-context denoising, which is a task that refines the connection between attention-based architectures and Dense Associative Memory networks (DAMs). Using a Bayesian framework, based on the minimization problem of $\\mathbb{E}_{X, \\tilde{X}} [X - f_\\theta(\\tilde{X})]$, where $\\tilde{X}$ is the perturbed version of $X$, this paper illustrates that certain context-dependent denoising problems can be solved with a single-layer transformer model. Furthermore, it also demonstrates that a context-aware DAM one-step update yields better performance than a non-context-aware DAM update. Altogether, the work further solidifies the connection between attention and associative memory (AM), while illustrating the relevancy of AM for in-context learning (ICL).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive assessment and the high score (4/5) provided. We concur with their observation that, collectively, we still don't know why the transformer block works so well. This shared curiosity drives our work to build theoretical foundations for understanding the mechanisms underlying the success and flexibility of attention mechanisms. We are thus deeply grateful for their positive feedback in this direction. \n\n**Re: Claims and Evidence**  \nWe respectfully disagree with point (2) in the evidence section. While our paper builds on the connection established by Ramsauer et al., our contribution goes significantly beyond this foundation. Reading Ramsauer et al., one gets the impression that the question is memory retrieval. In our formulation, the task could be quite distinct from memory retrieval. We set up an in-context learning task, whose one-layer transformer-based solution is not guaranteed to look like a one-step lowering of some simple energy function. Quite remarkably, it does so in the special cases we study.\n\n**Re: Essential References**  \nWe thank the reviewer for highlighting the Energy Transformer (ET) paper and regret missing the citation. It is distinct but complementary effort that will be discussed in our revised manuscript to better situate our contribution within the broader literature.\n\n**Re: Strengths and Weaknesses**  \nRegarding the weaknesses discussed:\n- (1) We agree that analyzing deeper transformer systems is an important direction. Our focus on single-layer transformers was deliberate to establish clear connections in this foundational case. The question of how weights in deeper networks might implement multiple iterative steps on energy landscapes represents an exciting avenue for future research.\n- (2) This paper was written as an invitation to those working in mainstream transformer theory towards the very interesting activity in the DAM community. The in-context denoising task was meant to be a bridge between these two worlds. We will include further discussion on connections to DAM in the final version, if accepted, such as a detailed comparison with the Energy Transformer and a discussion of how our approach relates to tasks of interest in the DAM community, which we initially omitted due to space limitations.\n- (3) Indeed, in specific retrieval tasks, one might very well benefit from multiple steps or varied step sizes. In our specialized denoising cases, the 'step size' gets fixed during training, where the transformer is trying to approximate the Bayes optimal answer. This highlights an interesting distinction between retrieval tasks and denoising tasks that could inform broader transformer design.\n\n**Re: Comments and Suggestions**  \n- (1) Thanks for the kind words! We are certainly open to refocusing the introduction to keep the interest of readers, but would prefer to lead the results section with in-context learning and the Bayesian setting + derivations, as we feel it sets the stage well for the emergence of the associative memory behavior in the optimal/trained networks.\n\n- (2) Thanks much for catching the incomplete sentence at the end of subsection C.2 in the Appendix. It was meant to be part of the header for Section C.3. We will correct this in the revised version.\n\n--\n\nWe thank the reviewer again for their positive feedback and valuable suggestions. We believe that addressing these points will further strengthen the paper while maintaining its core contributions to solidifying the connection between attention mechanisms and associative memory in a minimal in-context learning setting.", "claims_and_evidence": "### Claim\n\nThe paper has two fundamental claims: \n\n(1) They claim that there are certain denoising tasks that a single-layer transformer can effectively solve. \n\n(2) Once trained, the attention layer effectively performs an operation that is a single gradient descent step on a context-aware DAM's energy landscape. This update is better than an extract retrieval of a token or spurious sample. \n\n### Evidence \n\n(1) To support the first claim, the presents theoretical results on three elementary denoising tasks where the data comes from (a) linear manifolds, (b) non-linear manifolds and (c) Gaussian mixtures or clusters, see Figure 1. Specifically, it presents the Bayesian optimal function $f_\\text{opt}$ for each of these cases. For each case of data, the work makes an argument that the $f_\\text{opt}$ of each case is equivalent to the attention mechanism of different activation (e.g., linear or softmax attention). \n\nTo further support their statement on $f_\\text{opt}$ of each case being equivalent to the attention mechanism. They perform experiments illustrated in Figure (3), which highlights the MSE loss of the estimated $\\hat{x}$ and the expected loss obtained from $f_\\text{opt}$. This overall demonstrates that the 1-layer transformer model will converge to the expected loss obtained from the Bayesian optimal predictor. Note, this particular experiment is pretty convincing. \n\n(2) The connection between DAM and attention is rather straightforward since [Ramsaeur et al. (2020)](https://arxiv.org/abs/2008.02217) has already established this initially. To further establish this connection, this work connects the DAM energy to ICL where the context sequence $X_{1:L}$ serves as the memory in which the noised prompt $\\hat{x}$ uses for its alignment with the system. However, this connection is quite brief and obvious, since the formulation is already done [Ramsaeur et al. (2020)](https://arxiv.org/abs/2008.02217).\n\nTo support their second claim, they perform a simple experiment detailed in Figure 5. In this experiment, using the DAM energy equation, which involves the term $\\frac{1}{2 \\alpha} \\lVert s \\rVert^2$, they contrast the update rule \n$s(t + 1) = (1 - \\frac{\\gamma}{\\alpha}) s(t) + \\gamma \\, X_{1:L} \\, \\text{softmax}(\\beta X_{1:L}^\\top s(t))$. Specifically, they contrast when the update rule is equivalent to the attention mechanism, i.e., $\\gamma = \\alpha$, and when $\\gamma \\neq \\alpha$. As already shown in [Ramsaeur et al. (2020)](https://arxiv.org/abs/2008.02217), when $\\gamma = \\alpha$, the retrieval is fast and accurate and shown here by this work, when $\\gamma \\neq \\alpha$, the retrieval diverges away from the target point because the dynamics become query-independent.\n\n## Strength\n \n(1) The theoretical results regarding $f_\\text{opt}$ for the three data cases presented in this paper are concrete and easy to understand. The paper also presents the expected optimal weight for the single-layer transformer once it is trained. This weight is simply a scaled identity matrix.\n\n(2) Another interesting result presented in the paper is the effect of increasing context length on the convergence to the expected bound based on the Bayesian optimal predictor. With sufficient (or increasing) context length, it seems that the convergence rate is increased.\n\n(3) This paper provides an interesting in the dynamics of transformer regarding denoising problem and further connects such dynamics to DAM. \n\n## Weaknesses\n\n(1) The analyses are done on a single-layer transformer system. It is unclear on how normal models, i.e., those with a large depth, behaves. Specifically, how does their weights vary and what effects do they play in performing multiple iterative steps?\n\n(2) The connection to DAM feels a bit minimal, which is somewhat understandable since much of the results are focused on the theoretical section involving the three general denoising cases that a simple transformer can solve. \n\n(3) Although the one-step update in DAM (i.e., when $\\gamma = \\alpha$) is faster to converge or accurate in terms for the retrieval setting, it's not exactly better depending on the task. For example, in terms of denoising to generate new data, the update with $\\gamma \\neq \\alpha$\ncould be better, since you don't need exact retrieval in this case.", "methods_and_evaluation_criteria": "(1) The methods or experimentation which evaluated claim (1) make a lot of sense. The authors trained a single-layer transformer and analyze whether its convergence is bounded by their theoretical value. I don't see faults on their experimentation.\n\n(2) The experimentation for claim (2) is straightforward and quite simple. Once again, there are no faults in their approach here.", "theoretical_claims": "I (the reviewer) primarily paid much of my attention in sections (2 and 3) since they contained the theoretical proof and results. I don't see any problems with the proofs nor the claims. My general issues are laid out in the weaknesses above.", "experimental_designs_or_analyses": "As mentioned prior, I see no flaws in the experimental designs and their analyses. My general issues are laid out in the weaknesses above.", "relation_to_broader_scientific_literature": "In general, we now know that the attention mechanism from [Vaswani et al. (2017)](https://arxiv.org/abs/1706.03762) is connected to the Modern Hopfield Model, developed by [Krotov and Hopfield (2016)](https://arxiv.org/abs/1606.01164) and [Demircigil et al. (2017)](https://arxiv.org/abs/1702.01929), shown by [Ramsaeur et al. (2020)](https://arxiv.org/abs/2008.02217). Moreover, we know how powerful the applications developed with the Transformer architecture can be, e.g., ViTs [Dosovitskiy et al.](https://arxiv.org/abs/2010.11929). ***But we still don't know why the transformer block works so well and what exactly does the mechanism of the block means***.\n\nWith the perspective of DAM, we know the attention mechanism is simply a gradient step on the DAM's energy landscape as shown by [Ramsaeur et al. (2020)](https://arxiv.org/abs/2008.02217) and further reinforced and explained in this work. With simple data cases, this work demonstrates that a trained single-layer attention model will converge to become a gradient step on a Hopfield model's energy landscape. This is quite significant.", "essential_references_not_discussed": "A nice reference that could be mentioned is the Energy Transformer (or [ET](https://arxiv.org/abs/2302.07253)). It is an extension of [Krotov and Hopfield (2021)](https://arxiv.org/abs/2008.06996) and [Ramsaeur et al. (2020)](https://arxiv.org/abs/2008.02217) which converts the non-dynamical Transformer block into a dynamical (Energy) Transformer block.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "(1) I think there are a lot of results presented in the paper, which is a strength. However, these results can be quite complex for some readers. I think leading the presentation of the paper with the narrative of DAM can help keep the interests of the readers.\n\n(2) In section (C.2) of Appendix C, there is an incomplete sentence at the end of the sub-section:  \n\n\"The clustering case\"", "questions_for_authors": "See my weaknesses above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper received strong reviews. In-context denoising task elegantly highlights the connection between Dense Associative Memories (DAM) and attention for tasks extending beyond precise recall of a memorized pattern. This work significantly adds to the growing literature connecting attention operation and DAM networks. Reviewers emphasize correctness, novelty, and high quality of the results. Several reviewers suggested incorporating additional references, which the authors will include in the final camera-ready version. The paper also offers novel perspectives on the one step energy-decent dynamics in DAMs. I share reviewers’ enthusiasm about this work and recommend acceptance.
Decision: Accept (oral)