Elliptical Attention

Authors: stefan k nielsen
laziz u abdullaev
rachel s y teo
tan m nguyen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 04:34:07

Abstract: Pairwise dot-product self-attention is key to the success of transformers that achieve state-of-the-art performance across a variety of applications in language and vision. This dot-product self-attention computes attention weights among the input tokens using Euclidean distance, which makes the model prone to representation collapse and vulnerable to contaminated samples. In this paper, we propose using a Mahalanobis distance metric for computing the attention weights to stretch the underlying feature space in directions of high contextual relevance. In particular, we define a hyper-ellipsoidal neighborhood around each query to increase the attention weights of the tokens lying in the contextually important directions. We term this novel class of attention Elliptical Attention. Our Elliptical Attention provides two benefits: 1) reducing representation collapse and 2) enhancing the model's robustness as the Elliptical Attention pays more attention to contextually relevant information, rather than focusing on some small subset of informative features. We empirically demonstrate the advantages of Elliptical Attention over the baseline dot-product attention and state-of-the-art attention methods on various practical tasks, including object classification, image
segmentation, and language modeling across different data modalities.

Main paper: 
 introduction: 
 Attention mechanisms and transformers [82] have achieved state of the art performance across a wide variety of tasks in machine learning [27, 35, 75] and, in particular, within natural language processing [1, 2, 13, 65, 12] , computer vision [16, 39, 78, 66, 62] , and reinforcement learning [25, 5] . They have also demonstrated strong performance in knowledge transfer from pretraining tasks to various downstream tasks with weak or no supervision [63, 64, 15] . At the core of these models is the dotproduct self-attention mechanism, which learns self-alignment between tokens in an input sequence by estimating the relative importance of each token with respect to all others. The mechanism then transforms each token into a weighted average of the feature representations of the other tokens with weights proportional to the learned importance scores. The relative importance scores capture contextual information among tokens and are key to the success of the transformer architecture [83, 76, 8, 59, 36, 55] . Organization. We structure this paper as follows: In Section 2, we present preliminaries on selfattention and non-parametric kernel regression. In Section 3, we illustrate the theoretical benefits of hyper-ellipsoidal neighborhoods, demonstrate how we build the required transformation, and provide the full technical formulation of Elliptical Attention. We empirically validate the advantages of the Elliptical Attention in Section 4. Related work is discussed in Section 5 before presenting concluding remarks in Section 6. Proofs, technical details, and further experiments are provided in the Appendix. 
background: self-attention and non-parametric regression: 
 We first provide preliminaries on the self-attention mechanism followed by background on its connection to the Nadaraya-Watson (NW) estimator in non-parametric regression [48] .Given an input sequence X = [x 1 , . . . , x N ] ⊤ ∈ R N ×Dx of N feature vectors, the self-attention mechanism transforms the input to H := [h 1 , . . . , h N ] ⊤ ∈ R N ×Dx as follows:EQUATIONThe vectors q i , k j , and v j are the query, key, and value vectors, respectively. They are computed as [q 1 , . . . ,q N ] ⊤ := Q = XW ⊤ Q ∈ R N ×D , [k 1 , . . . , k N ] ⊤ := K = XW ⊤ K ∈ R N ×D , and [v 1 , . . . , v N ] ⊤ := V = XW ⊤ V ∈ R N ×Dv where W Q , W K ∈ R D×Dx , W V ∈ R Dv×Dxare the weight matrices. Eqn. 1 can be expressed in matrix form as:EQUATIONwhere the softmax function is applied row-wise to the matrix QK ⊤ / √ D. We refer to transformers built with Eqn. 2 as standard transformers or just transformers.We now present the connection between self-attention as described in Eqn. 1 and non-parametric regression. We first assume key and value vectors {k j , v j } j∈ [N ] are obtained from the following data generating process:EQUATIONwhere ϵ is random zero-mean noise E[ϵ] = 0, and f is the unknown function to be estimated. We consider the random design setting where the keys {k j } j∈[N ] are i.i.d samples drawn from the marginal distribution p(k). We use p(v, k) to denote the joint distribution of pairs (v, k) as obtained according to Eqn. 3. At any given new query q, we aim to estimate the unknown function f (q). The NW estimator is a non-parametric estimator of the unknown f described byEQUATIONwhere we apply zero-mean noise for the first equality and the definitions of conditional expectation and density for the second and final. Then, it can be shown that by estimating the joint density p(v, k) and marginal density p(k) using isotropic Gaussian kernels with bandwidth σ and evaluating the NW estimator at a new query q i , we obtainEQUATIONwhere choosing σ 2 = √ D as the isotropic variance recovers the full attention mechanism. We present the full derivation in Appendix A.Limitation of self-attention. We see in Eqn. 5 that standard self-attention computes the relative importance scores between queries and keys via Euclidean distance. Euclidean distances are spherically invariant and therefore fail to consider coordinate-wise significance in the feature space, meaning the proximity of k j from q i influences its contextual relevance equally regardless of direction.3 Elliptical Attention: Leveraging Hyper-Ellipsoids to Pay More Attention Without Losing FocusIn this section, we first present how NW regression obtains a lower MSE by taking hyper-ellipsoidal neighborhoods around queries. We then construct the required hyper-ellipsoidal transformation via a Mahalanobis metric. We present the framework relating robustness and representation collapse to the geometry of the query neighborhoods and show how our proposed scheme offers improvements in both areas. We then provide an efficient estimator of the coordinate-wise relevance before finally giving the full technical formulation of Elliptical Attention. Technical details on the implementation procedure are in Appendix E. 
