PaperlayerView exact revision on arXiv

Source-linked AI summary

Infinite attention: NNGP and NTK for deep attention networks

Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, Roman Novak

arXiv:2006.10540v1stat.MLcs.LG

TL;DR

The paper addresses the gap in infinite-width theory for neural networks containing attention layers, where standard Gaussian-process arguments do not directly apply. It develops NNGP and NTK characterizations for attention, analyzes scaling and architectural modifications, and finds Gaussian limits for sufficiently many heads alongside empirical gains and variable-length sequence support.

  • Problem

    Standard Gaussian-process proofs for wide neural networks do not directly extend to attention layers, whose single-head behavior can remain non-Gaussian.

  • Method

    The paper rigorously characterizes NNGP and NTK limits under d^-1/2 and d^-1 scalings and studies positional encodings, layer normalization, and modified attention mechanisms.

  • Results

    Multi-head attention becomes Gaussian as the number of heads tends to infinity, while attention kernels moderately improve CIFAR-10 GP performance and support variable-length sequence applications.

  • Takeaways & Limitations

    Attention can be incorporated into infinite-width NNGP/NTK analysis and applied to variable-length sequences, with practical gains from positional encodings and attention modifications.

  • Takeaways & Limitations

    Under d^-1 scaling, no attention-coefficient initialization avoids reduction to average pooling, limiting the expressivity of the resulting infinite-width attention behavior.

Abstract

from arXiv · show

There is a growing amount of literature on the relationship between wide neural networks (NNs) and Gaussian processes (GPs), identifying an equivalence between the two for a variety of NN architectures. This equivalence enables, for instance, accurate approximation of the behaviour of wide Bayesian NNs without MCMC or variational approximations, or characterisation of the distribution of randomly initialised wide NNs optimised by gradient descent without ever running an optimiser. We provide a rigorous extension of these results to NNs involving attention layers, showing that unlike single-head attention, which induces non-Gaussian behaviour, multi-head attention architectures behave as GPs as the number of heads tends to infinity. We further discuss the effects of positional encodings and layer normalisation, and propose modifications of the attention mechanism which lead to improved results for both finite and infinitely wide NNs. We evaluate attention kernels empirically, leading to a moderate improvement upon the previous state-of-the-art on CIFAR-10 for GPs without trainable kernels and advanced data preprocessing. Finally, we introduce new features to the Neural Tangents library (Novak et al., 2020) allowing applications of NNGP/NTK models, with and without attention, to variable-length sequences, with an example on the IMDb reviews dataset.

1. Introduction

The paper studies infinite-width behavior in neural networks with attention, where standard Gaussian-process proofs fail because single-head attention creates persistent statistical dependence. It rigorously characterizes the limits, showing that multi-head attention restores Gaussian behavior and supports practical kernel improvements.

  • Motivation: Wide-network theory derives NNGP and NTK kernels, but standard asymptotic Gaussianity arguments do not directly apply to attention layers.The difficulty arises because attention outputs depend on stochastic attention coefficients across spatial positions.
  • Motivation: Finite spatial dimension makes attention coefficients stochastic while increasing width, leaving output columns statistically dependent under d^-1/2 scaling.Each output column is a linear combination of the same stochastic matrix, preventing the independence assumptions used in conventional proofs.
  • Gaussianity: Single-head attention significantly deviates from Gaussian behavior at d=1000, whereas the 100-head architecture is compared against a Gaussian density.The figure examines outputs at initialization under d^-1/2 scaling.
  • Limitations of prior solutions: Replacing d^-1/2 with d^-1 enables Gaussian-conditioning arguments but forces attention toward average-pooling behavior and reduces expressivity.The paper therefore treats the scaling change as an incomplete solution.
  • Contributions: The paper rigorously characterizes infinite-width behavior under both d^-1/2 and d^-1 scalings, while evaluating positional encodings and modified attention mechanisms.It also reports moderate CIFAR-10 improvement for GP models and releases variable-length NNGP/NTK functionality with an IMDb example.

