Query-Oblivious Coresets for Softmax Attention: Improved Bounds and Efficient Constructions
Ofek I. Cohen
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^{ρ+\frac12\logρ+o(\log\logρ)}/\varepsilon)$ exist, $ρ$ being the query radius times the centred key radius, against a lower bound $Ω(\sqrt d\,e^ρ/\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^ρ(\sqrt{d_v}+\sqrt{d_k\log(1+ρ)})/\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+ρ)}$. A dimension-free sampling cap $O(e^{2ρ}/\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^ρ/\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^{-ρ}$ 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^ρ)$, dimension-free.