BLEND: Behavior-guided Neural Population Dynamics Modeling via Privileged Knowledge Distillation

## Abstract

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 % percent 50 50\% 50 % improvement in behavioral decoding and over 15 % percent 15 15\% 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 .

## 1 Introduction

Large-scale population-level 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, ; Jun et al., ; Pei et al., ). 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., ; Ye & Pandarinath, ; Le & Shlizerman, ).

Alongside the recorded neural population activity, the observed behavior signals provide crucial context and complementary information during neural dynamics modeling (Sani et al., ). 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 & Wei, ; Hurwitz et al., ; Schneider et al., ), bridging the gap between neural activity and its real-world manifestations.

Figure: Figure 1: Schematic illustration of neural population dynamics modeling mechanisms. In this paper, we 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 as 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 a prior. (c) Our BLEND framework, which considers behavior information as privileged knowledge for distillation.
Refer to caption: extracted/5895791/schematic.png

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 or receiving external stimuli. The absence of structured tasks means there’s no clear temporal alignment between neural events and behavioral events (Drew et al., ; Nozari et al., ).
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, ). 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 (behavior), represents a critical research question in neural dynamics modeling,
especially in bridging the gap between controlled experimental settings and real-world applications where privileged knowledge is limited or unavailable.
Maximizing the utility of existing privileged-regular feature pairs to enhance the regular-feature-based network performance is a 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. 1(c), BLEND constitutes a student and a teacher model, in which the teacher trains on both behavior observations and neural activity recordings, then distills knowledge to guide the student which takes only neural activity as input. This ensures the student model can make predictions using only recorded neural activity during the deployment/inference stage, but also benefits from the guidance from behavior information during the training stage, making it more versatile for settings lacking behavioral data. Our main contributions are summarized as follows:

- •
Built upon the privileged knowledge distillation framework, we introduce a 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 our framework, we conduct extensive experiments on two benchmarks, i.e., Neural Latents Benchmark’21 for neural activity prediction, behavior decoding, and matching to peri-stimulus time histograms (PSTHs), as well as a multi-modal calcium imaging dataset for transcriptomic identity prediction.
Results show that our framework elevates the performance of baseline methods by a large margin (>$50\%$ improvement in behavioral decoding and >$15\%$ improvement in neuronal identity prediction) and significantly outperforms the state-of-the-art (SOTA) models.
- •
We present a comprehensive analysis of our framework across base models and privileged knowledge distillation strategies, revealing key insights into the interaction between neural activity and behavior. Our results demonstrate that behavior-guided distillation not only improves model performance 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.

Figure: Figure 2: Illustration of the proposed BLEND framework, exemplified using neural spiking activity data. (a) Behavior-guided 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. (b) During the inference phase, the distilled student model is employed for neural population activity analysis and transcriptomic identity prediction tasks.
Refer to caption: extracted/5895791/methodology.png

## 2 Related work

Neural dynamics modeling (NDM). NDM refers to a category of methods that aim to capture the dynamics of neural activity by using the activity recordings as inputs. A common and effective choice is the latent variable model (LVM), which leverages low-dimensional latent factors to interpret these dynamics. Various LVMs are developed, ranging from simple non-temporal models such as principal components analysis (PCA) (Cunningham & Yu, ) and its variants (Kobak et al., ), to linear dynamical systems (Macke et al., ; Gao et al., ), and to complex state space models like LFADS (Pandarinath et al., ). Since the advent of Transformer, its ability to capture temporal dependencies and long-range interactions makes it appropriate for neural data, with representative works such as NeuralDataTransformer (NDT) (Ye & Pandarinath, ), STNDT (Le & Shlizerman, ), and EIT (Liu et al., ). As shown in Fig. 1(a), methods in this category purely depend on neural activity recordings for neural population dynamics modeling yet ignore using the paired behavior information as guidance.

Neural dynamics modeling with behavior as a prior. In recent years, an emerging research direction has focused on jointly modeling neural population dynamics and behavioral signals. As illustrated in Fig. 1(b), these methods can be categorized into two types. The first inputs both behavior and neural activity and utilizes behavior signals to guide the learning of neural dynamics: pi-VAE (Zhou & Wei, ) considers the behavior variables as constraints for the construction of latent space of LVMs; CEBRA (Schneider et al., ) utilizes behavior signals to construct contrastive learning samples for label-informed neural activity analysis. However, these approaches normally need to develop delicately designed modules or complex training strategies to achieve their goal. The second line of work aims to decompose neural activity into behavior-relevant and behavior-irrelevant dynamics and reconstruct both behavior signals and neural activity signals. To that end, a linear state-space model, PSID (Sani et al., ), is developed to decode population dynamics from motor brain regions. Nonlinear state-space models such as TNDM (Hurwitz et al., ) and SABLE (Jude et al., ) are proposed to further capture the nonlinear dynamics of neural 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.

