Mechanistic Interpretability

BlueDot TAIS — Puzzle 1

Mechanistic Interpretability: Finding and Building Unusual Representations

8 binary features 5-layer MLP 9 new encodings

A frozen sentence-transformer encoder (all-MiniLM-L6-v2, 384-dimensional) feeds a 5-layer MLP trained to predict 8 binary features from natural language sentences. The MLP has architecture 384→64→64→64→64→8. One feature — country — is encoded in an unusual, nonlinear way at the hidden layer h2. We identify it by comparing linear and nonlinear probe accuracies across all layers, characterise the exact encoding mechanism (an absolute-value code with Z/2 symmetry along the food axis), and then deliberately train nine models with geometrically richer representations spanning XOR parity, SO(2) rotational codes, JEPA residuals, helical geometry, superposition, a multi-class circular code, SAE decomposition, a bilinear model, and an emergent unit-norm bottleneck.

Setup

The Model & Task

Architecture overview and the 8 binary features the MLP is trained to predict from frozen sentence embeddings.

Model Architecture

Embedding 384

H0 64

H1 64

H2 64 ★

H3 64

Output 8

★ Unusual encoding introduced at H2

All embeddings pre-computed and frozen. Only the MLP is trained.

8 Binary Features

0 number

1 question

2 color

3 food

4 sentiment

5 country ★

6 person

7 body_part

★ country is encoded in an unusual nonlinear way

Each sentence in the dataset has exactly one binary label per feature. The model must learn all 8 simultaneously from the shared 64-dimensional hidden layers. Features share representational capacity — the country encoding exploits this by latching onto the food axis.

TASK 1

Finding the Unusual Feature

Linear and nonlinear probes trained at every layer reveal a striking divergence for exactly one feature.

Chart A — Probe Accuracy at H2 by Feature

Chart B — Country Probe Accuracy Across Layers

Finding: country is the only feature where linear and nonlinear probes diverge. The gap of 0.495 at H2 is uniquely diagnostic — every other feature sits between 0.97 and 1.00. The gap appears only at H2: before it the embedding still carries a linearly decodable country signal, after it the model has re-encoded it into a nonlinear code. The encoding is introduced by the H2 transformation. Leave-one-template-out cross-validation confirms this is not an artefact (linear = 0.486, nonlinear = 0.995).

TASK 2

The Encoding Mechanism

How exactly is country stored at H2? The answer involves the food axis, absolute values, and Z/2 symmetry.

The food and country labels are statistically independent in the dataset (label correlation = −0.005). Yet the model representationally entangles them — it co-encodes both features on the same activation axis at H2. Projecting H2 activations onto the direction that best separates food=0 from food=1 reveals four distinct intervals. The key observation: country=1 clusters near zero on this axis regardless of food value, while country=0 is pushed to the extremes — large negative (food=0) or large positive (food=1). In other words: country = 1 if and only if |projection on food axis| is small. This is an absolute-value (magnitude) code with Z/2 symmetry.

Chart C — Distribution Along the Food Axis (H2)

Chart D — Decoder Accuracy Comparison

The Z/2 symmetry: Both (country=0, food=0) and (country=0, food=1) map to a large |projection|, just with opposite signs. A linear probe sees the mean of positive and negative extremes — which is near zero — and cannot decode the feature. Replacing the raw projection with its absolute value jumps accuracy from 0.502 (chance) to 0.946. Any even function works: projection² achieves identical accuracy.

Minimum Complexity: How Many Hidden Units Are Needed?

Effectively 1D: A 2-hidden-unit MLP achieves 0.950. 1 unit achieves only 0.721. A 2-ReLU circuit suffices to decode it: |x| = ReLU(x) + ReLU(−x). A single ReLU cannot recover an even function, so two units is a natural lower bound for this Z/2-symmetric code — consistent with the empirical jump from 1 to 2 units.

Why did the model choose this representation?

