Paper: Multi-modal Gaussian Process Variational Autoencoders for Neural and Behavioral Data (ICLR-2024)

Reviewer comment (Theory):
Is it possible to put some shrinkage priors on the $W$ matrices in (4) so that unneeded dimensions are pruned away? Can the model distinguish between "true" and "false" shared latents? Is there some notion of parsimony here?

Author response:
The introduction of shrinkage priors is a good idea, and an avenue we are pursuing for follow-up work. For now, we avoid unneeded dimensions through our process of cross-validation. As we indicated in our other responses, we have added an additional section in the appendix devoted to explaining this process (see "Latent dimensionality selection" on page 9 in appendix). Briefly, we assure we have no more dimensions than needed by first assessing the dimensionality per modality, and then increasing the number of shared dimensions one by one while evaluating cross-validated performance. In all tested datasets, using at least 1 shared dimension performed better than 0 shared dimensions, suggesting some cross-modality structure. We increased the number of shared dimensions sequentially until we no longer saw an increase in cross-validated performance, and used that to identify the latent dimensionality. This process, however, is somewhat intensive and a shrinkage prior would be very helpful in speeding up this procedure.

Reviewer comment (Theory):
The assumption of linearity between latents in eqn (4) seems over-simplified, other nonlinearity variants could be further investigated.

Author response:
Equation 4 simply functions to combine the shared and modality specific latents in the first step of the decoder. This is akin to the latent space in probabilistic CCA models (see for example equation 2 in Gunderson et al. 2019) and so it acts to partition the variability between the shared and independent latent subspaces. This linear operation is therefore an essential feature of our multi-modal model. 

- Gundersen, Gregory, et al. "End-to-end training of deep probabilistic CCA on paired biomedical observations."

Reviewer comment (Experiments):
The construction of simulated dataset is confusing, it is not clear to see how the rotation angle of MNIST image is inherently related to spiking distribution. And it is puzzling that how it is tightly linked with real-world datasets.

Author response:
We apologize for this confusion. Please see our response to reviewer YQRh, where we clarify this in detail. If you have any additional questions on this, please let us know.

> [...]  linked with real-world datasets. 

We draw the reviewers attention to the final figure, where it is precisely the case that the latent variables that describe the neural activity are a linear combination of latents that are shared with a high dimensional moving stimulus (closely related in structure to the MNIST image) and independent wing-flapping (like our other 1d GP neural-only latent in our simulated data). By validating our approach in this synthetic setting first, we can be more confident in understanding these uncovered latents in figure 5. 
To re-emphasize this real-world finding, the neural activity is tuned from a two dimensional latent space, one dimension of which tracks the movement of the 2d visual stimulus, and the other dimension corresponds to wing-flapping. The overall Poisson rate for each neuron is thus a linear combination of these two experimental features, in a similar form to how we construct our synthetic data.

Reviewer comment (Experiments):
The paper proposed parameterize the latents in Fourier space to constrain them more distinguishable, and attempt to prune the high-frequency components. While it didn't provide solid ablation studies and compared to the original temporal space on real datasets, and it is problematic to claim no information preserved in high-frequency domains.

Author response:
Thank you for this suggestion. We have included the Fourier approach without pruning the frequencies in the appendix (see figure 6). As you can see, not pruning the frequencies prevents us from identifying the correct smoothly evolving latent structure, and so we did not continue with an unpruned approach in this work.

Importantly, we do not claim that there is no high-frequency information in the latent space (there very likely is) but our pruning regularizes our model in such a way to be able to uncover meaningful latents with clear temporal structure. As the reviewer suggests, our model is thus limited in its ability to identify fast moving temporal structure, and we will be sure to emphasize this limitation more clearly in the final manuscript.

Reviewer comment (Reproducibility):
The authors didn't provide clear explanation or guidance about important hyperparameters of latent dimensions.

Author response:
We apologize for this not being more clear. We mentioned the latent dimension selection process in the appendix but agree our treatment was cursory. We clarify our approach to reviewer Kdv4, so please read this response for more information. We have additionally clearly described our dimensionality-selection method in the appendix, now it can be found under ‘latent dimensionality selection’ section.