In contrast, our approach offers a model-agnostic learning paradigm that can be directly applied to existing NVM models without relying on strong assumptions, thereby circumventing the issues present in current NVM models that integrate behavior information as a prior.

Learning under privileged information (LUPI). Firstly proposed in Vapnik & Vashist (), 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 available 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 to different machine learning problems such as recommendation system (Yang et al., ), medical image analysis (Chen et al., ), emotion recognition (Aslam et al., ), and semantic segmentation (Liu et al., ).

In computational neuroscience, considering behavior information as privileged information to guide neural dynamics modeling remains understudied. This work develops a novel learning paradigm that represents an inaugural step toward addressing this underexplored research question and advancing the field.

## 3 Methods

This section introduces the details of our proposed BLEND framework, i.e., the behavior-guided neural population dynamics modeling via privileged knowledge distillation. We start with problem formulation of behavior-guided neural dynamics modeling in Section 3.1, followed by an exposition of the BLEND algorithm in Section 3.2.

### 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 $\mathbf{x}\in\mathcal{X}=\mathbb{N}^{N\times T}$ represent the input spike counts, where $\mathbf{x}_{i}^{t}$ denotes the spike count for neuron $i$ at time $t$. Let $\mathbf{b}\in\mathcal{B}=\mathbb{R}^{B\times T}$ be the corresponding behavior signal, with $\mathbf{b}_{i}^{t}$ 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.(^1^11BLEND can be readily extended to accommodate scenarios wherein behavioral information is discrete or non-temporal in nature.) As illustrated in Fig. 1, the model $f_{\theta}$ aims to reconstruct the randomly masked neural activity solely based on the unmasked portions of $\mathbf{x}$ or together with the auxiliary behavior signal $\mathbf{b}$. The optimization objective for masked time-series modeling (MTM) can be formulated as follows:

$$ $\theta^{*}=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{\mathbf{x}\sim p( \mathbf{x}),\mathbf{b}\sim p(\mathbf{b}),\mathbf{m}\sim p(\mathbf{m})}\left[ \frac{1}{|\mathbf{m}|}\sum_{i,t}\mathbf{m}_{i}^{t}\cdot\mathcal{L}_{\text{rec} }\left(f_{\theta}\left(\mathbf{x}_{\bar{m}},\mathbf{b}\right)_{i}^{t},\mathbf{ x}_{i}^{t}\right)\right]\triangleq\mathcal{L}_{\text{MTM}}$ (1) $$

where $\mathbf{m}\in\{\mathbf{0},\mathbf{1}\}^{N\times T}$ is a binary mask with $\mathbf{1}$ indicating masked elements, $\mathbf{x}_{\bar{m}}\in\mathcal{X}_{\bar{m}}$ represents the unmasked portions of $\mathbf{x}$, $|\mathbf{m}|$ is the number of masked elements, and $\mathcal{L}_{\text{rec}}$ is either Poisson or Cross-Entropy 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/inference, but also benefits from behavior data during training. This subsection details the algorithm of BLEND, including the formulation of privileged knowledge distillation under the context of neural dynamics modeling, teacher-student architecture, and distillation strategies.

We first conceptualize the behavior-guided neural population dynamics modeling problem under the framework of LUPI:

###### Definition 1 (Privileged Knowledge (Yang et al., 2022 ) ) .

Consider a general learning problem with input $\mathbf{x}\in\mathcal{X}$ and label $\mathbf{y}\in\mathcal{Y}$. If there is an additional source of information $\mathbf{b}\in\mathcal{B}$ that exists during training but not inference, we say $\mathbf{b}$ is the privileged knowledge if and only if $I(\mathbf{y};\mathbf{b}|\mathbf{x}):=H(\mathbf{y}|\mathbf{x})-H(\mathbf{y}|
\mathbf{x},\mathbf{b})>0$.

In Definition 1, $I(\cdot|\cdot)$ and $H(\cdot|\cdot)$ represent conditional mutual information and conditional entropy, respectively. According to Definition 1, the behavior information $\mathbf{b}$ serves as the privileged knowledge and offers additional predictive capabilities of label $\mathbf{y}$, which is the unmasked neural activity $\mathbf{x}$ in our scenario.
Next, we introduce the proposed privileged knowledge distillation framework, which is designed to incorporate behavioral guidance and facilitate neural dynamics modeling.

###### Definition 1 (Privileged Knowledge (Yang et al., 2022 ) ) .

#### 3.2.1 Privileged knowledge distillation with behavior as guidance

