Desbloqueando la Generalización Fuera de Distribución en Transformers mediante Razonamiento Recursivo en el Espacio 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

Resumen

La generalización sistemática y composicional más allá de la distribución de entrenamiento sigue siendo un desafío fundamental en el aprendizaje automático, y un cuello de botella crítico para las habilidades de razonamiento emergentes de los modelos de lenguaje modernos. Este trabajo investiga la generalización fuera de distribución (OOD, por sus siglas en inglés) en redes Transformer utilizando como banco de pruebas una tarea de aritmética modular en grafos computacionales al estilo de GSM8K. Introducimos y exploramos un conjunto de cuatro mecanismos arquitectónicos diseñados para mejorar la generalización OOD: (i) recurrencia adaptativa a la entrada; (ii) supervisión algorítmica; (iii) representaciones latentes ancladas mediante un cuello de botella discreto; y (iv) un mecanismo explícito de corrección de errores. En conjunto, estos mecanismos dan lugar a un enfoque arquitectónico para el razonamiento nativo y escalable en el espacio latente de las redes Transformer, con capacidades robustas de generalización algorítmica. Complementamos estos resultados empíricos con un análisis detallado de interpretabilidad mecanicista que revela cómo estos mecanismos dan lugar a habilidades robustas de generalización 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.