prabad/brainlm-sae

BatchTopK sparse autoencoder over BrainLM token activations

The trained sparse dictionary used in the accompanying submission. Anonymised for double-blind review.

It is fitted to final-layer token activations of the BrainLM old_13M encoder (a masked autoencoder pretrained on resting-state fMRI), applied frozen to the Human Connectome Project Young Adult cohort. The paper decomposes the model's own prediction of fluid intelligence over this dictionary in closed form.

Specification

Architecture	BatchTopK sparse autoencoder
Input dimension	512 (BrainLM old_13M final-layer token activation)
Dictionary width	4,096
Sparsity k	32
Alive features	3,781 of 4,096
Training activations	499,000 tokens
Epochs	30
Validation FVU	0.081

Fitted on the training split only — 656 of 936 subjects. The 280-subject holdout was drawn before any interpretability step and the dictionary never saw it. This matters: an SAE fitted on the full cohort has seen the holdout, and any attribution built on it inherits that contact.

Contents

sae.pt is a torch.save dictionary:

Key	Type	Meaning
state_dict	OrderedDict	encoder/decoder weights and b_pre
width	int	4096
k	int	32
mu	np.ndarray (1, 512)	activation mean, subtracted before encoding
scale	float	activation scale, divided before encoding
thresh	float	BatchTopK threshold, frozen over the whole training matrix

thresh is stored deliberately. BatchTopK's cut is otherwise a property of whichever batch encode() was handed, so an SAE re-run on new data in different batches will not reproduce its own training-time activations. Use the stored value rather than recomputing per batch.

Loading

mu is a NumPy array, so on PyTorch 2.6 and later — where torch.load defaults to weights_only=True — the default call raises UnpicklingError. Either:

import torch

ckpt = torch.load("sae.pt", map_location="cpu", weights_only=False)
print(ckpt["width"], ckpt["k"], ckpt["thresh"])

# Activations are normalised with the stored constants before encoding:
#   x_norm = (x - ckpt["mu"]) / ckpt["scale"]

or, to keep the safe loader, allowlist the NumPy reconstructors:

import numpy as np, torch

with torch.serialization.safe_globals(
    [np._core.multiarray._reconstruct, np.ndarray, np.dtype, np.dtypes.Float32DType]
):
    ckpt = torch.load("sae.pt", map_location="cpu")

On PyTorch below 2.6 the plain torch.load("sae.pt", map_location="cpu") works as written.

The model class and the encode path are in src/04_interpretability/train_sae_hcp.py in the code repository below.

Code

Analysis code and the aggregate result tables: https://anonymous.4open.science/r/fmri-model-interpretability-05D7

What is not here

Subject-level activation matrices are not redistributed. They are regenerable from the released code given access to HCP-YA, which is available from the open-access tier under the WU-Minn HCP Open Access Data Use Terms.