improving nw regression with hyper-ellipsoids: 
 Distance-based estimators, such as the NW estimator, can obtain a lower MSE by taking hyperellipsoidal neighborhoods around queries [29, 30] . The key idea is that we wish to stretch the axes of the underlying space in directions for which the true f in Eqn. 3 varies least. Figure 2 shows a situation in which f does not vary equally in all directions. This is actually a limiting case in which the function is sparse in the x 2 direction. In the left sub-figure, we show the result of stretching the Euclidean circular neighborhoods around each query in the x 2 direction for which the function does not vary. The right sub-figure then shows how the resulting ellipse in the x 2 direction can include additional data points without adding additional bias into the model. It is a well-established result that the variance of non-parametric estimates at a point is inversely proportional to the number of samples in that point's neighborhood, as the additional samples smooth out the effect of noise. As a result, stretching the neighborhood, as shown in the right sub-figure, decreases the variance. Crucially, including these additional samples does not cause the estimate to miss the true variation in the function, as there is no variation in the x 2 direction. By including points in this direction, we do not introduce bias into the estimate. Hence, we lower variance without the introduction of bias, obtaining a lower MSE estimator. This intuition is formalized in Theorem 1 in Appendix C, which shows that the best achievable rate of convergence for estimators of non-sparse Lipschitz functions is of the order O(n -2/(2+d) ) for a d dimensional feature space. However, when the function only depends on R ⊆ [d] coordinates, the rate improves to O(n -2/(2+|R|) ). In the case of approximate sparsity, when coordinate directions exhibit differing variability, the same intuition carries over as shown by the improvement in convergence rates in Theorem 2 in Appendix C.We leverage this analysis from non-parametric regression to motivate our Elliptical Attention. From the regression perspective, the self-attention mechanism, which performs NW regression, is able to learn a lower MSE estimator of the true underlying f by reducing the variance of the estimator without (or with minimal) introduction of bias. From the attention perspective, this means queries pay higher attention to more relevant keys, producing more contextually meaningful attention scores and better, more robust learned representations.We measure the variation in f in the i th coordinate direction by the expectation of the L 1 norm of the i th directional derivative taken over all k ∈ X k , where X k ⊆ R D denotes the feature space. Roughly speaking, this quantity corresponds to the average absolute gradient of f in the i th direction throughout the space. Formally, this quantity is defined as Definition 1 (Coordinate-wise Variability of f : R D → R Dv ) The coordinate-wise variability of f : R D → R Dv with Jacobian matrix J f ∈ R Dv×D in the i th direction is given by the quantity∥f ′ i ∥ 1,µ := E k∼µ ∥J f (k)e i ∥ 1 , i ∈ [D], where e i is an all-zero vector with a single 1 in the i th coordinate and µ is the marginal distribution of k over support X k .Remark 1 This definition is one of many possible. One could also take the supremum rather than the expectation or consider second derivatives. We select this definition as averages over first derivatives are more easily estimated and the definition still captures the intuitive properties of variability.i ∥ 1,µ by m i , we can then incorporate these quantities into a distance function of the formEQUATIONwhereM = diag(m 1 , m 2 , . . . , m D ) is a diagonal matrix whose diagonal elements are the estimates of ∥f ′ i ∥ 1,µ for i ∈ [D].Remark 2 The metric described in Eqn. 7 is a form of Mahalanobis distance metric, which can be interpreted as first applying a transformation to the underlying space in which we stretch the coordinate axes by the diagonal elements of M . Therefore using this metric within the self-attention computation produces the desired hyper-ellipsoidal neighborhoods around queries.In practice, we maxscale the estimates to obtain m i ← m i /m max where m max ≥ m i for all i ∈ [D]. This is because we care about the relative magnitudes of the direction-wise variability as opposed to the absolute magnitudes. Under this interpretation, we identify the most variable dimension and stretch all others relative to this direction.Before providing the technical procedure for estimating M and the full technical formulation of Elliptical Attention in Section 3.5, we first theoretically analyze in Propositions 1 and 2 how the hyper-ellipsoidal transformation in Eqn.7 improves robustness and alleviates representation collapse.Dimension-wise input sensitivity of Elliptical Attention and robustness. In Lemma 1, we show that when each input component is weighted according to the Mahalanobis transformation in Eqn. 7, the impact of perturbing the i th input coordinate on any coordinate of the output is proportional to the corresponding weighting parameter with proportionality coefficient depending on the indices i and j.Lemma 1 Let M : R D → R N denote the transformed Elliptical softmax operator for a given set of keys as M(x) :=1 j∈[N ] exp(x ⊤ M kj ) exp(x ⊤ M k 1 ), exp(x ⊤ M k 2 ), . . . , exp(x ⊤ M k N )⊤ for weight matrix M as in Eqn. 7. Then, the achievable rate of change ofM(x) in i th input dimension is proportional to m i , that is, sup x∈X |J M (x) ji | ∝ m i , for all i ∈ [D] and j ∈ [N ]where J M is the Jacobian matrix of M.By virtue of Lemma 1, which is proven in Appendix B.1, we show in Proposition 1 that choosing the weights as properly scaled estimates of the underlying function variability, as in Elliptical Attention, the output vectors become less prone to large errors caused by noisy input while simultaneously respecting the dimension-wise variability pattern of the true self-attention function.Proposition 1 (Robustness of Elliptical Attention) Let f : R D → R Dv be the true self-attention function, fd be the Elliptical Attention estimator with metric d as described in Eqn. 7. Then for any index i ∈ [N ] and noise ϵ ∈ R D , the following error bound holdsEQUATIONwhere {K j } j∈[N ] are constant diagonal matrices that depend only on the key vectors.Note that when the estimates are maxscaled so that m i ≤ 1, the achievable output error of Elliptical Attention is lower than that of standard self-attention where m i = 1 for all i ∈ [D]. Besides, when the true function exhibits approximate sparisity in some number of dimensions (i.e. m i → 0 + for majority of indices), the error bound in Eqn. 8 becomes significantly tighter for Elliptical Attention. The proof of Proposition 1 is provided in Appendix B.2.Input smoothing and representation collapse. In each layer, the standard self-attention mechanism fits a noisy estimate of the true function f , which is then fed into subsequent layers and iteratively refit. The input to each attention layer is then partially composed of noise, which is equivalently the common regularization method of random input smoothing. We show that by reducing the noise component in each layer, Elliptical Attention maintains expressive power and resists representation collapse. This is formalized in the following proposition:Proposition 2 (Elliptical Attention maintains expressive power by reducing noise) Let h ℓ d denote the output of a transformer using Elliptical Attention with metric d as described in Eqn. 7 and h ℓ denote the output of a transformer using standard self-attention at layer ℓ. Let D be the sampling distribution of the data and let c ∈ R D . Then, for any h, h d and layer ℓ, in expectation a standard self-attention transformer attenuates towards c faster than Elliptical Attention. Formally, we have:EQUATIONProof is provided in Appendix B.3. Proposition 2 shows Elliptical Attention maintains better expressive power than standard self-attention. We find this empirically supported as shown in Fig 3.We propose a simple difference-based estimator that effectively captures the coordinate-wise variability of the underlying function. Our estimator is easily and efficiently computed. It requires no additional learnable parameters and demands negligible additional memory. Let E n denote empirical mean over n samples, v ℓ (i) denote the i th component of the vector v at the ℓ th layer, andX ℓ,ℓ+1 v = {(v ℓ+1 , v ℓ ) : v ℓ = f (k ℓ )+ ϵ} be the value feature space at neighboring layers ℓ and ℓ + 1 where values are generated according to the process described in Eqn. 3. Then, our approach to estimating the i th coordinate-wise variability is described in the following proposition.Proposition 3 (Coordinate-wise Variability Estimator) Given a function f : R D → R Dv with i th directional variation ∥f ′ i ∥ 1,µ , i ∈ [D] and some δ > 0, the directional variation can be estimated by the quantityEQUATION)Remark 4 For the purposes of improving the performance of transformers by stretching the feature space according to the direction-wise variability of f , we note that consistent estimators of ∥f ′ i ∥ 1,µ for all i ∈ [D] are sufficient but not necessary. Instead, we require only the weaker objective of accurately estimating the relative magnitudes of the direction-wise variability. That is, if ∥f ′ i ∥ 1,µ ≥ ∥f ′ j ∥ 1,µ , we need only that m i ≥ m j . This is because the theory requires us only to identify coordinate directions of more or less variability and shrink or stretch the space accordingly.The intuition behind our estimator in Eqn. 10 lies in prior lines of research studying transformers as an Euler discretization of a continuous-time dynamic, usually as a system of first-order ordinary differential equations (ODEs) [40, 21, 53] . In fact, our estimator resembles the absolute value of a forward Euler discretization of the variability of the i th component of a value vector over time ∂v(i, t)/∂t, where the layers ℓ and ℓ+1 represent consecutive time points in an interval partition with the step size δ. We prove that our estimator in Eqn. 10 effectively estimates the relative magnitudes of the coordinate-wise variability of f in Appendix B.5.We now present the full formulation of Elliptical Attention. Given the distance function d(•, •) as in Eqn. 7, where M = diag(m 1 , . . . , m D ) is a diagonal matrix with elements m i as in Prop. 3, the M -norm can be defined as ∥x∥ M := √x T M x, which produces hyper-ellipsoidal stretching in the feature space. Then, Elliptical Attention is defined as follows.Definition 2 (Elliptical Attention Computation) Let φ d,σ : R D → R denote the Gaussian density kernel with variance σ 2 I equipped with the M -norm as defined above. Then the corresponding NW estimator at q i becomes fd,D (q i ) :EQUATION)Then, the Elliptical Attention output for the i th query q i given keys {k i } N i=1 and values {v i } N i=1 corresponding to the NW estimator (11) with σ 2 = √ D is given byEQUATIONwhere M = diag(m 1 , . . . , m D ) with m i defined as Eqn. 10 for all i ∈ [D].Eqn. 12 is equivalently expressed in matrix form asEQUATIONRemark 5 We see from the form of Eqns. 12, 13 that standard self-attention is recoverable by setting M = I D . Under our framework, this implies that standard self-attention assumes the underlying regression function to have exactly equal variability in all coordinate directions.Pseudocode for the Elliptical Attention computation is provided in Appendix F.12. 
