 |
ML Newsnew | past | best | rss | submit | about |
| Out-of-Distribution Generalization in Transformers via Latent Space Reasoning (arxiv.org) | |
| 9 points by marojejian 320 days ago | hide | past | pdf | 1 comment on HN | |
|
In plain words: A Transformer that repeatedly revises its own internal steps, trained to supervise those steps and check for mistakes, learns to reason in its internal states rather than pattern-match. On arithmetic over unfamiliar structures it stayed accurate where one-pass training broke down. Abstract · Unlocking Out-of-Distribution Generalization in Transformers via Recursive Latent Space ReasoningSystematic, 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. | |
| add comment on HN |
I found this paper both really interesting and clear. No one part is very novel, but It composes disparate threads to obtain what looks like strong results in OOD length generalization. Even for the toy task, and using a DSL (vs. being an LM), length-generalizing on simple math >4x is impressive, from what I've read.
This also fits my priors for the key elements of unlocking better OOD compositional generalization: variable recurrence, step-wise curriculum training to build depth-invariant algorithms, discrete bottlenecks. Finally, it's very interesting to compare this to the below recent article arguing for the benefits of continuous latent spaces: Reasoning by Superposition: A Theoretical Perspective on Chain of Continuous Thought (https://arxiv.org/abs/2505.12514)
My take is both papers are right, and that continuous spaces are more expressive and can handle tougher problem spaces (e.g. shortest graph path), whereas discrete spaces will provide a better inductive bias for elegant algorithms that can scale OOD. And I bet the two can be combined / balanced.