Forward Self-Models

What does a layer do?

A transformer layer takes activations in and produces activations out. Existing tools study this indirectly:

Probing — what information is linearly accessible
Ablation — what is causally necessary
SAEs — what features are present

These characterize what information exists and where it matters — but none directly approximates the transformation itself.

Forward self-models: the idea

Core intuition

Train a small auxiliary network to predict a model's later-layer activations from its earlier-layer activations — learning how the intervening layers transform their inputs.

What the forward model captures

The compressible component of the layer's computation
An executable, causally faithful approximation

What the residual reveals

Computational novelty — the genuinely hard-to-compress part
Tracks complexity, not prediction difficulty

Dissociating representation from computation

Model compression

Must jointly solve:

● Representation (discover good features)
● Computation (learn useful circuits)

→
condition on
activations

Forward self-model

Representation is given. Only needs to learn:

● Computation (transform \(a_i \to a_j\))

⇒ 1–3% of parameters

This dissociation explains both the parameter efficiency and the interpretability of the prediction errors.

Architecture

Loss is MSE between predicted and actual activations, with no gradient to the main model: \(\;\mathcal{L} = \frac{1}{T}\sum_t \lVert \hat{a}_{j,t} - a_{j,t}\rVert^2\)

Setup

Main models

GPT (ours) — 28.9M params
4 layers, 4 heads, \(d{=}256\), FineWeb-Edu 10M tokens

Llama 3.2 1B — 1.24B params
16 layers, 32 heads, \(d{=}2048\), frozen, 100M tokens cached

Forward models

1-layer gap — 330K params (1.1%)
1 layer, 1 head, \(d_{kqv}{=}64\)

3-layer gap — 660K params (2.3%)
2 layers, 1 head, \(d_{kqv}{=}64\)

Llama — 26.2M params (2.1%)
1 layer, 1 head, \(d_{kqv}{=}128\), SwiGLU hidden 4096

All forward models use causal attention, pre-norm, AdamW. Main model weights frozen during forward model training.

A 1%-parameter model captures the layer's function

0.972

cosine similarity

94%

KL recovery

330K

params (~1%)

Causal substitution

Replace block 1 with the forward model's prediction:

Mode	Acc.	KL
Normal	0.222	—
Substituted	0.217	0.062
Ablated	0.146	1.094

Recovers 94% of the layer's contribution. Costs 0.5pp accuracy; ablation costs 7.6pp.

Comparison

A per-position MLP of similar size:

	Cosine
Forward model (transformer)	0.972
Per-position MLP	0.788

The MLP is structurally blind to cross-position effects.

Same function, different weights

The forward model's single head replicates three of block 1's four attention patterns at \(>0.98\) cosine — while all weight cosines \(\approx 0\).

Head	Attn cos	Q cos	V cos
0	0.989	−0.02	−0.01
1	0.998	+0.01	−0.01
2	0.995	−0.03	−0.01
3	0.916	+0.02	+0.01

Head 3: the outlier. 37× higher KL than heads 0–2.

\((QR)(KR)^\top = QK^\top\)

The residual tracks computational complexity, not prediction difficulty

Computation type	Error vs. avg
Large errors (hard to compress)
Delimiter tracking	+0.84σ
After opener	+0.62σ
Small errors (easy to compress)
Sentence start	−0.85σ
Focused attention	−0.61σ
No effect
Easy vs. hard tokens	<0.1σ

σ = std. deviations from mean error

Errors reflect the type of computation (distributed vs. focused attention), not prediction difficulty.

Prediction depth: graceful degradation

Same 1.5%-parameter forward model, predicting across 1–3 layers on a frozen 77M GPT:

A proportionally larger forward model would close this gap (cf. capacity scaling).

Llama 3.2 1B: the key findings replicate

0.937

cosine similarity

74%

KL recovery

26.2M

params (2.1%)

Causal substitution

Mode	Acc.	KL
Normal	0.470	—
Substituted	0.456	0.093
Ablated	0.419	0.353

Recovers 74% of layer 8's contribution. Costs 1.4pp accuracy; ablation costs 5.1pp.

Uniform degradation

Category	KL	Recovery
High confidence	0.067	78%
Low confidence	0.108	72%
Function words	0.089	73%
Punctuation	0.061	75%

Recovery ranges 72–78% across all categories — no behavior-specific failure at either scale.

Per-head decomposition: error maps onto the attention/MLP boundary

	cos	corr
32 heads	0.019	0.07
MLP	0.410	0.577

Attention: 84% of output norm — but the forward model captures it almost perfectly.

Prediction error aligns with the MLP: pointwise nonlinear computation that resists compression.

Recovered via MSE on frozen activations alone.

Why the decomposition falls along this boundary

Attention

Gauge symmetry
\((QR)(KR)^\top = QK^\top\)

One compressed head learns a “principal attention pattern” — compresses 32 heads at 16×.

cos ≈ 0.02 with residual

MLP (SwiGLU)

Pointwise nonlinearity
Hidden 8192 → forward 4096

Each neuron: independent computation. No rotation trick — 2× compression leaves information behind.

cos = 0.41 with residual

What kind of object is a forward self-model?

Probe

What information is accessible?

linear readout

Ablation

What is causally necessary?

counterfactual

Forward self-model

What does the layer do?

executable approximation

The forward self-model approximates the transformation itself — not what information exists, not what is necessary, but how the layer transforms its inputs.

The representation–computation dissociation

Neural networks must simultaneously:

Discover representations
Find a basis encoding the data domain's regularities
Perform computation
Use that basis to produce valid outputs

Capacity limitation is helpful for computation (prevents memorization) but harmful for representation (limits abstraction).

Forward self-models dissociate these:

The main model's activations provide the optimal representation. The forward model only needs to learn the computational function \(a_i \to a_j\).

Consequence

This explains both:

The parameter efficiency (1–3%)
The interpretable residuals

Forward self-models as a primitive

A compact, causally faithful, computation-specific model of a neural network's dynamics.

For mechanistic interpretability

The residual reveals what is genuinely hard to compress
Automatic decomposition along architecturally meaningful boundaries
Complements probes and ablations — a different lens on the same object

For future architectures

An explicit model of a network's own computational dynamics
Predictions about what its own layers will do
Primitive for self-monitoring, metacognition, or adaptive computation

Motivation