Het ontgrendelen van Out-of-Distribution Generalisatie in Transformers via Recursief Redeneren in de Latente Ruimte

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

October 15, 2025

Auteurs: Awni Altabaa, Siyu Chen, John Lafferty, Zhuoran Yang

cs.AI

Samenvatting

Systematische, compositionele generalisatie buiten de trainingsdistributie blijft een kernuitdaging in machine learning – en een kritieke bottleneck voor de opkomende redeneervaardigheden van moderne taalmodelen. Dit werk onderzoekt out-of-distribution (OOD) generalisatie in Transformer-netwerken met behulp van een GSM8K-achtige modulaire rekenkunde op computationele grafieken als testomgeving. We introduceren en onderzoeken een set van vier architecturale mechanismen gericht op het verbeteren van OOD-generalisatie: (i) input-adaptieve recurrentie; (ii) algoritmische supervisie; (iii) verankerde latente representaties via een discreet bottleneck; en (iv) een expliciet foutcorrectiemechanisme. Collectief leveren deze mechanismen een architecturale benadering op voor native en schaalbare latente ruimte-redenering in Transformer-netwerken met robuuste algoritmische generalisatie mogelijkheden. We vullen deze empirische resultaten aan met een gedetailleerde mechanistische interpretatieanalyse die onthult hoe deze mechanismen leiden tot robuuste OOD-generalisatievaardigheden.

English

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.