experimental results: 
 In this section, we empirically justify the advantage of Elliptical Attention over baseline transformers that take hyper-spheres around queries. We evaluate our method on robust Wikitext-103 modeling under Word Swap contamination [45] , ImageNet classification under a wide range of attacks [14, 67] , the LRA benchmark [74] , and ADE20K image segmentation [87] . We compare Elliptical Attention with state-of-the-art (SOTA) clean and robust models, including Performer [9] , FourierFormer [54] , Robust Vision Transformer [44] , Fully Attentional Network (FAN) [89] , Mixture of Gaussian Keys (MGK) [52] , Mixture-of-Expert (MoE) based transformers, such as Switch transformer [18] and Generalist Langauge Model (GLaM) [17] , and robust kernel density estimation (KDE) based transformers, such as Median of Means (MoM) and Scaled Projected KDE (SPKDE) [23] . We aim to show that i) Elliptical Attention offers substantive improvements over baseline models across tasks on both clean and contaminated data; ii) Elliptical Attention attains these improvements on contaminated data while reducing memory requirements and increasing computational speed compared to comparative robust models; iii) Elliptical Attention can be combined with SOTA robust transformers to further improve robustness with negligible increase in computational overhead. We compare Elliptical Attention with baselines of the same configuration. Results are averaged over 5 runs with different seeds. Additional results and full details on experimental setup are in Appendix F.Experimental setup. We adopt the experimental setup in [54, 23] . We pretrain and evaluate our models on the WikiText-103 benchmark in comparison with the standard baseline Transformer [82] , Performer [9] , Transformer-MGK [52] , FourierFormer [54] , and the robust kernel density estimationbased Transformers including Transformer-SPKDE and Transformer-MoM [23] . All models use the 44M-parameter Transformer backbone. We pretrain all models on clean data for 125 epochs before attacking only the test set using a Word Swap Attack, which substitues random words with a generic 'AAA' token at a 2.5% swap rate. We report test perplexity (PPL) as the performance metric. Results. Table 1 shows our Elliptical Transformer (Elliptical) achieves top test perplexity in clean data while also achieving second top test perplexity under data contamination by Word Swap [47] , illustrating that the Elliptical Attention is highly robust and offers substantial advantages on clean data as well.Experimental setup. We adopt the experimental setup in [23] . We train and evaluate Elliptical on ImageNet-1K against standard vision transformers, including DeiT [78] and Distill [78] , as well as the FourierFormer [54] . We also compare Elliptical with robust vision transformers, including DeiT-KDE [23] , DeiT-MoM [23] , RVT [44] , and FAN [89] . The DeiT backbone is the tiny configuration of 5.7M parameters. We train all models on clean ImageNet-1K for 300 epochs before evaluating their top-1 and top-5 accuracy on the test dataset under fast gradient sign method (FGSM) [22] , projected gradient descent (PGD) [42] , and simultaneous perturbation stochastic approximation (SPSA) [81] . We also present results for performance against Auto Attack [11] , which is an ensemble of auto PGD-Cross Entropy (APGD-CE), auto PGD-targeted (APGD-T), fast adaptive boundary-targeted (FAB-T), and Square. We display results for attacks individually and in default sequential mode.Results. Table 2 shows Elliptical attains top robustness in PGD and SPSA and second top in FGSM while achieving highly competitive clean accuracy. DeiT-Elliptical is particularly impressive under black box attack SPSA, improving over the next best model, RVT, [44] , by 10%. Table 4 shows results on Auto Attack [11] , where we see DeiT-Elliptical substantially outperforms standard DeiT in each attack individually and sequentially. We again see strong performance against black box attack Square with an 8.5% improvement. When combining with SOTA robust transformer, FAN [89] , Elliptical Attention improves robustness to sequential Auto Attack and all individual attacks except FAB-T, for which it still remains highly competitive. This shows Elliptical Attention can further boost robustness when combined with SOTA robust models.Experimental setup. We adopt the setup in [7] . For each of the 5 tasks, equation calculation (ListOps) [50] , review classification (Text) [41] , document retrieval (Retrieval) [61] , image classification (Image) [32] , and image spatial dependencies (Pathfinder) [37] , we compare Elliptical with standard Transformer [82] , Linformer [26] , Reformer [28] , Performer [9] , and Longformer [3] .Results. Elliptical Attention achieves top or second top test accuracy in every task and top overall performance. This shows Elliptical Attention learns superior representations across a wide range of modalities in long-range contexts. Experimental setup. We adopt the setup in [71] . The encoder is pretrained on ImageNet-1K following the same specification described in 4.2. In particular, the encoder is a DeiT-tiny backbone of 5.7M parameters pretrained for 300 epochs. After pretraining, we then attach a decoder that contains 2-layer masked transformer and finetune the full encoder-decoder model for 64 epochs on the ADE20K [88] image segmentation dataset.Results. Table 7 reports pixel accuracy, mean accuracy, and mean intersection over union (IOU). Elliptical Attention boosts performance across all metrics, with intersection over union, in particular, improving by a substantive 4.7%.Experimental setup. For experiments using Switch Transformer [18] and GLaM [17] backbones, we adopt the setup in [60] . In particular, we integrate Elliptical into small (70M parameters) and medium (220M parameters) GlaM backbones and train the models on WikiText-103 for 80 and 120 epochs, respectively. We consider Switch backbones at medium (220M parameters) and large (388M parameters) configurations, both trained for 80 epochs. All models use top-2 expert routing. For the standard transformer experiments, we continue with the setup of [54] and additionally present results for Elliptical in a medium configuration with 90M parameters trained for 100 epochs.Results. We present in Tables 5 and 6 the performance of Elliptical in MoE backbones. We see moving from smaller to larger configurations, Elliptical maintains strong, consistent improvements in test PPL. We note particularly substantive improvements with scale in the GLaM backbone, where at the small configuration Elliptical attains a 2.7% improvement, but at the medium configuration this performance improvement almost doubles to 5.0%. Table 8 further shows that in the standard transformer backbone, Elliptical maintains its substantive 6.8% improvement when scaling up to a larger configuration. These results show that Elliptical Attention scales well with model size. Theoretical Frameworks for Attention. Attention mechanisms have been studied from a range of perspectives. [80] shows that self-attention can be derived from kernel similarity functions, and [77] points out that self-attention projects its query vectors onto the principal component axes of its key matrix in a feature space. [56] formulates self-attention as the support vector expansion derived from a support vector regression problem, while [73] explains attention through nonlinear singular value decomposition of asymmetric kernels. Attention has also been explained through ordinary/partial differential equations, Gaussian mixture models, and graph-structured learning [40, 68, 51, 72, 20, 52, 31, 86] . [54, 23] show that self-attention performs Nadaraya-Watson regression with Gaussian isotropic kernels. This paper leverages this viewpoint and proposes modifying the Gaussian isotropic kernels to include a Mahalanobis metric which can be interpreted as stretching the hyper-spherical neighborhoods of the kernel to hyper-ellipsoids.Robust Transformers. In vision, [43] proposes an ensemble defense strategy to white-box 12 of Appendix F.attacks while [44] proposes position-aware attention scaling and patch-wise augmentation. Recently, [89] proposes a fully-attentional network to attain state-of-the-art accuracy on corrupted image data. In language, [85] proposes structurally aware table-text encoding, [38] proposes a robust end-to-end transformer for crisis detection, and [33] proposes duration-based hard attention. [6, 4] integrate a Gaussian process into attention for out-of-distribution detection, and [79] develops equivariant neural functional networks for transformers. These methodologies are motivated by their respective domain and tend to have limited generalizability to differing domains. Our approach, by contrast, proposes a general framework that makes no assumption on the downstream task and requires no additional parameters and negligible computational overhead.Mahalanobis Metrics. Mahalanobis metrics have been used predominantly in classical machine learning algorithms. In nearest-neighbor (NN) classification and regression, [84, 49] learn the metric through backpropagation. In NN KL divergence estimation, [58] learns a Mahalanobis metric from density approximation. In kernel regression, [57] takes eigenvalues of the estimated Jacobian while [29, 30] estimate coordinate-wise variability of the true function. Our model similarly uses coordinate-wise variability of the unknown function to form the Mahalanobis transformation but instead uses a more efficient estimator that does not require materializing the prediction function and accommodates the self-attention setting. In general, our method is among the early work in incorporating Mahalanobis metrics into the self-attention mechanism. 
