Resumo Técnico: FAR (Substituição de Atenção Preservadora de Função para Inferência Compatível com IMC)

1. Declaração do Problema

As arquiteturas Transformer, particularmente aquelas que dependem de autoatenção, enfrentam gargalos significativos de eficiência quando implantadas em hardware emergente de Computação em Memória (IMC), como matrizes cruzadas ReRAM. Embora a autoatenção se destaque em processadores digitais paralelos (GPUs/TPUs) ao permitir a integração global de tokens, suas características computacionais conflitam fundamentalmente com o modelo de execução IMC:

• Multiplicações de Ativação para Ativação: A autoatenção depende fortemente de operações como $QK^\top$ e normalização softmax, que exigem mistura dinâmica e não local de ativações, em vez das multiplicações vetor-matriz (VMM) estacionárias em peso que o hardware IMC otimiza.
• Sobrecarga de Memória: A complexidade quadrática ($O(T^2)$) das interações entre tokens e a necessidade de ler repetidamente ativações espacialmente distribuídas levam a operações analógicas fragmentadas, extensas conversões Analógico-Digital/Digital-Analógico (ADC/DAC) e baixa utilização da matriz cruzada.
• Limitações Atuais: Aceleradores IMC existentes frequentemente descarregam a atenção para unidades digitais, deixando o gargalo quadrático sem resolução, ou tentam mapear a atenção diretamente para arrays analógicos sem abordar os custos inerentes de movimentação de dados.

Concomitantemente, análises recentes sugerem que as camadas de atenção exibem redundância substancial durante a inferência, com mapas de atenção tornando-se mais suaves e comprimíveis através da profundidade. Isso motiva a hipótese de que o papel funcional da atenção pode ser aproximado por módulos sequenciais que se alinham melhor à localidade IMC e à estacionariedade de pesos.

2. Metodologia: O Framework FAR

Os autores propõem o FAR (Substituição de Atenção Preservadora de Função), um framework que substitui todos os blocos de autoatenção em um Vision Transformer pré-treinado (especificamente a família DeiT) por módulos sequenciais, nomeadamente LSTMs Bidirecionais Multicabeça (BiLSTMs). A abordagem é projetada como uma conversão arquitetural "plug-and-play" que não requer treinamento do zero.

2.1 Estratégia de Substituição

• Substituição por Bloco: Cada bloco de atenção $A_l$ em um transformer pré-treinado é substituído por um substituto aprendível $A'_l$ (uma BiLSTM Multicabeça). Todos os outros componentes (blocos MLP, embeddings de patches, normalização de camadas) permanecem congelados.
• Equivalência Funcional via Destilação: A substituição é supervisionada através de destilação camada por camada. Para cada camada $l$, uma perda de similaridade ($L_{sim}$) minimiza a diferença entre a saída do cabeçote de atenção original e o novo cabeçote BiLSTM:
  $$L_{sim}^{(l)} = ||A_l(x_l) - A'_l(x_l)||_2^2$$
  Isso garante que o substituto replique a estrutura representacional do professor.
• Protocolo de Treinamento: O processo envolve duas etapas:
  1. Fase de Destilação: Apenas os módulos substitutos BiLSTM são treinados usando uma combinação de perda de similaridade e perda de tarefa (classificação), enquanto o restante do modelo permanece congelado.
  2. Fase de Ajuste Fino: Todo o modelo é ajustado finamente em conjunto para recuperar qualquer queda residual de precisão.

2.2 Arquitetura BiLSTM Multicabeça

Para preservar o viés indutivo da atenção multicabeça, o substituto é estruturado como uma BiLSTM Multicabeça:

• Projeção e Divisão: O embedding de entrada é projetado e dividido em $N$ subespaços (correspondendo ao número de cabeçotes de atenção do professor).
• Processamento Bidirecional: Cada subespaço é processado por um cabeçote BiLSTM separado, capturando dependências de tokens tanto para frente quanto para trás.
• Concatenação e Projeção: As saídas são concatenadas e projetadas de volta para a dimensão original.
• Compatibilidade com IMC: Diferentemente da atenção, esta estrutura depende de reutilização recorrente de pesos e fluxo de dados sequencial, eliminando interações de ativação entre todos os pares. Isso mapeia naturalmente para matrizes cruzadas ReRAM onde os pesos são estacionários e os dados fluem sequencialmente.

2.2 Poda Estruturada para Restrições IMC

Para adaptar ainda mais aos arrays IMC com recursos limitados, os autores aplicam poda estruturada às dimensões ocultas dos cabeçotes BiLSTM:

