Resumo Técnico: Atenção Esparsa Dinâmica e Mistura de Especialistas para Classificação Visual Iterativa

Problema

Classificadores visuais padrão normalmente operam em uma única passagem, transformando uma imagem em um rótulo sem a capacidade de revisar uma previsão incerta uma vez formada. Embora o refinamento iterativo tenha se mostrado eficaz em redes recorrentes e predição estruturada, a aplicação disso a imagens naturais apresenta dois desafios específicos:

1. Representação: O modelo requer uma representação visual compacta o suficiente para ser revisitada repetidamente sem excesso de sobrecarga de memória.
2. Eficiência Computacional: A atenção repetida sobre um histórico em expansão de estados latentes aumenta o custo computacional quadraticamente, a menos que o modelo possa atender seletivamente a estados e parâmetros relevantes em cada etapa.

Modelos recursivos existentes frequentemente dependem de decodificação autorregressiva ou sequências de comprimento fixo, que podem não escalar eficientemente ou oferecer traços de raciocínio transparentes para tarefas visuais.

Metodologia

O artigo propõe um classificador visual recursivo que refina uma hipótese de classe ao longo de múltiplas etapas usando uma arquitetura transformer compacta. O design central integra três componentes principais:

1. A Formulação Recursiva x/y/z

O modelo mantém três fluxos de informação distintos:

• $x_{feat}$: Um embedding de imagem global gerado por um codificador convolucional leve (três estágios de convolução-pooling seguidos por pooling de média global).
• $y_t$: Um token de predição de classe atualizado a cada etapa.
• $Z_t$: Um buffer de histórico latente contendo a sequência de estados latentes anteriores $(z_0, \dots, z_t)$.

Em cada etapa de recursão $t$, a entrada do transformer é a concatenação do embedding global, o token de predição anterior e todo o histórico latente acumulado: $h_t = [x_{feat}; y_{t-1}; Z_{t-1}]$. Consequentemente, o comprimento da sequência cresce linearmente com a profundidade da recursão ($L_t = t + 2$).

2. Atenção Esparsa Dinâmica (DSA)

Para gerenciar o crescimento do comprimento da sequência, o autor introduz a DSA, um mecanismo que seleciona um subconjunto de tokens para atenção com base em pontuações de relevância.

• Mecanismo: Uma rede indexadora leve pontua cada token na sequência com base em uma representação de query média por cabeça. Os top-$k$ tokens com as maiores pontuações são selecionados para o cálculo da atenção.
• Esparsidade Dinâmica: Como o comprimento da sequência $L_t$ cresce, o mecanismo de atenção é denso (selecionando todos os tokens) quando $L_t \le k$. Ele torna-se genuinamente esparso apenas quando o histórico excede o orçamento de atenção $k$. Isso permite o refinamento denso nas etapas iniciais e o refinamento seletivo nas etapas posteriores.
• Implementação: A DSA utiliza projeções de chave e valor compartilhadas entre as cabeças de atenção (estilo Multi-Query Attention) para reduzir ainda mais os custos de projeção.

3. Mistura de Especialistas (MoE)

A subcamada feed-forward dos blocos transformer é substituída por uma camada de Mistura de Especialistas para reduzir a computação ativa.

• Configuração: A camada contém $E=4$ especialistas, cada um com uma largura oculta igual à dimensão de embedding $d$. Um mecanismo de roteamento direciona cada token para os top $m=2$ especialistas.
• Balanceamento de Carga: Uma perda auxiliar é aplicada para evitar o colapso do roteamento, garantindo a utilização uniforme dos especialistas.

4. Variantes Avaliadas

O estudo compara três configurações no CIFAR-10, CIFAR-100, MNIST e FashionMNIST:

1. SA-MLP: Atenção densa padrão com uma camada feed-forward densa (baseline).
2. SA-MoE: Atenção densa padrão com feed-forward MoE.
3. DSA-MoE: Atenção Esparsa Dinâmica combinada com MoE.

Principais Contribuições