proof of proposition 3: 
 The lemma below encapsulates the necessary calculations that will then be used in the following proofs.Lemma 2 Given a normally distributed zero mean random variable ξ ∼ N (0, σ 2 ), the expectation of a random variable obtained by its absolute value is E|ξ| = 2σ 2 /π.Proof. Since ξ ∼ N (0, σ 2 ), by definition of expectation, we haveE|ξ| = ∞ -∞ |x| √ 2πσ 2 exp - x 2 2σ 2 dx = 0 -∞ -x √ 2πσ 2 exp - x 2 2σ 2 dx + ∞ 0 x √ 2πσ 2 exp - x 2 2σ 2 dx (44) = 2 √ 2πσ 2 ∞ 0 x exp - x 2 2σ 2 dx (45) = 2 πσ 2 -σ 2 exp - x 2 2σ 2 ∞ 0 = 2σ 2 π ,where we used the variable change x ← (-x) in the first integral of 44 to obtain 45. □We derive the bounds for the impact of noise in 3, with respect to its variance, on our estimator 10 in Lemma 3. Henceforth, we omit the factor δ in Eqn. 10 since it does not affect the further analysis.Lemma 3 Given that the noise term in 3 follows a normal distribution with zero mean and variance σ 2 , the following inequalityEQUATIONholds for all i ∈ [D], where f i denotes the i th component of f(k ℓ ) = (f 1 (k), f 2 (k), . . . , f D (k)) ⊤ .Proof. Since all value vectors are taken from the data generating process 3, we have EQUATIONEQUATIONRecall that ϵ ℓ i ∼ N (0, σ 2 ) and independent. Now we have that ϵ ℓ+1 i -ϵ ℓ i ∼ N (0, 2σ 2 ) as the mean value does not change while variance accumulates when subtracting two zero-mean normal variables. Therefore, the Lemma 2 gives thatE|ϵ ℓ+1 i -ϵ ℓ i | = 2 √ π σ.Plugging this back into the inequalities 48, we getE|f i (k ℓ+1 ) -f i (k ℓ )| - 2 √ π σ ≤ m i ≤ E|f i (k ℓ+1 ) -f i (k ℓ )| + 2 √ π σ,which is equivalent to 46 as desired. □ Remark 7 Note that in Lemma 3, we may also take into account the possible noise in the value vectors. Let ϵ ℓ i,v ∼ N (0, σ 2 v ) be the noise in the values vectors as Proof of Proposition 3. We shall first make the following assumptions on the data generating process 3:m ϵ i = E|v ℓ+1 i -v ℓ i + ϵ ℓ+1 i,v -ϵ ℓ i,v |. 
