🤖 machine learning

AdaSplash-2: Faster Differentiable Sparse Attention

O artigo apresenta o AdaSplash-2, uma implementação eficiente de atenção esparsa diferenciável baseada em $\alpha$ -entmax que utiliza inicialização por histograma e otimizações em GPU para reduzir a sobrecarga computacional, alcançando desempenho de treinamento comparável ou superior ao FlashAttention-2 em contextos longos e melhorando significativamente o desempenho em tarefas de longo alcance.

Autores originais: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

Publicado 2026-04-17

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está organizando uma festa gigante com milhares de convidados (os dados) e precisa decidir quem vai conversar com quem. No mundo das Inteligências Artificiais modernas (os Transformers), essa decisão é feita por um mecanismo chamado "atenção".

O problema é que, na versão atual, todo mundo tenta conversar com todo mundo. Se você tem 100 pessoas, são 10.000 conversas possíveis. Se tem 1 milhão de pessoas, o número de conversas explode e o computador fica sobrecarregado, lento e gasta muita energia. É como tentar organizar uma festa onde cada convidado precisa apertar a mão de todos os outros antes de começar a música.

Os cientistas tentaram resolver isso criando "atenção esparsa" (sparse attention), onde cada pessoa só conversa com um grupo pequeno e relevante. Mas havia um problema: o método mais inteligente para escolher esses grupos (chamado α-entmax) era muito lento para calcular, porque exigia um "cálculo de normalização" complexo, como se fosse tentar adivinhar o número exato de bolinhas de gude em um pote sem abri-lo, fazendo várias tentativas.

Aqui entra o ADASPLASH-2.

A Grande Ideia: O "Mapa de Tráfego" em Tempo Real

O ADASPLASH-2 é como um organizador de festas superinteligente que não tenta adivinhar o número de bolinhas de gude. Em vez disso, ele usa um truque genial:

O Histograma (O Mapa de Tráfego):
Imagine que, em vez de contar cada convidado individualmente, o organizador olha para a sala e cria um "mapa de calor" rápido. Ele divide a sala em 8 ou 16 zonas (caixas) e conta rapidamente quantas pessoas estão em cada zona.
- Na linguagem do papel: Eles criam um histograma compacto direto na memória rápida do chip (SRAM), enquanto os dados passam por ele. Isso é feito em uma única passada, sem precisar voltar atrás.
A Adivinhação Perfeita (Inicialização):
Com esse mapa de zonas, o organizador consegue fazer uma estimativa quase perfeita de onde está o "ponto de corte" para decidir quem conversa com quem.
- A mágica: Antes, o computador precisava fazer 10 ou 20 tentativas (iterações) para achar esse número. Com o ADASPLASH-2, ele precisa de apenas 1 ou 2 tentativas. É como se, em vez de tentar adivinhar o peso de um elefante chutando, você olhasse para a sombra dele e dissesse: "Ah, deve ser por aqui".
O Pulo do Gato (Pular o que não importa):
Uma vez que o organizador sabe quem é relevante, ele cria uma "lista de convidados VIP" (uma máscara de bits). Durante a festa, ele ignora completamente os grupos de pessoas que não vão conversar.
- Na prática: O chip da placa de vídeo (GPU) pula blocos inteiros de dados que são zero. Isso economiza tempo e energia, especialmente quando a festa é muito grande (contextos longos).

Por que isso é importante?

Velocidade: Em festas pequenas, o ADASPLASH-2 é tão rápido quanto os métodos antigos. Mas, em festas gigantes (longos textos, vídeos longos, documentos inteiros), ele é muito mais rápido do que a tecnologia atual (FlashAttention-2), porque aproveita o fato de que, em textos longos, a maioria das palavras não tem relação direta entre si.
Inteligência: Ao contrário de métodos que cortam aleatoriamente, o ADASPLASH-2 decide dinamicamente quem é importante. Isso ajuda a IA a entender melhor contextos longos, como ler um livro inteiro e lembrar de um detalhe da primeira página.
Resultados: Nos testes, os modelos treinados com essa técnica foram tão bons quanto os modelos antigos em tarefas curtas, mas muito melhores em tarefas longas, entendendo nuances que os outros perdem.

Resumo da Ópera

O ADASPLASH-2 é uma nova maneira de fazer a IA "ler" textos longos sem ficar cansada. Em vez de tentar processar tudo de uma vez (o que é lento e caro), ele cria um mapa rápido para saber exatamente onde focar, pulando o que é irrelevante.

É como trocar um sistema de entrega de cartas onde o carteiro precisa entregar uma carta em cada porta da cidade, por um sistema onde ele só entrega nas casas que realmente têm encomendas, usando um mapa inteligente que ele mesmo desenha enquanto anda pela rua. O resultado? A entrega é mais rápida, mais barata e a IA consegue "ler" livros inteiros sem se perder.

1. O Problema

O mecanismo de atenção em Transformers é um gargalo computacional significativo devido à sua complexidade quadrática ( $O(n^2)$ ) em tempo e memória para sequências longas. Embora o FlashAttention tenha resolvido o problema de memória para a atenção softmax (densa) através de otimizações de E/S e fusão de operações, a atenção softmax é inerentemente densa, atribuindo probabilidade não nula a todos os tokens.

Para contextos longos, essa densidade é indesejável, pois dilui a atenção em tokens irrelevantes. Alternativas esparsas adaptativas, como a transformação $\alpha$ -entmax, permitem esparsidade dependente da entrada (gerando zeros exatos), mas historicamente foram lentas. O principal obstáculo é o cálculo do normalizador $\tau$ , que requer métodos iterativos de busca de raízes (como bisseção ou Halley), exigindo múltiplas passagens sobre os dados e impedindo que superem o desempenho do FlashAttention-2.

