Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TRA: Better Length Generalisation with Threshold Relative Attention

Mattia Opper, Roland Fernandez, Paul Smolensky, Jianfeng Gao

TL;DR

Threshold Relative Attention (TRA) addresses length generalisation failures in decoder‑only transformers by masking irrelevant keys through a threshold on the semantic scores and deriving a contextualised distance that only considers the remaining keys, then modulating this distance with a learnable forget gate to produce final attention weights via a masked softmax. The approach yields substantial gains on synthetic length generalisation benchmarks and demonstrates robust long‑context perplexity improvements on WikiText‑103, often outperforming a wide range of baselines including RoPE, FoT, and CoPE. TRA’s core contribution is the synergy between selective sparsity and contextualised relative distance, enabling a more robust integration of content and position without retraining for longer sequences. The findings suggest a practical path toward more robust, long‑context transformers with potential avenues for efficiency optimisations and memory‑oriented extensions.

Abstract

Transformers struggle with length generalisation, displaying poor performance even on basic tasks. We test whether these limitations can be explained through two key failures of the self-attention mechanism. The first is the inability to fully remove irrelevant information. The second is tied to position, even if the dot product between a key and query is highly negative (i.e. an irrelevant key) learned positional biases may unintentionally up-weight such information - dangerous when distances become out of distribution. Put together, these two failure cases lead to compounding generalisation difficulties. We test whether they can be mitigated through the combination of a) selective sparsity - completely removing irrelevant keys from the attention softmax and b) contextualised relative distance - distance is only considered as between the query and the keys that matter. We show how refactoring the attention mechanism with these two mitigations in place can substantially improve the generalisation capabilities of decoder only transformers.

TRA: Better Length Generalisation with Threshold Relative Attention

TL;DR

Threshold Relative Attention (TRA) addresses length generalisation failures in decoder‑only transformers by masking irrelevant keys through a threshold on the semantic scores and deriving a contextualised distance that only considers the remaining keys, then modulating this distance with a learnable forget gate to produce final attention weights via a masked softmax. The approach yields substantial gains on synthetic length generalisation benchmarks and demonstrates robust long‑context perplexity improvements on WikiText‑103, often outperforming a wide range of baselines including RoPE, FoT, and CoPE. TRA’s core contribution is the synergy between selective sparsity and contextualised relative distance, enabling a more robust integration of content and position without retraining for longer sequences. The findings suggest a practical path toward more robust, long‑context transformers with potential avenues for efficiency optimisations and memory‑oriented extensions.

Abstract

Transformers struggle with length generalisation, displaying poor performance even on basic tasks. We test whether these limitations can be explained through two key failures of the self-attention mechanism. The first is the inability to fully remove irrelevant information. The second is tied to position, even if the dot product between a key and query is highly negative (i.e. an irrelevant key) learned positional biases may unintentionally up-weight such information - dangerous when distances become out of distribution. Put together, these two failure cases lead to compounding generalisation difficulties. We test whether they can be mitigated through the combination of a) selective sparsity - completely removing irrelevant keys from the attention softmax and b) contextualised relative distance - distance is only considered as between the query and the keys that matter. We show how refactoring the attention mechanism with these two mitigations in place can substantially improve the generalisation capabilities of decoder only transformers.

Paper Structure

This paper contains 19 sections, 8 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Model
  3. Background: Relative Positional Encodings
  4. Contextual Relative Distance
  5. Parameterising $\bar{\boldsymbol{D}}$
  6. The Final Attention Weights
  7. Experiments
  8. Tasks
  9. Baselines and Experimental Setup
  10. Results
  11. Flip-Flops++
  12. Why does contextualised relative distance matter?
  13. Language Modeling:
  14. Related Work:
  15. Conclusion
  16. ...and 4 more sections

Figures (3)

  • Figure 1: Contrasting Attention Heatmaps between TRA and FoT. TRA synergises semantic and positional information while FoT must trade one for the other.
  • Figure 2: WikiText-103 Test Perplexity ($\downarrow$) on OOD sequences lengths. Models trained for 100k steps with a window size of 128, and evaluated with window size up to 4098. Results taken over four random seeds.
  • Figure 3: Attention Heatmaps from trained TRA language model. White signifies null attention.