Motivation

One of the most important conjectures in mechanistic interpretability is the notion of ‘universality’ - that is, that for any subtask which an LLM needs to perform, there exists some optimal, or as I like to think of it, “Platonic ideal”, collection of features/representations and circuits to solve it. Consequently, regardless of architectural details, any neural network should implement an imperfect version of that ideal representation. If this hypothesis holds, we should observe that independently trained models display similar behavior, and studying one model should allow us to understand any other. This would indeed be excellent news for the project of interpretability, to say nothing of its intriguing consequences on our understanding of intelligence.

Several lines of prior work have explored this idea by training small models, often on particular tasks. For example, toy models have been used to study superposition (Elhage et al., 2022) and small transformers trained on synthetic tasks have been shown to develop interpretable mechanisms for copying (Olsson et al., 2022), modular addition (Nanda et al., 2023), sorting and parsing (Weiss et al., 2021; Allen-Zhu & Li, 2023), and even representing board states (Li et al., 2023). While small models are often interpretable, work by (Chughtai et al., 2023) on group composition revealed that mechanisms across models may belong to a family of solutions even though particular features and circuits are arbitrary. My project investigates this weaker notion of universality on the synthetic task of predicting line breaks.

Specifically, this project was inspired by When Models Manipulate Manifolds: The Geometry of a Counting Task (Gurnee et al., 2026), a paper authored at Anthropic that outlines a full mechanistic breakdown of line break prediction in Claude 3.5 Haiku. One major finding was that character count was not represented linearly, as features are often assumed to be, but on a curved low-dimensional manifold. They found attention heads that acted on this feature manifold to estimate the distance to the line boundary, enabling the model to determine when to break the line instead of writing the next textual token.

This works trains a small transformer purely on the task of predicting line breaks in an entirely synthetic setting in order to discover whether the model will implement similar mechanisms. We find provide a mechanistic breakdown of the algorithm the model uses and find similar motifs - a low-dimensional non-linear representation of character count that is acted on primarily by attention heads - but with several important differences.

Model Architecture, Synthetic Data, and Training

The model is a small causal Transformer with three layers and four attention heads per layer. Its residual stream has width 128, each head has dimension 32, and the MLPs have width 512. The model uses GELU activations, layer normalization, standard learned positional embeddings, and a vocabulary of 82 tokens. The context length is 150 tokens. This architecture is intentionally small enough to inspect directly, while still providing enough capacity to represent and manipulate the counting variable.

The training data are generated procedurally rather than stored in a fixed dataset. The vocabulary contains two special tokens, BOS and _NEWLINE_, plus 80 ordinary tokens divided evenly among eight token lengths: there are ten distinct token identities of each length from 1 through 8 characters. To generate one example, I first sample 150 ordinary tokens uniformly at random and prepend BOS. I then sample a line-width limit uniformly from the integers 50 through 100. Scanning from left to right, I accumulate the character lengths of the tokens; whenever the current line has reached the sampled limit, I insert _NEWLINE_ before the next token and reset the line counter. The resulting sequence is cropped to the required context and shifted by one position to form the inputs and next-token targets. Thus, each example contains a fresh random token sequence and a single line-width limit shared by all of its lines, while the correct newline locations are determined exactly by the token lengths and the running count.

The model is trained with the standard autoregressive next-token objective: the mean cross-entropy between its predicted distribution and the target token at every position. Newline prediction is therefore learned as part of ordinary language modeling rather than through a separate binary newline objective. In the main configuration, batches contain 256 examples and are generated on the fly. I train for 10,000 optimization steps with AdamW, using a learning rate of , weight decay of , gradient-norm clipping at 1, and a cosine learning-rate schedule with 500 warmup steps and a minimum learning rate of . Evaluation uses a fixed randomly generated batch so that changes in performance can be compared across training checkpoints.

Activations at Each Layer

The following plots show the activation trajectory for one example in a separate three-dimensional PCA space at the residual stream immediately before each attention layer. Each point corresponds to a token position, with the trajectory colored by cumulative character count. The PCA basis is fit separately for each layer, so the plots are useful for comparing the geometry within a layer rather than the absolute orientation of the spaces across layers.

Activation Geometry Before Layer 2

The three-dimensional plots above suggest that the pre-layer-2 residual stream is highly structured. We therefore asked a more demanding question: how much of the activation can be predicted from a small number of interpretable variables, rather than from an arbitrary 128-dimensional vector? For each token position, we used the normalized within-line phase , clipped to the range so that positions beyond a line boundary were still represented. We expanded into a cubic B-spline basis with 12 knots and fit ridge regressions to the held-out activation vectors. In the first set of fits, the regressors also included the sampled line width, its interaction with the phase basis, the residual from the line boundary, and token position. In the decomposition experiments, we focused on the simpler and more interpretable alternatives: phase alone; phase plus a single learned line-number direction; phase plus a separate learned offset for each line; and a phase/line interaction model.