conclusion and future work: 
 In this paper, we present Elliptical Attention, a novel variant of attention that computes a Mahalanobis transformation to stretch the underlying feature space in directions of high contextual relevance. This transformation can be interpreted as modifying the hyper-spherical neighborhoods around queries to hyper-ellipsoids which upweight the attention paid to keys lying in important directions, enabling the transformer to learn better and more robust representations. This approach makes no assumptions on the downstream task, requires no learnable parameters, and can be applied to any transformer to boost clean and robust performance. A limitation of our work is that we use the values over layers to estimate the average direction-wise gradient of the true self-attention function, which makes the estimate prone to noise. For ongoing work, we are exploring more precise estimation methods with provable convergence guarantees that do not compromise efficiency.  
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: Our research involves no human subjects or privacy concerns. Our research also has no clear links to discriminatory, unsafe, or harmful outcomes. Rather, as it is in large part concerned with robustness, particularly to adversarial attack, we hope our research might be usable to defend vital AI systems from ill-intentioned actors. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss broader impacts in Appendix G. We see our improved accuracy and robustness as offering societal benefits, for example with self-driving cars by our improved image segmentation results or with our adversarially robust vision models to defend against ill-intentioned actors.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper contains no such models with high risk of misuse such as pretrained language models, image generators, or scraped datasets. We use widely accepted, standard benchmark datasets and propose a fundamental and general improvement to a core architecture.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A Full Derivation of Self-Attention as Non-Parametric RegressionRecall NW estimator is a non-parametric estimator of the unknown f at any given query q described byf (k) = E[v|k] = R D v • p(v|k)dv = R D v • p(v, k) p(k) dv,where the first equality comes from the noise being zero mean, the second equality comes from the definition of conditional expectation and the final equality comes from the definition of conditional density. Eqn. 3 implies that if we can just obtain good estimates of the joint density p(v, k) and marginal density p(k) then we can estimate the required f (q). The Gaussian isotropic kernels with bandwidth σ are given byEQUATIONwhere φ σ is the multivariate Gaussian density function with diagonal covariance matrix σ 2 I D . Given the kernel density estimators in Eqn. 14, the unknown function can be estimated asfσ (k) = R D v • pσ (v, k) pσ (k) dv = R D v • j∈[N ] φ σ (v -v j )φ σ (k -k j ) j∈[N ] φ σ (k -k j ) dv = j∈[N ] φ σ (k -k j ) v • φ σ (v -v j )dv j∈[N ] φ σ (k -k j ) = j∈[N ] v j φ σ (k -k j ) j∈[N ] φ σ (k -k j ).Then, using the definition of the Gaussian isotropic kernel and evaluating the estimated function at q i we havef (q i ) = N j v j exp -∥q i -k j ∥ 2 /2σ 2 N j exp (-∥q i -k j ∥ 2 /2σ 2 ) = N j v j exp -(∥q i ∥ 2 + ∥k j ∥ 2 )/2σ 2 exp(q ⊤ i k j /σ 2 ) N j exp [-(∥q i ∥ 2 + ∥k j ∥ 2 )/2σ 2 ] exp(q ⊤ i k j /σ 2 ) = N j v j exp(q ⊤ i k j /σ 2 ) N j exp(q ⊤ i k j /σ 2 ) = j∈[N ] softmax(q ⊤ i k j /σ 2 )v j .Remark 6 Note that relaxing the assumption of normalized keys, the standard unnormalized selfattention score can be written asexp(q ⊤ i k j /σ 2 ) = exp -∥q i -k j ∥ 2 /2σ 2 exp (∥q i ∥ 2 + ∥k j ∥ 2 )/2σ 2 ∝ exp -∥q i -k j ∥ 2 /2σ 2 ,which shows that the dot-product self-attention scores are proportional to the NW kernel value with Euclidean distance. Hence the assumption of key normalization is sufficient to recover exactly the correspondence between self-attention and NW kernel regression, but not necessary. Analogously, the unnormalized Elliptical Attention score takes the following form:exp(q ⊤ i M k j /σ 2 ) = exp -d(q i , k j ) 2 /2σ 2 exp (∥q i ∥ 2 M + ∥k j ∥ 2 M )/2σ 2 ∝ exp -d(q i , k j ) 2 /2σ 2 ,where d(•, •) is the Mahalanobis distance used in Eqn. 7 and ∥ • ∥ M is the norm in the transformed space with metric d. This observation justifies the use of the transformed dot product instead of the full Mahalanobis distance metric in Eqn. 12 as it preserves the proportionality relationship between the attention computation and the corresponding nonparametric regression estimator with chosen distance metric.In this section, we present the omitted theorem statements and technical proofs in the main body of the paper.Let M : R D → R N be the transformed softmax operator as defined in Lemma 1. We wish to find its Jacobian matrix given byJ M (q) =       ∂M1(q) ∂q 1 ∂M1(q) ∂q 2 . . . ∂M1(q) ∂q D ∂M2(q) ∂q 1 ∂M2(q) ∂q 2 . . . ∂M2(q) ∂q D . . . . . . . . . . . . ∂M N (q) ∂q 1 ∂M N (q) ∂q 2 . . . ∂M N (q) ∂q D      , to measure the sensitivity of each output dimension to a change in each input dimension. Let M j : R D → R denote the j th component of the output vector for j ∈ [N ], that is, for a vector q ∈ R D ,EQUATION)Let q i and k i j denote the i th coordinates of vectors q and k j , respectively. Then,∂ ∂q i ln(M j (q)) = ∂ ∂q i   q ⊤ M k j -ln   s∈[N ] exp(q ⊤ M k s )     = m i k i j - s∈[N ] ∂ ∂q i exp(q ⊤ M k s ) s∈[N ] exp(q ⊤ M k s ) = m i k i j -m i s∈[N ] k i s exp(q ⊤ M k s ) s ′ ∈[N ] exp(q ⊤ M k s ′ ) = m i   k i j - s∈[N ] k i s M s (q)   .Since the output of Eqn. 16 consists of only positive components, we have∂ ∂q i M j (q) = ∂ ∂q i ln(M j (q)) • M j (q) = m i   k i j - s∈[N ] k i s M s (q)   M j (q).Therefore, the triangle inequality givesEQUATION)We now bound each term individually. Consider the terms j ̸ = s first. Since 0 ≤ M s (q) ≤ 1, we can bound them asEQUATION) Now recall that the inequality ab ≤ (a + b) 2 /4 holds for any real numbers a and b with equality holding at a = b. Therefore, for the first term, we obtainEQUATIONCombining inequalities 17, 18 and 19, we finally arrive atEQUATIONfor all i ∈ [D] and j ∈ [N ], where κ ij ≥ 0 denotes the coefficient in the bracket. □Let us estimate the distance between two output vectors of Elliptical attention mechanism corresponding to clean and contaminated query inputs, namely:h = j∈[N ] softmax(q ⊤ M k j /σ 2 )v j = j∈[N ] M j (q)v j h ϵ = j∈[N ] softmax((q + ϵ) ⊤ M k j /σ 2 )v j = j∈[N ] M j (q + ϵ) v j ,where M is defined as in Lemma 1. We omit the keys and scaling parameter for convenience since they do not affect the analysis. Then,∥h -h ϵ ∥ = j∈[N ] (M j (q) -M j (q + ϵ)) v j ≤ j∈[N ] |M j (q) -M j (q + ϵ)| ∥v j ∥ ≤ j∈[N ] ∥∇M j (q)∥ ∥ϵ∥ ∥v j ∥ (21) = j∈[N ] i∈[D] (J M (q) ji ) 2 ∥v j ∥ ∥ϵ∥ ≤ j∈[N ] i∈[D] κ 2 ij m 2 i ∥v j ∥ ∥ϵ∥ (22) = j∈[N ] tr(K 2 j M 2 )∥v j ∥ ∥ϵ∥ ,where K j := diag(κ 1j , κ 2j , . . . , κ Dj ) and κ ij is defined as in Eqn. 20. Note that 21 follows from mean value theorem for some β ∈ [0, 1] and q := q + βϵ while 22 follows from Lemma 1. □It should be noted that Proposition 1 addresses the impact of noise exclusively on the query vectors. However, the resulting bound can be extended to account for noise in all tokens by employing the same technique utilized in the proof. For completeness, we also provide the extension. LetEQUATIONAgain, take the difference between output vectors calculated from clean and noisy tokens as followsEQUATIONEQUATIONLet ∥ε∥ := max{∥ϵ q ∥, ∥ϵ k ∥, ∥ϵ v ∥} denote the noise with the largest norm among query, key and value noises. Then,EQUATIONEQUATION)Following the same steps as Lemma 1, one can derive the boundEQUATIONfor s, j ∈ [N ]. Therefore, we obtainEQUATIONwhere C s,j are the diagonal matrices whose elements are the proportionality coefficients in the derived upper bounds.There are two avenues through which to see resistance to representation collapse. In this section, we provide a proof based on noise propagation through layers, which decreases representation capacity as representations in deeper layers are increasingly composed of uninformative noise. We refer the reader to Appendix B.4 for an additional lens on representation collapse, where we show that Elliptical Attention is more sensitive to the variation and local features of the underlying function.Let the output at layer ℓ be denoted as h ℓ , the standard self-attention estimator and Elliptical estimator fitted at layer ℓ be denoted f ℓ and f ℓ d respectively, where d is the Mahalanobis metric described in Eqn. 7, and f be the true underlying function described in Eqn. 3. By assumption, f is a higher variance estimator than fd for any layer. The output for either estimator at layer ℓ can be decomposed into ground truth and noise as follows:EQUATIONEQUATIONwhere η ℓ ∼ γ(0, V η ), ϵ ℓ ∼ γ(0, V ϵ ) are the noise components of the estimate at q ℓ and f (q ℓ ) is the ground truth. By assumption of fd being lower variance, V ϵ -V η is a positive semi-definite matrix.We first require the following Assumption 1, which is described as:Assumption 1 (Random Input Noise Causes Estimator Attenuation) . Let f be any estimator of true function f and let the input x ∼ µ drawn from marginal µ be randomly corrupted by random noise ϵ ∼ (0, V ) of some unknown distribution and variance matrix V . Let c be some constant. Then, random input noise attenuates the estimator as follows:EQUATIONAssumption 1 is a well-studied phenomenon in parametric regression, often referred to as attenuation bias [69] , regression dilution [19] , or errors-in-variables [34] . In parametric regression, it can be shown to have an exact form where the estimated gradients of the model are attenuated towards 0 proportional to the variance of the noise ϵ. In non-parametric regression, addition of input noise is often referred to as random smoothing or random input smoothing [46, 10] , and is well known to be used as regularization technique to introduce bias into the model. In non-parametric models, no exact closed forms exist to express the attenuation bias, but for our purposes we only note the attenuation exists and provide a general form of it in Assumption 1.The outputs of 30 and 31 then become the inputs to the following layer after being self-added, normalized, projected, and linearly transformed. For notational simplicity and because these operations do not change the analysis, we denote the input at the next layer as the previous layer output q ℓ+1 = h ℓ . We therefore have the following process:EQUATIONwhere we see the output h ℓ+1 is obtained by fitting f ℓ to input z ℓ which is composed of ground truth f (q ℓ ) and noise ϵ ℓ passed through from the previous layer.The result then follows directly from the fact that in any given layer, the standard self-attention estimator produces noisier estimates, where that noise is then passed into the subsequent layer as input noise. This isEQUATIONEQUATIONEQUATIONEQUATIONwhere line 35 follows from combining the fact that η ℓ is lower variance with Assumption 1 and line 36 follows from the fact that E∥X∥ ≈ E∥Y ∥ when X, Y have the same mean and roughly similar distribution.Therefore we obtain at any layer ℓ the followingEQUATIONas required. □To further substantiate our findings on the mitigation of representation collapse in transformers, we now present an additional proposition that examines this phenomenon from a different perspective.In Proposition 4, we show that Elliptical attention reduces representation collapse by retaining the important local features (bumps etc.) better than the standard self-attention in the case of sparse piece-wise constant functions.Proposition 4 (Representation Collapse) Let f : A → R D for A ⊆ R D be a piece-wise constant function with f | Ai (q) = f i ∈ R D where A = i∈I A i for some (possibly infinite) index I. Let q 1 and q 2 be the queries lying in any of the adjacent domain pieces with distant function values. Then, the Elliptical estimates at these queries retain the distance better than the standard self-attention estimates which is formulated asEQUATIONProof. Assume all output vectors are normalized. Then, the Euclidean distance between two vectors is determined by their dot product sinceEQUATIONWithout loss of generality, we take A 1 and A 2 to be the two adjacent pieces so that f (q i ) = f i for i = 1, 2. Denote f (q i ) = h i and fd (q i ) = h id . Then, Eqn. 39 is equivalent to provingEQUATIONwhere the expectation is taken over the whole sampling distribution D but the points q 1 ∈ A 1 and q 2 ∈ A 2 are fixed as described in the definition. We drop the subscript D as this will be the default distribution for computing expectation unless specified otherwise. Let r S = cossim(f 1 , f 2 ) = f ⊤ 1 f 2 be the cosine similarity of the two piece-wise values. By definition of q 1 and q 2 and since the estimates work by averaging the output vectors with a small amount of noise, we haver S ≤ min E[h ⊤ 1d h 2d ], E[h ⊤ 1 h 2 ]. We now decompose h 1d and h 2d in terms of components along and orthogonal to f 1 and f 2 , respectively:EQUATIONwhere f ⊤ i f ⊥ i = 0. Then, for their dot product, we haveEQUATION)The analogous decomposition of h ⊤ 1 h 2 can be obtained. By Theorem 2 we have that the Elliptical estimator is lower variance and so we have E∥f i -h id ∥ 2 ≤ E∥f i -h i ∥ 2 . This has the following implications:1. 1 ≥ E[h ⊤ id f i ] ≥ E[h ⊤ i f i ] i.e. the component of h id along f i is larger than that of h i , and hence h ⊤ id f i is closer to 1. 2. Due to the first implication above, the orthogonal component f ⊥ i becomes smaller in terms of magnitude so that f ⊤ j f ⊥ i and (f ⊥ i ) ⊤ f ⊥ j are closer to 0 for Elliptical compared to the standard self-attention.⊤ 1d h 2d is closer to 1 • (f ⊤ 1 f 2 ) + 0 = f ⊤ 1 f 2 = r S which, by definition, is the smallest dot product over S, and hence,r S ≤ E[h ⊤ 1d h 2d ] ≤ E[h ⊤ 1 h 2 ] as desired. □ B.E|v ℓ+1 -v ℓ | -E|ϵ ℓ+1 i,v -ϵ ℓ i,v | ≤ m ϵ i ≤ E|v ℓ+1 -v ℓ | + E|ϵ ℓ+1 i,v -ϵ ℓ i,v |.m ϵ i -E|f (k ℓ+1 ) -f (k ℓ )| ≤ 2 √ π (σ + σ v ).Assumption 2 The underlying coordinate system in the feature space X k is independent, implying that the function f : R D → R D in Eqn. 3 can be separated asf (k) = (f 1 (k 1 ), . . . f D (k D )) ⊤ Assumption 3The noise term in Eqn. 3 has independent components with each component ϵ ℓ j following a normal distribution N (0, σ 2 ) for small σ, for all j ∈ [D] and ℓ ∈ N Assumption 4 The magnitude of each component of key perturbations across consecutive layers, defined as |k ℓ+1 i -k ℓ i |, follows a distribution with small, layer-independent mean (δ) and variance (ν)The assumption of layer-independence in Assumption 4, especially for deeper layers, is supported well empirically, as shown in Figure 5 . Given the over-smoothing observed in transformers [24] , where token representations stabilize after initial few layers, it is also practical to assume that key perturbations across layers have relatively small mean and variance when modelled as a random process.Proof. Under the Assumptions 2, 3, 4, we show that ∥f ′ i ∥ 1,µ ≥ ∥f ′ j ∥ 1,µ implies m i ≥ m j with high probability where m i is defined as in (10) .Directly from the Lemma 3, we havem i -E|f i (k ℓ+1 ) -f i (k ℓ )| ≤ 2 √ π σ.Letting σ → 0 in this inequality, which is feasible under the Assumption 3, one can get with a small error thatEQUATIONwhich in turn implies that the impact of the noise in ( 10) is negligible and the error of ignoring them can be controlled by the bounds given by (46) . Now according to the theorem statement,EQUATIONwhere we used the separability of f as given in Assumption 2 which simplifies the Jacobian matrix asJ f (k) =       ∂f1(k) ∂k1 ∂f1(k) ∂k2 . . . ∂f1(k) ∂k D ∂f2(k) ∂k1 ∂f2(k) ∂k2 . . . ∂f2(k) ∂k D . . . . . . . . . . . . ∂f D (k) ∂k1 ∂f D (k) ∂k2 . . . ∂f D (k) ∂k D       =       ∂f1(k1) ∂k1 ∂f1(k1) ∂k2 . . . ∂f1(k1) ∂k D ∂f2(k2) ∂k1 ∂f2(k2) ∂k2 . . . ∂f2(k2) ∂k D . . . . . . . . . . . . ∂f D (k D ) ∂k1 ∂f D (k D ) ∂k2 . . . ∂f D (k D ) ∂k D       =     f ′ 1 (k 1 ) 0 . . . 0 0 f ′ 2 (k 2 ) . . . 0 . . . . . . . . . . . . 0 0 . . . f ′ D (k D )     , so that [J f (k)] ii = f ′ i (k i ).Using the definition of derivative, the inequality 50 is equivalent toEQUATIONNext, we note that for a small δ, the limits in ( 51) can be approximated withEQUATION) Let us choose δ = E|k ℓ+1 i -k ℓ i |.Then, by Chebyshev's inequality, we have for anyEQUATION) Given that the variance ν is sufficiently small as in the Assumption 4, the inequality (53) implies that k ℓ+1 i ≈ k ℓ i ± δ with high probability. Therefore, it follows from ( 52) with high probability thatE|f i (k ℓ+1 i ) -f i (k ℓ i )| δ ≥ E|f j (k ℓ+1 j ) -f j (k ℓ j )| δ, which, due to 49, is equivalent to m i ≥ m j as desired. □ B.6 Lipschitz smoothness in (X , d) Below we show how Lipschitz smoothness of f changes when moving from Euclidean to the Mahalanobis transformed space. We shall follow similar steps to [29] and [30] but for a more general class of functions.Proposition 5 (Change in Lipschitz smoothness for f ) Suppose there exists a positive constant G i such that ∥∇f i (k)∥ ≤ G i for any k ∈ X k and m i > 0 for all i ∈ [D]. Then for any q, k ∈ X k , the following inequality holds:∥f (q) -f (k)∥ ≤   i∈[D] G i √ m i   d(q, k).Proof. Let ω := q-k ∥q-k∥ denote the unit vector pointing from k to q. The fundamental theorem of calculus implies thatf (q) -f (k) = ∥q-k∥ 0 d dt f (k + tω) dt = ∥q-k∥ 0 ω ⊤ J f (k + tω) dt,where J f is the Jacobian matrix of f as usual. Starting with the distance between outputs f (q) and f (k) we haveEQUATIONwhere, as for all other vectors, q i denotes the i th component of vector q. Now note thatEQUATIONCombining 54 and 55, we finally attainEQUATIONwhich completes the proof. □The following Theorem 1 is a classic result from [70] . We refer the reader to their work for details.Theorem 1 (Minimax rate for functions of bounded variability [70] )EQUATIONTheorem 2 (Improvement in MSE for approximately sparse functions [30] ) Let the norm of the largest gradient be λ := sup i∈[D] ∥∇f i (q)∥ sup and fd be an estimator in metric space (X q , d) where d is defined as Eqn. 7. Then,EQUATIONProof. We provide an abridged proof for completeness. We refer the reader to [30] for the full details.First, the full bound is described as follows:EQUATIONwhere d(X ) is the d-diameter of X defined assup x,x ′ ∈X d(x, x ′ ), R ⊂ [D], 1 ≤ C κ R ≤ C ′ (4κ R ) |R| , C and C 1 are universal constants and λ d ≥ sup i ∥f ′ i ∥ sup / √ m i . Let r(ϵ) ≤ |R| if ϵ ≥ ϵ ̸ R /d(X ) D -(D -|R|) log(d(X )/ϵ ̸ R ) log(1/ϵ) if ϵ < ϵ ̸ R /d(X ) .For bandwidth ϵ n , r = r(ϵ n ) and let |R| ≤ r ≤ D. Let ϵ > 0, c be defined as the same c in Theorem 1, and n ∈ N, define the function ψ n,d = Cϵ -r(ϵ) /n and ψ n,¡ d (ϵ) = C ′ 1 ϵ -D /n where C 1 = c (λ/Cλ d d(X )) D . Also define ϕ(ϵ) = C 2 D 2 λ 2 d d(X ) 2 • ϵ 2 . For any fixed n, let ϵ n,¡ d be a solution to ψ n, ¡ d (ϵ) = ϕ(ϵ). Solving for ϵ n, ¡ d obtains the following lower bound on the minmax rate ofEQUATIONFor any n ∈ N there exists a solution ϵ n,d to the equation ψ n,d (ϵ) = ϕ(ϵ) since r(ϵ) is nondecreasing. Therefore it is possible to obtain the following:EQUATION)Since ϕ is independent of n, and both ψ n,d and ψ n, ¡ d are strictly decreasing functions of n, we have that ϵ n,d and ϵ n, ¡ d both tend to 0 as n → ∞. Therefore we can define n 0 such that, for all n ≥ n 0 , both ϵ n,d and ϵ n,¡ d are less than ϵ & R /d(X ). Thus, ∀n ≥ n 0 , we have ϵ n,d < ϵ n, ¡ d if, for all 0 < ϵ < ϵ & R /d(X ), ψ n,d (ϵ) < ψ n, ¡ d (ϵ), which completes the proof □.In this section, we present a consistent centered difference-based quotient estimator of the coordinatewise variability obtained by perturbing the estimated function in the i th direction and measuring the L 1 norm of the difference. Similarly, this estimator requires no learnable parameters or gradients. The estimator is described in the following proposition.Proposition 6 (Consistent Estimator) Given a function f : R D → R Dv with ith directional varia- tion ∥f ′ i ∥ 1,µ , i ∈ [D], the directional variation can be estimated by the quantityEQUATION)where t is a hyperparameter controlling the degree of locality of the estimator and E n denotes the empirical expectation for n samples.Despite m i in proposition 6's simple formulation, it is nonetheless a consistent estimator of the coordinate-wise variation in the underlying function. We utilize a simplified version of a theorem from [30] , adapted to suit our specific needs, as the original formulation is more detailed than necessary for our purposes.Theorem 3 (Consistency of Centered Difference-based Estimator for Scalar Function [30] ) Let φ : R D → R be a smooth scalar function and ∥φ ′ i ∥ 1,µ := E x∼µ |e ⊤ i ∇φ| be the coordinate-wise variability for that scalar function. Then, for any direction i and any 0 < δ < 1/2, the following bound holds with probability of at least 1 -2δ:EQUATIONNote that the Theorem 3 is different from our setting by studying a scalar function as opposed to a vector valued function. However, we show that the result can be generalized to the latter case in Corollary 1 below via the estimator 62.Corollary 1 (Consistency of the Estimator (62) for Vector-valued Function) Let f : R D → R Dv be a vector valued function and ∥f ′ i ∥ 1,µ be defined as in Definition 1. Then, for any direction i and any 0 < δ < 1/2, the following bound holds with probability of at least 1 -2δ:EQUATIONProof. We first derive the relation between the left hand side of 64 and its coordinate-wise differences as follows:EQUATION=EQUATION) where line 65 follows from the definition of the ℓ 1 norm, line 66 follows from the linearity of expectation, the superscript j indicates that the case is reduced to the scalar function case for each j th summand individually. Note that the probability of the last bound is at least (1 -2δ/D) D since each component-wise bound holds with probability at least 1 -2δ/D. However, since we can choose δ small enough such that 2δ < 1, by Bernoulli's inequality (1 Remark 9 Despite the proven consistency of this estimator, we opt for the efficient estimator presented in our main body described in Eqn 10. This is because the consistent estimator requires materialising the prediction function -that is, computing a forward pass of the self-attention mechanism -twice per dimension. This makes the consistent estimator unusable in most problem settings. We present results for the consistent estimator in Appendix F.11.-2δ/D) D ≥ 1 -2Dδ/D = 1 -2δ. □Training and Inference. Given Elliptical Attention requires keys and values from the previous layer in order to compute the required transformation, we can only implement Elliptical Attention from the second layer on. We incorporate our Elliptical Attention into both training and inference stages. This is because, firstly, Elliptical Attention is designed to offer improvements to both clean and contaminated data, and so even in the presence of completely clean train and test data, it is advantageous to incorporate Elliptical Attention into both stages. Secondly, it is commonplace to encounter data contamination in test data and indeed also highly possible to encounter it in train data as well. Therefore, in the interest of robustness as well, we also incorporate Elliptical Attention into both stages.Computational Efficiency. Computing the required transformation requires no learnable parameters and is obtained simply by averaging absolute differences in values over layers. These operations are therefore just of the order O(bhnD) = O(n) for batch size b, head number h, key/value length n, and dimension D. Hence upper-bound time complexity of the overall Transformer is unaffected. We provide efficiency analysis in terms of computation speed and max GPU memory allocated (calculated by CUDA max_memory_allocated in Figure 4 , which shows that compared with baseline robust models, Elliptical is the fastest and most memory efficient. Elliptical exhibits no perceptible slowdown versus DeiT of the same configuration and only a 0.99% increase in max memory allocated, which is why Elliptical and DeiT are shown as the same data point in the Figure 4 . ImageNet-C consists of 15 types of algorithmically generated corruptions with 5 levels of severity (e.g blurring, pixelation, speckle noise etc). Given that Elliptical Attention learns attention weights dependant on the transformation M , which is itself dependant on the train data distribution, our proposed model is not designed for situations in which the test distribution is substantially different from the train distribution. This then includes OOD robustness and robustness to heavy corruption to the point where the underlying data distribution is fundamentally different. We nonetheless evaluate Elliptical Attention on ImageNet-A,R,C to assess these important forms of robustness as well. Table 9 shows that Elliptical Attention is still able to offer improvements over baseline DeiT in terms of OOD robustness, while maintaining approximately the same performance as the baseline for ImageNet-C. Figure 7 shows for Fog and Pixelate corruptions how Elliptical compares with DeiT over the 5 severity levels, where we see that at low severity levels Elliptical improves over DeiT, however as the severity level gets too high Elliptical falls behind. This agrees with our expectation that as the severity level grows, the distribution is further shifted relative to the train distribution and so Elliptical Attention is unable to improve performance.We provide in Figure 6 additional representation collapse results for ImageNet and ADE20K, showing that across modalities Elliptical Attention resists representation collapse. Corruption. Word Swap Text Attack3 corrupts the data by substituting random words with a generic token 'AAA'. We follow the setup of [23] and assess models by training them on clean data before attacking only the evaluation set using a substitution rate of 2.5%.Model, Optimizer & Train Specification. We adopt the training regime of [54] . To this end, the small backbone uses 16 layers, 8 heads of dimension 16, a feedforward layer of size 2048 and an embedding dimension of 128. We use a dropout rate of 0.1. We trained with Adam using a starting learning rate of 0.00025 and cosine scheduling under default PyTorch settings. We used a batch size of 96 and trained for 120 epochs and 2000 warmup steps. The train and evaluation target lengths were set to 256.The medium backbone uses 16 layers, 8 heads of dimension 32, a feedforward layer of size 2048 and embedding dimension of 256. We use a dropout rate of 0.1. We trained with Adam using a starting Compute Resources. All models are trained and evaluated on a single NVIDIA A100 SXM4 40GB GPU.Ablation Models. We consider the following models in our ablation studies:• Random Ablation. To validate the efficacy of our proposed estimator given in Eqn. 10, we consider an alternate model in which M is populated by weights uniformly drawn from the [0, 1] interval followed by the same maxscaling as in Elliptical.• Elliptical-Meanscale. We ablate the effect of maxscaling by considering meanscaling of the estimates m i . That is, each m i ← m i / m is scaled by the mean variability estimate m = E D [m i ].• Elliptical-Consistent. We consider also the performance of Elliptical when using the consistent estimator of ∥f ′ i ∥ 1,µ described by Equation 62.Language Modelling. Results are shown in Table 19 . Amazingly, the random ablation model performs extremely well on contaminated data. In general, this most likely suggests that training a model with randomness injected into the attention matrix can generate some robustness benefits, which is intuitive. It does, less surprisingly, come at the cost of clean data performance, where Random Ablation performs almost 10% worse than baseline transformer.Algorithm 1 presents a pseudocode for implementing Elliptical Attention as given by Eqn. 13 on top of conventional self-attention.ImageNet Classification and Attack. Require:1: Tensor Q ∈ R N ×D ▷ current layer queries 2: Tensor K ∈ R N ×D ▷ current layer keys 3: Tensor V ∈ R N ×D ▷ current layer values 4: Tensor V prev ∈ R N ×D ▷ previous layer values 5: float δ ∈ R + ▷ step size 6: integer D ∈ N ▷ head dimension 7: function ELLIPTICAL_ATTENTION(Q, K, V , V prev , δ, D) 8: M ← ELLIPTICAL_WEIGHTS(V , V prev , δ) ▷ compute weight matrix M 9: logits ← Q × M × K ⊤ × 1 √ D▷ modify the dot-product computation N ← V .size(0) ▷ sequence length 17: value_diff ← (V -V prev )/δ 18: M ← 1 N × NORM(value_diff, p = 1, dim = 0) ▷ column-wise average of L 1 norms 19:M ← DIAG_EMBED(M ) ▷ embed the vector into a diagonal matrix 20:return M 21: end function Our research offers benefits to both clean data and robust performance. We in particular show improved results in domains with wide social applicability. These include image segmentation, with benefits to self-driving cars, and language modeling, with benefits to AI chatbot assistants. We in particular show strong improvements against contamination by adversarial attack, which we hope can protect vital AI systems from malicious actors, and competitive performance in contaminated language modeling, which we hope can improve language models evaluated on imperfect data as is often the case in the real world. There is always possibility of misuse of AI systems, however our research shows substantive improvements in fundamental architectures and theory which we hope can spur further socially beneficial outcomes. Justification: We provide an efficiency analysis figure (Figure 4 ) containing computation time and memory resources. We also provide in Appenxix F the exact GPU resources used. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).www.salesforce.com/products/einstein/ai-research/the-wikitext-dependency-language-modeling-dataset/Implementation available at github.com/QData/TextAttack 