As is illustrated in Fig. 2(a), BLEND employs a teacher-student architecture to realize our goal: by leveraging the privileged information $\mathbf{b}$ at the training phase, to learn an NVM model for inference phase that outperforms those built on the regular feature $\mathbf{x}$ alone. This procedure is divided into two sequential stages:

Stage 1: Train a teacher model with regular information and privileged information. The teacher model
$f_{\theta_{T}}\in\{f|f:\mathcal{X}_{\bar{m}}\times\mathcal{B}\mapsto\mathcal{X}\}$
takes both masked neural activity recordings $\mathbf{x}$ and the corresponding behavior signals $\mathbf{b}$ as inputs. The model parameter $\theta_{T}$ is optimized by minimizing the MTM loss $\mathcal{L}_{\text{MTM}}$ formulated in Eq. (1).

Stage 2: Train a student model with regular information by distillation. By exploiting all available sources of information, the teacher model acquires a rich understanding of the neural population dynamics. We then aim to distill the learned knowledge to the student model $f_{\theta_{S}}\in\{f|f:\mathcal{X}_{\bar{m}}\mapsto\mathcal{X}\}$, which takes the neural activity as input. The parameter of student model $\theta_{S}$ is updated by minimizing the following objective:

$$ $\mathbb{E}_{\mathbf{x},\mathbf{b},\mathbf{m}}\Bigg{[}\alpha\cdot\underbrace{ \frac{1}{|\mathbf{m}|}\sum_{i,t}\mathbf{m}_{i}^{t}\cdot\mathcal{L}_{\text{rec} }\left(f_{\theta_{S}}\left(\mathbf{x}_{\bar{m}}\right)_{i}^{t},\mathbf{x}_{i}^ {t}\right)}_{\text{MTM Loss}}+(1-\alpha)\cdot\underbrace{\sum_{i,t}\mathcal{L} _{\text{distill}}\left(f_{\theta_{S}}\left(\mathbf{x}_{\bar{m}}\right)_{i}^{t} ,f_{\theta_{T}}\left(\mathbf{x}_{\bar{m}},\mathbf{b}\right)_{i}^{t}\right)}_{ \text{Distillation Loss}}\Bigg{]},$ (2) $$

where $\alpha\in(0,1)$ 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. 1. After training the student model with privileged knowledge distillation for MTM, it is applied to downstream tasks where only neural activity recordings are available for inference (illustrated in Fig. 2(b)).

While the above algorithm outlines the general framework of BLEND, different implementations of the distillation loss $\mathcal{L}_{\text{distill}}$ can capture various aspects of the teacher’s knowledge. In this study, we investigate the following four main distillation strategies in our BLEND framework, each designed to transfer certain aspects of the teacher’s knowledge to the student model.

Hard distillation. As a baseline, we first implement the $\mathcal{L}_{\text{distill}}$ with hard distillation strategy, which directly minimizes the mean squared error (MSE) between the teacher and student outputs:

$$ $\mathcal{L}_{\text{distill}}\left(f_{\theta_{S}}\left(\mathbf{x}_{\bar{m}} \right)_{i}^{t},f_{\theta_{T}}\left(\mathbf{x}_{\bar{m}},\mathbf{b}\right)_{i} ^{t}\right)=\left\lVert f_{\theta_{S}}\left(\mathbf{x}_{\bar{m}}\right)_{i}^{t }-f_{\theta_{T}}\left(\mathbf{x}_{\bar{m}},\mathbf{b}\right)_{i}^{t}\right \rVert_{2}^{2}$ (3) $$

Soft distillation. This approach distills knowledge by matching the softened probability distributions of the teacher and student models. We use a temperature parameter $\tau$ to soften the logits before applying the softmax function:

$$ $\mathcal{L}_{\text{distill}}\left(f_{\theta_{S}}\left(\mathbf{x}_{\bar{m}} \right)_{i}^{t},f_{\theta_{T}}\left(\mathbf{x}_{\bar{m}},\mathbf{b}\right)_{i} ^{t}\right)=\tau^{2}\cdot\text{KL}\left(\sigma\left(\frac{f_{\theta_{T}}( \mathbf{x}_{\bar{m}},\mathbf{b})_{i}^{t}}{\tau}\right)\Big{\|}\sigma\left( \frac{f_{\theta_{S}}(\mathbf{x}_{\bar{m}})_{i}^{t}}{\tau}\right)\right),$ (4) $$

where $\sigma$ is the softmax function, $\text{KL}(\cdot\|\cdot)$ is the Kullback-Leibler divergence, and $\tau$ is the temperature parameter.

Feature distillation. This method aims to align all the intermediate representations of the student model with those of the teacher model.

