Unlocking Out-of-Distribution Generalization in Transformers via Recursive Latent Space Reasoning
This addresses a critical bottleneck for reasoning in language models, though it appears incremental as it builds on existing Transformer architectures.
The paper tackled the challenge of out-of-distribution generalization in Transformers by introducing architectural mechanisms like input-adaptive recurrence and algorithmic supervision, achieving robust algorithmic generalization on a modular arithmetic task.
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.