Decision: This paper introduces the use of a Mahalanobis distance metric for computing attention weights, enabling the stretching of the underlying feature space in directions of high contextual relevance. The authors define a hyper-ellipsoidal neighborhood around each query to enhance the attention weights of tokens located in contextually significant directions.

While there are potential issues, such as the lack of detailed experimental setups, unclear model formulation, and areas needing improvement in writing and organization, this paper demonstrates strong theoretical originality. The shift to hyper-ellipsoidal attention regions is particularly novel. Additionally, the experiments are extensive, covering object classification, image segmentation, and language modeling across different data modalities. All reviewers have unanimously recommended this paper for acceptance.

Reviews: review: paper_summary: paper_summary: This paper propose a novel attention mechanism, named Elliptical Attention.
Elliptical Attention use a Mahalanobis distance metric to stretch the underlying feature space in directions of high contextual relevance.
The Elliptical Attention pays more attention to contextually relevant information, rather than focusing on some small subset of informative  features.
The Elliptical Attention can reduce the representation collapse and enhance the model’s robustness.
This method has theoretical support and extensive experimental results on object classification, image segmentation, and language modeling validate its superiority.

This is a fundamental research.
strength_weakness: strength_weakness: 1. The experimental table does not provide the parameter count and FLOPS.
The reviewer thinks these comparative data will provide a clearer understanding of the differences in method performance.