2. Metodologia: ADASPLASH-2

O ADASPLASH-2 é uma implementação eficiente e consciente de hardware para a atenção $\alpha$ -entmax. A inovação central reside na redução drástica do custo computacional do cálculo do normalizador $\tau$ e na exploração eficiente da esparsidade dinâmica em GPUs.

A. Inicialização Baseada em Histograma (O Núcleo da Otimização)

Em vez de calcular $\tau$ diretamente sobre todos os scores de atenção (o que exigiria múltiplas passagens), o método propõe:

Construção de Histograma On-Chip: Enquanto os blocos de chaves (key blocks) são transmitidos (streamed) através da memória SRAM on-chip (rápida), o algoritmo calcula um histograma compacto dos scores de atenção.
Aproximação Conservadora: O histograma discretiza os scores no intervalo unitário. A solução do problema de normalização sobre este histograma fornece um limite inferior provável para o verdadeiro $\tau$ .
Convergência Rápida: Esta estimativa inicial de alta qualidade permite que o solver de raízes (híbrido) convirja para a solução exata em apenas 1 a 2 iterações, eliminando a necessidade de múltiplas passagens completas sobre os dados.

B. Implementação Consciente de Hardware (GPU/Triton)

Codificação Bit-packed: O algoritmo utiliza uma máscara de blocos binária compacta (codificada em int32) para rastrear quais blocos de atenção contêm pesos não nulos. Isso permite pular blocos zero com sobrecarga mínima.
Solvers Híbridos: Baseado na suavidade da função $\alpha$ -entmax (que depende do valor de $\alpha$ ), o método utiliza um solver híbrido seguro: passos de Halley para $\alpha \le 1.5$ , Newton para $1.5 < \alpha \le 2$ , e secante para $\alpha \ge 2$ , sempre com fallback para bisseção se necessário.
Fusão de Operações: A implementação em Triton funde o cálculo do histograma, a solução inicial e a refinação em um único kernel, minimizando o tráfego de dados entre a memória HBM (lenta) e a SRAM (rápida).

3. Principais Contribuições

Normalização via Histograma On-Chip: Derivação de limites teóricos para o normalizador $\tau$ usando um histograma compacto na SRAM, permitindo avaliação rápida sem materializar intermediários densos.
Refinamento em Uma Passada: Demonstração de que a inicialização baseada em histograma reduz a necessidade de iterações para convergência, tornando o processo quase de "passada única" na prática.
Exploração Eficiente de Esparsidade Dinâmica: Design de um kernel otimizado que ignora blocos zero com sobrecarga negligenciável, superando estratégias de mascaramento ingênuas.
Resultados Empíricos Sólidos: Validação em benchmarks de modelagem de linguagem e tarefas downstream, mostrando ganhos de velocidade e precisão.

4. Resultados Experimentais

Desempenho de Treinamento (Velocidade)

Regimes de Esparsidade Moderada a Alta: O ADASPLASH-2 iguala ou supera o FlashAttention-2 (implementado em CUDA e Triton) quando a esparsidade de blocos é superior a ~60%.
Ganhos em Contextos Longos: À medida que o comprimento do contexto aumenta, a esparsidade natural emerge. O ADASPLASH-2 transforma essa esparsidade em ganhos de eficiência, superando o FlashAttention-2 em até 2x em cenários de alta esparsidade.
Passada de Retropropagação (Backward): O ganho é particularmente substancial na etapa de backward, que domina o tempo de treinamento, pois o algoritmo ignora eficientemente os blocos com gradiente zero.

Desempenho em Modelagem de Linguagem (Precisão)

Contextos Longos: Modelos treinados com $\alpha$ -entmax (usando ADASPLASH-2) superam significativamente as bases softmax em benchmarks de contexto longo (RULER e HELMET), especialmente em tarefas de rastreamento de variáveis e extração de palavras.
Contextos Curtos: Os modelos mantêm desempenho competitivo ou superior em tarefas de contexto curto, demonstrando que a esparsidade não prejudica a capacidade geral do modelo.
Codificação Posicional: A combinação de $\alpha$ -entmax com o esquema de codificação NAPE (No Positional Encoding + ALiBi) mostrou-se superior ao uso de RoPE, alinhando-se com descobertas anteriores de que o RoPE pode ser subótimo para atenção esparsa.

5. Significado e Impacto

O ADASPLASH-2 remove a principal barreira de eficiência que impedia a adoção generalizada da atenção $\alpha$ -entmax em larga escala. Ao tornar o cálculo do normalizador tão eficiente quanto o do softmax em regimes de esparsidade relevantes, o trabalho permite:

Treinamento Eficiente de Longo Contexto: Viabiliza o treinamento de modelos com janelas de contexto muito grandes (ex: 32K, 128K tokens) sem o custo quadrático proibitivo da atenção densa.
Melhoria na Generalização: Aproveita a capacidade da $\alpha$ -entmax de focar em tokens relevantes e ignorar ruído, melhorando a generalização em tarefas complexas de raciocínio e recuperação de informação.
Futuro da Arquitetura: Estabelece um novo padrão para kernels de atenção esparsa diferenciáveis, sugerindo que a combinação de inicializações inteligentes (histogramas) e exploração de hardware (bit-packing) é o caminho para superar os limites atuais dos Transformers.

Em resumo, o ADASPLASH-2 não é apenas uma otimização de kernel, mas um passo fundamental para tornar a atenção esparsa adaptativa uma alternativa viável e superior ao softmax para a próxima geração de modelos de linguagem de grande escala.