Resumo Técnico: Similaridade de Tensores para Interpretabilidade Mecanística

Declaração do Problema

A interpretabilidade mecanística busca decompor redes neurais em partes constituintes significativas para compreender seus cálculos internos. Um pré-requisito crítico para este campo é a capacidade de verificar rigorosamente se dois componentes decompostos (ou dois modelos diferentes) implementam a mesma função subjacente. Os métodos de avaliação atuais sofrem de duas limitações primárias:

1. Similaridade Empírica/Comportamental: Métodos que comparam saídas ou ativações de modelos em conjuntos de dados específicos são cegos a mecanismos fora da distribuição (OOD). Eles não conseguem detectar problemas como backdoors que são acionados apenas em entradas específicas ou falham em generalizar se a decomposição depender de heurísticas aproximadas em vez da função verdadeira.
2. Similaridade Baseada em Parâmetros: Métodos que comparam matrizes de pesos (por exemplo, similaridade de cosseno de matrizes) são altamente sensíveis a simetrias no espaço de pesos (como permutação e escala). Eles falham em reconhecer equivalência funcional quando modelos são treinados de forma diferente ou quando mecanismos estão distribuídos entre camadas em vez de concentrados em neurônios únicos.

Existe uma falta de uma métrica rigorosa e independente de conjuntos de dados que seja invariante a essas simetrias e capaz de capturar a equivalência funcional global.

Metodologia

Os autores propõem a Similaridade de Tensores, uma métrica baseada em pesos projetada para redes neurais baseadas em tensores (especificamente modelos multilineares como transformadores bilineares). Esta abordagem substitui o uso de funções de ativação não lineares padrão por arquiteturas polinomiais (multilineares) para permitir análise algébrica dos pesos.

Mudança Arquitetural Central

O método depende de modelos multilineares onde as funções de ativação são substituídas por portões multiplicativos (por exemplo, camadas bilineares e atenção bilinear).

• Camadas Bilineares: Mapeiam uma entrada elevada $\tilde{x} = (1, x)$ para uma saída via um polinômio de grau 2: $A(x) = D((L\tilde{x}) \odot (R\tilde{x}))$.
• Representação Multilinear: Essas camadas podem ser expressas como uma contração de um tensor de pesos $A$ com $n$ cópias da entrada: $A(x) = A \cdot \text{in}^{\otimes n} \tilde{x}$.
• Modelos Profundos: Redes multilineares profundas são construídas como decomposições em estrutura de árvore de tensores globais, permitindo computação eficiente sem materializar o tensor global exponencialmente grande.

A Métrica de Similaridade de Tensores

A métrica é definida como um produto interno generalizado entre dois tensores de pesos $A$ e $B$, normalizado por suas normas induzidas:
$$\text{sim}_M(A, B) = \frac{\langle A | M | B \rangle}{\|A\|_M \|B\|_M}$$
onde $M$ é um tensor métrico.

Para garantir que a métrica satisfaça os critérios de invariância a simetrias e independência de conjuntos de dados, os autores derivam uma forma específica para $M$:

1. Simetrização: Com base no isomorfismo de polarização, os autores estabelecem que dois modelos multilineares são funcionalmente equivalentes se e somente se seus tensores de pesos tiverem o mesmo representante no subespaço simétrico. A métrica projeta tensores neste subespaço simétrico para remover graus de liberdade relacionados a simetrias de permutação de entrada.
2. Métrica Gaussiana: Para refinar ainda mais a métrica e relacioná-la à similaridade comportamental, os autores definem $M$ usando a expectativa do tensor de entrada sob uma distribuição Gaussiana padrão ($x \sim \mathcal{N}(0, I)$). Usando o teorema de Isserlis, isso resulta em uma métrica de forma fechada que calcula o produto interno das ativações sobre todo o espaço de entrada Gaussiano, em vez de um conjunto de dados amostrado.
3. Computação Eficiente: Um algoritmo recursivo baseado em Gram permite o cálculo desta similaridade para redes profundas processando tensores locais camada por camada, evitando a necessidade de instanciar o tensor global completo.

Contribuições Principais