2. The state-of-the-art attention methods are not complete enough.
There are many improvements to classical attention mechanisms that were not compared in this paper, such as B-Attention in NeurIPS 2022. B-Attention utilizes the relationship between neighbours and improves the weight distribution of the attention mechanism.
questions: questions: The large language model is very popular now.
The reviewer is curious whether this Elliptical Attention mechanism complies with scaling laws.
Does the Elliptical Attention have the potential to replace self-attention in LLM?
limitations: limitations: No.
paper_summary: paper_summary: The paper proposes a new class of self-attention mechanism for transformers. It uses Mahalanobis distance to form hyper-ellipsoidal attention regions around queries, aiming to improve model robustness and reduce representation collapse. This approach is demonstrated to be effective across various tasks like language modeling, image classification, and segmentation.
strength_weakness: strength_weakness: (1) This paper fails to provide detailed setups of experiments, making it challenging to refer to other researchers and developers.

(2) The proposed method is unclear for readers. I have not found the exact formulation of elliptical attention. Moreover, it would be helpful to include pseudo codes of core algorithms.

(3) The Fig 3 and Fig 4 are in low resolution and would be blurred after zooming in.

(4) This paper is not well written, with disconnected organization and many abbreviations. 

(5) Fig 2 is pretty confusing. What does x1 and x2 mean here? What is the insight of these two subfigures? The authors failed to provide enough explanation of this figure.
questions: questions: N/A
limitations: limitations: N/A
paper_summary: paper_summary: This manuscript introduces Elliptical Attention, a new approach employing the Mahalanobis distance metric to calculate attention weights. This method delineates a hyper-ellipsoidal neighborhood around each query, amplifying the attention weights of tokens situated in contextually pivotal directions. When compared with conventional self-attention mechanisms, Elliptical Attention exhibits a reduction in representation collapse and enhances the model's robustness.
strength_weakness: strength_weakness: My primary reservations pertain to the experimental framework, which necessitates a more meticulous comparative analysis. Firstly, as Elliptical Attention is posited as an alternative to traditional attention computations, its juxtaposition with the standard Euclidean distance-based self-attention should be more exhaustive. For instance, within the ImageNet-1K classification task, it would be instructive to present results for both the vanilla Vision Transformer (ViT) and a modified ViT-Elliptical, wherein the standard self-attention is supplanted by Elliptical Attention. Further empirical comparisons between DeiT and DeiT-Elliptical across various model sizes and image resolutions on the ImageNet-1K classification task would substantiate the proposed method's advantages over the conventional self-attention.

