Can Cross-Layer Transcoders Replace Vision Transformer Activations? An Interpretable Perspective on Vision

Resumo Técnico: Cross-Layer Transcoders (CLTs) para Vision Transformers

1. Problema e Motivação

Os Foundation Models (FMs), especialmente os Vision Transformers (ViTs), demonstraram desempenho excepcional em diversas tarefas de visão computacional. No entanto, sua natureza "opaca" limita a interpretabilidade, o controle e a confiabilidade prática.

• Limitação das Abordagens Atuais: Métodos existentes, como Sparse Autoencoders (SAEs), operam localmente em camadas individuais. Eles falham em capturar a estrutura computacional entre camadas (cross-layer) e não conseguem quantificar a contribuição relativa de cada camada na formação da representação final.
• Desafio Específico: A transposição de técnicas de interpretabilidade de Linguagem Natural (LLMs) para Visão Computacional é complexa devido às diferenças estruturais (tokens espaciais, granularidade de patches, e o token global [CLS]).
• Objetivo: Investigar se os Cross-Layer Transcoders (CLTs) podem servir como substitutos funcionais e interpretáveis para os blocos MLP (Multi-Layer Perceptron) dos ViTs, permitindo uma decomposição linear e profunda das representações internas.

2. Metodologia

Os autores propõem o uso de CLTs como modelos proxy esparsos e conscientes da profundidade (depth-aware) para os blocos MLP dos ViTs.

• Arquitetura do CLT:

  • Codificação Esparsa: Para cada camada $i$, uma codificação linear aprendida transforma as ativações pré-MLP ($x_i$) em vetores esparsos ($z_i$). São utilizadas três funções de não-linearidade para induzir esparsidade: JumpReLU, ReLU-Top-k e Abs-Top-k.
  • Decodificação Triangular (Cross-Layer): Ao contrário dos transcoders tradicionais (que operam apenas dentro de uma camada), o CLT reconstrói a saída do MLP de uma camada alvo $j$ ($\hat{y}_j$) como uma soma linear das contribuições esparsas de todas as camadas anteriores ($i \leq j$):
    $$\hat{y}_j = \sum_{i=1}^{j} z_i W_{dec}^{i \to j}$$
    Isso cria uma estrutura causal estrita, prevenindo vazamento de informação de camadas futuras.

• Treinamento e Substituição:

  • Os CLTs são treinados para reconstruir as ativações pós-MLP de um ViT congelado (professor) usando apenas as ativações das camadas anteriores.
  • Modelo de Substituição (Replacement Model): Durante a inferência, os blocos MLP originais são substituídos pelas reconstruções dos CLTs. Isso permite criar um modelo "substituto" que pode ser usado para análise sem re-treinar o modelo principal.

• Métricas de Atribuição:

  • Introduz-se uma pontuação de contribuição baseada em projeção ($C_{i \to j}^{proj}$) para quantificar quanto cada camada fonte $i$ contribui para a representação da camada alvo $j$. Isso permite decompor a representação final em termos de camadas específicas.

3. Principais Contribuições

1. CLTs para Visão: Adaptação bem-sucedida dos Cross-Layer Transcoders para o domínio de visão, demonstrando a capacidade de reconstruir representações pós-MLP de forma esparsa e entre camadas.
2. Substituição Funcional Fiel: Demonstração de que os CLTs podem substituir os blocos MLP (especialmente nas camadas finais e para o token [CLS]) mantendo ou até melhorando a precisão de classificação zero-shot.
3. Interpretabilidade via Atribuição Cruzada: Revelação de uma assimetria estrutural fundamental:
   • Tokens de Patch: Apresentam atribuição diagonal-dominante (cada camada explica principalmente sua própria saída).
   • Token [CLS]: Apresenta atribuição distribuída ao longo de toda a profundidade, agregando sinais semânticos de muitas camadas anteriores.
4. Validação de Atribuição: Experimentos de ablação mostram que a representação final é concentrada em um pequeno subconjunto de camadas dominantes. Manter apenas as 4 camadas com maior pontuação preserva a precisão, enquanto remover a camada de maior pontuação degrada significativamente o desempenho.

4. Resultados Experimentais

Os experimentos foram conduzidos em modelos CLIP ViT-B/32 e ViT-B/16 nos conjuntos de dados CIFAR-100, COCO e ImageNet-100.

• Fidelidade de Reconstrução:

  • Os CLTs alcançaram alta fidelidade na reconstrução das ativações, com similaridade de cosseno entre 0,92 e 0,97 e $R^2$ entre 0,89 e 0,95.
  • O ViT-B/16 foi mais fácil de aproximar que o B/32, provavelmente devido à maior granularidade dos patches (menos informação por token).
  • A variante ReLU-Top-k geralmente apresentou o melhor desempenho, especialmente nas camadas intermediárias e tardias.

• Desempenho de Classificação (Substituição em Cascata):

  • Token [CLS]: A substituição de todos os blocos MLP apenas para o token [CLS] manteve a precisão quase idêntica à do modelo original (ex: 61,74% vs 61,65% no CIFAR-100).
  • Tokens de Patch: A substituição completa (todas as camadas e todos os tokens) causou degradação, especialmente em camadas iniciais. No entanto, substituir apenas as camadas finais (ex: 10→11) preservou a precisão.
  • Conclusão: Os CLTs são proxies confiáveis para o token global e para as transformações de camadas tardias.

• Análise de Atribuição:

  • A ablação baseada nas pontuações de contribuição confirmou que a representação final depende criticamente de poucas camadas. A remoção da camada com maior pontuação causou quedas drásticas de precisão (ex: -5,6% no ImageNet-100), validando que as pontuações identificam fielmente as camadas mais importantes.

5. Significado e Impacto

Este trabalho estabelece os Cross-Layer Transcoders como uma ferramenta fundamental para a interpretabilidade de modelos de visão.

• Novo Paradigma: Move-se além da análise de camadas isoladas para uma compreensão estruturada de como a informação flui e se integra através da profundidade da rede.
• Descoberta Arquitetural: A descoberta de que o token [CLS] funciona como um agregador global de informações de todas as camadas, enquanto os patches operam de forma mais local, oferece insights cruciais sobre o design de ViTs.
• Aplicabilidade: A capacidade de substituir partes do modelo por proxies interpretáveis sem perder desempenho abre caminho para:
  • Descoberta hierárquica de conceitos (de bordas/texturas a abstrações semânticas).
  • Análise de circuitos de características em visão.
  • Modelos mais controláveis e verificáveis para tarefas críticas.

Em suma, o artigo demonstra que é possível "abrir a caixa preta" dos Vision Transformers mantendo sua funcionalidade, revelando que a representação final é uma construção aditiva e esparsa de contribuições profundas, onde poucas camadas dominam a decisão final.