2. Definitions and notation

The paper defines the wide-network quantities used to derive attention NNGP and NTK kernels. It specifies network, attention, initialization, positional-encoding, and kernel-recursion notation for increasingly wide architectures.

  • Neural networks: For an input x, f^ℓ(x) denotes the ℓth-layer output and g^ℓ(x)=φ(f^ℓ(x)) its elementwise post-nonlinearity.The network has L hidden layers, with f^(L+1)(x) as the output.
  • Attention notation: Attention uses spatial and embedding dimensions, query, key, and value matrices, and multi-head variants indexed by head-specific weight matrices.Multiple spatial dimensions are flattened, and Gaussian weight initialization uses inverse-input-dimension variance scaling.
  • NNGP and NTK: As hidden widths grow, randomly initialized networks can converge to NNGP behavior, while gradient-trained networks can converge to NTK behavior.The corresponding zero-mean kernel functions are denoted κ^(L+1) and Θ^(L+1).
  • Kernel notation: The notation includes positional-encoding interpolation, optional operators, scaling-specific query-key conventions, and layer-normalization kernel transformations.Table 1 summarizes the discussed kernels under d^-1 and d^-1/2 query-key scalings.
  • Kernel recursion: Each layer maps preceding NNGP and NTK kernels to new kernels, with intermediate transformations representing nonlinearities before the layer update.The paper then derives these mappings for attention architectures.

3. Attention and Gaussian process behaviour

The paper shows that single-head attention can have non-Gaussian infinite-width behavior, while suitable multi-head limits yield Gaussian-process behavior. Under d^-1 scaling, attention collapses to global average pooling, whereas d^-1/2 scaling can recover convolutional kernels and Gaussian limits under stated assumptions.

  • Single-head attention: Single-head attention under d^-1/2 scaling can produce scale-mixture, non-Gaussian output behavior in the infinite-width limit.This complicates integration with layers whose infinite-width limits are Gaussian processes.
  • Limitations of the d^-1 scaling: Under d^-1 scaling, attention coefficients converge uniformly, making the resulting kernel equivalent to global average pooling.The same asymptotic distribution arises after flattening and applying dense layers.
  • Limitations of the d^-1 scaling: Proposition 2 shows that no attention-coefficient choice can simultaneously reproduce all positive semidefinite kernels under the d^-1 scaling.This establishes an expressivity limitation beyond the pooling reduction.
  • Infinite width limit under the d^-1/2 scaling: Under d^-1/2 scaling, multi-head attention converges to a Gaussian process when the number of heads or embedding dimensions per head grows.The proof proceeds by propagating Gaussian-process convergence through attention layers.

4. Beyond the vanilla attention definition

The paper replaces analytically intractable softmax expectations with alternative attention functions, then studies normalization, positional encodings, and modified infinite-width kernels. These changes improve finite and infinite attention models, while structured and residual kernels incorporate spatial or token relationships.

  • 4.1. Alternatives to softmax in attention networks: Replacing softmax with ReLU or identity makes the infinite-width attention expectation analytically tractable.ReLU preserves positivity without normalization, whereas identity permits unconstrained linear combinations of values.
  • 4.1. Alternatives to softmax in attention networks: Both alternative attention functions and layer normalization significantly improve finite-network performance, with identity attention and output normalization performing best across many settings.The experiments compare normalization applied after each head with normalization applied after the output projection.
  • 4.1. Alternatives to softmax in attention networks: Identity attention is used for the remaining experiments, while output layer normalization is retained because layer normalization preserves Gaussianity in the infinite-width limit.The paper contrasts identity with softmax where the corresponding kernels are available.
  • 4.2. Positional encodings: Positional encodings inject relationships between spatial dimensions, motivating covariance modifications that encode distances between image pixels or sequence tokens.The proposed covariance uses horizontal, vertical, and relative token distances to increase attention to neighboring positions.
  • 4.2.2. STRUCTURED POSITIONAL ENCODINGS: The resulting structured and residual kernels modify attention covariance, with the residual attention kernel arising as the infinite-width limit of a skip-connected architecture.Using the unscaled interpolated kernel can improve empirical performance, with further gains from the related residual kernel.

