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

ICLR 2026 Conference Submission20527 Authors

Keywords: algorithmic generalization, out-of-distribution, OOD generalization, compositional, systematic generalization, transformer, transformer architecture

TL;DR: This paper explores a set of architectural mechanisms to support algorithmic generalization in Transformers through a recurrent latent space reasoning approach.

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 algorithmic 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.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Submission Number: 20527