Unlocking Out-of-Distribution Generalization in Transformers via
  Recursive Latent Space Reasoning
Authors: name: Awni Altabaa; name: Siyu Chen; name: John Lafferty; name: Zhuoran Yang
Published: 2025-10-15 21:03:59

Abstract: 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.