Exploring Behavior-Relevant and Disentangled Neural Dynamics with Generative Diffusion Models

Authors: yule  wang
chengrui  li
weihan  li
anqi  wu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-04-30 02:31:27

Abstract: Understanding the neural basis of behavior is a fundamental goal in neuroscience. Current research in large-scale neuro-behavioral data analysis often relies on decoding models, which quantify behavioral information in neural data but lack details on behavior encoding. This raises an intriguing scientific question: "how can we enable in-depth exploration of neural representations in behavioral tasks, revealing interpretable neural dynamics associated with behaviors". However, addressing this issue is challenging due to the varied behavioral encoding across different brain regions and mixed selectivity at the population level. To tackle this limitation, our approach, named ("BeNeDiff"), first identifies a fine-grained and disentangled neural subspace using a behavior-informed latent variable model. It then employs state-of-the-art generative diffusion models to synthesize behavior videos that interpret the neural dynamics of each latent factor. We validate the method on multi-session datasets containing widefield calcium imaging recordings across the dorsal cortex. Through guiding the diffusion model to activate individual latent factors, we verify that the neural dynamics of latent factors in the disentangled neural subspace provide interpretable quantifications of the behaviors of interest. At the same time, the neural subspace in BeNeDiff demonstrates high disentanglement and neural reconstruction quality.

Main paper: 
 introduction: 
 Understanding and elucidating the complex interrelationships between behavioral data and neural population activity is a long-standing goal in systems neuroscience [Batty et al., 2019; Gomez-Marin et al., 2014; Krakauer et al., 2017; Berman, 2018] . Exploring the neural basis of behavior not only deepens our basic knowledge of brain functions but also establishes a foundation for developing improved treatments for psychiatric and neurological conditions [Vieira et al., 2017; Ibáñez et al., 2018] . Significant progress has been achieved in developing computational toolkits for neuro-behavioral decoding by using behavior video data [Whiteway et al., 2021; Batty et al., 2019; Musall et al., 2019; Stringer et al., 2019] . These methods perform region-based behavior decoding to map neural activity across multiple brain regions of the dorsal cortex to the behaviors from the videos. However, these methods only quantify how much behavioral information is encoded in neural populations, but do not reveal the details of such encoding. There has been markedly less focus, with cortex-wide signals, on enabling in-depth exploration of neural activities during behavioral 38th Conference on Neural Information Processing Systems (NeurIPS 2024).Figure 1 : Empirical study across multiple brain regions of dorsal cortex neural recordings of a mouse in a visual decision-making task. (A) The Brain Atlas map [Lein et al., 2007] . (B) Neural signals in various brain regions (SSp, MOs, and VIS) exhibit mixed selectivity in behavior of interest decoding. "Levers", "Spouts", "Paw-(x)", and "Jaw" are four behaviors of interest. cvR 2 is short for cross-validation coefficient of determination. The higher, the better.tasks, where specific neural patterns reveal dynamic evolutions corresponding to distinct behaviors of interest. However, empirically addressing this scientific question is challenging due to neural population activities in various brain regions exhibiting mixed selectivity [Sani et al., 2021; Hasnain et al., 2023] , responding robustly to multiple behaviors of interest. We further verify this finding through an empirical study across three brain regions on the dorsal cortex of head-fixed mice [Musall et al., 2019] (shown in Figure 1 ).To tackle this issue, we propose a method -Exploring Behavior-Relevant and Interpretable Neural Dynamics with Generative Diffusion Models -("BeNeDiff"). We first employ a neural latent variable model (LVM) to identify orthogonal and disentangled neural latent subspace. This is achieved through a semi-supervised variational autoencoder, which integrates behavioral labels to rotate the subspace. Subsequently, our main idea is to explore the neural dynamics of each latent factor in the learned subspace for distinct quantifications of the behaviors of interest. However, such a workflow is non-trivial since naïve latent manipulation produces samples not conform to the original distribution, leading to mapped video-based behavioral data that loses its validity (we further detail this part in Method Section 3.2.1).Notably, we aim to investigate the behavioral-specificity of neural latent factors in a generative fashion. We leverage state-of-the-art video diffusion models (VDMs) to generate behavior videos predicted to activate individual latent factors along the single-trial trajectory. Technically, the VDMs are capable of capturing the overall temporal dynamics and synthesizing behavior videos in a classifier-guided manner [Dhariwal and Nichol, 2021] . Inspired by Noise-Contrastive Estimation [Gutmann and Hyvärinen, 2010] , the guidance objective is formulated to amplify the variance of the selected latent factor along its neural trajectory while suppressing the variance of the neural trajectories of the other latent factors.We conduct experiments to verify the efficacy of BeNeDiff on a widefield calcium imaging dataset, where a head-fixed mouse performs a visual decision-making task across multiple sessions [Musall et al., 2018; 2019] . The neural subspace in BeNeDiff exhibits high levels of disentanglement and neural reconstruction quality, as evidenced by multiple quantitative metrics. By guiding the diffusion model to activate individual latent factors, we verify that the neural dynamics within the disentangled subspace provide interpretable and selective quantifications of the behaviors of interest (e.g., paw movements) across multiple brain regions. These results advance our understanding of neuro-behavioral relationships through the identification of fine-grained behavioral subspaces and the uncovering of disentangled neural dynamics.To highlight our major contributions: (1) This is the first work to explore wide-field imaging across multiple brain regions of the dorsal cortex of head-fixed mice during a decision-making task using neural subspace analysis, rather than merely performing neuro-behavior decoding. We uncover disentangled neural representations for various behaviors. (2) To visualize the behavior dynamics within a disentangled neural subspace of each brain region, we develop a novel VDM-based interpretation tool that faithfully reflects behavior-related neural dynamics. It is essential to interpret the meaning of each neural latent dimension as well as the behavior dynamics it encodes.Figure 2 : Schematic diagram of neural dynamics interpretation with BeNeDiff. We first employ a neural LVM to identify a disentangled neural latent subspace (the left part). Then, we train a linear neural encoder to map behavior video frames to neural trajectories. We use video diffusion models (VDMs) to generate behavior videos guided by the neural encoder, based on the objective of activating the variance of individual latent factors along the single-trial trajectory. This approach provides interpretable quantifications of neural dynamics in relation to the behaviors of interest.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: In Section 1 Introduction.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
