Research map / Architectures and efficiency / Attention and transformers
In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces in-context denoising, showing single-layer (linear or softmax) attention can approximate Bayes-optimal denoisers on linear, spherical, and Gaussian-mixture token distributions, and that one attention step equals one gradient step on a context-aware dense associative memory energy.
Problem
Understanding how transformer attention relates to dense associative memory (modern Hopfield) networks beyond exact retrieval, and whether in-context learning can solve unsupervised denoising tasks with few layers.
Approach
Frame denoising of a query given L clean context tokens drawn from a random task distribution as in-context learning; derive Bayes-optimal posterior-mean denoisers for linear manifolds, spheres, and clusters, show one-layer attention matches them, and interpret the attention step as one gradient-descent step on a context-dependent DAM energy with the query as initial state.
Key results
- Trained one-layer linear attention with n=16, d=8, L=500 learns scaled identity weights with product 0.327 approximately 1/3, matching the Bayes-optimal scale for linear manifold denoising.
- Softmax attention achieves performance similar to linear attention on the linear manifold task, with the learned beta about 0.194 and alpha about 1.607.
- For Gaussian mixtures the trained network attains beta about 5.127 versus the optimal 1/sigma_Z^2 and lies in the same valley of the loss landscape as the Bayes predictor.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase context length (ablation) | One-layer attention denoiser · Linear manifolds | MSE | Helped |
| Use softmax attention instead of linear attention (ablation) | One-layer denoiser · Linear manifolds | MSE | No clear effect |
| Evaluate on subspace dimensions different from the training dimension d=8 (ablation) | Linear attention denoiser trained on random d=8 subspaces · Linear manifolds; inference-time dimension shift | denoising performance | Hurt |
| Iterate the denoising update instead of taking one step (ablation) | Attention-derived dense associative memory · Nonlinear manifold denoising | denoising performance | Hurt |
Surprising findings
- Naively iterating the one-step attention denoising update can degrade performance by driving the state to a query-independent distant minimum.
- Trained weights on the spherical and Gaussian mixture tasks are not exactly the Bayes-optimal scales, yet the network lies in the same loss valley as the optimal predictor.
Limitations
- Analysis restricted to simple synthetic denoising tasks with known Bayes-optimal baselines and a single attention layer.
- Convergence guarantees are asymptotic in context length L and rely on bounded or Gaussian token assumptions.
- Multi-layer or multi-iteration extensions are left open.
Experimental setup
Datasets: synthetic linear manifold tokens synthetic sphere tokens synthetic Gaussian mixture tokens
Benchmarks: synthetic in-context denoising tasks (linear, sphere, Gaussian mixture)
Metrics: mean squared error comparison to Bayes optimal MSE
Baselines: Bayes optimal denoiser zero predictor linear attention
Builds on: Ramsauer et al. hopfield is all you need linear transformers emulate gradient descent mimetic initialization
Closest papers by meaning
- In-Context Learning as Conditioned Associative Memory Retrieval (ICML 2025 · similarity 0.90)
- Outlier-Efficient Hopfield Layers for Large Transformer-Based Models (ICML 2024 · similarity 0.86)
- Denoising Task Routing for Diffusion Models (ICLR 2024 · similarity 0.86)
- Attention-Only Transformers via Unrolled Subspace Denoising (ICML 2025 · similarity 0.86)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.86)
- Denoising Task Difficulty-based Curriculum for Training Diffusion Models (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.