1. Uma Métrica Baseada em Pesos Principiada: O artigo introduz uma métrica que é invariante a simetrias no espaço de pesos (permutações, escala) e independente do conjunto de dados de entrada. Ela verifica formalmente a equivalência funcional verificando se dois modelos implementam a mesma lógica polinomial.
2. Análise de Forma Fechada e Livre de Dados: Ao contrário de métricas comportamentais, a similaridade de tensores não requer amostragem de dados. Ela integra sobre todo o espaço de entrada Gaussiano, tornando-a capaz de detectar mecanismos OOD (como backdoors) que são invisíveis para métricas de desempenho padrão.
3. Capacidades de Localização: A métrica pode ser "fatia" para focar em dimensões de saída específicas ou "diferenciar" tensores para isolar mudanças funcionais entre checkpoints, permitindo a atribuição de mecanismos específicos a saídas específicas.
4. Escalabilidade: O algoritmo recursivo garante que a métrica seja computacionalmente tratável para redes profundas, evitando limites de convergência exponenciais ou quadráticos associados à amostragem de Monte Carlo.

Resultados

Os autores validam a similaridade de tensores em quatro domínios distintos, demonstrando sua superioridade sobre métricas existentes (CKA, cosseno de matrizes, similaridade comportamental):

1. Esquecimento Catastrófico (Visão): Em uma configuração de treinamento progressivo no SVHN, a similaridade de tensores localiza precisamente a perda de conhecimento para dígitos específicos (por exemplo, dígito '9') quando o conjunto de dados é reduzido. Ela mostra uma queda acentuada na similaridade para a classe afetada enquanto permanece estável para outras, whereas métricas comportamentais mostram sinais ambíguos.
2. Grokking (Aritmética Modular): Em tarefas de adição modular, a similaridade de tensores rastreia a reorganização contínua do modelo em frequências de Fourier específicas. Ela revela que a reorganização estrutural continua mesmo após a precisão saltar para a perfeição, fornecendo uma visão mais legível da dinâmica de treinamento do que curvas de precisão isoladas.
3. Detecção de Backdoor (OOD): Quando um backdoor é injetado no SVHN (acionando em um patch de losango preto), métricas comportamentais padrão em dados limpos falham em detectar a mudança. A similaridade de tensores, no entanto, sinaliza o backdoor apenas a partir de dados limpos e localiza a mudança para a classe atacada específica ('9'), demonstrando sua capacidade de detectar mecanismos OOD.
4. Modelagem de Linguagem: Aplicada a um transformador de atenção bilinear de duas camadas treinado no The Pile, a similaridade de tensores revela estruturas de blocos claras e transições nítidas em regimes de treinamento que métricas existentes embaçam. Ela rastreia a formação de estatísticas complexas de n-gramas mais efetivamente do que proxies comportamentais.

Significado e Afirmações

O artigo posiciona a similaridade de tensores como uma ferramenta fundamental para a interpretabilidade mecanística. Seu significado primário reside em reduzir o problema de medir similaridade e verificar fidelidade de um problema de aproximação empírica para um problema algébrico resolvido.

• Fidelidade: Ao depender das propriedades algébricas de modelos multilineares, a métrica garante que decomposições capturem a função subjacente verdadeira em vez de heurísticas aproximadas.
• Robustez: Ela fornece afirmações de robustez mais fortes ao aplicar-se à distribuição completa, efetivamente descartando mecanismos que só são acionados em entradas específicas (como backdoors).
• Limitações: Os autores modestamente reconhecem que o método atualmente depende de modelos multilineares, o que exclui a maioria das arquiteturas state-of-the-art que usam não linearidades não polinomiais. No entanto, eles argumentam que modelos multilineares servem como um campo de testes vital para interpretabilidade e que evidências sugerem que eles podem igualar o desempenho de arquiteturas padrão.
• Perspectiva Futura: O trabalho sugere que a restrição a modelos multilineares é uma troca necessária para análise algébrica tratável, mas abre caminho para futuros métodos de atribuição (por exemplo, calcular similaridade entre caminhos em transformadores profundos) e apoia a visão de que arquiteturas multilineares são mais transparentes do que redes profundas comuns.