In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval

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.

Re: Claims and Evidence
We 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.

Re: Essential References
We 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.

Re: Strengths and Weaknesses
Regarding the weaknesses discussed:

• (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.
• (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.
• (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.

Re: Comments and Suggestions

• (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.

• (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.

--

We 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.

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.

Re: Weaknesses
The 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.

We 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.

Re: References
We 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.

Re: Specific Questions

• (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.
• (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.

--

We 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.

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.

Re: Positioning and Contribution
Our 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:

• For certain denoising tasks, one-layer transformers can represent Bayes-optimal predictors
• Trained attention layers correspond to single gradient descent steps on a context-dependent energy landscape
• 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.

Re: Relation To Broader Scientific Literature
We 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.

Re: Equation referencing
We 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.

Re: Strengths and Weaknesses
Regarding 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.

As 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.

Re: Why not higher score?
We respectfully suggest reconsidering our contribution:

• 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.
• Our focus on simplified synthetic scenarios establishes rigorous theoretical foundations, a common approach in theoretical ML papers.
• 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.

Re: Questions for Authors

• (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.
• (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.

--

The 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.

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.'

Re: Strengths and Weaknesses
We 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.

• 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).

• 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.)

• 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}$.

• 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.

We 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.

Re: References
We regret missing the Energy Transformer (ET) citation and are grateful to the reviewer for highlighting it. We will certainly discuss it.

Briefly, 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.

We likewise will cite Ambrogioni alongside Hoover et al. in the revision (thank you!).

Re: Typo in Appendix E
We reviewed and did not find a typo there, but we will change $\mathbb{1}$ to $\mathbb{1}_L$ for notational consistency.

Re: Questions for Authors

• (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.
• (Q2a) The apparent outperformance in Fig. 3a (Case 2) represents mild training set overfitting. Test set performance remains bounded by the theoretical optimum.
• (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.
• (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.

--

We 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.