Multi-modal Gaussian Process Variational Autoencoders for Neural and Behavioral Data

Methods:

1. 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."

The construction of simulated dataset is confusing,

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.

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

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

Evaluations:

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.

2. 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.)

3. 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).

Writing:

We apologize for any confusion in our work. If the reviewer can be more specific about which important contents they feel are insufficiently addressed beyond what is discussed above, please let us know.

Questions:

1. A nonlinear operation is used as the next step in each of our decoders (an exponential nonlinearity to generate Poisson rates for the GPFA description of neural data and a deep neural network to describe the behavior/stimulus), so the ultimate decoding mapping to data is a nonlinear operation. If this first step were nonlinear, it would be highly unusual, a more difficult to train and we would lose the ability to interpret our shared/independent latent partitioning. (see also Gunderson et al 2019 for additional context).

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

We hope that the revisions and responses in our updated paper have sufficiently addressed all of the concerns and questions raised by the reviewer. If this is the case, we hope the reviewer could consider updating their score accordingly. If there are still any questions or concerns, we would be happy to clarify them and engage in further discussion.

We thank the reviewer for their positive appraisal of our work. Below we address each of the questions and points raised by the reviewer.

Weaknesses:

The synthetic experiments tend to focus on smooth, continuous variation [...]

One of the main assumptions that we make with our model is that we assume these multimodal measurements have an underlying low-d smooth structure. Therefore, we agree with the reviewer that MM-GPVAE might not be suitable for multi-modal data that does not smoothly evolve in time (indeed, this is a limitation of GPVAE and GPFA models as well). We talk about limitations of our prior assumptions in the conclusion section of our manuscript to compare other approaches and limitations of other prior approaches. However, comparing multi-modal neuroscience models with different prior assumptions remains an important avenue for exploration.

Questions:

As noted above, the GP method would seem to need data to be broken into snippets of reasonable size. [...]

That is correct. For our evaluations on real datasets, we need to artificially split our data into trials of equal length, which we describe in more detail in the appendix, page 8 and 9.

We find that using the Fourier parameterization aids with this process, and the model is able to handle longer-time series trials if processed in the fourier domain (likely due to the lack of the need to invert the prior GP covariance matrix, and the overall reduction in the number of variational parameters)

Is it possible to put some shrinkage priors on the matrices in (4) [...]

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.

We again thank the reviewer for their thoughtful comments and suggestions, and for their time reviewing the paper. We hope that the revisions and responses in our updated paper have sufficiently addressed all of the concerns and questions raised by the reviewer. If there are still any questions or concerns, we would be happy to clarify them and engage in further discussion.

We thank the reviewer for highlighting the relevance of this work to the neuroscience community and praising the approach and experiments. Below we address each of the questions and points raised by the reviewer.

Weaknesses

a.

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.

b.

This is a good point, and is discussed in some detail in the appendix (see page 10 in appendix). To highlight the important point – our learned $\sigma^2$ parameter did not vary much from its initialized value, and so our choice of the initialization of $\sigma^2$ reflects exactly the kind of trade-off the reviewer describes. We therefore initialized the $\sigma^2$ parameter to a value proportional dimensionality of the observed behavioral modality, and this assured a balancing across modalities for any of the real-world datasets we tested.

c.

Firstly [...]

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

Secondly [...]

Thank you for this suggestion. We have added a figure, fig 7, to the appendix showing exactly this. We see that increasing the number of neurons only improves the identification of shared and neural latents, and, expectedly, does not affect the ability of the model to identify the image-only latent structure. We also see improvement in reconstructions across both modalities as the number of neurons increases

Finally [...]

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.

d.

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.

Minor weaknesses

Thank you for the improvements to our text. All the changes will be implemented in the final manuscript. We address some of the more important minor points below:

3.

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.

8.

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

We again thank the reviewer for their thoughtful comments and suggestions, and for their time reviewing the paper. We hope that the revisions and responses in our updated paper have sufficiently addressed all of the concerns and questions raised by the reviewer. If this is the case, we hope the reviewer could consider updating their score accordingly. If there are still any questions or concerns, we would be happy to clarify them and engage in further discussion.

To the authors,

Thank you for your response and clarifications, and for swiftly updating the paper.

