Resumo Técnico: Autoencoders Esparsos Residualizados para Intervenções em Transformadores de Múltiplas Camadas

Declaração do Problema

Autoencoders Esparsos (SAEs) são uma ferramenta primária para interpretabilidade mecânica, decompondo estados ocultos de transformadores em características esparsas e interpretáveis. No entanto, a prática padrão treina SAEs independentemente nas ativações brutas de cada camada. Essa abordagem ignora uma propriedade estrutural fundamental dos transformadores: o fluxo residual é fortemente acoplado através da profundidade. Uma fração significativa da ativação de qualquer camada consiste em informações transportadas adiante ou levemente transformadas a partir de camadas anteriores.

Treinar SAEs independentes em ativações brutas cria dois problemas práticos para intervenções de múltiplas camadas:

1. Redundância de Capacidade: Dicionários por camada desperdiçam capacidade reaprendendo as mesmas informações previsíveis e transportadas adiante.
2. Interações Imprevisíveis: Quando múltiplas camadas são substituídas simultaneamente (por exemplo, via ablação ou reconstrução de características), as interações entre as camadas não são aditivas. Se a mesma informação estiver representada em múltiplos dicionários, a substituição simultânea pode causar a duplicação ou cancelamento de características, ou levar as ativações para fora da distribuição de treinamento, resultando em erros compostos e desempenho degradado do modelo (entropia cruzada).

Metodologia: Autoencoders Esparsos Residualizados (ReSAEs)

Os autores propõem Autoencoders Esparsos Residualizados (ReSAEs) para desacoplar informações específicas de camada de estruturas cruzadas entre camadas previsíveis. O método envolve três etapas principais:

1. Residualização Afinada Pares a Pares

Para um subconjunto selecionado de camadas $L = (\ell_1, \ell_2, \dots, \ell_M)$, o método ajusta um mapa afim entre camadas selecionadas consecutivas. Para um par $(\ell_m, \ell_{m+1})$, uma transformação afim $A_m h^{(\ell_m)} + c_m$ é aprendida via mínimos quadrados ordinários regularizados por crista para prever a ativação da camada posterior a partir da anterior.
A ativação residualizada $r^{(\ell_{m+1})}$ é então definida como o componente da ativação da camada posterior não explicado por este mapa afim:
$$r^{(\ell_{m+1})} = h^{(\ell_{m+1})} - (A_m h^{(\ell_m)} + c_m)$$
A primeira camada âncora é centralizada por sua média empírica.

2. Treinamento Independente de SAE por Camada

Cada SAE é treinado no resíduo em vez da ativação bruta.

• Entrada: A ativação residual é normalizada (via normalização RMS por bloco) e alimentada em um SAE padrão BatchTopK.
• Objetivo: O SAE minimiza a perda de reconstrução do alvo residual.
• Independência: Os gradientes não fluem entre SAEs de blocos diferentes; cada SAE aprende apenas a informação "nova" introduzida em sua camada específica.

3. Reconstrução no Espaço Original e Intervenção Online

Para avaliar os ReSAEs usando protocolos padrão, o resíduo reconstruído é mapeado de volta para o espaço de ativação original através da cadeia afim ajustada.

• Reconstrução: $\hat{h}^{(\ell_{m+1})} = A_m \hat{h}^{(\ell_m)} + c_m + S^{-1}_{m+1} \hat{z}^{(m+1)}$, onde $\hat{z}$ é a reconstrução do SAE do resíduo.
• Intervenção Online: Durante uma passagem forward com ganchos (hooks), a entrada residual para a camada $\ell_m$ é computada subtraindo a contribuição prevista da camada anteriormente reconstruída $\hat{h}^{(\ell_{m-1})}$ da ativação de entrada atual. Isso garante que o SAE na camada $\ell_m$ veja uma distribuição de entrada semelhante ao seu alvo de treinamento, mesmo que as camadas a montante tenham sido modificadas.

Contribuições Principais