$$ $\displaystyle\mathcal{L}_{\text{distill}}\left(f_{\theta_{S}}\left(\mathbf{x}_ {\bar{m}}\right)_{i}^{t},f_{\theta_{T}}\left(\mathbf{x}_{\bar{m}},\mathbf{b} \right)_{i}^{t}\right)=\sum_{l=1}^{L}\left\lVert f_{\theta_{S}}^{l}(\mathbf{x} _{\bar{m}})_{i}^{t}-f_{\theta_{T}}^{l}(\mathbf{x}_{\bar{m}},\mathbf{b})_{i}^{t }\right\rVert_{2}^{2},$ (5) $$

where $f_{\theta_{S}}^{l}$ and $f_{\theta_{T}}^{l}$ denote the outputs of the $l$-th layer of the student and teacher models, respectively.

Correlation distillation. This approach focuses on preserving the correlation structure of the teacher’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:

$$ $\mathcal{L}_{\text{distill}}\left(f_{\theta_{S}}\left(\mathbf{x}_{\bar{m}} \right)_{i}^{t},f_{\theta_{T}}\left(\mathbf{x}_{\bar{m}},\mathbf{b}\right)_{i} ^{t}\right)=\frac{1}{B}\sum_{j=1}^{B}\Big{\|}\text{Corr}(f_{\theta_{S}}( \mathbf{x}_{\bar{m}})_{j})-\text{Corr}(f_{\theta_{T}}(\mathbf{x}_{\bar{m}}, \mathbf{b})_{j})\Big{\|}_{2}^{2},$ (6) $$

where $B$ denotes the batch size, and $\text{Corr}(\cdot)$ computes the correlation matrix of the outputs for the $j$-th sample in the batch (see Appendix A.1.2 for details).

## 4 Experiments

As shown in Fig. 2(b), this work includes two benchmarks for evaluating our proposed BLEND framework. The first is a public benchmark for neural latent dynamics model evaluation from Pei et al. (), named Neural Latents Benchmark’21 (NLB’21).
The second is a recent, public multi-modal neural dataset from Bugeon et al. (), which contains calcium imaging recordings of neural population activity as well as single-cell spatial transcriptomics of the recorded tissue.
Details of these two benchmarks, tasks, and included baselines are introduced in Section 4.1 and 4.2.

### 4.1 Neural Population Activity Analysis on NLB’21 Benchmark

We include three sub-datasets from NLB’21, i.e., MC-Maze, MC-RTT, and Area2-Bump, for a series of neural dynamics modeling evaluations.
These datasets contain neural recordings from monkeys performing various reaching tasks: MC-Maze features delayed reaches through virtual mazes, MC-RTT involves continuous reaches to random targets without delays, and Area2-Bump includes reaches with occasional mechanical perturbations (see Appendix A.2 for details).
To evaluate the capability of our proposed BLEND on neural population dynamics modeling, we adopt three tasks from NLB’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 as the log-likelihood of held-out neurons’ activity (see Appendix A.3 for detailed computation steps).

Behavior decoding task requires the model to relate neural activity to observed behavior. To evaluate 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^{2}$ score (named $\text{Vel-}R^{2}$).

