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