Structured Recognition for Generative Models with Explaining Away
Abstract
A key goal of unsupervised learning is to go beyond density estimation and sample generation to reveal the structure inherent within observed data.
Such structure can be expressed in the pattern of interactions between explanatory latent variables captured through a probabilistic graphical model.
Although the learning of structured graphical models has a long history, much recent work in unsupervised modelling has instead emphasised flexible deep-network-based generation, either transforming independent latent generators to model complex data or assuming that distinct observed variables are derived from different latent nodes.
Here, we extend amortised variational inference to incorporate structured factors over multiple variables, able to capture the observation-induced posterior dependence between latents that results from “explaining away” and thus allow complex observations to depend on multiple nodes of a structured graph.
We show that appropriately parametrised factors can be combined efficiently with variational message passing in rich graphical structures.
We instantiate the framework in nonlinear Gaussian Process Factor
Analysis, evaluating the structured recognition framework using synthetic data from known generative processes.
We fit the GPFA model to high-dimensional neural spike data from the hippocampus of freely moving rodents, where the model successfully identifies latent signals that correlate with behavioural covariates.
Introduction
A central challenge of unsupervised learning is to identify and model patterns of statistical dependence in high-dimensional data.
One approach to this challenge exploits latent-variable generative models, which capture the statistical structure of data through the modelled influence of the latent variables on observations, and through interactions between the latents themselves.
Recent developments in deep learning have enabled deep generative models (DGM), with remarkable success in density estimation and high-fidelity image or text generation [1–5].
However, much of the DGM development has emphasised the expressiveness and accuracy of the latent-to-observation generative process, with the latents themselves often assumed independent a priori .
Variational learning requires inferring or approximating the posterior distribution of the latents conditioned on observations [6].
In a variational autoencoder [VAE; 1, 2], inference for a DGM is amortised by training a recognition network to return parameters of the variational posterior.
The structured variational autoencoder [SVAE; 7] connects the DGM framework with richer models that include structured prior dependence amongst the latent variables.
In an SVAE the structured latent prior is combined with a DGM link between each latent and a corresponding observation.
Variational inference is amortised using recognition networks that return parameters of a factor associated with each link, which are combined with the structured prior to obtain the full posterior.
However, in 36th Conference on Neural Information Processing Systems (NeurIPS 2022).
many generative models of interest, observations depend on more than one latent variable.
Such interactions induce joint potentials in the likelihood, coupling latents in the posterior even if they are independent a priori , a phenomenon sometimes called “explaining away” [8].
To capture such observation-induced dependence in the DGM setting we develop the structured recognition VAE (SRVAE), where recognition potentials that incorporate joint factors induced by the observations are learnt in a (structured) variational autoencoding framework.
We instantiate the SRVAE framework for a range of graphical structures, including non-linear latent
Gaussian process (GP) models.
In the latter case, we bring together the GP-prior VAE [9, 10], and (sparse) Gaussian process factor analysis (GPFA)
[11, 12] models, yielding a novel autoencoding model (SR-nlGPFA).
Experiments show that SR-nlGPFA outperforms alternatives that lack structured recognition.
We apply SR-nlGPFA to data collected from a population of neurons in the hippocampal complex, and show that the unsupervised approach automatically captures dimensions underlying neural firing that reflect relevant behavioural correlates.
Background 2.1 Variational Inference and Structured Variational Autoencoders
Consider a general generative process with latent variables z and observations y ,
[Equation]
where θ and γ are the parameters of the prior and conditional likelihood distributions, respectively (and we write Θ = θ ∪ γ ).
We assume deterministic θ for simplicity, but distributions over θ (with conjugate hyperpriors) can be incorporated by variational Bayes (VB) [13].
The posterior distribution, p ( z | y, Θ) ∝ p ( z | θ ) p
( y | z, γ ) , is often analytically intractable.
In such cases, a common approach is to seek an approximation q ( z ) constrained to a tractable class
Q by variational inference (VI); that is, by minimising the KL-divergence with respect to the true posterior distribution or, equivalently, maximising a variational free energy [6] F VI = (cid:10) log
[Equation]
In the variational autoencoder (VAE) architecture, the parametric optimisation of q ( z ) implicit in (1) is amortised by a recognition network that takes as input the observed values y and returns parameters of q .
The parameters of this recognition network are then trained jointly with the generative model by stochastic optimisation of the free energy [1, 2].
In both standard and amortised VI, the approximate distribution q is often constrained to factor over the latent variables, a so-called mean-field constraint.
While tractable, such mean-field approximations may be too restrictive to capture the complexity of the true posterior [14, 15].
Many approaches have been proposed for improving the expressiveness of variational approximation [16–19].
Here we consider the structured VAE [SVAE; 7].
Unlike standard VAEs, the SVAE assumes the generative prior distribution to be specified by a structured probabilistic graphical model (PGM), p ( z | θ ) ∝ (cid:81) c ∈ C ψ c ( z c ) , where { ψ c } c ∈ C correspond to C clique potentials.
In addition, the amortised inference network outputs recognition factors, r ( z | y, ϕ ) , that approximate the generative likelihood function rather than the full variational posterior.
These recognition factors are combined with the structured prior distribution to obtain the amortised variational posterior,
[Equation]
where ϕ ∗ minimises the averaged KL over the data observations Y .
Johnson et al.
[7] considered recognition factors that are local (singleton) evidence potentials, chosen to be conjugate to p ( z | θ ) .
Even in this case, the form of (2) allows the dependency structure established by the prior to be carried over to the variational posterior distribution.
Further details of SVAE appear in Appendix A.
2.2 Gaussian Process Factor Analysis
Gaussian Process Factor Analysis (GPFA) is a model used in neural data analysis to infer the dynamic latent structure underlying high-dimensional population spike trains [11, 12].
Standard GPFA assumes the following generative model:
[Equation]
where m kθ ( · ) and κ kθ ( · , · ) parametrise the GP prior, c jk and d n define an affine mapping from the latent space to the observation space (or a transform of the observations for general likelihoods), and g ( · ) is a smooth scalar link function appropriate for the observation distribution.
We take m kθ ( · ) =
unless stated otherwise.
Correlations between observed neurons are captured by dependence on the common set of latents, while temporal correlations in the high-dimensional observations are modelled by the latent temporal correlations of the GPs.
The PGM describing the GPFA generative model (with sparse variational approximation, see below) is shown in Fig 1a.
Although maximum-likelihood parameters for GPFA can be found using (variational) expectationmaximisation [20, 11] exact GP inference scales cubically in the number of observation times [21].
Sparse variational GP (svGP) inference based on auxilliary inducing points [22] reduces the time complexity of learning in GPFA [23], and also facilitates efficient extensions to (pointwise) nonconjugate likelihoods [12].
For k = 1 , . . .
, K , we introduce inducing values u k representing the evaluations of the latent process f k at M k inducing locations, z k .
For simplicity, we assume M 1 = · · · = M K = M unless stated otherwise.
The GPFA generative model (3) can be augmented to include these auxiliary variables:
[Equation]
where F k ( · ) = κ kθ ( · , z k )( K k z k z k ) − 1 is the linear operator that maps u k to f k ( · ) ; κ kθ ( · , z k ) is a vector-valued function with κ kθ ( x ; z k ) =
[ κ kθ ( x, z k 1 ) , κ kθ ( x, z k 2 ) , . . .
, κ kθ ( x, z kM k )] ; and K k z k z k is the covariance matrix obtained by evaluation of the kernel function κ kθ ( · ) at the inducing locations z k .
Introducing a variational distribution over the inducing points and the latent functions, and utilising the generative model (Fig 1a), the free energy of the sparse variational GPFA (svGPFA) can be expressed as following:
[Equation]
where y t =
[ y n ( t )] ∈ R N , f t =
[ f k ( t )] ∈ R K , F = [ f t ] T ∈ R T × K , C
= [ c nk ] ∈ R N × K , d = [ d n ] ∈ R n , U =
[ u 1 , . . .
, u K ] T ∈ R K × M , and Z =
[ z 1 , . .
, z K ] T ∈ R K × M .
The svGP approach constrains q ( F , U ; ϕ ) to the form q ( U ) p ( F | U , θ ) with Gaussian q ( U ) [22].
In the multi-GP case, the distribution is typically also taken to factorise over processes [23], q ( F , U ) =
(cid:81)
Kk =1 p ( f k | u k ) q ( u k ) , with f k =
[ f k (1) , . . .
, f k ( T )]
∈ R t , q ( u k ) =
N ( m k , S k ) .
Under these constraints on q , the svGPFA free energy simplifies to
[Equation]
where q ( h t ) =
(cid:82) d U p ( h t | U ) q
( U ) is itself a GP with mean and kernel functions given by
[Equation]
Hence the computational cost for sparse variational GP inference reduces from O ( T 3 ) down to
O ( M 3 + TM 2 ) .
Structured Recognition and Explaining Away 3.1 Structured Recognition in VAEs
A general non-linear likelihood may induce additional latent dependency structure in the posterior beyond that of the prior (a simple illustrative example appears in Appendix C).
Although the SVAE u k f k ( x n ) h n y n
F k u k
Cf 1: K + d g ( · ) k
=1 ,...,K n =1 ,...,N (a)
u k f k ( x n ) h n y n
F k u k
Cf 1: K + d r ( u | y ,ϕ ) k
=1 ,...,K n =1 ,...,N (b) u k f k ( x n ) h n y n
F k u k r ( h | y ,ϕ ) p ( u )
Cf 1: K + d r
h | y ,ϕ ) k
=1 ,...,K n =1 ,...,N (c) approach developed by Johnson et al.
[7] composed amortised inference with a structured prior distribution, the recognition models these authors discussed all contributed amortised potentials that factored over latent variables.
Such an approach cannot accurately model posteriors with dependence structure that differs from that of the prior PGM.
Here we adopt a recognition network that outputs structured factor potentials over the latents, providing amortised estimates of the additional dependencies induced by “explaining away”.
We denote the resulting model the Structured Recognition VAE (SRVAE).
The SRVAE variational approximation takes the form
[Equation]
[Equation]
where C r is the set of recognition factors and the ξ c s the associated factor potentials.
We assume that the ξ c s are chosen to be conjugate to the prior factors unless otherwise stated.
Hence the analytical (approximate, if necessary) form of q SRVAE can be computed with (variational) message passing.
In the most general form, we could let ξ ( z | y ; ϕ ) be a single joint factor potential over the complete set of latent variables.
The free energy objective takes the same form as the standard autoencoding free energy objective [1, 2], but with the structured amortised variational approximation, q SRVAE , as the variational distribution.
[Equation]
The training of the model follows standard VAE-style stochastic optimisation to update the parameters of the recognition and generative networks, as well as (optionally) the parameters of the prior PGM.
The following proposition supports the use of SRVAE framework (the proof appears in Appendix E).
Proposition 3.1.
The SRVAE objective function provides a tighter lower bound to the free energy than the SVAE objective function.
[Equation]
Below we describe two instantiations of the SRVAE framework, the first based on a latent Gaussian mixture model, and the second on a latent GPFA model, both with DGM outputs.
Structured amortised inference facilitates scalable inference of the posterior latent distribution with full covariance structure, allowing more accurate learning than with factored recognition approaches.
However, the SRVAE framework is more general-purposed, and can be combined with many different latent variable models.
Further examples appear in Appendix E.
3.2 Structured Recognition with Latent Gaussian Mixture Model
We instantiate SRVAE with a classical latent variable model, the latent Gaussian mixture model (GMM) [24].
We assume the following generative process (Figure 5a):
[Equation]
This model was studied by Johnson et al.
[7], using fully factorised recognition potentials, followed by variational message passing[VMP; 25] to obtain the GMM variational posterior (details can be found in Appendix A).
However, the coupled conditional model, parametrised by a neural network, induces posterior correlations between the latent variables that do not exist a priori .
Hence in order to implement the structured amortised inference, we assume a full covariance Gaussian recognition potential on h , and the SRVAE-GMM variational approximation is computed with VMP given the recognition and prior factors.
3.3 Structured Recognition Variational Autoencoding Nonlinear GPFA
The GPFA model (3) incorporates a generalised-linear likelihood, with the link function g ( · ) acting separately on each affine embedding value h n .
We are now in a position to use the SRVAE framework to extend GPFA to include a DGM likelihood, greatly increasing the expressiveness of the generative model.
Consider observations { ( x t , y t ) }
Tt =1 , defining f t ∈ R K and h t ∈ R N to be the corresponding vectors of latent process values and embeddings at x t .
[In neural applications, inputs x t are usually taken to be the timestamp t ; see (3).
Here we consider the more general case of arbitrary inputs.]
In the generative model, we retain the affine mapping from f t to h t , replacing the link function by a non-linear multivariate DGM: y t ∼ p ( y | g ( h t , γ )) .
Although the generative affine mapping could, in principle, be subsumed within the general deep network, the embeddings h t will play a valuable role in parametrising structured recognition.
For amortised svGP inference, our goal is to define a variational distribution q ( U | Y , ϕ ) .
Ashman et al.
[10] have considered a similar latent GP model called the SGP-VAE.
There, the variational distribution over the latents was found by combining the prior GP distribution with an amortised approximation of the likelihood function factored over latent processes [22]:
[Equation]
The graphical model for SGP-VAE inference corresponds to that of Figure 1b, with C
=
I and d = 0 .
While this approach captures induced correlation amongst the inducing points, it fails to capture correlations between latent processes that arise in the posterior through “explaining away", potentially leading to sub-optimal inference and learning.
Our solution is to recast the output of the amortised inference to the likelihood of h rather than U , and propose the following structured variational distribution.
[Equation]
where r ( h t | y t ) ∝ N ( h t | m ( y t ; ϕ ) , Ψ ( y t ; ϕ )) =
N (cid:0) µ ht , Ψ ht (cid:1) , with Ψ ht diagonal, and H =
[ h ( x 1 ) , . . .
, h ( x T )] .
Given the linear-Gaussian relationship between U and h (Fig 1a), the integral in (11) can be computed in closed form.
Furthermore, even though the recognition potential on h , r ( h | y t ) , is assumed to be fully factorised over the dimensions of each h t , the variational distribution on U includes coupling between latent processes induced by the affine mixing coefficients.
Hence q ( U ) captures the correlations both between the latent processes and through time (through combination with the GP prior under the structured autoencoding formulation).
Specifically, given the general expression of the GPFA generative model and the sparse approximation with inducing points, the linear Gaussian relationship between h and U at any x is given by
[Equation]
where
[Equation]
 with K x = diag [ κ kθ ( x, x )] .
