Steering Sparse Autoencoder Latents to Control Dynamic Head Pruning in Vision Transformers (Student Abstract)

Título: Direcionamento de Latentes de Autoencoders Esparsos para Controlar a Poda Dinâmica de Cabeças em Vision Transformers

1. Problema Identificado

Os Vision Transformers (ViTs) utilizam mecanismos de atenção multi-cabeça para capturar interações entre tokens. No entanto, muitas dessas cabeças são redundantes, aumentando o custo computacional sem ganhos proporcionais de desempenho.

• Solução Existente: Frameworks de poda adaptativa (como o AdaViT) utilizam redes auxiliares para selecionar quais cabeças remover com base em embeddings residuais.
• Limitação Crítica: Embora eficientes, as políticas de poda existentes são frequentemente opacas e difíceis de controlar. Como as decisões dependem de embeddings densos e polissêmicos, o processo de seleção de cabeças carece de interpretabilidade mecânica. Não é possível explicar ou manipular facilmente por que ou como certas cabeças são mantidas ou removidas em nível latente.

2. Metodologia Proposta

O trabalho propõe um novo framework que integra Autoencoders Esparsos (SAEs) à poda dinâmica de cabeças em ViTs, visando tornar as decisões de poda interpretáveis e controláveis no nível latente.

• Base de Poda Dinâmica: O modelo utiliza o AdaViT como baseline. Redes de decisão leves recebem o token CLS (que codifica o contexto global) da camada residual e geram logits de importância para cada cabeça. Máscaras binárias são aplicadas via amostragem Gumbel-Sigmoid para podar cabeças.
• Treinamento do SAE:
  • Um SAE do tipo k-sparse é treinado sobre o embedding residual da camada final do ViT (token CLS).
  • O objetivo é decompor o embedding denso em um espaço latente esparsamente ativado, onde cada dimensão tende a codificar conceitos monossêmicos e interpretáveis.
  • A arquitetura usa uma ativação Top-K para manter apenas as $k$ maiores dimensões ativas, garantindo esparsidade.
• Mecanismo de Direcionamento (Steering):
  • Após o treinamento, dimensões latentes específicas do SAE são amplificadas (adicionando um valor $\alpha$).
  • O vetor latente amplificado é reconstituído no espaço de embeddings e alimentado novamente na rede de decisão de poda.
  • Isso altera a distribuição de probabilidade das cabeças, forçando a rede de decisão a escolher um subconjunto diferente de cabeças.

3. Contribuições Principais

1. Framework de Controle Latente: Introdução de um método que utiliza SAEs para "direcionar" (steer) as decisões de poda dinâmica, permitindo controle explícito sobre quais cabeças são ativadas ou desativadas.
2. Interpretabilidade Mecânica: Demonstra que os latentes esparsos capturam conceitos discriminativos por classe, revelando subconjuntos de cabeças específicos para cada categoria.
3. Descoberta de Relações Semânticas: O método revela que classes semanticamente relacionadas (ex: "tigela" e "prato") compartilham os mesmos subconjuntos de cabeças quando direcionadas, validando a interpretabilidade do modelo.
4. Eficiência e Precisão: Prova que é possível reduzir drasticamente o uso de cabeças computacionais mantendo ou até melhorando a acurácia para classes específicas.

4. Resultados Experimentais

Os experimentos foram realizados no ViT-Small (pré-treinado no ImageNet) ajustado para CIFAR-100.

• Baseline: O modelo poda 30% das cabeças mantendo 89,79% de acurácia.
• Estratégias de Direcionamento: Foram testadas três estratégias para selecionar quais latentes amplificar:
  1. Per-class frequent (latentes mais frequentes por classe).
  2. Global frequent (latentes mais frequentes globalmente).
  3. Random (aleatório).
• Desempenho da Estratégia "Per-class":
  • Esta estratégia foi a mais eficaz, reduzindo o uso de cabeças enquanto mantinha a acurácia.
  • Exemplo Concreto (Classe "Bowl"): Ao amplificar os latentes frequentes para esta classe, a acurácia aumentou de 76% para 82%, enquanto o uso de cabeças caiu de 0,72 para 0,33. O modelo passou a depender quase exclusivamente das cabeças h2 e h5.
  • Exemplo (Classe "Pine tree"): Acurácia subiu de 79% para 84%, com uso de cabeças caindo de 0,93 para 0,35, focando nas cabeças h2 e h3.
• Comparação: Estratégias globais ou aleatórias resultaram em quedas significativas de acurácia à medida que a força de amplificação ($\alpha$) aumentava, indicando que a especificidade por classe é crucial.
• Validação de Interpretabilidade: A baixa sobreposição (0,1641) entre os latentes frequentes globais e por classe confirma que o SAE aprendeu conceitos específicos de classe, não apenas características gerais.

5. Significado e Impacto

Este trabalho preenche uma lacuna importante entre a eficiência computacional e a interpretabilidade mecânica em Vision Transformers.

• Ao demonstrar que a poda dinâmica pode ser controlada através de latentes esparsos, o estudo oferece uma nova ferramenta para entender como os ViTs tomam decisões de arquitetura interna.
• A capacidade de "sintonizar" o modelo para ser mais eficiente em classes específicas sem sacrificar a precisão abre caminho para ViTs mais adaptáveis e explicáveis.
• O trabalho sugere que a esparsidade induzida por SAEs não é apenas uma técnica de compressão, mas uma via para desvendar a semântica interna das redes de atenção.

Trabalhos Futuros: Os autores planejam estender o framework para camadas anteriores do ViT e aplicá-lo em modelos foundation de maior escala.