Nearly Optimal Attention Coresets

Com base no artigo fornecido, segue um resumo técnico detalhado de "Nearly Optimal Attention Coresets".

Definição do Problema

O artigo aborda o desafio de comprimir o cache de Chaves-Valores (KV) em modelos generativos baseados em transformadores. À medida que as janelas de contexto se expandem, o tamanho do cache KV torna-se um gargalo primário para a implantação de grandes modelos de linguagem. Os autores formalizam o problema da Seleção de Coreset de Atenção: dado um conjunto de $n$ pares chave-valor $(K, V)$ em $\mathbb{R}^d$ (onde chaves e valores têm norma unitária), o objetivo é encontrar um pequeno subconjunto $(K', V')$ de tamanho $S$ tal que a saída de atenção seja aproximada dentro de um erro $\epsilon$ para todos os vetores de consulta $q$ com norma limitada $\|q\| \le \rho$.

O mecanismo de atenção é definido como:
$$\text{Attn}(q, K, V) = \frac{A(q, K, V)}{B(q, K)} = \frac{\sum_i e^{q^T k_i} v_i}{\sum_i e^{q^T k_i}}$$

Os autores observam que, sem limitar a norma da consulta $\rho$, a compressão com perdas é impossível. Uma consulta com norma suficientemente grande pode "escolher seletivamente" um único par chave-valor, forçando o esboço a armazenar todo o conjunto de dados para manter a precisão. Portanto, o problema é restrito a consultas onde $\|q\| \le \rho$.

Metodologia

A abordagem técnica central baseia-se na teoria da discrepância, explorando especificamente o teorema de balanceamento vetorial de Banaszczyk. A metodologia envolve três componentes principais:

1. Divisão por Metade via Discrepância: Os autores empregam um "truque de divisão por metade" onde atribuem sinais $\sigma_i \in \{-1, 1\}$ aos termos na soma de atenção. Ao selecionar sinais para minimizar a discrepância (o termo de erro $E(q) = \sum \sigma_i e^{q^T k_i} v_i$), eles podem construir um subconjunto $(K', V')$ contendo os termos onde $\sigma_i = 1$. Este subconjunto tem no máximo metade do tamanho do conjunto original, e o erro é limitado pela discrepância da soma sinalizada.
2. Balanceamento de Norma Tensorial: Para limitar o erro para o numerador de valor vetorial $A(q, K, V)$, os autores expandem o termo exponencial $e^{q^T k_i}$ usando uma série de Taylor. Isso transforma o problema no balanceamento de uma soma de produtos tensoriais ($k_i^{\otimes m} \otimes v_i$). Eles definem um fator de escala específico $z(m)$ e constroem um vetor de alta dimensão $x_i$ contendo esses componentes tensoriais escalados.
3. Construção de Corpo Convexo: Eles aplicam o teorema de Banaszczyk para mostrar que existem sinais $\sigma_i$ tais que a soma $\sum \sigma_i x_i$ reside dentro de um corpo convexo específico $K$. Este corpo é definido por restrições nas normas duais das somas tensoriais. Ao escolher cuidadosamente a geometria deste corpo convexo, eles garantem que o erro permaneça pequeno em todas as ordens tensoriais $m$.
4. Balanceamento Simultâneo: Para lidar com o mecanismo de atenção completo (tanto o numerador $A$ quanto o denominador $B$), bem como o centralização das chaves, os autores provam um lema de balanceamento simultâneo. Isso permite que eles satisfaçam os limites de discrepância para múltiplos objetivos (numerador, denominador e centralização de chaves) simultaneamente usando um único conjunto de sinais, com apenas um aumento de fator constante no limite de erro.
5. Compressão Iterativa: O processo de divisão por metade é aplicado iterativamente ($\ell$ vezes) até que o erro acumulado esteja dentro da tolerância desejada $\epsilon$. Isso reduz o tamanho do conjunto de dados de $n$ para $O(\sqrt{d} e^\zeta / \epsilon)$.

Principais Contribuições e Resultados

1. Limite Superior (Teorema 1 e 4):
O artigo estabelece a existência de um coreset de atenção de tamanho:
$$S = O\left( \frac{\sqrt{d}}{\epsilon} e^{\zeta} \right)$$
onde $\zeta = \rho + \frac{1}{2}\log(\rho) + o(\log \log \rho)$.
Este resultado melhora o limite superior anterior mais conhecido (da literatura de aproximação de atenção em fluxo, e.g., [30]) por um fator de aproximadamente $\log(n) \cdot e^{\rho - o(\rho)}$. Os autores alcançam isso modificando e simplificando cirurgicamente argumentos de discrepância existentes para lidar com funções de valor vetorial de forma mais rigorosa.

2. Limite Inferior (Teorema 2 e 7):
Os autores provam um limite inferior correspondente, demonstrando que qualquer $\epsilon$-coreset para vetores de norma unitária e consultas de norma no máximo $\rho$ deve ter um tamanho de:
$$S = \Omega\left( \frac{\sqrt{d}}{\epsilon} e^{\rho} \right)$$
Este limite inferior é derivado via uma redução do problema de Indexação Unidirecional, utilizando códigos esféricos para construir instâncias difíceis onde distinguir entre bits requer um tamanho de esboço grande.

3. Near-Optimalidade:
Comparando os limites superior e inferior, a lacuna é um fator de $O(e^{\frac{1}{2}\log(\rho) + O(\log \log \rho)}) = O(\rho^{1/2 + o(1)})$. Os autores afirmam que isso estabelece um tamanho de coreset "quase ótimo", pois os limites coincidem até um fator subpolinomial em $\rho$.

Significado e Alegações

O artigo afirma resolver a questão do tamanho ótimo de coresets para mecanismos de atenção com garantias de erro.

• Avanço Teórico: Fornece os primeiros limites existenciais rigorosos para coresets de atenção, indo além de métodos de compressão heurísticos (como agrupamento ou destilação) para garantias prováveis baseadas na teoria da discrepância.
• Implicação Prática: Os resultados sugerem que, para consultas de norma limitada, o cache KV pode ser comprimido para um tamanho que depende da dimensão $d$, da tolerância de erro $\epsilon$ e da norma da consulta $\rho$, mas é independente do número total de tokens $n$ (uma vez que $n$ seja suficientemente grande).
• Limitações: Os autores observam que, sem a restrição na norma da consulta $\rho$, nenhuma compressão é possível. Além disso, reconhecem que a lacuna remanescente entre os limites superior e inferior (aproximadamente $\sqrt{\log n}$ no pior caso) provavelmente requer abordagens fundamentalmente diferentes para ser fechada.

O artigo não propõe novas implementações experimentais ou mudanças arquitetônicas específicas para transformadores, mas foca estritamente nos limites teóricos da compressão de dados para o mecanismo de atenção.