Why are Sensitive Functions Hard for Transformers? (2024)

Abstract: Empirical studies have identified a range of learnability biases and limitations of transformers, such as a persistent difficulty in learning to compute simple formal languages such as PARITY, and a bias towards low-degree functions. However, theoretical understanding remains limited, with existing expressiveness theory either overpredicting or underpredicting realistic learning abilities. We prove that, under the transformer architecture, the loss landscape is constrained by the input-space sensitivity: Transformers whose output is sensitive to many parts of the input string inhabit isolated points in parameter space, leading to a low-sensitivity bias in generalization. We show theoretically and empirically that this theory unifies a broad array of empirical observations about the learning abilities and biases of transformers, such as their generalization bias towards low sensitivity and low degree, and difficulty in length generalization for PARITY. This shows that understanding transformers’ inductive biases requires studying not just their in-principle expressivity, but also their loss landscape.

## introduction given dramatic advances in machine learning applications powered by transformer models, there has been substantial interest in understanding which functions are easier or harder to learn and represent using transformers. empirical research on both formal languages and synthetic functions has uncovered an intriguing array of learning biases, but theoretical understanding is lacking. for instance, @xcite experimentally argued that heldout generalization is biased towards low-degree polynomials and @xcite provided empirical evidence that transformers prefer to represent functions of low sensitivity, that is, functions that do not strongly depend on many input bits. perhaps the most prominent example of such learning biases is a consistent difficulty in learning the parity function, mapping bitstrings to their parity. this function is extremely sensitive, in the sense that flipping any bit flips the string's parity. empirical studies have consistently found that training transformers ## model of transformers we will focus on boolean functions. following the conventions in the analysis of boolean functions literature (o'donnell, 2014) in modeling bitstrings as elements of {-1, 1} n , we assume the alphabet σ = {-1, 1}, with word embeddings e(-1), e(+1) ∈ r d . there further are positional encodings p 1 , p 2 , p 3 , • • • ∈ r d . at the zero-th layer, token and positional encodings are added: y (0) i := e(x i ) + p i (@xmath0, . . . , n), where @xmath1. a transformer has a fixed number l of layers; the activations y (k) @xmath2, . . . , l) are defined as follows. each layer has a set of h attention heads; we first compute attention scores for the h-th head: of key importance will be the the normalization factor: our theoretical results will link n (k) i both to input-space sensitivity and parameter-space sharpness: we will find that large values of n (k) i can increase expressive capacity, but at the price of increased brittleness. finally, we assume that predictions are made by t := v t out • y (l) n for some parameter v @xmath3. throughout, we will add the input string @xmath4, e.g., writing t (x) for the overall prediction made on x. ## average sensitivity our results are centered around average sensitivity, a simple but foundational complexity metric for functions on the boolean cube (e.g. @xcite @xcite : definition 3. for a bitstring @xmath5: {±1} @xmath6, we write if f maps to {±1}, then s(x, f ) is the number of hamming neighbors of x on which f flips. this definition of as n ( f ) corresponds to the "total influence" from o'donnell (2014, def. 2.27). we explicitly define average sensitivity relative to input length n, as we will investigate the behavior of transformers performing a single function f across varying input lengths. the use of squared distances, rather than simple absolute distances, ensures that results about as n ( f ) transfer to results about degree profiles @xcite , which we will later investigate @xcite . average sensitivity is a general complexity metric with wide-ranging applications in theoretical computer science (e.g. @xcite . it is closely linked to the fourier analysis on the boolean cube @xcite , and is an average-case version of a family of sensitivity measures, closely related to other natural metrics such as decision tree depth and polynomial degree @xcite . both average sensitivity itself @xcite and the fourier structure @xcite have been empirically linked to transformers' generalization behavior. we will ground these empirical findings by relating average sensitivity to loss landscapes for the transformer architecture. ## lower bounds for sensitive functions we first prove that representing sensitive functions with transformers requires large parameter norms and, when inputs get longer and longer, highly unbounded normalization factors n (k) i in layer norm (5). we start from the global lipschitzness bounds developed by @xcite ; edelman et al. (2022); li et al. (2023) , but obtain more fine-grained average-case and highprobability bounds. these will then form the basis of our characterization of loss landscape sharpness around sensitive transformers. our bounds will include a constant c that is the product of and a term polynomial in h, d, the spectral norms of all parameter matrices appearing in the transformer, and the maximum norm of any positional or word embedding. see (32) in the appendix for formal definition. existing lipschitzness bounds (fact 2) imply a bound along the lines of foot_1 for 1 -hn -2 of inputs. we show this using a concentration bound argument, combining a chernoff bound applying to each attention weight individually, with a union bound over all attention weights (appendix, lemma 11). next, we address the role of layer norm. @xcite showed that, at ε → 0 in (3), layer norm enables transformers to represent par-ity. lipschitzness bounds in terms of ε (10) cease being meaningful in this limit. we thus study the layernorm induced blowup in more detail. in each layer, we consider the maximum blowup, given an input string @xmath7: the addition of 1 is for technical reasons (appendix, lemma 16); it has little impact in the cases relevant to our results, which will be when τ (k) (x) = ω(1). we then write blowup(x) := ∏ l @xmath8, an upper bound on the product of the successive layer-norm-induced blowups when going from the input to the output. when ε = 0 in (3), blowup(x) can be arbitrarily large. foreshadowing theorem 6, we will find that large values of blowup(x) create very sharp minima. our first theorem localizes the layer norm blowup to the hamming neighborhoods of sensitive inputs: theorem 4 (local bounds on layer norm blowup). consider a transformer with layer norm at arbitrary ε ≥ 0. with probability 1 -h n 2 over the choice of @xmath9, we have the proof is in appendix b.5. this permits us to state a bound on average sensitivity, in terms of the average layer norm blowup: corollary 5 (first main result). consider a transformer with layer norm at arbitrary ε ≥ 0. the proof is in appendix b.6. note that h n is small when n is large, and the bound is dominated by as n ( f ) √ n log n . this result thus shows a tradeoff between parameters and ln blowup: at least one of them needs to be large to represent a sensitive function. when the sensitivity depends on the input length and grows faster than √ n log n, this tradeoff changes with the input length, requiring larger parameters or larger layer norm blowup as the input length increases. let us investigate the implications for the functions introduced in section above. for parity, as n ( f ) = n, and (13) predicts for the other functions, the bound is o(1): for first, s( f , x) = as n ( f ) = 1, and the rhs of (13) is o(1). indeed, sparse functions can be modeled well by a family of transformers where the logits in the input layer scale with log n @xcite . unlike the prior lipschitzness bounds (9), these scaled logits do not contribute to c -our new bound is thus consistent with the ease with which transformers learn sparse functions. for mean, s( f , x) ∼ 1 n 2 ; again no blowup is predicted. for majority, as as n ( f ) ∼ √ n, no nontrivial average blowup is predicted. however, s( f , x) = n whenever the ones and zeros in x are almost fully balanced; on such strings or their hamming neighbors, (12) predicts @xmath10. ## sensitive transformers are brittle leveraging corollary 5, we now show that transformers expressing sensitive functions must be very sensitive to small perturbations of model parameters. that is, sensitivity in input space entails sensitivity in parameter space. an important consequence is that, for any highly sensitive function, any interpolating minima will be very sharp. this property is nontrivial, as seen from the fact that it disappears when adding a scratchpad (see below). given a parameter vector θ defining the transformer t θ , the average direction sharpness is (e.g. @xcite @xcite ) in defining the parameter vector θ in (15), we exclude positional encodings, both because they are often frozen, and because their number depends on n, hindering fair comparison across n. our next result lower-bounds l ρ,n in terms of the average sensitivity, provided the transformer is sufficiently wide in relation to its depth: theorem 6 (second main result). let t θ be a transformer where @xmath11. assume t θ (x) ∈ [-1, 1] for each x. then: informally, this theorem says that, when as n (t θ ) ∼ n, then even tiny perturbations to the parameters will, in expectation, lead to a substantial change in predictions on long inputs. this means that, as the input gets longer, the hessian of the mean-squared loss at the minimizer fitting a sensitive function has unboundedly large entries. see appendix c for the proof. the key idea of the proof is that, if t θ is very sensitive in input space, small perturbations to the parameters usually lead to a large drop in sensitivity when n is large, because they lead to large changes in the layer-norm induced blowup that is needed to represent high-sensitivity functions. as a consequence, transformers computing sensitive functions are isolated in parameter space. the theorem applies when d is substantially larger than l, as is indeed true of typical transformers (e.g., llama 7b has @xmath12. the convergence of the limits is slower when the parameternorms, as summarized by c, are larger, as larger parameters can increase sensitivity. however, remarkably, for any fixed transformer, c becomes irrelevant for l ρ,n in the limit where @xmath13. while we stated theorem 6 for an individual transformer, the same statement holds for families of transformers where weights may depend on n, as long as c remains bounded. a consequence is that scaling attention logits with log n (used to represent sparse functions by @xcite ; chiang and cholak (2022)) will, at least in the input layer, not mitigate the difficulty of sensitive functions. ## implications we discuss how theorem 6 unifies a range of diverse empirical findings about the behavior of transformers. difficulty of parity for parity, l ρ,n converges to 1 for large d, showing that arbitrarily small perturbations to a transformer computing parity will lead to a high loss for sufficiently long inputs. previously, @xcite noted that, for their handconstructed transformer representing parity, small changes to a specific parameter led to a large increase in loss, suggesting that this made the solution impossible to reach with sgd. theorem 6 shows that this phenomenon is unavoidable for any transformer representing a high-sensitivity function. for functions with sublinear average sensitivity, theorem 6 entails no nontrivial lower bound on sharpness, and no such phenomenon is predicted. ## experiments 7.1 setup we conducted experiments to test the predictions made by our theory, specifically assessing predictions regarding loss landscapes and sharpness of the minima. in all experiments, we use the transformer encoder architecture, using the default implementation in pytorch @xcite . each model is trained to fit a function f on a specific sequence length n. each training input x is generated uniformly from {±1} n , and each input bit, treated as a separate token, is embedded using learned token and positional encodings. 3 the representation of the last token is passed to a linear layer to generate the prediction t θ (x). parameters θ of the transformer are optimized for mse loss between t θ (x) and f (x) using adamw @xcite . for full details on hyperparameters and training setup, refer to appendix e.1. in implementation, we assumed versions of the functions outputting to {0, 1}; we rescaled sharpness values accordingly for comparability with the theory. we analyzed models using the following metrics: as we are interested in properties of models that compute given functions, runs that did not converge (evaluation mse higher than 10 -3 ) were discarded. ## discussion we have provided a rigorous theoretical explanation of the inductive bias towards low sensitivity observed in empirical research on transformers. theorem 6 describes a fundamental inductive bias of the transformer architecture: for long inputs, fitting sensitive functions is only possible in sharp minima of the loss. this holds without assumptions often made in previous work about expressive capacity, such as noninfinite precision (e.g. @xcite , hard attention (e.g. hahn, 2020), or lipschitz-continuous variants of layer norm @xcite . we speculate that theorem 6 reflects a fundamental limitation of parallelized differentiable computing with bounded depth. our results show that it is overcome by scaling the number of computation steps with the input length. while we focused on functions outputting a single label; an interesting direction for future research is the extension of this theory to sequence-to-sequence transductions, for which some empirical data has been reported (e.g. @xcite , but theoretical understanding remains wide open. due to the relation between average sensitivity and fourier analysis on the boolean cube (equation 17 ), a low-sensitivity bias can be viewed as a sequence modeling analogue of a spectral bias towards low frequencies observed in other neural architectures (e.g. @xcite . we note that, while our results show that sensitive transformers are very brittle, these results do not by themselves have implications for real-world generalization, as a low-sensitivity bias need not always be beneficial in real-world setups. indeed, the relationship between sharpness and real-world generalization is not straightforward (e.g. @xcite . our theory suggests that transformers generalize well to the extent that real-world data has bounded sensitivity (e.g. @xcite . ## conclusion we have proven that, under the transformer architecture, high sensitivity in input space can only be achieved in very sharp minima. empirical results confirm the predictions of the theory. taken together, our results explain a diverse set of empirical observations about transformers not explained by previous theoretical work. they suggest shifting theoretical research from in-principle expressiveness considerations to studying quantitative bounds and the shape of the loss landscape in order to understand the abilities of transformers.