Combining this result with the fully factorised recognition potential on h at x t (11)
[Equation]
we obtain the structured variational distribution on U :
[Equation]
This variational posterior on U leads to a corresponding posterior on h at each x :
[Equation]
and we write q ( h ( x t ))
=
N (cid:0) m h t , S h t (cid:1) .
Then the complete (reparametrised) Monte Carlo estimate of the free energy objective given a mini-batch of data, { ( x b , y b ) } Bb =1 , takes the following expression.
[Equation]
where ϵ s ∼ N ( 0 , I ) , θ is the set of kernel parameters, L hb is the lower-triangular Cholesky component of S hb
such that S hb =
L hb ( L hb ) T .
We note that we have retained the scalability of amortised inference by choosing an amortised diagonal-Gaussian potential on h .
However, the linear-Gaussian relationship between the inducing points U and h ( x ) (12) leads to a full-covariance variational Gaussian approximation for U , allowing amortised inference to capture the observation-induced posterior correlations between latent processes (commonly referred to as the “explaining away” effect).
Thus, the retention of the affine GPFA mapping introduces a key extenstion to SGP-VAE.
We will refer to the new model as structured recognition non-linear GPFA (SR-nlGPFA).
The generative model corresponds to that of svGPFA (Figure 1a) but with the nonlinearity g ( · ) generalised to a flexible form modelled by a neural network.
The complete inference procedure for SR-nlGPFA is graphically illustrated in Figure 1c.
The implicit computation of the variational distribution over U given the recognition potentials on h also makes it possible to carry out svGP inference with changed inducing locations (13).
This is particularly useful in situations that require inference over test datasets of different durations to those seen in training.
In particular, it may be possible to learn an amortised inference network using short sub-sequences drawn from a longer dataset, and then infer a posterior over latent GPs for the complete data sequence efficiently by optimising the placement of inducing points along its full length.
See further details on the free-form svGP inference step in Appendix F.2.
Related work
Structured Deep Generative Models.
A number of studies have considered latent graphical structure within the DGM framework [26, 27, 7, 28–30].
Our work builds on and generalises the prominent
SVAE proposal of Johnson et al.
Both SVAE and SRVAE combine the structure of a prior
PGM with the flexibility of neural network-based recognition.
The difference lies in the form of the recognition potentials.
These were taken to extend over single latent variables in the earlier study.
Here, we consider joint potentials that capture dependence induced by the generative likelihood.
In a closely related study, Lin et al.
[28] proposed a structured inference network, which approximates the variational distribution as the combination of the recognition potential and a separate structured latent distribution independent of the prior, the independent structured latent distribution can have non-conjugate factors to improve the expressiveness of the model.
svGPFA and Extensions.
Since its introduction as a model to identify linear low-dimensional struction in neural population data [11], GPFA has been extended to incorporate non-linear link
[Figure]
functions and non-conjugate (notably Poisson count or point-process) noise models, and combined with sparse variational inference for efficiency [23, 31, 12, 32].
Related work has considered GP latents in the context of DGMs, particularly using variational autoencoding [9, 33, 10].
In a sense,
SR-nlGPFA combines both approaches.
It incorporates multiple latent GPs (as in GPFA and some
DGM GP models) with an affine mapping that feeds into a nonlinear DGM.
However, this structure means that the affine map contributes little to the generative process.
Instead it provides a target for amortised inference, which combines with sparse variational GP inference to yield a full structured posterior on inducing point values.
As such, the affine map may is better seen as an element of the structured recognition model.
Results
We evaluated SRVAE methods for both latent mixtures, and nlGPFA, comparing to relevant baselines on synthetic and real datasets.
Empirical results on the instantiations of the SRVAE framework with other latent variable models can be found in Appendix F 1 .
5.1 Experiments with SRVAE-GMM
To assess the empirical performance of structured amortisation in a standard latent-variable model, we evaluated SRVAE-GMM on the classic pinwheel dataset (Figure 2a)
We compare SRVAE-GMM with the structured inference network [SIN; 28], which is a more flexible instantiation of the SVAE framework.
From Figure 2, we observe that SRVAE-GMM outperforms SIN-GMM in terms of the training variational free energy, reconstruction mean-squared error, and generation fidelity.
The models are identical apart from the recognition stage, providing support for the idea that by capturing the “explaining away” posterior factors, structured recognition leads to better training and stronger generation.
[Table]
5.2 Experiments with SR-nlGPFA 5.2.1 Synthetic and EEG Dataset
We compared SR-nlGPFA to methods that do not capture the inter-latent posterior correlation, using synthetic data and a small-scale EEG dataset used previously [34, 10].
Unless stated otherwise, we employed an exponentiated quadratic kernel, κ ( x, x ′ ) =
λ exp( − || x − x ′
||
2 τ) , where λ and τ are the marginal variance and length-scale parameters, respectively.
Our main baseline model is SGP-VAE [10] (described in Sections 3.3 and 4).
We also compared to a “vanilla” VAE with fully factorised Gaussian variational distribution [1, 2] [Table] and SVAE with a linear dynamical system (SVAE-LDS) latent prior [7].
All models were implemented with the same recognition and generative network architectures (see Appendix G for implementation details).
Synthetic Dataset We generated data from the GPFA generative model (Eq. 3), with g ( · ) =
σ (Φ( · )) , where Φ( · ) represents the functional mapping through a fixed, randomly-initialised 2-layer MLP with ReLU hidden non-linearity, and σ ( · ) is the sigmoid function.
EEG Dataset
We follow the experimental procedure described by Requeima et al.
[34], and consider an EEG measurement dataset spanning 1 second at 256 Hz sample frequency, taken during image viewing [35].
Each datapoint consists of the voltage readings from 7 electrodes positioned on the participant’s scalp.
Here we report results when all data are observed, contrary to the settings in [34, 10] (we include results and further discussion on partial observability in Appendix F.2).
SR-nlGPFA achieved lower standardised mean squared error (SMSE) and lower test negative loglikelihood (NLL) than the alternative methods on both data sets (Table 1), providing evidence for the benefits of structured recognition.
The gains come despite the computational complexity of
SR-nlGPFA being of the same order as SGP-VAE (see appendix F.2).
5.2.2 Population Neuronal Firing Data
The firing of place cells in Hippocampal area CA1 and grid cells in the medial Entorhinal Cortex (mEC) is known to be modulated by the animal’s location [36, 37], speed and direction of locomotion [38–40], with the mapping between these behavioural covariates and neural activity expressed in non-linear mixed tuning curves. The behavioural signals are continuous and often mutually dependent, and so create temporal and spatial structure in the time series of population activity.
We asked whether SR-nlGPFA and related methods would be able to identify and extract this structure without supervision; that is, without direct access to the behavioural covariates.
We used single-cell spiking data from neurons in the hippocampal CA1 and mEC regions of rats recorded during exploration of a Z-shaped track, as reported by Ólafsdóttir et al.
The data comprised 28 experimental sessions, each spanning 10 minutes.
Example neural firing patterns of the population are shown in Figure 11 of the Appendix.
For SR-nlGPFA and SGP-VAE, we adopted the GPFA generative model Eq.
3 with DGM non-linearity and Poisson observation likelihood.
[Equation]
where x are discrete times, N is the number of recorded neurons, g ( · ) represents the neural-network mapping from the GPFA features, h , to the rate of Poisson-distributed firing counts for each neuron over contiguous 100 ms bins.
SR-nlGPFA obtained a lower SMSE of prediction than SGP-VAE and other methods tested (Table 1). Computational constraints meant that the SR-nlGPFA model was trained using short batches of data and 64 inducing points per batch.
Latent trajectory estimates derived from such batches will not necessarily be continuous at the boundaries between them. Thus, once the model was fit, we performed svGP inference over complete sessions with increased numbers of inducing points (see
Section 3.3 and Appendix F.3.2).
To relate these recovered latent time-series to behavioural covariates, we performed two-dimensional Canonical Correlation Analysis [CCA; 42].
We present results for one session here, emphasising qualitative effects. Further qualitative and quantitative results across sessions can be found in Appendix F.3.2.
Figure 3a shows a heatmap of the correlation coefficients between the canonical correlates of the posterior means ( CCX { 1 , 2 } ) and the individual behavioural covariates (distance from one end of the track, speed, direction of travel, head direction, and ‘unfolded’ position along a full lap of the track).
Many correlations are high, indicating the low-dimensional manifold parametrised by the conjunctive set of behavioural correlates can be accurately captured by the posterior latent variables learned with SR-nlGPFA solely from neural spikes.
Please refer to Appendix Figure 13 for numerical values of the correlations (and for other sessions).
To see whether the learned latent structure contained decodable information about behavioural covariates, we extracted the direction-modulated neurons predicted by the trained model (details in Appendix F.3.2).
Figure 3b compares the direction modulation of the model-predicted neurons against that of the rest of the neurons, where direction modulation is defined as the correlation between raw single-cell spike counts and the direction values.
Neurons predicted by the model to be direction modulated exhibited significantly greater direction modulation than the other neurons (p-value= 5 .
56 × 10 − 4 ) [43].
Figure 3c shows the firing profile for travel in each direction for the neuron with the strongest predicted direction modulation. The pattern of firing fields exhibits clear directional dependence.
Similar comparisons hold for spatial and speed modulation, indicating that SR-nlGPFA is able, in a purely unsupervised fashion, to learn a latent space that contains linearly decodable information about the behavioural covariates associated with individual neuronal firing, even for neurons whose activity exhibits conjunctive coding.
Figure 3d shows the two canonical correlates obtained from the latent trajectories as a function of the animal’s spatial location.
These reflect both location and direction of movement (colours), indicating the latent dimensions learned by SR-nlGPFA disentangle direction from spatial location in a simple linear projection, despite the conjunctive coding mechanism exhibited in the CA1 neurons.
A similar plot derived from SGP-VAE (Figure 14b), shows less clear disentanglement.
Figure 3e shows the temporal evolution of the latent CCs for SR-nlGPFA and SGP-VAE, as well as the spatial location of the rat.
Both CCX 1 and CCX 2 of SR-nlGPFA consistently track the dynamics of the distance (with perfect phase alignment and half-cycle offset, respectively), whereas neither CCX 1 or CCX 2 of SGP-VAE could accurately reflect the distance.
Moreover, given the additional svGP inference step over the entire trajectory (with increased number of inducing points), latent signals of SR-nlGPFA exhibit significantly greater smoothness than that of SGP-VAE.
Discussion
We have developed a framework of structured recognition for variational autoencoders, which can be viewed as a strict generalisation of the SVAE [7], capturing posterior correlations induced by the “explaining away" effect.
Correspondingly, the SRVAE framework yields a tighter bound on the log-likelihood than the SVAE free energy objective.
Applying structured recognition and non-linear variational autoencoding to the GPFA framework, we introduce SR-nlGPFA, facilitating scalable free-form variational Gaussian approximation.
SRnlGPFA supports a fully non-linear, non-conjugate extension to standard GPFA [11, 23, 12], which is particularly well-suited to settings where posterior inference must be scaled to additional data.
We show that SR-nlGPFA outperforms the baseline on all presented tasks, both quantitatively and [Figure] qualitatively.
Studying population spiking data of hippocampal neurons, we show that SR-nlGPFA is able to identify latent signals that exhibit strong correlation with the behavioural covariates that are well-known to be conjunctively encoded by the recorded neurons, but without direct access to these covariates.
While here we focused on temporal correlation exhibited in time-series data, we note that the sparse amortised variational approximation presented here can also be straightforwardly applied to other GP-latent models, such as GP-LVM [44, 45].
While we have considered spiking data alone, these could be substituted by or combined with EEG, calcium imaging, or indeed behavioural data recorded during the same session.
Such joint analysis of latent structures given different recordings within the same session could potentially facilitate new findings with respect to, e.g., replay, phase coding, etc.
Acknowledgement
We thank the anonymous reviewers for helpful comments and discussions, and Zilong Ji and Dan
Bush for helpful discussions regarding the experiments with the neural data.
This work is funded by the UKRI, the Gatsby Charitable Foundation, the Simons Foundation and the Wellcome Trust.
A Details on Structured VAE
The variational free energy for VAE takes a similar formulation as the non-amortised version:
[Equation]
where we utilise a variational distribution parametrised by a recognition network, q ( z | y, ϕ ) .
The parametric assumptions in standard VAE formulation limits the expressiveness of the latent space modelling, hence leading sub-optimal training and variational inference.
In order to improve the flexibility of both the generative modelling and the variational approximation, Johnson et al.
[7] proposed structured VAE, combining PGM-parametrised latent prior distribution with amortised inference such that the flexibility of neural network modelling and the conditional dependence structure can be integrated to derive the variational distribution.
Specifically, the structured autoencoding framework considers the following free energy:
[Equation]
[Equation]
It is clear to see that the partial optimisation on q ( z | λ ) yields the following (partially) optimal variational approximation:
[Equation] [Equation]
Due to the apparent intractability, a recognition network is used to generate amortised approximation of the generative likelihood:
[Equation]
(18) where r ( z | y, ϕ ) is the recognition potential parametrised by ϕ .
With conjugate mean-field recognition potential, the resulting q ( z | λ ∗ ) can hence be computed analytically via VMP, and thus contains the factored structure inherent in the prior PGM, whilst exhibits flexibility due to the neural network parametrised recognition potential.
The generative and recognition parameters can be trained following standard stochastic optimisation by maximising the surrogate variational free energy:
[Equation]
Note that in the original SVAE formulation, Johnson et al.
[7] consider hyperpriors on θ PGM , and are updated given natural gradient descent.
Here we assume deterministic θ PGM for simplicity (except for the SRVAE-GMM model), and the structured amortisation framework can be easily extended to adapt to the variational Bayes setting.
B More Instantiations of Structured Recognition Framework with Latent Variable Models
Here we provide two additional instantiations of the SRVAE framework.
The corresponding empirical evaluations can be found in Appendix F.
B.1 Tree-Structured Latent PGM
We instantiate the SRVAE framework with discrete tree-structured PGM (Figure 4a), which we term as TreeSRVAE.
A general tree-structured PGM takes the following density function.
[Equation]
Non-linear generative likelihood functions introduces joint factors that does not exist a prior, known as “explaining away" (Section 3).
Below we provide two potential solutions
The optimal variational distribution should have the same PGM structure as the posterior, which often contains joint factors over many (possibly all) latent variables.
However, as the number of states grows exponentially with the number of latent variables, in practice it is usually not scalable to have amortised inference output a joint factor that approximates the exact posterior joint factor potentials.
We hence seek scalable alternatives.
Here we design the recognition network to output the amortised factor potential to be of the same structure as the prior PGM; in this case, a tree-structured distribution with the same set of nodes and pairwise connections (Figure 4b).
[Equation]
Note that we do not require the recognition network to output a “proper" density function, but only the factored potentials up to some normalising constant (singleton and pairwise in TreeSRVAE).
By utilising a amortised potential of the same structure of the prior distribution, we could easily perform the partial optimisation step of combining the recognition output with the prior distribution, by performing updates only on the existing factored structures.
[Equation]
where ψ ′ i ( z i )
=
ψ i ( z i ) ξ i ( z i ) , ∀ i, and
[Equation]
The training procedure follows the general VAE formulation, i.e., we need to generate reparametrised samples from the variational distribution to compute the Monte Carlo estimate of the free energy objective.
Thus far we only have the density function up to the normalising constant.
Given the tree-structured posterior latent variables, the normalising constant can be computed exactly with
Belief Propogation (BP; [6]).
At each iteration k , the message propagation between variable i and factor a are given by
[Equation]
Due to the tree structure (i.e., no loops), belief propagation converges after a single inward-outward pass (similar to forward-backward propagation in HMM models).
Upon convergence, we can compute the marginal distributions of the variables and factors as products of messages and factor potentials.
[Equation]
The singleton and pairwise marginal beliefs can be used to sample from the relevant components of q ∗ ( z | x, θ, ϕ ) .
Due to the tree-structure, the order of sampling the latent variables is irrelevant.
That is, we can initiate ancestral sampling from an arbitrary node i in the tree, computing the conditional distribution q ∗ ( z j | z i ) from the relevant marginals.
To sample categorical latent variables we employ the Gumbel-Softmax relaxation [46, 47] so that gradients of the expection can be back-propagated through the samples.
The free energy then takes the following form.
[Equation]
zzθ pgm zzzzzzx nθ x n... x nD ψ ij ( z i ,z j )
N ψ i ( z i ) (a) zzzzzzzzx nx n... x nD ψ ij ( z i ,z j )
N ξ i ( z i | x n ) ξ ij ( z i ,z j | x n ) ψ i ( z i ) (b)
Given the tree structure
[Equation]
we can compute the KL-divergence analytically with the singleton and pairwise beliefs computed with BP.
[Equation]
We note that it is also possible to use other dependency structures of the amortised factor potentials, such as a joint factor potential over all latent variables, or to employ latent structure discovery techniques to infer the latent dependency structure as well as the variational parameters [13].
B.1.1 Gaussian Factors in Tree-Structured Latent PGM
We note that treeSRVAE can also be implemented with continuous latent variables, and here we provide a simple example with Gaussian-distributed latent variables for illustration (note that the joint distribution is also Gaussian).
When working with continuous latent variables, summation is replaced with integration in the message passing steps (Eq. 22).
Since Gaussian distributions are fully parametrised by the natural parameters (same holds for other exponential family distributions), we are able to perform variational message passing by only propagating the messages of the sufficient statistics [25].
Assume the latent prior distribution specified by a Gaussian MRF has a density function of the following format.
[Equation]
where A is the precision matrix and
h =
A µ , with µ being the mean parameter of the joint Gaussian distribution.
Note that A is also known as the information matrix, which specifies the dependency structure of the latent distribution.
Assume each message at time k is parametrised by a Gaussian distribution, m ( k ) i
→ j ( x j ) ∼ N ( x j | µ ( k ) i → j , ( λ ( k ) i
→ j ) − 1 ) .
At each iteration k , the message updates are shown as following.
[Equation]
where λ ( k ) i
→ j = A 2 ij λ ( k ) i
[Equation]
; µ (
[Equation]
) i →
[Equation]
=
[Equation]
ij µ ( k ) i \ j λ ( k ) i
→ j and λ ( k ) i \
j =
A ii + (cid:88) l ∈ nbr ( i ) λ
( k ) l → i ;
[Equation] [Equation]
Then the singleton and pairwise beliefs (used for sampling and computing the KL divergence) can be computed as following.
[Equation]
where
[Equation]
where Λ ij = (cid:104)
A ii + λ ′ i
j
[Equation]
Given the prior tree structure (i.e., the information matrix, A ), the recognition network outputs the sufficient statistics (mean and precision matrix) of a Gaussian distribution with the same structure as the tree-based prior distribution, N ( z | µ ′ , ( A ′ ) − 1 ) .
Note that instead of parametrising a full-rank precision matrix, which requires O ( D 2 ) outputs
( D is the number of latent variables), by constraining the dependency structure as the prior distribution, the recognition network only need to output 3 D −scalars (for each batch), where 2 D − 1 contribute towards the non-zero entries in the precision matrix, and D contribute towards the linear term.
Such type of parametrisation is similar to the modelling of off-diagonal elements in the covariance matrix with low-rank decomposition, but we note that the off-diagonal entries in the covariance matrix do not convey information about the dependency structure of the latent variables, hence our framework provides stronger interpretability.
Dorta et al.
[48] proposes a direct low-rank parametrisation of the precision matrix in a VAE setting, which again lacks interpretation in terms of the latent dependency structure.
Moreover, we note that they require the matrix inversion to convert the precision matrix into the covariance matrix for inference and sampling, which requires O ( D 3 ) complexity, whereas we apply the BP, which on tree-structured
PGMs only require O ( D ) iterations to converge.
B.2 Gaussian Mixture Model Latent PGM
The PGMs for the generative and inference models of SRVAE-GMM is shown in Figure 5.
C Joint Factor Potentials Induced by Non-Linear Likelihood (“Explaining Away")
Here we show that even a simple non-linear likelihood function can induce non-trivial posterior correlations that do not exist a priori .
Consider a Bernoulli generative likelihood function, with the π z n
[Equation]
n = 1 ,...,N (a)
π z n x n θ y n n = 1 ,...,N
[Figure]
logits being an affine transformation of z .
[Equation]
where θ = { W, b } .
Given this simple generative model
[Equation]
and the tree-structured latent distribution (Eq. 20)
[Equation]
we can derive the true posterior distribution analytically.
[Equation]
where ϕ i ( z i )
=
ψ i ( z i ) exp( z i (cid:88)
[Equation]
We observe that the true posterior distribution involves the singleton and pairwise potentials that preserve the structure of the prior distribution, and at the same time the normalising constant of the conditional likelihood function introduces a joint factor over all latent variables, ϕ z ( z ) , which cannot be captured by the fully factorised amortised potential, but can be partially captured by the tree-structured potentials.
D Further Details of SR-nlGPFA
The derivation of the variational free energy objective for svGPFA (and SR-nlGPFA; Eq. 5) is shown as following Titsias [22].
[Equation]
where H [ q ] is the entropy of the variational distribution q ( u ) .
E Proof of Proposition 3.
We re-state the proposition.
Proposition E.1.
Consider the following latent structured prior.
[Equation]
where we assume θ 0 is the set of trainable deterministic prior parameter.
Consider the free energy objective.
[Equation]
Both the SVAE objective and the AEA-objective take the following expression.
[Equation]
For SVAE, q ∗ ( z ) is derived through partial optimisation given conjugate amortised inference outputs.
[Equation]
where (cid:81)
i l
i ( z i | x, ϕ ) represents the approximate local evidence potentials (corresponding to the product of singleton potentials).
For SRVAE, instead of being fully factorised, the recognition network outputs structured factor potentials, which in principle, could contain factors of arbitrary set of latent variables (e.g., a joint factor for all latent variables).
[Equation]
Then the SRVAE objective function provides tighter lower bound to the free energy than the SVAE objective function.
[Equation]
We note that throughout the paper we assume both the SRVAE and SVAE models employ identical generative models, hence we have the stronger results that the inequality for all θ .
For SR-nlGPFA, despite introducing the additional affine transformation (parametrised by C and d ) in the generative process, the linear operation can be completely subsumed into the neural-network generative model, hence leading to the identical generative model as SGP-VAE.
Proof.
The first inequality is trivial.
For the second inequality, it is easy to see that the set of fully factorised amortised potentials, L = { l ( z ) | l ( z ) ∝ (cid:81)
i l i ( z i | x ; ϕ ) } , is a strict subset of the set of all structured amortised potentials, H =
[Equation]
z c | x ; ϕ ) } , where C ∗ is the set of factors in the amortised approximation to the likelihood.
Hence for arbitrary p ( x | z, θ ) , we have that,
[Equation]
with equality if and only if p ( x | z, θ ) ∈ L .
Hence it is trivial to derive that,
[Equation]
with equality if and only if p ( x | z, θ ) ∈ L .
We note that having a tighter free energy lower bound does not necessarily lead to stronger generative modelling [49].
However, by exploring an alternative formulation of the free energy objective (relative to Eq. 1):
[Equation]
we see that a tighter free energy bound leads to smaller KL-divergence between the variational approximation and the ground-truth posterior distribution, hence achieving more accurate posterior inference.
This is indeed reflected in our empirical evaluation of SR-nlGPFA on the spiking dataset (that the posterior latent processes is significantly more coherent with the underlying behavioural covariates than that of SGP-VAE [10]), and also reflected by the better generation quality in the bar dataset (F.1).
F Further Experimental Results F.1 TreeSRVAE on “Bar" Dataset
We first evaluate tree-structured recognition potential on a synthetic “Bar" dataset, consists of squaregrid observations with horizontal and vertical bars.
[-where pa ( z ) denote the parent node of node z , σ ( · ) = 
-]
[Equation]
where pa ( z ) denote the parent node of node z , σ ( · ) = 1 1+exp( · ) is the sigmoid function, and B ( p ) denotes the Bernoulli distribution with rate p .
Each binary z i denotes the activation of one bar in the square-grid (e.g., z 1 = 1 indicates the presence of the first horizontal bar).
The generative parameters
W and b are defined as
[Equation]
if pixel i is on the bar j otherwise , b i = [Equation] i , for i = 1 , . .
[Equation]
where { ω i } D 2 i =1 are the temperature parameters that control the degree of randomness in the data generation.
Note that in current implementations we set ω
i = ω for all i .
Exemplary samples from the “Bar" dataset if shown in Figure 6 ( D = 8) . [Figure]
Note that here we only perform variational inference, instead of variational Bayes optimisation, hence we assume the parameters of the prior distribution (Eq. 20) to be deterministic quantities (which can still be optimised via gradient descent, but do not follow any natural-gradient updates).
Given the tree-structured prior distribution (Eq. 32) [Figure] we have that the singleton and pairwise potentials take the following expression.
[Equation]
We use the Gumbel-Softmax trick for relaxed reparametrised sampling of the discrete latent variables.
However, we can also apply a hard-transformation to the relaxed samples to generate the binary-valued latent samples as following.
[Equation]
In practice we find using the hard-sampling trick improves performance for all models considered (VAE, SVAE, TreeSRVAE) on the “Bar" dataset.
We set the prior distribution to be a "uninformative" prior, where p 0 = 0 .
5 and p iij for all ( i, j ) (Eq. 34).
In Figure 7 we show the training curves of the free energy (Eq. 7), reconstruction loss ( log p ( x | z ) ) and
KL-divergence (with respect to the prior distribution, KL [ q ∗ ( z, ϕ ) ||
p ( z, θ pgm )] ) for the “Bar" dataset generated with varying sampling temperature ( ω ). , where the observation is over an 8 × 8 square grid, hence the latent variables are z ∈ { 0 , 1 } 16 .
We evaluate on three models: i) the latent variable follows a fully factorised Bernoulli distribution, and the recognition network outputs mean-field inference, which is equivalent to standard VAE with Bernoulli latents [46]; ii) the latent variable follows a
[Figure]
VAE
[Figure] TreeSVAE [Figure] distribution characterised by a tree-structured probabilistic graphical model, and the recognition network outputs fully factorised potentials, which is equivalent to the original SVAE setup with tree-structured Bernoulli latent [7]; iii) the latent variable follows a distribution characterised by a treestructured probabilistic graphical model, and the recognition network outputs tree-structured factor potentials (containing both the singleton and pairwise factor potentials), which is the TreeSRVAE.
Note that we use the “hard-transformation" trick for all models presented in Figure 7.
From Figurewe see that TreeSRVAE consistently outperforms both VAE and SVAE on the “Bar" dataset with varying level of sampling noise, in terms of the overall free energy (both the sample efficiency and asymptotic performance), hence providing a tighter lower bound to the true likelihood.
What if the underlying latent structure deviates largely from being tree-structured?
We consider the following “side-dependent" prior distribution ( p ( y | z , W , b ) remain the same).
[Equation]
(36) where Cat ( p ) denotes the categorical distribution with probabilities p .
Intuitively, sampling from such prior distribution means that we will observe one and only one bar on each side of the square.
Clearly, such prior distribution cannot be modelled by any tree-structured distribution.
We apply the same set of models to the data from the new generative model, and the results are shown in Figure 8.
We observe that the performances for all three methods are mostly similar.
Hence in situations where there exists a large mismatch between the structures of the amortised potentials and the true posterior distribution, the additional structure present in the amortised inference will not interrupt learning, but instead allow the learning process be at least on the same level of performance as standard VAE and SVAE in terms of free energy.
Despite having achieved similar free energy through training (Figure 8), we argue that the structured recognition framework should still be preferred.
Namely, we examine the generation quality of the generated samples.
Firstly, visual inspection of the randomly samples from the three trained models in Figure 9 indicates TreeSRVAE generates more “cross" patterns than SVAE and VAE.
The qualitative indicates
TreeSRVAE has learned a better generative model than the other two models.
Now for the quantitative comparison.
Given the relatively small latent state space ( 16 binary latents, leading to 2 16 = 65536 possible latent configurations).
We exhaustively generate samples from all possible latent configurations given the three models.
For each generated sample, we compute the most similar “cross" pattern (in terms of squared error) with respect to the sample, and their squared distance.
We then compute the averaged smallest squared error given the generated samples of all possible 2 16 latent configurations for the three models.
The statistics is reported in Table 2.
Namely, for each z ∈ R 16 ,
[Equation]
(37) where Y denotes the set of all possible cross patterns ( 64 different patterns in an 8 × 8 square-grid.
Hence we compare ˜ d = 1 |Z| (cid:80) z ∈Z d z , where Z is the set of all possible latent configurations. [Figure]
Hence we have shown that, despite having achieved similar free energy, TreeSRVAE enables stronger generation both qualitatively and quantitatively. [Figure]
[Table]
F.1.1 Tree-Structured Latent PGM with Gaussian Factors
We note that the tree-structured recognition network generalises beyond discrete latents.
Namely, we developed a Gaussian-TreeSRVAE model (details in appendix B.1.1) and applied to the MNIST dataset [50].
Table 3 show that Gaussian-TreeSRVAE outperforms the two baseline models.
[Table]
F.2 SR-nlGPFA F.3 Full Quantitative Evaluation Comparison
The complete evaluation comparison between SR-nlGPFA and the selected baselines in terms of both the SMSE and the negative log-likeihood (NLL) is shown in Table 1.
F.3.1 Complexity Analysis of SR-nlGPFA
We note that SR-nlGPFA, despite allowing the inference of full-covariance structure over all latent processes, does not incur higher-order complexity than standard SVAEs with GP-latents (e.g., SGPVAE [10]), hence allowing scalable application.
The major difference between SR-nlGPFA and
SGP-VAE lie in the inference step, where SGP-VAE requires only inverting the covariance matrices for the latent-specific inducing points, and SR-nlGPFA needs to numerically invert the covariance matrices over all inducing points across latent dimensions.
Hence the complexities for such operation is O
( K 3 ¯ M 3 ) and O ( K ¯ M 3 ) for SR-nlGPFA and SGP-VAE, respectively.
Hence SR-nlGPFA incurs additional computations by a factor of K 2 .
However, we note that we usually seek a low-dimensional
GPFA latents that characterise the manifold upon which the high-dimensional neural trajectories lie within, hence K is usually chosen to be small (e.g., K = 4 for the single-cell spiking dataset from Section 5.2.2 and K = 2 for the EEG dataset considered in Section 5.2.1).
Hence overall speaking, we expect considerably small increase in computation time.
We show the comparison between the running times with varying latent dimensions and number of inducing points for SRnlGPFA and SGP-VAE in Figure 10, which is in accordance with our hypothesis that small additional computational cost is incurred with SR-nlGPFA.
[Figure]
F.3.2 Further Analysis of SR-nlGPFA on Neural Population Spiking Dataset
The firing fields of some exemplary neurons on the Z-shaped track in shown in Figure 11.
Visual inspections show that the majority of the recorded cells show spatial modulation [Figure] and resembling the firing patterns of place cells and (potentially) grid cells, amongst the recorded neurons there also exists ones whose firing patterns resemble that of interneurons (e [Figure]
g., third neurons from the left on the third row)
[Figure]
Extraction of behavioural-covariate-modulated neurons predicted by the trained SR-nlGPFA model. Given the behavioural covariate of interests, we wish to predict the neurons whose firing pattern exhibit strongest modulation with respect to the behavioural covariate using the trained SR-nlGPFA model (in a completely unsupervised fashion).
Specifically, given the projection matrices from the CCA fitting, we could qualitatively identify the latent processes that are most correlated with the target behavioural covariate.
Given the generative likelihood, we could compute the gradient of the output Poisson rate parameters for each neuron with respect to the selected latent variable (the Jacobian matrix) using AutoDiff as found in most deep learning libraries [51].
We define the relevance score of a latent process, f , with respect to the firing of neuron s , ∆ f ( s ) as the average squared-norm of the gradient of the associated mapping (given the likelihood function, Eq. 14).
[Equation]
We categorise the top 20% of the neurons with the highest (or lowest if negatively correlated) relevance score as the model-predicted neurons that are modulated by the target behavioural covariate.
Additional svGP Inference Step with Changed Inducing Locations
We note that since the recognition network outputs local evidence potentials on h , the computation of q ( U ) is not constrained to a single set of inducing locations Z , rather we are free to choose the set of inducing points to work with (13).
During training, due to the computational constraints of stochastic mini-batch training, it is impractical to use a large number of inducing points.
However, the small number of inducing points leads to temporal chunking artifacts in the variational posterior distribution of the latent processes ( q ( f ) ) over the entire sequence.
Hence SR-nlGPFA enables an additional svGP inference step given the trained recognition network and an expanded set of inducing points, leading to smoother temporal interpolation, hence alleviating the temporal chunking effects
. Quantitative Analysis of SR-nlGPFA across Experimental Sessions
Here we quantitatively evaluate various aspects of the posterior latent GPs given the trained models.
We train independent models of SR-nlGPFA with the same architectures and training procedures (see appendix G for implementation details) on the single-cell population spiking datasets from each of the 28 experimental sessions considered.
Firstly, we examine if full-covariance structure is a necessary assumption.
We compute the ratio between the magnitudes of the off-diagonal entries and that of all entries in the posterior covariance matrix ( S fn , for all n , Section 5.2.2), (cid:80) n (cid:80) i,j,i ̸ = j [Equation]
S fn
[Equation]
i,j | S fn
[Equation]
| .
From Figure 12a we observe that the posterior latent covariance matrices of all sessions have non-trivial off-diagonal elements, indicating the existence of posterior correlations induced by “explaining away" and the necessity of posterior inference with full covariance structure for capturing the induced correlations.
In order to assess the quality of the learned posterior latents, we perform two-dimensional Canonical
Correlation Analysis (CCA, [42]) on the learned posterior means and the behavioural covariates.
In
Figure 12b we observe that the correlations between the extract canonical correlates (CC) of the posterior means and the behavioural covariates exhibits high correlation over all sessions, hence indicating the extracted posterior latent dimensions captures large proportion of the information in the behavioural variables.
We further show the correlations between the CCX { 1 , 2 } (CCs of the latent posterior means) and each of the behavioural covariates for all sessions in Figure 13.
We observe strong correlations between one or both of the CCXs with the (unfolded) distance along the track quantity, showing that the extracted latent variables given the population spiking data is strongly indicative of the spatial location of the rat along the Z-shaped track, and corresponds nicely with prior knowledge that the majority of the recorded cells are located in the CA1 region of hippocampus and are well-known to exhibit significant spatial modulation [36].
Moreover, in most sessions we observe high correlations (in terms of magnitude) between the CCXs and the direction of travelling, again conforming with experimental evidence that CA1 place cells firing are modulated by direction along the linear track (the Z-shaped track environment is usually interpreted as a 1D linear track rather than a 2D environment [36]).
We additionally observe that the CCXs of the learned latent variables in most sessions show strong correlation with respect to speed of travelling, which also has nice experimental correspondence [38, 40]. [Figure] [Figure] [Figure]
[Figure]
[Figure]
Comparison with SGP-VAE
By applying SGP-VAE to the data from the same session studied in Section 5.2.2, we perform similar
CCA analysis on the learned latent variables of SGP-VAE.
We observe that the correlations between
CCX { 1 , 2 } and CCY { 1 , 2 } of SR-nlGPFA are significantly larger than that of SGP-VAE (with standard t-test, p-values shown in the figure; Figure 14a).
We additionally examine the same session we looked at in Section 5.2.2, and plot the CCX { 1 , 2 } of the learned latent variables of SGP-VAE against the x- and y-positions.
By comparing to the similar plot with SR-nlGPFA (Figure 3d), we observe that the resulting plot does not show clear clustered structure with respect to spatial and direction covariates compared to that of SR-nlGPFA, hence illustrating that SR-nlGPFA learns qualitatively better latent dimensions than SGP-VAE with respect to the underlying behavioural covariates. G Implementation Details “Bar" Dataset
The details of the generation of the “Bar" dataset can be found in appendix ?? .
We use the same neural network architectures for all models considered (VAE, SVAE, TreeSRVAE): both the recognition and generative networks are MLPs with two hidden layers of 50 hidden units, with ReLU non-linear activation function, and the latent dimension is 16 .
All models and sessions are trained with Adam optimiser with learning rate of 5 × 10 − 4 over 3500 epochs with a batch-size of 256 [52]. Neural Population Spiking Dataset We use the same neural network architectures for both SR-nlGPFA and SGP-VAE: both the recognition and generative networks are MLPs with two hidden layers with 256 hidden units, with ReLU activatieon function. The latent dimensions for both model is 6 , which we choose via cross-validation.
Note that the latent dimension of 6 corresponds to the results from other existing unsupervised latent feature extraction works on CA1 neuron firing patterns (e.g., Nieh et al.
[53] reported 5 − 7 is optimal for the latent dimension with their manifold inference model).
The dimension of the neurla feature ( h ) in the GPFA generative model (Eq. 3) is 20 .
Both models are trained with Adam optimiser with learning rate of 1 × 10 − 4 over 400 epochs with a batch-size of 128 [52].
The number of inducing points for each latent dimension is 64 for both models.
For the additional SVGP inference step after training with SR-nlGPFA (SGP-VAE does not allow the additional SVGP inference step as discussed in Section 5.2.2), we set the number of inducing points to be 1000 for all latent dimensions to smooth out the temporal chunking artifacts caused by the small number of inducing points in training due to scalability concerns with stochastic optimisation.
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 5 .
3 e0 .
C o rr e l a t i o n h e a t m a p 20151125 _ R 2337 _ t
r
a c k 1 ( 97 )
CCX
[Table] [Table]
_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 1 . 6 e0 .
C o rr e l a t i o n h e a t m a p 20140110 _ R 2192 _ t r a c k 1 ( 100 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 - 8 .
9 e - - 0 .
C o rr e l a t i o n h e a t m a p 20141213 _ R 2217 _ t
r
a c k 1 ( 88 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 3 .
7 e - - 0 .
C o rr e l a t i o n h e a t m a p 20151030 _ R 2336 _
t
r a c k 1 ( 77 )
[Table]
_1 CCX_ trackDist
[Table]
dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 1 .
9 e - - 0 .
C o rr e l a t i o n h e a t m a p 20151027 _ R 2335 _ t r a
[Table]
X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 - 9 . 2 e0 .
C o rr e l a t i o n h e a t m a p 20151201 _ R 2337 _ t
r
a c k 1 ( 101 )
CCX_1 CCX_ trackDist speed dir hd
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 3 . 7 e0 .
C o rr e l a t i o n h e a t m a p 20141210 _
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 9 .
2 e - - 0 .
[Table] [Table]
C o rr e l a t i o n h e a t m a p 20141218 _
[Table]
_ t r a c k
[Table] [Table]
[Table] [Table]
dir hd unfolded
[Table]
[Table]
_ 2 t r a c k D i s t s p ee d d
i r h d un f o l d e d
[Table]
[Table] [Table]
n
[Table]
a t
[Table]
a p
[Table] [Table]
unfolded
[Table] [Table]
t r a c k D
i s t s p ee d d i r
[Table]
d
[Table]
f o l d
[Table]
d
[Table]
[Table] [Table]
[Table] [Table]
n
[Table]
p
[Table]
C o rr e l a t i o n h e a t m a p 20140915 _ R 2192 _ t r a c k 1 ( 66 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 1 .
6 e - - 0 .
C o rr e l a t i o n h e a t m a p 20151124 _ R 2337 _ t
r
a c k
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 6 . 1 e0 .
C o rr e l a t i o n h e a t m a p 20140918 _ R 2198 _ t r a c k 1 ( 79 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 1 .
1 e - - 0 .
C o rr e l a t i o n h e a t m a p 20151127 _ R 2337 _ t
r
a c k 1 ( 125 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 - 8 .
7 e - - 0 .
C o rr e l a t i o n h e a t m a p 20151101 _ R 2336 _
t
r a c k 1 ( 64 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 5 e0 .
C o rr e l a t i o n h e a t m a p 20151031 _ R 2336 _
t r a c k 1 ( 58 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 - 1 . 3 e0 .
[Table]
[Table]
C o rr e l a t i o n h e a t m a p 20151024 _ R 2335 _ t r a c k 1 ( 68 )
[Table] [Table]
_ trackDist
[Table]
dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 - 4 .
8 e - - 0 .
[Table] [Table]
[Table]
C o rr e l a t i o n h e a t m a p 20151128 _ R 2337 _ t
r
a c k 1 ( 97 )
CCX_1 CCX_ trackDist speed dir hd unfolded
[Table]
t r
[Table]
c k D i s t s p ee d d
i r h d un f o l d
e d 1 - 7 . 9 e - - 0 .
13 6 . 4 e0 .
C o rr e l a t i o n h e a t m a p 20151026 _ R 2335 _ t r a c k 1 ( 96 )
[Table] [Table]
unfolded
[Table]
p
[Table]
i
[Table]
f
[Table]
d 1 - 1 . 3 e
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
C o rr e l a t i o n h e a t m a p 20151126 _ R 2 337 _ t
r
a c k 1 ( 100 )
CCX_
[Table] [Table]
dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 7 . 3 e0 .
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
C o rr e l a t i o n h e a t m a p 20151025 _ R 2335 _ t r a c k 1 ( 69 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 - 2 .
7 e - - 0 .
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
C o rr e l a t i o n h e a t m a p 20140920 _ R
[Table]
_ t r a c k 1 ( 74 )
[Table]
[Table] [Table]
dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 2 .
8 e - - 0 .
[Table] [Table]
[Table] [Table]
[Table]
C o rr e l a t i o n h e a t m a p 20151104 _ R
[Table]
_ t r a c k 1 ( 102 )
[Table] [Table]
_ trackDist speed dir hd unfolded
CC X _
[Table]
X _ 2 t r a c k D
i s t s p ee d d i r h d un f o l d e d 1 2 e0 .
C o rr e l a t i o n h e a t m a p 20140921 _ R 2198 _ t r a c k 1 ( 121 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d
e d 1 - 5 . 1 e0 .
C o rr e l a t i o n h e a t m a p 20140916 _
R 2192 _ t r a c k 1 ( 74 )
CCX_1 CCX_ trackDist speed dir hd unfolded
CC X _ 1 CC X _ 2 t
r a c k D i s t s p ee d d
i r h d un f o l d e d 1 - 2 . 9 e0 .
C o rr e l a t i o n h e a t m a p 20140917 _ R 2192 _ t r a c k 1 ( 111 ) 0 . 0 0 .
[Figure]