Débloquer la généralisation hors distribution dans les Transformers via un raisonnement récursif dans l'espace latent

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

October 15, 2025

papers.authors: Awni Altabaa, Siyu Chen, John Lafferty, Zhuoran Yang

cs.AI

papers.abstract

La généralisation systématique et compositionnelle au-delà de la distribution d'entraînement reste un défi central en apprentissage automatique — et un goulot d'étranglement critique pour les capacités de raisonnement émergentes des modèles de langage modernes. Ce travail étudie la généralisation hors distribution (OOD) dans les réseaux Transformer en utilisant une tâche d'arithmétique modulaire sur des graphes computationnels de style GSM8K comme banc d'essai. Nous introduisons et explorons un ensemble de quatre mécanismes architecturaux visant à améliorer la généralisation OOD : (i) récurrence adaptative aux entrées ; (ii) supervision algorithmique ; (iii) représentations latentes ancrées via un goulot d'étranglement discret ; et (iv) un mécanisme explicite de correction d'erreurs. Collectivement, ces mécanismes constituent une approche architecturale pour un raisonnement natif et évolutif dans l'espace latent des réseaux Transformer, dotée de capacités robustes de généralisation algorithmique. Nous complétons ces résultats empiriques par une analyse détaillée d'interprétation mécaniste qui révèle comment ces mécanismes engendrent des capacités robustes de généralisation 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.