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

STRATStratified Registers And Types

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

A Transformer architecture that separates data from control, inducing semantic structure and improving generalization on logical tasks with limited training data.

Cristina V. Lopes^1,* Yuangang Li^1,*,‡ Justin Tian Jin Chen² Alberto Krone-Martins¹ Iris Ma¹ Md Rakib Hossain Misu¹

¹ University of California, Irvine² University of California, Davis

*Denotes equal contribution.

Read the paper PDF Code Cite this work

THE CENTRAL IDEA

Give values and control their own space.

STRAT partitions the residual stream into orthogonal Data and Type subspaces.

TRANSFORMER BASELINE

A shared representation

Data + control

AttentionFeed-forward

Routing and numerical values occupy the same representation.

STRAT

Two strata. Distinct roles.

TypeAttention+ residualGateg

DataLinear transformv′Gated updateg ⊙ v′

Data_out=Data_in+g ⊙ v′

Updated Type controls the gate. The gated transform is added back to Data.

Simplified STRAT block from Paper §5.2. Attention updates Type only; the gate scales the Data transform elementwise (⊙). Both updates are residual. Follow the computation

01 / Overview

Structure shapes
what a model learns.

Can separating values from the types that control them make logical mechanisms easier to learn?

STRAT makes this separation an architectural primitive. It partitions the residual stream into orthogonal Data and Type subspaces. Type-based attention and gating govern Data transformations, keeping numerical values separate from routing decisions.

We study small Transformer-based models on arithmetic and logical tasks. The resulting architecture learns interpretable computational structure and generalizes from limited training data.

Read the 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 out-of-distribution (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.

02 / Architecture

Data carries values.
Types govern their transformation.

Type / controlData / values

Simplified STRAT architecture for 12 plus 3 minus 5: separate Data and Type dimensions, Type-only attention, a Type-conditioned GLU gate, Data residual updates, and global sum readout. — **STRAT.** The arithmetic example 12 + 3 − 5 illustrates the simplified block from Paper §5.2, from stratified embeddings to the final readout.Full figure

Separate the representation

Values and types occupy orthogonal subspaces. Numerical inputs enter the Data stratum; type embeddings carry control information.

Attend through types

Attention reads and updates the Type stratum. Data values remain unchanged during this phase.

Gate the data update

A Type-conditioned gate controls a learned Data transformation. The result is added back through a residual connection.

Inside the representation

Semantic structure
can emerge.

When Type embeddings are learned, the model discovers structure tied to functional roles.

In arithmetic, addition and subtraction embeddings become approximately antipodal. In selective summation, two distractor categories cluster together, apart from the target category.

PCA of learned Type embeddings · Paper §5.4

Two PCA plots: plus and minus embeddings occupy opposing directions; the two distractor categories cluster together away from target items. — Learned Type geometry in arithmetic and selective summation.

03 / Results

Generalization across
10 logical tasks.

Eleven datasets. Ten paired runs per dataset.
The same base examples and task-specific augmentations for both models.

STRAT has higher mean accuracy than the Transformer on all 11 datasets, both in distribution and under distribution shift. Macro-average accuracy: STRAT 93.5 to 91.1 percent; Transformer 67.5 to 55.8 percent. — Mean accuracy across 11 datasets, with 10 base training examples per dataset. Error bars show the standard error over 10 paired runs.Full figure

Under rule-preserving distribution shifts, STRAT's macro-average accuracy drops by 2.39 percentage points, compared with 11.75 for the Transformer baseline.

Evaluation uses task-specific Type encodings. Training augmentations, shift constructions, model configurations, and per-task results are detailed in the paper and appendix.

04 / Citation

Cite this work.

Questions about the paper?
yuanganl@uci.edu

@unpublished{lopes2026strat,
  title  = {The Geometry of Logic: Stratification Induces
            Semantic Structure and Robust Reasoning},
  author = {Lopes, Cristina V. and Li, Yuangang and
            Chen, Justin Tian Jin and Krone-Martins, Alberto and
            Ma, Iris and Misu, Md Rakib Hossain},
  year   = {2026},
  note   = {Preprint}
}