Match to PSTHs task aims to evaluate how well models can capture stereotyped features of neuronal responses across repeated trials of the same condition. To evaluate it, we calculate the $R^{2}$ between model-predicted trial-averaged rates and true PSTHs for each neuron across all conditions, then average these $R^{2}$ values across neurons (named $\text{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., ), LFADS (Pandarinath et al., ), NeuralDataTransformer (NDT) (Ye & Pandarinath, ), STNDT (Le & Shlizerman, ), and MINT (Perkins et al., ). The second type utilizes behavior information during training as a prior to facilitate the learning process of neural dynamics modeling, including pi-VAE (Zhou & Wei, ), PISD (Sani et al., ), and TNDM (Hurwitz et al., ). We choose NDT and LFADS as the base model for privileged knowledge distillation using behavior information, with best-distilled model performance reported in Tab. 1. See details of model configurations in Appendix A.4.

### 4.2 Transcriptomic Neuron Identity Prediction on Multi-modal 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 transcriptomic profiles that are used to categorize neurons as 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 on the following two tasks under both multi-animal scenario (all 4 mice) and single-animal scenario (only mouse SB025):

Excitatory/inhibitory (EI) neuron identity prediction task is a binary classification task that requires the model to predict the neuron identity, i.e., excitatory or inhibitory.

Subclass label prediction task is a 5-class classification task, where the model is required to predict the subclass label of inhibitory neurons, i.e., Lamp5, Pvalb, Vip, Sncg, or Sst.

Following the settings in Mi et al. (), we include a random model, PCA (Cunningham & Yu, ), UMAP (McInnes et al., ), LOLCAT (Schneider et al., ) and its variants, and NeuPRINT (Mi et al., ) 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. 2. Detailed model configurations are in Appendix A.4.

## 5 Results and Analysis

### 5.1 Benchmark 1: Neural Population Activity Analysis

Tab. 1 summarizes the results of neural activity prediction, behavior decoding, and match to PSTHs tasks of baseline methods and our proposed method on 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.

Table: Table 1: Comparison of the proposed BLEND framework with other neural dynamics modeling methods on NLB’21 Benchmark (Pei et al., ). Bold values denote the best performance for the corresponding metric.

Fit to neural activity. For neural activity prediction of held-out neurons, NDT-Distill-Best achieves $0.310$ Co-bps score on the MC-Maze dataset, reporting $12.7\%$ improvement over the base model NDT. LFADS-Distill-Best achieves a $0.321$ Co-bps score, reporting the best performance for this task against other baseline models. These observations indicate that our models gain a better understanding of neural population dynamics and yield improved predictions of held-out neural activity.

Fit to behavior. For behavior (monkey hand velocity) decoding, NDT-Distill-Best achieves $\text{Vel-}R^{2}$ of $0.891$, $0.372$, and $0.788$ respectively, outperforming the compared baselines. A $51.8\%$ improvement is observed on the Area2-Bump dataset over the base model NDT. While for LFADS-Distill-Best, $\text{Vel-}R^{2}$ 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 LFADS over the base model on MC-Maze and MC-RTT is not as significant as 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 additional knowledge effectively. Meanwhile, LFADS performs dimensionality reduction as 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 $\text{PSTH-}R^{2}$ scores of $0.592$ and $0.483$ on the MC-Maze and Area2-Bump datasets, respectively. These results represent improvements of 7.4% and a remarkable 66.6% over the baseline model performance. The distilled LFADS model achieves $\text{PSTH-}R^{2}$ 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, respectively. These findings indicate that leveraging behavioral data as privileged knowledge enhances the models’ capacity to reproduce stereotyped neural responses, as 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 Benchmark 2: Transcriptomic Neuron Identity Prediction

Tab. 2 shows the experimental results of transcriptomic identity prediction tasks, including EI neuron identity prediction and inhibitory neuron subclass label prediction. These two tasks are conducted in both multi-animal and single-animal settings. By incorporating our BLEND framework, the distilled model shows remarkable improvement over the base model and outperforms all baseline methods.

Table: Table 2: Top-1 accuracy of excitatory/inhibitory neuron identity prediction and inhibitory neuron subclass label prediction tasks on the multimodal population activity dataset from Bugeon et al. (). Bold values denote the best performance for the corresponding metric.

EI prediction. In the excitatory/inhibitory 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. Notably, 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 involves a five-class classification problem, the model’s performance metrics are comparatively lower than those observed in the EI prediction task, which is a binary classification problem.

As illustrated in Table 2, the model’s performance metrics in the single-animal setting are comparatively 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 et al. () points out a data-limited regime in this benchmark 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.

### 5.3 Ablation Studies on Behavior-guided Knowledge Distillation Strategies

Table: Table 3: Ablation study on privileged knowledge distillation strategies evaluated by neural dynamics modeling tasks of NLB’21 Benchmark (Pei et al., ) and cell class prediction tasks of Bugeon et al. () dataset. Bold values denote the best performance for the corresponding metric.

In this work, we explore different privileged knowledge distillation strategies, including hard distillation, 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-Correlation-Distill model generally performs best across most metrics, showing significant improvements over the baseline NDT. For MC-Maze, NDT-Correlation-Distill achieves the highest $\text{Vel-}R^{2}$ ($0.891$). While for LFADS and its distilled variants, LFADS-Hard-Distill shows improvements over the baseline LFADS, particularly in $\text{PSTH-}R^{2}$ for MC-Maze and $\text{Vel-}R^{2}$ for Area2-Bump. Notably, except for LFADS-Hard-Distill, other distillation methods lead to performance degradation for LFADS on MC-Maze (hence, only results for Soft-Distill are presented in Table 3). 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-Distill and NeuPRINT-Soft-Distill show consistent improvements over the baseline NeuPRINT across all tasks. NeuPRINT-Correlation-Distill performs best for Multiple Animals EI prediction. NeuPRINT-Soft-Distill excels in Single Animal predictions for both EI and Subclass.

Overall, the distillation strategies, particularly Hard-Distill and Correlation-Distill, tend to outperform 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 comparison 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-Distill (our distilled model). And subfigure (b) illustrates the 1D hand velocity over time for both the X and Y axes separately. For a more comprehensive visualization and qualitative analysis, please refer to Appendix A.5 and A.6.

2D hand movement trajectory decoding. For NDT the base model, the predicted trajectory follows the general U-shape of the ground truth but shows significant deviations, especially at the bottom of the U and the upper right corner. While our distilled model generates a trajectory that aligns much 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 substantial 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.

Figure: Figure 3: Visualization of behavior decoding on MC-Maze dataset. (a) Prediction and ground truth of 2D hand movement trajectory. (b) Prediction and ground truth of X and Y velocities, respectively.
Refer to caption: extracted/5895791/qualitative.png

## 6 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 behavioral 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 modeling architectures. This versatility enables the enhancement of a wide range of computational approaches in the field of computational neuroscience. Comprehensive empirical evaluations, encompassing neural population activity analysis and transcriptomic identity prediction tasks, substantiate the superior effectiveness of our proposed framework. BLEND demonstrates a remarkable capacity for capturing implicit patterns within neural activity, consistently outperforming state-of-the-art 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, we employ a straightforward approach to implement the behavior guidance by concatenating 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 a query in the cross-attention mechanism). However, this simplistic guidance strategy may not fully capture the intricate relationship between neural activity and behavior. Future research should investigate advanced integration methods to better utilize the complex interactions between these data types, potentially improving the model’s ability to represent subtle neural patterns. Moreover, 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. This expansion could encompass non-temporal or discrete behavioral signals, thereby broadening the framework’s applicability to diverse neuroscientific domains. Such extensions could prove valuable in investigating sleep stages, categorizing social behaviors, and exploring other scenarios where behavior is not continuously paired with neural activity. This generalization would significantly enhance the versatility and utility of BLEND across a wider spectrum of neuroscience research paradigms.

## Reproducibility Statement

To enhance the reproducibility of this study, we provide an Appendix section comprising six subsections 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-guided knowledge distillation strategies. Additionally, Appendix A.2 provides detailed descriptions of the datasets utilized, namely MC-Maze, MC-RTT, Area2-Bump, and the multi-modal neural activity datasets. To facilitate a deeper understanding of the tasks conducted in this study, Appendix A.3 elucidates the specifics of neural activity prediction, behavior decoding, and match to PSTHs tasks. Subsequently, Appendix A.4 delineates the model configurations, including architecture and hyperparameters. To further elucidate the superior performance of our proposed method, Appendices A.5 and A.6 offer additional visualizations comparing different base models for distillation and various distillation strategies, enabling qualitative assessment. Our source code will be made publicly accessible upon acceptance of this paper.

## Appendix A Appendix

### A.1 Supplementary contents of blend algorithm

#### A.1.1 BLEND Algorithm

Figure: Algorithm 1 Behavior-guided Teacher-Student Knowledge Distillation Framework (BLEND)

#### A.1.2 Distillation strategies.

Correlation-Distill. As shown in Section 3.2.1, this distillation strategy aims to preserve the correlation structure of the teacher’s outputs. Concretely, the correlation matrix is computed as:

$$ $\text{Corr}(\mathbf{Y})=\frac{(\mathbf{Y}-\bar{\mathbf{Y}})(\mathbf{Y}-\bar{ \mathbf{Y}})^{T}}{\sqrt{\text{diag}((\mathbf{Y}-\bar{\mathbf{Y}})(\mathbf{Y}- \bar{\mathbf{Y}})^{T})\cdot\text{diag}((\mathbf{Y}-\bar{\mathbf{Y}})(\mathbf{Y }-\bar{\mathbf{Y}})^{T})^{T}}}$ (7) $$

Here, $\mathbf{Y}\in\mathbb{R}^{N\times T}$ 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. $\bar{\mathbf{Y}}$ is the mean of $\mathbf{Y}$ along the time dimension, and $\text{diag}(\cdot)$ extracts the diagonal of a matrix. The resulting correlation matrix has dimensions $N\times N$.
Note that we use the Frobenius norm $\left\lVert\cdot\right\rVert_{F}$ to compute the difference between correlation matrices, as it provides a scalar measure of the overall difference between these matrices.

### A.2 Supplementary contents of datasets included

MC-Maze dataset (Churchland et al., ; ) comprises recordings from the primary motor and dorsal premotor cortices of a 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 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., ; Makin et al., ) features motor cortical recordings during a random target task, comprising continuous, point-to-point 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 to its non-stereotyped movements, lack of trial repetitions, and the unpredictability 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., ) contains neural recordings from area 2 of the somatosensory cortex, which processes proprioceptive information, as a monkey performed a visually-guided reaching task using a manipulandum. It comprises 462 trials with data from 65 neurons, along with hand kinematics and perturbation information, where in 50% of random trials, the monkey’s arm was unexpectedly bumped before the reach cue. This dataset challenges models to infer inputs describing activity after sensory perturbations and to perform robustly with low neuron counts, offering insights into the distinct dynamics of somatosensory areas compared to motor areas.

