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

Cristina V. Lopes; Yuangang Li; Justin Tian Jin Chen; Alberto Krone-Martins; Iris Ma; Md Rakib Hossain Misu

cs.LGSep 26, 2026

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

Authors: Cristina V. Lopes, Yuangang Li, Justin Tian Jin Chen, Alberto Krone-Martins, Iris Ma, Md Rakib Hossain Misu

Organizations: University of California, Irvine · University of California, Davis

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 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.

Figures & tables

Figure 1: Mean accuracy across 11 datasets spanning 10 tasks, with both STRAT and the Transformer baseline trained from only 10 base examples per dataset: (a) ID and OOD accuracy for each dataset; (b) macro-averaged accuracy under distribution shift.

Index	Content	Functional Role
D0	Value	Numerical magnitude (Payload)
D1–D4	Flags	Type Indicators (IsNum, IsOp, etc.)
D5	Accu	Conditional Storage

Table 1: Dimension Assignments in the Residual Stream

			ID MAE (Train)			OOD MAE (Test)
Cfg	Strategy	Hypothesis	25%	Med	75%	Med	75%	Diagnosis / Mechanism
Ref	Hand-Coded	Existence Proof	0.0	0.0	0.0	0	0	Geometric Logic
A	Simplified Xformer	Standard Training	3.5	15.5	38.2	5,977	10,809	Linear Trap
A s	A Scaled	Scale Invariance	97.1	98.7	100.1	6,202	6,265	Linear Trap
A ln	LayerNorm	Internal Stability	9.6	16.9	25.9	8,378	9,856	Feature Collapse
A curr	Curriculum	Complexity Sched.	212.2	600.3	1736.1	2306	2,682	Linear Trap

Table 2: The Failure Taxonomy (N=50 runs/config). We report 25%-ile / Median / 75%-ile errors.