Regarding other model hyperparameters, we have a section devoted to the observation variance parameter in the supplement, as well as a thorough description of our selection of length scales for the GP latents and neural network architectures. In general, model performance is not sensitive to specific network architectures and length scale initializations can encourage smoother or less smooth after inference is complete to some extent (also see page 4, 6, 7, 9, and 10 of the appendix for more information).
Please let us know if you need any further clarification on any specific model hyperparameters.

Reviewer comment (Experiments):
No quantitative measurements on real datasets only with a few qualitative samples, and not compared with other baselines.

Author response:
> No quantitative measurements on real datasets 

Please see our response to reviewer bxWt, as we emphasize our quantitative measurements about the contributions of the shared and independent subspaces on the hawkmoth dataset, and propose the addition of additional reconstruction examples of data from each modality as we ablate the shared and independent latents in the appendix. 

> Not compared with other baselines

As we mention to other reviewers, we do compare to two time-series multi-modal models used in neuroscience today in the appendix, both on a synthetic dataset as well as a real neural dataset (fig 1 and 2). However, it is clear we did not highlight this enough in the main body of the manuscript, and so we will be sure to update this accordingly.

Reviewer comment (Presentation):
In Fig 4(d), different behavior states are not clearly separable in these subspaces.

Author response:
That is correct, the MMGPVAE does not separate behavioral conditions particularly cleanly across fly behaviors. The separation of behavioral categories of the latents is an avenue for further scientific exploration, and may require tuning additional model features or adjusting the geometry of the latent prior. However, figure 4 is meant to demonstrate that the MMGPVAE is able to generate good predictive performance on this multi-modal dataset in held-out trials, and that each of the latent subspaces are able to be visualized and analyzed in this setting. (as far as we know, we are the first multi-modal model flexible enough to describe simultaneous 16-dimensional limb-tracking and calcium recordings.)

Reviewer comment (Evaluation):
Blurriness and mismatch of reconstruction outputs and latents in the results.

Author response:
The blurriness of the 3 in our model reconstructions is expected given our model. To be clear, our goal is not to best reconstruct the three, but rather to identify the smoothly evolving angle and scale of the three in the latent space. The reconstructed threes are thus provided only angle and scale information from the latents, and so they otherwise reflect an average of all handwritten threes in the training set (hence the blurry, canonical 3 looking the same at all angles and scales). This is an essential point in our work, and we are happy to clarify further if the reviewer has additional questions or concerns. (note also that the blurry canonical 3 is seen in other GPVAEs of Casale et al, Fortuin et al. and Ramchandran et al).

Reviewer comment (Experiments):
How to decide dimensions for separate and joint latent?

Author response:
We thank the reviewer (and other reviewers) for pointing this out and have more clearly highlighted this procedure in the Appendix, page 9. Please let us know if further clarification is needed.

Reviewer comment (Experiments):
**Fixed time horizon**: The use of an MLP to convert the per-timestep embeddings into per-sequence Fourier coefficients means that you can only consider fixed-length sequences. This seems to be a real limitation, since often neural/behavioral data – especially naturalistic behavior – is not of a fixed length. This could be remedied by using an RNN or neural process in place of the MLP, so this is not catastrophic. However, this should at least be noted as a limitation of the method, and preferably, an RNN or neural process should be substituted for the MLP in one of the examples to concretely demonstrate that this is not a fundamental limitation.

Author response:
We agree with the reviewer that this is a limitation of the model, and we will be sure to highlight this in the main body of the final manuscript. We do, in fact, need to split our datasets into ‘trials’ of fixed length to fit the model (as is briefly mentioned in the appendix sections on pages 8 and 9). Fortunately, this did not practically prevent our model from providing insight in the datasets in 4 and 5.

Reviewer comment (Experiments):
**Missing experiments – baseline comparison**: In Figure 3, there should be a model that uses the data independently to estimate the latent states and reconstruction. It seems unfair to compare multimodal methods to methods that use just one channel. An acceptable baseline would be averaging the predictions of image-only and neuron-only models (co-trained with this loss) so that all models have access to the same data, and the performance improvement can be attributed to the novel structure.

