Resumo Técnico: Transporte Ótimo Semântico para Correspondência de Características de Autoencoders Esparsos e Compressão de Circuitos

1. Definição do Problema

Autoencoders Esparsos (SAEs) emergiram como uma ferramenta primária para interpretar Modelos de Linguagem de Grande Escala (LLMs) ao decompor espaços latentes em características interpretáveis. No entanto, dois desafios críticos na escalabilidade da análise de SAEs permanecem:

1. Correspondência de Características entre Camadas: Identificar características semanticamente similares em diferentes camadas de um LLM. Métodos existentes geralmente dependem da comparação de vetores de decodificador. Essa abordagem falha ao correspondecer características entre camadas não consecutivas porque os vetores de decodificador são treinados para reconstruir distribuições de ativação distintas, tornando a comparação direta de vetores imprecisa.
2. Compressão de Circuitos de Características: Circuitos de características (grafos causais conectando características entre camadas) frequentemente contêm centenas de nós, tornando a inspeção manual impraticável. Métodos atuais de compressão dependem do agrupamento manual de nós semanticamente similares em "supernós", um processo que não escala.

O problema central identificado pelos autores é a estimativa da distância semântica entre características que residem em diferentes variedades de ativação. Abordagens existentes tratam características como vetores únicos (por exemplo, pesos de decodificador), o que descarta o rico contexto de onde e quão fortemente uma característica ativa, levando à perda de informação e à incapacidade de lidar com comparações multi-variedade.

2. Metodologia

Os autores propõem um framework distribucional que reformula a comparação de características como um problema de Transporte Ótimo (OT).

2.1 Representação de Características como Distribuições

Em vez de representar uma característica por um único vetor de decodificador, o método representa cada característica $i$ na camada $\ell$ como uma distribuição empírica ponderada por ativação sobre os estados ocultos onde ela ativa.

• Dado um corpus de posições de tokens $T$, o método seleciona as top-$K$ posições de tokens onde a característica $i$ ativa mais fortemente.
• Cada token selecionado $t$ é atribuído um peso $w_{i,t}$ proporcional à sua magnitude de ativação.
• A característica é representada como uma distribuição $\hat{\mu}_{i,T} = \sum w_{i,t} \delta_{x_t}$, onde $x_t$ é o estado oculto no token $t$.

2.2 Espaço de Referência Compartilhado

Para comparar características entre diferentes camadas (que residem em variedades diferentes), o método projeta essas distribuições específicas de camada em um espaço de referência compartilhado $\mathcal{R}$.

• Para correspondência entre camadas, o espaço de referência é tipicamente o espaço de estado oculto da camada de destino. Ativações da camada de origem são projetadas para as posições de token correspondentes na camada de destino.
• Para compressão de circuitos, o espaço de referência é a concatenação de estados ocultos através de todas as camadas.

2.3 Distância Semântica via Distância de Wasserstein

Uma vez projetadas no espaço compartilhado, a distância semântica entre duas características é computada usando a distância de Wasserstein (Distância do Movimentador de Terra).

• A distância $W_c(\hat{\mu}, \hat{\nu})$ mede o custo mínimo para transportar a massa de probabilidade da distribuição de ativação de uma característica para a outra.
• O custo base $c(z, z')$ mede a despesa de correspondência de duas representações de contexto no espaço compartilhado.
• Essa abordagem captura a estrutura geométrica dos padrões de ativação, em vez de apenas a similaridade vetorial ponto a ponto.

2.4 Tarefas a Montante

• Correspondência de Características: Para uma característica de destino, o método encontra a característica de origem com a menor distância de Wasserstein (vizinho mais próximo).
• Compressão de Circuitos: O método trata nós de circuitos como pontos no espaço distribucional e aplica agrupamento (por exemplo, agrupamento aglomerativo) para agrupar nós com pequenas distâncias de Wasserstein parciais em supernós, minimizando a distância semântica dentro do cluster.

3. Contribuições Principais e Garantias Teóricas

O artigo fornece três garantias teóricas principais para o framework proposto:

1. Invariância de Redimensionamento de Ativação: A distribuição de característica projetada é invariante à escalonamento global dos valores de ativação. Como a representação depende do ranking relativo e da distribuição proporcional das ativações (seleção Top-K e pesos normalizados), deslocamentos escalares arbitrários nas magnitudes de ativação não alteram a distância semântica.
2. Estabilidade sob Perturbações: A distância de Wasserstein entre distribuições de características é estável sob perturbações suaves das distribuições. A mudança na distância é limitada pelas constantes de Lipschitz do custo base e pela distância de Wasserstein entre as distribuições perturbadas e originais. Isso garante robustez contra ruído na seleção de tokens ou no retreinamento de SAE.
3. Recuperação Baseada em Margem: Os autores estabelecem condições sob as quais a correspondência empírica de vizinho mais próximo e o agrupamento semântico recuperam perfeitamente as correspondências verdadeiras (oráculo). Se a "margem" (a lacuna entre a melhor correspondência e a segunda melhor correspondência) exceder o dobro do erro de estimativa, a decisão empírica é garantida como correta.

4. Resultados Experimentais

O método foi avaliado em SAEs treinados em GPT-2 e Gemma-2-2B, comparando-se com baselines incluindo SAE Match (L2 de vetor de decodificador), Feature Flow (cosseno de vetor de decodificador), Patching de Atribuição e seleção baseada em LLM.

• Correspondência de Características: O método proposto superou todas as baselines em Avaliação de LLM (classificações de similaridade semântica) e Precisão de Correspondência em configurações de camadas consecutivas, do meio para o final, e do início para o final. Crucialmente, manteve desempenho competitivo em métricas a montante (Perda de Entropia Cruzada e Variância Explicada), demonstrando que identifica características com funções semelhantes, e não apenas tokens de ativação similares. Notavelmente, teve sucesso em correspondecer características sutis de "adição de dígitos" onde outros métodos falharam.
• Compressão de Circuitos: O método alcançou maior precisão na compressão de circuitos de características em supernós intuitivos, comparado ao agrupamento manual e a outras baselines. Identificou com sucesso grupos modulares de características (por exemplo, "cabeças de indução" e "detecção de tokens") que consistentemente desempenham o mesmo papel funcional em circuitos diferentes.
• Estudo de Ablação: Experimentos confirmaram que representar características como distribuições preserva mais informação semântica do que colapsá-las em um único vetor centróide, mesmo quando projetadas na mesma variedade.

5. Significado e Alegações

O artigo alega fornecer uma métrica semântica unificada que resolve tanto a correspondência de características entre camadas quanto a compressão de circuitos como instâncias do mesmo problema fundamental: estimar a distância semântica através de variedades.

• Escalabilidade: Ao automatizar o agrupamento de características em supernós, o método remove a dependência de rotulagem humana, tornando viável a análise de circuitos de características em grande escala.
• Robustez: A abordagem distribucional captura o "contexto" da ativação de características, permitindo correspondência precisa mesmo quando as características evoluem significativamente entre camadas (distribuições dissimilares), um cenário onde métodos baseados em vetores falham.
• Interpretabilidade: A capacidade de comprimir automaticamente circuitos complexos em supernós interpretáveis aprimora a compreensão dos mecanismos do modelo, conforme demonstrado em estudos de caso sobre circuitos de "Dallas" (geografia), "Rima" e "Oposto".

Os autores posicionam este trabalho como um passo fundamental rumo à interpretabilidade mecânica escalável e automatizada, superando as limitações da comparação de características centrada em vetores.