You can deploy this model and start earning money today!

jspr/
forward-self-model-checkpoints

mechanistic-interpretability

forward-self-models

activation-prediction

mit

Views

No views yet

Forward Self-Model Checkpoints

Checkpoints from Forward Self-Models Learn an Empirical Approximation of Neural Network Computation.

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

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. 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.

jspr/forward-self-model-checkpoints