Author response:
We apologize but we are a little unclear as to this point. Can the reviewer clarify further what they mean here? Our ablations in figure 3 do use the datasets independently, and the latent dimensionality is kept the same per-modality (two dimensions for each unimodal evaluation). In other words, each unimodal piece of our model has identical flexibility whether or not the model is trained jointly or separately. Therefore, it is the cross-modality structure that increases the performance (the neural latents are able to use the structure of the image data to better reconstruct rates, and vice-versa).

Reviewer comment (Experiments):
**Missing experiments – image occlusion**: An experiment with image occlusion (half of the image randomly blacked out) is suggested. This would simulate irregularities often present in neural/behavioral data (e.g., failed keypoint detection in some frames) and allow inspection of the model's long-range inference capacity, as opposed to a nearly-supervised reconstruction task.

Author response:
An excellent suggestion. We have added it to the appendix. Our smooth gp-based approach allows for better identification of an uncorrupted image than a standard VAE while retaining the ability to reconstruct neural rates, rendering the MMGPVAE robust to measurement errors.

Reviewer comment (Evaluation):
**Slightly lacking analysis**: The analysis of the inferred latents is somewhat lacking. More incisive analysis is requested to determine what the individual and shared features extract from the data – for example, whether shared latent states indicate “speed” or if this is confined to the individual behavioral latent, whether stimulus type could be decoded from the continuous latent states, and how decoding accuracy differs between the three different $z$ terms. This type of analysis is a key point of training and deploying such models and would help highlight the benefits of the method.

Author response:
Incidentally, since submitting the manuscript, we have follow-up work that begins to answer this question in the context of the hawkmoth dataset. As we mention in the manuscript in our evaluation of the hawkmoth dataset, the shared latent accounts for about 30% of the variability of image reconstructions, and about 5% of the variability of neural data. As we point out, this is expected as the neural data is dominated by wing-flapping. We have since looked at what features of the stimulus and neural activity are encoded in the shared subspace by removing this shared component and plotting the data reconstructions. We find that removing the shared component removes the slowly varying features of neural reconstructions and blurs and distorts the overall image reconstruction, particularly at the times when the stimulus changes direction. If the reviewer feels strongly, we could add these figures to an appendix, but currently we feel it is more appropriate for follow-up work.

Reviewer comment (Novelty):
The discussion of existing work is incredibly limited (basically two citations). There is a plethora of work on computational ethology, neural data analysis, and interpretable methods. This notably weakens the paper by painting an incomplete picture of the field and obscuring why this method is appealing. A much more thorough literature review is expected in the final version.

Author response:
We will change this section title in the discussion, to “competing multi-modal approaches in neuroscience”, and thank the reviewer for pointing out this confusion. While we survey the field significantly more broadly in the introduction, in this discussion section, we wanted to simply point out two existing multi-modal time-series LVMs that have been used in similar settings in neuroscience. We find it important to point out these models as we compare our approach to them in two figures in the appendix. We will be sure to clarify the writing here and make sure we conclude with a more accurate picture of the field.

Reviewer comment (Writing):
The use of “A” and “B” is confusing – A is Behavior, and B is Neural? Better naming is suggested, such as B for Behavioral (and A for neural) or A for behavioral (currently referred to as “other”) and consistently referring to it as Auxiliary or Alternative data, and B as “Brain” data.

Author response:
We apologize for this confusion. We propose flipping the notational indices, where A reflects “activity” and B reflects “behavior”

Reviewer comment (Experiments):
The synthetic data seems weak and unnecessary. Why not use a simulator like NEST or augment one and create "behavioral" data with known connections to this. MNIST images have a different structure because they are a two dimensional image.

