Representational Strengths and Limitations of Transformers

Clayton Sanford; Daniel Hsu; Matus Telgarsky

Source-linked AI summary

Representational Strengths and Limitations of Transformers

Clayton Sanford, Daniel Hsu, Matus Telgarsky

arXiv:2306.02896v2Jun 5, 2023cs.LGstat.ML

TL;DR

The paper asks for a mathematical account of when attention layers are more or less expressive than recurrent and feedforward architectures. It analyzes sparse averaging and matching tasks using approximation-theoretic, communication-complexity, and geometric arguments, finding efficient attention representations for sparse averaging and pair matching but strong limitations for unrestricted triple matching.

Problem
The paper addresses the lack of a mathematical description of attention layers’ representational benefits and deficiencies relative to other neural architectures.
Method
The paper studies sparse averaging, pair matching, and triple matching to analyze transformer expressivity as a function of width, depth, embedding dimension, and related complexity parameters.
Results
Attention efficiently represents sparse averaging and pair matching, whereas standard self-attention faces polynomial-complexity barriers for unrestricted triple matching and can solve structured variants efficiently.
Takeaways & Limitations
Sparse averaging characterizes the role of embedding dimension, while pairwise communication and contextual clues help explain why attention remains useful for practical sequential tasks.
Takeaways & Limitations
The paper conjectures that multi-layer transformers also require polynomial growth in width, depth, embedding dimension, or bit complexity for Match3, but addresses this only heuristically.

Abstract

from arXiv · show

Attention layers, as commonly used in transformers, form the backbone of modern deep learning, yet there is no mathematical description of their benefits and deficiencies as compared with other architectures. In this work we establish both positive and negative results on the representation power of attention layers, with a focus on intrinsic complexity parameters such as width, depth, and embedding dimension. On the positive side, we present a sparse averaging task, where recurrent networks and feedforward networks all have complexity scaling polynomially in the input size, whereas transformers scale merely logarithmically in the input size; furthermore, we use the same construction to show the necessity and role of a large embedding dimension in a transformer. On the negative side, we present a triple detection task, where attention layers in turn have complexity scaling linearly in the input size; as this scenario seems rare in practice, we also present natural variants that can be efficiently solved by attention layers. The proof techniques emphasize the value of communication complexity in the analysis of transformers and related models, and the role of sparse averaging as a prototypical attention task, which even finds use in the analysis of triple detection.

1 Introduction

The paper develops approximation-theoretic tasks that expose both the strengths and limitations of transformer self-attention, especially the role of embedding dimension. Sparse averaging favors attention over recurrent and fully connected networks, while triple matching reveals limits of standard attention that structured variants can avoid.

Research tasks: The paper uses sparse averaging, pair matching, and triple matching to relate transformer expressivity to embedding dimension and expose self-attention limitations.These tasks aggregate structural information encoded in embeddings and support comparisons with recurrent and fully connected architectures.
Sparse averaging: m ≳ q is necessary and sufficient for one self-attention unit to approximate q-sparse averaging.The result characterizes how embedding dimension controls the representation of sparse averaging.
Sparse averaging: Any fully connected network approximating qSA requires first-layer width at least Ω(Nd), while recurrent networks require hidden state at least Ω(N) bits.The attention construction instead has parameter dimension polynomial in q, d, and log N.
Pair and triple matching: A single O(d)-dimensional self-attention unit computes Match2, but Match3 requires H or m to grow polynomially in N for one standard multi-headed layer.This contrasts efficient pairwise interaction with the difficulty of purely triple-wise matching.
Pair and triple matching: Standard transformers efficiently compute modified Match3 variants using embedding-structure or locality assumptions, and third-order attention computes Match3 with one unit.The paper also conjectures that multi-layer transformers require polynomial growth in width, depth, embedding dimension, or bit complexity for unrestricted Match3.
Implications: The authors connect these results to practical sequence modeling by proposing that many sequential problems decompose into pairwise interactions and contextual clues.They state that purely triple-wise problems may not be necessary for language modeling, while Match2 may model primitives such as co-reference.

2 Preliminaries

The preliminaries formalize self-attention, multi-headed attention, transformer composition, and the notation used for fixed-precision vector computations. The framework can incorporate positional embeddings and residual connections while treating element-wise MLPs as arbitrary fixed-precision functions.

