Fast and Exact Relative-Position Attention via Structure-Aware Kernels

Fast and Exact Relative-Position Attention via Structure-Aware Kernels. By Delyan Boychev and two others, from INSAIT, Sofia University. Published on the thirty-first of October.

Relative-position attention mechanisms are a compelling alternative to RoPE across various transformer architectures, and in domains such as audio they deliver superior performance. They typically work by adding terms to the attention scores, either content-position or bias-position terms. That is where the trouble starts. RoPE encodes position with a rotation applied to each token, directly on the queries and keys. Relative-position terms instead generally need computation or lookups that depend on the offset between every pair of tokens. That makes them hard to integrate into fused attention kernels, and as sequences grow it brings severe overheads in memory, in bandwidth, and in compute.

The authors propose an exact attention algorithm that brings additive relative-position terms inside FlashAttention-style tiled computation. Nothing is approximated: the attention computed is the same, and the gain comes from exploiting the structure of the offsets.

Three ingredients do the work. First, offset-aware tiling, which avoids redundant indexing and redundant memory accesses. Second, the method computes only the key and query relation products that are actually needed, and reuses them across token pairs, in both the forward and the backward pass. Third, the gradients are sharded and aggregated by offset, which eliminates atomic contention. All of this comes as a single, unified implementation in Triton that covers diverse relative-position attention mechanisms, evaluated against both naive and specialized implementations.

The headline figures are for sequences of eight thousand tokens, measured against the strongest baseline for each mechanism. The forward pass speeds up by between one point eight and six point four times. The backward pass speeds up by between two and twenty-one times.

Peak memory falls as well. In the forward pass the reduction reaches between two point nine and one hundred and twenty-eight times. In the backward pass it reaches between one point four and fifty-four times. Carried through to a full system, this enables long-context end-to-end transcription that runs one point eight times faster and uses two point four times less memory.

Relative-position attention had a quality argument in its favour and an engineering argument against it. This work goes after the engineering side directly, without changing the mathematics, so that these mechanisms become a drop-in, hardware-efficient option for scaling transformers to long contexts.