5. Experiments

The experiments evaluate attention NNGP/NTK kernels on CIFAR-10 and IMDb, testing architectural choices and comparisons with pooling-based baselines. Attention improves CIFAR-10 performance moderately, while IMDb gains are limited with high-quality embeddings.

  • Evaluation: Attention NNGP/NTK kernels are evaluated on CIFAR-10 and IMDb reviews.CIFAR-10 provides an established benchmark, while IMDb represents a typical language setting for attention models.
  • CIFAR-10 kernel choices: Except for Flatten, kernel choices achieve their best performance at smaller preceding depths.This pattern is consistent with observations reported in earlier literature.
  • CIFAR-10 kernel choices: Struct and Residual attention kernels significantly outperform the Vanilla kernel.The comparison indicates benefits from positional embeddings and layer normalisation in infinite-width models.
  • CIFAR-10 benchmark: Attention performs significantly better than GAP and moderately better than recent LAP results on the full CIFAR-10 experiment.The larger-scale comparison uses eight-layer models with Struct and Residual kernels and validation-selected hyperparameters.
  • Software and scope: The released Neural Tangents extension supports NNGP/NTK models on variable spatial extent and arbitrary dimensionality.This enables applications to inputs such as text, time series, and video.
  • IMDb reviews: Struct kernels outperform GAP-only on IMDb but provide marginal benefit over GAP-FCN.The authors conjecture that high-quality word embeddings already incorporate part of the relevant inductive bias.
  • IMDb reviews: IMDb sample results are not competitive with the state of the art, which reaches up to 97.4%.The authors position these results as a baseline for future infinite-width sequence-model research.

6. Conclusion

The conclusion establishes the asymptotic behavior of attention under different scalings and head counts, then summarizes empirical and software contributions. Multi-head attention can recover Gaussian limits, while positional encodings and attention modifications improve performance in finite and infinite settings.

  • Conclusion: Under d−1/2 query/key scaling, single-head attention may have non-Gaussian asymptotic behavior.The paper contrasts this with the d−1 scaling proposed in prior work.
  • Conclusion: Taking the number of heads to infinity can produce a Gaussian limit for attention outputs.The paper also reports improved performance from positional encodings and softmax replacements in finite and infinite architectures.
  • Conclusion: Attention moderately improves CIFAR-10 results over the previous state of the art for GPs without trainable kernels and advanced preprocessing.The paper additionally releases code for variable-length sequence inputs and demonstrates it on IMDb reviews.

A. Experimental details

The experimental details describe preprocessing, regression encoding, hyperparameter search, numerical settings, and kernel-comparison conventions. These choices define how the reported NNGP/NTK experiments were implemented and evaluated.

  • Data preprocessing: CIFAR-10 inputs are standardised per pixel and data point, with 10−15 added to standard deviations to avoid division by zero.This preprocessing is equivalent to using LayerNorm as the first layer.
  • Inference: Classification targets are encoded as C-dimensional vectors for closed-form NNGP and NTK regression inference.The inference uses a Gaussian likelihood with mean-squared-error loss.
  • Architectural settings: The experiments use specified weight, bias, and attention variances together with convolutional padding, stride, and filter-size settings.The exact configurations differ across the reported experimental setups.
  • Hyperparameter selection: Hyperparameters are selected through staged searches, with the second stage refining candidate kernels and including the Residual attention kernel.The best settings used in the main figures and tables are reported in bold in the search tables.
  • Implementation: Computations use 32-bit precision and up to eight NVIDIA V100 GPUs with 16GB of RAM each.These resources are stated for the experimental configurations described here.
  • Kernel evaluation: Kernel-matrix distances compare empirical and theoretical covariance matrices using the logarithm of their reported distance measure.The empirical and predicted matrices are defined on the training set.

