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

TL;DR

This paper presents a measure-theoretic framework for Transformers as associative memory, providing theoretical analysis and minimax optimality results for their ability to recall and predict from distributional contexts.

Contribution

It introduces a novel measure-theoretic perspective on Transformers, proving their optimality and providing a framework for analyzing long-context recall with guarantees.

Findings

Transformer with softmax attention learns recall-and-predict maps under spectral assumptions.

A minimax lower bound matches the upper bound, showing optimal convergence rates.

Framework enables principled design and analysis of Transformers for distributional contexts.

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 $\nu = I^{-1} \sum_{i=1}^I \mu^{(i^*)}$ and a query $x_{\mathrm{q}}(i^*)$, the task decomposes into (i) recall of the relevant component $\mu^{(i^*)}$ and (ii) prediction from $(\mu_{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.