Attention units: A self-attention unit maps an input matrix through softmax(XQK^TX^T)XV using query, key, and value matrices.The unit has input dimension d, output dimension d′, embedding dimension m, and precision p.
Attention units: Multi-headed attention runs self-attention units in parallel and combines their outputs into an m-dimensional representation.The formal class is indexed by head count H and the constituent attention units.
Transformer architectures: Transformer models compose multi-headed attention layers with element-wise MLPs, which are modeled as arbitrary fixed-precision vector functions.The formal transformer contains D layers of H-headed self-attention with interspersed MLPs.
Notation and assumptions: The framework represents real-valued inputs and outputs using p-bit fixed-precision numbers, while some results allow infinite bit complexity.The model may also append a constant <END> token for a specific theorem’s representation setup.
Transformer architectures: Positional embeddings can encode an index as an input coordinate, and residual connections can be added as input to an MLP layer.These architectural features are implementable within the formal framework despite being omitted from its basic definition.

3 Sparse averaging with attention units

Sparse averaging shows that self-attention can aggregate arbitrary sparse subsets with embedding dimension scaling logarithmically in sequence length, while this construction is nearly optimal and supports communication-based lower bounds.

Task definition: Sparse averaging represents subset-dependent averages of input vectors and exposes how self-attention embedding dimension controls expressivity.The task models interactions between subset indicators y_i and vectors z_i.
Positive result: m = O(q log N) suffices for fixed-precision self-attention to approximately solve qSA.The construction uses keys, queries, and values arranged so selected subset elements receive larger attention scores.
Empirical evidence: Self-attention trained on qSA recovers an attention pattern similar to the theoretical construction.The visualization uses q = 3 and compares selected subset cells with the learned attention matrix.
Lower bound: m = Ω(q) is necessary in the infinite-precision setting, making the construction nearly optimal in its dependence on sparsity.The lower bound follows from a communication-complexity reduction to set disjointness.
Lower-bound method: Communication complexity reduces insufficient embedding dimension to set disjointness, using an mp-bit protocol for qSA.The reduction encodes one party’s input in a subset and uses the query embedding to recover disjointness.

4 Standard transformer models can only efficiently represent intrinsically pairwise functions

Standard self-attention efficiently represents pairwise matching but faces a sharp triple-wise limitation: Match3 requires polynomial growth in heads, embedding dimension, or precision, while structured variants remain tractable.

Pairwise matching: A single self-attention unit with embedding dimension m = 3 exactly computes Match2 for every input size N.The construction uses a blank token and trigonometric positional embeddings to focus attention on matching values.
Triple-wise matching: Match3 cannot be computed exactly by one multi-headed attention layer when mpH ≤ cN / log log N.Thus heads, embedding dimension, or precision must grow polynomially in N.
Scope: The general limitation concerns standard pairwise attention, whereas third-order self-attention can efficiently compute Match3.The paper conjectures that multiple layers of multi-headed standard attention share the impossibility.
Structured variants: Match3Bigram is efficiently computed with depth D = 2 and embedding dimension m = 3.One layer forms neighboring-pair sums, and the next applies the Match2 construction.
Structured variants: Match3Local is efficiently computed with embedding dimension m = O(K log N) when triples are restricted to a neighborhood of width K.Sparse averaging gathers the local window before an output MLP checks for matching triples.

A.2 Only high-memory recurrent neural networks can approximate qSA

Recurrent and other sequential memory-bounded algorithms require linear hidden memory to approximate sparse averaging, even on a causal subclass of inputs.

Model: The lower bound applies to recurrent neural networks and related sequential algorithms represented by arbitrary state-update functions.Each step computes an output and updated m-bit memory state from the current input and previous state.
Theorem: Any memory-bounded algorithm ε-approximating causal qSA with q = 1 requires m ≥ (N − 1) / 2 bits.The causal instances leave early subset fields empty and restrict later subsets to the early positions.
Proof strategy: The proof simulates the algorithm as a communication protocol for set disjointness, passing only the intermediate memory state.Alice sends the state after the first input block, and Bob continues the simulation on the second block.
Proof strategy: The resulting protocol implies m ≥ n = (N − 1) / 2 because disjointness requires at least n bits of communication.The algorithm’s output signs determine whether the encoded sets intersect.
Scope: Using set disjointness rather than INDEX extends the argument to multiple-pass and bidirectional recurrent models.The paper notes that INDEX suffices for single-pass algorithms, while disjointness handles broader sequential settings.

B.1 Proof of Theorem 2

Theorem 2 constructs fixed-precision self-attention for sparse averaging by encoding subset membership into queries and keys, using compressed-sensing embeddings to separate selected indices.