For each of the above datasets, neural activity recordings are counted in 5ms bins, and behavior signals are also measured at 5ms intervals. Other preprocessing information can be found in Pei et al. ()(^2^22[https://github.com/neurallatents/nlb_tools/tree/main](https://github.com/neurallatents/nlb_tools/tree/main)) or Ye & Pandarinath ()(^3^33[https://github.com/snel-repo/neural-data-transformers](https://github.com/snel-repo/neural-data-transformers)).

Multi-model neural activity dataset collected by Bugeon et al. () 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), encompassing 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 a portion of inhibitory neurons in one mouse (SB025). Complementing 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 a rich resource for investigating the interplay between neural activity, cell types, and behavior.

For the preprocessing details of this dataset, please refer to Mi et al. ().(^4^44[https://github.com/lumimim/NeuPRINT](https://github.com/lumimim/NeuPRINT))

### 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 to predict firing rates $\lambda$ for the held-out test neurons. Performance is measured using log-likelihood under a Poisson observation model:

$$ $p(\mathbf{\hat{y}}_{n,t})=\text{Poisson}(\mathbf{\hat{y}}_{n,t};\lambda_{n,t}),$ $$

where $\mathbf{\hat{y}}_{n,t}$ is the held-out spike count for neuron $n$ at time $t$. The log-likelihood $\mathcal{L}(\lambda;\mathbf{\hat{y}})$ is 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:

$$ $\text{bits/spike}=\frac{1}{n_{sp}\log 2}(\mathcal{L}(\lambda;\mathbf{\hat{y}}_ {n,t})-\mathcal{L}(\bar{\lambda}_{n};\mathbf{\hat{y}}_{n,t})),$ $$

where $\bar{\lambda}_{n,:}$ is the mean firing rate for neuron n and $n_{sp}$ 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 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., ).

Behavior decoding. As mentioned in Section 4.1, following the settings in Pei et al. (); Ye & Pandarinath (), 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 at 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 responses captured by peri-stimulus time histograms (PSTHs). PSTHs are computed by averaging neuronal responses across trials within a given condition, revealing consistent features of neural activity. For datasets with clear trial structures (MC-Maze and Area2-Bump), the task computes the $R^{2}$ 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., ) is a deep learning method designed to 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 a 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 generate de-noised 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\times 10^{-3}$ with $5000$ warm-up iterations and weight decay to $5\times 10^{-5}$. Mask ratio is set to $0.25$.

NDT (NeuralDataTransformer) (Ye & Pandarinath, ) is a non-recurrent architecture designed to model neural population spiking activity, based on the BERT encoder. The core of the NDT consists of a stack of Transformer layers (typically 6 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 is 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 $2$. For model training, we set the batch size to $64$, learning rate to $1\times 10^{-3}$ with $5000$ warm-up iterations and weight decay to $5\times 10^{-5}$. Mask ratio is set to $0.25$.

NeuPRINT (Mi et al., ) is designed to assign time-invariant representations to individual neurons based on population recordings of neural activity. The model uses a transformer architecture to implement an implicit dynamical system that predicts neural activity based on the past activity of the neuron itself and permutation-invariant statistics of the population activity. NeuPRINT learns both the dynamical model and time-invariant representations for each neuron by minimizing the prediction error of future neural activity. The learned representations can then be used for downstream tasks such as predicting transcriptomic cell types.

As the base model for our behavior-guided knowledge distillation in transcriptomic identity prediction, we follow the same configurations in Mi et al. () and set the dimension of time-invariant embedding to $64$. For model architecture, $1$ transformer layer with $2$ attention heads is used. For model training, the batch size is set to $1024$ and the learning rate is set to $1\times 10^{-3}$.

### A.5 Additional qualitative analysis on NDT and LFADS

This section provides more qualitative analysis based on visualization of model predictions in the behavior decoding task. As shown in Fig. 4, a comprehensive visualization comparing the performance of base models (NDT and LFADS) with their behavior-guided distilled versions (NDT-Hard-Distill and LFADS-Hard-Distill) on the MC-Maze dataset is presented.

Figure: Figure 4: Visualization of predicted hand velocity on 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.
Refer to caption: extracted/5895791/appendix_qualitative.png

#### A.5.1 1D Hand Velocity Decoding (X and Y axes)

Fig. 4(a) shows the decoded X and Y hand velocities compared to the ground truths.

NDT vs. NDT-Hard-Distill: The distilled NDT model demonstrates significant improvements in capturing both X and Y velocity components simultaneously. The base NDT model shows considerable deviations from the ground truth, particularly in the latter half of the time series. In contrast, NDT-Hard-Distill tracks both velocity components much more accurately throughout the entire sequence. It better captures the magnitude and timing of velocity changes in both dimensions, resulting in a much closer match to the ground truth trajectory.

LFADS vs. 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-Hard-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 a 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(b) presents a comparison of 2D hand movement trajectories for four models: NDT (base model), NDT-Hard-Distill (our distilled model), LFADS (base model), and LFADS-Hard-Distill (our distilled model), alongside the ground truth trajectory.

For NDT, NDT-Hard-Distill, and LFADS, there’s a 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 LFADS-Hard-Distill, the predicted trajectory closely follows the ground truth, suggesting much better performance for this model. Moreover, LFADS-Hard-Distill appears to have successfully learned the system’s 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 models’ ability to capture fine-grained details and reduce erratic predictions, resulting in smoother, more accurate velocity profiles. Meanwhile, LFADS-Hard-Distill demonstrates superior performance in accurately predicting the trough region, as illustrated in Fig. 4(a), and in capturing the trajectory trend, as shown in Fig. 4(b). 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 LFADS-Hard-Distill consistently outperforms the other three models in behavior decoding.

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 are 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 quantitative 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. 4), 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, we compare the performance of NDT and its different behavior-guided distilled variants in decoding hand movement trajectories and velocities.

Figure: Figure 5: Visualization of predicted hand velocity on MC-Maze dataset of different behavior-guided distilled models, including based model NDT, NDT-Hard-Distill, NDT-Soft-Distill, NDT-Feature-Distill, and NDT-Correlation-Distill. (a) Prediction and ground truth of X and Y velocities, respectively. (b) Prediction and ground truth of 2D hand movement trajectory.
Refer to caption: extracted/5895791/appendix_qualitative2.png

#### A.6.1 1D 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-guided distilled models demonstrate better tracking of the ground truth, with reduced oscillations and improved accuracy in predicting direction changes. Moreover, the distillation approaches are particularly effective in capturing the synchronous changes in X and Y velocities, suggesting better prediction of overall motion dynamics. NDT-Correlation-Distill 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 deviation from the ground truth, especially in the lower part of the trajectory. In contrast, all distillation methods demonstrate improved trajectory prediction, with closer alignment to the ground truth path. Specifically, NDT-Correlation-Distill appears to provide the most accurate trajectory prediction, closely matching the ground truth’s shape and endpoints.

#### A.6.3 Model-specific Observations

For each distillation strategy employ in this study, the following key observations can be discerned from Figure 5:

- •
NDT-Hard-Distill: Shows substantial improvement over base NDT in both velocity and trajectory prediction.
- •
NDT-Soft-Distill: Offers improved performance, though with some oscillations in velocity predictions.
- •
NDT-Feature-Distill: Demonstrates good overall performance, with smooth velocity curves and accurate trajectory prediction.
- •
NDT-Correlation-Distill: 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 predictive capabilities of the NDT model. The improvements are particularly noticeable in handling complex motions and maintaining accuracy over longer time horizons. The correlation-based distillation method seems to be the most effective, suggesting that preserving relational information during distillation is crucial for accurate predictions.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
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?

Could 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?
soundness: 4
strengths: Originality
- 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.

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

Clarity
- The idea and the main point of the paper are well explained. 

Significance 
- The authors provide new perspectives on how behavioral observations can be leveraged to guide the complex modeling of neuronal population dynamics.
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. 

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

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 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.
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).
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. 

