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 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) has the form
[Equation]
where (with [ x i ] denoting an array obtained by iterating over the index i ), y t =
[ y n ( t )] ∈ R N , f t =
[ f k ( t )] ∈ R K , F = [ f Tt ] ∈ R T × K , C
= [ c nk ] ∈ R N × K , d = [ d n ] ∈ R N , U =
[ u k ] T ∈ R K × M , and Z =
[ 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]
and [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]
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
The first experiments in Section 5 apply SRVAE to the setting of a latent Gaussian mixture model (GMM) [24] of the form (Figure 5a)
[Equation]
This model was studied by Johnson et al.
[7], who used an SVAE with fully factorised recognition potentials and variational message passing[VMP; 25] to obtain the GMM variational posterior (details can be found in Appendix A).
However, the neural-network defined conditional model induces posterior correlations between the latent variables beyond those imposed by the prior.
Thus here we used an SRVAE model, with a full covariance Gaussian recognition potential on h , combined with the prior factor using VMP.
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 and define 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, the input x t is usually taken to be identical to the timestamp t ; see (3).
Here we consider the more general case of arbitrary time-dependent inputs.]
The generative model we consider retains the affine mapping from f t to h t , but replaces 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]
and 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 arguably be 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 pinwheel dataset of [7] (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 a more accurate learnt model.
[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.
Baselines
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.
These 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, and g ( · ) represents the neuralnetwork 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 GPFA, we introduce SR-nlGPFA, facilitating scalable free-form variational Gaussian approximation.
SR-nlGPFA 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 qualitatively.
[Figure]
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, DeepMind, the Gatsby Charitable Foundation, the Simons Foundation and the Wellcome
Trust.