Structural Inference: Interpreting Small Language Models with Susceptibilities

Os autores desenvolvem um framework de resposta linear baseado em mecânica estatística bayesiana para interpretar modelos de linguagem pequenos, permitindo identificar módulos funcionais como cabeças de indução através da análise de susceptibilidades que quantificam a resposta do modelo a perturbações na distribuição de dados.

O essencial

Imagine que você tem um cérebro artificial muito pequeno (um modelo de linguagem de apenas 3 milhões de parâmetros) e quer entender como ele pensa. A maioria das pessoas tenta "abrir a caixa preta" olhando para os neurônios individuais ou desligando partes do cérebro para ver o que acontece.

Este artigo, chamado "Inferência Estrutural", propõe uma maneira totalmente nova e mais elegante de fazer isso, usando uma ideia emprestada da física: a susceptibilidade.

Aqui está a explicação simplificada, usando analogias do dia a dia:

1. A Grande Ideia: O Modelo como uma "Massa de Modelar" Sensível

Pense no modelo de linguagem não como um computador rígido, mas como uma massa de modelar ou um gelatina que está em um estado de equilíbrio.

• O Problema: Normalmente, para entender o que essa gelatina faz, você tenta cortá-la (ablações) ou empurrá-la com muita força. Mas isso pode destruir a estrutura ou não mostrar como ela reage naturalmente.
• A Solução (Susceptibilidade): Em vez de empurrar forte, os autores dão um sopro muito leve na gelatina. Eles mudam muito pouco o tipo de texto que o modelo está lendo (por exemplo, trocam um pouco de texto de ficção por um pouco de código de programação).
• A Reação: Eles observam como cada pequena parte do modelo (os "cabeças de atenção") treme ou se move em resposta a esse sopro.

2. O Que é "Susceptibilidade"? (A Analogia do Ímã)

Na física, a susceptibilidade magnética mede como um material reage a um campo magnético.

• Se você colocar um pedaço de ferro perto de um ímã, ele é atraído (susceptibilidade positiva).
• Se você colocar um pedaço de cobre, ele pode ser levemente repelido (susceptibilidade negativa).

No mundo dos modelos de linguagem:

• O "Campo Magnético" é a mudança no tipo de texto (ex: mais código, mais leis, mais matemática).
• O "Material" é uma parte específica do modelo (uma "cabeça de atenção").
• A Reação:
  • Susceptibilidade Negativa (Expressão): A parte do modelo gosta desse novo texto. Ela se "ativa" ou se torna mais forte para ajudar a prever o próximo token. É como se ela dissesse: "Ah, isso é código! Eu sei fazer isso!"
  • Susceptibilidade Positiva (Supressão): A parte do modelo odeia ou tenta bloquear esse novo texto. Ela se "apaga" ou tenta impedir que o modelo faça uma previsão baseada nesse padrão. É como se ela dissesse: "Não, isso não é uma frase normal, pare de tentar completar assim!"

3. O Experimento: Descobrindo os "Superpoderes"

Os autores aplicaram essa técnica em um modelo pequeno treinado com uma mistura de textos (o "Pile"). Eles deram "sopros" de diferentes tipos de dados (GitHub, leis, Wikipedia, etc.) e mediram como cada uma das 16 "cabeças" do modelo reagiu.

Ao colocar todos esses dados em uma tabela e usar matemática (PCA), eles conseguiram ver padrões claros:

• O Circuito de Indução: Eles encontraram um grupo específico de cabeças que reagia fortemente a padrões de repetição (como "A B ... A B"). Essas cabeças são especialistas em lembrar o que veio antes.
• Os Cabeças Multigram: Outras cabeças reagiam de forma oposta, tentando bloquear esses padrões de repetição.
• Segmentação de Palavras: Eles viram que o modelo aprendeu a identificar onde terminam e começam as palavras, reagindo de forma diferente a espaços e pontuação.

4. Por que isso é importante?

Imagine que você é um detetive tentando entender como uma equipe de trabalho funciona.