NVM acronym is not defined

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See above
soundness: 3
strengths: - Writing is clear with adequate literature review;
- Solid technical/experimental contribution, the experiments are adequately done;
- 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.
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.
weaknesses: - Intuitions & clearness of section 3.2.1
1. 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.
2. 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? 

- Fairness and more details of model selection and evaluation protocol. 
1. 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?
2. 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.)

- The paper can benefit from more experimental case studies.
1. 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.

- Lack of theoretical & analytical results that supports the experimental results.

- Minor formatting suggestions:
1. Do the authors mind putting Table 1 and 2 at either the bottom or top of the pages?

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


Overall, 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.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: See Weaknesses.
soundness: 3
strengths: - The organization of this paper is well and the presentation of method is clear.

- A novel approach of BLEND, which conbines privilege knowledge (training teacher model) and knowledge distillation (training student model), is proposed for neural activity reconstruction.

- BLEND has achieved a good performance on two benchmarks.
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.
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.

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

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


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

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

Meta-review: The paper presents a novel approach where a teacher model, trained using both neural activity and behavioral signals, distills its knowledge to a student model that relies solely on neural activity. This method is applied to neural population activity analysis and transcriptomic identity prediction tasks, showing improvements over existing models. The paper is well-structured, clearly written, and effectively demonstrates the potential of behavioral data to enhance the modeling of neuronal dynamics, outperforming state-of-the-art methods.

Decision: Accept (Poster)