Multi-session, multi-task neural decoding from distinct cell-types and brain regions
Authors: Mehdi Azabou
Krystal Xuejing Pan
Vinam Arora
Ian Jarratt Knight
Eva L Dyer
Blake Aaron Richards
Venue: ICLR 2025 Spotlight (2025)
Keywords: neural population
multi-task
transformer
tokenization
two-photon calcium imaging
visual stimuli
cell types
Primary area: applications to neuroscience & cognitive science
TL;DR: This paper introduces a multi-session, multi-task framework for building large-scale, interpretable models of calcium activity recorded across visual areas and distinct cell-types.

Abstract: Recent work has shown that scale is important for improved brain decoding, with more data leading to greater decoding accuracy. However, large-scale decoding across many different datasets is challenging because neural circuits are heterogeneous---each brain region contains a unique mix of cellular sub-types, and the responses to different stimuli are diverse across regions and sub-types. It is unknown whether it is possible to pre-train and transfer brain decoding models between distinct tasks, cellular sub-types, and brain regions. To address these questions, we developed a multi-task transformer architecture and trained it on the entirety of the Allen Institute's Brain Observatory dataset. This dataset contains responses from over 100,000 neurons in 6 areas of the brains of mice, observed with two-photon calcium imaging, recorded while the mice observed different types of visual stimuli. Our results demonstrate that transfer is indeed possible -combining data from different sources is beneficial for a number of downstream decoding tasks. As well, we can transfer the model between regions and sub-types, demonstrating that there is in fact common information in diverse circuits that can be extracted by an appropriately designed model. Interestingly, we found that the model's latent representations showed clear distinctions between different brain regions and cellular sub-types, even though it was never given any information about these distinctions. Altogether, our work demonstrates that training a large-scale neural decoding model on diverse data is possible, and this provides a means of studying the differences and similarities between heterogeneous neural circuits.

Published as a conference paper at ICLR 2025
MULTI-SESSION, MULTI-TASK NEURAL DECODING FROM DISTINCT CELL-TYPES AND BRAIN REGIONS
Mehdi Azabou
Krystal Xuejing Pan" Mila
Vinam Arora
Georgia Tech
Georgia Tech
Ian Knight Georgia Tech
Eva L. Dyer
Blake Richards Mila, McGill University
Georgia Tech
ABSTRACT
Recent work has shown that scale is important for improved brain decoding, with more data leading to greater decoding accuracy. However, large-scale de- coding across many different datasets is challenging because neural circuits are heterogeneous-eecch brain region contains a unique mix of cellular sub-types, and the responses to different stimuli are diverse across regions and sub-types. It is unknown whether is possible to pre-train and transfer brain decoding mod- els between distinct tasks, cellular sub-types, and brain regions. To address these questions, we developed a multi-task transformer architecture and trained it on the entirety of the Allen Institute's 's Brain Observatory dataset. This dataset contains responses from over 100,000 neurons in areas of the brains of mice, observed with two-photon calcium imaging, recorded while the mice observed different types of visual stimuli. Our results demonstrate that transfer is indeed possible- combining data from different sources is beneficial for a number of downstream decoding tasks. As well, we can transfer the model between regions and sub- types, demonstrating that there is in fact common information in diverse circuits that can be extracted by an appropriately designed model. Interestingly, we found that the model's latent representations showed clear distinctions between different brain regions and cellular sub-types, even though it was never given any informa- tion about these distinetions. Altogether, our work demonstrates that training a large-scale neural decoding model on diverse data is possible, and this provides means of studying the differences and similarities between heterogencous neural circuits.
INTRODUCTION
All behaviors, perceptions, and thoughts arise from the activity of neural circuits in the brain. In principle, decoding this neural activity could allow us to infer what someone is doing, perceiving. or thinking. Effective decoding of neural activity holds the potential to unlock transformative ap- plications, such as brain-computer interfaces for thought-comttooleed systems and advanced clinical
interventions like targeted deep brain stimulation (Saha et al. 2021 Merk et al. 2022).
Recent advancements in transformer architectures (Radford et al., 2018) have introduced a new paradigm in brain decoding (Azabou et al., 2024 Ye et al. 2024] Zhang al., 2024). These mod- els demonstrate that increasing the scale of neural data, by training on many recording sessions and across multiple animals, leads to significant improvements in decoding performance. This trend mir- rors the success of large-language models, where pretraining on large corpora has resulted in models capable of generalizing across wide range of tasks. In the case of brain decoding, larger datasets allow models to capture more nuanced patterns of neural activity, enabling better generalization to new behaviors (Azabou et al. 2024 Tiang et al., 2024). This scaling effect is especially evident when small, previously uninformative datasets are paired with pretrained models, unlocking new decoding abilities and overcoming limitations imposed by insufficient data (Azabou al. 2024).
"These authors contributed equally to this work These senior authors contributed equally this work. Contact: mazahou @aaouehaaattttuuuuuo Project page and code: https::/pooo-plll -urooyy-pppuuuu... www.plgisttii.... github. 10

---

Published as a conference paper at ICLR 2025
However, large-scale training for brain decoding is challenging due to the diversity of neural record- ings. Though there is a lot of neural data available in aggregate, it is mostly small, highly hetero- geneous individual datasets [Ferguson et al. 2014). Each individual dataset can contain recordings that involve distinct cellular sub-types and brain regions- as well as different stimuli or behaviors (Estahany al. 2018: Dipoppa et al., 2018 Scala al. (021). This diversity in data complicates efforts to train large-scale models on large corpora. Prior work has demonstrated that scaling data within neural populations that are relatively similar across sessions can improve decoding accuracy Azzaoou et al., 2024 Ye et al. 20244] Zhang et al. 2024), but it is unknown whether similar improve- ments can be achieved when training models on heterogeneous datasets involving different cellular sub-types, brain regions, and tasks. Can neural decoding models be trained on multiple decoding tasks at once? Can neural decoding models trained on one cellular sub-type or brain region trans- fer to another? Can pretraining on diverse data yield models that generalize across distinct neural circuits, or does this diversity introduce too much noise for effective learning?
To address these open questions, we developed multi-task, multi-modal transformer architecture and trained it on the Allen Institute for Brain Science Observatory dataset. This dataset includes recordings from over 110,000 units recorded from mice using two-photon calcium imaging while they observed various types of visual stimuli (de Vries et al., 2020). These recordings span 256 animals, 6 brain regions, and 13 genetically defined cellular sub-types, making it one of the most diverse neural data archives available. By leveraging this large and heterogeneous corpus of data, we aimed to test whether transfer possible across distinct brain regions and cellular sub-types, and whether pretraining on diverse data can improve decoding performance for a variety of tasks.
Our results demonstrate that training across brain regions and cellular sub-types-desple functional differences- not only feasible but also beneficial. We show that combining data from heterogeneous recordings leads to improved performance on downstream decoding tasks, indicating that there is shared information across diverse neural circuits that can be captured by an appropri- ately designed model. Notably, even in cases where single-session baselines fail to decode stimulus information, pretraining on diverse datasets significantly enhances decoding capabilities. This find- ing suggests that neural circuits, despite their heterogeneity, share common computational principles that can be harnessed by large-scale model for decoding.
their
We then examined our model's latent embeddings and discovered that they revealed meaningful structure. Even though the model was not trained to differentiate brain regions cellular sub- types (and received no information about this), its latent space naturally organized neurons based on their anatomical and physiological similarities. When we clustered the latent representations, we observed that cellular sub-types and brain regions with more similar anatomy and physiology were more closely aligned in the latent space. This emergent organization provides new insights into the relationships between different neural circuits and suggests that large-scale models can capture higher-order patterns of neural activity that reflect the underlying functional architecture of the brain.
In summary, the main contributions of this work include:
Multi- Malti-TaskKaae and Multi-Modal Decoder for Flexible Querying: We introduce a novel multi-task decoder that can query the learned latent space to engage in a wide range of decoding tasks with high fidelity. This decoder is capable of handling both continuous and diserete decoding tasks, demonstrating versatility across different neural recordings and experimental setups.
Cross-Regiin and Cross-Cell-Type Transfer: Our model shows successful transfer learn- ing across brain regions and genetically distinct cellular sub-types. This cross-area and cell line transfer demonstrates that the model can generalize representations beyond specific regions and sub-types, capturing shared computational principles across diverse circuits.
Across-Region Scaling: We provide evidence that increasing the diversity of brain regions in the training data leads to improved decoding performance. As more regions are included in training, the model's ability to decode and generalize across previously unseen regions significantly improves.
Meaningful Latent Representations: We show that the model's 's latent representations of the data carry meaningful information that reflect the anatomy and physiology of different regions and sub-types. This illustrates that it is possible for large-scate scale decoding models to identify biologically meaningful distinctions from neural activity alone.
2

---

