SoftmaxLift reduces key retrieval for long-context attention

As contexts grow, exact attention compares each new query with every preceding token, so computation grows with the context.

Sparse methods can retrieve keys with large query scores, but the key question is whether they can estimate the full attention result accurately with a bounded search and retrieval budget.

It develops priority sampling for maximum inner product search (MIPS), which finds high-scoring keys, then adds a dimension so one index can estimate attention with O(1/ε²) retrieved keys.

Priority sampling gives keys random priorities based on their attention weights, then uses a threshold to reweight the selected keys. This lets a small retrieved set estimate the weighted values and total weight that exact attention would compute over the whole context.

On 60 matched 128K examples, SoftmaxLift more closely matched the full attention result and achieved a higher score on RULER’s synthetic long-context tasks than TopK and TopK+Sampling. The methods used the same exact search and a 128-key retrieval budget. This supports the comparison under those evaluated conditions.

The theoretical guarantees assume an exact search oracle and do not account fully for the cost of building or implementing search. The 1M-context comparison covers only eight examples and changes both the attention procedure and hardware path, so it is preliminary.