Source-linked AI summary
Are Transformers universal approximators of sequence-to-sequence functions?
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, Sanjiv Kumar
TL;DR
The paper asks what sequence-to-sequence functions parameter-shared Transformers can represent. It formalizes contextual mappings, analyzes the distinct roles of self-attention and feed-forward layers, and proves universal approximation results, while also evaluating simpler alternatives and identifying practical scope boundaries.
Problem
The expressive power of widely used Transformers is not well-understood despite their parameter sharing and limited token interactions.
Method
The paper formalizes contextual mappings, proves that fixed-width self-attention can compute them, and analyzes how they combine with feed-forward layers and positional encodings.
Results
Transformers universally approximate continuous permutation equivariant sequence-to-sequence functions with compact support, and trainable positional encodings extend this to continuous functions on compact domains.
Takeaways & Limitations
The universal approximation property depends on contextual-mapping capability, so components other than dot-product attention can preserve it when they provide that capability.
Takeaways & Limitations
The authors note that fixed-depth practical Transformers might not exactly implement contextual mappings as defined in the paper.
Abstract
from arXiv · showhide
Despite the widespread adoption of Transformer models for NLP tasks, the expressive power of these models is not well-understood. In this paper, we establish that Transformer models are universal approximators of continuous permutation equivariant sequence-to-sequence functions with compact support, which is quite surprising given the amount of shared parameters in these models. Furthermore, using positional encodings, we circumvent the restriction of permutation equivariance, and show that Transformer models can universally approximate arbitrary continuous sequence-to-sequence functions on a compact domain. Interestingly, our proof techniques clearly highlight the different roles of the self-attention and the feed-forward layers in Transformers. In particular, we prove that fixed width self-attention layers can compute contextual mappings of the input sequences, playing a key role in the universal approximation property of Transformers. Based on this insight from our analysis, we consider other simpler alternatives to self-attention layers and empirically evaluate them.
1 INTRODUCTION
The paper addresses the unresolved expressive power of parameter-shared Transformers and proves universal approximation results for sequence-to-sequence functions, including arbitrary continuous functions when trainable positional encodings are used.
- Motivation: Transformers are widely used for NLP tasks that map input token sequences to output token sequences.The cited applications include machine translation, language modeling, and question answering.
- Motivation: Parameter sharing and limited token interactions leave the representable class of Transformer sequence-to-sequence functions theoretically unclear.The paper frames this uncertainty around the roles and necessity of self-attention and feed-forward layers.
- Main results: Multi-head self-attention can compute contextual mappings, while feed-forward layers provide value mappings needed for universal approximation.The paper formalizes contextual mappings and identifies their computation as a key part of the approximation argument.
- Alternatives and experiments: Experiments find that combining simpler contextual-mapping architectures with Transformers performs better than standalone Transformers.The alternatives include bi-linear projections and separable convolutions, whose practical utility is evaluated empirically.
- Main results: Transformers universally approximate continuous permutation equivariant sequence-to-sequence functions with compact support.This is stated as Theorem 2 in the paper’s contribution summary.
- Main results: Trainable positional encodings remove the permutation equivariance restriction, enabling universal approximation of continuous sequence-to-sequence functions on compact domains.This result is stated as Theorem 3.
2 TRANSFORMER NETWORKS
A Transformer block maps sequences of token embeddings through self-attention and token-wise feed-forward layers, both using skip connections and shared parameters across token positions.
- Block structure: A Transformer block is a sequence-to-sequence function mapping R^d×n to R^d×n.Its input and output contain n token embeddings of dimension d.
- Block structure: Self-attention transforms each token using weighted combinations of all token embeddings, with weights generated from pairwise dot-products.The block definition uses multiplicative or dot-product attention.
- Block structure: The token-wise feed-forward layer independently processes modified token embeddings without interaction among tokens.Both layers include skip connections.
- Network definition: Transformer networks are compositions of blocks with h attention heads of size m and feed-forward hidden size r.The family T^h,m,r consists of compositions of these parameterized blocks.
- Structural property: Because each token undergoes the same transformations, Transformer blocks and their compositions are permutation equivariant.Permuting input columns permutes output columns in the same way.
- Structural property: Interactions between tokens in a block are captured only through pairwise dot-products in the softmax operator.This restriction, together with parameter sharing, makes the approximation capability non-obvious.
3 TRANSFORMERS ARE UNIVERSAL APPROXIMATORS OF SEQUENCE-TO-SEQUENCE FUNCTIONS
The paper proves that fixed-width Transformers universally approximate continuous permutation equivariant sequence-to-sequence functions, and that trainable positional encodings extend this to arbitrary continuous functions on compact domains. The proof separates self-attention's contextual-mapping role from feed-forward layers' value-mapping role.
- Universal approximation: Transformers with constant heads, head size, and feed-forward width approximate every continuous permutation equivariant sequence-to-sequence function with compact support.Theorem 2 gives the explicit construction g ∈ T 2,1,4 with dp(f, g) ≤ ϵ.
- Universal approximation: Trainable positional encodings remove the permutation equivariance restriction, enabling approximation of every continuous sequence-to-sequence function on a compact domain.Theorem 3 applies to the class FCD of continuous functions defined on compact domains.
- Representation power: The universal-approximation results use fixed widths whose parameters do not depend on sequence length n or embedding dimension d.The function classes become richer as the numbers of heads, head size, and hidden-layer size increase, while the stated fixed-width construction remains independent of n and d.
- Contextual mappings: Token-wise feed-forward layers map the distinct contextual values to outputs of arbitrary sequence-to-sequence functions.Contextual mappings let subsequent token-wise functions realize arbitrary outputs while preserving the sequence structure.
- Contextual mappings: Self-attention layers compute contextual mappings that assign token values based on the entire input sequence, despite using shared projections and pairwise interactions.A series of self-attention layers implements these mappings, overcoming the apparent limitation of pairwise interaction.
- Proof strategy: The approximation proof proceeds by approximating continuous targets with piecewise constant functions, approximating those with modified Transformers, and then recovering standard Transformers.The proof bounds the final error using the approximation steps and a term O(δ^d/p).
4 PROOF SKETCH OF PROPOSITION 4: DIFFERENT ROLES OF TWO LAYERS
The proof separates Transformer computation into feed-forward quantization and output mapping, with self-attention supplying contextual mappings that distinguish input sequences up to permutation.
- Contextual mapping: Self-attention layers then compute a contextual mapping whose entries distinguish grid inputs that are not permutations of one another.The construction uses a composition of fixed-width self-attention layers and separates distinct contexts through pairwise interactions.
- Function-value mapping: Feed-forward layers finally map contextual embeddings to the desired function values independently at each token.This division reflects the token-wise role of feed-forward layers after contextual information has been computed.
- Quantization: Feed-forward layers first quantize each input into a grid element, mapping inputs outside the compact support to a designated value.The quantization maps every input in cube S_L to grid point L and marks out-of-support elements separately.
- Contextual mapping: The contextual mapping is constructed with selective shifts that move entries in chosen quantization intervals while leaving other entries unchanged.Two attention heads of size 1 implement each selective shift, and stacking 1/δ layers processes the quantized values.
- Approximation complexity: The construction requires O(n(1/δ)^(dn)/n!) Transformer blocks for permutation-equivariant approximation and O(n(1/δ)^(dn)) layers without that restriction.The paper states that these parameter orders are optimal in the worst case and that the equivariant construction is more efficient than the compared residual-network construction.
5 DISCUSSION AND EXPERIMENTS
The discussion argues that contextual mappings are central to Transformer approximation and evaluates cheaper alternatives, finding that limited replacement of self-attention can improve performance while reducing cost.
- Discussion: Replacing dot-product attention with another component capable of contextual mappings should preserve the universal approximation property.This motivates studying alternatives with different computational and memory costs.
- Bi-linear projection: Bi-linear projection updates token embeddings through BProj(X) = X + W_O · X · W_P and uses fewer matrix multiplications than dot-product attention.Its contextual weights are input-independent, unlike self-attention weights.
- Separable convolution: Depth-wise separable convolution has sequence-length-independent parameters and can provide a cheaper contextual-mapping mechanism when multiple layers are stacked.A single small-filter layer is insufficient, but stacking layers can expand the receptive field.
- Experiments: The experiments evaluate BProj and SepConv on SQuAD and MNLI under a BERTBASE-style pre-training and fine-tuning setup.Standalone simpler models perform more weakly than self-attention models, motivating partial rather than complete replacement.
- Experiments: 1 or 2 convolution layers replacing early self-attention layers improve BERTBASE performance, while replacing more layers hurts performance.The hybrid models also reduce computational cost and parameter count; the paper suggests early broad attention may be handled more efficiently by convolutions.
A PROOF OF CLAIM 1
The proof shows that the Transformer function class is permutation equivariant because both self-attention and token-wise feed-forward operations commute with token permutations.
- Self-attention: Permuting the input tokens permutes the self-attention output by the same permutation.The derivation uses the permutation-matrix identity P P^T = I through the softmax attention computation.
- Feed-forward layer: The token-wise feed-forward layer is also permutation equivariant because ReLU(XP) = ReLU(X)P.Independent processing of each token preserves the ordering transformation rather than introducing token-specific positions.
- Consequence: Consequently, the Transformer function class T h,m,r is restricted to permutation-equivariant functions.This restriction explains why positional encodings are needed for arbitrary sequence-to-sequence functions.
B PROOF DETAILS OF THEOREM 2
The proof details approximate continuous permutation-equivariant functions by grid quantization, piecewise-constant approximation, and replacement of modified operations with standard Transformer components.
- Grid approximation: Uniform continuity permits a sufficiently fine grid and piecewise-constant approximation within ε/3 under the entry-wise ℓp metric.The argument uses equivalence of finite-dimensional ℓp and ℓ∞ norms to select the grid width δ.
- Grid approximation: Permutation equivariance of the target function makes the piecewise-constant grid approximation permutation equivariant as well.Permuting a cube maps it to the correspondingly permuted cube, preserving the approximation relation.
- Operation replacement: The resulting modified Transformer approximation can therefore be converted to a standard Transformer approximation with the stated width parameters.The proof concludes the conversion by choosing sufficiently large parameters for the approximating operations.
- Operation replacement: Softmax self-attention can approximate hardmax self-attention arbitrarily closely by scaling the attention parameters.Thus the modified hardmax construction can be transferred to standard self-attention with the same head count and head size.
- Operation replacement: Piecewise-linear activations with at most three pieces can be approximated arbitrarily closely using four ReLU units.This lets the modified feed-forward construction be represented by original ReLU feed-forward layers with four hidden nodes.
B.3 FINISHING PROOF OF PROPOSITION 4
The proof composes quantization, contextual mapping, and value mapping to approximate a piecewise constant permutation-equivariant function, except on a small-measure set.
- Proof construction: The modified Transformer first quantizes each input X to a grid element L, then computes contextual mappings q(L), and finally maps them to target values.These steps are formally established in Lemmas 5, 6, and 7.
- Approximation: The composition g = gv ◦ gc ◦ gq approximates the underlying piecewise constant function f over its support except on a set of measure O(δ^d).The exceptional set arises in the approximation construction over the function’s compact support.
- Approximation: For points X in the relevant grid cells, quantization gives gq(X) = L and the final composition satisfies g(X) = f(X).The equality is established for X ∈ S.
B.4 PROOF OF LEMMA 5
The construction quantizes inputs to an extended grid and uses selective shifts to create permutation-equivariant contextual identifiers for distinct input contexts.
- Quantization: Token-wise feed-forward layers quantize each input row into increments of δ, producing a map from R^d×n to the extended grid G+_δ.The construction repeats the row-wise quantization procedure across all d rows.
- Quantization: Each quantization layer maps values in [kδ, kδ + δ) to kδ while leaving other intervals unchanged.The activation functions used are piecewise linear with three pieces, with at least one constant piece.
- Selective shifting: A multiplicative attention construction selectively shifts the first coordinate of entries whose projected values fall within a chosen interval, leaving other coordinates untouched.The shift equals the range between the maximum and minimum projected values.
- Column identification: The projection u^T L_:,j acts as a bijective column identifier for valid grid columns and separates columns containing the extended-grid sentinel value.Distinct valid columns receive distinct projected values, while sentinel-containing columns project below zero.
- Contextual mapping: Stacking (1/δ)^d attention layers applies shifts in increasing column-id order, producing strictly increasing projected values across columns.A final global shift is then applied using the largest projected value.
- Contextual mapping: The resulting map gc is a permutation-equivariant contextual mapping on eG_δ, with distinct contexts producing separated projected outputs.The construction uses (1/δ)^d + 1 attention layers.
B.5.4 PROOF OF LEMMA 10
Lemma 10 establishes bounds and injectivity for the final projected column identifier after the sequence of selective shifts.
- Bounds: The proof derives lower and upper bounds on the final projected value eℓ_n after all selective shifts.These bounds constrain the possible values produced by the construction.
- Injectivity: The map from increasing projected input identifiers [ℓ_1, ℓ_2, …, ℓ_n] to eℓ_n is one-to-one.The proof compares the identifiers successively from the largest downward.
- Injectivity: The successive comparison argument forces ℓ_n = ℓ′_n and then ℓ_{n−1} = ℓ′_{n−1}, establishing injectivity across the ordered identifier sequences.The remaining terms have insufficient resolution to cancel the leading differences.
B.6 PROOF OF LEMMA 7
The value-mapping construction distinguishes valid contextual outputs from invalid-grid cases and maps each valid contextual column to its desired output.
- Case separation: Invalid-grid inputs are mapped to a zero matrix, while valid inputs L ∈ eG_δ retain gc(L).Negative entries are also mapped to zero before the value-mapping stage.
- Case separation: The construction first separates contextual outputs for valid grid inputs from outputs associated with invalid extended-grid points.A feed-forward layer tests whether projected values lie inside the interval [t_l, t_r].
- Value mapping: For valid inputs, distinct contextual columns have unique projected values separated by at least δ, enabling feed-forward value mapping.The method assigns each unique projected value to the corresponding output column.
- Value mapping: The feed-forward layers map each contextual column gc(L)_:,j to the desired output column (A_L)_:,j without affecting other columns.This mapping handles columns that coincide under permutations consistently.
- Value mapping: The construction requires O(n(1/δ)^(dn/n!)) layers because this bounds the number of unique projected contextual values.There is one layer for each unique value of u^T gc(L)_:,j.
C PROOF OF THEOREM 3
The proof constructs a Transformer that quantizes positionalized inputs, computes distinct contextual representations, and then maps each context to the desired output.
- Positional encoding: Positional encodings place token columns in distinct coordinate intervals, making their coordinates monotonically increasing.The first through nth columns lie in [0,1]^d, [1,2]^d, through [n−1,n]^d.
- Quantization: The proof quantizes positionalized inputs into a discrete grid using feed-forward layers.After dn/δ layers, each input entry in [kδ, kδ + δ) is quantized to kδ.
- Contextual mapping: Attention layers apply selective shifts that produce an injective mapping from each quantized context L to a unique ordered number.The construction yields e^l1 < e^l2 < ··· < e^ln, and different contexts map to distinct numbers separated by at least δ.
- Value mapping: A final value-mapping construction maps the distinct contextual representations to the desired outputs.Because permutation equivariance is absent, this construction requires O(n(1/δ)^(dn)) layers.
D EXPERIMENTAL SETUP
The experiments follow BERTBASE pretraining and fine-tuning settings while replacing self-attention with simpler alternatives, including bi-linear projection, separable convolution, and average attention.
- Baseline and pretraining: The baseline is BERTBASE, a 12-layer Transformer with 768-dimensional embeddings, 12 self-attention heads, and 3072-wide feed-forward layers.Models are pretrained using masked language modeling and next sentence prediction on English Wikipedia and BooksCorpus.
- Evaluation: The study evaluates models on SQuAD and MNLI after masked language model pretraining and task-specific fine-tuning.The reported experimental setting follows Devlin et al. (2018).
- Average attention: Average attention replaces input-dependent attention weights with uniform averaging over projected tokens.The model retains skip connections and feed-forward layers.