1. Arquitetura ReSAE: Introdução de um paradigma de treinamento onde SAEs são treinados no componente da ativação não linearmente previsível a partir de uma camada selecionada anterior, com reconstruções compostas via uma cadeia afim.
2. Intervenção de Múltiplas Camadas Aprimorada: Demonstração de que, apesar de reconstruir menos variância bruta, os ReSAEs recuperam a entropia cruzada do transformador de forma mais fiel sob substituição simultânea de múltiplas camadas em comparação com SAEs brutos.
3. Redução de Redundância: Evidências de que a residualização produz direções de decodificador mais ortogonais, indicando que a capacidade é direcionada para informações locais à camada em vez de estruturas cruzadas entre camadas duplicadas.
4. Avaliação Abrangente: Benchmarks extensivos em Pythia-1.4B e Gemma-2-9B cobrindo qualidade de reconstrução, degradação de entropia cruzada, similaridade de decodificador, sondagem esparsa, perturbação de sonda direcionada (TPP) e remoção de correlações espúrias (SCR).

Resultados

Os experimentos foram conduzidos em Pythia-1.4B (camadas 0, 6, 12, 18) e Gemma-2-9B (vários intervalos de camadas).

• Reconstrução vs. Funcionalidade: Os ReSAEs exibem menor Variância Explicada (EV) no espaço de ativação original em comparação com SAEs brutos, pois descartam intencionalmente estruturas linearmente previsíveis. No entanto, essa EV menor correlaciona-se com menor similaridade de cosseno do decodificador (menos redundância).
• Recuperação de Entropia Cruzada: Sob substituição sequencial-online de múltiplas camadas (onde erros se propagam), os ReSAEs superam significativamente os SAEs brutos na recuperação da entropia cruzada do modelo, particularmente em níveis de esparsidade mais altos ($k \ge 160$). Em configurações forçadas pelo professor (sem propagação de erros), os ReSAEs consistentemente alcançam a melhor recuperação de entropia cruzada.
• Sobreinteração: SAEs brutos mostram grande "sobreinteração" negativa (degradação não aditiva) quando múltiplas camadas são substituídas. Os ReSAEs reduzem essa métrica para próximo de zero, indicando que suas substituições por camada compõem-se de forma mais previsível.
• Qualidade das Características:
  • Sondagem Esparsa e TPP: Os ReSAEs geralmente superam os SAEs brutos em precisão de sondagem esparsa e perturbação de sonda direcionada, sugerindo que as características aprendidas são mais úteis para leitura linear e ablação seletiva.
  • Remoção de Correlações Espúrias (SCR): Os ReSAEs não superam uniformemente os SAEs brutos aqui; em algumas configurações (particularmente ablações conjuntas por camada no Pythia), os SAEs brutos obtêm pontuações mais altas. Isso sugere que características residualizadas podem ser menos seletivas para remover atributos espúrios específicos sem danificar sinais de tarefa.
• Escala e Espaçamento: Os benefícios dos ReSAEs persistem no modelo maior Gemma-2-9B através de diferentes intervalos de camadas, com reduções consistentes na redundância de decodificador e melhorias na recuperação de entropia cruzada.

Significado e Afirmações

O artigo argumenta que tratar o fluxo residual do transformador como um objeto acoplado em vez de uma coleção de camadas isoladas é crucial para intervenções eficazes de múltiplas camadas.

• Alocação de Capacidade: Ao remover estruturas cruzadas entre camadas linearmente previsíveis, os ReSAEs forçam o dicionário esparsos a focar nos componentes da ativação mais relevantes para o cálculo a jusante do modelo naquela camada específica.
• Estabilidade da Intervenção: O método fornece uma base mais estável para intervenções de múltiplas camadas, reduzindo os erros compostos e interações imprevisíveis que afligem implantações padrão de SAE.
• Compromisso: Os autores reconhecem que os ReSAEs não são uma melhoria universal para todas as tarefas de interpretabilidade (por exemplo, os resultados de SCR são mistos). No entanto, para o objetivo específico de substituição de múltiplas camadas e intervenção baseada em características, o compromisso de menor variância bruta por maior fidelidade funcional e menor redundância é benéfico.

O trabalho sugere que a remoção de estruturas previsíveis via regressão afim é uma etapa de pré-processamento útil para SAEs em configurações de múltiplas camadas, embora a escolha ótima do preditor (afim vs. não linear) e a seleção de camadas permaneçam um espaço de design aberto.