On the Long Range Abilities of Transformers (2311.16620v1)

Abstract: Despite their dominance in modern DL and, especially, NLP domains, transformer architectures exhibit sub-optimal performance on long-range tasks compared to recent layers that are specifically designed for this purpose. In this work, drawing inspiration from key attributes of long-range layers, such as state-space layers, linear RNN layers, and global convolution layers, we demonstrate that minimal modifications to the transformer architecture can significantly enhance performance on the Long Range Arena (LRA) benchmark, thus narrowing the gap with these specialized layers. We identify that two key principles for long-range tasks are (i) incorporating an inductive bias towards smoothness, and (ii) locality. As we show, integrating these ideas into the attention mechanism improves results with a negligible amount of additional computation and without any additional trainable parameters. Our theory and experiments also shed light on the reasons for the inferior performance of transformers on long-range tasks and identify critical properties that are essential for successfully capturing long-range dependencies.

Standard Transformer architectures, despite their widespread success, often exhibit suboptimal performance on tasks requiring the modeling of long-range dependencies compared to architectures specifically designed for this purpose, such as state-space models (SSMs) or models incorporating global convolutions (Zimerman et al., 2023 ). The work "On the Long Range Abilities of Transformers" (Zimerman et al., 2023 ) investigates the underlying reasons for this performance gap and proposes minimal modifications to the standard attention mechanism to enhance its long-range capabilities, particularly evaluated on the Long Range Arena (LRA) benchmark.

Analysis of Transformer Limitations in Long-Range Tasks

The paper posits that the limitations of standard Transformers on long-range tasks are not primarily due to fundamental issues with expressiveness or optimization. Theoretically, it is shown that a single Transformer layer possesses sufficient capacity to express any state-space layer, and by extension, any global convolution kernel (Appendix B, Theorem 1) (Zimerman et al., 2023 ). Furthermore, architectural features like parallel token processing, layer normalization, softmax activation, and residual connections contribute to stable optimization dynamics, mitigating issues like vanishing or exploding gradients often associated with recurrent architectures.

Instead, the core deficiency identified is related to generalization. Standard Transformers appear to lack an appropriate inductive bias for long-range sequential data, leading them towards hypothesis classes that perform poorly on unseen long sequences, often manifesting as overfitting. This contrasts with successful long-range models which typically incorporate strong structural priors. The observation that Transformers can achieve high accuracy on LRA tasks when trained directly on validation data further supports the idea that the issue lies in generalization from the training set rather than an inherent inability to model the required functions (Zimerman et al., 2023 ).

Key Principles from Specialized Long-Range Models

Analysis of architectures excelling at long-range tasks, including SSMs (like S4), linear RNN layers (like EMA), and global/long convolution layers (like SGConv, Mega), reveals two recurring principles contributing to their effectiveness:

1. Locality Bias (Exponential Decay Structure): Many successful long-range models implicitly or explicitly enforce a structure where the influence between tokens decays, often exponentially, with distance. While seemingly counterintuitive for long-range modeling, this local bias allows models to capture complex dependencies hierarchically by combining local interactions effectively. The exponential decay appears particularly well-suited compared to other decay forms like linear decay.
2. Smoothness Regularization: The operators or kernels employed in these models often exhibit smoothness properties. This can arise from specific parameterizations (e.g., SSMs mapping inputs through smooth functions) or explicit regularization techniques that favor smoother transformations, potentially acting as a form of implicit regularization that improves generalization on structured sequential data.

The LaS-Attention Mechanism

To integrate these principles directly into the Transformer's self-attention mechanism with minimal disruption, the paper proposes Local and Smooth (LaS) Attention. This modification adjusts the standard scaled dot-product attention calculation without introducing any new learnable parameters or significant computational overhead.

The standard attention computation is:

$$Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$$

The LaS-Attention modifies this as follows:

$$LaS-Attention_c(Q, K, V) = AP(softmax(exp(-\alpha_c D_L) \odot \frac{QK^T}{\sqrt{d_k}})) V$$

Where:

• $c$ denotes the attention head index.
• $\odot$ represents element-wise multiplication (Hadamard product).
• Locality Incorporation (ELD Operator): The term $exp(-\alpha_c D_L)$ introduces the locality bias.
  • $D_L$ is a matrix representing pairwise distances between token positions, potentially incorporating a causal mask (for autoregressive tasks) where $D_L(i, j) = |i - j|$ if $i \ge j$ and $\infty$ otherwise.
  • $\alpha_c$ is a non-learnable hyperparameter specific to head $c$, controlling the rate of exponential decay. Different heads employ different $\alpha_c$ values (e.g., initialized uniformly in $[0, B]$ on a log scale, with one head often having $\alpha_0=0$ to retain a standard global attention view), allowing the model to capture dependencies across multiple locality scales simultaneously. This Exponentially Locally Decay (ELD) operator biases the attention mechanism to prioritize nearer tokens before the softmax normalization.