1. Modelo Recursivo Dimensionalmente Explícito: Uma arquitetura unificada conectando uma representação de imagem global, um token de predição e um histórico latente crescente, herdando a recorrência $x/y/z$ de trabalhos anteriores (TRM), mas adaptando-a para embeddings visuais.
2. Formulação de Atenção Esparsa Dinâmica: Um novo mecanismo de atenção onde a esparsidade é uma função da profundidade da recursão, com uma descrição analítica do ponto exato de cruzamento onde o mecanismo transita de denso para esparso.
3. Análise de Eficiência Unificada: Uma comparação abrangente entre atenção densa, roteamento de especialistas e sua combinação, distinguindo entre parâmetros armazenados, parâmetros ativos, contagens teóricas de multiplicações-adições e tempo de treinamento medido.
4. Interpretabilidade de Trajetórias Internas: Uma análise das interações de papéis de tokens e dinâmicas de roteamento, esclarecendo que o modelo refina as predições com base nos papéis dos tokens (resumo da imagem, hipótese, histórico) em vez de regiões espaciais da imagem devido ao uso de pooling global.

Resultados Experimentais

Os experimentos foram conduzidos usando uma única semente aleatória (42) em uma GPU RTX 3090.

• Acurácia: O ponto de operação selecionado do DSA-MoE (16 etapas de recursão, embedding hard, $k=6$) alcançou a maior acurácia de semente única em todos os datasets:

  • CIFAR-10: 92,3%
  • CIFAR-100: 69,2%
  • MNIST: 98,7%
  • FashionMNIST: 92,3%
  • Nota: Essas margens sobre o baseline denso (SA-MLP) são pequenas (0,2–0,7 pontos percentuais) e os pontos de operação diferem em profundidade e estratégia de embedding.

• Eficiência:

  • Parâmetros Ativos: O DSA-MoE reduziu os parâmetros idealmente ativos de 13,13M (SA-MLP) para 7,10M.
  • Computação: A razão teórica de multiplicação-adição para o bloco transformer foi reduzida para 0,52× em relação ao baseline denso.
  • Tempo de Treinamento: O tempo de treinamento medido diminuiu de 180 segundos/época (SA-MLP) para 165 segundos/época (DSA-MoE), uma redução de aproximadamente 8,3%.

• Descobertas de Ablação:

  • Orçamento de Esparsidade ($k$): Um orçamento moderado de $k=6$ rendeu o melhor equilíbrio, proporcionando 36% de esparsidade média. Reduzir $k$ para 2 aumentou a esparsidade para 75%, mas causou uma queda significativa na acurácia (3,2 pontos percentuais no CIFAR-10).
  • Profundidade de Recursão: O desempenho geralmente melhorou de 4 para 8–12 etapas, com o DSA-MoE continuando a melhorar até 16 etapas.

Significância e Alegações

O artigo faz uma alegação modesta, porém focada: A computação condicional recursiva pode preservar uma acurácia de classificação competitiva enquanto reduz a computação ativa.

O autor enfatiza que:

• Os ganhos em acurácia são pequenos e requerem validação com múltiplas sementes para serem estatisticamente confirmados.
• A principal vantagem é a redução na computação ativa (parâmetros e FLOPs) em vez de um salto massivo na acurácia bruta.
• O design de "sequência crescente" é crítico; sem o histórico em expansão, o mecanismo top-$k$ permaneceria denso e não ofereceria ganhos de eficiência.
• Os traços internos revelam interações de papéis de tokens (ex: dependência do resumo da imagem vs. histórico), mas não fornecem atribuição espacial devido ao front-end de pooling global.

O estudo conclui que, embora a configuração específica do DSA-MoE tenha mostrado resultados favoráveis, a magnitude dos ganhos depende fortemente da profundidade da recursão, do orçamento de esparsidade e da configuração de roteamento. Trabalhos futuros devem focar em validação estatística com múltiplas sementes, tokenização espacial para permitir atenção de nível de região e mecanismos de parada adaptativos.