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

arXiv: https://arxiv.org/abs/2609.32927 · HTML: https://mlsanity.com/paper/2609.32927/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Mathematical reasoning, one of 17 directions in LLM reasoning.
- Builds on 32 of its 43 references that are in the map.
- First posted on arXiv 26 Sep 2026.

## Placed in

- LLM reasoning / [Mathematical reasoning](https://mlsanity.com/lr/mathematical-reasoning/)

## Builds on (most cited first, 30 of 32)

- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [Show Your Work: Scratchpads for Intermediate Computation with Language Models](https://mlsanity.com/paper/2112.00114/) (2021)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [The Linear Representation Hypothesis and the Geometry of Large Language Models](https://mlsanity.com/paper/2311.03658/) (2023)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [GSM-Symbolic: Understanding the Limitations of Mathematical Reasoning in Large Language Models](https://mlsanity.com/paper/2410.05229/) (2024)
- [Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task](https://mlsanity.com/paper/2210.13382/) (2022)
- [Analysing Mathematical Reasoning Abilities of Neural Models](https://mlsanity.com/paper/1904.01557/) (2019)
- [SATNet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver](https://mlsanity.com/paper/1905.12149/) (2019)
- [PMLB: A Large Benchmark Suite for Machine Learning Evaluation and Comparison](https://mlsanity.com/paper/1703.00512/) (2017)
- [Large Language Models for Automated Data Science: Introducing CAAFE for Context-Aware Automated Feature Engineering](https://mlsanity.com/paper/2305.03403/) (2023)
- [Teaching Arithmetic to Small Transformers](https://mlsanity.com/paper/2307.03381/) (2023)
- [Randomized Positional Encodings Boost Length Generalization of Transformers](https://mlsanity.com/paper/2305.16843/) (2023)
- [Arithmetic Without Algorithms: Language Models Solve Math With a Bag of Heuristics](https://mlsanity.com/paper/2410.21272/) (2024)
- [Transformers Can Do Arithmetic with the Right Embeddings](https://mlsanity.com/paper/2405.17399/) (2024)
- [Copy Suppression: Comprehensively Understanding an Attention Head](https://mlsanity.com/paper/2310.04625/) (2023)
- [Neural Arithmetic Logic Units](https://mlsanity.com/paper/1808.00508/) (2018)
- [OpenFE: Automated Feature Generation with Expert-level Performance](https://mlsanity.com/paper/2211.12507/) (2022)
- [SATBench: Benchmarking LLMs' Logical Reasoning via Automated Puzzle Generation from SAT Formulas](https://mlsanity.com/paper/2505.14615/) (2025)
- [Neural Arithmetic Units](https://mlsanity.com/paper/2001.05016/) (2020)
- [Fine-tuned In-Context Learning Transformers are Excellent Tabular Data Classifiers](https://mlsanity.com/paper/2405.13396/) (2024)
- [Learning to Solve Constraint Satisfaction Problems with Recurrent Transformer](https://mlsanity.com/paper/2307.04895/) (2023)
- [Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count](https://mlsanity.com/paper/2410.15787/) (2024)
- [It's Hard for Neural Networks To Learn the Game of Life](https://mlsanity.com/paper/2009.01398/) (2020)
- [LifeGPT: Topology-Agnostic Generative Pretrained Transformer Model for Cellular Automata](https://mlsanity.com/paper/2409.12182/) (2024)
- [FAMOSE: A ReAct Approach to Automated Feature Discovery](https://mlsanity.com/paper/2602.17641/) (2026)
- [Mitigating Label Shift in Tabular In-Context Learning via Test-Time Posterior Adjustment](https://mlsanity.com/paper/2605.04363/) (2026)

## Abstract (arXiv)

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 STRAT (STratified Registers And Types), 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.