Resumo Técnico: Indução em Ambas as Direções

Declaração do Problema

Embora os mecanismos internos de transformadores autorregressivos (AR) tenham sido extensivamente estudados através da lente da interpretabilidade mecanística, a mecânica operacional dos Modelos de Linguagem de Difusão (DLMs) permanece amplamente inexplorada. Os DLMs geram texto via denoising iterativo em vez de predição sequencial do próximo token, utilizando atenção bidirecional e objetivos de treinamento distintos. Existe uma lacuna crítica em entender se os DLMs implementam a indução — um mecanismo canônico de aprendizado em contexto onde um modelo identifica um contexto repetido e copia o token subsequente — da mesma maneira que os modelos AR. Especificamente, não está claro como os DLMs lidam com a indução quando o token alvo está oculto atrás de um token de máscara [M] e cercado por contexto bidirecional, em oposição ao contexto estritamente causal disponível para os modelos AR.

Metodologia

Os autores conduziram um estudo mecanístico controlado comparando transformadores AR de atenção-apenas pareados e DLMs de máscara de absorção (absorbing-mask).

• Arquitetura do Modelo: O estudo utilizou modelos de pequena escala (profundidades $L \in \{1, 2, 3\}$, dimensão oculta $d_{model}=512$, 16 cabeças de atenção) treinados em sequências FineWeb-Edu. Ambas as famílias de modelos compartilharam arquiteturas e receitas de treinamento idênticas, diferindo apenas em suas máscaras de atenção (causal para AR, bidirecional para DLM) e objetivos de treinamento (predição do próximo token vs. difusão de máscara de absorção).
• Tarefa de Avaliação: As capacidades de indução foram medidas usando sequências sintéticas de comprimento 512, onde a primeira metade foi repetida como a segunda metade. A tarefa envolveu a predição de tokens mascarados na segunda metade baseando-se no conteúdo não mascarado da primeira metade (indução direta/forward) e vice-versa (indução reversa/reverse).
• Técnicas de Análise:
  • Ablações Causais: Ablações médias foram realizadas em cabeças de atenção individuais para determinar sua necessidade causal para a indução.
  • Decomposição de Pesos: O estudo empregou decomposições QK (Query-Key) e OV (Output-Value) para analisar como os logits de atenção são formados e como a informação flui da fonte para o alvo.
  • Probing Linear: Um probe linear foi treinado no fluxo residual para detectar a presença de informação da fração global de máscara.
  • Experimentos de Patching: Técnicas de "patching" foram usadas para intervir causalmente na representação da taxa de máscara global para observar mudanças na entropia de predição.
• Variantes de Modelo: Para a análise mecanística, os autores utilizaram DLMs "dobrados sem LayerNorm" (folded no-LayerNorm), convertendo modelos treinados com LayerNorm para mapas lineares para facilitar a decomposição aditiva do fluxo residual, verificando simultaneamente que o desempenho permanecia comparável aos modelos originais com LayerNorm.

Contribuições Principais e Resultados

1. Circuito de Indução Bidirecional

O estudo identifica que os DLMs aprendem um circuito de indução bidirecional que é simétrico em direção. Diferente dos modelos AR, que dependem de uma única via causal, os DLMs utilizam duas vias distintas:

• Cabeças de Token Anterior e Próximo Token: Na primeira camada (Camada 0), cabeças específicas atendem aos vizinhos imediatos à esquerda e à direita do token mascarado, escrevendo essas características de contexto local no fluxo residual.
• Cabeças de Indução: Em camadas subsequentes (Camada 1+), as cabeças de indução utilizam pesos QK para fazer o correspondimento das características escritas pelas cabeças da Camada 0 na posição do mask com as características correspondentes na posição da fonte da resposta.
• Mecanismo: A decomposição QK revela que as cabeças de indução utilizam Q-composição (query construída a partir do contexto local da máscara) e K-composição (key construída a partir do contexto local da fonte) para localizar a fonte correspondente. O caminho OV então copia o token da fonte para a posição da máscara.

2. Acesso ao Contexto vs. Força do Mecanismo

Os autores demonstram que o desempenho superior de indução dos DLMs em comparação aos modelos AR não se deve a um mecanismo de indução fundamentalmente mais forte, mas sim ao acesso ao contexto bidirecional.

• Quando restritos apenas ao contexto à esquerda (imitando o cenário AR), os DLMs performam no mesmo nível ou abaixo dos modelos AR.
• Os DLMs superam os modelos AR apenas quando os vizinhos à esquerda e à direita do token mascarado estão visíveis.
• A pontuação de indução no cenário bidirecional total aproxima-se da soma das pontuações das configurações de apenas esquerda e apenas direita, confirmando que a vantagem advém da utilização de informações de ambas as direções.

3. Timestep Implícito via Taxa de Máscara Global

O artigo fornece evidência causal de que os DLMs computam e utilizam a fração global de tokens mascarados como um timestep implícito, apesar de carecerem de embeddings de timestep explícitos.

• Representação: Um probe linear consegue recuperar a fração global de máscara do fluxo residual após a primeira camada de atenção com alta precisão ($R^2 \approx 0.82-0.91$).
• Utilidade Causal: O patching da "direção da taxa de máscara" de um estado de alta corrupção para um estado de baixa corrupção aumenta causalmente a entropia de predição na posição da máscara. Isso confirma que o modelo utiliza o nível de corrupção global para modular sua incerteza, tratando efetivamente a taxa de máscara como um timestep.

4. Robustez e Emergência do Circuito

• Emergência: A indução emerge abruptamente durante o treinamento (por volta de 8–24k passos) e requer pelo menos duas camadas para funcionar, consistente com a necessidade de uma camada para escrever características e uma camada subsequente para recuperá-las.
• Janela Local: O mecanismo de indução depende fortemente de uma janela local de aproximadamente cinco tokens (o token mascarado mais dois vizinhos de cada lado). O desempenho cai significamente se esses vizinhos imediatos forem mascarados, mesmo que a taxa de corrupção global seja baixa.

Significância e Alegações

Este trabalho representa um passo fundamental para entender os mecanismos internos dos DLMs, indo além da suposição de que eles simplesmente replicam comportamentos AR. Os autores alegam que:

1. Divergência Mecanística: Os DLMs implementam a indução através de um circuito bidirecional simétrico que difere estruturalmente dos circuitos de indução causais em modelos AR, especificamente na forma como as características de query e key são construídas a partir das posições mascaradas.
2. Vantagem Contextual: A lacuna de desempenho entre DLMs e modelos AR no aprendizado em contexto é atribuível à disponibilidade de contexto bidirecional, e não a uma superioridade inerente no algoritmo de indução em si.
3. Gestão de Estado Implícito: Os DLMs derivam internamente a informação temporal (timestep) a partir do estado global da sequência (fração de máscara) sem embeddings explícitos, um mecanismo que está causalmente ligado à entropia de predição.

O artigo conclui que estas descobertas fornecem uma base para entender o aprendizado em contexto em modelos não-autoregressivos, embora reconheça limitações quanto ao uso de modelos pequenos, apenas de atenção, e tarefas sintéticas, observando que pesquisas adicionais são necessárias para determinar se estes mecanismos escalam para DLMs maiores, contendo MLPs, e tarefas de linguagem natural.