Query-Oblivious Coresets for Softmax Attention: Improved Bounds and Efficient Constructions

Abstract

A query-oblivious coreset for a softmax-attention head is a subset $S$ of the key--value pairs such that attention computed from $S$ alone is within $\varepsilon$ of the full output, in $\ell_2$, simultaneously for every query in a ball. Liberty, Andoni and Kleiner proved that unweighted coresets of size $O(\sqrt d\,e^{\rho+\frac12\log\rho+o(\log\log\rho)}/\varepsilon)$ exist, $\rho$ being the query radius times the centred key radius, against a lower bound $\Omega(\sqrt d\,e^{\rho}/\varepsilon)$, and conjectured that closing the gap needs new techniques. We show it does not. A spherical lift of both balls into one exponential-kernel instance lets the chaining bound of Bozzai and Rothvoss apply directly, and Chevet's inequality splits key from value dimension: unweighted coresets of size $O(e^{\rho}(\sqrt{d_v}+\sqrt{d_k\log(1+\rho)})/\varepsilon)$ exist and are computable in randomised polynomial time, the first with a whole-ball guarantee at the existential size up to $\sqrt{\log(1+\rho)}$. A dimension-free sampling cap $O(e^{2\rho}/\varepsilon^{2})$ completes the envelope. In fixed dimension the logarithm disappears: completing the key ball to a sphere makes the kernel an unweighted Gaussian one, so Tai's diameter-free bound gives $O_{d_k,d_v}(e^{\rho}/\varepsilon)$, ruling out a matching logarithmic lower bound there and answering the Gaussian-restriction case of a question of Bozzai and Rothvoss for the exponential and Hellinger kernels. We restate the Liberty--Andoni--Kleiner bound in the centred convention with a full proof, and show that the one-way communication bounds of Chen et al.\ transfer to query-oblivious coresets, where for $\varepsilon\ll e^{-\rho}$ they are the strongest floors known. The dimensional factor is the price of one signing for all queries: for a single query the discrepancy is $O(e^{\rho})$, dimension-free.