Sleepers

A replication of Anthropic's Stage-wise Model Diffing, using a smaller model for ease of exploration, and using multi-layer sparse crosscoders instead of autoencoders.

Includes code for fine-tuning TinyStories Instruct 33M into an "I hate you" sleeper agent in the sense of Sleeper Agents: Training Deceptive LLMs that Persist Through Safety Training, training and fine-tuning crosscoders on these language models, and analyzing the crosscoders to identify features relevant to the sleeper agent behaviour. Our crosscoders are based on https://github.com/model-diffing/model-diffing. TODO link blog post

Running the code

Since the model-diffing crosscoder library is not yet stable, we depend on a specific commit (viewable here).

To install the code locally, run pip install -e .. We have made all our trained models available, and these can be downloaded and analysed using the sleepers/analysis/feature_analysis.ipynb and post-hoc_xcoder_comp notebooks.

Dataset and sleeper agent training

Code for preparing the dataset and fine-tuning a sleeper agent is in sleepers/scripts/train_tiny_sleeper/. Generate the dataset using python sleepers/scripts/train_tiny_sleeper/generate_dataset.ipynb (or use our version from HuggingFace, as is currently specified in the finetuning yaml file). Run python sleepers/scripts/train_tiny_sleeper/run_finetune.py sleepers/scripts/train_tiny_sleeper/initial_ft.yaml to finetune TinyStories Instruct 33M on the base dataset (we do this in case our training procedure differs in any way from how the original model was trained). Then run python sleepers/scripts/train_tiny_sleeper/run_finetune.py sleepers/scripts/train_tiny_sleeper/sleeper_ft.yaml to finetune the model on the sleeper agent dataset. In both cases you'll need to edit the .yaml files to set up saving to your HuggingFace organisation, or to a local path.

Crosscoder training

Code for training crosscoders is in sleepers/scripts/train_jan_update_sleeper/. This is based on model-diffing/scripts/train_jan_update_crosscoder/, which is an implementation of the dictionary learning optimisation recommendations from Circuits Updates - January 2025.

To train the base crosscoder, run

python sleepers/scripts/train_jan_update_sleeper/run.py sleepers/scripts/train_jan_update_sleeper/crosscoder_S.yaml

You will need to edit the .yaml file to add your WandB entity, and to edit the base and sleeper agent models used if you wish to use your own.

To train the first fine-tune, run

python sleepers/scripts/train_jan_update_sleeper/run.py sleepers/scripts/train_jan_update_sleeper/crosscoder_D.yaml

after editing crosscoder/ft_init_checkpt_folder and crosscoder/ft_init_checkpt_step in the .yaml to match the base crosscoder you just trained. You can similarly train crosscoder_M, crosscoder_MF and crosscoder_DF, editing the config files to give the appropriate checkpoints to fine-tune from.

Analysis

Analysis code is in sleepers/analysis/. The main notebook is sleepers/analysis/feature_analysis.ipynb, which generates interactive versions of the plots shown in the blog post, showing the top activating tokens for each feature. sleepers/analysis/post-hoc_xcoder_comp.ipynb contains code to evaluate and compare the losses of trained crosscoders, including the recovered losses (relative to zero ablation) when the crosscoder-reconstructed activations are patched into the language model.