Theorem 2: m ≥ Ω(d′ + q log N) and p = Ω(log(q log N)) suffice for an ε-approximation of qSA.The construction uses fixed-precision arithmetic and produces an attention matrix concentrated on each requested subset.
Construction: Values copy z_i, while keys encode index-dependent vectors and queries encode subset-dependent separating vectors.The MLP maps each input to its value, an approximate subset vector, and an index key; Q, K, and V copy these components.
Precision: Finite precision is handled by approximating the subset vectors while controlling the resulting inner-product error.The proof bounds the perturbation after replacing ideal vectors with p-bit floating-point representations.
Compressed-sensing ingredient: Restricted isometry and orthogonality embeddings make selected keys distinguishable from nonselected keys.A random sign matrix satisfies the required property when m is at least proportional to q log N.

B.2 Proof of Theorem 3

The proof constructs attention units for sparse averaging using neighborly polytopes and a lookup-table MLP, with cyclic polytopes supplying the required geometry.

Construction: Theorem 3 uses a q-neighborly polytope to construct an attention unit that ϵ-approximates qSA.The construction follows from a geometric lemma for fixed key vectors forming a q-neighborly convex hull.
Geometric ingredient: A polytope is q-neighborly when every subset of at most q vertices forms a face.
Geometric ingredient: Cyclic polytopes generated by distinct real parameters provide the q-neighborly polytopes used in the construction.
Construction: The attention construction uses key and value embeddings from a lookup-table MLP, while query vectors are chosen through supporting hyperplanes of polytope faces.A sufficiently large scaling parameter α makes the attention output approximate qSA within ϵ.

B.3 Proof of Theorem 4

The proof reduces sparse averaging to two-party set disjointness, showing that a low-dimensional transformer would yield an insufficiently communicative protocol.

Lower bound: Therefore, if mp ≤ cq, no transformer in the stated class ϵ-approximates qSA.
Reduction: The proof embeds a DISJ instance of size q into a qSA instance with N ≥ 2q + 1.The constructed inputs encode Alice’s and Bob’s bits in the selected subset and value vectors.
Reduction: qSA(X) at the designated position equals −e1 exactly when DISJ(a, b) = 0.
Communication protocol: Assuming a transformer approximates qSA, Alice sends the query embedding to Bob using O(mp) bits.Bob combines this message with his encoded values to evaluate the designated attention output.
Lower bound: The resulting protocol computes DISJ, but communication complexity requires Ω(q) bits, forcing mp to scale at least linearly with q.

B.4 Optimality of Theorem 3 under restricted architectures

The paper proves near-optimality of the sparse-averaging construction for a restricted attention family by combining fixed-key geometry with a communication lower bound.

Scope: Theorem 16 establishes near-optimality for a restricted family resembling cross-attention rather than general self-attention.
Lower bound: For q < N and m ≤ q(1 − C log_N q), no attention unit in this restricted family can approximate the stated sparse-averaging variant.
Scope: The result assumes value embeddings reproduce target values, keys depend only on indices, and queries depend only on subset indices and positions.The authors leave extension to unrestricted infinite-precision cross-attention or self-attention architectures open.
Geometric argument: The lower bound uses fixed key embeddings whose convex hull lacks neighborliness, preventing separation of every size-q subset of value embeddings.
Geometric argument: The proof derives the missing neighborliness from a bound on how many column partitions m′-dimensional linear thresholds can separate.

C.1 Proof of Theorem 6

Theorem 6 constructs a single self-attention unit that solves pair matching with constant embedding dimension by using positional encodings and modular inner-product separation.

Theorem: For input range M = N^O(1) and precision p = O(log M), a single self-attention unit with embedding dimension m = 3 computes Match2 exactly.
Construction: The construction appends a blank token and uses positional encodings to define three-dimensional input, query, key, and value representations.
Matching test: The query-key inner product reaches its maximum exactly when two inputs sum to zero modulo M, while nonmatching pairs are separated by a fixed margin.O(log M)-bit fixed-precision arithmetic distinguishes the matching and nonmatching cases.
Readout: A sufficiently large softmax scale makes the attention output distinguish whether each input has a modular match, and an output MLP maps that distinction to Match2.The output MLP can be implemented with two ReLU gates.

C.2 Proof of Theorem 7

The proof reduces a restricted Match3 computation to two-party set disjointness, yielding a lower bound on the product of embedding dimension, precision, and head count. The same discussion contrasts this limitation with a two-layer construction and a third-order attention construction that computes Match3 efficiently.