Author response:
Assessing performance for GPFA models in neuroscience to synthetic poisson or gaussian spiking data generated from a low-dimensional subspace is nearly ubiquitous for this class of models (e.g. see Yu et al 2009, Lakshmanan et al 2015, Zhao and Park 2017, Dunker and Sahani 2018, Keeley et al 2020, Jensen et al 2021, Gokcen et al 2022) A simulator like NEST is not designed to generate spikes with true low-dimensional generative latent structure so unless we’ve misunderstood the reviewer we would have no way of validating our LVM approach on such synthetic data. Similarly, GPVAE models each use rotating digits as benchmarks for their model performance (see Casale et al 2018, Fortuin et al 2020, Ramchandran et al 2021) generated from a low-dimensional latent space. Because, in each of these cases in these works, true generative low-dimensional GP latents give rise to very different kinds of data (MNIST and Poisson spikes), we consider this precise type of synthetic data vital in validating our inference approach, and necessary to relate our model pieces to the GPFA and GPVAE literature upon which we draw.  By using rotating digits and GPFA-generated spike trains, we are able to demonstrate that we maintain state-of-the-art performance on the relevant datasets for each unimodal model (GPFA and GPVAE). This would be more difficult to do if we generated a completely new type of synthetic dataset. 
MNIST images have a different structure than our experimental validation in figure 4 but quite similar structure to our experimental validation in figure 5, where we assess the model on spiking motor neurons and a smoothly moving 2d visual stimulus.

Reviewer comment (Experiments):
The authors lack comparison to anything but an ablation of their own model, and that too on the above artificial dataset.

Author response:
We apologize if we did not highlight these clearly enough in the main manuscript, but we do, in fact, compare MM-GPVAE to an RNN based model (Hurwitz et al.) as well as an ODE based model (Sani et al.). Both of these existing models are multi-modal models used in neuroscience . We evaluate these existing multi-modal models on both synthetic and real-world datasets in the figures 1 and 2 of the appendix. 

- Sani et al. (2021) "Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification"

- Hurwitz et al. (2021) “Targeted Neural Dynamical Modeling”

Reviewer comment (Experiments):
At least the authors should compare to RNN/transformers/ODE or other sequential models for this data.

Author response:
We apologize if we did not highlight these clearly enough in the main manuscript, but we do, in fact, compare MM-GPVAE to an RNN based model (Hurwitz et al.) as well as an ODE based model (Sani et al.). Both of these existing models are multi-modal models used in neuroscience . We evaluate these existing multi-modal models on both synthetic and real-world datasets in the figures 1 and 2 of the appendix. 

- Sani et al. (2021) "Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification"

- Hurwitz et al. (2021) “Targeted Neural Dynamical Modeling”

Reviewer comment (Theory):
The authors seem to criticize DNNs for transforming the data too much such that the latent variables are "not obviously related" to external variables, and yet they too use DNNs.

Author response:
We are of course not critical of DNNs in general. As the reviewer points out, there are many DNNs used in our work, including in the encoding networks in all models, and in the decoding networks for behavioral data modalities. These networks are crucial for model inference and greatly enhance the flexibility of our approach. However, DNNs are limited in their abilities to provide scientific insight, and this is precisely what motivates the use of a linear decoder for the neural data. By having a linear decoder, our map from latents to rates is akin to GPFA, allowing the MMGPVAE to capitalize on its wide successes (see the GPFA citations below). In addition, it  allows us to perform the analyses in figure 5 of the manuscript, where we can evaluate how individual neurons are independently tuned to shared (stimulus-tracking) and independent (wing-flapping) subspaces. We are pursuing follow-up work which uses these linearly-identified tuning features to better understand the hawkmoth motor system.

Reviewer comment (Theory):
There is no theory or study of identifiability of the dynamics or conditions/noise under which dynamics are recovered.

Author response:
This is correct. We will be sure to highlight this limitation in the final version of the manuscript, and we thank the reviewer for pointing this out.

Reviewer comment (Novelty):
Methodological contributions may be limited. Many papers have done similar things, such as extending the latent space into frequency domain [1,3] and using the Poisson model for spiking data [1,2]. The method proposed in this paper is combination of existing methods, not fundamentally different.

