A Measure-Theoretic Analysis of Reasoning: Structural Generalization and Approximation Limits

2026-05-19

AI Summary

This work addresses the limited out-of-distribution (OOD) generalization of large language models in structured tasks and the lack of theoretical understanding behind it. By leveraging optimal transport theory, the authors embed discrete reasoning trajectories into a continuous metric space and quantify domain shift using the Wasserstein-1 distance. Combining Kantorovich duality, they analyze OOD generalization bounds through the lenses of Lipschitz continuity and approximation limits in Barron space. They establish, for the first time, that positional encoding affects translation invariance and prove that rotary embeddings outperform absolute encodings. Furthermore, they link Dyck-k languages to backtracking reasoning and derive a circuit-depth lower bound for TC⁰ Transformers, demonstrating the irreplaceability of physical layers. Experiments across 54 Transformer configurations confirm that generalization risk monotonically increases with Wasserstein shift and underscore the critical role of depth scaling in mitigating representational collapse.

Abstract

While empirical scaling laws for LLM reasoning are well-documented, the theoretical mechanisms governing out-of-distribution (OOD) generalization remain elusive. We formalize reasoning via optimal transport, projecting discrete trajectories into a continuous metric space to quantify domain shifts using the Wasserstein-1 distance. Invoking Kantorovich duality, we bound OOD generalization via architectural Lipschitz continuity and functional approximation limits. This exposes two primary constraints. First, position-dependent attention (e.g., Absolute Positional Encoding) fails to preserve shift invariance, yielding an $Ω(1)$ Lipschitz constant and expected risk, whereas shift-invariant mechanisms (e.g., Rotary Embeddings) preserve equivariance and bound the error. Second, by mapping sequential backtracking to a Dyck-$k$ language, we establish a strict circuit depth lower bound for $\text{TC}^0$ Transformers. Scaling physical layer depth is necessary to avert representation collapse -- a constraint that scaling representation width cannot bypass due to irreducible approximation bounds in Barron spaces. Evaluations across 54 Transformer configurations on combinatorial search corroborate these bounds, demonstrating that generalization risk degrades monotonically with the Wasserstein domain shift.

Problem

• out-of-distribution generalization
• structural generalization
• approximation limits
• reasoning
• domain shift

Innovation

• optimal transport
• Wasserstein distance
• Lipschitz continuity
• shift invariance
• circuit depth lower bound