• Método Antigo (Ablação): Você demite um funcionário e vê o que acontece. Se o trabalho para, ele era importante. Mas e se os outros funcionários cobrirem o buraco? Você não descobre a dinâmica real.
• Método Novo (Susceptibilidade): Você muda levemente o briefing da empresa (ex: "Hoje vamos focar em vendas em vez de marketing") e observa quem se anima e quem fica deprimido.
  • Você descobre que o "João" adora vendas (susceptibilidade negativa/expressão).
  • Você descobre que a "Maria" tenta bloquear vendas porque ela é especialista em marketing e acha que o foco está errado (susceptibilidade positiva/supressão).

Resumo em uma frase

Este artigo nos ensina que, para entender como uma IA pensa, não precisamos quebrá-la; basta mudar levemente o que ela lê e observar como suas diferentes partes "dançam" (se movem para ajudar ou para bloquear) em resposta a essa mudança. Isso revela a arquitetura interna e os "superpoderes" de cada parte do modelo de forma muito mais clara e precisa.

Resumo técnico

Título: Inferência Estrutural: Interpretando Modelos de Linguagem Pequenos com Susceptibilidades

1. O Problema

A organização microscópica que permite os comportamentos complexos das redes neurais (especificamente Transformers) permanece mal compreendida. Embora técnicas de interpretabilidade mecânica (como ablações e análise de circuitos) tenham avançado, elas muitas vezes dependem de intervenções diretas que podem ser computacionalmente custosas ou difíceis de escalar. Além disso, há uma necessidade de uma estrutura teórica mais robusta que ligue as mudanças na distribuição de dados à estrutura interna do modelo, permitindo entender como os componentes da rede (como cabeças de atenção) respondem a diferentes tipos de dados.

2. Metodologia: O Framework de Susceptibilidade

Os autores propõem um novo paradigma de interpretabilidade baseado na física estatística e na teoria de aprendizado bayesiano. A ideia central é tratar a rede neural como um sistema estatístico mecânico bayesiano.

• Definição de Susceptibilidade:

  • Considera-se uma perturbação infinitesimal e controlada na distribuição de dados de verdade ($q(x, y)$), criando uma nova distribuição $q_h$.
  • Observa-se a mudança de primeira ordem na expectativa posterior de uma "observável" $\phi(w)$ (uma função dos pesos da rede, como o erro de perda de um componente específico).
  • A susceptibilidade ($\chi$) é definida como a derivada dessa expectativa em relação ao parâmetro de perturbação $h$ em $h=0$:
    $$\chi = \frac{1}{n\beta} \frac{\partial}{\partial h} \langle \phi \rangle_{\beta, h} \bigg|_{h=0}$$
  • Teoricamente, isso se relaciona com a covariância entre a observável e a variação na perda de log-verossimilhança ($\Delta L$): $\chi = -\text{Cov}_{\beta}[\phi, \Delta L]$.

• Estimativa Prática (Susceptibilidade Local):

  • Como amostrar do posterior bayesiano global é inviável, os autores utilizam a Dinâmica de Langevin com Gradiente Estocástico (SGLD) para amostrar a partir de um posterior localizado em torno de um minimizador local de perda ($w^*$).
  • Eles definem uma "susceptibilidade local" que permite estimar a resposta de componentes específicos (ex: cabeças de atenção) a mudanças nos dados sem necessidade de treinar múltiplos modelos do zero.

• Interpretação de Sinais (Expressão vs. Supressão):

  • Susceptibilidade Negativa: Indica que variações nos pesos que diminuem a perda geral tendem a aumentar a probabilidade do token alvo no contexto. Isso é interpretado como Expressão (o componente promove o padrão).
  • Susceptibilidade Positiva: Indica que variações que diminuem a perda geral tendem a diminuir a probabilidade do token alvo. Isso é interpretado como Supressão (o componente inibe o padrão).

