File size: 2,434 Bytes
302675f 86cd27c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | ---
license: mit
tags:
- mechanistic-interpretability
- forward-self-models
- activation-prediction
---
# Forward Self-Model Checkpoints
Checkpoints from [Forward Self-Models Learn an Empirical Approximation of Neural Network Computation](https://jagilley.github.io/forward-self-models.html).
A *forward self-model* is a small auxiliary network trained to predict a neural network's later-layer activations from its earlier-layer activations, learning an empirical approximation of the computational function that the intervening layers implement.
Code: [github.com/jagilley/forward-self-models](https://github.com/jagilley/forward-self-models)
## Checkpoints
### `llama-3.2-1B-layer8-forward-model/`
A 26.2M-parameter forward model (2.1% of Llama) predicting layer 7 → layer 8 of [Llama 3.2 1B](https://huggingface.co/meta-llama/Llama-3.2-1B). Achieves 0.937 cosine similarity with the target activations and 74% KL recovery in causal substitution.
- **Architecture**: 1-layer transformer, 1 attention head, d_head=128, SwiGLU MLP (hidden 4096)
- **Training**: MSE on 100M tokens of frozen Llama activations (FineWeb-Edu), lr=1e-4
- **Files**: `fwd_model.pt` (weights), `config.json` (training config and metrics)
### `toy-gpt-30M-layer1-forward-model/`
A 330K-parameter forward model (~1% of main model) predicting block 0 → block 1 of a 28.9M-parameter GPT-2. Achieves 0.972 cosine similarity and 94% KL recovery.
- **Architecture**: 1-layer transformer, 1 attention head, d_head=64, GELU MLP (hidden 512)
- **Main model**: 4-layer, 4-head, 256-dim GPT-2 trained on FineWeb-Edu (10M tokens)
- **Files**: `fwd_model.pt` (forward model weights), `main_model.pt` (main GPT weights), `config.json` (training config and metrics)
## Loading
```python
import torch
from forward_model import TransformerForwardModel
# Llama forward model
fwd = TransformerForwardModel(
d_model=2048, d_head=128, n_head=1, n_layer=1,
mlp_mult=2.0, block_size=2048, use_swiglu=True,
)
fwd.load_state_dict(torch.load("llama-3.2-1B-layer8-forward-model/fwd_model.pt", map_location="cpu"))
# Toy GPT forward model
fwd_toy = TransformerForwardModel(
d_model=256, d_head=64, n_head=1, n_layer=1,
mlp_mult=2, block_size=128,
)
fwd_toy.load_state_dict(torch.load("toy-gpt-30M-layer1-forward-model/fwd_model.pt", map_location="cpu"))
```
`forward_model.py` and `gpt_model.py` are included in the repo root for convenience.
|