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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Altabaa, Awni, Chen, Siyu, Lafferty, John, Yang, Zhuoran
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915557067980800
author Altabaa, Awni
Chen, Siyu
Lafferty, John
Yang, Zhuoran
author_facet Altabaa, Awni
Chen, Siyu
Lafferty, John
Yang, Zhuoran
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2510_14095
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Unlocking Out-of-Distribution Generalization in Transformers via Recursive Latent Space Reasoning
Altabaa, Awni
Chen, Siyu
Lafferty, John
Yang, Zhuoran
Machine Learning
Artificial Intelligence
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.
title Unlocking Out-of-Distribution Generalization in Transformers via Recursive Latent Space Reasoning
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2510.14095