Author response:
We agree with the reviewer that the core contribution of this work represents a combination of GPFA, and GP-VAEs along linear partitioning of a latent space for datasets into shared and independent components of two modalities, and so is in a sense a combination of existing methods. However, we believe that we are the first time-series latent variable model developed that uses this partitioning to separate both shared and independent latent trajectories in multi-modal settings in neuroscience, and we do so in a way that is more flexible and scalable compared to competing multi-modal methods (see appendix). We feel that this is therefore a novel approach with wide applicability across a range of multi-modal neuroscientific experimental set-ups that we demonstrate in the manuscript. We also would like to highlight our novel amortization method in the Fourier domain for a deep GP model, as well as our derivation of a model specific ELBO in the appendix that precludes the need to approximate two of the terms in the objective.

Reviewer comment (Experiments):
The experimental results can be more comprehensive: (1) More baselines, especially similar methods such as [1,4] need to be compared.

Author response:
[1] and [4] are not multi-modal models, and so we do not believe that they are fundamentally similar to our approach. We do, however, recognize that [1] and [4] each relate to one piece of our model, and we agree with the reviewer that comparisons of our multi-modal model to these unimodal variants will help us understand model performance. However, our comparisons of GPFA and GP-VAE in figures 1 and 3 are identical to the relevant variant of the method outlined in [1] and closely related to that of [4]. ELBO from [4] takes a similar general form to the GP-VAE ELBO we adapted which was presented in Casale et al 2018, barring a different observation likelihood and differently factorized variational distribution. Therefore, we consider our comparison to Casale et al 2018 to be a very close approximation to the comparison in [4]. We thank the reviewer for pointing us to this work and will add the citation of [4] and discussion as to this point in our revised version.  
> More baselines  
We benchmark against two prominent multi-modal models used for neuroscience in the appendix - one of these uses ODEs to characterize the latents, and the other uses an RNN. If the reviewer finds these benchmarks insufficient, please let us know.

Reviewer comment (Experiments):
An ablation study between the models with Frequency prior and with other priors (e.g., a temporal form) should be added.

Author response:
> Ablation study with fourier frequencies  
Thank you for this suggestion, we now address this in the global response.

Reviewer comment (Evaluation):
More discussions on the claimed interpretability should be provided.

Author response:
> Interpretability  
We are very specific about interpretability in this work. In the synthetic examples, the "interpretability” of our model refers to the fact that we can recover latents that identify two features of the data that are meaningful to us, in this case, the digit scaling factor and angle. In the case of the real data examples, interpretability means the identified latent in the data corresponds to something scientifically meaningful - for example, the position of a visual stimulus, the beating pattern of a hawkmoth wing (Fig 5) , the behavioral state of a fly (Fig 4) , or the reaching position of a monkey (Fig 2 in appendix). We will edit our wording to be more specific about our use of the term ‘interpretability'

Reviewer comment (Experiments):
How sensitive is the performance with respect to the setting of several hyperparameters, including $F$ and the dimension of shared and independent latent subspace, etc.?

Author response:
For the real neural dataset examples in the main paper, the dimensionality was chosen for the neural dataset using PCA (selected for capturing greater than 97% of the variance) and the dimensionality was chosen for the behavioral data by increasing the number of latent dimensions until cross-validated performance did not show significant improvement. To determine the shared dimensionality, we similarly systematically increased the number of dimensions shared between datasets until predictive performance no longer improved. However, the results are not particularly sensitive to latent dimensionality within a few dimensions of either modality.

Reviewer comment (Reproducibility):
How optimization was done with FFT/iFFT involved in MM-GPVAE?

Author response:
The encoder networks first identified latent variables per time-point, and was then passed through a FFT and reduced in dimensionality across time before being used to sample for the ELBO. Thus, inference was done in the Fourier domain. When decoded, the latents were mapped back out of the Fourier domain using an inverse transform. Because each of these transformations are linear operations, they are simply parts of the encoder and decoder networks. Network schematics for the exact mappings for each of the analyses can be found in figures 9-12 of the appendix. Please let us know if the reviewer needs any further clarification.