From Sparsity to Simplicity: Enabling Simpler Sequential Replacements via Sparse Attention Distillation

Resumo Técnico: De Esparsidade a Simplicidade (S2S)

Declaração do Problema
Os mecanismos de autoatenção são a base do pré-treinamento de transformadores em grande escala, permitindo a mistura flexível de tokens e a modelagem de dependências de longo alcance. No entanto, sua complexidade quadrática em relação ao comprimento da sequência resulta em alta latência de inferência e sobrecarga de memória. Embora substituir a atenção por módulos sequenciais mais simples (por exemplo, RNNs, Modelos de Espaço de Estado) seja uma via atraente para eficiência, a substituição ingênua frequentemente leva a degradação significativa da precisão, particularmente em escalas maiores. Abordagens de eficiência existentes frequentemente modificam o próprio cálculo da atenção ou podam tokens, mas mantêm a atenção como o operador central. Este artigo aborda o desafio de substituir completamente camadas de atenção por módulos sequenciais sem perda substancial de desempenho.

Metodologia
O artigo propõe um framework chamado De Esparsidade a Simplicidade (S2S), que postula que transformadores pré-treinados decompõem dependências complexas de tokens em mapeamentos bloco a bloco de complexidade efetiva variável. Os autores hipotetizam que camadas que exibem padrões de atenção esparsos são inerentemente mais fáceis de aproximar com módulos sequenciais mais simples.

A metodologia consiste em três componentes principais:

1. Substituição Unificada de Camada por Camada:

   • Os autores introduzem um framework plug-and-play onde camadas de atenção específicas em um Vision Transformer (ViT) pré-treinado são trocadas por misturadores de tokens sequenciais (especificamente BiMamba e BiLSTM) enquanto se preserva a interface da camada original (LayerNorm, conexões residuais e MLPs).
   • Os módulos sequenciais adotam uma estrutura bidirecional multi-cabeça para corresponder à natureza não causal e de contexto completo da atenção do ViT.
   • Supervisão alinhada por camada é usada durante a destilação: a saída do misturador sequencial substituto é comparada diretamente à saída da camada de atenção original na mesma profundidade usando uma perda de erro quadrático médio (MSE).

2. Destilação Guiada por Esparsidade:

   • Para controlar e estudar o impacto da esparsidade, os autores integram A-ViT (ViT Adaptativo), que usa pontuações de parada para gerar máscaras binárias de retenção de tokens.
   • Etapa 1 (Destilação): O modelo estudante (com substituições sequenciais) é treinado sob as mesmas máscaras de tokens por camada geradas pelo professor (A-ViT). Isso garante que o estudante e o professor operem em tokens ativos idênticos, isolando a dificuldade de aproximar a função de mistura de tokens de diferenças no roteamento de tokens. A perda inclui perda de similaridade, um termo de alinhamento de parada e perda de classificação.
   • Etapa 2 (Ajuste Fino): O estudante é ajustado finamente usando seus próprios módulos de parada aprendidos para gerar máscaras, permitindo que ele opere de forma autônoma enquanto mantém o desempenho da tarefa.

3. Estratégia de Substituição Esparsa Consciente de Camada:

   • Em vez de substituição uniforme, a estratégia visa grupos específicos de camadas com base em seus perfis de esparsidade natural. Os autores substituem apenas as camadas mais profundas (que exibem a maior esparsidade natural) e aplicam o framework de destilação guiada por esparsidade a esses blocos específicos.

Contribuições Principais

• A Hipótese S2S: O artigo estabelece que a esparsidade da atenção é uma assinatura de complexidade de mapeamento. Blocos esparsos em transformadores pré-treinados podem ser substituídos por módulos sequenciais mais simples com perda mínima, enquanto blocos densos são mais difíceis de aproximar.
• Framework de Substituição Controlada: Um método unificado para trocar atenção por módulos sequenciais (Mamba/LSTM) enquanto mantém compatibilidade de interface e permite análise seletiva por camada.
• Controle Explícito de Esparsidade: A introdução de um pipeline de destilação onde a esparsidade induzida pelo professor (via máscaras A-ViT) é explicitamente alinhada com o estudante, demonstrando que o aumento da esparsidade do professor reduz consistentemente a lacuna de desempenho entre estudante e professor.
• Compromisso Otimizado: Uma estratégia que combina esparsidade natural por camada com esparsificação explícita para alcançar melhores compromissos precisão-eficiência do que bases de atenção esparsa.

Resultados Experimentais
Experimentos foram conduzidos no ImageNet usando backbones DeiT (Tiny, Small, Base) com BiMamba e BiLSTM como substitutos.

• Esparsidade Natural e Substituibilidade: Visualizações de mapas de interação de tokens mostram que, enquanto camadas rasas exibem mistura global densa, camadas mais profundas tornam-se marcadamente mais esparsas e seletivas. Experimentos controlados substituindo diferentes grupos de camadas revelam que substituir camadas mais profundas (mais esparsas) incorre em quedas de precisão significativamente menores do que substituir camadas mais rasas (mais densas).
• Efeito da Esparsidade Explícita: Quando a taxa de retenção de tokens é reduzida via A-ViT, a lacuna de precisão entre o professor baseado em atenção e o estudante sequencial encolhe consistentemente. Isso confirma que a esparsificação explícita simplifica o mapeamento alvo, tornando mais fácil para os módulos sequenciais aprenderem.
• Compromisso Precisão-Eficiência: A substituição esparsa consciente de camada proposta (substituindo apenas as 3 camadas mais profundas com BiMamba sob orçamentos de esparsidade correspondentes) alcança:
  • Precisão comparável ou ligeiramente superior às bases de atenção esparsa (A-ViT) em taxas de retenção de tokens semelhantes.
  • Ganhos significativos de eficiência: Até 1,71× de aceleração estimada na resolução 384² em comparação com a base densa, com maior taxa de transferência de mistura de tokens do que modelos de atenção esparsa em contagens de tokens maiores.

Significado e Afirmações
O artigo afirma que a substituição de atenção deve ser tratada como um problema estruturado e dependente de camada em vez de uma troca de arquitetura uniforme. O significado do trabalho reside em demonstrar que:

1. A esparsidade é um sinal de substituibilidade: Modelos pré-treinados decompõem naturalmente a complexidade de modo que camadas esparsas são as mais passíveis de simplificação.
2. A esparsificação explícita auxilia a destilação: Forçar esparsidade no professor torna a tarefa de aprendizado para o estudante sequencial mais fácil, reduzindo a lacuna de desempenho.
3. A substituição estratégica supera abordagens uniformes: Ao focar a substituição em camadas naturalmente esparsas e guiar o processo com esparsidade controlada, pode-se alcançar um equilíbrio superior de precisão e latência em comparação com simplesmente podar a atenção ou substituir todo o modelo.

Os autores concluem que essa perspectiva encoraja pesquisas futuras em modelos híbridos que aproveitam o conhecimento de transformadores pré-treinados enquanto integram módulos de sequência mais simples e eficientes onde a estrutura interna do modelo permite.