Figure 2: Overview of the STRAT-Lite architecture with an arithmetic example (

12+3-5

Table 5

Figure 3: Antipodal discovery. Left: arithmetic. Blue circle: plus; red circle: minus; grey circle: numbers. Right: selective sum. Green cuircles: target category; purple squares: distractor category A; blue triangles: distractor category B.

	OOD MAE Distribution (Lower is Better)			Success Rates (Higher is Better)
Model	10%-ile	25%-ile	Median	Strict ( $<$ 20)	Logic ( $<$ 100)
Baseline (Config A)	$5,072$	$6,268$	$9,046$	$0\%$	$0\%$
STRAT Flat (2D/1H)	$33$	$80$	$293$	$5\%$	$31\%$
STRAT Flat MH (2D/2H)	$38$	$97$	$314$	$2\%$	$27\%$
STRAT Compact (3D/1H)	$34$	$77$	$266$	$3\%$	$31\%$
STRAT Std (4D/1H)	$\mathbf{33}$	$\mathbf{77}$	$\mathbf{258}$	$4\%$	$\mathbf{31\%}$

Table 5: Topological Locking Results (

N=10

, 400 Seeds). Strict Acc (

<20

) denotes perfect circuit execution. Logic Acc (

<100

) denotes functional gating with Sigmoid tax. The Baseline fails (

0\%

acc). STRAT achieves

\approx 31\%

logic locking, with the top 10% achieving MAE

\approx 33

Task	Dataset	STRAT (Ours)		Transformer
Task	Dataset	ID Acc. (%)	OOD Acc. (%)	ID Acc. (%)	OOD Acc. (%)	$\Delta$ ID Acc.	$\Delta$ OOD Acc.
Constraint satisfaction Sudoku solving	SATNet Sudoku ( Wang et al., 2019 )	$\mathbf{86.58\pm 1.08}$	$\mathbf{83.80\pm 1.30}$ ( -2.78 )	$27.92\pm 4.97$	$10.18\pm 2.36$ ( -17.74 )	$\mathbf{+58.66\pm 5.22}$	$\mathbf{+73.62\pm 3.24}$
Formula satisfaction classification Given-assignment satisfaction	SATBench ( Wei et al., 2025 )	$\mathbf{92.33\pm 1.36}$	$\mathbf{90.50\pm 2.15}$ ( -1.83 )	$65.64\pm 5.38$	$64.48\pm 5.40$ ( -1.16 )	$\mathbf{+26.69\pm 5.58}$	$\mathbf{+26.02\pm 5.54}$
Formula satisfaction classification Given-assignment satisfaction	SATLIB ( Hoos and Stützle, 2000 )	$\mathbf{99.98\pm 0.02}$	$\mathbf{99.96\pm 0.04}$ ( -0.02 )	$62.67\pm 5.58$	$56.18\pm 4.91$ ( -6.49 )	$\mathbf{+37.31\pm 5.58}$	$\mathbf{+43.78\pm 4.91}$
Board classification Winner detection	Tic-Tac-Toe Endgame ( Aha, 1991 )	$\mathbf{98.19\pm 0.00}$	$\mathbf{94.64\pm 0.00}$ ( -3.54 )	$80.91\pm 2.53$	$60.49\pm 5.47$ ( -20.42 )	$\mathbf{+17.28\pm 2.53}$	$\mathbf{+34.15\pm 5.47}$
State-transition prediction Next-state prediction	Game of Life ( Strandgaard, 2023 )	$\mathbf{88.71\pm 2.72}$	$\mathbf{79.26\pm 12.99}$ ( -9.45 )	$80.54\pm 2.64$	$69.08\pm 12.89$ ( -11.46 )	$\mathbf{+8.16\pm 3.61}$	$\mathbf{+10.18\pm 17.94}$

Table 6: We report accuracy (%) across 11 dataset evaluations as mean

\pm

standard error of the mean (SEM) over 10 seeds. The difference columns use STRAT minus Transformer: ID accuracy (

\Delta

ID Acc.),and OOD accuracy (

\Delta

OOD Acc.). The OOD Acc. columns include the OOD performance and the difference relative to the ID accuracy in parentheses ( green for increase, red for decrease).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Index	Identifier	Functional Role
D0	Value	Numerical magnitude $n$ (Payload)
D1	Is_Num	Binary flag ( $1.0$ ) for numerical literals
D2	Is_Add	Binary flag ( $1.0$ ) for the ‘+’ operator
D3	Is_Sub	Binary flag ( $1.0$ ) for the ‘-’ operator
D4	Is_Op	Generic flag ( $1.0$ ) for any operator
D5	Accu	Reserved register for conditional storage

Appendix

Table 7: Dimension Assignments in the Residual Stream.

Token	D0	D1	D2	D3	D4	D5
Literal ( $n$ )	$n$	$1$	$0$	$0$	$0$	$0$
Plus ( $+$ )	$0$	$0$	$1$	$0$	$1$	$0$
Minus ( $-$ )	$0$	$0$	$0$	$1$	$1$	$0$

Appendix

Table 8: Sparse Token Embeddings

E(t)

Config	Name	Mechanism / Modification	Hypothesis Tested
A	Baseline	Standard Xavier initialization; Raw inputs.	Can standard GD solve it?
A $s$	Scaled	Inputs/Targets scaled (max-normalized).	Is magnitude the problem?
A ln	LayerNorm	Internal LayerNorm	Feature Collapse (OOD Fail)
A curr	Curriculum	Range $[0,100]\to[0,5000]$	Linear Trap (Locked Early)
B fl	Frozen Logic	Logic weights fixed to ideal; others learned.	Can gradients flow through logic?
B fa	Frozen Attn	Attention weights fixed to Oracle (Perfect).	Is the failure in Routing or Logic?

Appendix

Table 9: Detailed hyperparameters for the 12 Ablation Configurations. All models share the same 6-dim architecture.

Feature	Standard (Run 1)	Mirror (Run 2)
Readout Sign	Positive ( $+$ )	Negative ( $-$ )
Default Operation	Addition ( $+x$ )	Subtraction ( $-x$ )
Active Trigger	MINUS Token	PLUS Token
Attention Role	Pass-through	Signal Inversion
Resulting Logic	$y=\text{Default}-\text{Routed}$	$y=\text{Default}+\text{Routed}$

Appendix

Table 10: Comparison of the two learned algorithms.

Dataset	STRAT parameters	Transformer parameters
SATNet Sudoku	396,335	805,349
SATBench	1,648	5,310
SATLIB	1,648	5,310
Tic-Tac-Toe Endgame	33,474	58,498
Game of Life	17,993	30,761
PMLB Parity5	35,314	60,594

Appendix

Table 11: Trainable parameter counts for STRAT and the Transformer baseline across the 11 dataset evaluations.

	OOD MAE $\downarrow$			Runs below threshold $\uparrow$
Model	P10	P25	Median	MAE $<20$	MAE $<100$
Baseline (Config A)	5,072	6,268	9,046	0%	0%
LSTM-controlled NAC	4,120	5,083	5,819	0%	0.25%
LSTM-controlled NALU	4,973	5,327	5,770	0%	0%
STRAT Flat (2D/1H)	33	80	293	5%	31%
STRAT Flat MH (2D/2H)	38	97	314	2%	27%

Appendix

Table 12: Comparison with specialized arithmetic baselines. The Transformer and four STRAT-Lite rows reproduce Table 5 ; NAC and NALU are evaluated over 400 runs each. P10/P25 denote the 10th/25th percentiles of per-run OOD MAE; the last two columns give the percentage of runs below each error threshold.

LLM Reasoning

Large language model reasoning studies how models solve multi-step problems, verify intermediate steps, and use computation or tools reliably.

See all 931 papers

Transformers Provably Learn to Internalize Chain-of-Thought
May 27, 2026Yixiao Huang, Hanlin Zhu, Zixuan Wang +4Transformer CoT Reasoning
UC Berkeley · Princeton University
Feature Resemblance: Towards a Theoretical Understanding of Analogical Reasoning in Transformers
Mar 5, 2026Ruichen Xu, Wenjing Yan, Ying-Jun Angela ZhangAnalogical Reasoning Analogical Reasoning for Language Models
Department of Information Engineering, The Chinese University of Hong Kong, Hong Kong.
Quantifying Logical Consistency in Transformers via Query-Key Alignment
Sep 10, 2026Eduard Tulchinskii, Anastasia Voznyuk, Laida Kushnareva +4Transformer Interpretability LLM Evaluation
Skolkovo Institute of Science and Technology · AI Foundation and Algorithm Lab · Moscow Institute of Physics and Technology +2

May 27, 2026J/K move · Enter open · S save

Transformers Provably Learn to Internalize Chain-of-Thought

Yixiao Huang, Hanlin Zhu, Zixuan Wang, Jiantao Jiao+3

UC Berkeley · Princeton University

Transformer CoT Reasoning Latent CoT Reasoning

Open paper AskPDF arXiv

Chain-of-Thought (CoT) prompting substantially improves the sample efficiency of transformers, reducing the complexity of tasks like parity learning from exponential to polynomial in the input length. However, generating explicit reasoning steps at inference is computationally expensive. Implicit Chain-of-Thought (ICoT) has emerged as a promising empirical remedy that trains models to internalize intermediate steps within their hidden states, but its theoretical foundations remain poorly understood. We give the first theoretical analysis of ICoT, proving that an

L

-layer transformer trained under our proposed Log-ICoT curriculum learns

k

-parity with

\mathsf{poly}(n)

samples and

L = \log_2 k

training stages. This matches the sample efficiency of explicit CoT while eliminating its inference overhead, and extends prior one-layer parity guarantees to multi-layer architectures. Compared to standard ICoT, which removes thinking tokens one at a time, Log-ICoT removes them in geometric chunks, reducing the number of stages from linear in

k

to logarithmic. Experiments on multi-layer transformers confirm the theory and visualize how reasoning is progressively absorbed into deeper layers.

Abstract

LLM Reasoning

Explore similar work

Transformers Provably Learn to Internalize Chain-of-Thought

Feature Resemblance: Towards a Theoretical Understanding of Analogical Reasoning in Transformers

Quantifying Logical Consistency in Transformers via Query-Key Alignment