Disentangling MLP Neuron Weights in Vocabulary Space

Título: Desemaranhando os Pesos dos Neurônios MLP no Espaço de Vocabulário

Autores: Asaf Avrahamy, Yoav Gur-Arieh, Mor Geva (Tel Aviv University)
Método Proposto: ROTATE (Rotation-Optimized Token Alignment in weighT spacE)

---

1. O Problema

A interpretabilidade mecânica de Grandes Modelos de Linguagem (LLMs) enfrenta um desafio fundamental: entender como a informação é codificada nos pesos do modelo, independentemente de dados de entrada específicos.

• Polissemantismo: A maioria dos componentes internos (como neurônios em camadas MLP) é "polissêmica", ou seja, um único neurônio codifica múltiplos conceitos entrelaçados, dificultando a interpretação direta.
• Limitações das Abordagens Atuais: Métodos existentes, como Autoencoders Esparsos (SAEs), geralmente requerem grandes volumes de dados de ativação (forward passes) para aprender representações monossêmicas. Isso é computacionalmente custoso e introduz viéses relacionados à distribuição dos dados de treinamento.
• Necessidade: É necessário um método que disente os pesos dos neurônios diretamente no espaço de pesos, sem depender de dados de ativação, para revelar direções interpretáveis e monossêmicas.

2. Metodologia: ROTATE

O artigo introduz o ROTATE, um método livre de dados (data-free) que não requer passes forward através do modelo. A abordagem baseia-se em uma observação estatística chave e em otimização de rotação.

A. Observação Central: Curtose no Espaço de Vocabulário

Os autores observam que neurônios que codificam conceitos coerentes e monossêmicos exibem alta curtose quando seus vetores de peso são projetados no espaço de vocabulário do modelo (via multiplicação pela matriz de unembedding $U$).

• Curtose: Mede o "peso das caudas" de uma distribuição. Uma alta curtose indica que a distribuição é fortemente concentrada em poucos tokens (picos agudos) com caudas pesadas, em vez de ser uma distribuição gaussiana difusa (que indicaria polissemantismo ou ruído).
• Hipótese: Maximizar a curtose da projeção de um peso no vocabulário revela direções monossêmicas.

B. O Algoritmo ROTATE

Dado um vetor de peso de um neurônio $w$, o ROTATE aprende uma matriz de rotação $R$ para transformar $w$ em um vetor $v = Rw$ que maximize a curtose no espaço de vocabulário ($z = vU$).

1. Função de Perda: Otimiza-se uma função composta por dois termos:
   • Perda de Curtose ($L_{kurt}$): Maximiza a curtose da projeção $z$ (empurrando $w$ para direções monossêmicas).
   • Perda de Regularização ($L_{reg}$): Penaliza a distância cosseno entre o vetor rotacionado $v$ e o original $w$, garantindo que os canais descobertos sejam fiéis ao neurônio original e não direções arbitrárias.
2. Parametrização: Utiliza matrizes de Householder para garantir que as rotações sejam ortogonais.
3. Descoberta Iterativa (Mascaramento): Como um neurônio pode conter múltiplos conceitos, o algoritmo é executado iterativamente. Após descobrir um canal (vetor $v_i$), os tokens que contribuem mais para a curtose desse canal são mascarados (excluídos) na próxima iteração. Isso força o algoritmo a encontrar novos canais de alta curtose que representam conceitos distintos, evitando a convergência para a mesma direção dominante.
4. Canais de Vocabulário: Os vetores resultantes $\{v_i\}$ são chamados de "canais de vocabulário", pois são projeções do neurônio original alinhadas com a base de vocabulário do modelo.

3. Contribuições Principais

1. Observação Estatística: Demonstração de que a distribuição de alta curtose no espaço de vocabulário é um sinal confiável de direções monossêmicas em LMs.
2. Método ROTATE: Introdução de uma técnica de decomposição de pesos livre de dados que isola componentes interpretáveis sem necessidade de dados de ativação.
3. Superioridade sobre Baselines: Evidência experimental de que os canais descobertos pelo ROTATE superam métodos baseados em SAEs (como Gemma Scope e Llama Scope) em termos de fidelidade e cobertura do espectro de ativação do neurônio.
4. Descrições de Neurônios: Demonstração de que agregar as descrições dos canais de vocabulário gera descrições naturais de neurônios mais precisas e abrangentes do que as abordagens baseadas apenas em ativações máximas.

4. Resultados Experimentais

Os experimentos foram realizados em modelos Gemma-2-2B-it e Llama-3.1-8B-Instruct.

• Fidelidade (Input e Output):

  • Input-side: Os canais descobertos pelo ROTATE preveem com maior precisão quando e por que um neurônio se ativa. Em testes de ablação, remover um canal específico suprime seletivamente a ativação do neurônio em exemplos relacionados a esse conceito, sem afetar outros.
  • Output-side: A ablação de canais de saída remove causalmente a promoção de conceitos específicos pelo neurônio, confirmando que os canais controlam funções independentes.
  • Comparação: O ROTATE obteve pontuações de fidelidade significativamente mais altas (ex: 0.71 vs 0.49 em Llama-3.1, camada 18) comparado a baselines baseadas em SAEs.

• Completude:

  • O conjunto de canais descobertos cobre uma fração maior do espaço de ativação do neurônio. O ROTATE conseguiu mapear corretamente a descrição do canal para a entrada ativa em 49-60% dos casos, superando os SAEs (36-49%).

• Descrição de Neurônios (Head-to-Head):

  • Ao agregar as descrições dos canais, o ROTATE superou as melhores abordagens baseadas em ativação (MaxAct++) e métodos híbridos (MaxAct+VocabProj) em 2-3 vezes nas taxas de vitória em comparações pareadas. O ROTATE foi particularmente superior em descrever comportamentos em níveis de ativação moderados, onde métodos baseados em picos de ativação falham.

5. Significado e Impacto

• Escalabilidade: Por ser livre de dados e operar apenas nos pesos, o ROTATE é altamente escalável e pode ser aplicado a qualquer modelo de linguagem sem a necessidade de coletar e processar terabytes de dados de ativação.
• Interpretabilidade Fina: Permite uma compreensão mecânica mais profunda, decompondo neurônios polissêmicos em seus componentes constituintes (canais) que correspondem a conceitos humanos legíveis (ex: "tempo", "negação", "programação").
• Fundação para Futuras Pesquisas: Oferece um bloco de construção fundamental para a descoberta de circuitos mais precisos e para o estudo de interações entre componentes da rede, movendo a interpretabilidade de uma análise baseada em dados para uma análise baseada na estrutura intrínseca do modelo.

Em resumo, o ROTATE demonstra que a estrutura semântica dos modelos de linguagem está codificada diretamente em seus pesos e pode ser acessada através da otimização de propriedades estatísticas (curtose), oferecendo uma alternativa superior e mais eficiente aos métodos baseados em ativação.