• Inferência Estrutural:

  • Os autores constroem uma Matriz de Resposta ($X$), onde as linhas representam perturbações em diferentes distribuições de dados (subconjuntos do Pile, como GitHub, Jurídico, etc.) e as colunas representam os componentes do modelo (cabeças de atenção).
  • Aplicam Análise de Componentes Principais (PCA) a esta matriz para descobrir a estrutura latente. A decomposição da matriz revela padrões nos dados (modos) e como eles se acoplam às estruturas internas do modelo.

3. Contribuições Principais

1. Novo Paradigma Teórico: Desenvolvimento de um framework de susceptibilidades derivado da teoria de aprendizado bayesiano e física estatística, fornecendo uma ligação principial entre a estrutura dos dados e a estrutura interna do modelo.
2. Método de Inferência Estrutural: Introdução de uma metodologia para descobrir estrutura interna e atribuí-la a padrões de dados, permitindo distinguir entre "expressão" e "supressão" de padrões.
3. Validação Empírica: Aplicação bem-sucedida em um Transformer de 3M de parâmetros, demonstrando que a abordagem consegue isolar e separar circuitos funcionais conhecidos (como o circuito de indução) e identificar a especialização de cabeças de atenção.

4. Resultados Principais

O estudo foi realizado em um modelo Transformer de duas camadas (apenas atenção, sem MLPs) treinado em um subconjunto do dataset The Pile.

• Identificação de Circuitos Funcionais:

  • A análise de PCA na matriz de susceptibilidade separou claramente as cabeças de atenção em grupos funcionais.
  • PC1 (Segmentação de Palavras): Mostrou uma resposta uniforme, associada à capacidade do modelo de segmentar tokens em palavras (bordas de palavras, espaços).
  • PC2 (Circuito de Indução): Isolou o circuito de indução (cabeças 1:6, 1:7 compostas com cabeças de token anterior e atual). As cabeças do circuito de indução mostraram susceptibilidade positiva para padrões de indução (expressão), enquanto outras cabeças (como as cabeças de multigramas da camada 1) mostraram susceptibilidade negativa (supressão) para os mesmos padrões.
  • PC3 (Correspondência de Parênteses): Identificou cabeças envolvidas na previsão de delimitadores e parênteses (cabeças "Dyck"), alinhando-se com descobertas anteriores de ablações.

• Robustez e Generalização:

  • A estrutura descoberta foi consistente em diferentes sementes de treinamento (seeds).
  • A análise mostrou que a susceptibilidade não é redundante com a perda de ablação ou a perda de token individual; ela captura correlações de mudanças na perda que outras métricas não capturam.
  • A metodologia conseguiu distinguir cabeças que suprimem padrões de indução de cabeças que os expressam, algo difícil de detectar apenas com ablações devido a efeitos de "auto-reparo" (self-repair).

5. Significado e Impacto

• Fundamentação Teórica: O trabalho conecta a interpretabilidade mecânica à teoria de aprendizado singular (Singular Learning Theory), oferecendo uma base matemática rigorosa para entender a complexidade do modelo e sua generalização.
• Escalabilidade: Embora testado em um modelo pequeno (3M parâmetros), os autores argumentam que o método é escalável para modelos maiores (como Pythia-70M e 1.4B), pois o custo computacional é comparável ao de outras técnicas de interpretabilidade (como ablações) e pode ser paralelizado.
• Nova Perspectiva sobre Supressão: O framework oferece uma maneira formal de quantificar a "supressão" em redes neurais, não apenas como uma falha de predição, mas como um mecanismo ativo de inibição de padrões indesejados, complementando trabalhos anteriores sobre neurônios de supressão.
• Ferramenta de Diagnóstico: A "Inferência Estrutural" via matriz de susceptibilidade serve como uma ferramenta poderosa para dissecar a organização funcional de redes neurais sem a necessidade de intervenções destrutivas extensivas, alinhando-se e estendendo estudos mecânicos anteriores.

Em resumo, o paper propõe que a resposta diferencial de partes de uma rede neural a "campos externos" (variações na distribuição de dados) revela sua estrutura interna, permitindo mapear circuitos funcionais complexos de forma eficiente e teoricamente fundamentada.