FlashAttention tornou a atenção exata mais rápida movendo menos dados

Transformers de sequências longas eram lentos e consumiam muita memória porque a autoatenção padrão materializava grandes matrizes intermediárias e as movia repetidamente entre a memória de alta largura de banda da GPU e as unidades de computação.

A correção tentadora era a atenção aproximada: alterar a matemática do modelo para reduzir a complexidade. O FlashAttention seguiu o caminho oposto. Manteve a atenção exata, mas processou blocos na SRAM e recalculou as estatísticas correntes do softmax para que a matriz de atenção completa nunca precisasse ser armazenada.

O artigo relatou um ganho de 15% no tempo de execução de ponta a ponta no BERT-large em relação ao recorde de treinamento do MLPerf 1.1, ganho de 3x no GPT-2 com comprimento de sequência de 1K e ganho de 2,4x nas tarefas do Long Range Arena.

Otimizar o tráfego de memória, não a matemática

A restrição central era a E/S entre a HBM e a SRAM, não apenas as operações de ponto flutuante.

A divisão em blocos permitiu que cada bloco fizesse trabalho útil enquanto os dados já estavam no chip.

Manter o algoritmo exato evitou as regressões de qualidade que tornavam muitos métodos aproximados pouco atraentes.

A mesma primitiva abriu caminho para contextos mais longos e novas tarefas de sequência sem alterar a abstração do Transformer.

Quando a movimentação de hardware é o gargalo, um algoritmo melhor pode ser aquele que calcula a mesma resposta atravessando a hierarquia de memória menos vezes.

O FlashAttention tornou-se um bloco de construção padrão para o treinamento e a inferência eficientes de Transformers. Sua lição mais ampla se espalhou além da atenção: algoritmos modernos para aceleradores muitas vezes vencem projetando em torno da hierarquia de memória antes de inventar uma nova matemática de modelo.