A.1.3. DETAILS FOR FIGURE 3

The Figure 3 details specify the CIFAR-10 and IMDb configurations used to compare finite and infinite attention architectures. They cover data splits, preprocessing, embeddings, model families, and evaluation settings.

  • CIFAR-10 setup: CIFAR-10 experiments use a 45K/5K train/validation split and report validation accuracy after 1000 Adam epochs.The attention layers use d−1/2 query/key scaling and append positional encodings.
  • CIFAR-10 setup: Each CIFAR-10 configuration is run with three random seeds, with results included in Figure 3.The experiments use fixed learning rates and Adam for adaptive step sizing.
  • IMDb setup: IMDb experiments use GloVe embeddings, with sentences truncated or padded to 1000 tokens and binary targets mapped to {−0.5, 0.5}.Inference regularisation is selected from a validation-based range of diagonal regularisers.
  • Evaluation: IMDb models are selected using validation splits before evaluation on held-out test data.The reported configurations include 3200/1600/1600 splits and a 25K/25K train/test evaluation setup.
  • IMDb model families: The IMDb comparison includes GAP-only, GAP-FCN, CNN-GAP, and Struct model families.Struct models place one or two attention layers before pooling, while CNN-GAP varies the number of preceding convolutional layers.
  • Evaluation: The reported finite-attention results average performance over three random seeds.This convention is stated in the table describing the finite architectures.

B. Proofs

The proofs establish attention-layer convergence using exchangeability, moment control, and an infinite-width, finite-fan-out construction under explicit regularity and parametrisation assumptions.

  • Assumptions: The assumptions include a countable input set, polynomially bounded continuous nonlinearities, and additional regularity for NTK derivatives.The NTK arguments assume derivatives are continuous and bounded almost everywhere, including the stated convention for ReLU-like nonlinearities.
  • Proof strategy: The infinite-width, finite-fan-out construction adds unused heads to place networks on one probability space without changing their outputs.This analytical construction allows convergence claims for the standard finite-width networks.
  • Assumptions: The proofs use NTK parametrisation, with independently Gaussian-initialised weights scaled by layer width to normalise forward and backward passes.The stated parametrisation treats each weight as Gaussian with variance σ^2/d, while the associated scaling normalises both passes.
  • Proof strategy: The analysis inductively proves that attention layers preserve Gaussian-process convergence for NNGP and establishes the corresponding NTK result.The NNGP induction is treated in Theorem 3, while the NTK proof is treated in Theorem 18.
  • Probabilistic structure: Attention outputs are exchangeable across spatial indices and, for fixed layers, across heads, enabling conditional independence arguments.The proofs invoke de Finetti’s theorem after conditioning on suitable random variables.

B.1. d−1/2 NNGP convergence proof

The d−1/2 analysis proves that multi-head attention outputs converge to Gaussian-process limits by combining exchangeability, moment bounds, and finite-dimensional convergence arguments.

  • Gaussian-process limit: Under Theorem 3’s assumptions, the attention-layer output converges in distribution to a centred Gaussian process with the covariance specified by the theorem.The proof establishes convergence of finite-dimensional marginals and then uses product-topology arguments for the full process.
  • Gaussian-process limit: Head contributions are exchangeable and become conditionally i.i.d., allowing a Gaussian limit to be obtained from their averaged sums.The argument verifies exchangeability, zero mean, covariance convergence, variance convergence, and third-moment growth conditions.
  • Proof details: The proof reduces convergence to scalar projections of finite-dimensional marginals using the Cramér–Wold device and continuous mapping arguments.Uniform integrability then permits convergence of expectations through the stated theorem on weak convergence.
  • Proof details: Exchangeability, zero means, covariance identities, variance convergence, and third absolute-moment bounds supply the hypotheses of the Gaussian convergence lemma.These ingredients are established across Lemmas 8–17 before the theorem is concluded.

