Reproducing When Models Manipulate Manifolds

The Geometry of a Counting Task — on an open-weight model

An independent, from-scratch reproduction of the mechanistic-interpretability findings in Anthropic's When Models Manipulate Manifolds: The Geometry of a Counting Task (Gurnee et al., 2025), carried out on the open-weight model Gemma-3-4B-it with Gemma Scope 2 SAEs, on a single consumer GPU.

_{The headline result. To wrap text the model must track two things — how many characters it has written (count) and how wide the line should be (width). Neither is stored as an integer: each lives on a smooth, low-dimensional curved manifold in the residual stream. Here the per-value mean activations at layer 19 trace clean, ordered arcs in their top-3 principal components (left: character count; right: line width). These are the open-weight analogues of the paper's central finding, recovered with nothing but probes + PCA.}

---

TL;DR

The paper reverse-engineers how Claude 3.5 Haiku decides where to put line breaks when producing fixed-width text — a task that secretly requires counting characters even though the model only sees tokens. Their finding is geometric and surprising: the model represents the count as position on a curved manifold, discretizes it with place-cell-like sparse features, and computes distance to the line limit by having attention heads geometrically "twist" the count representation against a learned line-width representation.

Because Haiku's weights are proprietary, this project reproduces the methodology on an open model and asks: do the same geometric phenomena appear? They do. Across four phases I find — and causally validate — a low-dimensional character-count manifold, a sparse place-cell code that tiles it, a separate line-width manifold, and a late-layer newline decision driven by the characters-remaining quantity k − count.

	Phenomenon from the paper	This reproduction (Gemma-3-4B-it)
§0	Model can do the implicit-wrap task	within-line newline AUC 0.94 (0.95–0.99 after line 3)	✅
§1	Count lives on a ~6-D curved manifold	top-6 PCs 93% var, R²(count|6PC) = 0.99, PC1↔count +0.94	✅
§1	Sparse "place cells" tile & dilate over count	6 SAE cells, 93% tiling, dilation ρ +0.83, decoder↔manifold duality +0.99	✅
§1	Count subspace is causally used for newlines	ablation hits newlines only (+2.6 nats); rank-6 steering flips the decision (P 0.01→0.51)	✅
§2	Separate low-D line-width manifold	R²(width|6PC) = 0.999, width readout RMSE 3.8 chars	✅
§2	Count & width are related-but-distinct axes	raw cosine +0.19, localized to a mid-network band (the angle §2 heads would "twist")	✅
§2	Newline = a linear readout near the output	binary is_break probe AUC 0.980 at L32	✅
§2	k − count ("chars remaining") drives the break	ablation +5.5 nats on newlines; steering P 0.78→0.01; rank-6 suffices	✅
§1/§2	How the count is built (additive accumulation)	scoped as future work (Phase X)	⏳
§2	Boundary heads & the QK "twist" (0.25→1.0)	raw geometry established; head analysis deferred (3b/3c)	⏳

Every result is gated by pre-committed pass/fail thresholds and stress-tested with controls (random-subspace, variance-matched, shuffled-label, within-fixed-width). Numbers differ from Haiku — a different, smaller, open model — and the README is explicit about where and why.

---

Table of contents

• Background: what the paper found
• Why a methodology reproduction
• Approach & design rules
• Results
  • Phase 1 — The model can do the task
  • Phase 2 — Representing character count (§1)
  • Phase 3 — Sensing the line boundary (§2)
• Reproduction scorecard vs. Haiku
• What is not (yet) reproduced
• Repository layout
• Reproducing the results
• References & citation
• Disclaimer & license

---

Background: what the paper found

When a model generates pre-wrapped fixed-width prose, it must decide where the next newline goes. To do that it has to (a) learn each token's character length, (b) accumulate those lengths into a running character count since the last newline, (c) infer the target line width k from context, (d) compute characters remaining = k − count, and (e) emit a newline when the next word would overflow. The paper's discoveries:

• Counts live on low-dimensional curved manifolds. Instead of an integer counter, "how many characters into the line am I" is a position along a smooth spiraling curve in the residual stream, discretized by sparse feature families that behave like biological place cells.
• The count is built additively. Early layers accumulate per-token character lengths into a running count manifold, and (separately) a line-width manifold for the target width.
• Attention heads "twist" the manifolds. "Boundary" heads use their W_QK geometry to rotate/align the count manifold against the line-width manifold, so their inner product spikes when the remaining characters fall in a target range — effectively computing distance to the limit.
• The newline decision is a linear readout. Estimates are arranged so a simple decision boundary fires the "predict newline" feature when the next word would overflow.
• Findings are validated with causal interventions and "visual illusions" — crafted sequences that hijack the counter.

Why a methodology reproduction

Claude 3.5 Haiku is proprietary; its weights and activations are unavailable, so the exact numbers cannot be re-derived outside Anthropic. This project instead reproduces the methodology and tests whether the same geometric phenomena appear in an open-weight model whose internals — and whose own SAEs (Gemma Scope 2, trained on Gemma's own weights) — are public.

Gemma-3-4B-it is the closest open analogue to Haiku: instruction-tuned, ~Haiku-class scale, with public same-model SAEs, and small enough (~8–9 GB bf16) to cache every layer at once, hold a wide SAE, and run causal interventions on a single RTX 4090 (24 GB).

Approach & design rules

The backbone is deliberately simple and falsifiable: linear/logistic probes + PCA on cached residual-stream activations, plus causal interventions (mean-ablation and activation patching). Pretrained SAEs are used only as optional, second-angle confirmation — never as the primary evidence. The hard part is not the machinery but avoiding the many ways to accidentally fake the result:

Design rule	Why it matters
Wrap width k is implicit — inferred in-context from preceding lines. We never instruct "wrap at N chars"; we feed raw pre-wrapped text as a continuation.	Telling the model the width changes the task into something that no longer requires counting.
"Count" = characters since last newline, measured on detokenized text (not token indices).	The entire phenomenon is about integrating variable-length tokens. Wrong labels invalidate every probe.
Teacher-force the analysis over correctly-wrapped strings.	Free generation compounds errors and makes per-position count labels ambiguous.
No private feature dictionary. Backbone = probes + PCA; SAEs are optional.	Anthropic's 10M-feature cross-layer transcoder is unavailable; the paper shows probes+PCA recover the same manifold.
Re-discover layers/heads empirically.	Haiku's layer indices don't transfer; layers are found by variance-in-6-PCs / readout-AUC criteria.
Control for the count↔width confound. A count-c token only occurs at width k ≥ c.	Pooled per-count means secretly mix widths; every key claim is re-checked at a fixed width.

A recurring open-model wrinkle that the paper (working in transcoder space) never sees: Gemma has a handful of massive-activation "scale" dimensions that dominate raw PCA and masquerade as PC1. Z-scoring the residual dims before PCA is decisive at the featured mid-network layer (raw PC1↔count +0.57 → z-scored +0.94) and is applied uniformly for cross-layer comparability.

---

Results

Layers are re-discovered, not assumed. The featured count/width layer is L19 (mid-network, where the count manifold is sharpest); the newline decision is read out at L32 (late, where the AUC peaks). All metrics below are held-out unless noted, and each phase carries pre-committed gates that all currently PASS.

Phase 1 — The model can do the task

📓 Full notebook: notebooks/phase1_summary.ipynb

Before any geometry, a behavioral gate confirms Gemma-3-4B-it actually performs implicit fixed-width wrapping. Teacher-forcing 4,200 wrapped passages across widths k ∈ {20,…,140} (≈ 0.97M within-line decision positions):

• Within-line newline AUC 0.942 overall; 0.952–0.988 once the model has seen ≥3 lines to infer the width.
• Mean P(newline) is 0.68 at true breaks vs 0.036 mid-line — a 19× separation.
• P(newline) rises smoothly as the line fills and spikes in the last few characters — exactly the boundary-sensing signature.

_{Mean P(newline) vs. characters remaining on the line: flat when there's room, sharply rising as k − count → 0.}

Gate: PASS → the behavior exists, so the geometry is worth dissecting.

Phase 2 — Representing character count (§1)

📓 Full notebook: notebooks/phase2_summary.ipynb · joint count↔width geometry in notebooks/phase2_ck_manifold.ipynb

2a · The count manifold. Grouping residuals by character count, averaging, z-scoring, and running PCA at L19 yields a ~6-D curved manifold: the top-6 PCs hold 93% of the variance, count is recovered from them at R² 0.99, and PC1 alone tracks count at +0.94. The count is already present but diffuse at an early layer (L2) and sharpens into a clean manifold by L19.

_{2a · The character-count manifold (PCA at L19). Left: PC1–3 of the z-scored per-count means trace a smooth, count-ordered arc — a clean low-dimensional manifold whose PC1 is itself the count axis (corr +0.94). Right: PC4–6 hold the residual, scrappier structure. Color = character count.}

The featured layer is not hand-picked — it is exactly where the count probe is sharpest:

_{2a/2b · Why layer 19. Character-count information by layer. The count-classifier RMSE (navy, the main metric, lower = better) bottoms out at L19 (≈10 chars vs. chance ~30), where the top-6 PC variance is high and the manifold is sharpest — present but diffuse early (L2), sharp mid-network. The featured layer is the data-driven RMSE minimum, not hand-picked.}

2b · Probes. An n-way multinomial logistic count classifier reads the count off a single token's residual with RMSE ≈ 10 chars (chance ≈ 30.5) and within-±5 accuracy 0.79; predictions concentrate tightly around the truth rather than scattering, and the readout directions themselves are geometrically ordered along the count axis:

_{2b · The count probe is sharp and geometrically ordered. Left (accuracy): the classifier's predictions concentrate on the diagonal (color = P(pred | true)) — counts are read off cleanly, not memorized. Right: the cosine-similarity matrix of the per-count readout-weight vectors forms a smooth count-ordered band (nearby counts aligned in red, far-apart counts opposed in blue) — the probe directions rotate continuously along the manifold, the readout-space echo of the 2a geometry.}

2c · Token length is linearly decodable from the static embeddings at R² 0.96, and the per-length mean embeddings trace a length-ordered arc with an oscillating component — the raw material the count is built from.

2d · Causal validation — the centerpiece, reproducing the paper's two "Prediction" experiments:

Necessity — ablate the count subspace. Mean-ablating the top-6 count PCs raises the loss only on newline targets (+2.6 nats at k=6, vs ~0 on non-newlines and ~0 for a random 6-D subspace), and the damage saturates at ~6 dimensions — the manifold's dimensionality. Newline AUC collapses 0.977 → 0.809, while variance-matched, massive-axis, and shuffled-count controls leave it untouched. Reproduces "ablating the count subspace has a large effect only when the next token is a newline."

Sufficiency — steer the perceived count. Patching a − μ_c0 + μ_c across the count-bearing band (L17–21) within a fixed width drives P(newline) monotonically (0.01 → 0.51 at k=60) and a rank-6 patch tracks the full-vector patch at corr ≈ 1.0. The model breaks the line based on the manipulated count, not the real one. Reproduces "intervening on a rank-6 subspace is sufficient to change the line-breaking behavior."

_{Steering the perceived count toward the line boundary raises P(newline); the rank-6 subspace patch (dashed) matches the full patch, and P(true next word) falls as the model decides to break instead.}

2e · SAE place cells (optional confirmation). Encoding the same L19 residuals with a Gemma Scope 2 SAE reveals a small family of place-cell-like features that discretize the manifold — the dictionary-space view of the same geometry:

6 place cells tile the count axis with localized, dilating receptive fields (93% coverage; dilation ρ +0.83, slope CI excludes 0 — a Weber-like widening at higher counts).

Feature↔geometry duality (the decisive, non-circular test): a cell's decoder direction predicts its preferred count at corr +0.99 (median localization error 4 chars vs 56 for random decoders). The decoder side never enters selection, so this is independent confirmation.

Gates: PASS (2a manifold, 2b probe + concentration, 2c token-length, 2d necessity + sufficiency, 2e all seven place-cell gates). SAEs are the most dictionary-dependent result; the place-cell family survives the strictest sampling tried (count-balanced, width-mixed, source-balanced, +1-bit specificity margin) for the denser dictionary.

Phase 3 — Sensing the line boundary (§2)

📓 Full notebook: notebooks/phase3_summary.ipynb

3a · The line-width manifold. On a finer width grid (k ∈ {20,25,…,150}), the newline-token residuals carry a smooth low-dimensional width manifold — the structural twin of the count helix: R²(width|6PC) = 0.999, PC1↔width +0.99, and a width classifier with RMSE 3.8 chars (chance 34.9). Width is held across an early/mid plateau and then discarded late — anti-phased with the newline readout that rises late as the width is "consumed."

The line-width manifold at L19 — a clean width-ordered arc, the analogue of the count helix.

Count and width are related-but-distinct axes. In a joint PCA they form two separate arcs; their cosine is just +0.19, and that alignment is localized to a narrow mid-network band (≈0.03 elsewhere) — exactly the raw geometry §2's boundary heads are reported to "twist" toward 1.

3d · The newline readout. A binary is_break logistic probe reads the decision off the residual stream at AUC 0.980, at a re-discovered late layer L32. AUC is low early (~0.65), jumps at L16–18, and plateaus near 0.98 — the decision is computed late, precisely where width information is consumed.

_{Left: newline-decision AUC by layer — flat early, then a sharp rise to 0.98 at the readout layer. Right: the residual dimensions that most push toward (predictors) and away from (suppressors) a newline.}

3e · Causal validation of "characters remaining." Two experiments parallel to 2d, but on the boundary quantity rem = k − count. Necessity: ablating the rem subspace raises newline loss by +5.5 nats (AUC 0.981 → 0.655) and nothing on non-newlines. Sufficiency: steering the perceived remaining distance drives P(newline) from 0.78 (line full) down to 0.01 (plenty of room), monotonically, with a rank-6 patch matching the full patch.

_{Steer the perceived "characters remaining": as room → 0 the model breaks the line (P→0.78); as room grows it doesn't (P→0.01). Together with 2d, this shows the model compares its running count against the inferred width — it breaks on k − count.}

Gates: PASS (3a width manifold + probes + joint count/width, 3d readout, 3e necessity + sufficiency). Honest caveat: since rem = k − count, the rem axis anti-aligns with the count axis (cos −0.70), so 3e extends 2d to the explicit boundary framing rather than being fully count-independent. The boundary heads and the QK "twist" (3b/3c) are not yet analyzed — 3a establishes the raw count/width geometry those heads would operate on.

---

Reproduction scorecard vs. Haiku

Haiku anchors are the paper's reported numbers on a different, proprietary, larger model — targets, not expectations. The point is whether the phenomenon reproduces, not whether the digits match.

Phenomenon	Haiku (paper)	This repro (Gemma-3-4B)	Verdict
Newline-decision AUC	≈ 0.91	0.980 (readout L32)	✅ exceeds
Count manifold dimensionality	≈ 6-D, top-6 ≈ 95% var	top-6 93% var, R²(count|6PC) 0.99, PC1 +0.94	✅
Count probe quality	R² ≈ 0.985 (continuous regression)¹	classifier RMSE 10 chars, within-±5 0.79	⚠️ partial¹
Place-cell family	~10 cells, dilating, 2–3 co-active	6 cells, 93% tiling, dilation +0.83, duality +0.99	✅
Ablate count subspace → newline-specific	yes	+2.6 nats newline-only, saturates at 6-D	✅
Rank-6 steering is sufficient	yes	P(nl) 0.01→0.51, rank-6 ≈ full	✅
Line-width manifold	low-D	R²(width|6PC) 0.999, RMSE 3.8 chars	✅
Raw count/width alignment	cos ≈ 0.25	cos +0.19 (mid-network band)	✅
QK "twist" (0.25 → ~1.0)	yes	deferred (3b)	⏳
Additive accumulation (R² 0.93→0.97)	yes	deferred (Phase X)	⏳

¹ Haiku's 0.985 is a continuous-regression R² and is not on the within-±5 scale used here; a data-coverage study found the open-4B within-±5 readout saturates at ~0.82 and is not closed by more data. This is reported as a genuine quantitative gap between a 4B open model and Haiku — not papered over — and does not by itself establish an essential mechanistic difference (larger-model / layer-sweep ablations remain untested).

What is not (yet) reproduced

Kept explicit for honesty — these are scoped as future work, not silently omitted:

• Boundary heads & the QK "twist" (3b/3c). The marquee §2 mechanism (heads rotating count vs. width via W_QK so the inner product spikes near the limit). The raw count/width axes it operates on are established (3a); the per-head W_QK analysis — which on Gemma must fold RMSNorm + QK-norm + RoPE and handle GQA — is the next milestone.
• Additive accumulation (Phase X). How the count is built by summing token lengths across early-layer head sets (Haiku's R² 0.93 → 0.97 build-up). Phase 2 characterizes the count; this would reconstruct how it is computed.
• Visual illusions (3f, stretch). Crafted two-character inserts that hijack the counter.

Repository layout

src/
  data/        corpus streaming · newline-stripping · greedy word-wrap · per-token labels
  manifold/    per-count means · z-scoring · PCA · helix geometry
  probes/      logistic count probe · token-length probe · binary newline readout
  causal/      subspace ablation · activation patching · interventions
  saes/        Gemma Scope 2 loading · per-feature place-cell tuning curves
  cache_util.py
scripts/       00–04 env/smoke/tokenizer/SAE/attention checks · 10–11 data + behavioral
               20–33 Phase 2 (geometry, probes, causal, SAEs) · 40–42 Phase 3 (width, causal)
notebooks/     phase0–phase3 summaries (one per phase, mirroring the paper's figures)
results/       per-phase write-ups (.md) + metrics (.json) + cached activations (.npz)
figures/       saved plots — phase1 / phase2 / phase2e / phase3
ROADMAP.md     the full phase-by-phase plan, design rules, and Haiku anchors

Each phase ends in a notebook + saved figures + a results/phase*.md recording our numbers against the Haiku anchors, the pre-committed gates, and exactly where the open model diverges.

📓 Notebooks (each reproduces its phase's figures from that phase's cached results):

Notebook	Contents
phase0_summary.ipynb	Environment, model loading & correctness checks (logits-match, tokenizer/offset, SAE reconstruction)
phase1_summary.ipynb	Behavioral gate — implicit fixed-width wrapping
phase2_summary.ipynb	Count manifold (2a), probes (2b/2c), causal ablation + steering (2d), SAE place cells (2e)
phase2_ck_manifold.ipynb	Joint count↔width manifold geometry
phase3_summary.ipynb	Line-width manifold (3a), newline readout (3d), characters-remaining causal (3e)

Reproducing the results

Environment — single RTX 4090 (24 GB), Python 3.11, conda env manifold. Install CUDA-pinned PyTorch first (transformer-lens pulls in torchvision, which can otherwise silently clobber CUDA with a CPU build):

conda create -n manifold python=3.11 && conda activate manifold
pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu126
pip install -r requirements.txt
huggingface-cli login            # then accept the Gemma-3-4b-it license on its HF page

Verify, then run a phase (Gemma-3 is gated and uses newer attention variants; the 00–04 checks gate everything downstream — including a logits-match correctness check and tokenizer/offset round-trips):

# 0 · verify the harness (gates everything downstream)
python scripts/00_env_check.py          # torch +cu126, CUDA, bf16, VRAM
python scripts/01_smoke_test.py         # load model (text tower) + one Gemma Scope 2 SAE
# optional deeper checks: 02_tokenizer_check · 03_sae_check · 04_attn_heads_check

# 1 · data + behavioral gate
python scripts/10_build_dataset.py      # -> data/wrapped_eval.jsonl (+ per-token count/width labels)
python scripts/11_behavioral_eval.py    # Phase 1 gate            -> figures/phase1/

# 2 · representing character count (§1)
python scripts/20_count_geometry.py             # 2a/2b/2c -> results/phase2_cache.npz, figures/phase2/
python scripts/30_causal_validation.py          # 2d control battery (selectivity / AUC-collapse)
python scripts/31_causal_paper_experiments.py   # 2d ablation + intervention -> figures/phase2/causal_*
python scripts/32_sae_place_cells.py            # 2e SAE place cells -> figures/phase2e/

# 3 · sensing the line boundary (§2)
python scripts/40_build_width_dataset.py        # -> data/wrapped_width.jsonl   (REQUIRED before 41)
python scripts/41_width_geometry.py             # 3a/3d -> results/phase3_width_cache.npz, figures/phase3/
python scripts/42_causal_chars_remaining.py     # 3e -> figures/phase3/causal_rem_*

# README composite figures (re-render from the caches written by 20 and 41 — no model pass)
python scripts/50_headline_figure.py    # -> figures/headline_manifolds.png            (needs 20 + 41)
python scripts/51_probe_2b_figure.py    # -> figures/phase2/probe_accuracy_cosine.png  (needs 20)

The per-phase notebooks in notebooks/ reproduce that phase's figures from its cached results, so run the phase's scripts above first to generate the caches; the headline and combined figures are rendered by scripts/50 and scripts/51. Large caches/datasets are .gitignored and regenerated by the scripts above.

---

References & citation

Wes Gurnee, Emmanuel Ameisen, Isaac Kauvar, Julius Tarng, Adam Pearce, Chris Olah, Joshua Batson. "When Models Manipulate Manifolds: The Geometry of a Counting Task." Transformer Circuits Thread, 2025. https://transformer-circuits.pub/2025/linebreaks/index.html · arXiv:2601.04480

@article{gurnee2025manifolds,
  title   = {When Models Manipulate Manifolds: The Geometry of a Counting Task},
  author  = {Gurnee, Wes and Ameisen, Emmanuel and Kauvar, Isaac and Tarng, Julius and
             Pearce, Adam and Olah, Chris and Batson, Joshua},
  journal = {Transformer Circuits Thread},
  year    = {2025},
  url     = {https://transformer-circuits.pub/2025/linebreaks/index.html},
  note    = {arXiv:2601.04480}
}

Model: google/gemma-3-4b-it · SAEs: google/gemma-scope-2-4b-it · Tooling: TransformerLens, SAELens, nnsight.

Disclaimer & license

This is an independent reproduction created for research and educational purposes. It is not affiliated with or endorsed by Anthropic or Google. All credit for the original findings belongs to the paper's authors; all results here are on a different, open-weight model and differ from the proprietary originals.

MIT © 2026 Guanchun Li