To See the Unseen: on the Generalization Ability of Transformers in Symbolic Reasoning

1. O Problema Investigado

O artigo investiga a capacidade de modelos transformer (especificamente os do tipo decoder-only) de realizar raciocínio simbólico abstrato e generalizar para tokens (símbolos) que não foram observados durante o treinamento.

• Contexto: Modelos de Linguagem (LLMs) modernos têm sucesso em tarefas de raciocínio, mas falham quando os nomes das variáveis ou valores numéricos são alterados (falha de generalização fora da distribuição).
• Desafio Específico: O foco é em problemas de lógica proposicional onde o modelo deve executar algoritmos in-context (como encadeamento direto/forward-chaining) usando símbolos arbitrários e novos.
• Falha Conhecida: Trabalhos anteriores mostraram que os modelos falham ao copiar ou gerar tokens não vistos. A hipótese comum era a dificuldade de "cópia", mas este artigo identifica uma causa raiz adicional e fundamental: o colapso das representações (unembeddings).

2. Metodologia e Análise Teórica

A. Análise Teórica do Colapso

Os autores provam teoricamente que, ao treinar transformadores com Descida de Gradiente Estocástica (SGD) regularizada por $L_2$ (weight decay) e LayerNorm, ocorre um fenômeno inevitável:

• Colapso de Pesos: Os vetores de unembedding (pesos da última camada) de tokens não vistos (que nunca aparecem nos dados de treinamento) tendem a colapsar para o mesmo vetor.
• Mecanismo: O Lemma 4.1 demonstra que, se os rótulos $i$ e $j$ não estão presentes nos dados, a distância entre seus pesos $W_i$ e $W_j$ diminui a cada passo de atualização, desde que as entradas tenham norma limitada (garantido pelo LayerNorm) e o modelo se ajuste bem aos dados.
• Consequência: Isso torna os tokens não vistos indistinguíveis para o modelo, impedindo-o de diferenciar variáveis diferentes em problemas de raciocínio.

B. Arquitetura e Técnicas Propostas

Para mitigar esse problema, os autores propõem uma combinação de três técnicas:

1. Arquitetura com Atenção de Cópia (Copy Attention): Uma modificação na arquitetura que adiciona um "atalho" direto das embeddings de entrada para a saída. Isso facilita a cópia explícita de tokens do contexto, superando a dificuldade de propagar embeddings através de várias camadas de MLPs.
2. Diversidade de Dados: Aumento da diversidade simbólica nos dados de treinamento (variando o tamanho do vocabulário de predicados) para forçar o modelo a aprender a estrutura lógica em vez de memorizar tokens específicos.
3. Gerenciamento de Embeddings:
   • Congelamento (Freezing): Manter as embeddings/unembeddings fixas em sua inicialização (que são ortogonais e distinguíveis).
   • Esquecimento Ativo (Active Forgetting): Reinicializar periodicamente as embeddings durante o treinamento para evitar que elas colapsem.

C. Experimentos

• Dados Sintéticos: Problemas de lógica proposicional (cláusulas definidas/Horn clauses) gerados sinteticamente com cadeias de pensamento (CoT) geradas via algoritmo de encadeamento direto.
• Configuração: Modelos pequenos (10 camadas, ~60-75M parâmetros) treinados do zero.
• Variáveis: Tamanho do vocabulário de treinamento, presença de tokens não vistos (apenas na consulta ou em todos os predicados), e uso de copy attention.
• Modelos Reais: Análise dos modelos abertos Gemma 3 (família de modelos com tokens deliberadamente não utilizados).

3. Resultados Principais

Experimentos Sintéticos

• Falha Padrão: Transformadores padrão falham completamente quando tokens não vistos são introduzidos, mesmo com alta diversidade de dados. O modelo aprende a não gerar esses tokens.
• Eficácia da Solução Combinada:
  • A Atenção de Cópia sozinha melhora a generalização para um único token não visto, mas falha quando todos os símbolos são novos (devido ao colapso das embeddings).
  • Congelar as embeddings ou usar Esquecimento Ativo (reinicialização periódica) permite que o modelo generalize para tokens não vistos, pois mantém a distinção entre os vetores.
  • A combinação de Arquitetura de Cópia + Diversidade + Congelamento/Reinicialização alcança generalização robusta.
• Dinâmica Temporal: A generalização in-context (ICL) é transitória em arquiteturas padrão; sem intervenções, o modelo perde a capacidade de generalizar conforme o treinamento avança.

Experimentos com Gemma 3

• Colapso em Modelos Grandes: Os autores analisaram os 99 tokens deliberadamente não utilizados no vocabulário do Gemma 3.
• Correlação Alta: As embeddings desses tokens não utilizados apresentam alta similaridade de cosseno entre si (ex: 0.78 no modelo de 1B), enquanto tokens usados têm baixa similaridade. Isso confirma o colapso em modelos pré-treinados reais.
• Impacto no Fine-tuning: Ao fazer fine-tuning do Gemma 3 1B em problemas de lógica:
  • Símbolos comuns (letras, adjetivos) atingem ~90% de precisão em ~500 passos.
  • Símbolos não utilizados (com embeddings colapsadas) levam ~5000 passos para atingir o mesmo desempenho, indicando que o colapso inicial prejudica severamente a inicialização para tarefas downstream.

4. Contribuições Chave

1. Identificação do Colapso de Unembedding: Demonstração teórica e empírica de que o colapso das representações de tokens não vistos é um mecanismo fundamental que impede a generalização simbólica, explicando por que a simples mudança de arquitetura (como sugerido por trabalhos anteriores) não é suficiente.
2. Explicação Mecanística: Oferece uma explicação mecânica para a eficácia de heurísticas como "esquecimento ativo" (reinicialização de embeddings), mostrando que elas previnem o colapso.
3. Solução Prática: Propõe uma combinação de técnicas (cópia, diversidade e gestão de embeddings) que resolve o problema de generalização para tokens não vistos em raciocínio simbólico.
4. Validação em Modelos Reais: Estende as descobertas para modelos de grande escala (Gemma 3), mostrando que o colapso afeta tokens reservados para uso futuro e impacta negativamente a eficiência do fine-tuning.

5. Significado e Implicações

Este trabalho sugere que a fragilidade dos LLMs em raciocínio simbólico não é apenas uma questão de falta de dados ou capacidade de cópia, mas sim uma falha estrutural na dinâmica de otimização que homogeniza representações de tokens raros ou inexistentes.

• Para Pesquisa: Destaca a necessidade de monitorar a diversidade de embeddings durante o treinamento e sugere que a "diversidade de dados" sozinha não é suficiente sem mecanismos para preservar a distinção entre tokens.
• Para Aplicações: Alerta que tokens reservados em vocabulários de modelos grandes (como o Gemma) podem não ser uma boa inicialização para novas tarefas devido ao colapso prévio, exigindo estratégias de reinicialização ou fine-tuning mais longo.
• Para Arquitetura: Reforça a importância de mecanismos explícitos de cópia (como copy attention) combinados com estratégias de regularização de embeddings para tarefas que exigem manipulação de símbolos abstratos.