• Smoothness Incorporation (AP Operator): The $AP(\cdot)$ operator denotes a 1-D Average Pooling operation applied independently to each row of the attention matrix after the softmax normalization. This operation uses a small pooling kernel (e.g., size 3) with padding to maintain the matrix dimensions. Applying average pooling to the attention weights $A_{ij}$ for each query $i$ promotes smoother attention distributions, effectively regularizing the attention mechanism by reducing sharp discontinuities in how attention is allocated across keys.

These modifications are designed to be computationally inexpensive, primarily involving element-wise multiplication and a simple average pooling step, adding negligible overhead compared to the dominant matrix multiplications in standard attention.

Empirical Performance on the Long Range Arena (LRA)

Experiments conducted on the LRA benchmark demonstrate the effectiveness of LaS-Attention (Zimerman et al., 2023 ):

• LaS-Attention vs. Other Transformers: LaS-Attention significantly outperforms the baseline vanilla Transformer and various efficient Transformer approximations (Reformer, Linformer, Performer, Luna). It achieves an average LRA accuracy of 73.99%. This compares favorably to the baseline Transformer (~61.5%) and other efficient variants, which typically score in the range of 50-62%. The improvement is particularly pronounced on tasks involving structured signals like the Image task (CIFAR10 classification based on pixel sequences), where LaS-Attention surpasses the next best Transformer variant by over 22%, indicating the benefit of the locality/smoothness bias for such data modalities.
• LaS-Attention vs. Specialized Layers: While LaS-Attention substantially narrows the performance gap, it does not surpass the top-performing specialized architectures like S4, MEGA, and SGConv, which report average LRA accuracies exceeding 80-88%. The paper notes that many of these leading models utilize bidirectional contexts, whereas the presented LaS-Attention implementation is causal/unidirectional, potentially contributing to the remaining performance difference. Nonetheless, LaS-Attention is highlighted as the first reported layer not based on 1D long convolutions to achieve an average LRA score above 70%.
• LaS-Chunk Variant: A linear complexity variant, LaS-Chunk, applies the LaS mechanism within fixed-size local attention windows (chunks of size 128). This variant achieves an average accuracy of 65.73%, still outperforming most other Transformer variants (both quadratic and linear complexity ones) but lagging behind the full LaS-Attention, particularly on tasks like Pathfinder that demand integration over very long ranges. This indicates that while the LaS biases are beneficial even in a localized context, the ability to attend over the full sequence (as in the quadratic version) provides additional advantages for certain LRA tasks.

Ablation Studies and Analysis

Further analyses provide insights into the contributions of the individual components and the nature of the inductive biases (Zimerman et al., 2023 ):

• Component Contributions: Ablation studies confirm that both the ELD (locality) operator and the Average Pooling (smoothness) operator are crucial for the observed performance gains. Removing either component ("L-Attention" with only ELD, or "S-Attention" with only AP) results in significantly lower average LRA accuracy compared to the full LaS-Attention.
• ELD vs. Alibi: The ELD operator (exponential decay) is directly compared against the Alibi positional bias (which uses linear decay). On a subset of LRA tasks, ELD consistently outperforms Alibi. Furthermore, combining ELD with the AP smoothness operator (LaS-Attention) yields substantially better results than combining Alibi with the AP operator (average 68.52% vs. 59.92%), suggesting that an exponential decay bias is more effective than a linear one for these long-range tasks.
• Context Length and Data Scaling: Performance degrades markedly when the attention context window (chunk size) is reduced, confirming that LaS-Attention effectively leverages longer contexts when available. Additionally, performance consistently improves as the amount of training data increases, supporting the hypothesis that Transformers' long-range challenge is partly a generalization issue that can be alleviated by more data, especially when combined with a more suitable inductive bias like that provided by LaS.

Conclusion

The research presented in "On the Long Range Abilities of Transformers" (Zimerman et al., 2023 ) demonstrates that the performance of Transformer models on long-range sequence tasks can be significantly enhanced by incorporating inductive biases for locality (via an exponential decay operator) and smoothness (via average pooling of attention weights). These parameter-free modifications, embodied in the LaS-Attention mechanism, substantially improve performance on the challenging LRA benchmark compared to standard Transformers, narrowing the gap with specialized long-range architectures. The findings suggest that the perceived weakness of Transformers in this domain may stem less from fundamental architectural limitations and more from the absence of appropriate biases needed for effective generalization on long sequences.