📄 Adaptive Attention Span in Transformers

Summary

Self-attention in Transformers has memory and compute costs that grow with how far back each head can look. As a result, standard models rarely use more than about a thousand tokens of context, which hurts character level language modeling, where relevant information may lie thousands of steps back. The authors observe that heads differ: some attend only to nearby characters, while others spread attention across the whole context. They therefore let each head learn its own context length. A piecewise-linear soft mask with a learnable span parameter z (ramp width R) is applied to the attention weights. An L1 penalty on z keeps spans short unless longer context helps. A dynamic variant predicts the span from the current input. The model is a sequential Transformer with relative position embeddings and Transformer-XL-style caching, evaluated on text8 and enwik8. With a span limit of 8192 characters, it matches or beats earlier Transformers in bits per character while using far fewer FLOPS. This shows that long contexts become affordable when they are allocated only where needed.

Key points

• Datasets: text8 and enwik8 (Mahoney, 2011), each with 100M tokens. Results are reported in bpc on the dev and test sets.
• text8 results:
  • The small Adaptive-Span model (12 layers, 38M params) reaches 1.11 test bpc, versus 1.18 for T12, with an average span of only 314 despite S = 8192.
  • The large model (24 layers, 209M params) reaches 1.07, versus 1.08 for T-XL. It uses 179M FLOPS compared with 438M for T-XL.
• enwik8 results: The large adaptive model reaches 0.98 test bpc (T-XL 24 layers: 0.99) with 209M params and 181M FLOPS. The small model reaches 1.02.
• Efficiency:
  • Inference FLOPS drop by up to 70% at large spans.
  • The largest fixed-span model that fit in training memory used S = 2048 (about 550 ms per batch). An adaptive model with S = 8192 fit in memory at a similar time per batch.
• Learned spans:
  • The lowest 5 layers shrink to the minimum span (R = 32), while a few upper-layer heads reach several thousand characters.
  • Dynamic spans match adaptive spans on text8 (1.08 dev bpc at S = 1024). Spans grow at word beginnings.
  • Code is released at facebookresearch/adaptive-span.