Lower bound: If mpH ≤ cN/log log N, no f ∈ T 1,H satisfies Match3 on all X ∈ [M]^N for sufficiently large N.Theorem 7 states the lower-bound condition for any M ≥ N + 1.
Communication reduction: The proof encodes set-disjointness instances into restricted Match3 inputs and uses the assumed transformer to obtain a communication protocol.Alice and Bob split the input, exchange approximated normalization terms and partial convex combinations, and compute the transformer output collaboratively.
Communication reduction: The resulting protocol communicates O(mpH log log N) bits, so avoiding contradiction with the disjointness lower bound requires mpH ≥ Ω(N/log log N).The communication argument directly produces the theorem’s asymptotic requirement.
Depth separation: The restricted Match3 variant is substantially easier than the full-domain problem and can be computed by a two-layer single-headed transformer with constant embedding dimension.The first layer computes Match2 outputs, while the second takes their maximum.
Higher-order attention: A third-order transformer computes full Match3 with one self-attention unit and embedding dimension m = 5 for polynomial input range and logarithmic precision.The construction uses higher-order tensor products to represent dense three-wise interactions, though the model has steep evaluation cost.

C.5 Heuristic argument for Informal Conjecture 1

The heuristic argues that bounded-width, bounded-precision transformers cannot reliably detect planted triples unless their effective information capacity grows polynomially with sequence length. A matching construction shows this scaling is sufficient, while related graph-transformer results use Congest reductions to establish sharper separations.

Conjecture: The conjectured lower bound states that computing Match3 requires width, depth, embedding dimension, or bit complexity at least N^Ω(1).The formal version applies for sufficiently large N, M ≥ N + 1, and mpHD ≤ N^Ω(1).
Heuristic lower bound: A single self-attention output can incorporate information about at most mp inputs per layer, so after D layers with H heads its effective reach is bounded by (mpH)^D.The argument applies this information-growth heuristic to randomly generated inputs containing either no matching triple or one planted triple.
Heuristic lower bound: With mpHD = O(N), a randomly chosen triple fails to jointly appear in one self-attention output with probability at least 0.9.This makes the transformer unlikely to distinguish the planted-triple and no-triple cases under the heuristic distribution.
Tightness construction: A matching construction computes Match3 when M = N^O(1), p ≥ log(M), H = 1, m ≥ 4, and mD ≥ CN^2.Each layer checks a partition of candidate pairs, and the final MLP marks an element exactly when it belongs to a three-way match.
Graph-transformer separations: For graph subproblems, standard graph transformers require mpHD = Ω(N / log^2 N) for Cycle5 and DirectedCycle3, whereas higher-order graph transformers admit efficient constructions.Cycle5 is computed by a fourth-order transformer with one self-attention unit, and DirectedCycle3 by a third-order transformer with one unit, when p = O(log N).
Communication-complexity method: The graph lower bounds are obtained by simulating transformer computation on a bounded-degree Congest graph and reducing the result to two-party communication across a partition.The simulation takes O(HD(m + log N)) rounds, while the partition cuts O(N log N) edges and yields an O(pRN log N)-bit communication bound.

D Experiment details

The experiments compare attention, MLP, and LSTM architectures on synthetic sparse averaging data, examining training dynamics, alignment patterns, and generalization. Attention is the only architecture reported to generalize well across the tested attempts.

Experimental setup: Synthetic qSA experiments use 1,000 training and testing examples with sequence length N = 20 and q = 3.Each sequence element has total dimension 32, and attention uses width and embedding dimension m = 100.
Experimental setup: The data combine a four-dimensional data vector, positional encoding, and q additional positional encodings identifying elements to average.The positional encoding dimension is d0 := ⌈1 + 2 ln(N)⌉, and the data vector has d1 = 4.
Architectures: The comparison includes an MLP with a flattened input and one width-256 hidden ReLU layer, plus a two-layer LSTM with hidden state size 800.The MLP reshapes its final output to the desired sequence shape; the LSTM hidden state is 200 times the target output dimension 4.
Figures 6–7: Figure 6 plots training and testing error against thousands of training iterations, with error bars representing five random-initialization runs.The plotted regression objective compares the attention architecture with the additional MLP and LSTM architectures.
Figures 6–7: Only the attention layer achieves good generalization across all tested architectural variations, whose curves otherwise did not qualitatively change.Figure 7 extends alignment analysis to examples with minimum, median, and maximum loss, whereas Figure 2 uses only the median-loss example.