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

• URL: http://arxiv.org/abs/2510.14095v1
• Date: Wed, 15 Oct 2025 21:03:59 GMT
• Title: Unlocking Out-of-Distribution Generalization in Transformers via Recursive Latent Space Reasoning
• Authors: Awni Altabaa, Siyu Chen, John Lafferty, Zhuoran Yang,
• Abstract summary: This work investigates out-of-distribution (OOD) generalization in Transformer networks using a GSM8K-style modular arithmetic on computational graphs task as a testbed.<n>We introduce and explore a set of four architectural mechanisms aimed at enhancing OOD generalization.<n>We complement these empirical results with a detailed mechanistic interpretability analysis that reveals how these mechanisms give rise to robust OOD generalization abilities.
• Score: 50.99796659680724
• License: http://creativecommons.org/licenses/by/4.0/
• 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 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.

Related papers

• Layer Specialization Underlying Compositional Reasoning in Transformers [6.581088182267414]
  Transformers exhibit compositional reasoning on sequences not observed during training.<n>We investigate this phenomenon using the Random Hierarchy Model (RHM), a probabilistic context-free grammar.<n>Mechanistically, we identify a progressive emergence of layer specialization during training that correlates with generalization performance.
  arXiv  Detail & Related papers  (2025-10-20T12:08:22Z)
• Propositional Logic for Probing Generalization in Neural Networks [3.6037930269014633]
  We investigate the generalization behavior of three key neural architectures (Transformers, Graph Convolution Networks and LSTMs) in a controlled task rooted in propositional logic.<n>We find thatTransformers fail to apply negation compositionally, unless structural biases are introduced.<n>Our findings highlight persistent limitations in the ability of standard architectures to learn systematic representations of logical operators.
  arXiv  Detail & Related papers  (2025-06-10T16:46:05Z)
• Partial Transportability for Domain Generalization [56.37032680901525]
  Building on the theory of partial identification and transportability, this paper introduces new results for bounding the value of a functional of the target distribution.<n>Our contribution is to provide the first general estimation technique for transportability problems.<n>We propose a gradient-based optimization scheme for making scalable inferences in practice.
  arXiv  Detail & Related papers  (2025-03-30T22:06:37Z)
• Analyzing the Inner Workings of Transformers in Compositional Generalization [15.599899071518545]
  We investigate the inner workings of a Transformer model by finding an existing subnetwork that contributes to the generalization performance.<n>We find that the model depends on syntactic features to output the correct answer, but that the subnetwork with much better generalization performance than the whole model relies on a non-compositional algorithm.
  arXiv  Detail & Related papers  (2025-02-21T08:07:53Z)
• GRAM: Generalization in Deep RL with a Robust Adaptation Module [62.662894174616895]
  In this work, we present a framework for dynamics generalization in deep reinforcement learning.<n>We introduce a robust adaptation module that provides a mechanism for identifying and reacting to both in-distribution and out-of-distribution environment dynamics.<n>Our algorithm GRAM achieves strong generalization performance across in-distribution and out-of-distribution scenarios upon deployment.
  arXiv  Detail & Related papers  (2024-12-05T16:39:01Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)
• Principled Understanding of Generalization for Generative Transformer Models in Arithmetic Reasoning Tasks [5.522116934552708]
  Transformer-based models excel in various tasks but their generalization capabilities, especially in arithmetic reasoning, remain incompletely understood.<n>This paper develops a unified theoretical framework for understanding the generalization behaviors of transformers in arithmetic tasks.
  arXiv  Detail & Related papers  (2024-07-25T11:35:22Z)
• Understanding the Language Model to Solve the Symbolic Multi-Step Reasoning Problem from the Perspective of Buffer Mechanism [68.05754701230039]
  We construct a symbolic multi-step reasoning task to investigate the information propagation mechanisms in Transformer models.<n>We propose a random matrix-based algorithm to enhance the model's reasoning ability.
  arXiv  Detail & Related papers  (2024-05-24T07:41:26Z)
• Convexifying Transformers: Improving optimization and understanding of transformer networks [56.69983975369641]
  We study the training problem of attention/transformer networks and introduce a novel convex analytic approach. We first introduce a convex alternative to the self-attention mechanism and reformulate the regularized training problem of transformer networks. As a byproduct of our convex analysis, we reveal an implicit regularization mechanism, which promotes sparsity across tokens.
  arXiv  Detail & Related papers  (2022-11-20T18:17:47Z)
• I-BERT: Inductive Generalization of Transformer to Arbitrary Context Lengths [2.604653544948958]
  Self-attention has emerged as a vital component of state-of-the-art sequence-to-sequence models for natural language processing. We propose I-BERT, a bi-directional Transformer that replaces positional encodings with a recurrent layer.
  arXiv  Detail & Related papers  (2020-06-18T00:56:12Z)
• Target-Embedding Autoencoders for Supervised Representation Learning [111.07204912245841]
  This paper analyzes a framework for improving generalization in a purely supervised setting, where the target space is high-dimensional. We motivate and formalize the general framework of target-embedding autoencoders (TEA) for supervised prediction, learning intermediate latent representations jointly optimized to be both predictable from features as well as predictive of targets.
  arXiv  Detail & Related papers  (2020-01-23T02:37:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.