BLEND: Behavior-guided Neural Population Dynamics Modeling via Privileged Knowledge Distillation
Authors: Zhengrui Guo
Fangxu Zhou
Wei Wu
Qichen Sun
Lishuang Feng
Jinzhuo Wang
Hao Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Computational Neuroscience
Neural Dynamics Modeling
Behavior as Guidance
Privileged Knowledge Distillation
Primary area: applications to neuroscience & cognitive science

Abstract: Modeling the nonlinear dynamics of neuronal populations represents a key pursuit in computational neuroscience. Recent research has increasingly focused on jointly modeling neural activity and behavior to unravel their interconnections. Despite significant efforts, these approaches often necessitate either intricate model designs or oversimplified assumptions. Given the frequent absence of perfectly paired neural-behavioral datasets in real-world scenarios when deploying these models, a critical yet understudied research question emerges: how to develop a model that performs well using only neural activity as input at inference, while benefiting from the insights gained from behavioral signals during training?

To this end, we propose **BLEND**, the **B**ehavior-guided neura**L** population dynamics mod**E**lling framework via privileged k**N**owledge **D**istillation. By considering behavior as privileged information, we train a teacher model that takes both behavior observations (privileged features) and neural activities (regular features) as inputs. A student model is then distilled using only neural activity. Unlike existing methods, our framework is model-agnostic and avoids making strong assumptions about the relationship between behavior and neural activity. This allows BLEND to enhance existing neural dynamics modeling architectures without developing specialized models from scratch. Extensive experiments across neural population activity modeling and transcriptomic neuron identity prediction tasks demonstrate strong capabilities of BLEND, reporting over 50% improvement in behavioral decoding and over 15% improvement in transcriptomic neuron identity prediction after behavior-guided distillation. Furthermore, we empirically explore various behavior-guided distillation strategies within the BLEND framework and present a comprehensive analysis of effectiveness and implications for model performance. Code will be made available at https://github.com/dddavid4real/BLEND.

Published as a conference paper at ICLR 2025
BLEND: BEHAVIOR-GUIDED NEURAL POPULATION DYNAMICS MODELING VIA PRIVILEGED KNOWLEDGE DISTILLATION
Zhengrui Guo
Fangxu Zhou
The Hong Kong University of Science and Technology Beijing Institute Collaborative Innovation zguobal connect uus .hk
Peking University
zhoufangxu@pku. edu en
Wei Wu
Qichen Sun
Peking University
Peking University 2000010820@stu.ppm edu cn
weiwuêstu. pku. edu.
Lishuang Feng
Jinzhuo Wang Peking University
Beihang University
Beijing Institute Collaborative Innovation Lenglishuang@buaa edu.cc cn
Sangjinzhuu@pku.edu on
Hao Chen X
The Hong Kong University of Science and Technology jhc@cse. ust
ABSTRACT
Modeling the nonlinear dynamics of neuronal populations represents a key pur- suit in computational neuroscience. Recent research has increasingly focused on jointly modeling neural activity and behavior to unravel their interconnec- tions. Despite significant efforts, these approaches often necessitate either intri- cate model designs or oversimplified assumptions. Given the frequent absence of perfectly paired neural-behavioral datasets real-world scenarios when deploy- ing these models, critical yet understudied research question emerges: how develop model that performs well using only neural activity as input at inference, while benefiting from the insights gained from behavioral signals during training?
To this end, we propose BLEND, the Behavior-guided neural population dy- namics modElling framework via privileged kNowledge Distillation. By con- sidering behavior as privileged information, we train a teacher model that takes both behavior observations (privileged features) and neural activities (regular features) as inputs. A student model is then distilled using only neural activ- ity. Unlike existing methods, our framework is model-agnostic and avoids mak- ing strong assumptions about the relationship between behavior and neural ac- tivity. This allows BLEND enhance existing neural dynamics modeling ar- chitectures without developing specialized models from scratch. Extensive ex- periments across neural population activity modeling and transcriptomic neuron identity prediction tasks demonstrate strong capabilities of BLEND, reporting over 50% improvement in behavioral decoding and over 15% improvement transcriptomic neuron identity prediction after behavior-guie guided distillation. Fur- thermore, we empirically explore various behavior guided distillation strategies within the BLEND framework and present comprehensive analysis of effective- ness and implications for model performance. Code will be made available at http:::///ituuu.ccom///ee
X Co-corresponding Author

---