• Extensão do DeepHoyer: Eles estendem o framework DeepHoyer para atingir unidades ocultas de LSTM. Como as portas da LSTM (entrada, esquecimento, célula, saída) compartilham unidades ocultas, podar uma única unidade requer remoção coordenada em todas as matrizes de peso associadas ($W_{ih}$, $W_{hh}$ e pesos de projeção).
• Esparsidade em Grupo: Uma medida de esparsidade Hoyer diferenciável é aplicada a grupos de pesos correspondentes a cada unidade oculta, incentivando a rede a aprender quais unidades são redundantes.
• Compressão Adaptativa à Camada: A poda é realizada independentemente por cabeçote e direção, permitindo que o modelo adapte sua capacidade à redundância específica de cada camada, resultando em módulos compactos que se encaixam nos tamanhos de arrays IMC.

3. Resultados Chave

O framework foi avaliado na família DeiT (Tiny, Small, Base) no ImageNet e em várias tarefas downstream (CIFAR-10/100, Cars, Flowers, iNaturalist).

• Preservação de Precisão:
  • DeiT-Tiny: O FAR-Tiny na verdade superou o DeiT-Tiny de base em +1,2% no ImageNet (73,3% vs 72,2%), sugerindo que o substituto sequencial atua como um viés indutivo benéfico para modelos menores.
  • Modelos Maiores: Para DeiT-Small e Base, o FAR manteve precisão competitiva com uma lacuna modesta de aproximadamente 1–2% em comparação com o professor, demonstrando que o papel funcional da atenção é amplamente reproduzível via recorrência sequencial.
  • Transferência Downstream: Os modelos FAR mantiveram forte transferibilidade para tarefas de classificação downstream, indicando que os mapeamentos sequência-para-sequência aprendidos generalizam efetivamente.
• Robustez à Poda:
  • A aplicação de poda estruturada reduziu a dimensão oculta média em 40–60% com apenas uma queda adicional de 1–2% na precisão.
  • Os padrões de poda foram não uniformes, com maior retenção nas camadas intermediárias e poda mais agressiva nas extremidades da rede, refletindo heterogeneidade interna na redundância representacional.
• Especialização por Cabeçote: A visualização dos mapas de interação de tokens mostrou que os cabeçotes FAR herdaram com sucesso os papéis semânticos distintos dos cabeçotes de atenção originais, concentrando-se em regiões específicas de objetos em vez de colapsar em padrões redundantes.

Análise de Eficiência de Hardware

Usando um modelo analítico de IMC ReRAM, os autores estimaram latência e energia:

• Latência: O FAR em IMC ReRAM é estimado em ~18x mais rápido que o DeiT em IMC ReRAM (devido à evitação da sobrecarga quadrática da atenção) e ~400x mais rápido que o DeiT em uma GPU.
• Energia: O FAR em IMC ReRAM é estimado em ~3x mais eficiente energeticamente que o DeiT em IMC ReRAM e ~150x mais eficiente que o DeiT em uma GPU.
• Mecanismo: Esses ganhos derivam da eliminação de multiplicações de ativação para ativação, redução de conversões ADC/DAC e habilitação de execução estacionária em peso. Os autores observam que excluir os custos de softmax (que são gargalos significativos para a atenção em IMC) torna essas estimativas conservadoras para o FAR.

4. Significado e Alegações

O artigo afirma que o FAR fornece uma alternativa prática e eficiente em hardware para implantar transformers pré-treinados em aceleradores de borda baseados em IMC, sem a necessidade de treinar modelos do zero.

• Viabilidade de Substituição Total: O trabalho demonstra que a atenção pode ser substituída de ponta a ponta no momento da inferência, mantendo precisão prática, desafiando a noção de que interações de tokens entre todos os pares são estritamente necessárias para tarefas de visão de alto desempenho.
• Alinhamento com Hardware: Ao substituir a atenção por módulos sequenciais e estacionários em peso, o FAR alinha o fluxo de dados computacional com as forças físicas das matrizes cruzadas ReRAM, reduzindo significativamente o tráfego de memória e a latência.
• Preservação de Conhecimento Pré-treinado: O framework prova que as ricas representações semânticas aprendidas durante o pré-treinamento em grande escala podem ser preservadas e transferidas para arquiteturas estruturalmente diferentes e compatíveis com IMC através da destilação por bloco.

Em resumo, o FAR oferece um caminho para reestruturar a inferência de transformer em torno de módulos sequenciais compatíveis com IMC, fechando a lacuna entre a dominância de modelos baseados em atenção e as restrições do hardware emergente centrado em memória.