problem formulation: 
 We first provide the notations of the paired neuro-behavioral observations. The single-trial neural population activities are denoted as X = [x 1 , . . . , x L ] ⊤ ∈ R L×N ,Y = [Y 1 , . . . , Y L ] ⊤ ∈ R L×H×W, where H, W are the height and width of the compressed behavior video frames. We extract behavior labels U = [u 1 , . . . , u L ] ⊤ ∈ R L×B from the video frames using a behavior LVM [Whiteway et al., 2021] . B is the number of the behavior.We build a variational autoencoder (VAE) [Kingma and Welling, 2013] to infer the neural latent trajectories Z = [z 1 , . . . , z L ] ⊤ ∈ R L×D , which are also informed by behavioral labels. D is the latent factor number. We denote its probabilistic encoder and decoder as q ψ (Z | X, U) and p ϕ (X, U | Z), respectively. We denote the neural trajectory of a single latent factor as z (d) = Z :,d , where d ∈ {1, 2, • • • , D}. Our primary goal is to investigate the neural dynamics of z (d) through selectivity quantifications of its corresponding single-trial behavioral video data Y.Diffusion models have also achieved impressive results in video synthesis over recent years [Ho et al., 2022b; a; Harvey et al., 2022] . VDMs process a fixed number of frames and factorize them over the temporal dimension via a deep neural network [Ho et al., 2022a; Harvey et al., 2022] . The training of VDMs starts from a forward process with a variance schedule {β 1 , . . . , β T }, the noised sample Y t follows the Gaussian conditional:q (Y t | Y 0 ) := N (Y t ; √ ᾱt Y 0 , (1 -ᾱt ) I), where α t := 1 -β t and ᾱt := t s=1 α s . A denoising model εθ (•) is trained to reverse the forward process using a weighted mean squared error loss:EQUATIONin which time-steps t are uniformly sampled and w(λ t ) is the weighting ratio. This loss function can be justified as optimizing a weighted variational lower bound on the data log-likelihood. In the sampling phase, we start from Y T ∼ N (0, I L×H×W ) and perform step-by-step denoising,EQUATIONwhere random noise perturbation ϵ t ∼ N (0, I L×H×W ) for timesteps t > 1, ϵ t = 0 when t = 1, andσ 2 t = 1-ᾱt-1 1-ᾱt β t .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: In Section F Discussion.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methods: 
 Then, we train a linear neural encoder from the behavior video frames to the neural trajectories. We leverage video diffusion models (VDMs) to generate behavior videos guided by the neural encoder, based on the objective of activating the variance of individual latent factors along the single-trial trajectory, providing interpretable quantifications of neural dynamics with respect to the behaviors of interest.In this section, we first detail the process by which BeNeDiff infers a disentangled neural latent subspace. We then discuss the approach that BeNeDiff interprets the selectivity of neural dynamics of latent factors using the video diffusion model.Drawing inspiration from recent progress in the field of neural LVMs [Kingma et al., 2014; Klys et al., 2018] , we employ a VAE to learn a disentangled neural subspace. The neural data X usually contains a good amount of information other than behavior [Hasnain et al., 2023] , thus an unsupervised disentangled VAE won't effectively discover disentangled subspace with behavior only. Therefore, we introduce behavior labels U to inform the VAE to learn a latent subspace that better accounts for the variance related to behavior. We note that this technique is widely adopted in previous neuro-behavioral analysis works [Wang et al., 2024; Schneider et al., 2023; Gondur et al., 2023] . Notably, to enforce the disentanglement in the latent subspace, we incorporate a total-correlation (TC) penalty term [Chen et al., 2018] to enforce the VAE to find statistically independent latent factors in the semi-supervised setting. The VAE optimizes the following evidence lower bound (ELBO) [MacKay, 2003] :log p ϕ (X, U) ≥E q ψ (Z|X,U) log p ϕ (X | Z) Neural Reconstruction + log p ϕ (U | Z)Behavior Info.-D KL q ψ (Z | X, U) p(Z) Regularization Term -β D KL q ψ (Z | X, U) d q ψ (z (d) | X, U) Total Correlation =: -L VAE (ϕ, ψ)(3) in which z (d) denotes the neural trajectory of the d-th latent factor, and the value of β controls the strength of disentanglement penalty. However, the factorial density in this term is untractable in practice, so here we use the minibatch-weighted sampling estimator [Chen et al., 2018] to approximate the TC penalty term. We note that the variational autoencoder employs a sequential architecture [Fabius and Van Amersfoort, 2014] to capture the overall temporal dynamics along the single-trial trajectory {x l } L l=1 , plugging bi-directional recurrent units [Schuster and Paliwal, 1997 ] into both the probabilistic encoder q ψ (•) and decoder p ϕ (•).As for testifying the neural dynamics of a single disentangled latent factor z (d) on the behavioral videos Y, a straightforward attempt is to train a neural-net model to approximate the posterior distribution p(Y | Z) and then perform latent manipulation on each single latent factor. There are two major techniques to perform latent manipulation. The first is a naïve manipulation. This method manipulates a single subspace z (d) while keeping the non-target latent factors fixed at arbitrary values. It then observes how the manipulation affects Y. The induced changes in the videos reveal the dynamics encoded by z (d) . The second method uses classifier-free guidance [Ho and Salimans, 2022] , where we allow the activated latent factor z (d) to evolve while fixing non-target latent factors to arbitrary values. However, setting arbitrary values without knowing the true distributions of non-target subspaces can lead to unnatural distortions in generated videos, complicating the visualization and interpretation of genuine animal behavioral dynamics.So here we employ the video diffusion models (VDMs) to explore factor-wise neural dynamics through a generative manner, which is capable of maintaining temporal consistency and behavioral dynamics across frames. The primary goal is to perform behavior data generation conditioned on activating a single latent factor along the neural trajectory. Thus the resulting behavior video can provide interpretable quantifications of the neural dynamics of factor z (d) . Specifically, we implement classifier guidance [Kawar et al., 2022] . By Bayes rule, we obtain the following posterior density and gradient [Mardani et al., 2023] :EQUATIONEQUATIONin which θ, λ are the parameter sets for the classifier and the denoising model, respectively. Note that t indicates the time step in the diffusion model. Our goal is to estimate the two terms on the RHS of Eq. ( 5) to perform conditional denoising in each step. We first approximate the density of the behavior video data through a standard denoising model εθ (Y t , t) according to Eq. ( 1) since the first unconditional gradient term can be derived through it:EQUATIONFor the calculation of the guidance term, we first train a linear neural encoder as the classifier from the behavior video data to the neural latent variables of the learned semi-supervised VAE subspace.We denote the estimated neural latent trajectories as Ẑt = [ẑ t,1 , . . . , ẑt,L ] ⊤ ∈ R L×D , in which:EQUATIONwhere 1 ≤ l ≤ L, ẑt,l ∈ R D denotes the estimated value of latent factors at time bin l and diffusion step t. The parameter set of the linear encoder λ = {W, Q}. W ∈ R D×M is the linear transformation matrix, Q ∈ R D×D is the covariance matrix and vec(•) represents vectorizing the two-dimensional video frame into column vector. After training the encoder, we fix all parameters and use it to construct the densityp λ (c | Y t ).The class labels c ∈ c (1) , c (2) , . . . , c (D) , in which c (d) is a one-hot column vector with a one at the d-th dimension and zeros elsewhere. Drawing inspiration from Noise-Contrastive Estimation [Gutmann and Hyvärinen, 2010] , our guidance objective of the activation of latent factor d-th is formulated as maximizing the variance of the trajectory ẑ(d) while minimizing the variance of the other latent factor trajectories in Ẑ:EQUATIONwhere d) calculates the variance of the selected latent factor and j) , j ∼ Uniform({1, 2, . . . , D}\{d}) calculates the variance of another sampled latent factor's trajectory. Var Ẑ ∈ R 1×D is a row vector where each element is the variance of every latent factor along the neural trajectory. τ is the temperature parameter. K is a hyperparameter controlling the number of sampled negative samples at each iteration.f + λ Ẑ, c (d) = Var Ẑ c (f - λ Ẑ, c (d) = Var Ẑ c (The gradient ∇ Yt log p λ c (d) | Y t is computed using automatic differentiation [Paszke et al., 2017] . Algorithm 1 describes the guided behavior video generation steps of our proposed framework BeNeDiff.Input: Condition label c (d) for interpreting the neural dynamics of the d-th latent factor Initiate Y T ∼ N (0, I L×H×W )for t = T to 1 do ε′ θ,λ (Y t , t) = εθ (Y t , t) - √ 1 -ᾱt ∇ Yt log p λ c (d) | Y t ϵ t ∼ N (0, I) if t > 1, else ϵ t = 0 Y t-1 = 1 √ αt Y t -1-αt √ 1-ᾱt ε′ θ,λ (Y t , t) + σ t ϵ t end Output: Generated behavior video Y 0 4 Related Works Disentangled Latent Subspace Learning.Neural LVMs is a fundamental framework which posits that single-trial neural population activities rely on low-dimensional "neural manifolds" [Gallego et al., 2018; Mitchell-Heggs et al., 2023; Li et al., 2023a] and their extracted latent variables are successful in describing single-trial neural activities [Li et al., 2024c; 2022; 2024d; Liu et al., 2021; 2022; Li et al., 2024a] . Learning disentangled latent variables that uncover statistically independent latent factors [Chen et al., 2018] can provide enhanced robustness, interpretability, and controllability. Typically, this type of work involves adding auxiliary regularizer terms to enhance orthogonality [Mathieu et al., 2019] and reduce the total correlation [Chen et al., 2018] among the latent factors. In the field of neuroscience, there have been studies focusing on the disentanglement of latent subspace within rich behavioral data [Whiteway et al., 2021; Shi et al., 2021] . However, our work is the first to discover interpretable and disentangled latent subspaces of wide-field imaging data.Generative Diffusion Models. In recent years, diffusion models have achieved great success in generating high-quality images due to their expressivity and flexibility [Ho et al., 2020; Song et al., 2020a; b; Vahdat et al., 2021] . Moreover, for the more challenging task of video generation, there have been several explorations using diffusion models to address it. From a modeling perspective, the key concern is how to maintain temporal dynamics and consistency across frames. Most existing works [Ho et al., 2022b; a] extend the 2D U-Net architecture [Ronneberger et al., 2015; Song et al., 2024] to a 3D framework by considering the time axis. In this 3D framework, convolutions are performed in both spatial and temporal dimensions. Additionally, recent studies in neural computation have leveraged generative diffusion models to tackle domain-specific tasks, such as neural distribution alignment [Wang et al., 2024] and decoding visual stimulus from brain activities [Sun et al., 2024b; a] . Our work is the first to employ generative diffusion models for analyzing neuro-behavioral data relationships. A head-fixed mouse performed a visual decision-making task while neural activity across the dorsal cortex was optically recorded using widefield calcium imaging [Musall et al., 2019; Churchland et al., 2019] . The mouse's behavior included both instructed and uninstructed movements. For behavioral data acquisition, two cameras captured video frames from both a side view and a bottom view. The dataset comprises 1126 trials conducted over two sessions, with 189 frames per trial at a frame rate of 30 Hz. Concurrently, neural activity was recorded at the same frame rate. The grayscale video frames were downsampled to 128×128 pixels. We extract 275 dimensions of neural signals from the high-dimensional widefield imaging data using the open-sourced LocaNMF decomposition toolkit [Saxena et al., 2020] . As shown in Figure 3 , the behaviors of interest include the moving lick spouts, moving levers, the single visible right paw trajectories, and the movement of the jaw and chest, all tracked using DeepLabCut [Mathis et al., 2018] .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
disentangled neural latent subspace investigation: 
 We note that we train a unique neural LVM for each individual brain region (single-region), and we evaluate both the behavior decoding and neural reconstruction performance of each brain regionspecific neural latent trajectories.Single Latent Factor Behavior Decoding. In order to verify the disentanglement of the learned neural subspace in BeNeDiff, we evaluate the behavior label decoding performance of each individual latent factor. Specifically, we train a unique linear regressor for each latent factor from the VAE and plot the decoding accuracy as the R-squared value (R 2 %). The results of VIS-Right region (the right visual region) are shown in Figure 4 . The main observation is that each latent factor is specific to a unique behavior of interest, confirming the orthogonality and clear disentanglement of the inferred latent trajectories from a quantitative perspective.Neural Observation Signals Reconstruction. To prevent the VAE from overfitting to the behavioral labels, BeNeDiff also aims to maintain a low reconstruction error for neural activity. Table 1 presents the quantitative reconstruction results compared to baseline methods, including Semi-Supervised Learning (SSL) [Kingma et al., 2014] , CEBRA [Schneider et al., 2023] , and pi-VAE [Zhou and Wei, 2020] . The table records the R-squared values (R 2 , in %) and RMSE for each method. Additionally, we plot the ground-truth neural signals and the reconstructed signals of several methods in a single trial in Figure 5 . The main observation is that the neural reconstruction is well-preserved given the behavioral priors. One possible explanation is that the behavioral labels rotate the latent subspace while preserving the necessary information for reconstructing the neural data. The neural signals can be hardly recovered from the behavior labels only. It indicates that the behavior-informed latent does encode significant neural information that is not contained in the behavior labels. Furthermore, we evaluate the disentanglement quality of the latent subspace using the widely-adopted MIG (Mutual Information Gap) metric [Chen et al., 2018] , also listed in Table 1 . We observe that the learned latent subspace of BeNeDiff significantly enhances disentanglement compared to the vanilla VAE. We observe that the neural reconstruction quality from the latent subspace of BeNeDiff is maintained given the behavioral labels. "Self-Supervised" denotes the VAE w/o behavior labels.From the quantitative experiments in the previous subsection, we obtained information about the decoding and disentanglement quality within the subspace. However, these metrics have limitations in interpreting single-trial neural dynamics, especially the complex temporal structures over time. Here, we visualize the generated videos from BeNeDiff and the baseline latent manipulation methods, demonstrating that BeNeDiff provides interpretable quantifications of the behaviors of interest.Latent Manipulation Methods for Comparison. We compare the neural dynamics exploration performance of BeNeDiff against the following two latent manipulation methods:• Naïve Latent Manipulation: the standard manipulation method discussed in Section 3.2.1, which approximates the posterior of behavioral videos given the neural latent trajectories p(Y | Z), using a neural network that incorporates recurrent units and spatio-temporal convolutional layers.• Classifier-free Guidance [Ho and Salimans, 2022] : a method that approximates the posterior p(Y | Z) with diffusion models. It co-trains a conditional and an unconditional diffusion model together, combining the resulting conditional and unconditional scores at each diffusion step. In the conditional model, the entire neural latent trajectory Z is set as the condition, formulating the denoiser as ε (Y t , Z, t). For the manipulation of the latent, we keep the activated latent factor z (d) to evolve while setting the values of the other latent factors to those in the first frame of the trial.Setup. To verify the neural dynamics interpretation capability of BeNeDiff, we generate behavioral video data given the activation of each behavior of interest (generated trials with the activation of Jaw and Paw-(y) are shown in Figure 6 and Figure 9 , respectively). For visualization and video analysis, we plot frames at intervals of five and compute their frame differences. The conditional module of the classifier-free guidance method is trained with an auxiliary convolutional head. Compared to general video synthesis [Harvey et al., 2022; Esser et al., 2023] , our behavioral video data are more focused on maintaining the temporal dynamics and consistency across video frames, thus in BeNeDiff, we tailor the standard 3D U-Net architecture [Çiçek et al., 2016] from temporal self-attention layers to temporal convolutions layers [Li et al., 2023b; 2024b] to maintain local temporal consistency. While we keep the spatial self-attention layers the same. The diffusion model is trained on an Nvidia V100, using approximately 20 computer hours.Results Analyses of the Generated Videos. As shown in Figure 6 , for the naïve latent manipulation method, the distribution of neural signals often falls outside the original distribution after manipulation, resulting in blurred generated frames. The frame differences are entangled, and the "Jaw" latent factor affects the entire head movement of the mouse, particularly in the first four frames shown.Classifier-Free GuidanceInter-Frame DifferenceFigure 6 : Generated Single-trial Behavioral Videos with Latent Factor Guidance from the side view. Compared to baseline methods, we observe that the neural dynamics of latent factor in the results of BeNeDiff show specificity to the "Jaw" movements.On the other hand, for classifier-free guidance, the generated videos maintain coherent consistency between frames. However, it does not interpret neural dynamics well in this context, resulting in a trajectory with small movements in the "Jaw". This is because the overall latent trajectory is used as the input to the model and the other latent factors are kept fixed, making it difficult to discriminate the evolution of a single factor effectively. In contrast, the results of BeNeDiff show more specificity to the targeted behavior of interest. The inter-frame differences in BeNeDiff's results are clearly specified to the "Jaw" movements, and the structure of the neural dynamics is well-preserved and consistent with ground-truth "Jaw" behavior trajectories. A similar pattern is evident with the other latent factors, as shown in Figures 9, 10, and 11 in the appendix.Besides the capability of revealing interpretable neural dynamics of each latent factor associated with behaviors, here we further investigate the neural dynamics differences across brain regions through BeNeDiff. As shown in Figure 7 and Figure 12 in the appendix, we present the 2D neural latent trajectories of two latent factors, specifically related to "Paw-(x)" and "Paw-(y)", across six brain regions for two randomly selected trials. From the starting point of the trial, we observe that the latent trajectories corresponding to the left and right hemispheres of the VIS both show a noticeable change starting earlier. Next, the SSp regions show a large shift in activity, followed by a similar change in the MOs regions. However, it is difficult to clearly visualize the specific motion encoded by each region and to distinguish how different the motions are encoded solely based on neural trajectory plots. This further highlights the need for using a video diffusion model for visualization and interpretation.In contrast, in the generated behavior video samples of BeNeDiff (as illustrated by the frame differences in Figure 8 and Figure 13 in the appendix), where the "Paw-(x)" and "Paw-(y)" latent factors are activated, the behavioral dynamics encoded by these two latent factors are observed across different brain regions. First, paw movements are detected in the VIS regions before the "Levers" come in. This early activity in VIS could reflect its role in the predictive coding of behaviors, indicating that this region may predict motor movements before they happen. Next, the SSp regions exhibit paw movements that are synchronized with the onset of the "Levers", indicating a potential role for SSp in processing somatosensory feedback. Subsequently, in the MOs regions, paw movements are observed following the "Levers" onset, which is consistent with MOs' role in motor execution and control, occurring slightly after SSp. Inter-Frame DifferenceInter-Frame Difference"Levers" Coming In Point [2014] for optimization and the learning rate is set as 0.001. The batch size is uniformly set to 32. The latent subspace factor number is fixed at 6, which is the same as the number of behaviors of interest. We employ the dropout technique Srivastava et al. [2014] and the ReLU activation function Rasamoelina et al. [2020] between layers in our probabilistic encoder and decoder neural networks.Training Details of Video Diffusion Models. We adopt the architecture of the VDM of 3D-UNet [Ho et al., 2022b] with the ϵ-parameterization. We use both spatial attention and spatial convolutions.The temporal convolutions are used to maintain consistency between frames. The embedding input size to the UNet architecture is set as 32 and the UNet has three downsampling and upsampling layers. The diffusion timestep is set as 200. The training batch size is set as 64, with a learning rate of 0.001. We use Group Normalization.B In-depth Investigation on the neural LVM module across brain regions Using Figure 9 as an example, for the naïve latent manipulation method, the generated frames are in a reasonable form. Nevertheless, the frame differences are still intertwined, and the latent factor of "Paw-(y)" heavily affects the "Spout" movement. Meanwhile, for classifier-free guidance, the trajectories focus on the mouse movements, but they are still entangled with the "Chest" movements. In contrast, the results of BeNeDiff show more specificity to the targeted behavior of interest. The inter-frame differences in BeNeDiff's results are clearly specified to the "Paw-(y)" movements, and the temporal evolution of the neural dynamics is coherent with real-world mouse paw trajectories. The generated results in Figures 10 and 11 show a similar trend, demonstrating specificity to the Paw-(x)" and Spout" factors.Classifier-Free GuidanceInter-Frame DifferenceFigure 9 : Generated Single-trial Behavioral Videos with Latent Factor Guidance from the bottom view. Compared to baseline methods, we observe that the neural dynamics of a latent factor in the results of BeNeDiff show specificity to the "Paw-(y)" movements.Inter-Frame DifferenceClassifier-Free GuidanceInter-Frame DifferenceFigure 10 : Generated Single-trial Behavioral Videos with Latent Factor Guidance from the bottom view. Compared to baseline methods, we observe that the neural dynamics of a latent factor in the results of BeNeDiff show specificity to the "Paw-(x)" movements.Inter-Frame DifferenceClassifier-Free GuidanceInter-Frame DifferenceFigure 11 : Generated Single-trial Behavioral Videos with Latent Factor Guidance from the bottom view. Compared to baseline methods, we observe that the neural dynamics of a latent factor in the results of BeNeDiff show specificity to the "Spout" movements. Inter-Frame DifferenceInter-Frame Difference Our study introduces BeNeDiff, a novel approach leveraging behavior-informed latent variable models and generative diffusion models to uncover and interpret neural dynamics. Through empirical validation, we demonstrate that BeNeDiff effectively identifies a disentangled neural subspace and synthesizes behavior videos that provide interpretable insights into neural activities associated with distinct behaviors of interest. However, for the neural latent variable model (LVM) module, there exists a balance between disentangling the neural subspace with behavior semantics and maintaining neural reconstruction performance. For each brain region and session, at this stage, a careful hyperparameter search is necessary to balance the weight between these two components. For the generative video diffusion module, we implement the neural encoder (classifier for guidance) as a linear regressor for interpretability. This linear assumption can be relaxed later for improved guidance performance.NeurIPS Paper Checklist 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: The paper presents a two stage approach for investigating latent representations derived from neural activity.

First, the authors present a TC-VAE with additional supervision (akin to pi-VAE), which they claim provides additional disentanglement of the feature space. Reviewers agree that the claims made regarding disentanglement are too strong, and **should be tuned down in the revision of the paper**. In particular, the standards of empirical evaluation are below what would be usually expected at NeurIPS on a standardized ML benchmark, and the authors should put care into the wording for the camera-ready version.

Regarding experiments, the additional evaluations made during the rebuttal phase are great (especially adding more methods like pi-VAE and CEBRA, and adding the TC regularization there) and should be added to the main manuscript and discussed in the context of disentanglement. Ideally, more disentanglement metrics beyond MIG should be reported in the tables to get a more holistic impression of the model performance.

Finally, the authors should make it clearer that the system is a combination of multiple potentially independent methods, vs. an end-to-end learnable system where the diffusion model informs the latent space -- this seems to be suggested by wording and e.g. the diagram in Figure 2. Better highlighting that the method of estimated the latents is separate from the "visualization" method would improve the clarity.

I recommend acceptance on the paper conditional on these modifications for the camera-ready revision of the paper.

Reviews: review: paper_summary: paper_summary: The authors implement a VAE with behavioral labels to identify behavior-specific latent variables. They use these latents to then create a generative model for experimental video. They focus on the application of their model to a wide-field cortical recordings of a mouse during an experimental visual task. They show that their model identifies disentangled behaviorally relevant latents and can create experimental video where identifiable behavioral elements are generated from specific neural latents.
strength_weakness: strength_weakness: Each of the primary contributions to neural LVM type models (the use of total correlation and the diffusion modeling) are not sufficiently benchmarked against competing approaches. It is unclear how this model compares to the others used in neuroscience. If the authors wanted to strengthen the case that their model is an important novel contribution for behaviorally-relevant latent disentangling, they could compare how total correlation VAE compares to other methods where behavioral labels are used to help disentangle the latent states. One important comparison that I would be curious to see that isn't cited would be the pi-VAE1, but the others the authors cite could be used. Specifically, these other models could be added to figures 4 and 5.

The authors could also focus on how using video diffusion is a novel application in this setting. Again, some additional benchmarking with existing deep LVMs in neuroscience and some further discussion and emphasis as to how the diffusion methodology solves an important problem in neuroscience that perhaps existing methods cannot solve would strengthen this paper.

Though the authors highlight the application of the model to this specific dataset as a contribution, this alone would be more appropriate for a more experimentally-focused venue. For neurips, the focus should be on the model advancements and their applications, not the significance of the dataset.

1 - Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE - Zhou and Wei
questions: questions: NA
limitations: limitations: The authors only minimally discussed the limits of the ability of their model.
paper_summary: paper_summary: The paper proposed BeNeDiff to learn disentangled neural trajectories together with a generative diffusion mode for behavior. BeNeDiff leverages beta-VAE for learning of disentangled space space Additional behavior generation module is applied for interpretation of the neural subspace. It is interesting to see that the learned latent latent trajectory is separated for different behavior types. The qualitative results and quantitative results demonstrate the superior of the proposed method.
strength_weakness: strength_weakness: 1.It is unclear to me if the behavior video generation part is trained together with the disentangled subspace learning. It seems true based on the figure, but didn't see how they are optimized together from text.

2.it is interesting to learn the disentangled neural space. Some visualization of the learned latent space is beneficial. 

3.Additional ablation study on different contribution of the loss function could be performed.
questions: questions: 1.How training and testing data are split? Are they session wise or random? How the method can be generalized to different sessions or animals?

2.Does the learned latent trajectory show some clustering property, or somehow organized, could you visualize it?

3.For session 3.2, $c$ is better to be described as it is first mentioned in eq.5. Does the number of class is equals to feature dimension D and the same as the number of the behavior classes"? 

4.To get figure 4, which latent variable do you use, $\hat{Z}$ or $Z$?

5.For fig.6, and fig.7, could you show the difference of the ground truth frame with the same behavior type as comparison.
limitations: limitations: It could be better the paper include more analyze on the learned latent trajectory and see if it could provide some scientific insight.
paper_summary: paper_summary: This study presents BeNeDiff that integrates behavior-informed latent variable models with state-of-the-art generative diffusion models to investigate neural dynamics during behavioral tasks. The methodology involves identifying fine-grained and disentangled neural subspaces and synthesizing behavior videos to provide interpretable quantifications of neural activity related to specific behaviors. The authors validate their approach using multi-session widefield calcium imaging datasets, demonstrating high levels of neural disentanglement and reconstruction quality.
strength_weakness: strength_weakness: 1. This work primarily relies on a behavior LVM for neural activity disentanglement, which can be influenced by the newly introduced total correlation (TC) penalty term with $\beta$ weighted. The impact of this penalty should be considered through ablation studies.


**Not a Weakness**: There is significant research combining behavioral labels with latent variables of neural activity [1,2], considering the effects of different disentanglement algorithms on behavior encoding would further enhance this study.

[1] Zhou, Ding, 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 (2020): 7234-7247.

[2] Schneider, Steffen, Jin Hwa Lee, and Mackenzie Weygandt Mathis. "Learnable latent embeddings for joint behavioural and neural analysis." Nature 617.7960 (2023): 360-368.
questions: questions: I have no question.
limitations: limitations: see WeaKnesses part.

Comments: title: Thank you
Synthesizing Videos with Behavior Classifier Latents
Thank you
Thank you
Thank you
comment: Dear Reviewer CLKJ,

We appreciate your positive feedback and the adjustment to your score. As mentioned in our previous response, we will include the video results synthesized from behavior classifier latents in the supplementary materials of the revised manuscript. Thank you once again for your valuable suggestions and comments.
Dear Reviewer CLKJ,

We thank you for the valuable response. We agree that incorporating behavior labels can be helpful for learning a disentangled neural subspace. However, we note that our neural LVM module in BeNeDiff maintains a high neural data reconstruction rate (as shown in Table 1 of the manuscript), thereby preserving the neural dynamics in the latent trajectories.

It is true that we can train a probabilistic behavior classifier (a behavior decoder from neural data $\mathbf{X}$ to behavior labels $\mathbf{U}$) to approximate the conditional distribution $p(\mathbf{U} \mid \mathbf{X})$ and derive latent variables $\mathbf{Z}'$ from it. The $\mathbf{Z}'$ would also be disentangled. However, such $\mathbf{Z}'$ lacks physical or semantic meaning, as there is no reconstruction of the neural data, making it unable to interpret neural dynamics or provide any scientific insights. We further trained a diffusion model guided by $\mathbf{Z}'$ and presented the visualization results. We have sent the AC a comment linking to an anonymous repository containing these synthesized video results. In accordance with this year’s NeurIPS guidelines, please have the AC forward it to you.

We observe that while the synthesized videos by activating each latent factor does show some specificity to a particular behavior of interest, **the overall video results are quite chaotic and do not align with the ground-truth behavioral trajectories.** **The behavior movements are overly drastic and contain many unnatural distortions** compared to the videos generated by the neural latent trajectories in BeNeDiff. These videos will be included in the supplementary materials of the revised manuscript. Additionally, we have included the videos generated by BeNeDiff and the ground-truth behavior videos in the anonymous repository provided to the AC for comparison.

Looking ahead, our future research aims to generalize BeNeDiff to settings where the latent (feature) dimension is larger than the number of behavior classes. We hope that these responses and the above rebuttals address your concerns and will be considered when determining the final review score.
Dear Reviewer o3cJ,

We sincerely appreciate your positive evaluation of our method's contribution and additional benchmarking. We apologize for the confusion in the Introduction section and will clarify the model's generalizability across datasets when discussing the major contributions in the revised manuscript. We thank you once again for the valuable feedback and suggestions.
Thank you for the response. I think with this added detail the paper's contribution is clearer and I'd suggest emphasizing and clarifying the contribution similarly in a revised manuscript. I also appreciate the additional benchmarking with the total correlation measure. I have adjusted my score accordingly.

Regarding my comment about your emphasis on the data set - I did not mean to suggest that the analysis is specific to this dataset, but rather that your phrasing in the manuscript emphasizes the data set in-and-of-itself as a major contribution. The end of the introduction reads "To highlight our major contributions: (1) This is the first work to explore wide-field imaging of
the dorsal cortex of mice during a decision-making task using neural subspace analysis," and on like 167 "However, our work is the first to discover interpretable and disentangled latent subspaces of wide-field imaging data". I agree that your methods are generally applicable across datasets and I believe, therefore, that the model and it's generality should be what is primarily emphasized when discussing the major contributions in the manuscript.
Dear Reviewer o3cJ,

Thank you for your positive feedback and for adjusting your score. As mentioned in our rebuttal, we will incorporate the neural LVM module comparisons and the VDM module-generated video visualizations for multiple datasets into the appendix of the revised manuscript.
Thanks for the responses. The responses addressed most of my concerns.

But towards Q2, if the number of latent factor $D$ and the number of behavior of interest are the same, use the class label as guidance would obviously be helpful to learn a disentangled latent space. Directly learning a classifier and use the latent from the classifier would also be disentangled.

It would be good to provide the latent variables from the trained classifier to diffusion model to see if the behavior videos could be synthesized from these variables.
Dear Reviewer HCYS,

Thank you for your timely response. We will include the ablation study of the hyperparameter $\beta$ within the total-correlation term and present the results in the appendix of the revised manuscript. Additionally, we will assess the generalizability of the video diffusion modeling (VDM) module to the pi-VAE and CEBRA neural LVM baselines (with the total-correlation penalty term) to discover neural dynamics in their learned disentangled subspace. The generated video results (similar to Figures 6 and 7 of the manuscript) will be  added to the supplementary materials.

We again appreciate your evaluation and recognition of our work.
Thank you for the author's response. The new information meets my expectations, and I hope this content can be added to the supplementary materials. Thus, I maintain my original score.