Published as a conference paper at ICLR 2025
INTRODUCTION
Large-scale population-Ievel recordings of neural activity enable the understanding of how complex abilities of the brain in sensing, movement, and cognition emerge from the collective activity of grouped neurons (Stevenson & Kording 2011 Jun et al. 2017 Pei et al. (021). This insight has led to the development of various neural dynamics modeling methods to disentangle and interpret the hidden structure of neural population activity with large-volume neural recordings as inputs (Pandarinath et al. 2018 Ye & Pandarinath 202T; Le & Shlizerman 2022).
Alongside the recorded neural population activity, the observed behavior signals provide crucial context and complementary information during neural dynamics modeling (Sani et al. [202]] For example, behavior allows for the integration of neural data with other physiological measures (e.g., muscle activity, eye movements). By incorporating behavioral information, more comprehensive and functionally relevant models of neural dynamics have been proposed (Zhou & Weil 2020 Hur- witz et al. 2021 Schneider et al. 2023b), bridging the gap between neural activity and its real-world manifestations.
However, neural activity recordings with paired behavior signals are not always available in real- world settings, i.e., behavioral data might be partial, limited, or not available for all periods of neural recording. For instance, some studies focus on resting-state neural activity, where a subject is not engaged in any specific task receiving external stimuli. The absence of structured tasks means there's 's no clear temporal alignment between neural events and behavioral events (Drew et al.] 2020 Nozari al. 2024). This discrepancy between the availability of behavioral and neural data leads to a critical distinction in the types of information accessible for model development. These features that only exist during the training stage are called privileged features, and those that exist throughout training and inference stages are termed regular features (Vapnik & Vashist 2009). Thus, creating a model that can perform well using only regular features (neural activity) at inference time, while still benefiting from the insights gained from privileged features (bchavior), represents a critical research question in neural dynamics modeling, especially bridging the gap between controlled experimental settings and real-world applications where privileged knowledge is limited or unavailable. Maximizing the utility of existing privileged-regulat feature pairs to enhance the eegular-feature-basee network performance is promising strategy to answer this question. Yet in the field of computational neuroscience, these efforts remain underexplored.
To this end, in this paper, we propose BLEND, the Behavior-guided neuraL population dynamics modElling framework via privileged kNowledge Distillation. As shown in Fig. [(c) BLEND con- stitutes student and teacher model, in which the teacher trains both behavior observations and neural activity recordings, then distills knowledge to guide the student which takes only neural activity input. This ensures the student model can make predictions using only recorded neural activity during the deployment/infeeene information during the training stage, making it more versatile for settings lacking behavioral data. Our main contributions are summarized as follows:
stage, but also benefits from the guidance from behavior
Built upon the privileged knowledge distillation framework, introduce simple yet effective neural dynamics modeling paradigm, BLEND, relying on a fundamental assumption that behavior can serve as explicit guidance for neural representation learning. Notably, our approach is model- agnostic, allowing for its seamless integration with diverse existing neural dynamics modeling architectures, thus avoiding the need for developing specialized models from scratch.
To evaluate framework. Latents Benchmark 21 for neural activity prediction, behavior decoding, and matching to peri- stimulus time histograms (PSTHs), as well multi-modal calcium imaging dataset for transcrip- tomic identity prediction. Results show that our framework elevates the performance of baseline methods by large margin (>50% improvement in behavioral decoding and > 15% improvement in neuronal identity prediction) and significantly outperforms the state-of-the-art (SOTA) models.
conduct extensive experiments two benchmarks, i.e., Neural
We present a comprehensive analysis of our framework across base models and privileged knowl- edge distillation strategies, revealing key insights into the interaction between neural activity and behavior. Our results demonstrate that behavior- -uided distillation not only improves model per- formance but fundamentally enhances the quality of learned neural representations. This leads to more accurate and nuanced modeling of neural dynamics, offering new perspectives on how behavioral information can be leveraged to better understand complex neural patterns.
2

---

Published as a conference paper at ICLR 2025
Neuml Activity
Neuna NeurilActivity
Behavior Signal
Neural Activity
Neural Activity
Neural Activity
Model
Model
Model
Student Model
Guidance
Teacher Modal
Neurn Activity
Noural Activity
Behavior Signal
Neural Activity
Neural Activity
Noural Activityy
Behavior Signal
(a)
(b)
(c)
Figure 1: Schematic illustration of neural population dynamics modeling mechanisms. In this paper, benchmark all the methods under the framework of masked neural activity reconstruction, in which the model is firstly trained in an unsupervised manner to reconstruct the randomly masked neural activity and then applied to downstream tasks such neural activity prediction and behavior decoding. (a) Neural dynamics modeling methods that only use neural population activity as input. (b) Neural dynamics modeling methods that take behavior information as prior. (c) Our BLEND framework, which considers behavior information as privileged knowledge for distillation.
2 RELATED WORK
Neural dynamies modeling (NDM). NDM refers to category of methods that aim to capture the dynamics of neural activity by using the activity recordings as inputs. common and effective choice is the latent variable model (LVM), which leverages low-dimensional latent factors to inter- pret these dynamics. Various LVMs are developed, ranging from simple non-temporal models such as principal components analysis (PCA) (Cunningham & Yu 2014) and its variants (Kobak et al.) 2016), to linear dynamical systems (aake et al. 2011 Gao et al. (0066), and to complex state space models like LFADS (Pandarinath et al.) 2018). Since the advent of Transformer, its ability to cap- ture temporal dependencies and long- range interactions makes it appropriate for neural data, with representative works such NeuralData' TransaattaTrasrrmer (NDT) (Ye & Pandarinath, 2021). STNDT (Le) & Shlizerman, (022), and EIT (Liu et al., 2022). As shown in Fig. [Taa] methods this category purely depend on neural activity recordings for neural population dynamics modeling yet ignore using the paired behavior information as guidance.
prior. In recent years, an emerging research di- rection has focused on jointly modeling neural population dynamics and behavioral signals. As illustrated Fig. [[b]] these methods can be categorized into two types. The first inputs both be- havior and neural activity and utilizes behavior signals to guide the learning of neural dynamics: Pi-VAE (Zhou & Weil 2020) considers the behavior variables as constraints for the construction of latent space of LVMs; CEBRA (Schneider et al., (023b) utilizes behavior signals to construct contrastive learning samples for label-informed neural activity analysis. However, these approaches normally need to develop delicately designed modules complex training strategies to achieve their goal. The second line of work aims to decompose neural activity into behavior-relevant and schavior-irreleann dynamics and reconstruct both behavior signals and neural activity signals. To that end, a linear state-space model, PSID (Sani et al., (021), is developed to decode population dynamics from motor brain regions. Nonlinear state-space models such as TNDM (Hurwitz al.) 2021] and SABLE (Jude et al. 2022) are proposed to further capture the nonlinear dynamics of neu- ral activity and behavior observation. Nonetheless, these models assume a clear distinction between behaviorally relevant and irrelevant dynamics in input neural activity, which might not always be practical, potentially leading to oversimplification.
Neural dynamics modeling with behavior
In contrast, our approach offers model-agnostic learning paradigm that can be directly applied existing LVM models without relying on strong assumptions, thereby circumventing the issues present in current LVM models that integrate behavior information prior.
Learning under privileged information (LUPP). Firstly proposed in Vapnik & Vashist 2009).
LUPI refers to the setting where, alongside the primary data modality, the model has access to an additional source of information. This extra source of input, termed privileged, is exclusively avail- able during the training phase. The main objective of LUPI is to leverage this privileged information to learn a better model in the primary data modality than one would learn without the privileged information. This learning paradigm has been applied different machine learning problems such as recommendation system (Yang et al. (002),, medical image analysis (Chen et al. 202T), emotion recognition (Aslam al. 2023), and semantic segmentation (Liu al. 2024).
3

---

Published as a conference paper at ICLR 2025
Adapt
Student
Student
Model
Model
Masked Spikes
Privileged Knowledge Distillation
Predicted Rates
Reconstructed Spikes
Spikes
Neuron
Teacher Model
Excitotory
Behevior
Inbibitory
Predicted Rates
Recometructed Spikes
Neurat Population
Transcriptomic Identity Prediction
Activity Analysis
Figure 2: Illustration of the proposed BLEND framework, exemplified using neural spiking activity data. Left: Behavior-guicee guuided neural representation learning via privileged knowledge distillation. The teacher model is trained on a composite of neural activity and behavioral signals, subsequently distilling its knowledge to a student model that utilizes solely neural activity as input. Right: During the inference phase, the distilled student model is employed for neural population activity analysis and transcriptomic identity prediction tasks.
In computational neuroscience, considering behavior information as privileged information to guide neural dynamics modeling remains understudied. This work develops novel learning paradigm that represents an inaugural step toward addressing this underexplored research question and advancing the field.
METHODS
This section introduces the details of our proposed BLEND framework, i.e. the behavior- guided neural population dynamiis modeling via privileged knowledge distillation. We start with problem formulation of eehavior-guided neural dynamics modeling in Section[.1] followed by an exposition of the BLEND algorithm in Section 3.2 Detailed exploration of the effectiveness of the BLEND framework can be found in Appendix A.8 and A.9 including empirical studies implications of behavior guidance and how to choose distillation strategies for different models and data.
3.1 PROBLEM FORMULATION
We concretize the learning problem in this study with neural spiking data. For each trial of the input neural activity, let x E X = NXT represent the input spike counts, where x denotes the spike count for neuron at time Let b E 5 R BxT be the corresponding behavior signal, with by denoting the behavioral signal at time t. Without loss of generality, we assume that behavioral signals are temporally continuous and have the same number of time steps as the neural activity As illustrated Fig. μ the model fo aims to reconstruct the randomly masked neural activity solely based on the unmasked portions of x or together with the auxiliary behavior signal b. The optimization objective for masked time-series modeling (MTM) can formulated as follows:
0" = arg min
Σ m Lrec (Jo (xm. xx?) 1.t
LMTM
(1)
where m E (0. NxT is a binary mask with indicating masked elements, Xmm E Xmm represents the unmasked portions x, m the number of masked elements, and Lree is either Poisson or Cross-Entrope loss as specified in the model configuration.
3.2 BLEND
The goal of this work is to create a model that can make predictions using only neural activity data during deployment/inferrnce,, but also benefits from behavior data during training. This subsection details the algorithm of BLEND, including the formulation of privileged knowledge distillation un- der the context of neural dynamics modeling, teacher-student architecture, and distillation strategies.
BLENND can be readily extended to accommodate scenarios wherein hehavioral information is discrete non-temporal in nature, preliminary explorations in Appendix A.10
4

---

Published as a conference paper at ICLR 2025
We first conceptualize the behavior-guiee guiiiddrrggeddddd neural population dynamics modeling problem under the framework of LUPI:
Definition (Privileged Knowledge (Yag et al. 2022)). Consider a general learning problem with imput x X' and label y y. If there is an additional source of information b B that exists during training but not inference, we say b is the privileged knowledge if and only if [y;; b|x)= H (y(x,b)) O.
H Hyyx)
In Definition 1 7(-I-) and H(-I-) represent conditional mutual information and conditional entropy, respectively. According to Definition | the behavior information b serves as the privileged knowl- edge and offers additional predictive capabilities of label y, which is the unmasked neural activity in our scenario. Next, we introduce the proposed privileged knowledge distillation framework, which is designed incorporate behavioral guidance and facilitate neural dynamics modeling.
3.2.1 PRIVILEGED KNOWLEDGE DISTILLATION WITH BEHLVVIOR AS GUIDANCE
As is illustrated in Fig. 2 BLEND employs teacher-student architecture to realize our goal: by leveraging the privileged information the training phase, learn an LVM model for inference phase that outperforms those built on the regular feature x alone. This procedure is divided into two sequential stages:
Train teacher model with regular information and privileged information. The teacher model far E {f Xm x B X3 takes both masked neural activity recordings x and the corresponding behavior signals b inputs. The model parameter 0r is optimized by minimizing the MTM loss LMTM formulated in Eq. ①..
Stage 1:
Stage 2: Train student model with regular information by distillation. By exploiting all available sources information, the teacher model acquires a rich understanding the neural population dynamics. We then aim to distill the learned knowledge to the student model fas {sfs Xm + X). which takes the neural activity as input. The parameter student model 0g is updated by minimizing the following objective:
Ex,h.m
a m
Σ 1.6
m
free (Jos (xm)) x3) +(1 +(1- -x) (x) Σ
Ldissin (Jos (xm)'e : (xm, b)}
8,t
MIM
Distillation Tooss
where a E (0, denotes the mixing ratio between the MTM loss and the distillation loss. The entire
optimization procedure of the proposed BLEND framework is presented in Alg. [] After training the student model with privileged knowledge distillation for MTM, it is applied to downstream tasks where only neural activity recordings available for inference (shown in Fig. 2 right).
While the above algorithm outlines the general framework of BLEND, different implementations of the distillation loss Edisill can capture various aspects of the teacher's knowledge. In this study, investigate the following four main distillation strategies in our BLEND framework, each designed transfer certain aspects of the teacher's knowledge the student model.
Hard distillation. As a baseline, we first implement the Edistn with hard distillation strategy, which directly minimizes the mean squared error (MSE) between the teacher and student outputs:
Caistill (fos (xm))s for (Xm, b): = .fs (xm)) for (x b
(3)
Soft distillation. This approach distills knowledge by matching the softened probability distribu- tions of the teacher and student models. We use temperature parameter soften the logits before applying the softmax function:
Ldistill (Tos (xm)i Joy (Xm, b))) = KL (a (Sor (Xm, b)f (1
11o
for (Xm)
where o is the softmax function, KL(-II-) ) the Kullback-Leibler divergence, and 7 is the tempera- ture parameter.
5

---

Published as a conference paper at ICLR 2025
Feature distillation. This method aims align all the intermediate representations of the student model with those of the teacher model.
Ldissill (Jos (xm)'s for (Xm; b):)
Σ [| (xm)i Ser (Xm; b):ll
(5)
where fig and Soy denote the outputs of the I-th layer of the student and teacher models, respectively. Correlation distillation. This approach focuses on preserving the correlation structure of the teacher's 's outputs in the student model. For each sample in the batch, we compute the correlation matrices of the teacher and student outputs and minimize their difference:
Caisuil (Fos (xm)) for (xm: b)i)
Σ j=1
(o - Corr(fe, (Xm: b)>>((
= B
(6)
where B denotes the batch size, and Corr(.) computes the correlation matrix of the outputs for the j-th sample in the batch (see Appendix A. for details).
4 EXPERIMENTS
As shown in Fig. 2 this work includes two benchmarks for evaluating our proposed BLEND frame- work. The first a public benchmark for neural latent dynamics model evaluation from Pei et al. (021), named Neural Latents Benchmark'2 21 (NLB'21). The second is recent, public multi-modal neural dataset from Bugeon et al. (202)), which contains calcium imaging recordings of neural pop- ulation activity as well as single-cell spatial transcriptomies of the recorded tissue. Details of these two benchmarks, tasks, and included baselines are introduced in Section and
4.1 NEURAL POPULATION ACTIVITY ANALYSIS ON NLB'2J BENCHMARK
We include three sub-datasets from NLB'21, i.e., MC-Maze, MC-RTT, and Area2- Area2-Bump, for a se- ries of neural dynamics modeling evaluations. These datasets contain neural recordings from mon- keys performing various reaching tasks: MC-Maze features delayed reaches through virtual mazes, MC-RTT involves continuous reaches random targets without delays, and Area2-Bump includes reaches with occasional mechanical perturbations (see Appendix AA.2] for details). To evaluate the capability of our proposed BLEND neural population dynamics modeling, adopt three tasks from NL. 21, i.e., neural activity prediction, behavior decoding, and matching to peri-stimulus time histograms (PSTHs).
Neural activity prediction task aims to predict the neural activity of held-out neurons, measured by metric Co-bps, computed the log-likelihood of held-out neurons' activity (see Appendix A.3 for detailed computation steps).
Behavior decoding task requires the model relate neural activity observed behavior: To eval- uate this task, we first train a ridge regression model to predict behavioral data from neural firing rates in the training set. We then use this model to predict behavior from neural activity in the test set and measure the accuracy of these predictions using the R² score (named Vel-Rt2).
Match to PSTHs task aims to evaluate how well models can capture stereotyped features of neu- ronal responses across repeated trials of the same condition. To evaluate it, we calculate the R² between model-predicee
averagee rates and true PSTHs for each neuron across all conditions, then average these R²2 values across neurons (named PSTH-R(2)
We include two types of models as baselines for this benchmark. The first solely uses neural activity recordings as inputs, including LRNN (Stolzenburg et al. 2018). LEADS (Pandarinath et al.) 2018), TrannfDataTrannsfommrr (NDT) (Ye & Pandarinath 2021 l), STNDT (L Shlizerman 2022), and
MINT (Perkins et al.) 2023). The second type utilizes behavior information during training prior to facilitate the learning process of neural dynamics modeling, including pi- VAE (Zhou & Wei. 2020), PISD (Sani et al. 2021), and TNDM (Hurwitz et al. 2021). We choose NDT and LFADS as the base model for privileged knowledge distillation using behavior information, with best-distilled model performance reported in Tab. See details of model configurations in Appendix A.4

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of the proposed BLEND framework with other neural dynamics modeling
methods on NLB'2 21 Benchmark (Pei et al. 2021 Bold values denote the best performance for the corresponding metric.
Methods
MC-Maze Vel-It
MC-RIT Vei-It"
Area2-Bump
Co-bps
PSTH-R
Vel-R-
PSTH-R
Notaal Deenling
LRNN
Stokzenbary
2018
0:148
0.317
0.473
MINT
Parkins et al. 2023
0:181
0.646
0.166
0.159
0.370
0.103
STNDT
Le Shlizerman
2022
0.282
0.773
0.585
0.233
0.563
0,416
Belaviic Prio
Zhou
2020
0:214
0:434
RNN PSED (Sani el al.
2021
0.229
0.683
0.199
0.287
0.514
0.372
TNDM (Hurwitz et all
2021
0:218
0,730
0,509
0.345
0,677
0,402
NDT Ye & Pandarinath 2021
0.275 (+0.0%)
0.729 (+0.0%%
0.051 (+0.0%)
0.318 (+0.0%%
0.519 (+0.0%)
0.250 (+0.0%%
BLEND
NDTT-Disiil-Bett (Ours)
0.310 (+12.7%)
0.801 (+-1.44%)
0.302 [+7.445)
0:372 (+17.00%)
0.788 (+51.89))
0.483 (++06.65)
LFADS (Pandarinath et al. 2018
0.515 (+0.0%)
0.858 (+00.05)
0.579 (+0000%)
0:436 (+-0.095)
0.649 (+00.0%
0.425 (+-0.0%)
LFADS-Distil-Bee Best (Ours)
0.531 (+1.99)
0.877 (+2.256)
0.604 (+4.33%
0.429 (+3.15)
0.837 (+20 0%)
0.615 (+44.7%)
4.2 TRANSCRIPTOMIC NEURON IDENTITY PREDICTION ON MULATI-MDDAL NEURAL ACTIVITY DATASET
This multi-modal dataset combines calcium imaging from mouse primary visual cortex (V1) with single-cell transcriptomics. The dataset includes functional recordings from 9728 neurons across 17 sessions from 4 mice (SB025, SB026, SB028, SB030). Additionally, the dataset provides tran- scriptomic profiles that are used to categorize neurons either excitatory or inhibitory. For mouse SB025, about half of the neurons classified as inhibitory are further subdivided into specific subtypes (see details of this dataset in Appendix A.2) We evaluate models the following two tasks under both multi-animal scenario (all 4 mice) and single-animal scenario (only mouse SB025):
Excitatory/nhibitory (EI) neuron identity prediction task is binary classification task that quires the model to predict the neuron identity, i.e., excitatory or inhibitory.
Subclass label prediction task is 5-class classification task, where the model is required to predict the subclass label of inhibitory neurons, i.e., Lamp5, Pvalb, Vip. Sneg, or Sst.
Following the settings in Mi et al. 2023) we include random model, PCA (Cunningham & Yu 2014). UMAP (Mcinnes et al.) 2018), LOLCAT (Schneider et al., (022a) and its variants, and Ne-
uPRINT (Mi al., [202) as the baseline models for this benchmark. We choose NeuPRINT as the base model for privileged knowledge distillation using behavior information with best-distilled
model performance reported in Tab. Detailed model configurations are in Appendix A..4
RESULTS AND ANALYSIS
5.1 BENCHMARK 1: NEURAL POPULATION ACTIVITY ANALYSIS
Tab. summarizes the results of neural activity prediction, behavior decoding, and match to PSTHs tasks of baseline methods and our proposed method MC-Maze, MC-RTT, and Area2-Bump datasets. Our behavior-guided distilled NDT and LFADS models achieve the best performance and outperform SOTA methods on all tasks and datasets, demonstrating the effectiveness of our proposed BLEND mechanism.
Fit to neural activity. For neural activity prediction of held-out neurons, NDT-Distill-Be achieves 0.310 Co-bps score on the MC-Maze dataset, reporting 12.7% improvement over the base model NDT. LEADD-Disttll--Bes achieves 0.321 Co-bps score, reporting the best performance for this task against other baseline models. These observations indicate that our models gain better under- standing of neural population dynamics and yield improved predictions of held-out neural activity.
Fit to behavior. For behavior (monkey hand velocity) decoding, ND'T-DistilI-ect achieves Vel-R² of 0.891, 0.372, and 0. 788 respectively, outperforming the compared baselines. A 51.8% improve- ment is observed on the Area2- dataset over the base model NDT. While for LFADS-Distle Best, Vel-R² scores of 0.877, 0.429, and 0.837 are reported, showcasing better behavior decoding capabilities over baselines and the distilled NDT. Our findings demonstrate that the proposed method achieves a superior correlation between neural activity and behavioral signals compared to existing models. Note that the improvement of distilled LFADDS over the hase model MC-Maze and MC-
7

---

Published as a conference paper at ICLR 2025
Table 2: Top-I accuracy of exciaatryyinhibitory neuron identity prediction and inhibitory neuron subclass label prediction tasks the multimodal population activity dataset from Bugeon et al. 2002). Bold values denote the best performance for the corresponding metric.
Multiple Animals
Single Animal
Methods
EI (2 class)
Subclass (5 class)
EI (2 class)
Subclass (5 class)
Random
0.523
0.302
0.488
0,260
PCA Cunningham & Yu, 2014)
0.565
0.330
0.572
0.263
UMAP Mellnnes al. 2018
0.520
0.340
0.438
0.281
LOLCAT Schneider et al.
2023a,
0.600
0.404
0,608
0.474
LOLCATIsh Mi al., 2023
0.640
0.474
0.632
0.491
LOLCATRaw Mi et al. 2023
0.664
0.439
0.616
0.386
NeuPRINT (ii et al. 2023
0.748
(+0.0%%
0.495
(+000%%
0.667
(+-0.%%
0.508
(+-0.0%
NeuPRINT-Distil--en (Ours)
0.789 (+5.5%)
0.571 (+ 15.4)
0.722 (+8.2%)
0.588 (+15.7%%
RTT is not as significant distilled NDT. The potential reason behind this is that LFADS, being RNN-based, may reach its capacity limits with larger datasets, making it harder to incorporate addi- tional knowledge effectively. Meanwhile, LFADS performs dimensionality reduction part of its model, which might limit its ability to incorporate additional information from the teacher in larger datasets where the intrinsic dimensionality is already well-captured.
Match to PSTHs. In the PSTH matching task, our distilled NDT attains PSTH-R scores of 0.592 and 0.483 on the MC-Maze and Area2-Bump datasets, respectively. These results represent im- provements of 7.4% and remarkable 66.6% over the baseline model performance. The distilled LFADS model achieves PSTH-R scores of 0.604 and 0.615, demonstrating superior performance among all evaluated models and yielding improvements of 4.3% and 44.7% over the base model, re- spectively. These findings indicate that leveraging behavioral data as privileged knowledge enhances the models' capacity to reproduce stereotyped neural responses, quantified by PSTHs.
Overall, the BLEND framework provides a simple yet effective method that could directly improve the performance of existing NDM models in neural population activity analysis tasks.
5.2 BENCIMARK 2: TRANSCRIPTOME NEURON IDENTITY PREDICTION
Tab. shows the experimental results of transcriptomic identity prediction tasks, including El neu- ron identity prediction and inhibitory neuron subclass label prediction. These two tasks are con- ducted in both multi--nimal and single-animal settings. By incorporating our BLEND framework, the distilled model shows remarkable improvement over the base model and outperforms all baseline methods.
El prediction. In the exciaaooryinhiiitory neuron identity prediction task, our distilled NeuPRINT model attains the top-1 accuracy of 0.789 and 0.722 in multi-animal and single-animal settings, respectively. These findings demonstrate improvements of 5.5% and 8.2% over the base model. No- tably, our proposed method exhibits superior performance, surpassing all other baseline approaches included for comparison.
Subclass label prediction. For the inhibitory neuron subclass label prediction task, the distilled NeuPRINT model improves the base model by 15.4% and 15.7% in multi-animal and single-animal settings, respectively, achieving the best top-1 accuracy of 0.571 and 0.588. Given that this task in- volves a five-class classification problem, the model's performance metrics are comparatively lower than those observed in the El prediction task, which is binary classification problem.
As illustrated in Table the model's 's performance metrics in the single-animal setting are compara- tively lower than those observed in the multi-animal setting. We posit that the multi-animal scenario provides a more diverse and extensive set of neural activity inputs, thereby enhancing the model's generalization capabilities. Besides, Mi al. (2023) points data-limited regime in this bench- mark characterized by scarce labeled data for transcriptomic prediction, with only a small fraction of neurons having both calcium imaging and transcriptomic data, and an imbalanced distribution across subclasses. Our BLEND framework alleviates this problem and improves the model performance since the teacher model provides regularization and rich knowledge transfer to the student.
8

---

Published as a conference paper at ICLR 2025
Table 3: Ablation study on privileged knowledge distillation strategies evaluated by neural dynamics modeling tasks of NLB'21 Benchmark (Pei et al. 2021) and cell class prediction tasks of Bugeon ct al. (2022) dataset. Bold values denote the best performance for the corresponding metric.
Neural Population Activity Analysis
Methods
MC-Maze
MC-RIT
Area2-Bump
Co-hps
Vel-R
PSTH-R²
Vel-R²
Vel-R?
PSTH-R
NDT Ye & Pandarinath 2021
0.275 (+0.0%)
0.779 (-0.0%)
0.551 (+0.0%)
0.318 (+0.0%)
0.519 (+0.0%) 0.290 (+0.0%)
NDT-Hard-Distil (Ours)
0.290 (+5.5%)
0.890 (+1122%)
0.587 (+6.5%)
0.380 (+19.59)
0.717 (-19.9%
0.303 (+25.2%)
NDT-SofI-Diaae (Ours)
0,259 (-5.8%)
0.881 (+13.19)
0.191 (-10.9%)
0.333 (+11.69)
0.714
(-13.4%)
0.318 (+9.7%)
NDT-Feature-Distie (Ours)
0.303 (+10.2%)
0,888 (+111.0%)
0.596 (+8.29)
0.381 (++19.893)
0.787
(-51 6%)
0.481 (465.9%)
NDT-Correlation- Distill (Ours)
0.310 (+12.7%)
0.891 (+14.4%)
0.502 (+7-49)
0.372 (+117 (05)
0.788
(-51.8%)
0.483 (+66.6%)
LEADS Pandarinath nl. 2018
0.315 (1 (10.0%)
0.858
(0.0%)
0.579 0.0%%
0.416 (.00%)
0.649
(10.0%)
0:425 [ (.00%)
LFADS-Ilar--Disill (Ours)
0.321 (+1.9%)
0.877 ( 2.2%%
0.604 (++4.3%)
0.429 (..15)
0.837 (+29.00% 0.615 (+44.7%%
LEADS (Ours)
0.282 (- 10.5%)
0.500 6.8%)
0.516 (- 10.9%%
0.443 (++6.5%)
0.740 14.0%) 0.578 (135.5%)
EI Neuron Identity & Inhibitory Neuron Subclass Label Prediction
Multiple Animals
Single Animal
Methods
EI class)
Subclass (5 class)
EI (2 class)
Subclass (5 class)
NeuPRINT (Mi et al. 2023
0.748 (+0.0%)
0.495 (+0.0%)
0.667 (+0.0%
0.508 (-0.0%)
Nend (Ours)
0.789 (+5.55)
0.571 (+15.4%)
0.714 (+7.0%
0.581 (+ 14.4%)
NeuPRINT-Soft-Dish (Ours)
0.782 (+-4.5%)
0.541 (+9.3%)
0.722 (+8.2%)
0.588 (+15.7%)
NeuPPINNT-Faatur--iitttll (Ours)
0.781 (+-4.4%)
0.531 (+7.3%)
0.718 (+7.6%)
0.524 (+3.6%)
NeuPRINT-CCorrclatonn Distill (Ours)
0.792 (+5.9%)
0.537 (+ (+-2..5%)
0.705 (+5.7%)
0.570 (+12.2%%
5.3 ABLATION STUDIES ON BEILAVIOR-GUDDER KNOWLEDGE DISTILLATION STRATEGIES
In this work, we explore different privileged knowledge distillation strategies, including hard distil- lation, soft distillation, feature distillation, and correlation distillation, as detailed in Section 3.2.1
Neural population activity analysis on NLB'21 benchmark. We find that the NDT-Coreeltion- Distill model generally performs best across most metrics, showing significant improvements over the baseline NDT. For MC-Maze, NDDT-Coreeltiion-Dise
achieves the highest Vel-R (0.891). While for LFADS and its distilled variants, LFADS-Hare Hard- -istill shows improvements over the base- line LFADS, particularly in PSTH-R" for MC-Maze and Vel-R.2 for Arca2- -ump.. Notably, except for LLFADS-HarddDiin
other distillation methods lead to performance degradation for LFADS on MC-Maze (hence, only results for Soft-Distil are presented in Table This contrasts with the NDT results, where multiple distillation strategies show improvements, indicating that not all knowledge transfer techniques are universally beneficial and highlighting the importance of careful method selection when applying distillation techniques to different models or tasks.
Transcriptomic identity prediction on multi-modal calcium imaging dataset. It can be seen that NeuPRINT-Hard-Diss and NeuPRINT-Soft-Dine show consistent improvements over the base- line NcuPRINT across all tasks. El prediction. SeuPRINT-Soft-Distti excels in Single Animal predictions for both El and Subclass.
performs best for Multiple Animals
Overall, the distillation strategies, particularly Hard-Distill and Correlation- -iittill,, tend to outper- form their respective baselines. suggesting that privileged knowledge distillation can significantly enhance neural dynamics modeling. Different distillation strategies seem to be more effective for different tasks and metrics, indicating that the choice of distillation method should be task-specific.
5.4 QUALITATIVE ANALYSIS
We visualized the decoded behavior trajectories of the base model and our distilled model for com- parison in Fig. 3 Subfigure (a) shows the comparison between predicted and ground truth 2D hand movement on MC-Maze for NDT (base model) and NDT-Hard -Dstilll (our distilled model). And subfigure (b) illustrates the 1D hand velocity over time for both the X and Y axes separately. For more comprehensive visualization and qualitative analysis, please refer to Appendix A..5 and
2D hand movement trajectory decoding. For NDT the base model, the predicted trajectory follows the general U-shape of the ground truth but shows sigmificant deviations, especially at the bottom of the U and the upper right comer. While our distilled model generates trajectory that aligns much
9

---

Published as a conference paper at ICLR 2025
Thail 1
Velocity
Velcity
10000
Ground Trush
200
750
Ground Truth Prodicted
Predictad
Yookonn
250
NDT
-200
-10000
-100
-15000
-20000
-600
-500
10000 20000
-750
X Position
Time Step
Time Steg
NDT-Hard- Dlsii (Ours)
Trial
Tiol 4 Velecity
Tisl 4-Y Velocity
10000
5000
Prodicted
Ground Trush
200
302
750-
Distil
-200
Uecocity
290
-15000
-20000
-600
-500
10000
20000
-500
90
100
120
240
Time Step
(a) 2D Trajectories
(b) Velocity Trajectories
Figure 3: Visualization of behavior decoding on MC-Maze dataset. (a) Prediction and ground truth 2D hand movement trajectory. (b) Prediction and ground truth of X and Y velocities, respectively.
more closely with the ground truth. It captures the U-shape more accurately, especially at the bottom curve and the endpoints, showing a marked improvement in behavior decoding.
1D hand velocity decoding (X and Y axes). For both X and Y velocities, the base model struggles to accurately capture the magnitude and timing of velocity changes in both dimensions. Specifically, it underestimates peak velocities and shows erratic fluctuations in X velocity, and shows substan- tial deviations from ground truth in Y velocity. particularly after time step 60. While our distilled NDT significantly enhances velocity predictions in both X and Y directions. It more accurately captures the magnitude, timing, and overall pattern of velocity changes with improved smoothness, suggesting a better understanding of the hand movement dynamics.
CONCLUSION AND DISCUSSION
In this study, we introduce BLEND, a novel framework for neural population dynamics modeling that leverages privileged knowledge distillation. Predicated on the fundamental premise that be- havioral data can serve as an explicit guide for neural representation learning, we present a model- agnostic framework that facilitates seamless integration with diverse existing neural dynamics mod- cling architectures. This versatility enables the enhancement of wide range of computational ap- proaches in the field of computational neuroscience. Comprehensive empirical evaluations, encom- passing neural population activity analysis and transcriptomic identity prediction tasks, substantiate the superior effectiveness of our proposed framework. BLEND demonstrates a remarkable capac- ity for capturing implicit patterns within neural activity, consistently outperforming sate-of-thheard models by a significant margin, thereby providing new perspectives on how behavioral observations can be leveraged to guide the complex modeling of neuronal population dynamics.
Although the initial results of BLEND are promising, a few limitations remain to be addressed. In this work, employ straightforward approach to implement the behavior guidance by con- catenating the neural activity and behavior information along the feature dimension (assuming they have the same length of time steps). This methodological choice is motivated by the relatively low-dimensional nature of behavioral signals, which precludes their use as an independent guidance source (e.g., as query in the aateetttnnn mechanism). However, this simplistic guidance strategy may not fully capture the intricate re relationship between neural activity and behavior. Future research should investigate advanced integration methods better utilize the complex interactions between these data types. potentially improving the model's ability to represent subtle neural patterns. More- over, while the current study focuses on temporal behavioral signals that correspond directly with neural activity at each time step, future research should aim to extend BLEND to more generalized settings, as we initially explored in Appendix (. A0 This expansion could encompass non-temporal or discrete behavioral signals, thereby broadening the framework's applicability to diverse neuro- scientific domains. Such extensions could prove valuable in investigating sleep stages, categorizing social behaviors, and exploring other scenarios where behavior is not continuously paired with neu- ral activity. This generalization would significantly enhance the versatility and utility of BLEND across a wider spectrum of neuroscience research paradigms.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by the Hong Kong Innovation and Technology Commission (Project No. GHP/006/22GD and ITCPD/17-9), HKUST (Project No. FS111), and the Research Grants Council of the Hong Kong Special Administrative Region, China (Project No. T45-401/22-N).
REPRODUCIBILITY STATEMENT
To enhance the reproducibility of this study, we provide an Appendix section comprising six sub- sections that offer detailed supplementary information. Appendix A.1 presents the pseudo-code of our proposed BLEND framework, followed by a comprehensive explanation of the behavior-guide knowledge distillation strategies. Additionally, Appendix [A.2 provides detailed descriptions of the datasets utilized, namely MC-Maze, MC-RIT, Area2-Bump, and the multi-modal neural activity datasets. To facilitate deeper understanding of the tasks conducted in this study, Appendix AA.3 elucidates the specifics of neural activity prediction, behavior decoding, and match to PSTHs tasks. Subsequently, Appendix [A.4 delinnatss the model configurations, including architecture and hyper- parameters. To further elucidate the superior performance of our proposed method, Appendices A.5 and A offer additional visualizations comparing different base models for distillation and vari- ous distillation strategies, enabling qualitative assessment. Our source code will be made publicly accessible upon acceptance of this paper.
REFERENCES
Muhammad Haseeb Aslam, Muhammad Osama Zeeshan, Marco Pedersoli, Alessandro Koerich, Simon Bacon, and Eric Granger. Privileged knowledge distillation for dimensional emotion recog- nition the wild. In Proceedings of the IEEF/CVÉ conference on computer vision and pattern recognition, pp. 3338-3347, 2023.
Stephane Bugeon, Joshua Duffield, Mario Dipoppa, Anne Ritoux, Isabelle Prankerd, Dimitris Ni- coloutsopoulos, David Orme, Maxwell Shinn, Han Peng, Hamish Forrest, al. A transcriptomic axis predicts state modulation of cortical interneurons. Nature, (07((918)3330-1)
2022.
Cheng Chen, Qi Dou, Yueming Jin, Quande Liu, and Pheng Ann Heng. Learning with privileged multimodal knowledge for unimodal segmentation. JEEE transactions on medical imaging, (3):621-632) 2021.
Raeed H Chowdhury, Joshua I Glaser, and Lee E Miller. Area 2 of primary somatosensory cortex encodes kinematics of the whole arm. Elife, Pre48198, 2020.
Mark M Churchland, John Cunningham, Matthew T Kaufman, Stephen I Ryu, and Krishna V Shenoy. Cortical preparatory activity: representation of movement first cog in dynamical machine? Neuron, 68(3):387-400, 2010.
Mark Churchland, John Cunningham, Matthew Kaufman, Justin Foster, Paul Nuyujukian, Stephen Ryu, and Krishna V Shenoy. Neural population dynamics during reaching. Nature, 487 (7405):51-56, 2012.
John Cunningham and Byron M Yu. Dimensionality reduction for large-scale neural recordings. Nature neuroscience, 17(11):1500-155 2014.
Patrick J Drew, Celine Mateo, Kevin L Turner, Xin Yu, and David Kleinfeld. Ultra-slow oscil- lations fmri and resting-state connectivity: neuronal and vascular contributions and technical confounds. Neuron, 107 804, 2020.
Yuanjun Gao, Evan W Archer, Liam Paninski, and John P Cunningham. Linear dynamical neural population models through nonlinear embeddings. Advances neural information processing systems, 29, 2016.
Cole Hurwitz, Akash Srivastava, Kai Xu, Justin Jude, Matthew Perich, Lee Miller, and Matthias Hennig. Targeted neural dynamical modeling. Advances in Neural Information Processing Sys- tems, 34:29379-22399.. 2021.
11

---

Published as a conference paper at ICLR 2025
Justin Jude, Matthew G. Perich, Lee E. Miller, and Matthias H. Hennig. Robust alignment of cross- session recordings of neural population activity by behavior via unsupervised domain adaptation. International Conference Machine Learning, 2022.
James Jun, Nicholas A Steinmetz, Joshua H. Siegle, Daniel J Denman, Marius Bauza, Brian Barbar- its, Albert K Lee, Costas A Anastassiou, Alexandru Andrei, Çağatay Aydın, et al. Fully integrated silicon probes for high-density recording of neural activity. Nature, 55 551(7679):322-33, 2017.
Dmitry Kobak, Wieland Brendel, Christos Constantinidis, Claudia E Feierstein, Adam Kepecs, Zachary F Mainen, Xue-Lian Qi, Ranulfo Romo, Naoshige Uchida, and Christian K Machens. Demixed principal component analysis of neural population data. elife, 5:e10989, 2016.
Trung Le and Eli Shlizerman. Stndt: Modeling neural population activity with spatiotemporal trans- formers. Advances in Neural Information Processing Systems, 35: 19926-17999, 2022.
Liyang Liu, Zihan Wang. Minh Hieu Phan, Bowen Zhang. Jinchao Ge, and Yifan Liu. Bpkd: Boundary privileged knowledge distillation for semantic segmentation. In Proceedings of the IEEF/CVF Winter Conference on Applications of Computer Vision, pp. 1062-1072, 2024.
Ran Liu, Mehdi Azabou, Max Dabagia, Jingyun Xiao, and Eva Dyer. Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers. Advances in neural information processing systems, 35:2377-2391, 2022.
Jakob H Macke, Lars Buesing, John P Cunningham, Byron M Yu, Krishna V Shenoy, and Maneesh Sahani. Empirical models of spiking in neural populations. Advances in neural information processing systems, 24, 2011.
Joseph G Makin, Joseph E O' "Doherty, Mariana MB Cardoso, and Philip N Sabes. Superior arm- movement decoding from cortex with a new, unsupervised-leaming algorithm. Journal of neural engineering, 15(2)0026010, 2018
Leland Mclnnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint 180.0:180. 2018.
Lu Mi, Trung Le, Tianxing He, Eli Shlizerman, and Uygar Sümbül. Learning time-invariant rep- resentations for individual neurons from population dynamics. Advances in Neural Information Processing Systems, 36:46007-46026. 2023.
Erfan Nozari, Maxwell A Bertolero, Jennifer Stiso, Lorenzo Caciagli, Eli Combblath, Xiaosong He, Arun S Mahadevan, George J Pappas, and Dani S Bassett. Macroscopic resting-state brain dynamics are best described by linear models. Nature biomedical engineering, 8(1):68-844 (468-84,777 2024.
Joseph O'Doherty, Mariana MB Cardoso. Joseph G Makin, and Philip N Sabes. Nonhuman primate reaching with multichannel sensorimotor cortex electrophysiology. Zenodo http://doi. org/10.520 5281/zenodo, 583331, 2017.
Chethan Pandarinath, Daniel J O' O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg, et Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10):805-81) 815, 2018.
Felix Pei, Joel Ye, David M. Zoltowski, Anqi Wu, Raeed H. Chowdhury, Hansem Sohn, Joseph E. O'Doherty, Krishna V. Shenoy, Matthew T. Kaufman, Mark Churchland, Mehrdad Jazayeri, Lee E. Miller, Jonathan Pillow, I1 Memming Park, Eva L. Dyer, and Chethan Pandarinath. Neu- ral latents benchmark '21: Evaluating latent variable models of neural population activity. In Advances in Neural Information Processing Systems (NeurIPS), Track Datasets and Bench- http/////////./ org/ abs/ /atp://11...0 04463
marks, 2021. URL hattts://taaattt
Sean M. Perkins, John P. Cunningham, Qi Wang, and Mark M. Churchland. Simple decoding of behavior from complicated neural manifold. elifee, October 2023. doi: 10.7554/elife.89444
Omid G Sani, Hamidreza Abbaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification. Nature Neuroscience, 24(1): 140-1999, 2021.
12

---

Published as a conference paper at ICLR 2025
Aidan Schneider, Mehdi Azabou, Louis McDougall- Vigier, David F Parks, Sahara Ensley, Kiran Bhaskaran-Nair, Tomasz Nowakowski, Eva Dyer, and Keith B Hengen. Transcriptomic cell type structures in vivo neuronal activity across multiple timescales. Cell reports, 42(4), 2023a.
Steffen Schneider, Jin H Lee, and Mackenzie W Mathis. Learnable latent embeddings for joint be- havioural and neural analysis. Nature, 2023b. doi: https://ooi.org/ 10. 10338//41558--022--66633 1-6,
lan H Stevenson and Konrad P Kording. How advances in neural recording affect data analysis. Nature neuroscience, 14(2):139-1 2011.
Frieder Stolzenburg, Sandra Litz, Olivia Michael, and Oliver Obst. The power of linear recurrent neural networks. arXiv preprint arXiv:/802.033300 2018.
Vladimir Vapnik and Akshay Vashist. new learning paradigm: Learning using privileged infor- mation. Neural networks, 22(5-6):544-557, 2009.
Shuo Yang, Sujay Sanghavi, Holakou Rahmanian, Jan Bakus, and Vishwanathan SVN. Toward understanding privileged features distillation learning-to-rame Advances Neural Information Processing Systems, 55:26658-26677 2022.
Joel Ye and Chethan Pandarinath. Representation learning for neural population activity with Neural Data Transformers. Neurons, Behavior, Data analysis, and Theory, August 2021. ISSN 2690- 2664. doi: 10.51628/001c.27355..
Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of high- dimensional neural activity using pi-vae. Advances in Neural Information Processing Systems, 33:7234-7247 2020.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 SUPPLEMENTARY CONTENTS OF BLEND ALGORITHM
A.1.1 BLEND ALGORITHM
Algorithm Behavior Behavior-guided guided Maacher-Student Knowledge Distillation Framework (BLEND)
Require:
Training data = {(x;, b;)], Teacher model fo,, Student model fag
Ensure: Trained student model Jog
1: // Train teacher model
2: Initialize teacher model parameters 0r
3: for each epoch do
4: for each batch (x,, b,) D do
5: Generate random mask m, p(m)
6: Apply the generated mask m, original input X; Create unmasked portions Xms
7:
8:
No 4- [og (Xm; by)
Teacher's predictions
9:
Compute LMTM using Eq- (T)
10: Update 0r using V.IM
11: end for
12: end for
13: // /raiin student model
14: Initialize student model parameters 0s
for each epoch do for each batch (xs, by) in D do
16:
Generate random mask m, p(m)
18:
Apply the generated mask m, original input Create unmasked portions Xms
20:
2T
(x by)
Teacher's predictions Student's predictions
21:
Es f- fos (Xms)
22:
Compute Laistill using Eq.
23: LKD t aLMIM X + (1 a:) Laisaul 2g, Zr)
Distillation loss Lassill chosen from one of Eq.
(6)
24:
Update Os using LLD
25:
end for
26:
end for
return fes
A.1.2 DISTILLATION STRATEGIES.
Correlation-Diste As shown in Section 3.2 3.2.1 this distillation strategy aims to preserve the corre- lation structure of the teacher's outputs. Concretely, the correlation matrix is computed as:
(Y - Y)(Y Y)T
Corr diag((Y Y) Y)T) diag((Y DOYY-YTT
Here, Y E NN represents the output of either the student or teacher model for a single sample, with N being the number of neurons and T' the number of time steps. Y is the mean of Y along the time dimension, and diag(-) extracts the diagonal of matrix. The resulting correlation matrix has dimensions N x N. Note that we use the Frobenius norm ||=||, to compute the difference between correlation matrices, as it provides scalar measure of the overall difference between these matrices.
A.2 SUPPLEMENTARY CONTENTS OF DATASETS INCLUDED
MC-Maze dataset (Churchland et al. 2010 2012) comprises recordings from the primary motor and
dorsal premotor cortices of monkey performing instructed-delay reaching tasks in 108 different configurations, involving various target positions and virtual barriers. This dataset includes one full session with 2,869 trials and 182 neurons, along with hand kinematics. This dataset is notable for
14

---

Published as a conference paper at ICLR 2025
its behavioral richness, stereotyped repetitions, high trial counts, and clean separation of preparation and execution phases, making it valuable for studying neural population dynamics during movement.
MC-RTT dataset (O'Doherty et al. 2017 Makin et al. 2018) features motor cortical recordings during a random target task, comprising continuous, point-to-poin reaches with variable lengths and locations, without delay periods. It spans 15 minutes of continuous reaching, artificially divided into 1,351 600ms trials, and includes data from 130 neurons along with simultaneous hand kinematics. Unlike the MC-Maze dataset, MC-RTT introduces modeling challenges due its non-stereotyped movements, lack of trial repetitions, and the unpredietability of random data snippets, precluding simple trial-averaging approaches for de-noising. This dataset tests models' ability to infer latent representations from single-trial data and detect unpredictable inputs to population activity, offering a more naturalistic benchmark for latent variable models.
Area2-Bump dataset (Chowdhury et al., 2020) contains neural recordings from area 2 of the so- matosensory cortex, which processes proprioceptive information, as monkey performed a visually- guided reaching task using manipulandum. It comprises 462 trials with data from 65 neurons, along with hand kinematics and perturbation information, where in 50% of random trials, the mon- key's arm was unexpectedly bumped before the reach cue. This dataset challenges models to in- fer inputs describing activity after sensory perturbations and to perform robustly with low neuron counts, offering insights into the distinct dynamicc of somatosensory areas compared to motor areas.
For each of the above datasets, neural activity recordings are counted in 5ms bins, and behavior signals also measured 5ms intervals. Other preprocessing information be found in Pei et al. 2021 or Ye & Pandarinath (2021)
Multi-model neural activity dataset collected by bugune et al. (2022) ) combines neural recordings, genetic information, and behavioral data from mice. This comprehensive dataset includes calcium imaging from the primary visual cortex of four mice (SB025, SB026, SB028, SB030), encompass- ing over 9,700 neurons across 17 recording sessions. Each session captures roughly 500 neurons for about 20 minutes at a sampling rate of 4.3 Hz. The genetic component comprises expression data for 72 specific genes, enabling the classification of neurons into excitatory or inhibitory types, with further subtyping available for portion of inhibitory neurons in one mouse (SB025). Comple- menting the neural data are behavioral recordings such as running speed and pupil dilation, as well as overall brain state classifications. The dataset also provides spatial information for the recorded neurons, offering rich resource for investigating the interplay between neural activity, cell types, and behavior.
For the preprocessing details of this dataset, please refer to Mi al. 2023
A.3 SUPPLEMENTARY CONTENTS OF TASKS INCLUDED
Neural activity prediction. This task requires the model to predict neural activity for held-out neurons. Co-bps, also named Co-smoothing, is the primary evaluation metric used in the NLB'21 to assess how well models can predict held-out neural activity. The test data is split into held-in and held-out neurons, with models using the training data and held-in test neurons predict firing rates À for the held-out test neurons. Performance is measured using log-likelihood under a Poisson observation model:
p(fm.:) Poisson($n.. An,.),
where n,, is the held-out spike count for neuron n at time t. The log- log-liielihood L(}:)) summed over all held-out neurons and time points, then normalized to "bits per spike" by comparing to a baseline model that uses only the mean firing rate of each neuron:
bits/spike Nap log 2
(C(): - L(An;
where Xn, is the mean firing rate for neuron and Rap is its total number of spikes. A positive bits per spike (bps) score indicates the model predicts time-varying activity better than the mean firing rate baseline. The co-smoothing metric allows standardized comparison across different types of
15

---

Published as a conference paper at ICLR 2025
models and neural datasets, as it only requires models to output predicted firing rates rather than imposing constraints on model architecture or training (Pei et al. 2021).
Behavior decoding. As mentioned in Section 4.1 following the settings in Pei ct al. 2021); Ye & Pandarinath (202I). we evaluate the performance of behavior decoding by fitting a ridge regression model. This linear mapping is enforced for all models in the decoding process to prevent complex decoders from compensating for poor neural dynamics estimation. More sophisticated decoders could potentially achieve better performance but the cost of obscuring the quality of the underlying latent representations. The behavioral data used for decoding in MC-Maze, MC-RTT, and Area2- Bump is monkey hand velocity.
Match to PSTHs. This task evaluates how well models can reproduce the stereotyped neural re- sponses captured by peri-stimulus time histograms (PSTHs). PSTHs are computed by averaging neuronal responses across trials within given condition, revealing consistent features of neural ac- tivity. For datasets with clear trial structures (MC-Maze and trea2-Bumph the task computes the Re between trial-averaged model rate predictions for each condition and the true PSTHs, first for each neuron across all conditions and then averaged across neurons.
A.4 SUPPLEMENTARY CONTENTS OF MODEL CONFIGURATIONS
LFADS (Latent Factor Analysis via Dynamical Systems) (Pandarinath et al.) 2018) is a deep learning method designed model neural population dynamics from single-trial spiking data, which utilizes recurrent neural networks (RNNs) to model the underlying dynamics of neural populations. The model consists of an encoder RNN that compresses the input spike data into latent code, and a generator RNN that reconstructs the data from this latent representation. LFADS infers low-dimensional latent factors and initial conditions for each trial, which are then used to gener- ate de-noiser estimates of neural firing rates. This architecture allows LFADS to capture complex, non-linear dynamics in neural data while providing interpretable latent representations of neural population activity on single trials.
As the base model for our privileged knowledge distillation with behavior information in neural dynamics modeling, we set the hidden dimension to 64 and factor size to 32. For model training, we set the batch size to 64, learning rate to 1 10-3 with 5000 warm-up iterations and weight decay to 5 x 10 Mask ratio is set to 0.25.
NDT (NeuralData Transformrr))))) (Ye & Pandarinath, 2021 is non-recurrent architecture designed to model neural population spiking activity, based on the BERT encoder. The core of the NDT con- sists of a stack of Transformer layers (typically layers), each containing self-attention mechanisms, layer normalization, and feedforward neural networks. It uses masked modeling during training, where random portions of the input sequence are masked and the model trained to reconstruct the original input, encouraging it to leverage contextual information. The NDT processes neural data in parallel rather than sequentially, enabling faster inference times compared to recurrent models like LFADS, while achieving comparable performance in modeling autonomous neural dynamics and enabling accurate behavioral decoding.
As the base model for our behavior-guided knowledge distillation in neural dynamics modeling, we set the Transformer layer number to 4 for MC-Maze, MC-RTT, and Area2-Bump datasets, hidden dimension to 128, and number of attention heads to For model training, we set the batch size to 64, learning rate to 1 x 10-3 with 5000 warm-up iterations and weight decay to 5 x 10-5, Mask ratio is set to 0.25.
NeuPRINT (Mi et al. 2023) is designed to assign time-invariant representations to individual neu- rons based on population recordings of neural activity. The model uses transformer architecture to implement implicit dynamical system that predicts neural activity based on the past activity of the neuron itself and cermutation-invariin
statistics of the population activity. NeuPRINT learns both the dynamical model and time-invariant representations for each neuron by minimizing the pre- diction error of future neural activity. The learned representations can then be used for downstream tasks such predicting transcriptomic cell types.
As the base model for our rehavior-guides knowledge distillation in transcriptomic identity predic- tion, we follow the same configurations in Mi et al. (2023) and set the dimension of time-invariant
16

---

Published as a conference paper at ICLR 2025
800
Fal 5- Welsery
Fial Velocil
Trial
600
Rreditted
300
Freditted
20020
Ground Truth
NDT
803
200
200
Position
200
-10000
-20000
Time Stop
100
120
340
20
Time Stop
100
120
545
20020
Position
20000
NDT-Hard- Distil (Ours)
Tall 5 Wilcky
Vial 5- Velocty
Trial
820
Ground Truhh
303
Ground Truh
Peckkted
20000
Ground Predicted
033
200
Fosttoon
10000
200
100
-10000
-00000
-20000
20000
Tame Step
Time
Posittion
Vial S-X Velority
Fial 5 Welocity
Trial
000
Groad tuth Rredited
200
Grounal Tuth
Preditted
Ground Truth Predicted
20000
LFADS
400
Position
-10000
-203
-20000
-20000
20000
Time
Time
Positton
LFADS-Hard Disti--istl (Ours)
Fial
Velocity
S-Y Velocity
Trial
50
Precktedd
20000
800
Positio
420-
-100
0-
-150
-10000
Predicted
100
120
1/1
20
60
300
120
140
-20000
20000
Time Step
Time Step
(a) Velocity Trajectories
(b) 2D Trajectories
Figure 4: Visualization of predicted hand velocity MC-Maze dataset of base model NDT and LFADS, as well as their behavior-guided distilled counterparts. (a) Prediction and ground truth of X and Y velocities, respectively. (b) Prediction and ground truth of 2D hand movement trajectories.
embedding to 64. For model architecture, 1 transformer layer with 2 attention heads used. For model training, the batch size is set to 1024 and the learning rate is set to 1 10-3
A.5 ADDITIONAL QUALITATIVE ANALYSIS ON NDT AND LFADS
This section provides qualitative analysis based on visualization of model predictions in the be- havior decoding task. As shown in Fig. a comprehensive visualization comparing the performance of base models (NDT and LFADS) with their Selavior-guided distilled versions ANDT-Hard-Dislle and Hard-Distill)) on the MC-Maze dataset is presented.
A.
ID HAND VELOCITY DECODING (X AND Y AXES)
Fig.
shows the decoded X and Y hand velocities compared to the ground truths.
The distilled NDT model demonstrates significant improvements in capturing both X and Y velocity components simultaneously. The base NDT model shows consid- erable deviations from the ground truth, particularly in the latter half of the time series. In contrast, NDT--Had--Dishe tracks both velocity components much more accurately throughout the entire
NDT vs. NDT-Hard-Disttin
17

---

Published as a conference paper at ICLR 2025
quence. It better captures the magnitude and timing of velocity changes in both dimensions, resulting in much closer match to the ground truth trajectory.
LFADS LFADS-Hard-Distill Both LFADS and the distilled model perform well in tracking the X and Y velocities, but the distilled version shows noticeable improvements. LFADS- -Hadd- Distill more accurately captures the peaks and troughs of both velocity components, especially in the middle and latter parts of the sequence. The refinements are subtle but consistent, indicating a better overall representation of the hand movement dynamics. Notably, in the trough region of the Y velocity component, NDT, NDT-Hard-Distill and LFADS exhibit a significant mischaracterization of the velocity profile, predicting peak where the ground truth demonstrates a trough. In contrast, LFADS Hard- Distill accurately captures this critical feature, successfully predicting the trough trend in concordance with the ground truth.
A.5.2 2D HAND MOVEMENT TRAJECTORY DECODING
Fig. 4 bb) presents a comparison of 2D hand movement trajectories for four models: NDT (base model), NDD-Hard-Distill (our distilled model), LFADS (base model), and HardDs-ard-Disiill (our distilled model), alongside the ground truth trajectory.
For NDT, NDT--Had-Distill, and LFADS, there's significant discrepancy between the ground truth and predicted trajectories. While the ground truth curves downward, the predictions curve upward, indicating poor performance or a fundamental misunderstanding of the underlying pattern. For FEADS-Hard-o Distill, the predicted trajectory closely follows the ground truth, suggesting much better performance for this model. Moreover, LFADS Hoos--ar--JJ mast--l appears to have successfully learned the system's is behavior, demonstrating superior predictive capabilities.
A.5.3 COMPARATIVE ANALYSIS
Overall, we can conclude that the distillation process benefits both NDT and LFADS, with more pronounced improvements observed in the NDT model. The distillation process enhances the mod- els' ability to capture fine-grained details and reduce erratic predictions, resulting smoother, more accurate velocity profiles. Meanwhile, LFADD-Hard- accurately predicting the trough region, as illustrated in Fig. Ma), and in capturing the trajectory trend, as shown in Fig. ()). In both instances, the other three models exhibit significant deviations from the ground truth. This qualitative observation aligns with the quantitative results presented in Tab. 3 wherein decoding.
demonstrates superior performance in
consistently outperforms the other three models in behavior
These visualizations provide strong evidence for the effectiveness of the BLEND framework in improving neural dynamics modeling. Moreover, the distillation process appears to transfer valuable information from the behavior-guided teacher model to the student model, enhancing its ability to infer accurate hand velocities in multiple dimensions from neural activity alone. The improvements consistent across different aspects of the movement (2D movement trajectories, 1D X and Y velocities), suggesting a comprehensive enhancement of the models' predictive capabilities.
In summary, this analysis demonstrates that the BLEND framework not only improves quantita- tive metrics but also leads to qualitatively better predictions of complex behavioral outputs from neural activity. These visualizations provide intuitive and compelling evidence for the benefits of incorporating behavioral information through privileged knowledge distillation in neural dynamics modeling.
A.6 ADDITIONAL. QUALITATIVE ANALYSIS ON DIFFERENT DISTILLATION STRATEGIES
This section presents the qualitative analysis of different distillation strategies we employ in this work, including Hard Distillation (Eq. 3. Soft Distillation (Eq. Feature Distillation (Eq. 5 and Correlation Distillation (Eq. 6). This analysis is conducted based on the visualization of model predictions in the behavior decoding task of the MC-Maze dataset. As illustrated in Fig. 5 compare the performance of NDT and its different guiiddd distilled variants in decoding hand movement trajectories and velocities.
18

---

Published as a conference paper at ICLR 2025
Valocity
Wilccty
100
400
Predicted
Ground
20000
200
20000
NDT
400
Velechty
-100
-20000
Time Step
100
120
140
20
Time
300
120
140
-20000
Position
20003
'Pial 7 X Velocity
Tial Y Velocity
18
NDT-Hard- Dii-tiiitll
Ground
Ground
Pradictedd
20000
Ground Tuuh
200
10000
-00
Poshton
202
-20000
320
-20000
20000
Time
Time
Thaal 7- Velocity
Fiel 7-Y Welccity
NDT-Soft- Distill
000
400
20000
Prsittio
-200
-10000
-00000
-70003
70000
Time Step
Timee S1o9
Postbion
850
TPall 7-X Velocity
Pial 7- Velocty
Trial
NDT-Feeture-De Distill
Groand
400
Ground
Ground Trcc Predicted
608-
20000
Poiition
-0000
-400
-20000
22000
20000
Time Stop
Time 50ep
Position
Velacity
Trial 18
NDT-Correlation--ee Distill
602-
Precited
Grand
400
Groard Truth Predicted
20000
Ground Truth Predicted
200
10000
Voleck
-02
Velcctty
Position
200-
200
-400
20000
Time Step
120
120
340
20
Time Stop
1.00
120
340
-70000
Position
20000
(a) Velocity Trajectories
(b) 2D Trajectories
Figure 5: Visualization of predicted hand velocity MC-Maze dataset of different behavior-guided distilled models, including based model NDT, NDT-Hard -isttill, SoT-Soft-Distill NDT-Feature- Distill, and Der-Coorlaaiio-iistillll (a) Prediction and ground truth of X and Y velocities, respec- tively. (b) Prediction and ground truth 2D hand movement trajectory.
19

---

Published as a conference paper at ICLR 2025
A.6.1 ID HAND VELOCITY DECODING (X AND Y AXES)
For X velocities, all distilled models show improved alignment with ground truth compared to the base model NDT, especially during peak velocity periods (time steps 40-80). For the Y velocities, our behavior-guides distilled models demonstrate better tracking of the ground truth, with reduced oscillations and improved accuracy in predicting direction changes. Moreover, the distillation ap- proaches are particularly effective in capturing the synchronous changes in and Y velocities, sug- gesting better prediction of overall motion dynamies. NDT-Correlation-Disill stands out with its ability to accurately predict velocities in both dimensions simultaneously, indicating a strong grasp of the interdependencies between X and Y motions.
A.6.2 2D HAND MOVEMENT TRAJECTORY DECODING
For 2-dimensional hand movement trajectory decoding. the base model NDT shows significant de- viation from the ground truth, especially in the lower part of the trajectory. In contrast, all dis- tillation methods demonstrate improved trajectory prediction, with closer alignment to the ground truth path. Specifically, HDT-Coorelation-Disiti appears to provide the most accurate trajectory prediction, closely matching the ground truth's ss shape and endpoints.
A.6.3 MODEL-SPECIIFC OBSERVATIONS
For each distillation strategy employed in this study, the following key observations can be discerned from Figure 5
NDT-HHar-Distill: Shows substantial improvement over base NDT in both velocity and trajectory prediction.
JDT-Soff-Distill. Offers improved performance. though with some oscillations in velocity predictions.
COTT-Feature-Dicg Demonstrates good overall performance, with smooth velocity curves and accurate trajectory prediction.
JDDT-Correlation-Diitt Appears to be the best-performing method, showing excellent alignment with ground truth across all metrics.
Overall, we could conclude that the behavior-guided distillation approaches clearly enhance the pre- dictive capabilities of the NDT model. The improvements are particularly noticeable in handling complex motions and maintaining accuracy over longer time horizons. The correlation- dis- tillation method seems to be the most effective, suggesting that preserving relational information during distillation is crucial for accurate predictions.
A.7 CROSS-CORRELATION ANALYSIS
This subsection provides cross-correlation and mutual information analyses the MC-Maze dataset to quantify the relationship between neural activity and behavioral variables.
The cross-correlation heatmap was generated by computing correlations between each neuron's iiroo'sso ac- tivity and behavioral variables (X and Y velocities) across different time lags. For each neuron, the data was flattened across trials and timepoints, detrended, and normalized before computing corre- lations using Fast Fourier Transform (FFT) method. The resulting correlations were arranged into matrix where each row represents neuron, each column represents a time lag, and the color in- tensity indicates the correlation strength. As shown in Fig- 6. our analysis revealed robust temporal relationships between neural activity and movement kinematics. Nearly half of the recorded neurons (47.4%) exhibited leading relationships with horizontal (X) velocity, while slightly larger propor- tion (54.7%) led vertical (Y) velocity components. Importantly, we observe median leads of 20ms (X) and 10ms (Y), indicating that neural activity and behavior exhibit complex, complementary dynamics.
The strength of these neural-behavioral relationships was confirmed by highly significant correla- tions for both movement components (p < 0.001 for both X and Y velocities). This robust statistical coupling explains why behavior serves an effective privileged information source in our BLEND
20

---

Published as a conference paper at ICLR 2025
knowledge distillation framework, i.e., it provides reliable signal about the neural computations being performed. Crucially, we find that population-level mutual information (MI) is substantially higher than single-neuron MI, indicating that behavior captures coordinated neural dynamics that emerge at the population level. This difference demonstrates that behavior provides window into population-level neural representations, which is key motivation for our knowledge distillation approach.
Taken together, these analyses show that neural activity and behavior in the MC-Maze dataset strongly coupled, with behavior offering complementary contextual information about the under- lying neural computations. By effectively leveraging these BLEND framework is able to achieve significant performance improvements in learning better neu- ral population representations. The robust temporal structure and statistical significance of the ob- served cross-correlations further validate the importance of behavior as a privileged signal in our knowledge distillation approach.
relationships, our
Velocity
strength)
28
35
70
borto
105
20
112
119
-40
-100
100
Lag (ms)
Y Velocity Cross-correlation
(botte correlation strength)
14
21
20
56
20
91
105
119
126
-50
133
100
80
600
40
-20
0
20
40
60
80
100
Lag (ms)
Figure 6: Visualization of neural activity and behavior cross-correlation heatmap MC-Maze dataset. Upper: cross-correlation between independent neurons and horizontal movement; Below: cross-correlation between independent neurons and vertical movement.
A.8 EMPIRICAL STUDY ON BEHAVIOR GUIDANCE
To understand why incorporating behavior signals leads to better neural activity reconstruction, we conducted detailed analyses of the relationships between behavioral and neural data (exemplified with MC-Maze dataset).
A.8.1 NEURAL SPACE ORGANIZATION
We first examined how behavior signals influence the organization of neural representations using t-SNE visualization (Figure 7 The baseline model's representations (shown in red) appear scat- tered, disconnected clusters across the neural space. contrast, BLEND's representations (shown in blue) form more cohesive structures, suggesting that behavior signals help constrain neural activ-
21

---

Published as a conference paper at ICLR 2025
ity patterns into meaningful manifolds. This improved organization likely contributes to BLEND's superior reconstruction capability.
Noural
Distiled
new
Figure 7: t-SNE visualization of neural representations showing the organization of neural space. Red points represent baseline model representations, while blue points show BLEND model rep- resentations. The more cohesive clustering of BLEND's ' representations suggests better capture of underlying neural structure.
Neural Activity Patterns Across Behavior States
Reconstruction Error vs Behavior Coupling
Behavior
0.0006
Behavior 2/4
0.14
Behavior 3/4
Behavior
0.12
Reconstruction Enor
0.0004
Mean Activity
0.08
0.0003
8
0000
0.06
0.0002
0.04
8
0.0001
0.02
75
125
Coupling
High Coupling
Neuron Index
(a)
(b)
Figure 8: Behavior-neural analysis. (a) Neural activity patterns across different behavioral states show systematic variation. Different colors represent different behavior quartiles. (b) Box plots comparing reconstruction errors for neurons with low versus high behavior coupling, demonstrating improved reconstruction for behavior-coupled neurons.
A.8.2 BEHAVIOR- DEPENDDENT NEURAL ACTIVITY PATTERNS
We analyzed how neural activity patterns vary systematically with behavioral states (Figure g left). The behavioral states determined through the behavior data, which consists of two continu-
22

---

Published as a conference paper at ICLR 2025
variables (x and y coordinates of hand position) over time. We used the first behavior vari- able (x coordinate) and divided into quartiles using numpy's 's percentile function: behavior_bin percentile(behavior data, [0, 25, 50, 75, 100]). This created four behavior states (1/4 through 4/4), representing different ranges of hand positions. For each state, we computed the mean neu- ral activity pattern across all time points falling within that state. The resulting patterns show clear modulation of neural activity by behavioral state, with distinct activity profiles for different hand po- sitions. Some neurons (e.g., indices 25, 50, and 125) show particularly strong behavior-dependent modulation.
A.8.3 RECONSTRUCTION QUALITY ANALYSIS
We quantitatively assessed how behavior coupling influences reconstruction quality through com- prehensive analysis framework. First, we quantified behavior-neurg coupling using mutual infor- mation (MI): MI; B) -m, p(n, (2, b) log p(n.b)
p(n))[b)] where N; is the activity of neuron i and B is the behavior variable. We computed this for each neuron using scikit-learn' 's mutual _info_ regression
function. Neurons were then classified into high coupling (MI > median) and low coupling (MI < median) groups.
For each neuron i, we computed the reconstruction error as E, = the predicted firing rate and rf is the true firing rate time t. The box plots (Figure & right) show the distribution these errors for both groups Statistical analysis revealed significant differences between high and low coupling groups (Wilcoxon rank-sum test: < 0.001), with high-coupling neurons showing consistently lower reconstruction errors.
+E
r1)². where f is
These analyses collectively demonstrate that BLEND's ' improved performance stems from its ability to leverage inherent structure in behavior-neume relationships. The systematic variation of neural activity with behavior and the improved reconstruction for behavior-coupled neurons provide strong empirical evidence for the effectiveness of incorporating behavior signals in neural activity recon- struction.
A.9 DISTILLATION STRATEGY EVALUATION ON SYNTHETIC DATA
This subsection provides detailed exploration and comprehensive understanding of when and why different strategies excel. We conduct additional experiments using carefully designed synthetic datasets that mirror the structure of real neural recordings while allowing precise control over neural- behavioral relationships.
We created three types of synthetic datasets with distinct characteristics.
Simple: The first type implements simple linear relationships between neural activity and behavior, designed to test basic knowledge transfer. In this dataset, neural activity is directly derived from sinusoidal behavioral trajectories through a linear transformation. This creates clear and interpretable relationship between behavior and neural responses, with added Poisson noise to simulate realistic spike counts.
The second type features hierarchical relationships where different ron groups encode behaviors at varying timescales, simulating the layered processing of- ten observed in neural circuits. The hierarchical dataset introduces temporal complexity by dividing neurons into three functional groups, each processing behavior at different timescales. Fast neurons respond to immediate behavior, medium neurons integrate over timepoints, and slow neurons average over 10 timepoints, creating a rich temporal hierarchy that mirrors the multi-timescal processing observed in real neural systems.
Complex: The third type incorporates complex population-level corrllations with tem- poral dependencies, mimicking the intricate dynamics found in real neural populations. The complex population dataset represents the most sophisticated model, organizing neu- rons into five correlated assemblies with explicit population- structure. This dataset features decaying sinusoidal behavior patterns and maintains specific correlation structures (0.3 within groups) through multivariate normal noise, creating realistic population dynam- ics.
23

---

Published as a conference paper at ICLR 2025
Methods
Simple
Hierarchical
Comp ex
LFADS-Harre Hard-Distill
0.985
0.976
0.182
NDT-Hard-Distill
0.934
0.926
0.147
NDT--Soft-issill
0.927
0.932
0.146
NDT-Fcature-Distil
0.926
0.940
0.139
NDTT-Corelation-Disilll
0.935
0.944
0.180
Table 4: Results showing variants of BLEND framework on three synthetic datasets with distinct characteristics. Vel-R is used as the measurement metric.
All three datasets share common dimensions (2,869 trials, 140 timepoints, 182 neurons, 2 behavioral variables) and are systematically divided into train/eval splits and held-in/held-out neurons, enabling rigorous testing of neural analysis methods.
Our experiments with these synthetic datasets revealed clear patterns that explain the performance differences observed in Table LFADS consistently performed better with Hard Distillation across all synthetic cases, particularly excelling with simple linear relationships (around 6% improvement over NDJT-based variants V²)- This aligns with our understanding that LFADS's RNN-based architecture, with its inherent temporal continuity and built-in constraints, benefits most from direct knowledge transfer. The architecture's 's strong internal regularization makes well-suited for precise, deterministic knowledge transfer through hard distillation.
Conversely, NDT showed superior performance with Correlation Distillation, especially in cases with complex population-level correlations (around 20% improvement over other NDT-based vari- ants). This finding stems from NDT's transformer-bast architecture, which processes information in parallel and lacks inherent temporal constraints. The correlation distillation strategy provides valuable structural guidance that complements NDT's flexible architecture, helping it maintain im- portant population-level relationships in the learned representations.
These insights from synthetic data experiments provide concrete guide- lines for strategy selection in practice. For architectures with strong internal regularization like LFADS, Hard Distillation offers the most direct and effective knowledge transfer. For more flexible architectures like NDT, Correlation Distillation helps maintain complex population dynamics and temporal relationships. Soft Distillation, serving as a robust middle-ground option, can be bene- ficial for datasets where subtle variations in neural responses need be preserved, though it may not fully capture complex population-level dynamics. Feature Distillation, while generally showing more modest improvements, is particularly effective when the neural data exhibits clear hierar- chical structure, as it can capture relationships across different levels of neural representation. This understanding not only explains the performance patterns our original results but also offers a principled approach to selecting distillation strategies for future applications of BLEND different neural datasets and architectures.
Intuitions and guidelines:
A. EXTENDING BLEND TO DISCRETE OR NON-TEMPORAL BEHAVIOR LABELS
To verify the capability of BLEND with discrete or non-temporal behavior observations, ex- tend BLEND on the DMFC-RGS dataset from NLB'2 21 Benchmark Pei et al. (2021), which con- tains recordings from dorso-medial frontal cortex (DMFC) during a cognitive timing task known as Seady-Set-Go (RSG). The behavioral variables in this task are distinctly discrete in nature. The task incorporates five key behavioral features, represented as binary or categorical variables: "is eye' indicating whether the response was an eye movement (1) or joystick movement (0), 'theta' speci- fying the categorical direction of movement (left or right), and 'is_short' denoting whether the trial used the Short prior (1) or Long prior (0). The remaining variables 'ts' and 'tp' represent the sam- ple and produced time intervals respectively. These discrete behavioral choices create a structured experimental design with clearly defined conditions.
Given that this dataset has discrete and non-temporal behavior labels, we align the neural activity and behavior by introducing learnable temporal position encoding for discrete behavior variables across the length neural recording. We compare EDT-Hard-Disthe with the original NDT model on this dataset, and this simple extension shows improvement in both Co-Bps (0. 127 achieved by
24

---

Published as a conference paper at ICLR 2025
NDT-Hard-Distil and 0.118 achieved by NDT) and PSTH-R" (0.444 achieved by HDT-Hard-Distil and 0.338 achieved by NDT).
These promising results demonstrate the effectiveness of BLEND both continuous/discrete and cemporal/non-empoon behavior observations, showcasing a flexible solution we provide for learn- ing better neural population dynamics.
25


Reviews:
[{"review_id": "LQ2h3ihue5", "reviewer": "Reviewer_EUQu", "summary": "In this paper, the authors propose behavior-guided neural population dynamics modelling framework via privileged knowledge Distillation. They considering behavior as privileged information, and train a teacher model that takes both behavior observations (privileged features) and neural activities (regular features) as inputs, while a student model is distilled using only neural activity. The authors calm that the framework is model-agnostic, avoiding strong assumptions about the relationship between behavior and neural activity. \n\nThe authors present empirical evaluation that includes extensive experiments across neural population activity modeling and transcriptomic neuron identity prediction tasks. They report over 50% improvement in behavioral decoding and over 15% improvement in transcriptomic neuron identity prediction after behavior-guided distillation. Additionally, the authors also explore various behavior-guided distillation strategies within the proposed framework and present analysis of effectiveness and implications for model performance.", "questions": "Could you provide insights into which components of the BLEND framework (e.g., teacher-student architecture, distillation process, loss functions) contribute most significantly to the performance improvements? How do these components leverage the relationship between neural activity and behavior to enhance model performance?\n\nCould you discuss the potential for applying BLEND to other multimodal neuroscience datasets, such as combining neural activity with gene expression data or neuroimaging? What modifications, if any, would be necessary to adapt BLEND for these different data types?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Originality\n- Interesting approach. The idea to train a teacher model using neural activity and behavioral signals, while subsequently distilling its knowledge to a student model that utilizes solely neural activity as input seems novel. The inference phase is the common, where the authors use the distilled student model for neural population activity analysis and transcriptomic identity prediction tasks, but having embodied the privilege, behavioral information.\n\nQuality\n- The paper is well motivated, structured and presented, the problem is well introduces and connected to existing work. The writing is good. There is extensive and diverse evaluation. The proposed approach demonstrates a good capacity for capturing implicit patterns within neural activity. It shows improvements, outperforming state-of-the-art models.\n\nClarity\n- The idea and the main point of the paper are well explained. \n\nSignificance \n- The authors provide new perspectives on how behavioral observations can be leveraged to guide the complex modeling of neuronal population dynamics.", "weaknesses": "I would be interested and I believe it would be nice to see the cross-correlation within the data e.g., between the neural activity and behavior data, so that we can have a better understanding if there is a cross-correlation (such as time-lagged cross-correlations or mutual information analyses between neural features and behavioral variables). \n \nThis would help quantify the relationship between these data types (the extend of correlation and/or how much they complement) and potentially pinpoint and explain the benefits of the approach. It can also serve as additional insight to get to know more about why the proposed approach provides significant improvement in performance.", "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": "6NErUaXokY", "reviewer": "Reviewer_4WeH", "summary": "The authors proposed a two-step training process to create a model which is capable of creating a model that performs well using only regular features (neural activity) at inference time while benefiting from insights gained from privileged features (behaviour).", "questions": "My main concern is about distillation collapsing. Image to have two points B1 and B2 with the associated Xs. What happens while training (which I suggest decomposing to avoid overflow of equation number) when F_t(X1, B2) is similar to F_t(X2, B2) but X1 and X2 are not? In that case, the teaching model will collapse both X1 and X2 into the same point of the student model, since Distillation Loss push them to be similar. Do you think that this could happen? If so, I believe that such aspect must be addressed somehow.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and easy to follow, thanks to the clear structure. The method is introduced after the problem, and both are well formulated and discussed. The ablations studies are well performed, and the qualitative analysis are interesting to see. In general, the whole experimental section address the right points.", "weaknesses": "First two images must be reworked. Right now, those are single images divided into sub images within the draw. This lead to a weird referencing approach in the paper. I suggest dividing them into proper figures, and improve figure's 1 resolution. \n\nNVM acronym is not defined", "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": "qwTGxwGIWt", "reviewer": "Reviewer_WbXB", "summary": "This work focuses on developing a framework that incorporates behavior information during the training process, while using only the neural activities during inference stage for improved neural-behavior modeling. The proposed method is called BLEND, which uses the trick of privileged knowledge distillation, and distills a teacher model that is aware of behavioral information to a student model for inference. The model demonstrated strong empirical performance on both behavioural decoding tasks and neuron identity prediction tasks.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Writing is clear with adequate literature review;\n- Solid technical/experimental contribution, the experiments are adequately done;\n- Overall, an important topic for the community to consider using privileged knowledge distillation in such tasks. The performance improvement seems large, demonstrating opportunities for potential future works.", "weaknesses": "- Intuitions & clearness of section 3.2.1\n1. While I personally appreciate the comparison of different distillation techniques, it is unclear to me their intuitions (e.g., when to use which), and I'd appreciate if the authors can provide more about such intuitions in section 3.2.1. to improve clearness.\n2. It'd be even better if the authors can present a case study, ideally with synthetic datasets, showing why/when a distillation strategy is better. For example, in Table 3, it seems LFADS-Hard performs better than LFADS-Soft, yet for NDT NDT-Correlation in general performs better, did the authors check the reason behind it? Can the authors provide more reasoning/hypotheses based on the performance? \n\n- Fairness and more details of model selection and evaluation protocol. \n1. The models in Table 1 and 2 are named as \"Distill-Best\", which I assume is employing the best distillation strategy. What are the performance of the worst distillation strategies then? Can the authors be more specific about the model selection pipeline here?\n2. Has the models run across different random seeds? What is the standard deviation? Is the training stable? Are there any special techniques applied to ensure the distilled model is good? (disclaimer: I didn't read the appendix, if some of those info are available in appendix, please make more detailed co-references.)\n\n- The paper can benefit from more experimental case studies.\n1. While there are a bunch of experiments, some of the authors claims were not well justified. Especially, as stated in footnote 1 \"BLEND can be readily extended to accommodate scenarios wherein behavioral information is discrete or non-temporal in nature.\" - While the framework is easily extendable to discrete labels, can the authors actually attempt to perform an experiment given behavioral signals as discrete labels? e.g. The authors can mimic what is done in [1] and use the final reaching target direction to guide the learning, and show the corresponding performance of the model.\n\n- Lack of theoretical & analytical results that supports the experimental results.\n\n- Minor formatting suggestions:\n1. Do the authors mind putting Table 1 and 2 at either the bottom or top of the pages?\n\n[1] Liu, R., Azabou, M., Dabagia, M., Lin, C. H., Gheshlaghi Azar, M., Hengen, K., ... & Dyer, E. (2021). Drop, swap, and generate: A self-supervised approach for generating neural activity. Advances in neural information processing systems, 34, 10587-10599.\n\n\nOverall, I like the scope and the results presented in the paper and appreciate the technical contributions. In general, I think the paper could benefit from more intuitive case studies or synthetic experiments, aside from the presented tables and numbers, as the datasets used in computational neuroscience are in general small. My rating of the paper is a solid 6 despite the issues.", "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": "TWj1oo2TE9", "reviewer": "Reviewer_UDDD", "summary": "This manuscript propose a novel algorithm of BLEND to deal with neural activity reconstruction. Concretely, the behavior signals are used as the privileged knowledge and combined with masked neural activity data to train a teacher model. Then the teach model is distilled to a student model, which only use masked neural activity data as the input for reconstruction. The effectiveness of BLEND is verified on two benchmarks.", "questions": "See Weaknesses.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The organization of this paper is well and the presentation of method is clear.\n\n- A novel approach of BLEND, which conbines privilege knowledge (training teacher model) and knowledge distillation (training student model), is proposed for neural activity reconstruction.\n\n- BLEND has achieved a good performance on two benchmarks.", "weaknesses": "I am not an expert in computational neuroscience, and cannot give a rational judgement w.r.t. experiments. However, I still think there are some major weaknesses in terms of writing and the rationale of BLEND, please see details as follows.\n\n- **Q1:** The gap between BLEND and previous approaches is unclear, and the proposal of BLEND is suddenly. Is there other approaches that uses both masked neural activity data and behavior signals for training and only masked neural activity data for test? If it is, a detail comparison should be presented. Otherwise the significance of this scenario should be clear first.\n\n- **Q2:** The rationale of proposed BLEND. The superiorty of BLEND is only verified on two recent benchmarks with limited baselines. It will be more convincing if a theoreitcal or empirical analysis can be provided to show why a better neural activity reconstruction is achieved by incorporating behavior signals. Therefore, the internal struction between neural activity data and behaivor signal data should be investigated.\n\n\n- **Q3:** BLEND incorporates the technique of knowledge distillation. Is there a large performance drop by comparing teacher models and student models? In the experiments, many different distillation methods are used but only the performance of the best one is reported, which is not very fair for comparison.\n\nBased on the above questions, I think there lacks a investigation on the characteristic of neuron data, so the rationale of proposed method cannot be well verified, i.e., why is BLEND good for the specific type of neuron data.", "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 (Poster)