Towards Understanding Inductive Bias in Transformers: A View From Infinity

TL;DR

This work analyzes the inductive bias of transformers through the neural-network Gaussian process (NNGP) lens, showing a bias toward permutation-symmetric sequence functions when data are invariant under token permutations. By leveraging the representation theory of the symmetric group, the authors derive eigenstructure-based predictions for learnability and generalization, including a context-length scaling law, and validate these predictions using a simplified transformer block with linear activations and a mixture-of-HMM datasets. They demonstrate that the GP predictions extend to wide but finite networks and survive distributional shifts, including out-of-distribution settings, with qualitative support from experiments on linear and nonlinear activations. Additionally, analyses of WikiText-2 indicate that natural language data exhibit approximate permutation symmetry in principal components, aligning with the proposed framework and suggesting practical relevance for NLP model design and evaluation.

Abstract

We study inductive bias in Transformers in the infinitely over-parameterized Gaussian process limit and argue transformers tend to be biased towards more permutation symmetric functions in sequence space. We show that the representation theory of the symmetric group can be used to give quantitative analytical predictions when the dataset is symmetric to permutations between tokens. We present a simplified transformer block and solve the model at the limit, including accurate predictions for the learning curves and network outputs. We show that in common setups, one can derive tight bounds in the form of a scaling law for the learnability as a function of the context length. Finally, we argue WikiText dataset, does indeed possess a degree of permutation symmetry.

Figures (4)

• Figure 1: (Illustration of diagonalization using symmetries) The figure illustrates the direct sum (block) structure described in Prop. \ref{['prop:rep-theory']}. Each color-shaded block represents an irrep, and each solid color represents a multiplicity block within the irrep. All elements outside the multiplicity blocks vanish, both between different irreps and within an irrep. The symmetry actions $g \in G$ can mix multiplicity blocks as indicated by the arrows. Since all multiplicity blocks inside an irrep are linked by the symmetry actions they are all degenerate .
• Figure 2: Left: (theory vs. experiment) Two sections, along constant $N$ (in blue) and $L$ (in red), of the MSE loss as shown in the center. We find good agreement between our theoretical predictions (calculated for the train and test distributions) and exact inference with a GP, equivalent to inference with an infinitely wide NN. Stars indicate the experimental MSE loss calculated on the test dataset, where the majority of samples are OOD w.r.t to training dataset. Center: (learnability scaling law) Prediction for the generalization MSE. We see the learnability threshold as a diagonal valley (marked by the dashed line) of constant $N/L$ ratio as a consequence of having target eigenvalues of scaling $\lambda \sim L^{-1}$. This is the onset of the regime where there are enough datapoints to use the full potential of the context length. Right: (linear MLP predictions extend to non-linear MLP) Projections of the exact GP inference predictor with non-linear MLP on the basis vectors $\{ \varphi \}$ specified in equation \ref{['eq:phi_base_def']}, together with the theoretical prediction for the projection for a GP with linear MLP as a function of the number of training points. The predictions for the transformer block with linear attention and linear MLP are useful for non-linear MLP as well.
• Figure 3: Left: (Predicted OOD generalization matches experiment) Finite NN and exact GP inference (cyan and blue respectively), optimal predictions (red) and ground truth (dashed light red) as a function of the ground truth target for the sample. We see good agreement between predictions and experiment in distribution and OOD. The performance OOD is only slightly worse than in distribution, as indicated by the spread around the ground truth. Center: (the scaling bounds are tight for eigenvalues of softmax attention) The spectrum of the empirical kernel of a network with ${\rm softmax}$ attention as a function of the context length ($L$). The scaling with $L$ is bound tightly by the scaling deduced from the dimension of the corresponding irrep of the symmetric group. The light dashed lines serve only as a guide to the eye for the scaling law; they are not predictions for specific values. Right: (evidence for approximate permutation symmetry in WikiText) The triangle shows the cosine similarity induced by the Frobenius inner product between the linear features of WikiText $C^{kk}$ and $C^{k'k'}$ for the $k$'s indicated on the boundary. We see all sampled $k \neq 0$ are similar to one another but different from $k=0$ as predicted by the irrep decomposition. The Empirical CDF plot shows the CDF for the eigenvalues of those sampled matrices. Different $k$'s for $k\neq0$ are almost identical. $k=0$ has a distinct distribution.
• Figure 4: Left: (Similarity measure between $C^{kk}$ and $C^{k'k'}$ on WikiText) the cosine similarity induced by the Frobenius inner product between the linear features of WikiText $C^{kk}$ and $C^{k'k'}$ for the $k$'s indicated on the boundary. We see all sampled $k \neq 0$ are similar to one another but different from $k=0$ as predicted by the irrep decomposition. Center: (Similarity measure between $C^{kk}$ and $C^{k'k'}$ on a permutation symmetric baseline dataset) The baseline dataset is created by sampling words from WikiText with frequencies as in WikiText, but with to sequential order. The underlying distribution of baseline is therefore completely permutation symmetric in sequence space. We display the same quantity as the figure on the left, this time calculated on the baseline dataset. We see all sampled $k \neq 0$ are very similar, measurably more so then the same features of WikiText. Yet, comparing the differences between the datasets to the similarity gap of $\simeq 0.6$ between $k=0$ and all $k'\neq 0$, the results for WikiText and the baseline dataset are remarkably similar, suggesting an approximate permutation symmetry. Right: (Comparing the spectra of $C^{kk}$ between WikiText and the permutation symmetric baseline) The similarity between $k$s with $k\neq 0$ is again seen in the spectra. One notably difference between the baseline and WikiText is the spectrum of $C^{00}$ which differs along almost all the scale of eigenvalues, showing the principle components do capture information about sequence dependence, information that does not exist in the baseline.

Theorems & Definitions (29)

• Proposition 2.1
• Proposition 2.2
• Theorem 2.1
• Claim 1
• Claim 2
• Definition 3.1: Partition
• Theorem 3.2
• Definition 3.3: Young Diagram
• Definition 3.4: Young Tableau
• Definition 3.5: Standard Young Tableau