We evaluated both the geometry and the computation. The PCA-based models first predicted three global PCA coordinates and mapped them back into the full residual space; the full-space models predicted all coordinates directly. We also fit an unconstrained additive model and an orthogonal model in which the phase directions were projected to be orthogonal to the line-number direction. Every model was fit on one set of generated examples and evaluated on held-out examples. For the intervention metrics, we replaced the model’s pre-layer-2 activations with the fitted reconstructions and ran the remaining network. Activation MSE measures reconstruction error, while variance reduction is defined relative to the global-mean baseline as . Newline precision is the fraction of predicted newline tokens that are correct, and newline recall is the fraction of true newline positions recovered.

The results are summarized below. The clean row is the unmodified model, and the random-subspace row averages over 20 random three-dimensional subspaces. The latter provides a useful control: a low-dimensional representation is not automatically useful simply because it has three dimensions.

Approximation	Activation MSE	Variance reduction	Newline precision	Newline recall
Clean model (no replacement)	—	—	0.519	1.000
Global 3D PCA	0.558	88.4%	0.458	0.993
PCA-reconstructed phase-only spline	2.537	47.3%	0.219	1.000
PCA-reconstructed phase + one line vector	1.808	62.4%	0.219	0.951
PCA-reconstructed separate line offsets	1.364	71.6%	0.181	0.503
Full-space unconstrained phase + line	1.685	65.1%	0.235	0.977
Full-space orthogonal phase + line	1.761	63.6%	0.228	0.932
Random 3D subspace

The global PCA reconstruction captures most of the activation variance and preserves newline recall, but it is not behaviorally perfect: precision falls from 0.519 to 0.458. The more structured phase-and-line models recover a substantial fraction of the variance with only a few fitted degrees of freedom, and the additive full-space model preserves newline recall particularly well. At the same time, no single parameterization reproduces both the geometry and the behavior completely. In the phase-plus-one-line model, the learned phase coefficient matrix is itself approximately rank two: its first two singular directions capture about 95% of its variance. This is consistent with a picture of a curved, effectively two-dimensional phase trajectory accompanied by a separate line-number direction, while the random-subspace control suggests that this structure is task-relevant rather than an incidental consequence of projecting into three dimensions.

How Layer 1 Constructs the Character Count Manifold

The next experiment isolates the outputs of the four attention heads in layer 1. For each head, we record its raw hook_z output before the output projection, which lives in the head’s 32-dimensional space, and fit a separate three-dimensional PCA to the trajectory across token positions. The orange trajectory therefore shows how that head’s own output changes over the sequence after compression into three coordinates. Points are colored by cumulative character count, and the interactive plots expose the token position and count on hover.

We then measure the effect of each head on the residual stream by running the same example while replacing that head’s output with zero. The resulting trajectories are shown in a shared three-dimensional PCA basis so that the clean and zero-ablated paths can be compared in the same coordinates. These interventions change only the selected head’s hook_z; the rest of the network and the other heads remain active. The plots show the pre-attention trajectory, the clean post-attention trajectory, and the corresponding zero-ablated trajectory for each head.

Mean-ablation metrics

To measure the behavioral role of each head, notebook 13 replaces its output with a mean vector estimated from calibration sequences, then evaluates newline recall, newline precision, and newline prediction frequency. The three heatmaps report those quantities for every layer and head. This is a mean-ablation intervention, distinct from the zero-ablation trajectory plots above.

Where the character-count representation emerges

Notebook 09 fits held-out linear probes for cumulative character count at several layer-1 sites. The probe is trained separately for each site and reports and mean absolute error (MAE), excluding BOS and newline positions. The comparison tracks how linearly readable the count becomes through the layer.

The head-wise probe values are:

Head	Raw : held-out	Raw : MAE	Residual contribution: held-out	Residual contribution: MAE
0	0.0969	18.9593	0.0965	18.9556
1	0.9752	2.9050	0.9741	2.9576
2	0.1398	18.3749	0.1405	18.3515
3	0.1108	18.7098	0.1115	18.6876

A single count probe is then trained on clean data and applied without refitting after each layer-1 head is zero-ablated:

Condition	Held-out	MAE
Clean post-layer-1 activation	0.9969	0.9902
Head 0 zero-ablated	0.9896	1.9283
Head 1 zero-ablated	−1.2870	29.9623
Head 2 zero-ablated	0.9668	3.7937
Head 3 zero-ablated	0.9282	5.2966

