Transformers as Measure-Theoretic Associative Memory: A Statistical Perspective and Minimax Optimality

TL;DR

This work recast associative memory at the level of probability measures, treating a context as a distribution over tokens and viewing attention as an integral operator on measures, and establishes a matching minimax lower bound with the same rate exponent.

Abstract

Transformers excel through content-addressable retrieval and the ability to exploit contexts of, in principle, unbounded length. We recast associative memory at the level of probability measures, treating a context as a distribution over tokens and viewing attention as an integral operator on measures. Concretely, for mixture contexts $ν= I^{-1} \sum_{i=1}^I μ^{(i^*)}$ and a query $x_{\mathrm{q}}(i^*)$, the task decomposes into (i) recall of the relevant component $μ^{(i^*)}$ and (ii) prediction from $(μ_{i^*},x_\mathrm{q})$. We study learned softmax attention (not a frozen kernel) trained by empirical risk minimization and show that a shallow measure-theoretic Transformer composed with an MLP learns the recall-and-predict map under a spectral assumption on the input densities. We further establish a matching minimax lower bound with the same rate exponent (up to multiplicative constants), proving sharpness of the convergence order. The framework offers a principled recipe for designing and analyzing Transformers that recall from arbitrarily long, distributional contexts with provable generalization guarantees.

Figures (4)

• Figure 1: Associative recall at the level of measures (informal): the query $x_\mathrm{q}(i^*)$ selects the relevant component measure $\mu^{(i^*)}_{v^{(i^*)}}$ from the mixture $\nu\propto \sum_{i=1}^I \mu^{(i)}_{v^{(i)}}$, followed by prediction from $(\mu^{(i^*)}_0,x_\mathrm{q})$. Note that each $\mu^{(i)}_{v^{(i)}}$ is constructed by $v^{(i)} \in \mathbb{S}^{d_1}$ and a measure $\mu^{(i)}_0$ on $\mathbb{R}^{d_2}$.
• Figure 2: A geometric illustration of how query $x_\mathrm{q}$ and components $\mu^{(i)}_{v^{(i)}}$ are mapped by the (simplified) first layer $h_j((v,z)) = (v,e_j(z))\in\mathbb{R}^{d_1+d_2}$, where $e_j$ is the $j$th Mercer eigenfunction. The product of the first $d_1$ indices of $h_j(x_\mathrm{q})$ and $h_j(y) \sim {h_j}_\sharp \mu_{v^{(i)}}^{(i)}$ tells whether $y=(v,z)$ is sampled from $\mu^{(i^*)}_{v^{(i^*)}}$ or not.
• Figure 3: Geometric sketch of associative recall. Components $\mu_{v^{(i)}}^{(i)}$ are separated along a feature axis via the first MLP layer $f_1$; the query maps to $\psi_j(x_q)$ and aligns with anchor $\psi_j(v^{(i^*)}) \coloneq \psi_j(({v^{(i^*)}}^\top,\boldsymbol{0}^\top)^\top)$, thereby recalling $\mu_0^{(i^*)}$. The pushforward ${\psi_j}_\sharp\mu^{(i^*)}_{v^{(i^*)}}$ provides features used by $F^\star(\mu^{(i^*)},x_q)$.
• Figure 4: Empirical risk $L(n)$ for the synthetic measure-valued experiment, plotted on a transformed axis $(\log n)^{\alpha/(\alpha+1)}$ together with the fitted curves. Each risk was calculated with $2000$ unknown samples.

Theorems & Definitions (85)

• Definition 1
• Example 1: Heat-kernel RKHS grigor2006heat
• Definition 2: Measure-theoretic attention layer furuya2025transformers
• Definition 3: Attention hypothesis class
• Definition 4: MLP hypothesis class
• Definition 5: Composition of measure-theoretic mappings
• Remark 1
• Definition 6: Transformer hypothesis class
• Theorem 1: Sub-polynomial convergence; informal version of \ref{['appendix-theorem-subpoly-overview']}
• Theorem 2: Sub-Polynomial Convergence, A Simplified Version of \ref{['theorem-subpoly', 'theorem-poly-capacity-beyond']}