Cascade Token Selection for Transformer Attention Acceleration

Resumo Técnico: Seleção em Cascata de Tokens para Aceleração da Atenção em Transformers

Declaração do Problema
Os modelos Transformer enfrentam um gargalo computacional significativo no mecanismo de autoatenção, onde o custo de calcular o produto $QK^\top$ e a subsequente soma ponderada por softmax escala quadraticamente com o comprimento da sequência ($O(T^2 d_h)$). Embora métodos recentes como o FlashAttention tenham otimizado a largura de banda de memória, a complexidade algorítmica permanece uma barreira para sequências longas.

O artigo complementar, Atenção por Dessincronização de Ativação (ADA), propôs uma solução para reduzir essa complexidade de $O(T^2)$ para $O(r^2)$ selecionando um pequeno conjunto de $r \ll T$ tokens "representativos" em cada camada. O ADA identifica esses tokens usando um limiar de Gram: um token é considerado representativo se sua máxima correlação de cosseno com todos os tokens precedentes estiver abaixo de um limiar $1 - \tau^2$. No entanto, o próprio processo de seleção exige o cálculo de uma matriz de Gram completa $T \times T$ em cada camada, incorrendo em um custo de $O(T^2 d)$. Para comprimentos de sequência moderados (por exemplo, $T=2048$), essa sobrecarga de seleção é comparável ao custo da atenção que visa reduzir, anulando o potencial de aceleração. O problema central abordado por este artigo é o alto custo computacional da etapa de seleção de tokens representativos no ADA.

Metodologia: O Mecanismo em Cascata
O artigo introduz um mecanismo de "cascata" que explora a coerência de profundidade do conjunto de tokens representativos. Os autores observam que o conjunto de tokens informativos na camada $l$ é fortemente preditivo do conjunto na camada $l+1$. Isso é fundamentado teoricamente na continuidade de Lipschitz dos blocos de transformer: tokens que são quase-duplicatas (redundantes) ou dessincronizados (representativos) em uma camada tendem a manter essa relação na próxima camada, com apenas perturbações menores à medida que informações contextuais são adicionadas.

Em vez de recalcular a matriz de Gram completa $T \times T$ em cada camada, o mecanismo em cascata:

1. Herda o conjunto representativo $R^{(l)}$ da camada anterior.
2. Valida o conjunto herdado calculando dois produtos de Gram menores:
   • Um produto $r \times r$ entre os representantes herdados para verificar se algum se tornou redundante.
   • Um produto cruzado $(T-r) \times r$ entre não-representantes e os representantes herdados válidos para verificar se algum não-representante se tornou dessincronizado (ou seja, novos representantes).
3. Atualiza o conjunto com um pequeno número de adições e remoções baseadas nessas verificações.

Essa abordagem reduz a complexidade de seleção de $O(T^2 d)$ para $O(Trd)$ por camada, um fator de aceleração de $T/r$.

Principais Contribuições

• Inovação Algorítmica: A proposta de um algoritmo de seleção de tokens em cascata que substitui o cálculo completo da matriz de Gram por etapas incrementais de validação e atualização.
• Justificativa Teórica: Uma prova (Proposição 3) demonstrando que a coerência de profundidade do conjunto representativo é uma consequência da continuidade de Lipschitz dos blocos de transformer e da natureza aditiva das conexões residuais.
• Garantia Conservadora: Uma prova (Proposição 4) mostrando que o método em cascata é conservador; ele produz um superconjunto do conjunto representativo independente. Ele nunca perde um verdadeiro representante, garantindo que a qualidade da atenção não seja degradada, embora possa aumentar ligeiramente o número de representantes ($r$).
• Aceleração Ortogonal: O artigo posiciona este método como o análogo na dimensão de tokens ao Inferência em Subespaço Portão (GSI), que acelera camadas lineares via herança de bases coerentes em profundidade. Juntos, eles cobrem toda a passagem direta do transformer (tanto o eixo $d$ quanto o eixo $T$).

Resultados Experimentais
O método foi validado em três famílias de modelos (GPT-2 124M, GPT-J 6B, OPT 6.7B) usando hardware AMD MI300X com um comprimento de sequência de $T=512$ e um limiar de Gram $\tau=0,30$.

• Coerência de Profundidade: O índice de Jaccard entre conjuntos representativos em camadas consecutivas é alto (0,83 a 0,94) em camadas profundas, confirmando a estabilidade do conjunto de tokens.
• Economia de Operações de Gram: O mecanismo em cascata reduziu as operações totais de Gram em:
  • 22% para GPT-2 (12 camadas)
  • 52% para GPT-J (28 camadas)
  • 63% para OPT (32 camadas)
    As economias aumentam com a profundidade do modelo porque o custo inicial de Gram completo na camada 0 é amortizado sobre mais camadas.
• Taxa de Rotatividade: Em camadas profundas, a taxa de rotatividade (fração do conjunto modificada) cai para 3–18%, indicando que o conjunto herdado requer ajustes mínimos.
• Previsões de Escala: Os autores projetam que para sequências mais longas (por exemplo, $T=8192$), onde $r$ cresce sublinearmente, a aceleração de seleção poderia atingir $18\times$, reduzindo os custos de seleção de um termo dominante para um termo negligenciável.

Significado e Alegações
O artigo afirma que o mecanismo em cascata revela uma propriedade estrutural de transformers treinados: o conjunto de tokens informativos é uma propriedade estrutural coerente que se propaga através da profundidade da rede. Ao explorar essa coerência, o método torna a sobrecarga de seleção prática para o ADA, permitindo que a compressão de atenção alcance sua aceleração teórica sem ser gargalada pela etapa de seleção.

Os autores enfatizam que isso é uma aceleração complementar aos métodos existentes como FlashAttention (que otimiza a largura de banda de memória) e atenção esparsa (que usa padrões fixos). Diferentemente de métodos de poda ou fusão de tokens que alteram permanentemente as representações, o ADA em cascata seleciona dinamicamente representantes com base na estrutura de entrada sem alterações arquitetônicas ou re-treinamento. A combinação de GSI e ADA em cascata oferece uma redução multiplicativa no custo de inferência, reduzindo efetivamente a complexidade da passagem direta para a "complexidade efetiva" da entrada em vez do tamanho nominal do modelo.