B.2. NTK convergence proof

The NTK proof decomposes attention-layer gradients into direct and indirect contributions and shows that both converge in probability under the NNGP assumptions and derivative regularity.

  • NTK convergence: Theorem 18 establishes convergence of the attention NTK at initialisation under the assumptions of Theorem 3 and additional derivative conditions.The proof treats arbitrary spatial coordinates and input pairs, then combines the direct and indirect contributions.
  • Contribution decomposition: The direct contribution captures gradients with respect to the current attention-layer parameters, while the indirect contribution propagates effects through preceding layers.The proof analyses these terms separately before combining their limits.
  • Direct contribution: Head-wise averaging and exchangeability reduce the variance of gradient sums, while bounded moments and uniform integrability control their limits.The argument repeatedly applies Chebyshev’s inequality, continuous mapping, and expectation-convergence results.
  • Derivative limits: The limiting derivative terms follow the pushforward of the Gaussian-process distribution under the nonlinearity derivative for τ = 1/2, while the τ = 1 limit is almost surely constant.These conclusions use continuity of the derivative and the corresponding NNGP limits.
  • Indirect contribution: The indirect-term proof classifies the non-negligible index matchings and shows that their second moments converge to the required limit.The analysis identifies three relevant matching types after discarding lower-order groups of summands.

B.3. Expressivity of d−1 and d−1/2 induced attention kernels

This section characterises expressivity limits of attention kernels under d−1 and d−1/2 scaling, including restrictions imposed by universal attention coefficients.

  • d−1 scaling: Proposition 2 shows that no fixed set of attention coefficients can reproduce every convolutional kernel through the relevant diagonal terms.The argument restricts universal coefficients to signs and fixed magnitudes, leaving insufficient freedom for arbitrary kernels.
  • Constructions: The kernel argument depends on the convolutional receptive-field sets N_a and N_b and their ordered pixel correspondences.These sets determine which pixels contribute to the new values at positions a and b.
  • d−1/2 scaling: Under d−1/2 scaling, Proposition 4 constructs a distribution over attention variables that realises any positive semidefinite kernel simultaneously.The proposition is stated for arbitrary inputs, coordinates, and positive semidefinite kernels.
  • Constructions: The construction uses random scaling variables applied to query and key rows to coordinate attention weights across spatial positions.The proposed variables are jointly distributed to achieve the desired kernel behaviour.

B.4. Auxiliary results

This section supplies convergence and moment-control tools used to establish infinite-width attention limits. Exchangeability, uniform integrability, and bounded moments support convergence of empirical averages and inner products.

  • Weak convergence in R^N is characterized by weak convergence of every finite-dimensional projection.
  • Uniform integrability combined with weak convergence yields convergence of expectations.
  • For exchangeable triangular arrays, empirical averages converge in expectation and probability when the limiting pairwise moment factorizes.
  • Moment propagation establishes uniformly bounded moments across widths, units, heads, and layers under the theorem’s assumptions.
  • These bounds and convergence results establish convergence of attention inner products and support the subsequent infinite-width arguments.

C.1. NNGP limit

The NNGP analysis extends to positional encodings without destroying Gaussian-process behavior. Positional encodings modify the limiting covariance through an interpolation with their covariance, while related NTK terms acquire additional contributions.

  • Exchangeability arguments remain valid after conditioning on the positional encodings.
  • For additive positional encodings, the limiting covariance includes the positional covariance term (1 − α)R_ab.
  • Adding positional encodings does not prevent GP behavior in the infinite-width limit.
  • The limiting covariance replaces each attention kernel ˜κℓ_ab(x, x′) with its interpolated form I ◦˜κℓ_ab(x, x′).
  • The NTK limit gains additional gradient contributions from positional encodings, including factors determined by the encoding construction.
Loading 2006.10540v1…