Both country and food are linearly decodable from the raw sentence-transformer embeddings — country at 0.980 and food at 0.967 — so the model could have kept both as independent linear directions throughout. Instead, at H2 it spontaneously shares a single axis: food as sign, country as magnitude. This is representational superposition — two features on one axis, with country compressed into the magnitude channel. One plausible pressure is capacity: the 64-dimensional bottleneck must carry 8 features simultaneously, and co-encoding two features that happen to align in embedding space is a natural compression strategy, even when their labels are uncorrelated.

The H3 decoding circuit: After H2, the model must recover country from the magnitude code. Inspecting the H2→H3 weight matrix reveals the mechanism directly: two neurons have near-perfectly equal-and-opposite loadings on the food axis (+2.603 and −2.522, ratio 1.03). One neuron activates for positive food-axis projections, the other for negative projections. Together they implement |x| ≈ ReLU(x) + ReLU(−x) — a two-ReLU circuit for recovering magnitude from a signed code. These two neurons alone probe country at 0.895 — recovering most of the country signal, with full H3 reaching 0.961, so the remaining neurons add some meaningful cleanup. This lines up with the probe analysis: 2 hidden units suffice, consistent with the two ReLUs this magnitude-decoding circuit uses.

TASK 3

Nine New Representations

Nine models trained with geometrically richer encodings, each connecting to a different theoretical framework.

Before running experiments, "more interesting" had to be defined. Two definitions were used: (1) probe-resistant — linear probe stays near chance, as in the baseline; (2) geometrically richer — higher-dimensional symmetry, multiple features sharing one space, angular rather than magnitude coding. A representation can be geometrically richer even if it remains linearly decodable — that follows from the binary label structure, not the encoding.