Published as a conference paper at ICLR 2025
Caldium Traces
POYO+
Output tokens
POYO Encoder
Multi-modal Decoder
2Hz
Query tokens
DC Orientation
DS Tamproaal Beqpancy
www.
Task embedding Session embedding: Output timestamp
Pupil location
Running speed
Figure 1: Our approach to diverse multi-region, multi-task neural decoding. Our approach leverages task embedding query the neural data latents decode different contextual behavioral information from the neural activity.
METHODS
In this section, we describe POYO+, our multi-task, multi-modal model for neural decoding (illus- trated in Figure T.. The main innovation of our approach is the development of a multi-task decoder that allows for flexible training with different stimuli and behavioral measures as decoded outputs, across varying contexts, all through set of learned task embeddings.
2.1 TOKENIZATION
The calcium recordings consist of multiple time-series, one for each neuron that was identified in the calcium imaging session. The values the time-series correspond to the change in fluorescence over baseline value (i.e. AF/F), which reflects the in action potential firing of neuron. Let fij. denote the scalar value of the original fluorescence signal, AF F, at time-step i for neuron j.
To tokenize the activity of population of P neurons, adapt the POYO tokenizer (Azabou al. 2024) to regular timeseries data by creating a token for each timestep and for each neuron, i.c. over T time-stess
create sequence of P T tokens (Xij, fay) where E [1, [1,..., T). [1, [1,..., P]- These tokens carry information about the neuron's 's activity, which in turn can be used to learn about it other properties-e.. cell-type, cortical layer, and the brain-region Specifically, each token, is tuple with two components: (Xij, lig), where X,y [u, W is an embedding consisting of information about the unit and signal amplitude determined by learnable parameters W R.V. u, is N d-dienniional, neuron-level learned, embedding that is constant across sime-steps. Each neuron receives its own learned embedding, which is accessed via a look-up table. Importantly, uj does not contain any information about the brain region the recording was done in, nor the cellular sub-type of the neuron, prior to learning. liy a 2N -imeennsiaaal sinusoidal time-embeddingg which encodes the timing of the token relative to the start of the context window. In this way, the activity of any population of neurons can be represented flexibly through a single sequence of tokens, which is passed as input to the transformer model.
which it belongs.
2.2 POYO ENCODER
To build rich representations of neural activity, we use the POYO encoder (Azabou et al.) 2024) which employs a cross-attention layer to compress neuron-level spike tokens into a lower- dimensional latent space. In our case, we will learn set of M x N latent token tuples ((z)) Tman), m E [I,..., MJ. n E [1, [1,..., .N) that will used to generate queries for the neural data. Here, each 2 is learned latent vector, and Tm.n is sinusoidal rotary embedding. The index is used to indicate different latent tokens, and the index n is used to indicate the different time-steps across the context window, which are evenly spaced. The cross-attention operation is given by:
P
z
Σ
Σ
softmax (R (Tmm (armnngmmm (mm) IT (R(tz))k.))) )
1=1
X are computed from the neural activity tokens, and the queries qmn W:0 Zfm computed from the learned latent tokens, and R(t) is predefined rotation matrix whose values are determined by the time-step given (see Appendix C.T). Rotating keys and queries (Su et al. 2021) allows us to incorporate information about the relative
where the values Viz W Xij and keys k.j
3

---

Published as a conference paper at ICLR 2025
timing between two tokens into the attention mechanism. As mentioned above, our latent tokens positioned uniformly across the context window to provide "virtual timesteps" for each query.
After cross-attention, then apply L rounds of self-attention to compute sequence of final latents z Throughout, we use rotary time encoding with standard transformer blocks, pre-normalization layers and feed-forward nets.
2.3 MULTI-TASK, MULTI-MODAI. DECODER
Neural datasets often have multiple sources of behavioral information that are simultaneously ac- quired along with brain activity (i.c., pupil diameter, running speed) while different stimuli are presented to the animal. This diversity in potential decoding targets must be harnessed as we scale and build models that can generalize diverse datasets.
To accommodate this type of diversity, we developed a model that can be trained jointly on different measures of behavior as well as information about the stimuli presented context of the recording (Figure[T) In order to have a generalist model that can work with these different types of labels, we built a flexible decoder that uses a shared cross-attention layer to query from the latent space using different task-specific tokens. This allows us to train on heterogeneous neural datasets and decode any number of different variables without assuming the same behaviors stimuli are present in the recordings. Indeed, we can combine both dense, continuous regression decoding tasks (like pupil diameter prediction or running speed prediction), sequence-level classification tasks (like drifting gratings orientation classification or static gratings spatial frequency classification), and segmenta- tion tasks (like movie frame prediction or locally sparse noise pattern prediction).
Our multi-modal decoder uses shared cross-attention layer to query the latent space for labels relevant to the different downstream tasks. An output query token for decoding task k. in recording session l, at time-step t, is a tuple (Ork. .t). The component Ozk Φb Sy is defined by two quantities: (i) Task embedding: Φx E IR is a learned embedding corresponding to the decoding task of interest (e.g. running speed regression
drifting grating orientation classification, etc.). As with neurons, each task receives its own learned embedding that is accessed via a look-up table; (ii) Session embedding: S; E IE learned embedding corresponding to the recording session. Again, each recording session receives its own learned embedding accessed via a look-up table. Session embeddings are used to capture variability across sessions in the datasets.
We calculate the final output tokens for the ith time-step. Yeu. using a cross-attention mechanism:
N
Yik: = OEk Σ
softmax (R(t,)qek) (R(mnn) kmm) ) )
where we have queries qı Or. values V.n. = W. and keys k.j- W calculated from our final set of tokens produced by self attention. As before, R(t) is a matrix determined by a rotary positional embedding. After cross-attention, we apply a task-specific linear layer WR R." > RDx where D is the dimension of each output tokens year. and Dr is the dimension of the target for task k. For regression tasks, Dr is the dimensionality of the variable being decoded, for example, D2 = 1 for rumning speed and D2 2 for pupil location. For classification and segmentation tasks, D& will be the number of classes.
The entire network is trained end-to-end on all tasks simultaneously. We use mean-squared error loss for regression tasks, and cross-entropy loss for classification and segmentation tasks. The losses are summed and weighted based on the number of predictions for each task (i.e. the more predictions requested the lower the weight for any specific token).
With this set-up, we are able to decode multiple labels simultaneously during both training and inference. For example, a context window of one second where the stimulus presentation drifting gratings, we can decode on four different tasks: running speed regression, pupil location regression, orientation classification, and temporal frequency classification. Because running speed and pupil diameter are monitored at a frequency of 30 Hz we will have 30 outputs for each of these tasks over a s context window. For the other two tasks, we will have only one output, since these values constant throughout the context for drifting gratings. Thus, we will have total of 62 query tokens and outputs. But, notably, our molti decoder is able to handle any number of query tokens, from different tasks due its flexible design.
4

---

Published as a conference paper at ICLR 2025
3 RESULTS
3.1 TRAINING A MODEL ON DIVERSE MIXTURES OF NEURAL RECORDINGS
Datasets. To train our model, we leverage the Allen Brain Observatory (BO) calcium dataset (de Vries et al., 2020). The Allen BO is the largest collection of in vivo two- photon calcium imaging data, acquired from
30000
Exc. Inh.
Mial
25000
v155m
V15ri
20000
the visual cortex of awake mice. It con-
Tore Numbter
tains neuronal activity in six different corti- cal areas: Primary Visual Area (VISp), Pos- teromedial visual area (VISpm), Anterome- dial visual area (VISam), Rostrolateral vi- sual area (VISrl), Anteromedial visual area (VISal), Lateral visual area (VISI), in re- sponse to various visual stimuli, including both simple drifting gratings (eight different orientations), static gratings, natural scenes, and naturalistic movies (Appendix [B.. Each experiment consists of three one-hour long imaging sessions, one for each mouse, imag- ing depth, and visual area, recorded at various depths, ranging from 150 to 600 um from the surface of the brain, each of which can be classified into different layers (L22/3, L4, L5, or L6).
15000
10000
5000
Emc
Post
Nesal
Word
rapa
Rocb
Scrmza
sic17a7
md
Pual
Cheche
Figure 2: Data diversity in the Allen Brain Observa- tory. breakdown of the number of neurons for each Cre-line by brain region. The Cre-lines are organized by Excitatory and Inhibitory lines from left to right, and the visual areas are depicted with different colors.
The mice expressed genetically encoded calcium sensor (GCaMP6f) in specific genetic sub-types of neurons using Cre driver lines. There were two pan-excitatory lines where the calcium sensor was expressed in almost all excitatory neurons (Fmx1, SIc17a7), eight lines where it expressed only in specific sub-types of excitatory neurons (Cux2, Rorb, Scnnla, Nr5al, Rbp4, Fezf2, Tix3, Ntsr!) and three lines where it was expressed only in specific sub-types of inhibitory neurons (Vip, SST, PV). Each of these cellular sub-types show distinct response properties (de Vries et al. 2020). and very different densities of neurons in different brain regions (Figure 2 Our challenge here was to see if we could incorporate data from these diverse sub-types to improve decoding.
Experimental setup. Each recording session is divided into training. validation and test splits with 60/10/30 ratio. Given the unstructured nature of the recording, we tried to maintain these ratios across each stimulus epoch. When reporting accuracies, we perform evaluation on the test split of each session, and report the average across sessions. Note that there is high variability in the per- session decoding accuracy, likely because we are working with heterogeneous data, with half of the sessions having less than 62 neurons and some as few as 5. We note that in total we use 1335 3-hour sessions for training our model. During training, we leverage neuron dropout where random subset of the population is removed from the context window.
For the purposes of transfer, we holdout sessions from 30 animals that are never seen during training. We use the training split of each heldout session finetune POYo+, then report the test accuracy. We use gradual unfreezing where we first unfreeze the unit and session embeddings then the multi- modal decoder then the encoder. We find that in practice, the model already converges with just the unit and session embeddings unfrozen. More details can be found in Appendix C.2
The benefits of multi-session training. Given the inherent diversity of dataset, we first inves- tigated the scaling properties of the model when trained on specific context. For this experiment, we used the drifting gratings stimulus condition and trained a variant of POYO that we designed to work with our new tokenization strategy for regular time series data as input. The model was trained on all sessions that include drifting gratings, we call this model Poyo (DG) (DG for Drifting Grat- ings). For comparison, we trained two groups of single-session baseline models (i) MLP and (ii) TCN-on each of the 403 sessions collected under the DG condition. Each model was extensively tuned using validation set, we compared total of more than 80,000 models that were trained during hyperparameter tuning. When compared to these single-session baselines, we observed sig- nificant improvements on test trials, with the multi-session model providing nearly 10% perfor- mance gain on average relative to the MLP baseline (Figure (3)). We also compared POYO+ to a contrastive multi-session approach for training that uses labeled data to create positive views across
5

---

Published as a conference paper at ICLR 2025
B Drfting Oraling Orieratarr
Sporial
1.0
POYO
50.02
POTO (05) accurocy
POYD (DG)
48.96
focel
TCN
25
5 ROss
100 ROls
38.30
0.0
0.2
0.4
0.6
0.8
1.0
MLP Accuracy
Accurscy
Figure 3: The benefits of a multi-session, multi-task neural decoder. A. PoYo (DG) compared to an MLP on in-distribution datasets. B. PoYo+ and PoYo (DG) compared to MLP and TCN baselines on the drifting gratings orientation prediction task. C. Comparing task-wise variants of POYO PoYo+. POYO+ single model trained on tasks simultaneously, while we have 12 variants of POYO.
animals (CEBRA, (Schneider et al. 2023b), see Appendix D.2 for more details). We found that our approach again provides superior decoding performance for multi-session data. In conclusion, we find that our single-task variant of POYO consistently beats different baselines (Figure B)). This highlights the advantages of aggregating information across multiple sessions for this decoding task.
Incorporating multi-task learning across diverse contexts. After confirming the benefits of multi- session training the drifting gratings condition, we extended our model to include additional stimulus contexts well behavioral decoding tasks, such pupil diameter and running speed. In Figure BC, we compare POYO+ to the single-task POYO variants each trained on one of the 12 behavioral and contextual decoding tasks. We find that incorporating multiple tasks into the output space of the decoder improves performance on a number of tasks, especially the behavioral decoding tasks. We hypothesize that multi-task training provides dense prediction tasks (e.g., running speed and pupil location) provides auxiliary information to the decoder, improving overall predictions. Second, it engages a broader population of neurons in the model, as some neurons responded more strongly to behavioral changes than to specific stimulus contexts. This suggests that multi-tase learning not only enhances decoding performance but also better captures the rich, context- context-dependett modulation of neural activity.
key benefits. First, the inclusion of
3.2 TRANSFER ACROSS BRAIN REGIONS AND CRE-LINES
Across-region transfer. The visual cortex exhibits distinct activity patterns across its regions, rais- ing the question of whether models trained in one area can transfer effectively to others. Our re- sults show that many regions provide solid basis for transfer, with VISal and VISI delivering the strongest overall performance. Specifically, the POYO+ model trained on VISal achieved the high- est accuracy in VISal itself (73.84%) and demonstrated good transfer to VISam (52.19%), VISI (59.78%), and VISp (52. 14%). Similarly, the POYO+ model trained on VISI also performed well, achieving high transfer accuracy in VISal (72.22%) and VISpm (35.34%). Interestingly, VISp, primary visual area, exhibited strong transfer to VISal (74.01 %) and VISI (5%.7 suggesting its foundational role in cross-region transfer.
Table Transfer across held-out brain areas. The results from training on one region and finetuning on new held-out sessions in another region, are compared with Poyo+ trained on all brain areas. All of the across-line results that perform better than the base within-line model are highlighted in green. The best performing model for each testing condition in boldface and the number of sessions in each test condition are parentheses.
Training sessions
Average (30)
VISal
VISam (2)
VISI
VISp
VISpm (3)
Model MLP
(4)
(8)
(13)
43.36%
56.39%
35 68.
45.65%
42.17%
30.13%
PoYo+ (VISal)
119
54.12%
73,84%
52.19%
59.78%
52.14%
36.37%
POYO+ (VISam)
112
47.00%
63.20%
47.54%
49.58%
45.67%
28.15%
Poyo+ (VISI)
310
53.98%
72.22%
51.17%
58.01%
52.52%
35.34%
Payo+ (VTSp)
419
53.67%
74.01 %
48.89 48.89%
58.74%
52.15%
31.73%
Poyo+ (VISpm)
258
51.94%
69, 10%
51.83%
56.88%
50.72%
30.41%
Poyo+ (All)
1335
55.96%
73.96%
49.39%
61.46%
54.84%
36.01%

---

Published as a conference paper at ICLR 2025
A
D
Brain Area decodding accuracy
59.81%%
175
150
V15am
021
ViJ
581
PM
100
39.01%
36.74%
75
Vi5cm
810
613
0M
50
Vism
Baseline
POYO-
ViSal
VSp Vi5pm Visr!
-190
100
600
67.93%
Cie Cie-lire decoding accuracy
Eusiidea Distance
56.00%
2-33
200
0.00
$43.00
sp
(127) oot
End un tok anto cot
scciat
190.05
0.5
Scon 200oo
05
Tw305
Guz 040
Emit Eme
96727 ra
Voo
Worl 0.)
Rb54 (13)
Mov1 (.3)
Nes1 (13))
Baseline POVO m4 POYD+
heft
Emd
wathl
mode
Figure 4: Examining the organization of the latent embeddings by Cre-line and brain region. A.
Balanced accuracy for brain area classifcation based on hand-crafted features versus session-averaged latents from POYo+ B. Confusion matrix for brain area classification C. Dendrogram of average latents by brain area D. t-SNE projection of PoYO+ embeddings. E. Balanced accuracy for Cre-line classification based hand- crafted features session-averaged latents from POYo+ Confusion matrix for Cre-line classification G. Dendrogram of average latents by Cre-line.
These results highlight that training on higher-order areas like VISal and VISI facilitates superior transfer other regions, likely due their integrative roles in visual processing. Furthermore, the model trained on all regions, POYO+ (All), achieved the highest overall performance (55 96%)), demonstrating that exposure data from multiple brain areas improves generalization across unseen regions. The finding that transfer performance improves with more diverse training data emphasizes the importance of using data from different brain areas and animals, which enables the model to capture wider range of neural dynamics and improves its adaptability across regions.
Across-Cre-ling transfer. Next, we evaluated whether our mod- els could transfer effectively across excitatory and inhibitory Cre-lines. While our model trained on all lines surpasses the other Cre-line specific models,
Table Transfer across held-out Cre-lines. The decoding accuracy for the drifting gratings task for held-out animals for models trained on Excitatory lines only, Inhibitory lines only, and All lines. The performance is measured (16) Excitatory ses- sions and Inhibitory sessions spanning different lines. All of within-
across-region results that perform better than region model (purple) are highlighted in green. The best perform- ing model for each testing condition is in boldface.
find that the model trained on inhibitory Cre-lines does surprisingly well on the drift- ing gratings task. Specifically, the model trained inhibitory data, Poyo+ (Inh), achieved the best per- formance on inhibitory data and also surpasses the model trained on exci- tatory data alone.
Training sessions
Average (30)
Excitatory
Inhibitory
Model
(16)
(8)
MI.P (Base)
43.36%
48.47%
26.61%
PоYo+ (Exc)
987
54.74%
55.40%
30.01%
PoYo+ (Inh)
55.85%
62.89%
31.96%
PoYo+ (All)
1335
55.96%
62.93%
31.80%
3.3 ANALYSIS OF LATENT EMBEDDINGS
Although our model has been trained to decode the stimulus information or behavior, we wanted to understand the extent to which the model has leamned session-level differences that capture bio- logically meaningful information about the Cre-line and brain regions for each session. To test this, examined the latents at the output of the encoder in Poyo+ To obtain the session-level latent

---

Published as a conference paper at ICLR 2025
Vip [14)
Set (1.5)
Vip (1.23) Pvalb (1.2/3) Pualb (L.5)
Pvalb (L44
Set (L4)
Ni5al (14)
NrSal (L.))
Rbp4 (1.5
Nisrl (L.6) Sconla (14)
Sp Nit ni Sp. so Dis Sp. NM DG Sp. DG Sp. MM Sp. nis Sp. Nit so fip ti MM so DC Sp. NM SC NM so Sp.
Rorb (L4) Emx1 (L5) Sic17a7 (LS) Sic17a7kuni Scnnla (LS) Cux2 (L4) Cux2 (L2/3) Feaf2 (LS) Emx1 (L2/3)5ic17a7 (L4) Email (L4) Tixs (LS)
5p. N. Sp. 50 DG Sp. NM Sp. DG Sp. NN Sp. 50 00 Sp. No Sp. DG 5G Sp. 00.5p. 50 Sp.
Figure 5: Deviation of Cre-line averaged embeddings from baseline when considering one visual stim- ulus type at a time. We use I-dimensional PCA projection and plot the average latent embeddings for the Spontaneous (Sp), Natural Movies (NM), Static Gratings (SG), and Drifting Gratings (DG) for each Cre-line separated by layer (L2/3, L4, L5, L6).
embeddings, we average the latents across randomly sampled 1 context windows in the recording. See Appendix [E] for further details on this analysis.
Latent representations of distinet brain regions. First, we examined the latent embeddings for sessions from different brain regions. Using a linear decoder trained on the latent embeddings for different sessions, we were able to decode the brain region of the recording session with high accuracy (Figure [A). much higher than using standard POYO latents or a bascline using hand- crafted features used for cell-type analysis (see Appendix [2.2 for more details). Examination of the confusion matrix for decoding from the POYO+ latents showed that some regions were more likely to be confused with each other by the decoder, for example VISal was likely to be confused for VISI (Figure 4B). This suggests that the underlying latents were more or less similar for different groups of brain regions. To examine this, we conducted hierarchical clustering on the latents for sessions from different brain regions. Interestingly, we found that the hierarchical clustering distinguished medial regions clearly from lateral ones, and the lateral region which was most closely clustered to the medial regions (VISrl) is in fact the most medial of the lateral regions (Figure EC). This demonstrates that the learned latent embeddings reflect the actual anatomy of visual cortex, despite the model having received no information about this anatomy.
Latent representations of cellular sub-types. We then conducted similar analyses on the ses- sions for different cellular sub-types. First, we examined the projection of the latents to lower- dimensional space, using t-SNE. Interestingly, we observed clear separability for session latents from inhibitory and excitatory sub-types (Figure 4D), suggesting that the latents had Icarned to cap- ture cellular sub-type distinctions. Next, to quantify the separability of different Cre-lines we again trained a linear decoder on top of the extracted latents, this time to decode the Cre-lines. We find that the organization of Cre-line is quite impressive, with a overall 67.93% accuracy in decoding, compared 56% and 39.57% for transcriptomic baseline (Schneider et al.) (2023a) and standard POYO, respectively (Figure 4 b). These results suggest that to really understand the neural activity across different cellular sub-types it is crucial to consider multi-task view.
We also then examined the confusion matrix for the Cre-lines. As with brain regions, we observed that some cellular sub-types are more or less likely to be confused with each other, e.g. excitatory neurons are far more likely to be confused with other excitatory neurons than inhibitory neurons (Figure F To understand the structure in the latents better, we again conducted hierarchical cluser- ing for the latents from different sessions based on the Cre-lines (Figure 3T.. We observed several interesting results. First, as expected given the low-dimensional embeddings, the inhibitory neurons were more closely grouped together, and the clusters within the inhibitory types reflected true bi- ological distinctions (c.g. the different types of Pvalb sessions were clustered together). Second, the clustering had some relationship to the laminar position of the cellular sub-types, such that cells in the supragranular layers (L2/3) tended to cluster with each other or sub-types from L4. Finally, observed that the clusters appeared respect the axonal projections of the different cellular sub-types. For example, intratelencephatic (IT) neurons tended to be closer to each other in the hierarchy than to pyramidal tract (PT) neurons. The fact that the latent embeddings for the different sessions reflected these real biological differences in physiology and anatomy was surprising but encouraging, and suggests that our model can be used reveal meaningful biological relationships.
8

---

Published as a conference paper at ICLR 2025
Visualization of task-dependent modulation of embeddings. Finally, in examining the latent embeddings for the different sessions we wanted to see whether they reflected stimulus-specific dis- tinctions between cellular sub-types. In Figure 5 aggregate the latent embeddings within each stimulus context, and examine them across Cre-lines. We compute the first principal component (PC) and project the latents from each task onto the PC, and report the multi-task aggregated em- bedding as the dotted baseline. We are able to see the main source of variation in the latents and how maps onto the different stimulus contexts and Cre-lines.
First, we note outliers like Sst, an interneuron sub-type that behaves similarly to other interneurons (Vip. Pvalb) during all contexts except for the drifting gratings context, where it deviates from the baselines in the opposite direction from other interneurons. This makes sense based upon prior work which shows that VIP are suppressed during drifting gratings (de Vries et al., 2020). We also observe subtle differences across different layers in the same Cre-line, like Pvalb neurons in layer maintaining their baseline dynamies during drifting gratings, while we observe a shift in Pvalb neurons in L 2/3 and 4. In excitatory neurons, see overall similar modulation across stimuli, but an interesting stability in Nr5al neurons which appear to be consistent regardless the stimulus. These data demonstrate another advantage of multi-task model trained on heterogeneous data: it allows us to compare distinct cellular sub-types in a stimulus/context dependent manner.
3.4 TRANSFER TO OPENSCOPE DATASET
In this experiment, tested whether could transfer from our pretrained model to another open dataset with different recording conditions and stimuli, collected as part of The Allen Insti- tute Mindscope's 's OpenScope program (Gillon iet al.) 2023) (Table 3) This dataset consists of responses of neurons in VisP from two Cre driver lines targeting the soma and dendrites of distinct excitatory populations (Cux2, somatic; Rbp4, dendritic). When we finetune pretrained multi-task model on sessions from this new dataset, we find a near 9% improvement with POYO+ over the MLP baseline, and a 2% improve- ment over the DG-only pretrained model. If we break these numbers down by the type of recording, we find that the improvements on Dendritic sessions benefit the most, with a 10% improvement with PoYo+ despite the fact that did not pretrain dendritic data which varies considerably from somatic data (see Appendix D.I.
Table 3: Transfer from POYO+ to Openscope
dataset. We compare the decoding performance on a 3-class task for an MLP, single task POYO (DG) and PoYo+ both finetuned for the new task.
Avg (5)
Dend (3)
Soma (2)
MLP
72.94%
73.51%
72.10%
POYO (DG)
79.70%
81.74%
76.76%
POYO+
81.91 %
84.19%
78.49%
3.5 TRANSFER TO HIPPOCAMPUS DATASET
To understand how our pretrained model transfer to non-visual areas, we obtained a calcium imaging dataset from the hippocampus
Table 4: Transfer from POYO+ trained on visual areas to data recorded in hip- pocampus. We finctune PoYo+ to perdict the position of the mouse in the maze. We report the mean euclidean error.
mice explored differ- ently shaped 75 75 em open mazes Lee et al. (2025) (10 sessions). We used this data decode the animal's position. We compared POYO+ pretrained on the Allen Observatory data and one trained from scratch. Impor- tantly, for the pretrained model, we froze all weights ex- cept for the cross-attention layers at the input and output. As we now show in Table 4 the model pretrained on the visual data performs better than the model trained from scratch. Moreover, both models are better in accuracy than the Bayesian decoding model applied by Lee et al. (2025), which produced position decoding error of 12 cm. This con- firms that our model is capable of learning representations that transfer from data collected in one region to another, even if those regions are functionally very distinct.
Error
Poyo+ from scratch POYO+ pre-trained
10.33 cm
9.58 cm
4 RELATED WORK
Multi-session training for neural recordings. Neural decoding has been extremely challenging, especially when trying to transfer to recordings from different days or animals. This is due to a number of issues, one of the most basic being the fact that each recording session has different set

---

Published as a conference paper at ICLR 2025
of neurons associated with it Dabagia et al. (2022). ). Recently, there has been progress in demonstrat- ing that large transformer architectures can be trained on multiple recording sessions to successfully decode and transfer to new brains and tasks (Azabou et al. 2024; Schneider al.) 2023b Ye al, 2024). In particular, the POYO architecture and tokenization scheme introduced by (Azabou et al, 2024) provided a scalable solution to this problem by using latent embedding scheme for encod- ing individual neural dynamics, along with a PerceiverIO encoder-decoder architecture to efficiently decode high-resolution behavioral outputs from brain signals. However, existing approaches for multi-session training ( (Ye et al. 2024 Azabou et al. 2024) have still only been applied to sessions containing the same brain regions and neural cell types, and for the same decoding tasks. Here, we investigate whether it is possible to train this type of transformer architecture on increasingly hetero- geneous datasets from different areas of the brain and genetically distinct cell types. Additionally, previous methods have been applied to tens to a hundred sessions, POYO+ moves to the next order of magnitude with 1,335 sessions.
Deep learning methods for analysis of neural population activity in calcium data. As calcium imaging datasets are improving, there have been a number of works that have focused on building
neural population models for calcium data analysis (Schneider et al. 20235: Prince et al., [021]
Zhu et al. 2022 de Oliveira Fonseca et al., 2023: Park al. 2023). Many of these methods designed for dynamical systems analysis and forecasting and thus the decoding accuracy is not the focus. Most similar ours CEBRA (Schneider al., 2023b) as the method provides a way to incorporate supervised data for training a decoder. CEBRA A uses behavioral labels to define positive views to align labeled examples. This method requires either trial-aligned data to concatenate many animals into pseudopopulation or multi-data stitching approaches where you need to train a differ- ent encoder head for each dataset. contrast to these methods, we demonstrate flexible decoding of multiple behavior modalities and contexts. Additionally, due to the rich encoder and tokenization method, we also show that our method can be used to decipher different cell types and regions from the embeddings of population- neural activity.
CONCLUSION
In this work, we introduced novel -uuti-session, multi-task neural decoding approach and applied it to large-scale dataset from the Allen Institute for Brain Science Our results demonstrate that scaling up to incorporate multiple tasks can effectively extract meaningful insights into neural circuit diversity, relying solely on neural activity without prior knowledge of neuron sub-types or brain regions. The latent representations learned by our model revealed clear distinctions between cellular sub-types and brain regions, underscoring the power of multi-task learning to uncover biological structures embedded within neural data.
Our approach further highlights the value of integrating diverse neural datasets, offering powerful framework for comparing and contrasting activity patterns across different cell populations, brain regions, experimental conditions, and even across animals. This comparative analysis enables us to observe how distinct neural populations interact and coordinate their activity, yielding insights that would be difficult to capture through isolated experiments. By leveraging large collection of diverse recordings, we were able to learn patterns of neural population dynamics that generalize across subjects and conditions.
In developing this approach, we utilized the largest and most diverse dataset publicly available, covering multiple brain regions, Cre-lines, and experimental sessions. This comprehensive dataset allowed us to train models that generalize across varying conditions. However, despite the dataset's s richness, we have only tested our model in visual cortex and the current work does not incorporate other regions or species, which could further validate the generalizability of our model. Addition- ally, our model does not yet incorporate neuron activity prediction, which could help to further improve the model's ability to understand the dynamics of different cellular sub-types and regions. Addressing these limitations by expanding to additional brain areas, species, and neuron prediction tasks will be crucial areas in future research.
Looking forward, our framework has the potential to go beyond traditional limitations in neuro- science by algorithmically stitching together neural populations from different experiments. Apply- ing this model to even larger datasets, across species and brain regions, could lead to discoveries in how neural circuits are organized and how they change over learning and disease.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We would like to thank J. Quinn Lee and Mark P. Brandon for sharing the Lee et al. 2025 dataset and providing feedback on the experiment. We would also like to thank Saskia de Vries for provid- ing feedback the analysis. BAR acknowledges support from NSERC (Discovery Grant: RGPIN- 2020-05105; Discovery Accelerator Supplement: RGPAS-2020-0N8 2020-0032::; Arthur B. McDonald Fel- lowship: 366355-20233 CIFAR (Canada AI Chair; Learning in Machine and Brains Fellowship), and the Canada Excellence Research Chairs Program. KXP was supported by PhD Fellowship from Google DeepMind. ED acknowledges support from NIH award 1R01 IRO1E8029852-0, NSF award IIS- 116-2146072 as well as generous gifts from the CIFAR Azrieli Global Scholars Program. MA acknowledges support from the National Science Foundation and by DoD OUSD (R&E) under Co- operative Agreement PHY-2229929 (The NSF AI Institute for Artificial and Natural Intelligence).
REFERENCES
Mehdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael Mendelson, Blake Richards. Matthew Perich, Guillaume Lajoie, and Eva Dyer. A unified, scalable framework for neural population decoding, Advances in Neural Information Processing Systems, 36, 2024.
Max Dabagia, Konrad Kording, and Eva L Dyer. Aligning latent representations neural activity. Nature Biomedical Engineering, pp. 1-7, 2022.
Antonio Henrique de Oliveira Fonseca, Emanuele Zappala, Josue Ortega Caro, and David Van Dijk. Continuous spatiotemporal transformer. International Conference on Machine Learning (ICML), 2023.
Saskia EJ de Vries, Jerome Lecoq, Michael A Buice, Peter A Groblewski, Gabriel Ocker, Michael Oliver, David Feng, Nicholas Cain, Peter Ledochowitsch, Daniel Millman, et al. A large-scale standardized physiological survey reveals functional organization of the mouse visual cortex. Nature neuroscience, 23(1)1138-151 2020.
Daniel Deitch, Alon Rubin, and Yaniv Ziv. Representational drift in the mouse visual cortex. Current Biology, 91(19):427744339 00.1016/j.ubb2221..00.4 URL hnttps://iinking p11/s096099222100000
October 2021. ISSN 09609822. doi:
sseviee ere. cwm/ coonrrerieee.
Mario Dipoppa, Adam Ranson, Michael Krumin, Marius Pachitariu, Matteo Carandini, and Ken- neth D. Harris. Vision and locomotion shape the interactions between neuron types mouse visual cortex. Newron. 98(3):602 615.e8, 2018. ISSN 0896-6273 6273. doi: http:://bi..rr//001110//
neuron.2018.03.om URL https: www www.swwinncccieec....cmmmmm www com//ccience/artiile pii 0151700896627318304422
Kathleen Esfahany. Isabel Siergiej. Yuan Zhao, and I1 Memming Park. Organization of neural population code mouse visual system. eNeuro, 5(4), 2018 doi: 10.1523/NEEUR..044411, www. http://wwwww.....oo//o// htt.://ewww..ooooooooooee/////// 1TT.:/EEEEERRRO11111222200 0414-17.2018
2018. URL https: /tttp::/www.eett
Adam R Ferguson, Jessica L Nielson, Melissa H Cragin, Anita E Bandrowski, and Maryann Martone. Big data from small data: data-sharing in the the"long tail' neuroscience. Nature neuro- science, 17411):1422-144, 2014.
Colleen Gillon, Jérôme Lecoq, Jason E Pina, Ruweida Ahmed, Yazan N Billeh, Shiella Caldejon, Peter Groblewski, Timothy M Henley, India Kato, Eric Lee, et al. Responses of pyramidal cell somata and apical dendrites mouse visual cortex over multiple days. Scientific Data, 10(1): 287, 2023.
Weibang Jiang, Liming Zhao, and Bao liang Lu. Large brain model for learning generic represen- tations with tremendous EEG data in BCI. In The Twelfth International Conference on Learning Representations, 2024. URL https: httpsttt/ppteeeeeeeeeeee http://perreeeeeteeeeee.oorrnn//

---

Published as a conference paper at ICLR 2025
J. Quinn Lee, Alexandra T. Keinath, Erica Cianfarano, and Mark P. Brandon. Identifying rep- resentational structure in cal to benchmark theoretical models cognitive mapping. Neuron, 11302):307-320.e5, Jan 2025. ISSN 0896-6273. doi: 10.1016/j 10.1016/j.euurn.22200 nww.on.202...000..... 20.000
URL
https://dof . / C. wwwe//.eero...0222 10.224.1.0000 10.027,
Athena H Li and Shi-Bing Yang. Utilizing cre-lox-based mouse genetic tools in neuroscience search. Gene Reports, pp. 101825, 2023.
Timon Merk, Victoria Peterson, Richard Köhler, Stefan Haufe, R Mark Richardson, and Wolf-Julian Neumann. Machine learning based brain signal decoding for intelligent adaptive deep brain stim- ulation. Experimental Neurology, 351:1 113991 2022.
Seungbin Park, Megan Lipton, and Maria Dadarlat. Decoding multi-limb movements from low temporal resolution calcium imaging using deep learning. bioRxiv, pp. 2023-11, 2023.
Luke Y Prince, Shahab Bakhtiari, Colleen J Gillon, and Blake A Richards. Parallel inference of hicrarchical latent dynamics in two-photon calcium imaging of neuronal populations. bioRxiv, pp. 2021-03, 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, llya Sutskever, et al. Improving language under- standing by generative pre-training. 2018.
Simanto Saha, Khondaker A Mamun, Khawza Ahmed, Raqibul Mostafa, Ganesh R Naik, Sam Darvishi, Ahsan H Khandoker, and Mathias Baumert. Progress in brain computer interface: Chal- lenges and opportunities. Frontiers in systems neuroscience, 15:578875, 2021.
Federico Scala, Dmitry Kobak, Matteo Bernabucci, Yves Bernaerts, Cathryn René Cadwell, Je- sus Ramon Castro, Leonard Hartmanis, Xiaolong Jiang, Sophie Laturnus. Elanine Miranda, et al. Phenotypic variation of transcriptomic cell types in mouse motor cortex. Nature, 598(7879); 144 150, 2021.
Aidan Schneider, Mchdi Azabou, Louis McDougall- Viggee,,, David Parks, Sahara Ensley, Kiran Bhaikaran-Nair, Tom Nowakowski, Eva L Dyer, and Keith B Hengen. Transcriptomic cell type structures in vivo neuronal activity across multiple time scales. bioRxiv, pp. 2022-07, 2022.
Aidan Schneider, Mchdi Azabou, Louis McDougall- Vigier, David F. Parks, Sahara Ensley, Ki- ran Bhaskaran-Nair Tomasz Nowakowski, Eva L. Dyer, and Keith B. Hengen. Transcriptomic cell type structures in vivo neuronal activity across multiple timescales. Cell Reports, 42(4): 112318, April 2023a. ISSN 2211 1247. doi: 10.10066/..edrrep2223..1 122311
12318.
URL https:
//linkinghub. sevier.eee..eee com wtt.c///eeeerreeeee//////2222110
Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. Learnable latent embeddings for joint behavioural and neural analysis. Nature, pp. 1-9, 2023b.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. arXiv preprint 1rXiv:2104.09864 2021.
Rudi Tong, Ronan da Silva, Dongyan Lin, Arna Ghosh, James Wilsenach, Erica Cianfarano, Pouya Bashivan, Blake Richards, and Stuart Trenholm. The feature landscape of visual cortex. bioRxiv, pp. 2023-11, 2023.
Laurens der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal Ma- chine Learning Research, 9(86):2579-6033 2008. URL http:/ttp///// jmlr
varie rww.teenaaatt.tmmmmmmmmm........ html
Joel Ye, Jennifer Collinger, Leila Wehbe, and Robert Gaunt. Neural data transformer multi- context pretraining for neural spiking activity. Advances in Neural Information Processing Sys- tems, 36, 2024.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 1904.0096, 2019.
12

---

Published as a conference paper at ICLR 2025
Hongkui Zeng and Joshua R Sanes. Neuronal cell-type classificatton: challenges, opportunities and the path forward. Nature Reviews Neuroscience, 18(9):530- 546, 2017.
Yizi Zhang, Yanchen Wang, Donato Jimenez-Beneto, Zixuan Wang, Mehdi Azabou, Blake Richards, Olivier Winter, The International Brain Laboratory, Eva Dyer, Liam Paninski, et al. Towards a"" universal translator" for neural dynamics at single-cell, single-spike resolution. arXiv preprint FxXi::2407.14666 2024.
Feng Zhu, Harrison A Grier, Raghav Tandon, Changjia Cai, Anjali Agarwal, Andrea Giovannucci, Matthew T Kaufman, and Chethan Pandarinath. A deep learning framework for inference of single-trial neural population dynamics from calcium imaging with subframe temporal resolution. Nature neuroscience, (5(12):1224-77)) 2022.
13

---

Published as a conference paper at ICLR 2025
APPENDIX
A WHY BUILD A MODEL ACROSS DIVERSE CELL TYPES AND BRAIN REGIONS?
A.I CHALLENGES
Sensory processing occurs across multiple interconnected brain areas, which collectively transform their inputs into diverse, hierarchically organized representations. Consequently, the representation at the population level in different brain regions exhibit significant variation, particularly as we move up the hierarchy (Togg et al. 2023 Esfahany al. 2018). For example, when we examine the rep- resentations in primary visual cortex (VisP) and compare them to those in a higher-order visual area like rostro-lateral visual cortex (VisRL), the neurons appear to respond different pieces of infor- mation, e.g. VisP neurons show strong orientation tuning while VisRL neurons do not (Esfahany) al. 2018). This variability highlights the complexity of neural "conversations" in the brain, as neurons in different areas encode information in unique, complementary ways that are essential for comprehensive sensory processing.
In addition to understanding the variability across brain regions, understanding the information pro- cessing properties of different genetically defined cellular sub-types crucial. Luckily, modern neuroscience tools allow researchers to target their recordings to distinct cellular sub-types using Cre-lox technology (Li & Yang 2023). As a result, it now known that different cellular sub-types exhibit unique response patterns to the same stimuli (Zeng & Sanes] 2017). For example, differ- ent sub-types in visual cortex are distinctly modulated by locomotion differently, and show distinct response profiles to the same sensorimotor conditions (Dipoppa al. 2018). As well, some cel- lular sub-types are more numerous than others, with very high variance in the sparsity/density of recordings from different Cre-lines (de Vries et al., 2020). Training decoding model across these distinct cellular sub-types challenging due to the heterogeneous and context-dependen nature of their responses, but it is also critical for the ability to train across highly diverse neural recordings.
A.2 OPPORTUNITIES
What could we achieve with a neural decoding model that combines information across brain re- gions, cell types, and contexts? First, such a model could serve generalist framework- pretrained for any brain region or Cre-line-allowing for fine-tuning on new datasets regardless of the regions sub-types involved. By stitching together data from different experiments, could algorithmically link neural populations that are otherwise studied in isolation. For example, in single mouse or experiment, we typically record from only one sub-type of cell in a specific brain region. But by integrating multiple datasets, we could begin to understand how different cel- lular sub-types work together across regions in unified system, gaining a broader view of neural dynamics than would be possible using data from single experiment.
Combining diverse neural datasets offers a powerful opportunity compare and contrast activity patterns across different cell populations, brain regions, experimental conditions, and even across animals. This comparative approach allows us to observe how distinct neural populations interact and coordinate their activity in ways that individual experiments cannot capture. By integrating diverse datasets, we can explore how neural circuits change during learning, identifying patterns of adaptation that generalize across subjects. Similarly, we can track how neural activity evolves in the context of disease, revealing key differences between healthy and pathological states. Altogether, by training generalist models wealth of diverse data, can uncover neural circuit interactions that remain hidden when analyzing data from single source.
B DETAILS ON THE ALLEN BRAIN OBSERVATORY DATASET
Below we provide a breakdown of all the experiments in the Allen Brain Observatory (Table A1 where we see the number of sessions per brain region and the different tasks.
14

---

Published as a conference paper at ICLR 2025
TASK
TARGETED STRUCTURE
All
AL
AM
P
PM
RL
Drifting (DG)
456
41
38
106
144
87
40
Static (SG)
456
41
38
106
144
87
40
Scenes (NS)
456
41
38
106
144
87
40
Movies (NM)
1368
123
114
318
432
261
120
Sparse Noise (LSN)
456
41
38
106
144
87
40
Table Al: Session breakdown
C MODEL DETAILS AND EXPERIMENTAL SETUP
C.I ROTARY POSITION EMBEDDINGS (RoPE)
Position of tokens our model are determined by their timestamps. Thus, we allow the positions to be continuous values, which different from the conventional application of RoPE in domains such as language and vision. For this, we define a custom rotation matrix R() as follows:
R (t, Ti)
O (t, T2)
0
R(t)
:
0
Rax2(1, Topp)
With each 2 x sub-matrix being defined as:
(t, = [cos(2ml/T) sin (TT/ cos(2Tl/T)
- sin(221l/ (T)
Here, Ti, .., Top denote the "time-period" of different sinusoids. These decide what resolution and range the model is capable of resolving. We set these time time-periods to be logarithmically spaced from 0.lms to 4.0s.
C.2 TRAINING DETAILS
We the LAMB optimizer You al. (2019) with weight decay to train PoYo+ The learning rate held constant, then decayed towards the end training (last 25% epochs), using a cosine decay schedule. We train POYO+ for 300 epochs using 8 Nvidia H100s with global batch size of 4800.
POYO+ is trained with 8 transformer blocks: 2 cross-attention blocks at the input and output and attention
self attention blocks. We use 128 latent tokens with dimension of 128. We use with cross attention layers using heads, and self-attention layers using 8 heads. We use dropout when computing attention, and the feed-forward networks (dropout rate 0.2).
The weight for each task determined based on the scale of the loss from an initial model. The regression targets are z- scored based on the mean and standard deviation computed from the training trials across all training sessions.
POYO+ is trained on a total of 12 tasks, which are detailed in Table A2
C.3 MULTI-MODAL DECODER IN POYO+
In Figure A detail the architecture of the multi-modal decoder used in PoYo+ All tokens regardless of the task are processed through the same cross-attention layer, and a custom router groups together tokens from the same task and feeds them task-specific linear decoder.
15

---

Published as a conference paper at ICLR 2025
Table A2: Summary of tasks, dimensions, loss functions, and reported metrics used in POYO+
Task
Dim
Type
Loss
Weight
Metric
Drifting gratings orientation
Classification
bce
1.0
Accuracy
Drifting gratings temporal frequency
5
Classification
bce
1.0
Accuracy
Natural movie one frame
900
Segmentation
bce
0.25
Acc within Is
Natural movie two frame Natural movie three frame Locally sparse noise frame Static gratings orientation
900
Segmentation
bce
0.2
Acc within 1s
3600
Segmentation bce
0.2
Acc within Is
8000
Segmentation bce
1.0
Accuracy
Segmentation bce
1.0
Accuracy
Static gratings spatial frequency
Segmentation bce
1.0
Accuracy
Static gratings phase
5
Segmentation bce
1.0
Accuracy
Natural scenes
119
Segmentation
bce
0.3
Accuracy
Running speed
Regression mse
1.5
y²
Pupil location
Regression mse
8.0
22
D ADDITIONAL EXPERIMENTS
D.1 TRANSFER EXPERIMENTS TO OPENSCOPE DATASET
In another set of experiments, we examined our ability to transfer from our pretrained model to another open dataset with different recording conditions and stimuli, collected part of The Allen Institute Mindscope's OpenScope program (Gillon et al., 2023). This dataset consists of responses of neurons in VisP from two Cre driver lines targeting distinet excitatory populations (Cux2 and Rbp4) recorded at two different depths. The two different depths were used to target either the cell bodies or distal apical dendrites of the neurons. Thus, these recordings contained responses from dendritic processes (Figure A2 in the Appendix), not just neural cell-bodies, which allowed us to examine our ability transfer dendrites.
We tested our ability to finetune our pretrained model to decode the orientation of Gabor stimuli, from either the activity of cell bodies (Soma) or dendrites (Dend) (Table 3}. Our results suggest that there is indeed benefit in pretraining on another large dataset. When we compare to the MLP baseline, we find near 7% improvement performance with POYo (DG) and 9% improvement with POYO+ which is pretrained much more diverse tasks. When we break these numbers down by the type of recording, we find that the improvements on Dendritic sessions benefit the most, with
Route output tokens different task specific linear projectors
Multi-modal Decoder
latent
Attention
Linear
Output tolenns
90"
2Hz
Query takens
DG Oriemattinn
Tempooak Insuuetiyy
H CO -
W
w
Task embedding Session embeciding: Cutput timestamp
Pupil location
Running speed
Figure Al: Detailed architecture of the decoder in POYO+. All tokens from all tasks are processed together throught the same cross-attention layer. A router is used to group together tokens from the same task, for which a task-specific linear decoder is used to project the token into the target dimension.
16

---

Published as a conference paper at ICLR 2025
somatic
dendritic
Figure A2: Visualizations of somatic dendritic region-of-interese
over 10% improvement with PoYo+ despite the fact that we did not pretrain dendritic data which varies considerably from somatic data.
Table A3: Transfer from POYO+ to new Openscope dataset. Here, compare the decoding performance 3-class task for MLP baseline, the single task POYO (DG) baseline finetuned for the task, and POYO+ finetuned.
Avg (5)
Dend (3)
Soma (2)
MLP
72.94%
73.5 1%
72.10%
PΟYO (DG)
79.70%
81.74%
76.76%
POYO+
81.91
84.19%
78.49%
D.2 COMPARISON OF CEBRA AND PoYO+
We compare POYO+ to CEBRA (Schneider et al. 2023b), which is state-of-the-art neural latent model that supports multi-session training. CEBRA self-supervised method that uses auxiliary context variables to form representations that are guided by the behavior or stimulus labels. We test CEBRA two different settings. first setting, we replicate the drifting gratings experiment in our paper. In the second, we replicate the experimental setup in Schneider et al. (023b) using the pseudomouse method.
D.2.1 DRIFTING GRATINGS TASK
We pretrain multi-session CEBRA on the drifting gratings data from all sessions in the Allen Brain Observatory. We then fit linear decoder on top of the representations learned by CEBRA. We found that this achieves a performance of 18.66%. For comparaison, the MI.P baselines averaged across the same session reach performance 38.30%. POYO (DG) which is trained on the same data as CEBRA, in multi-session fashion, reaches performance of 48.96%.
D.2.2 NATURAL MOVIE FRAME PREDICTION TASK USING PSEUDOMOUSE
In a second attempt at comparing CEBRA and PoYo+ we replicate the CEBRA pseudomouse experiment using the same Allen Brain Observatory dataset natural movie subset used in the original CEBRA paper[Schnene [Schneider et al. (023b).
Following CEBRA'S experimental setup, we excluded sessions recorded from mice expressing three inhibitory Cre-lines and Sosrl-Cre.GN200 Instead using the traces provided by the Allen In- stitute, we also used the preprocessed and aligned version of the traces from Deitch et al. (2021) hatpss//gahub..co//i/llabbb/a wttpr//sssisaaditittee////aaaaaa driit//reee/main/daaa)a This dataset provides traces aligned with movie frames for natural movie stimulus, facilitating aggregation of neurons across animals. We note that we only compared CEBRA's VISp results here because the main findings and CE- BRA's best performance were reported from VISp out of the six brain regions.
17

---

Published as a conference paper at ICLR 2025
Natural Movie Frame Decoding
100
Average Accuracies (%)
80
60
40
20
CEBRA
POYO
200
400
600
800
1000
Number of Neurons
Figure A3: Decoding accuracy comparison between PoYo+ and CEBRA through neuron concate- nation. PoYO+ consistently outperforms CEBRA across all pseudomouse sizes on natural movie frame decoding task.
We used all sessions from VISp available in the dataset. In each session, the natural movic one video was played 10 times consecutively. We split the training in the exact same way as CEBRA, using the first eight repeats the training set, the ninth as the validation and the tenth as the test set.
We constructed a "pseudomouse" by first aggregating all neurons from all subjects, only including the neurons that were tracked across all three types of sessions, ensuring we have the same neuron pool as CEBRA. Since we could not identify the exact same neurons used the CEBRA experi- ments, we employed the same neuron pool and method, using a random sampler to sample different numbers of neurons. This allowed us to construct pscudomice of ten different population sizes: {10, 30, 50, 100, 200, 400, 600, 800, 900, 1000}. For each size, sampled randomly with five different seeds to ensure reproducibility, generating 50 pseudosubjects in total.
We tested the model's frame decoding performance with 900-class classification task. Since the movie frames are very similar within a small time window, we followed CEBRA's approach: the prediction was within one-second window, counted it as correct prediction.
We train the model for 3000 epochs for pseudomice of sizes {10, 30, 50, 100, 200, 400} neu- rons, 1500 epochs for sizes (600, 800) neurons and 1000 epochs for sizes (900,1000) neurons. We then computed the averages decoding accuracy for each size, as well as the standard er- ror of the mean (s.c.m.). We report CEBRA results from their repository for reproducing figures hatpa///mmmammeeeeoooooooooollaa
The results show that POYO+ outperformed CEBRA for pseudosubjects of all sizes (Figure Notably, PoYo+ outperformed CEBRA by a large margin for pseudosubjects containing fewer neu- rons. For example, PoYo+ outperformed pseudosubjects of 30 neurons by more than 10% and 200 neurons by almost 20%. This suggests that PoYo+ is able to extract latent information with much fewer neurons than CEBRA.
E ANALYSIS OF LATENT EMBEDDINGS
We provide details about the process of analysis to obtain Figure 4
E. EXPERIMENT DETAILS
To carry out the analysis in Figure we extracted the latents from random sample of approximately 500,000 Is windows. For a Is context window, PoYo+ produces 64 latent tokens, each with a 128- dimensional embedding. We concatenated these latents, resulting in combined 8 3192-dimensional
18

---

Published as a conference paper at ICLR 2025
latent representation for each sample. For part of the analyses, will compute, session and Cre-line averaged latent representations.
Dimensionality Reduction by t-SNE For the purposes of visualizing the latent space, down- sample the number of samples to around 110,000. We first used Principal component analysis (PCA) to reduce the dimensionality from 8,192 50. Next, performed t-SNE (van der Maaten & Hin- ton 2008 to further reduce the dimentionality from 50 to 2 for visualization.
Linear Decoding For all linear decoding tasks, used the seikit-learn implementation of logistic regression with a train-validation-test split of 60/20/20. We performed hyperparameter search for regularization parameter C between 0.1 and 500.. We determine the best model based on the per- formance on the validation set, and report the linear layer's performance on the test set. To ensure fair comparison with the Baseline metrics, we calculated the average latent representation for each session instead of using all samples individually.
E.2 CRE-LINE AND BRAIN REGION DECODING
Here, we report the accuracies for cell type and brain region classification, from our model, the single task (DG) model, and the the Baseline feature-based classifier in Schneider al. (2022) (Table A4).
worth noting that Poyo (DG) only contains roughly 400 samples in total, while the Baseline and PoYo+ have over 1000 samples for linear decoding.
Table A4:
Performance comparison of Baseline, POYo (DG). and POYO+ across different region
and cell type classification tasks. Num of classes
Baseline
Poyo (DG)
PoYo+
Cre-line
13
56.00%
39.57%
67.93%
Cre-line Layer
24
38.48%
16.86%
40.88%
Area
6
36.47%
39.01%
59.81%
Baseline Metries Inspired by Schneider al. (202a)), used the calcium event data provided by the Allen Institute to compute the inter-spike intervals (ISIs) for each neuron in each session. Using these ISIs, we calculated the following metrics for each neuron:
Minimum, maximum, and median of ISIs
Mean firing rate (mean FR)
Standard deviation (SD)
Coefficient variation (CV)
Local variation (LV)
Revised local variation (LVR)
Local alternative to the coefficient of variation (CV2)
Spectral density power four bands: delta, theta, alpha, beta, and gamma (PSD)
Two parameters of the gamma distribution (a and b) used to describe the ISI distribution (a and 3)
After calculating these metrics, averaged the results across all neurons in each session to obtain Baseline metrics at the session level.
is important to note that due to the sparse nature of calcium data, LV, LVR, and CV2 values cannot be calculated. As a result, these metrics are excluded from the subsequent analyses.
E.3 LEARNED TASK EMBEDDINGS
visualize the learned task embeddings in Poyo+. We find that related tasks typically cluster together. For example the various drifting gratings decoding tasks, or the decoding tasks corresponding to the three natural movies.
In Figure A4
19

---

Published as a conference paper at ICLR 2025
7.5
DRIFTING, GRATINGS, ORIENTATION
GRATINGS,
FREQUENCY
STATIC CRATTGGGAINGSS.OOOO ORIENTATION
STATIC_GRATINGS GRATINGS, LAISSSSAAAAALLRROLALAA FREQUENCY
STATIC
NATURAL
-25
NATURAL_ DOVIE_ON__PATE
ONE_FRAME
NATURAL_ MOVIE_TVOOO TIRE__WO__FAARR
FRAME
-500
NATURAL MOVIE THREE FRAME
THREE_FRAME
-7.
LOCALLY SPASE FRAME
RUNNING SPEED
-10.0
RUPIL LOCLLLOCAIION
Figure A4: PCA projection of learned task emheddings
E.4 VISUALIZATION OF MORE FINE-GRAINED CRE-LINE CONFUSION MATRIX
We provide a visualization of the confusion matrix for a more fine-grained classification task with POYo+ (Figure A5
Pvolb Pvalb_1//3 12/3
0:33
0.8
0.200.20
0.62
0.05
0:33
$4
0.40
0.200.200.20
$55_55
0.45
0.04
0.090.188
0.09
Vb.14
0.75
0.25
0.6
12-3
0.05
0.050.85
Emx1, L2/3
0.04
0.04 O OR 0.160.288..00
0.12
0.5
True label
Emx]_L4 L4
0.930.270.0
Percentage
0.14
0.140.000
0.005.055
0.05
Feaf2
0.29
0.14
0:43 0.14
0.4
Nr5a]_ue L4
0.200.10
0:30
Ntsr3_L6
0.50...5
0.3
Rbp4_13
0.3G
0.57
0.07
Rorb L4
0.08
0b 0.46
Scnnla L5
0.34
0.57
Sc17e7_12/
0.10
0.200.30
0.1
Sk1177_L5
0.20
0.200.20
0.200.20
Tik 15
Pvolb L22
huai L5
S9.s
Wi 1220
Cae L4
Emid 1225
rest2 15
NSS1 L4
Nur1 L5
Ripa L5
Scnnda Ls
Ste3777.4 14
L5
SFt
Vip.
Emai
NiSal
Rorb
Sonda
Sic3707
Predicted label
Figure A5: Examining the organization the latent embeddings from POYO+ by Cre-line. Confusion
matrix for Cre-line classification, including the layer for each Cre-line.
E.5 VISUALIZATION OF BASELINE AND PoYo (DG)
We provide the confusion matrices for Baseline and PoYo (DG) in classification tasks (Figure A6 Figure A7).
20

---

Published as a conference paper at ICLR 2025
ViSal
0.21
0.28
0.12
0.19
0.03
0.18
Visal
0.25
0:25
0.00
0.25
0.25
0.00
0.10
0.34
0.13
0.10
0.03
0.30
V.5am
0.00
0:33
0.00
0.17
0.50
0.00
True label
VIS1
0.25
0.00
0.60
0.10
0.00
0.05
True label
VISI
0.16
0.00
0.32
0.37
0.16
0.00
V15p
0.06
0.13
0.17
0.30
0.09
0.26
V1Sp
0.09
0.00
0.22
0.58
0.12
0.00
V15pm
0.11
0.04
0.11
0.21
0.36
0.28
V15pm
0.00
0.15
0.23
0.23
0.31
0.06
ViSri
0.19
0.15
0.15
0.00
0.11
0.41
VI5i
0.00
0.14
0.00
0.00
0.29
0.57
ViSai
ViSam
Vi5p Predicted label
VISI
Vi5pm
ViSri
ViSal
ViSam
VSp Predicted label
VISI
V15pm
0.662..551
0.12
0.00 0.90.100.00
Vpp
0.83
0.17
Cux2
0.020.02
0.43
0. 0.02 05
020.02
0.200.00.00
Oux2
0.44
0.110.33
0.11
Emx3
0.220.57
0. 0.040.13
0.04
Emx11
0.430.44
0.14
True label
Feaf2
0:33
0.67
label
Fezf2
1.500.50
Nr5al
0.03
0.060.066
0...8
0.120.150
.600..3
Nr5a1
Mar2
0.09
0.08
0.73
0.08
Nis1
Rbo4
0.230..05 .5
0.05
0.37
0.21
0.13
Rbp4
0.80
Rarb
0.430.14
0.140.29
Semla
0.75
Scnnla
1.00
0.740.06
Stc17e7
0.320.07
0.07
0.400.13
0.330.67
Tha
0.33
0.33
0.33
Pybb.
vp.
Rop4
Roft.
Soniaa
SIc1787
Tba
Peelb
5at
Wp
Cuv2
Fort
Nrat
nbpa
Robb
Sonnt
Sta7777
Predicted label
Predicted label
Figure A6: Baseline POYO (DG) confusion matrices for brain area Cre-line classification. A.
Brain area classification (Baseline). B. Brain area classification (PoYo (DG) latents). C. Cre-line classification
(Baseline). D. Cre-line classification (PoYo (DG) latents).
We also provide the dendrograms for detailed cre-lines and brain area for PoYo (DG) (Figure A8
E.6 TASK DEPENDENT MODULATION EMBEDDINGS CRE-LINE
We provide the extended version of Figure in this subsection (Figure A9
21

---

Published as a conference paper at ICLR 2025
A
Nab 1233
200-00000
P
0.10..10..40
Pvolo LS
5.50.10
0.250.330.08 0.17
0.05
0.08
0.8
Vo 12/3
0.38
0.650.27
0.09
M4.14
0.05
0.14
100.33
0.10
Cas2_44
0.00
0.28
0.05
0.180.00
0.00
0.05
0.270:10::445
0.6
Truee label
Emmc1 1.4
Percentage
EnvO_LS
0.20
0.67
14511,44
M5511L5
0.03
0.14
36.0.2
0.006.030.07
0.070.030.03
Mer1_
Rop4 L5
0.11
0.050.16
14
0.24
120.000
0.32
0.12
0400.0
04.000
L4
SundaaLS LS
0.330.35
Ske1767
0.20
0:350.27
Sc17aT_15
0.4005500.0
0.67
Past
Prabb
te 14
SXL
Viyu4
Emat L25
No50
Mest L
Ro04
Sant
Sonna L4
9e3707 L33
SCIN22 13
ma3
0.0
Predicted label
B
1.0
Posbs12/3
1.00
P4
Pvense
L4
0.8
L2
3500.50
0-330.33
Oux2-L//3
0-4G
Cur2-44
9.400
0:20
n-20
0.6
label
Emxc 14
0.70
Percentage
Eme1445
2.32
0:33
0:3,
Rea12-L5
M550
0:25
0.25
0.32
0.4
5.50
R5p4-L5
0.20 8:30
0.20
0.30o.a
0.40
0:20
Sonsta45
0.50
Sc1747444
Sc17w7455
0:33
0330.33
bralbu4
Praboss
vp.229
Vpsu
Go2 123
Curbu4
Emxb420
Kur2us
Madke
Rop4Lss
Sonta
Sontais
9c77.7429
Sc177774
Se177745
Tisus
Pradictad
Figure A7: Baseline and Poyo (DG) confusion matrices for fine-grained Cre-line layer classification. A. Baseline. Poyo (DG)
22

---

Published as a conference paper at ICLR 2025
A
175
150
125
Distance
100
50
25
B
VISam
VISpm
VISp
VISI
ViSrl
VISal
Region
300
250
200
100
50
Scnnla (L4)
Rorb (L4)
(L2/3)
Emx1 (L2/3)
Sic17a7 (L2/3)
Slc17a7 (L5)
Tix (L5)
Scnnla (L5)
Emx1 (L5)
Emx1 (L4)
Cux2 (L4)
Sic17a7 (L4)
Ntsr1 (L.6
Fezf2 (L5)
Nr5a1 (L.4)
Nr5a2 (L5)
Pvalb (L2/3)
Pvalb (L44
Vip (L2/3)
Vip (L4)
Pvalb (LS)
Sst (L4)
Sat (L5)
Rop4
Cre Depth
Figure A8: Visualizations of hierarchical clusterings from POYO (DG). A. Brain areas. B, Cre-lines.
23

---

Published as a conference paper at ICLR 2025
Sp. NM NS SG DG LSN
Sp. NM NS SG DG LSN
Sp. NM NS 5G DG LSN
DG LSN
Scnnla (L4)
Scnnla (L5)
Pvalb (L5)
e-g-
Tix3 (L5)
SG
NS
Sp. NM
Sp. NM NS SG DG LSN
Sp. NM NS SG DG LSN
LSN
LSN
SG DG
DG
Pvalb (L4)
Ntsr1 (L6)
Cux2 (L4)
Emx1 (L2/3)
SG
Sp. NM NS
NS
Sp. NM
Sp. NM NS 5G DG LSN
Sp. NM NS SG DG LSN
Sp. NM NS SG DG LSN
DG LSN
Slc17a7 (L2/3)
Slc17a7 (L4)
Vip (L2/3)
Rbp4 (L5)
SG
NS
Sp. NM
Sp. NM NS 5G DG LSN
Sp. NM NS SG DG LSN
Sp. NM NS SG DG LSN
Sp. NM NS SG DG LSN
Pvalb (L2/3)
Nr5a1 (L5)
2.
Slc17a7 (L5)
Emx1 (L4)
Sp. NM NS 5G DG LSN
Sp. NM NS 5G DG LSN
Sp. NM NS 5G DG LSN
Sp. NM NS 5G DG LSN
Nr5a1 (L4)
Emx1 (L5)
Cux2 (L2/3)
Sst (L5)
5p. NM NS SG DG LSN
5p. NM NS 5G DG LSN
5p. NM NS 5G DG LSN
NM NS SG DG LSN
Vip (L4)
Sst (L4)
Rorb (L4)
Fezf2 (L5)
Sp.
PC1
P
TDL
Figure A9:
Deviation of cre-line averaged embeddings from Baseline.
24


Reviews:
[{"review_id": "X3LhgzbbpJ", "reviewer": "Reviewer_1QiM", "summary": "The paper explores large-scale, multi-task neural decoding across different cell types and brain regions. It introduces a multi-task transformer model, named POYO+, trained on data from the Allen Institute’s Brain Observatory dataset. The paper aims to determine whether transfer learning and multi-task decoding are feasible in the paper of neural heterogeneity, where each cell type and brain region exhibits distinct response dynamics.", "questions": "* I’m really surprised that CEBRA does worse than an MLP in Figure 3B. Why would this happen?\n* Neuron Dropout: The paper uses a neuron dropout to remove some neurons during training. Is this normal? I couldn’t tell from CEBRA overall.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* Calcium Data: From what I can tell, neural decoding with calcium data is rare. Scaling calcium data could have benefits.\n* Interpretability: The paper includes several interpretability analyses to show learned features in the latent representation space. This is useful for neural decoding, which should be more than just performance driven but also have scientific value. \n* Transfer: The authors demonstrate that Poyo+ could be applied to other datasets as well, although I would have liked to see a dataset with different regions instead.", "weaknesses": "* Distinct Regions: I found the region coverage in the paper to be rather weak, only focusing on mice visual cortex. Although multiple sub-regions are considered as well, other papers use multi-region to mean coverage of the whole brain such as [1] or [2]. If the claim of novelty comes from handling diverse brain regions, then I would expect a dataset that spans multiple, functionally distinct areas of the brain. Could the authors explain why they only used the visual cortex? Is it because calcium data is difficult to collect across wider regions? If so, could the authors discuss the application of their findings to a broader set of regions? Also I’m not sure I quite buy the huge distinctions covered in Appendix A -- neural recording datasets can incorporate much larger amounts of spatial coverage. \n* Experimental Limitations\n    * More broadly, only using one dataset is a bit weak to clearly establish benefits from scaling. I would appreciate it if an additional more substantial experiment was studied on the OpenScope dataset or something manageable with new regions. \n    * Moreover, I found figure 3 to show that Poyo+ and Poyo were rather comparable despite needing to train 12 variants of Poyo. Could the authors discuss these findings? \n\nOverall, the paper is well written but I’m not entirely convinced of the claims. I think more discussion can go into the findings. I would raise my scores after seeing a deeper description of using distinct regions and some more focus on added experiments/limitations.\n\n[1] Chau, Wang, et. al. Population Transformer: Learning Population-level Representations of Intracranial Activity. arxiv, 2024.\n[2]  Zhang et. al. Brant: Foundation Model for Intracranial Neural Signal, NeurIPS, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fBo60VKScy", "reviewer": "Reviewer_feML", "summary": "The paper introduces a novel tokenization method and multimodal neural decoder to decode behavioral and stimulus information from calcium responses recorded from multiple regions of the mouse visual cortex. The proposed model achieved state-of-the-art decoding performance in the Allen Brain Dataset; it outperforms previous models in cross-session, cross-region, and cross-dataset settings. Moreover, the clustering of cell types and brain regions emerges from the latent representation of the model purely from response-task pairs. These results suggest that it is beneficial to train neural decoders on multi-task and multi-session recordings.", "questions": "- What is the difference between the query tokens and output tokens in Figure 1? My understanding is that the decoder takes latent representation from the POYO encoder and the query tokens. The query tokens inform/query the decoder which task(s) to predict (e.g. running speed, pupil dilation, etc.), and then a task-specific linear layer processes the output tokens to generate the prediction. However, in Figure 1, the color code suggests the output tokens are used as query tokens. Can the authors clarify this point (and perhaps update the caption of Figure 1)?\n- Following the previous question, does the model predict in an autoregressive fashion? i.e. predict one time-step at a time. Or does the model input the entire calcium trace and predict the whole segment at once? If it is the latter, is causal attention used so that the model doesn’t have access to future time-steps?\n- In Section 3.1 and Figure 3, the authors compare single-session trained baselines vs multi-session trained POYO+ (also POYO and CEBRA). Do the authors have a sense of how data-hungry POYO+ is? For instance, would POYO+ outperform the baseline models in a single-session setting?\n- When comparing the models' performance (Figure 3, Table 1, 2, and 3), can the authors report s.e.m. or standard deviation across (one of) trials/sessions/animals so that the readers can have a sense of how much variation is in the performance?\n- In Section C.2 Training details, can the authors clarify that “the weight for each task is determined based on the scale of the loss from an initial model”? Does this mean that an untrained (randomly initialized) model is used to make predictions for each task, and then their loss magnitude is used to determine the loss scale?\n- A very minor comment: I don’t think “POYO+” was introduced in the main text as the name of the method.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- This paper demonstrates training a neural decoder on large-scale multi-session, multi-region, and multi-task recordings can greatly improve model decoding performance across a wide range of scenarios, and is an important step towards a foundation model of neural data, a topic of great interest in NeuroAI.\n- The experiments are thorough, covering a wide range of realistic transfer learning scenarios. \n- The analysis of the latent embedding on the trained POYO encoder is very interesting. It shows that the model, purely trained on decoding tasks, without any information on the neurons (location, type, region, etc.), learns cell-specific properties.\n- Overall, the paper is clear and well-written.", "weaknesses": "I don’t think the paper has any major weaknesses, however, there are several points I hope the authors can clarify regarding the model setup. Please find them in the Questions section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "woHTBqcSV8", "reviewer": "Reviewer_VpA8", "summary": "The authors proposed a study that extends POYO, a method for decoding cell-types and brain regions according to neural recordings, to multi-session data and multiple tasks. The authors proposed a multi-task decoder and a training strategy that used multi-session data. They also demonstrated the merits of multi-session and multi-task training through extensive experiments and comparisons.", "questions": "1. It is better to provide a figure to illustrate the multi-task decoder, maybe in Appendix . \n2. How is the multi-task decoder proposed in this study different from a general multi-task transformer decoder? What are the unique designs in this multi-task decoder and how are those designs related to the unique characteristics of the decoding tasks?\n3. “The losses are summed and weighted...”, how to determine the weight? Is it inverse proportional to the number of predictions?\n4. The proposed model outperforms CEBRA in natural the movie frame decoding task, how about other tasks?\n5. There are 12 single-task variants of POYO, but it seems that Fig. 3C reports 9 cases. Why are the results for the other three cases missing?\n6. “ we find that the model trained on inhibitory Cre-lines does surprisingly well on the drifting gratings task.” How about other tasks?\n\nMinor issues:\nBoth POYOplus and POYO+ were used in the manuscript (e.g., Fig.3). Please try to make it consistent. \nTypos: “one for each mouse, imaging depth, and visual area. from various depths...”\nHow such relatively low decoding accuracies can support neuroscience research?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The idea of multi-session and multi-task decoding of cell-types and brain regions is not entirely new, however, the authors proposed several key components to realize such an idea. The study was clearly presented and the validations are extensive.", "weaknesses": "Some details can be further elaborated. Specifically, the architecture of the multi-task decoder can be presented with more details, e.g., the numbers of self-attention blocks and attention heads.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WuBHWrUb8f", "reviewer": "Reviewer_RcB8", "summary": "This study investigates the impact of data scale on the performance of brain decoding models, developing a multi-task transformer architecture trained on the brain observation dataset from the Allen Institute for Brain Science. This dataset includes response data from over 100,000 neurons across six brain regions in mice, demonstrating the effects of different stimuli on neuronal activity. The research finds that transfer learning across brain regions and cell subtypes is effective, and that integrating data from heterogeneous datasets can enhance decoding performance. Despite the heterogeneity of neural circuits, the model's latent representations can automatically distinguish between different brain regions and cell subtypes, indicating that the shared computational principles among them can be extracted. The main contributions of the research include: 1) the introduction of a novel multi-task decoder capable of flexibly handling various decoding tasks; 2) validation of successful transfer learning across brain regions and cell types; 3) provision of evidence that increasing the diversity of brain regions can improve decoding performance; and 4) demonstration of meaningful structures in the model's latent representations that reflect biological differences. Overall, this work offers new insights into the similarities and differences among heterogeneous neural circuits and lays the groundwork for training large-scale neural decoding models.", "questions": "The authors mention a data-driven feature selection process.  Could they elaborate on the specific criteria and methods used for selecting features from the neural data?  \n\nThe evaluation metrics presented in the paper primarily focus on accuracy and predictive performance.  Can the authors justify their choice of metrics, particularly in relation to the biological significance of the results?  Are there any additional metrics that could provide a more comprehensive assessment of the model's effectiveness, particularly in understanding the underlying neural mechanisms?\n\nWith the model's complexity and the relatively limited size of the dataset (six brain regions), what steps have the authors taken to mitigate the risk of overfitting?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "This paper presents a novel multi-task transformer architecture for decoding neural activity from a large dataset comprising over 100,000 neurons across six different brain regions.  The originality lies not only in the innovative architecture itself but also in the approach of leveraging transfer learning across heterogeneous datasets.  By demonstrating that insights gained from one brain region can enhance decoding performance in another, the authors push the boundaries of existing methodologies in neural decoding. The quality of the research is commendable, with a robust experimental setup that includes training on extensive neural response data.  The authors employ rigorous validation techniques and comprehensive analyses, providing a solid foundation for their claims.  The results are quantitatively backed by clear metrics that illustrate the performance improvements achieved through their proposed model, showcasing its practical applicability. Generally, this findings not only contribute to our understanding of neural dynamics but also propose a framework that can be utilized in various applications, ranging from neuroprosthetics to brain-computer interfaces.", "weaknesses": "While the authors provide a comprehensive analysis of their proposed multi-task transformer architecture, the experimental validation could be further strengthened.   Specifically, the comparison against state-of-the-art methods in the field is somewhat limited.   For a more robust assessment of the proposed model's efficacy, additional benchmarks should be included, such as comparisons with other established neural decoding frameworks.   This would provide clearer context for the reported performance improvements and demonstrate the advantages of the proposed approach more convincingly.  Meanwhile, this study relies heavily on a large dataset from six brain regions;   however, the diversity of this dataset may not sufficiently represent the variability found across different species or experimental conditions.   To enhance the generalizability of the model, the authors could consider including additional datasets that encompass various neural activity patterns from other organisms or contexts.   This would help validate the transfer learning claims and ensure that the model performs well across a broader spectrum of neural encoding scenarios.\nFurthermore, the discussion of limitations is relatively sparse.   A more thorough acknowledgment of potential shortcomings, such as the influence of noise in the data or the computational demands of the proposed model, would provide a more balanced perspective.   Including these considerations would demonstrate the authors' awareness of the challenges and offer avenues for future research to address these limitations effectively.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)