A Measure-Theoretic Analysis of Reasoning: Structural Generalization and Approximation Limits
For the ML theory community, this provides a rigorous theoretical framework for OOD generalization in reasoning, revealing fundamental architectural constraints that empirical scaling laws alone cannot explain.
This paper formalizes out-of-distribution generalization in LLM reasoning using optimal transport, proving that shift-invariant attention mechanisms (e.g., Rotary Embeddings) bound generalization error while position-dependent ones (e.g., Absolute Positional Encoding) incur Ω(1) risk, and that scaling depth is necessary to avoid representation collapse for Dyck-k reasoning. Experiments on 54 Transformer configurations confirm monotonic degradation of generalization risk with Wasserstein domain shift.
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.