Secondly, despite assertions in the introduction regarding Elliptical Attention's reduced memory requirements and accelerated computational speed, Figure 4 shows negligible efficiency gains over the original DeiT. I would suggest the authors to furnish detailed results such as memory usage, number of model parameters, FLOPs, and throughputs for a lucid comparison. This can be done using ImageNet. They may also provide a side-by-side assessment of ViT against ViT-Elliptical.
questions: questions: My question is to show more detailed experimental comparisons with the standard self-attention using different backbones.
limitations: limitations: The main limitation of this work is the need for a more extensive experimental comparison.

Comments: title: Thanks for your endorsement!
Response
Thanks for your endorsement!
Thanks for your endorsement!
After Rebuttal
Any Questions from Reviewer 5EPe on Our Rebuttal?
Any Questions from Reviewer ySYF on Our Rebuttal?
Any Questions from Reviewer EhBA on Our Rebuttal?
comment: Thanks for your response, and we appreciate your endorsement.
Thank you for your detailed answer !

I have read your reply and will consider it.
Thanks for your response, and we appreciate your endorsement.
Thanks for providing more detailed experimental comparisons in the response. My concerns are well addressed. I thus keep the accept rating.
Thanks for your response, and we appreciate your endorsement. As you suggested, we will include the algorithm for Elliptical Attention in our main manuscript.
Thank you for the detailed response and revisions from the authors. Points 1, 2, 3, and 5 have been well addressed. Please include the algorithm in the main manuscript to facilitate readers’ understanding of the method. I would like to increase the score to 5.
We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.
We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.
We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.