---
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](https://interpretability4discovery.github.io/)
workshop, NeurIPS 2026.
Paper and reviews: <https://openreview.net/forum?id=0KplE7FkqO>

```bibtex
@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:

```python
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:

```python
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).