Positional LSH: Binary Block Matrix Approximation for Attention with Linear Biases

TL;DR

This paper introduces a formal framework connecting positional bias, masks, and embeddings in transformers, proposing a method for efficient ALiBi-biased attention with theoretical guarantees and empirical validation.

Contribution

It presents a novel positional LSH scheme that approximates ALiBi bias matrices, enabling near-linear time attention with provable accuracy guarantees.

Findings

Spectral and max-norm approximation guarantees for the binary mask scheme.

High-probability uniform approximation of ALiBi-biased attention outputs.

Empirical validation on large language models supports theoretical results.

Abstract

Positional encoding in transformers is commonly implemented through positional embeddings, attention masks, or bias terms, but formal connections between these mechanisms remain limited. We study attention with positional bias through the lens of locality-sensitive hashing (LSH), focusing on Attention with Linear Biases (ALiBi). We show that the ALiBi bias matrix is the expectation of contiguous block-diagonal binary masks induced by a ``positional LSH'' scheme. The empirical mean of masks sampled from this scheme yields spectral norm and max-norm approximation guarantees with bounded block sizes with high probability. This structural theorem implies a uniform approximation theorem for ALiBi-biased attention: with high probability over the sampled masks, the approximate attention output is accurate simultaneously for all query-key-value inputs and can be computed in near-linear time in the context length, reducing long-context ALiBi to a collection of randomized short-context regular (positionally unbiased) attention operations. Conceptually, this connects positional bias, masks, and positional embeddings in a single formal framework and suggests an approach to efficient ALiBi-biased attention. Experiments on large language models validate our theoretical findings.