kaushikreddyxyz/toybench-models

ToyBench Models

Trained autoencoders for the ToyBench Distributions benchmark. Each autoencoder compresses a 1,296-dimensional feature distribution into a 200-dimensional embedding space, serving as a toy model of the LLM residual stream.

SAEs are trained to recover features from the embedding space of these autoencoders.

Model Summary

Property	Value
Architecture	TiedLinearRelu (tied encoder/decoder weights + ReLU + bias)
Input features (N)	1,296
Embedding dimensions (d)	200
Parameters	260,496 (W: 200x1296, b: 1296)
Dtype	torch.float32
Training epochs	50,000
Batch size	512
Learning rate	3e-4
Weight decay	0.05
Initialization seed	0
Distributions	8

Autoencoder Architecture

The autoencoder uses tied weights with a ReLU activation:

Encoder: h = W @ f
Decoder: x_hat = ReLU(W^T @ h + b)

Where W ∈ R^{200×1296} and b ∈ R^{1296}. Tying weights forces the decoder to read from where the encoder has written, making directions in embedding space unambiguous.

Models

Distribution	Description
zipfian	Baseline: independent sparse features with power-law activation frequencies
correlated_pairs	Pairwise correlations between jointly-firing feature pairs
hierarchical_pairs	Parent-child pairs with conditional firing and magnitude coupling
deep_hierarchy	DAG-structured features with random-walk-to-root activation
preferential_attachment	Power-law digraph with one-step causal propagation
simplicial_complex	Multi-dimensional features on 7-simplices
spherical	Features on S^4, activated by cosine bumps
toric	Features on T^4 (4-torus), periodic cosine bump activation

File Structure

toybench-models/
  {distribution_name}/
    weights/
      weights.json          # Training config, architecture, parameter shapes
      weights.safetensors   # Model weights (W and b)

Usage

Load with occhio:

from occhio import ToyModel
from occhio.distributions import HuggingFaceDistribution

# Load distribution samples
dist = HuggingFaceDistribution(
    repo_id="kaushikreddyxyz/toybench-distributions",
    filename="correlated_pairs/samples/samples.safetensors",
)

# Load trained autoencoder and create toy model
model = ToyModel.from_pretrained(
    repo_id="kaushikreddyxyz/toybench-models",
    filename="correlated_pairs/weights/weights.safetensors",
    distribution=dist,
)

# Get embeddings for SAE training
embeddings = model.sample_latent(batch_size=1024)  # shape: (1024, 200)

Alternatively, load weights directly:

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

path = hf_hub_download(
    "kaushikreddyxyz/toybench-models",
    "correlated_pairs/weights/weights.safetensors",
)
weights = load_file(path)  # {"W": tensor(200, 1296), "b": tensor(1296)}

Relationship to ToyBench Distributions

Feature samples (N=1296)  →  Autoencoder (this repo)  →  Embeddings (d=200)  →  SAE  →  Recovered features

The autoencoders were trained on the samples in toybench-distributions. SAEs are then trained on the 200-dimensional embedding space to recover the original 1,296 features from superposition.

Citation

@dataset{toybench,
  title={ToyBench Distributions},
  author={Kupper, Niclas and Siewke, Oliver and Reddy, Kaushik and Ayonrinde, Kola},
  year={2026},
  url={https://huggingface.co/datasets/kaushikreddyxyz/toybench-distributions},
  license={CC-BY-4.0}
}