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

Awni Altabaa · Siyu Chen · John Lafferty · Zhuoran Yang

Enhancing Transformer architectures with recursive latent space reasoning mechanisms for robust algorithmic generalization

---

💡 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 four architectural mechanisms aimed at enhancing OOD generalization:

1. 🔄 Input-adaptive recurrence - Recurrent architecture that scales computation through input-adaptive recurrence.
2. 📚 Algorithmic supervision - Structured learning objectives that encode algorithmic knowledge
3. ⚓ Anchored latent representations - Discrete bottlenecks for stable feature learning
4. 🔧 Explicit error-correction mechanism - Built-in self-correction capabilities

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.

---

📂 Organization of Codebase

• experiments/ - Contains the main experimental code for training and evaluating models with the proposed architectural mechanisms
• experiments/baselines/ - Baseline implementations for comparison including chain-of-thought models and standard transformers
• experiments/evaluation/ - Code for evaluating the algorithmic generalization capabilities of different methods.
• experiments/model_interp/ - Mechanistic interpretability analysis tools and visualizations
• Simtransformer/ - Helper framework implementing Transformer modules and related utilities

See experiments/readme.md for instructions on how to reproduce the experiments in the paper.