In response to "C. Firstly": As far as I understand, in Figure 3c, the baselines are only provided with one channel, whereas your proposed method is provided with two channels. Is the performance increase coming from simply getting to see both channels, whereas the baselines only see one channel? Are the authors saying that "image-only GP-VAE" is in some way accessing the neural data?

E.g., You're making a salad. It is unfair to compare a salad made by someone who just had access to tomatoes, to someone who just had lettuce, to someone who gets to use both tomatoes and lettuce. I want you to compare four salads (models): just tomato (neuron only GPFA), just lettuce (image only GP-VAE), lettuce and tomato side by side but no mixing (the baseline I'm requesting, e.g. averaging predictions/reconstructions of neuron only GPFA and image only GP-VAE), and a salad with lettuce and tomato all mixed together (MM-GPVAE).*

Please correct me if my understanding of the experiment presented is incorrect.

In response to "D.": I do feel reasonably strongly on this, actually. If you have such results, then I think at least a taster of them should be included here. If you are angling this paper as an interpretable neuroscience method, then I think some interpretations/neuroscience should be included. I don't think you should include so much analysis that you cannot publish another paper somewhere else, but I think this paper might've been "sliced a bit thin" in that context. I would very much like to see some simple analysis in this vein included in the final version.

Overall, I think this paper is of publication quality. Subject to the inclusion of a more thorough discussion of related techniques, and the comparisons suggested by Kdv4, i vote for the inclusion of this paper.

Thank you very much, and good work,

bxWt

*This is the most ridiculous thing I've written in a review. I hope you enjoyed it and/or found it useful...!

We thank the reviewer for finding MM-GPVAE useful for understanding the components of behavioral and neural data. Below we address each of the questions and points raised by the reviewer.

Weaknesses:

[...] simulator like NEST [...]

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

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.

The authors lack comparison [...] authors should compare to RNN/transformers/ODE [...]

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”

The authors seem to criticize DNNs [...]

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.

There is no theory [...]

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.

We thank the reviewer for the positive appraisal of the manuscript. We address some of the concerns and misunderstandings below.

The method proposed in this paper is combination of existing methods, not fundamentally different.

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.

More baselines, similar methods such as [1,4]

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

Ablation study with fourier frequencies

Thank you for this suggestion, we now address this in the global 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'

Questions

1

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.

2

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.

We again thank the reviewer for their thoughtful comments and suggestions, and for their time reviewing the paper. If there are still any questions or concerns, we would be happy to clarify them and engage in further discussion.

More baselines, especially similar methods such as [1,4] need to be compared

To elaborate further, the model in [1] is a unimodal model developed for neural data that explores trial-by-trial deviations in trail-based neural data with exact repeated stimulus conditions, so alone the model introduced in [1] is not directly applicable to our setting. However, we recognize that a “signal only” version of [1] is an equivalent model to the poisson GPFA model first introduced in Zhao and Park 2017 learned with a Fourier representation, similar to our approach in an exclusively linear setting without deep amortization. Therefore, this version of [1] is an identical model to the “GPFA only” column in figure 3 in our paper. So, while we did not make this connection explicit in our paper, we do have a direct comparison here to the relevant modified version of [1]. (the only difference is our model uses amortized inference, instead of a black-box VI approach used in [1], which improves out-of-sample performance and allows us to adapt this representation to the deep setting). We will make sure this comparison to this version of a unimodal poisson-GPFA is clearer upon revision.

Similarly, [4] is a unimodal GP-VAE model, and it uses Bernoulli likelihood and different factorized variational distribution but otherwise is closely related to the ELBO of Casale et al. Our primary contribution is comparing this same ELBO with and without the Fourier domain representation of the prior and latents (The prior and latents of [4] are identical to that of Casale et al.)

I would like to thank the authors for the effort to address some of the questions. I am wondering for a given new dataset, how much tuning time would be needed if latent dimension, frequency pruning, and other hyperparameters were involved with cross validation, etc.

Regarding the baseline comparison, I still felt that the model itself has been proposed or very similar to the existing ones. If the authors would like to emphasize their latent space partitioning idea, it may be important to compare with multi-modal feature fusion and other similar multi-modal embedding methods.

I am little bit confused with the authors' statement: "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." If the model has to be trained in the Fourier domain, is this what have to be done? It would be great if the authors can provide more details why the amortization methods in the Fourier domain is 'novel'.