Desbloqueando a Generalização Fora da Distribuição em Transformers via Raciocínio Recursivo no Espaço Latente

Unlocking Out-of-Distribution Generalization in Transformers via Recursive Latent Space Reasoning

October 15, 2025

Autores: Awni Altabaa, Siyu Chen, John Lafferty, Zhuoran Yang

cs.AI

Resumo

A generalização sistemática e composicional além da distribuição de treinamento permanece um desafio central no aprendizado de máquina — e um gargalo crítico para as habilidades emergentes de raciocínio dos modelos de linguagem modernos. Este trabalho investiga a generalização fora da distribuição (OOD, do inglês "out-of-distribution") em redes Transformer utilizando uma tarefa de aritmética modular em grafos computacionais no estilo GSM8K como um ambiente de teste. Introduzimos e exploramos um conjunto de quatro mecanismos arquitetônicos voltados para aprimorar a generalização OOD: (i) recorrência adaptativa à entrada; (ii) supervisão algorítmica; (iii) representações latentes ancoradas por meio de um gargalo discreto; e (iv) um mecanismo explícito de correção de erros. Coletivamente, esses mecanismos resultam em uma abordagem arquitetônica para o raciocínio nativo e escalável em espaços latentes em redes Transformer, com capacidades robustas de generalização algorítmica. Complementamos esses resultados empíricos com uma análise detalhada de interpretabilidade mecanicista que revela como esses mecanismos dão origem a habilidades robustas de generalização OOD.

English

Systematic, compositional generalization beyond the training distribution remains a core challenge in machine learning -- and a critical bottleneck for the emergent reasoning abilities of modern language models. This work investigates out-of-distribution (OOD) generalization in Transformer networks using a GSM8K-style modular arithmetic on computational graphs task as a testbed. We introduce and explore a set of four architectural mechanisms aimed at enhancing OOD generalization: (i) input-adaptive recurrence; (ii) algorithmic supervision; (iii) anchored latent representations via a discrete bottleneck; and (iv) an explicit error-correction mechanism. Collectively, these mechanisms yield an architectural approach for native and scalable latent space reasoning in Transformer networks with robust algorithmic generalization capabilities. We complement these empirical results with a detailed mechanistic interpretability analysis that reveals how these mechanisms give rise to robust OOD generalization abilities.