Functional form of the Head 1 output

Notebook 12 tests whether the layer-1 Head 1 output can be represented as an affine function of simple sequence-level variables. The count-only model uses cumulative character count and line width ; the length-histogram model adds the counts of tokens of each character length since the previous newline. A third model also includes position within the sequence. Each model predicts the raw 32-dimensional head output or its residual contribution and is evaluated on held-out ordinary-token positions.

Model	Raw : held-out	Raw : MSE	Residual contribution: held-out	Residual contribution: MSE
Count-only	0.8103	0.0551	0.9226	0.0526
Length histogram	0.8105	0.0550	0.9226	0.0526
Length histogram + position	0.8105	0.0550	0.9226	0.0526

Incremental output and token length

The accumulator test also compares adjacent Head 1 outputs within a line segment. For each ordinary token, it measures the norm of the output increment associated with adding that token and groups the increments by the new token’s character length. The plot shows all held-out increments, their mean and standard deviation at each length, and three simple fits: linear, logarithmic, and square-root dependence on token length.

The held-out values for the linear, logarithmic, and square-root fits are 0.6338, 0.5734, and 0.6176, respectively. The corresponding test of adjacent output increments gives .

Pre-MLP geometry

Finally, notebook 11 examines the post-attention, pre-MLP residual stream, blocks.1.hook_resid_mid. It splits each sequence at newline tokens, keeps newline-free segments with at least four ordinary tokens, and visualizes their activation trajectories in a shared three-dimensional PCA basis. For the line-geometry regression, each segment’s 128-dimensional activations are projected onto their first principal component. We measure the fraction of variance captured by that one-dimensional line, then repeat the fit using alternating token positions as training data and evaluate the reconstruction error on the held-out positions. Across 1,143 usable segments, the mean in-sample is 0.926 and the mean held-out is 0.892.

How Layer 2 Utilizes Character Counts for Prediction

The final-attention experiment examines how Head 3 forms its attention pattern from query and key embeddings. We take the residual-stream activation immediately before the final attention layer, apply that layer’s layer normalization, and project it through and . A joint PCA is fit to all query and key embeddings in the example, so the two sets of vectors appear in the same three-dimensional coordinate system. The normalized causal-attention heatmap shows the actual row-wise softmax weights cached from the model; entries above the diagonal are masked because a query cannot attend to future positions.

The two examples below use the same sampled sequence. At query position 58, the model’s top prediction is _NEWLINE_; at position 59, its top prediction is an ordinary token. For each query, the three plots show the joint 3-D Q/K geometry, the actual attention weights alongside attention reconstructed from the 3-D PCA coordinates, and the raw inner products between the query and key vectors in that PCA space. The reconstruction computes the PCA-space inner products, applies the usual scaling, masks future keys, and applies a row-wise softmax. The attention plots include the current key position as well as the earlier causal prefix.

Layer-2 Head 3 output interventions

Notebook 29 separately intervenes on the values and attention pattern of final-layer Head 3. Value interventions replace the head’s value vectors with their mean, shuffle them across positions, or set them to zero. Pattern interventions replace the attention distribution with attention concentrated on previous newline positions, a uniform distribution over those positions, or a distribution with newline attention removed. The model is then evaluated on newline precision, recall, and prediction frequency.

The attention-predictor analysis also reports the area under the ROC curve for using attention statistics to predict newline targets:

Attention statistic	Target	AUC
Total attention to previous newline tokens	True newline	0.8634
Attention to the most recent previous newline	True newline	0.8757

For the pattern interventions, the held-out newline metrics are:

Intervention	Precision	Recall	Frequency
Baseline	0.5220	0.9995	0.1029
Attention to previous newline	0.4936	0.9995	0.1088
Uniform attention over newline positions	0.5229	0.9209	0.0946
Remove newline attention	0.5165	0.9735	0.1013

The final table summarizes the same interventions on true newline targets, using changes in newline logit and newline margin. The absolute-change columns record the magnitude of the change without regard to sign.

Intervention	NL logit	NL logit	NL margin	NL margin
Baseline	0.000	0.000	0.000	0.000
Mean newline values	0.369	2.533	0.357	2.589
Shuffled newline values	−0.061	2.853	−0.089	2.919
Zeroed newline values	−1.137	2.754	−1.179	2.830
Previous-newline attention	0.552	0.632	0.562	0.646
Uniform newline attention	−2.091	2.256	−2.148	2.322
Removed newline attention	0.103	1.981	0.088	2.025

Conclusions

Further Directions

My Takeaways