Pattern Reconstruction and Interpretation of Structured Motifs in transformers
Department of Neuroscience · Kavli Institute for Neuroscience · Wu Tsai Institute · Yale University
TL;DR. Transformers can model neural recordings across regions, time, and modalities, but standard attribution collapses those axes into one importance map. PRISMt maps a transformer’s classification-token attribution back onto the native axes of the recording and decomposes it into spatial, temporal, and modality-resolved motifs that are faithful to what the model actually uses. In cortical calcium and acetylcholine imaging during learning, it revealed cholinergic motifs reorganizing toward frontal cortex.
Large-scale recordings increasingly combine signals across brain regions, time, and measurement modalities. Transformers can relate all of these in parallel, but interpreting them is a separate problem. Attention rollout describes how information is routed through the layers, and gradient- or relevance-based methods add class sensitivity, yet both characterize model tokens rather than the experimentally meaningful axes of the recording. Decomposing the observations alone finds dominant structure in the data, not the structure used for a particular prediction.
PRISMt is applied to a trained classification-token transformer. It combines attention-mediated routing with the signed value message carried by each token, maps the result back onto sample, region, time, and modality, reduces it across layers and heads, and factorizes it into structured motifs that can be aligned across model fits and experimental conditions. Masked autoencoding learns distributed neural structure; PRISMt identifies the structured signals that drive a specific prediction.
Region × time × modality, per trial, instead of one collapsed importance map.
Non-negative CP factors with separate spatial, temporal, and modality profiles, comparable across conditions.
Distinguishes pathways the classifier uses from correlated decoys present in the data.
A multimodal transformer is pretrained by masked autoencoding under block-causal attention and fine-tuned with a CLS token. PRISMt then turns each prediction into a region × time × modality attribution tensor and decomposes it into motifs.
Each trial is a tensor \(X_n \in \mathbb{R}^{R\times T\times M}\). Every scalar observation becomes a token: it is projected to the model dimension (\(d = 512\)) and combined with a learned token embedding. For the cortical dataset, \(R = 41\) regions, \(T = 10\) time bins, and \(M = 2\) modalities give 820 data tokens per trial.
A learned CLS token provides the sequence-level representation used for classification.
During pretraining, 90% of data tokens are replaced by a learned mask embedding and the model reconstructs them from the remaining context. The loss is evaluated only on masked tokens.
Try it: mask cells, then reconstruct. The estimate here is a toy that uses the information available under block-causal attention: the past of each region and other tokens at the same time.
A token at time \(t\) attends to all regions and modalities at the same time and to every token at earlier times, never to future bins. The same backbone models recordings whose modalities differ in scale, noise, and temporal structure.
For every layer \(\ell\), head \(h\), and token, PRISMt multiplies the residual-aware rollout coefficient \(\rho\) (how much the token is routed to CLS) by the token’s signed value message projected onto the direction of the target logit \(\zeta_n\).
Layer–head pairs are then weighted by how strongly their attribution differs between classes, removing the transformer-specific axes while keeping sample, region, time, and modality.
PRISMt factorizes the magnitude of the reduced attribution tensor with a non-negative canonical polyadic decomposition. Each motif has a sample factor \(a_k(n)\), a spatial profile \(d_k(r)\), a temporal profile \(e_k(t)\), and a modality allocation \(f_k(m)\) constrained to the simplex, so every motif sits on a common modality scale.
Motifs are calcium- or ACh-dominant when one modality weight exceeds 0.65 and shared otherwise, and they are aligned across cross-validated fits and bootstrap resamples before conditions are compared.
All synthetic systems come from the same 64-unit recurrent network trained on a NeuroGym two-alternative perceptual decision task, projected to 32 regions and two modalities, so they differ only in how task-relevant information is arranged. Recovery asks whether a method finds the pathways placed in the data; faithfulness asks whether attribution follows the pathways the classifier uses.
Unit-support AUROC, System 1 (one modality-specific pathway)
System 2 (two separated modality-specific pathways)
System 3 (adds a shared pathway with gain increasing across trials)
Decoy system, primary pathway (observation-level methods: 0.547–0.827)
The decoy test
In System 4 a decoy pathway is correlated with held-out labels but never learned by the classifier. Observation-level decompositions recover both pathways; PRISMt places essentially no attribution on the decoy (causal allocation agreement 1.000, false-attribution mass 0.000).
Operating range
Across a sweep of decoy informativeness and pathway gain, PRISMt outperformed the observation-level baseline in 14 of 20 settings; most failures occurred where occlusion showed the classifier had switched to the decoy.
Generalization
The same backbone modeled spontaneous activity in dual-color calcium/ACh3.0 imaging and in mesoscale calcium imaging combined with extracellular electrophysiology, exceeding a channel-mean floor in both.
Mice expressing the calcium indicator jRCaMP1b and the acetylcholine sensor GRAB-ACh3.0 throughout cortex were imaged through the intact skull while learning a visual detection task, first response timing (Stage 1) and then contrast detection (Stage 2).
The transformer reconstructed masked signals, predicted forward in time, and generalized to held-out mice; behavioral prediction was strongest when calcium and ACh3.0 were modeled together. Applying PRISMt to these models, calcium-associated motifs remained broadly distributed across frontal, somatosensory, and visual cortex, whereas ACh-associated motifs shifted from posterior visual cortex toward frontal cortex across Stage 2 learning, consistently across individual mice.
Two independent objectives agreed: ACh3.0 became more predictable in anterior cortex under both random masking and forward prediction, while jRCaMP1b changed little. Regions selected by PRISMt also became more strongly coordinated in the cholinergic signal after learning.
Distributed motifs; little learning-related change in predictability.
Motifs reorganize toward frontal subnetworks; anterior predictability and connectivity increase.
The model and its interpretation are separate: the transformer learns to predict; PRISMt explains the prediction.
Organize each trial as regions × time × modality in the standardized .mat format, using the MATLAB scripts or the GUI.
% In MATLAB
run_prismt_gui % launch the GUI
% or:
standardize_data('your_data.mat', 'output_standardized.mat')python train.py --data_path data/standardized.mat --task_name phase_classificationpython analyze_results.py \
--checkpoint results/best_model.pt \
--data_path data/standardized.mat \
--output_dir results/analysis/tokens (cortex)41 × 10 × 2 + CLSmodel dim d512mask ratio90%optimizerAdamWweight decay1e-2batch size128learning rate5e-4, cosineepochs≤ 200early stoppingpatience 20grad clip1.0CP restarts5 × 2,000 itermodality threshold0.65If you find PRISMt useful, please cite the preprint:
@article{ortegacaro2025cholinergic,
title = {Selective changes in cortical cholinergic signaling during learning},
author = {Ortega Caro, Josu{\'e} and Batchelor, Hannah M. and Lohani, Sweyta
and van Dijk, David and Cardin, Jessica A.},
journal = {bioRxiv},
year = {2025},
doi = {10.1101/2025.08.29.673096},
publisher = {Cold Spring Harbor Laboratory}
}