license: mit
tags:
- fmri
- neuroscience
- sparse-autoencoder
- interpretability
- brainlm
BatchTopK sparse autoencoder over BrainLM token activations
The trained sparse dictionary used in Interpretability nominates a testable hypothesis from an fMRI foundation model: a proof of concept with sparse-feature attribution (Payman Roghani, Interpretability for Discovery workshop, NeurIPS 2026).
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.
Paper
Payman Roghani. Interpretability nominates a testable hypothesis from an fMRI foundation model: a proof of concept with sparse-feature attribution. Interpretability for Discovery workshop, NeurIPS 2026. Paper and reviews: https://openreview.net/forum?id=0KplE7FkqO
@inproceedings{roghani2026interpretability,
title = {Interpretability nominates a testable hypothesis from an {fMRI}
foundation model: a proof of concept with sparse-feature attribution},
author = {Roghani, Payman},
booktitle = {NeurIPS 2026 Workshop on Interpretability for Discovery},
year = {2026},
url = {https://openreview.net/forum?id=0KplE7FkqO}
}
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://github.com/payman21/brainlm-interpretability
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.
Contact
Payman Roghani, Center for Brain and Cognition, Computational Neuroscience Group, Universitat Pompeu Fabra, Barcelona (proghani@gmail.com).