Nearly Optimal Attention Coresets

TL;DR

This paper proves the existence of nearly optimal size coresets for approximating the Attention mechanism in neural networks, improving previous bounds and establishing lower bounds.

Contribution

It introduces nearly optimal size coresets for Attention estimation, advancing the theoretical understanding of space-efficient neural network approximations.

Findings

Existence of coresets of size $O({\sqrt{d} e^{ ho+o( ho)}/\varepsilon})$ for Attention

Coresets approximate Attention within $\varepsilon$ for all bounded queries

Lower bounds show coresets must have size at least $\Omega({\sqrt{d} e^{ ho}/\varepsilon})$

Abstract

We consider the problem of estimating the Attention mechanism in small space, and prove the existence of coresets for it of nearly optimal size. Specifically, we show that for any set of unit-norm keys and values $(K,V)$ in $\mathbb{R}^d$, there exists a subset $(K',V')$ of size at most $O({\sqrt{d} e^{\rho+o(\rho)}/\varepsilon})$ such that \[ \left\| \operatorname{Attn}(q,K,V)- \operatorname{Attn}(q,K',V') \right\| \le \varepsilon \] simultaneously for all queries whose norm is bounded by $\rho$. This outperforms the best known results for this problem. We also offer an improved lower bound showing that $\varepsilon$-coresets must have size $\Omega({\sqrt{d} e^{\rho}/\epsilon})$.