The main lessons: probe resistance requires architectural pressure, not harder labels — depth linearises internally regardless of label complexity (XOR, helix), and binary labels are always linearly separable, so SO(2) and helical codes cannot resist a linear probe however clean the geometry. Genuine probe resistance needs multi-class structure — the MNIST 10-class circular code drives the linear probe to chance (0.51) while the nonlinear probe holds at 0.985, a 47-point gap and the strongest result here. Geometric richness can emerge unprompted — the unit-norm d-sphere bottleneck lets person develop a +0.42 linear/nonlinear gap with no target angles, which (with the superposition model's shared 2D space) connects directly to how neural networks store more features than they have dimensions.

XOR / Parity Code

Theory: Parity-pruning in state-tracking transformers — Li et al., arXiv:2503.02854, ICML 2025

A synthetic feature parity = sentiment XOR question is introduced. A 9th output head is added and trained to predict it. The model must represent both inputs jointly — no single linear direction encodes XOR. This is a controlled miniature version of the compositional state-tracking mechanism described by Li et al. for real transformers doing compositional state updates.

Probe	Accuracy
Linear (h2)	0.951
2-layer MLP	0.982

This experiment failed to produce a weirder representation. Linear probe = 0.951 means h2 is more linearly decodable than the baseline (0.471), not less. The model had enough depth to linearise XOR before h2. The interesting finding is what this reveals about the model — it composes parity internally, exactly like the Li et al. mechanism at scale — but the representation itself is not harder to probe. Forcing a weirder code requires architectural constraints, not just a harder label.

Rotational / SO(2) Code

Theory: LeJEPA identifiability — Klindt et al., arXiv:2605.26379 — transition operators must be linear and orthogonal

A circular regulariser pushes a learned 2D projection of H2 onto the unit circle: country=0 → 0°, country=1 → 90°. This is a phase code rather than a magnitude code. The update from country=0 to country=1 is a 90° rotation in the learned 2D subspace. LeJEPA's identifiability theorem states that a consistent world-model transition operator must be linear and orthogonal — exactly a rotation matrix.

Probe	Accuracy
Linear on full h2	0.989
Angular decoder (2D)	0.988

The angular structure is clean — country=0 clusters at mean 1.0° (std 8.5°), country=1 at mean 89.4° (std 8.7°). But points are not on the unit circle: country=0 mean radius = 6.24, country=1 = 3.67. The model mixed phase code (angle) and magnitude code (radius) simultaneously. Linear probe = 0.989 because two orthogonal clusters are always linearly separable — this is a binary-label limitation, not a regulariser failure. Rotational codes only resist linear probes for 3+ classes on a circle.

JEPA Residual Code

Theory: Direct implementation of the JEPA objective. P linearity = LeJEPA identifiability condition.

A predictor P takes the 7 other binary feature labels (all features except country) as input and predicts the full H2 activations. The residual H2 − P(other_labels) is what the other labels cannot explain — country's contribution. The empirical question: is that residual direction consistent across contexts, and is P linear? If P is linear, the LeJEPA identifiability condition is satisfied.

Probe	Accuracy
Linear on raw h2	0.471
Linear on residual	0.551
Other features (residual)	~1.000

Finding: A linear matrix W approximates P well — the unexplained fraction ||P(z)−Wz||² ≈ 0. P is essentially linear, empirically confirming the LeJEPA condition in a fully controlled ground-truth setting.

Helical Code

Theory: Transformer positional encodings and grid-cell theories of spatial encoding

A regulariser pushes a learned 3D projection of H2 onto a helix: country=0 → (1, 0, 0), country=1 → (0, 1, pitch). Motivated by how transformer positional encodings and grid cells in the mammalian hippocampus encode continuous location using phase-coded representations.

Decoder	Accuracy
1D projection	0.989
2D angular	0.993
3D helical fit	0.993

Finding: Helix collapses to 1D for a binary label. The two classes sit at the tips of two separate 1D manifolds rather than on a spiral. Helical structure requires multi-valued (ordinal or continuous) variables to be informative.

Superposition / Entanglement Code

Theory: Anthropic's superposition hypothesis — models under capacity pressure interleave features

Country and food are forced to share a 2D bottleneck with interleaved target positions: (c=0,f=0)→(1,0), (c=1,f=0)→(0,1), (c=0,f=1)→(−1,0), (c=1,f=1)→(0,−1). Neither feature occupies a dedicated dimension. This is a deliberate, controlled demonstration of the superposition failure mode — not mysterious, just geometry.

Feature	Accuracy	Baseline
country	0.978	0.503
food	0.738	0.501

The four groups approximately hit their target angles — (0,0) at 8.7°, (0,1) at 81.2°, (1,0) at 175.1°, (1,1) at 273.4° — close to the intended 0°/90°/180°/270°. But precision is uneven: (0,0) std 22°, (0,1) std 10°, (1,1) std 39.5°, and (c=1,f=0) collapses into a ray with std 71.9° and a larger mean radius (4.04 vs 1.4–2.1 for the others). The ray causes the cross shape in the scatter plot and explains why food probes at only 0.738 — the (c=1,f=0) ray bleeds angularly into adjacent regions. Country still probes at 0.978 because the left/right boundary remains clear.

MNIST Circular Code (10-class) — with unit-norm fix

Theory: Fix to Idea 2 — rotation works when there are 3+ classes interleaved on the circle. Unit-norm constraint eliminates the radius shortcut.

A small CNN trained on MNIST with a 2D circular bottleneck. Digit k is pushed toward angle k × 36°. Even digits (0,2,4,6,8) and odd digits (1,3,5,7,9) alternate perfectly around the circle. Unit-norm constraint (F.normalize on bottleneck) forces all representations onto the unit circle, removing the radius magnitude as a linear shortcut.

Probe	Accuracy
10-class digit accuracy	0.981
10-class linear probe on 2D	0.982
Even/odd linear probe on 2D	0.510
Even/odd nonlinear probe on 2D	0.985
Base rate	0.508

With unit-norm, the even/odd linear probe drops to 0.510 ≈ chance (base rate 0.508), while the nonlinear probe stays at 0.985 and digit accuracy improves to 0.981. The 47 percentage-point gap is the largest in this study. Without unit-norm, the linear probe was 0.624 because the model used radius as a shortcut — the constraint removes that pathway entirely.

SAE Decomposition of h2

Theory: Sparse dictionary learning / sparse autoencoders as a basis for mechanistic analysis of distributed representations.

A top-k sparse autoencoder (k=10, d_feats=256) is trained on the 64-dimensional h2 activations of the original puzzle model. Each sparse feature is then tested as a linear country probe, individually and jointly.

Decoder	Accuracy
Single best SAE feature (feat 100)	0.519
Second best SAE feature (feat 36)	0.592
Top-2 pair, linear probe	0.707
Top-2 pair, nonlinear probe	0.717
All 256 SAE features, linear	0.827
All 256 SAE features, nonlinear	0.958

Country's Z/2 encoding distributes across many sparse features — no single dictionary atom reaches above 60%. The top-2 features capture the two arms of the V-shape (positive and negative food projections), together achieving 0.707. All 256 features recover 0.958 nonlinearly, matching the original nonlinear probe. The representation is genuinely distributed: country is not monosemantically stored in any single SAE feature.

Bilinear Text Model + CPD Analysis

Theory: Bilinear layers are their own CPD decomposition. The original country encoding is a rank-1 bilinear form — a bilinear model is the natural architecture.

A single bilinear layer replaces the full ReLU MLP: h2 = (W_L · emb) ⊙ (W_R · emb). By construction, B[f,i,j] = Σ_r D[f,r] L[i,r] R[j,r] — the weight matrices ARE the CPD factors.

Feature	Linear probe	Nonlinear probe
country (baseline: 0.471 linear)	0.993	0.992
food	0.983	0.980
question	0.999	0.999
sentiment	0.981	0.977
XOR (bilinear vs ReLU)	0.963 vs 0.951 (Δ=+0.013)	0.977

Key finding: A bilinear layer encodes country linearly (0.993 vs 0.471 in the original). A bilinear layer naturally implements quadratic forms — the rank-1 form (wf·h2)² becomes a linear function of h2 in the bilinear basis. The original model's probe resistance for country arises from the ReLU nonlinearities, not from the bilinear product structure. XOR is slightly more linearly accessible in the bilinear model (Δ=+0.013) — bilinear directly computes the product of two linear features.

Bottleneck Superposition (Emergent)

Theory: Anthropic's superposition hypothesis — capacity pressure forces multiple features to share dimensions. Unit-norm eliminates the magnitude channel, so geometry must be angular.

A narrow d-dimensional bottleneck with unit-norm (F.normalize) forces all representations onto the d-sphere. Two variants: d=4 and d=2. No target positions are prescribed — geometry is fully emergent. Eight per-feature MLP heads (d→16→ReLU→1) achieve high accuracy regardless of geometry. The unit-norm constraint eliminates the magnitude channel used by the original model, so all information must be angular.

d = 4 probes (acc=0.846)

Feature	Linear	Nonlinear	Gap
number	0.583	0.599	+0.015
question	0.993	0.993	+0.000
color	0.679	0.787	+0.107
food	0.894	0.896	+0.002
sentiment	0.551	0.759	+0.208
country	0.937	0.936	-0.001
person	0.485	0.907	+0.421
body_part	0.866	0.866	+0.000

d = 2 probes (acc=0.697)

Feature	Linear	Nonlinear	Gap
number	0.514	0.536	+0.022
question	0.482	0.521	+0.039
color	0.521	0.531	+0.009
food	0.705	0.893	+0.189
sentiment	0.608	0.797	+0.189
country	0.677	0.858	+0.181
person	0.495	0.492	-0.003
body_part	0.949	0.948	-0.001

Angular frequency R² (d=2): R² of y ~ cos(k·θ)+sin(k·θ)

Feature	k=1	k=2	k=3	k=4	Peak k
number	0.002	0.001	0.004	0.004	k=4
question	0.003	0.005	0.002	0.001	k=2
color	0.002	0.004	0.004	0.004	k=2
food	0.228	0.318	0.085	0.078	k=2
sentiment	0.103	0.229	0.060	0.010	k=2
country	0.147	0.486	0.045	0.043	k=2
person	0.001	0.002	0.001	0.001	k=2
body_part	0.665	0.021	0.108	0.016	k=1

Key finding: In d=4, person has the largest nonlinear gap of any text-domain experiment (+0.421): linear probe = 0.485 (chance), nonlinear = 0.907 — emergent nonlinear encoding with no prescribed geometry. In d=2, body_part dominates the unit circle (k=1, R²=0.665 — clean half-arc split), while country, food, and sentiment use k=2 angular encoding (alternating twice around the circle). The remaining four features are squeezed out of the angular structure entirely, consistent with the 2D capacity limit. This is the first text-domain experiment where probe resistance emerges purely from architectural pressure.

Square Wave (Constructive) — the binary-country attack

Theory: k=2 on a circle is the |x| baseline bent into 2-D — an even-symmetric code. A dose-response over harmonic k tests whether imposing a harmonic readout defeats a linear probe.

emb → enc → Linear(64,2) → F.normalize → θ. The country logit is constrained to the k-th harmonic α·cos(kθ)+β·sin(kθ)+b; the other 7 features use free MLP heads on the circle and act as the spreader. Architecturally induced, not emergent. Sweep k=1..5.

k	Linear	Nonlinear	Freq peak	country=1 arcs	Overall acc
1	0.503	0.539	k=4	12	0.710
2	0.711	0.959	k=2	12	0.719
3	0.961	0.967	k=3	5	0.683
4	0.985	0.985	k=3	3	0.662
5	0.987	0.988	k=4	3	0.659

Partial — honest. k=2 is the sweet spot: the frequency peak lands on the imposed k=2 and the nonlinear probe hits 0.959 (+0.25 gap). But the linear probe stays at 0.711 — it does not beat the baseline's 0.471: the binary label leaks into the k=1 half-circle split. Constraining the readout to k=2 does not force the geometry to be k=2. At k=1 country collapses to chance on both probes; at k≥3 the model abandons the fast harmonic and country becomes linearly trivial (arcs collapse 5→3).

Fourier Comb (Constructive) — superposition showpiece

Theory: multiplexing more features than linear directions is controlled superposition. The linear-probe failure is the geometry; the harmonic-confusion matrix is the measurement.

Same 2-D circle, three features multiplexed onto distinct harmonics: country→k=2, food→k=3, sentiment→k=4, each with its own harmonic readout; the other five use free MLP heads. One scalar angle is asked to carry three independent bits.

Multiplexed-feature probes (acc=0.682)

Feature (target k)	Linear	Nonlinear
country (k=2)	0.529	0.525
food (k=3)	0.937	0.939
sentiment (k=4)	0.598	0.867

Harmonic-confusion R² (peak bold)

Feature	k=1	k=2	k=3	k=4	k=5
question	0.773	0.823	0.183	0.038	0.121
food	0.639	0.735	0.699	0.073	0.027
sentiment	0.056	0.014	0.070	0.550	0.496
country	0.005	0.005	0.002	0.001	0.002

Partial — one clean multiplex, a diagnosable failure. Only sentiment lands on its target harmonic (peak k=4, +0.27 gap). food leaks to k=2 and stays linearly readable; country collapses to chance on both probes. The confusion matrix shows why: question (unconstrained, the easiest feature) hijacks k1/k2 (R² 0.77/0.82), starving the multiplexed features of low-frequency capacity. Three bits on one angle is past the d=2 capacity ceiling Idea 9 already flagged at ~k=2.

Linked Rings (Constructive) — topological showpiece, strongest result

Theory: two linked rings are provably not linearly separable — any hyperplane that isolates ring A forces all of ring B onto the same side. Probe resistance by topology, not tuning.

emb → enc → Linear(64,3) 3-D bottleneck. A regulariser pushes country=0 onto ring A (xy-plane, centre origin) and country=1 onto ring B (xz-plane, centre (1,0,0)) so B threads A (linking number 1); food sets the within-ring angle (spreader). All 8 features decode through per-feature MLP heads. Loss = BCE + λ·(distance to target ring), λ=1.0.

Metric	Value
country linear probe	0.533 (near chance)
country nonlinear probe	0.940
linear/nonlinear gap	+0.41
net disc-crossings (linking number)	1
ring-B arc occupancy	9 / 12
overall 8-feature acc	0.806

Success — the open problem, answered constructively. The linear probe sits at 0.533 (near chance, vs baseline 0.471) while nonlinear holds 0.940 — a +0.41 gap — and the linking is real: a numerical disc-crossing count returns exactly one net crossing of the country=1 cloud through ring A's disc. The collapse risk did not materialise: even though the nominal spreader (food) is binary, ring B is populated across 9 of 12 arcs because BCE pressure from the other features spread country=1 around the ring. A binary feature can be made probe-resistant by construction — but it takes a topologically non-separable geometry, not a high-frequency readout (Idea 10) or multiplexing (Idea 11).

SUMMARY

Comparison of All Encodings

Ordered from clearest failure to clearest success. Most approaches failed to beat the baseline on probe difficulty — and those failures are the most instructive results.

Encoding	Linear probe	Verdict	Key lesson
Baseline (abs-value)	0.50	—	Z/2 symmetry, 1D
XOR / parity	0.95	Failed — linearised	Depth beats label complexity
Helical	0.99	Failed — collapsed to 1D	Helix needs multi-valued labels
Rotational SO(2)	0.99	Partial — orthogonal structure achieved	Binary labels always linearly separable
JEPA residual	0.47	Informative — erasure not encoding	Explicit erasure strips a feature from a layer
Superposition	0.98 / 0.74	Partial — geometrically richer	2D shared encoding; two features, one space
MNIST circular (10-class) + unit-norm	0.51	Best — genuinely probe-resistant	Multi-class interleaving + unit-norm constraint
SAE decomposition (h2, all 256 feats)	0.83	Analytical — distributed encoding	Country spreads across many sparse features
Bilinear model (country)	0.993	Informative — ReLU is the source	Bilinear linearises the rank-1 quadratic form
Bottleneck d=4 (unit-norm, emergent)	0.49 (person)	Emergent — nonlinear by capacity	person: linear=0.49, nonlinear=0.91 (+0.42 gap)
Bottleneck d=2 (unit-norm, emergent)	0.48 (question)	Emergent — angular superposition	body_part k=1; country/food/sentiment k=2
Square wave k=2 (constructive)	0.711 (country)	Partial — real k=2 code	Harmonic readout leaks to k=1; not sufficient
Fourier comb (constructive)	0.529 (country)	Partial — only sentiment multiplexed	Easy unconstrained feature hijacks low harmonics
Linked rings (constructive)	0.533 (country)	Success — probe-resistant + topological	Non-separable topology (linking #1); +0.41 gap

Chart F — Linear Probe Accuracy by Encoding

Key takeaway: The MNIST circular (with unit-norm) achieves linear probe = 0.51 (chance), and the text bottleneck d=2 achieves 0.48 — the lowest linear probe for any text-domain encoding tested. The baseline abs-value and JEPA residual also defeat linear probes. The bilinear analysis reveals why the original encoding works: the country encoding is a rank-1 bilinear form, and a bilinear layer naturally exposes this as a linear feature (0.993). The probe resistance in the original model comes from ReLU nonlinearities, not the multiplicative structure. The SAE analysis confirms country is distributed across many sparse features — no single dictionary atom decodes it alone.

Where the baseline stands after the capstone. The original country code sits at a linear probe of 0.471 (chance). For the first nine experiments nothing matched it: the superposition and bottleneck models make country more linearly decodable (0.978 and 0.677), and the only probe-resistant win there (MNIST, 0.51) abandoned the country feature for a 10-class digit task. That motivated the capstone, and Idea 12 (linked rings) closed most of the gap: a constructive binary-country code with linear probe 0.533 (near chance) and nonlinear 0.940 — probe-resistant and topologically weirder than |x| (verified linking number 1). The honest caveats: 0.533 is marginally above 0.471 rather than below it, and the code is constructed, not emergent. The cheaper constructions show why this is hard — a harmonic readout (square wave k=2 → 0.711) leaks into the k=1 half-circle split, and multiplexing (Fourier comb) collapses country when an easy unconstrained feature hijacks the low harmonics. The finding: a binary feature already at chance is near a local optimum for probe resistance, and the only thing that reliably beats a cleverer two-cluster geometry is a genuinely non-separable structure — multi-class interleaving (MNIST) or a linked topology (Idea 12). Driving binary country strictly below 0.471 while keeping it weird — and getting that geometry to emerge rather than be imposed — remains open.