The Geometry of Logic: Stratification Induces Semantic Structure and Robust Reasoning

2026-09-26

AI Summary

This study addresses the limitation of Transformers in symbolic reasoning, where reliance on statistical shortcuts leads to poor generalization. We propose STRAT, an architecture that introduces the separation of values and types as a fundamental architectural primitive for the first time. By orthogonally decomposing the residual stream into data and type subspaces, and regulating data transformations via type attention and gating mechanisms, STRAT effectively circumvents failure modes such as linear traps and gradient walls, thereby achieving decoupled logical structures and interpretable, robust reasoning. Experimental results demonstrate that STRAT reduces out-of-distribution error by 35× on arithmetic tasks and improves average accuracy by 26 percentage points across eleven datasets. Furthermore, under distribution shifts, its performance degrades by only 2.39%, confirming the efficacy of the proposed approach.

Abstract

Transformer-based language models perform well on symbolic tasks, yet it remains unclear whether they learn generalizable rules or rely on statistical shortcuts. Mechanistic studies link algorithmic behavior to structured internal representations, motivating the hypothesis that robust reasoning benefits from separating values from the types that control their manipulation. Can making this separation an architectural primitive improve the learnability and generalization of logical mechanisms? We introduce \textbf{STRAT} (\textbf{ST}ratified \textbf{R}egisters \textbf{A}nd \textbf{T}ypes), which partitions the residual stream into orthogonal Data and Type subspaces and uses Type-based attention and gating to govern Data transformations. Controlled arithmetic ablations identify three failure modes associated with data-control interference: the Linear Trap, Gradient Wall, and Open Gate Trap. Mechanistic analysis reveals interpretable logical structure, and in arithmetic, STRAT reduces median OOD error 35-fold relative to a Transformer baseline. On each of 11 datasets spanning 10 tasks, STRAT outperforms the Transformer baseline in mean accuracy, by 26 percentage points on average, with both models trained from 10 base examples per dataset using identical task-specific augmentation where applicable. Under distribution shift, STRAT's mean accuracy drops by only 2.39 percentage points, compared with 11.75 for the Transformer.

Problem

• logical reasoning
• generalization
• distribution shift
• Transformer
• mechanistic interpretability

Innovation

• Stratified Architecture
• Orthogonal Subspaces
• Type-based Attention
• Out-of-Distribution Generalization
• Mechanistic Interpretability