Amortised Inference in Structured Generative Models with Explaining Away
Anonymous Author(s)
AffiliationAddress email
Abstract
A key goal of unsupervised learning is to go beyond density estimation and samplegeneration to reveal the structure inherent within observed data.
Such structurecan be expressed in the pattern of interactions between explanatory latent variables captured through a probabilistic graphical model.
Although the learning ofstructured graphical models has a long history, much recent work in unsupervisedmodelling has instead emphasised flexible deep-network-based generation, eithertransforming independent latent generators to model complex data or assumingthat distinct observed variables are derived from different latent nodes.
Here,we extend the output of amortised variational inference to incorporate structuredfactors over multiple variables, able to capture the observation-induced posteriordependence between latents that results from “explaining away” and thus allowcomplex observations to depend on multiple nodes of a structured graph.
Weshow that appropriately parameterised factors can be combined efficiently withvariational message passing in elaborate graphical structures.
We instantiate theframework based on Gaussian Process Factor Analysis models, and empiricallyevaluate its improvement over existing methods on synthetic data with knowngenerative processes.
We then fit the structured model to high-dimensional neuralspiking time-series from the hippocampus of freely moving rodents, demonstratingthat the model identifies latent signals that correlate with behavioural covariates.
Introduction
A central goal of unsupervised learning is the extraction of low-dimensional latent structure thatcaptures dependence between the component variables of high-dimensional observations.
Recentdevelopments in deep learning have led to a family of deep generative models (DGM) that haveenjoyed remarkable success in density estimation and high-fidelity image or text generation [21, 42,13, 14, 35].
However, much of the focus with DGMs has been on increasing the expressiveness andaccuracy of generation, often overlooking the importance of uncovering the dependency structureinherent within the observed data
Structured variational inference, by including conditional dependence between the set of latentvariables, improves the expressiveness of the variational approximation for the latent posteriordistribution.
The structured variational autoencoder [SVAE; 18] connects the DGM framework withone element of such inference, combining a structured prior distribution with fully factored (andfunctionally conjugate) amortised likelihood potentials.
However many generative models of interestinvolve observations that depend on more than one latent variable.
This induces joint potentials inthe likelihood, coupling latents in the posterior even if they are independent a priori , a phenomenonsometimes known as “explaining away” [38].
Our goal here is to extend the amortised inferencesetting to capture such observation-induced dependence in the DGM setting.
We propose structured amortised inference where a structured recognition potential is utilised incombination within the structured autoencoding framework [18].
We instantiate the framework byextending Gaussian Process Factor Analysis (GPFA) models that are widely used in neuroscience,to incorporate deep network-based nonlinear mappings from the latent Gaussian process values toobservations.
The corresponding likelihood functions induce posterior “explaining-away” correlationsbetween the processes.
We develop “amortised explaining-away sparse variational GPFA” (AEAsvGPFA), where we introduced a novel variational approximation and achieve scalable amortisedinference of the full-covariance posterior distribution within the sparse inducing-point GP framework.
We empirically evaluate AEA-svGPFA on both synthetic and real datasets, including a spike-countdataset collected from a population of neurons in the hippocampal complex, and show that theunsupervised approach automatically captures dimensions underlying neural firing that reflect relevantbehavioural correlates.
Background2.1 Variational inference and structured variational autoencoders
We consider the variational inference (VI) problem of approximating the posterior distribution of thelatent variables z given observations y .
Consider a general generative process
[Equation]
where θ and γ are the parameters of the prior and conditional likelihood distributions, respectively.
Note that we assume deterministic θ for simplicity, but hyperpriors on θ can be incorporated inthe context of Variational Bayes (VB).
The posterior distribution, p ( z | y )
∝ p ( z | θ ) p ( y | z, γ ) , is 55 often analytically intractable.
In such cases,a common approach is to seek a tractable variationalapproximation to the true posterior, q ( z ) , learned by minimising the KL-divergence with respect tothe true posterior distribution or, equivalently, maximising a variational free energy [47]:
[Equation]
In the Variational Autoencoder (VAE) architecture, the iterative optimisation of q ( z ) implicit in (1) isamortised by a recognition network that takes as input the observed values y and returns parametersof q . The parameters of this recognition network are then trained jointly with generative model bystochastic optimisation of the free energy [21, 42].
Although the mean-field variational distribution is often utilised in both standard and amortised VIowing to its simplicity, the assumption of fully factorised latent distributions can be too restrictive tocapture the complexity of the true posterior [44, 16].
A number of approaches have been proposed forimproving the expressiveness of variational approximation [5, 46, 28, 26].
Here we specifically look atstructured VAE [SVAE; 18].
Unlike standard VAEs, SVAE assumes the generative prior distribution tobe specified by a structured probabilistic graphical model (PGM), p ( z | θ ) ∝ (cid:81) c ∈ C ψ c ( z c ) .
Moreover,the amortised inference network outputs factors, r ( z | y, ϕ ) , that approximate the generative likelihoodfunction rather than the variational distribution itself.
These are combined with the structured priordistribution to obtain the variational posterior.
[Equation]
Johnson et al.
[18] assumed that the amortised inference outputs fully factorised potentials (chosen tobe conjugate to p ( z | θ ) ).
Comparing to the mean-field approach in the standard VAE, Eq. 2 is able tocarry the dependency structure in the prior distribution to the variational posterior distribution.
Weprovide a more in-depth review of the SVAE framework in Appendix.
2.2 Gaussian Process Factor Analysis
Gaussian Process Factor Analysis (GPFA) is a model used in neural data analysis to infer the dynamiclatent structure underlying high-dimensional population spike trains [50].
Standard GPFA assumesthe following generative model.
[Equation]
where c jk and d n are parameters defining an affine mapping from the latent space to the neuralspace (or a transform of the neural space for general likelihoods), and g ( · ) is some (point-wise)smooth function.
We assume m kθ ( · ) =
unless stated otherwise.
The temporal correlation in thehigh-dimensional neural space are fully captured by the latent temporal correlations modelled bythe latent processes.
The PGM describing the GPFA generative model is shown in Fig 1a.
Theparameters of GPFA can be fitted with (variational) expectation-maximisation [9, 50].
However, the applicability of GPFA models is limited by the poor scalability of exact posteriorinference of Gaussian processes (GP), which scales cubically with the number of observations [40].
To this end, an extension of the sparse variational approximation for GPs is applied to the GPFAframework, leading to the svGPFA model [45, 1, 12].
Details of svGPFA formulations are includedin our discussion of AEA-svGPFA in Section 3.2.
Structured amortised inference and explaining away3.1 Structured amortised inference with SVAE
A general non-linear likelihood induces additional latent dependency structure in a posterior thatdoes not exist a priori (we provide a simple example in Appendix B illustrating such effect).
Despiteintroducing the structured prior distribution, the factored recognition potential assumption of SVAEonly augments the prior distribution with singleton potentials. This approach cannot accurately modelposteriors with dependence that differs from that of the prior PGM.
Instead of enforcing a fully factorised form, we assume the recognition network instead outputs structured factor potentials over the latents, providing amortised estimates of the additional dependenciesinduced by “explaining away”.
Following the SVAE formulation, the “Amortised Explaining Away”(AEA) variational approximation takes the following form.
[Equation]
where C r is the set of recognition factors and ξ c is the associated factor potential.
We assumeconjugate factors unless otherwise stated.
Hence the analytical (approximate) form of q AEA ( z | x, ϕ, θ )can be computed with (loopy, if necessary) belief propagation.
In the most general form, we could letr ( z | y, ϕ ) be a single joint factor potential over the complete set of latent variables.
The free energy objective has the identical functional form as the standard autoencoding free energyobjective [21, 42], but with the amortised structured variational approximation, q AEA , as the variationaldistribution (same applies with the SVAE free energy, with q SVAE (Eq. 2)).
[Equation]
The training of the model follows standard VAE-style stochastic optimisation to update the parametersof the recognition and generative networks, as well as (optionally) the parameters of the prior PGM.
The following proposition justifies the utility of using structured amortised inference (the proof canbe found in appendix D).
Proposition 3.1.
The AEA-SVAE objective function provides a tighter lower bound to the free energythan the SVAE objective function.
[Equation]
Below we describe one specific instantiation of the AEA framework based on the GPFA models,where we show that our novel variational approximation allows scalable inference of the posteriorlatent distribution with full covariance structure, hence enabling arbitrary posterior dependencystructure to be captured.
The AEA framework can be additionally combined with general latentvariable models, which we provide more examples in Appendix A.
3.2 Structured autoencoding nonlinear GPFA
Neural computations are thought to be implemented by the joint activity of large populations ofneurons.
Unsupervised methods have been developed to identify the dynamic covariates that underlie 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) population neuronal spiking trains, and relate these to behaviour or environmental variables so as togain understanding of the functional and cognitive aspects of neural computations.
One popular modelis GPFA [50].
In order to alleviate the scalability issue, sparse variational approximation inferencefor GPFA with inducing points has been proposed [45, 1, 12].
For k = 1 , . . .
, K , the inducing pointsu k represents the function values of the latent process f k at the M k inducing locations, z k .
We thenaugment the GPFA generative model (Eq. 3) with the inclusion of the inducing points.
[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 vectorvalued 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 Grammatrix of evaluation of the kernel function κ kθ at the inducing locations z k .
By introducing a variational distribution over the inducing points and the latent functions, and utilisingthe generative model (Fig 1a), the free energy of the sparse variational GPFA (svGPFA) can beexpressed as following.
[Equation]
It is common practice to use the mean-field variational approximation of the joint posterior distributionof the latent process and the inducing points [45, 1], q ( f 1: K , u 1: K ) =
(cid:81)
Kk =1 p ( f k | u k ) q ( u k ) , whereq ( u k ) =
N ( m k , S k ) .
Hence the mean-field free energy objective simplifies as following.
[Equation]
where q ( h n ) = (cid:82)
d U p ( h n | U ) q ( U ) .
It is easy to derive that q ( h n ) ∼ GP ( m hn ( · ) , ν hn ( · , · )) is itself a 138 GP with mean and kernel function given by the following expressions.
[Equation]
Hence the computational cost reduces from O ( T 3 ) down to O ( ¯ M 3 ) , where ¯ M = 1 K (cid:80) Kk
=1 M k , andscales linearly with respect to the number of observations.
By introducing a recognition network into the svGP framework, we could perform amortised inferenceover the inducing points, hence improving the scalability of the model.
In order to preserve thewithin-process latent correlations, Ashman et al.
[2] introduced the SGP-VAE, where the variationaldistribution over the latent processes is computed by combining the prior GP distribution with an amortised approximation of the likelihood function given a set of observation pair, { ( x b , y b ) } Bb =1 146 with the standard mean-field variational approximation [45].
[Equation]
where p ( U ) =
(cid:81) k N ( m kθ ( z k ) , k kθ ( z k , z k )) .
The graphical model for SGP-VAE inference is shownin Figure 1b, with C
=
I and d = 0 .
However, we note that this approach still fails to capturecorrelations between the latent processes that arise in the posterior through “explaining away”, henceleading 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 | y b ) = N ( h | m ( y b ; ϕ ) , ψ ( y b ; ϕ ))
=
N ( µ hb , Ψ hb ) , where Ψ hb is diagonal.
We note that 154 here we assume g ( · ) is some general non-linear mapping instead of being a point-wise functionas in the standard GPFA model, which could greatly increase the expressiveness of the generalmodelling.
Given the linear-Gaussian relationship between U and h (Eq. 11, Fig 1a), the integralcan be computed in closed form.
Furthermore, even though the recognition potential on h , r ( h | y b ) ,is assumed to be fully factorised over the dimensions of h for each y b , the variational distributionon U includes coupling between latent processes induced by the affine mixing coefficients.
Henceq ( U ) captures the correlations given both by between the latent processes and by the latent temporaldynamics (through combination with the GP prior under the SVAE formulation).
Specifically, given the general expression of the GPFA generative model and the sparse approximationwith inducing points, we have the linear Gaussian relationship between h and U .
[Equation]
where
[Equation]
with U =
[ u 1 , . . .
, u K ] being the collection of all inducing points for all latents, and K x =diag [ κ kθ ( x, x )] .
Given the fully factorised recognition potential on h b (Eq. 10), we could compute thestructured variational distribution on U given the linear-Gaussian relationship p ( h | U ) (Eq. 11).
[Equation]
b p ( U ) N ( CF ( x b ) U | µ hb , Ψ hb ) , with
[Equation]
Hence giving the b − th datapoint, x b = x , we have the following:
[Equation]
Then the complete (reparameterised) Monte Carlo estimate of the free energy objective given a batchof data, { ( x b , y b } Bb =1 , takes the following expression.
[Equation]
where ϵ s ∼ N ( 0 , I ) , θ is the set of kernel parameters, L hn is the lower-triangular matrix such that
S
hn
=
L hn ( L hn ) T , and Z =
[ z 1 , . .
z K ] is the set of all inducing locations.
We note that we have retained the scalability of amortised inference by outputting a diagonalamortised Gaussian potential on h , and due to the relationship between the inducing points U and h
( F -projection and affine transformation, Eq. 11), we are able to compute a full-covariance variational
Gaussian approximation for U , allowing the computation of posterior correlations between latentprocesses induced by “explaining away”.
We term our model as Amortised Explaining Away svGPFA(AEA-svGPFA), which has the same generative model as svGPFA (Figure 1a) but replacing g ( · ) withsome general non-linear mapping modelled by a neural network.
The complete inference procedurefor AEA-svGPFA described above is graphically illustrated in Figure 1c.
Instead of directly outputting a recognition potential on the inducing points, the variational distributionover U in AEA-svGPFA is computed implicitly given the recognition potential on h and the linearGaussian relationship between the two variables.
Hence it is possible to change the inducing locationsand the number of inducing points for inference over varying timescales upon training complete,given by an addition structured svGP inference step (Eq. 12).
This is particularly useful in situationswhen inference over testing dataset of varying length upon training is required, and when the latentprocesses underlying a particularly long time-span are to be inferred.
The latter situation is moreimportant since in practice, the number of inducing points is usually set small relative to the entire setof timepoints for scalability.
For instance, when we wish to examine the inferred latent dynamics overthe entire time-span, it is very likely we would observe local inconsistencies and abrupt variationsthat are not reflected in the data, and inference with more inducing points could alleviate such issue.
Related Works
Structured Deep Generative Models.
Ample literature has studied the inclusion of latent graphicalstructures into DGM [22, 10, 18, 29, 27, 37].
The prominent work by Johnson et al.
[18] proposed
SVAE, where the recognition network outputs fully factorised conjugate factor potentials withrespect to the structured latent prior, and the variational distribution is computed through partialoptimisation of the surrogate free energy given the recognition potentials.
The model parametersare jointly optimised by maximising the free energy.
In another closely related work, Lin et al.[29] proposed structured inference network, which approximates the variational distribution as thecombination of the recognition potential and a separate structured latent distribution independent ofthe prior, the independent structured latent distribution can have non-conjugate factors to improve theexpressiveness of the model.
Our model can be interpreted as a strict generalisation of the SVAEframework, where we allow arbitrary structured amortised potential, broadening the flexibility ofvariational approximation given SVAE.
However, the current framework cannot handle models withlatent priors containing non-conjugate factors, which we leave as future work.
In practice, we extendthe recognition potential to contain factor potentials over multiple latent variables whereas Johnsonet al.
[18], Lin et al.
[29] have only exhibited success with mean-field recognition potentials.svGPFA and Extensions.
High-dimensional time-series data, such as population neural firingpattern, usually lie on a low-dimensional manifold that can be characterised by the small numberof dynamical covariates that fully explain the temporal correlation exhibited in the observations.
Extraction of the latent covariates is particularly essential in analysing neural data, where the set ofexternal measurable variables and internal mental states are potentially entangled and go through ahighly non-linear process to influence the neuronal firing.
Hence latent variable models such as GPFAare usually adopted for extracting the relevant latent covariates that contains non-trivial temporalcorrelations [50, 8, 12].
Recent advances have proposed sparse variational inference and methodsfor dealing with non-conjugate likelihood function to promote the scalability and applicability of
GPFA models [1, 12, 19].
There also exists prior works of combining DGM (e.g., VAEs) with GPlatents [7, 6, 2].
However, AEA-svGPFA differs from all prior works as it is able to combine svGPFAwith amortised inference to allow scalable variational inference of a “free-form” GP posterior (i.e.,with full-covariance).
The closest model to AEA-svGPFA is the SGP-VAE [2], which similarlyincorporate the latent temporal dynamics into posterior inference through the SVAE formulation, andassumes the same generative model as AEA-svGPFA (consider setting C
=
I and d = 0 in Eq.
3),but with different inference procedure from AEA-svGPFA as discussed in Section 3.2.
However,since SGP-VAE directly imposes a fully factorised recognition distribution on the inducing points, iti) cannot capture the posterior correlations between the latent processes; ii) does not allow additionalsvGP inference step with changed inducing points.
Results
Here we empirically examine the performance of AEA-svGPFA and compare to relevant baselineson various synthetic and real datasets.
Empirical results on the instantiations of the AEA frameworkwith other latent variable models can be found in Appendix E.
5.1 Synthetic and EEG Dataset
We firstly empirically evaluate whether AEA-svGPFA outperforms methods that do not capturethe intra-latent posterior correlation.
We use a synthetic dataset and a small-scale EEG dataset forcomparison.
Note that we assume Squared Exponential (SE) kernel, κ ( x, x ′ ) =
λ exp( − || x − x ′
||
2 τ) ,unless stated otherwise, where λ and τ are the scale and length-scale parameters, respectively.
Baselines Our main baseline model is SGP-VAE [2], which we introduce and provide discussionsrelative to AEA-svGPFA in Section 3.2 and 4, respectively.
For completeness, we additional compareagainst a vanilla VAE with fully factorised Gaussian variational distribution [21, 42], and SVAE withlinear dynamical system (SVAE-LDS) latent prior [18].
All models are implemented with the samerecognition and generative network architectures (see Appendix F for implementation details).
Synthetic Dataset We generate from the GPFA generative model (Eq. 3), with g ( · )
= σ (Φ( · )) , whereΦ( · ) represents the functional mapping through a fixed, randomly-initialised 2-layer MLP with ReLUhidden non-linearity, and σ is the sigmoid function.
EEG Dataset
We follow the experimental procedure as described in Requeima et al.
[41], andconsider an EEG measurement dataset spanning 1 second at 256 Hz frequency when images areshown to the patients [51].
Each datapoint consists of the voltage readings from 7 electrodespositioned on the patient’s scalp.
We assume all data are observed, contrary to the settings in [41, 2](we include results and further discussion on partial observability in appendix E.3).
[Table]
Given the evaluation comparison in Table
[Table]
we see that comparing to all selected baseline models,
AEA-svGPFA achieves stronger performance on both the synthetic and EEG dataset, with respect toboth the SMSE and NLL metric.
Hence the empirical results provide evidence for the advantages ofthe inclusion of structured amortised inference for capturing the additional posterior correlations thatdo not exist a priori .
Moreover, the computational complexity of AEA-svGPFA is on the same orderas SGP-VAE, see appendix E.3 for further analysis and discussion.
5.2 Population Neuronal Firing Data
It is well-known that the hippocampal CA1 place cells and medial Entorhinal Cortex (mEC) gridcells are spatially modulated [34, 32], and previous work has also established speed and directionalmodulation of CA1 and mEC neurons [31, 43, 23].
Moreover, there are often non-trivial correlationsbetween the behavioural covariates.
For instance, one might reasonably expect correlation betweenspeed and spatial location as an animal runs fastest in the central area of the environment and slowernear the boundaries.
Neither svGPFA or SGP-VAE can capture the posterior correlations betweenlatent variables given their mean-field assumptions, whereas AEA-svGPFA, given its free-formvariational approximation with full covariance structure, is able to capture arbitrary correlationbetween the latents whilst utilising the same mutually independent GP priors, and hence is an idealmodel for unsupervised learning of the latent structure that correctly reflect the behavioural covariatesunderlying the neural firing pattern.
We use the single-cell spiking data of neurons in the hippocampal CA1 and mEC regions of ratsrecorded during exploration in a Z-shaped track reported by ´Olafsd´ottir et al.
The datasetcontains 28 experiment sessions, each spanning 10 minutes.
Exemplary neural firing patterns of thepopulation are shown in Figure 11.
Note that due to the nature of the task, we can interpret the spatiallocation of the rat being standard and “unfolded” distances along the track (“unfolded” meaning the sign of the distance value is dependent on the direction of movement).
Other considered behaviouralcorrelates are speed, direction of movement (inbound v.s. outbound), and head-direction.
We again use the same generative model in Eq. 3, but with Poisson observation likelihood.
[Equation]
where λ NN ( · ) is a neural-network mapping from the GPFA features, h , to the rate of Poissondistributed firing counts for each neuron over a bin-size of 100 ms .
We apply AEA-svGPFA and SGP-VAE to the dataset, In Table 1 we show that AEA-svGPFAquantitatively outperforms SGP-VAE (and other baselines) with respect to SMSE of prediction.
Having trained the model, and given the firing pattern observations, ( x n , y n ) , we could immediatelycompute the variational posterior latent distribution over the latent processes.
[Equation]
where q ( U ) is computed as described in Section 3.2 (Eq. 12).
Note that due to computationalconstraints and the stochastic mini-batch nature of the training process, the number of inducing pointsis much smaller than the total number of observations ( 64 v.s. 6000 in this case), hence causingtemporal chunking artifacts in the posterior estimation.
As discussed in Section 3.2, we note that
AEA-svGPFA can resolve the temporal chunking with an additional svGP inference step over a muchlarger time window (i.e., increasing the number of inducing points) with the weights of the recognitionnetwork fixed at the final state of the training process (see further discussion in Section E.3.2).
We perform the two-dimensional Canonical Correlation Analysis (CCA) on the variational posteriormeans and the behavioural correlates over the entire trajectory [15].
Here we choose one session forin-depth investigation and mostly show the qualitative predictions and comparison, more qualitativeand quantitative results across sessions can be found in Appendix E.3.2.
The heatmap of the correlation coefficients between the canonical correlates of the posterior means( CCX { 1 , 2 } ) and each of the behavioural correlates is shown in Figure 2a.
We observe an overallpattern of high-correlations, indicating the low-dimensional manifold parameterised by the conjunctive set of behavioural correlates can be accurately captured by the latent variables learned with
AEA-svGPFA solely from neural spikes.
In order to verify that if the learned latent structure indeed contain decodable information withrespect to the behavioural covariates, we extract the direction-modulated neurons predicted by thetrained model (details in appendix E.3.2).
In Figure 2b, we compare the direction-modulationof the such model-predicted neurons against that of the rest of the neurons, where the directionmodulation is defined as the correlation between raw single-cell spike counts and the directionvalues.
We observe that the model-predicted direction-modulated neurons exhibits significantlygreater direction-modulation than the other neurons (p-value= 5 .
56 × 10 − 4 ).
For the neuron with thestrongest predicted direction-modulation, we show the firing profiles of the neuron when the animal istravelling along both directions, Figure 2b shows that the model-predicted neuron has different firingfields corresponding to different travel direction, hence indeed exhibiting strong direction-modulation.
Similar comparison holds for spatial and speed modulations as well, indicating AEA-svGPFA is ableto learn the latent structure that coheres well with the ground-truth behavioural covariates, givensolely the population spike trains, even for neurons that exhibit conjunctive coding mechanism.
We plot CCX { 1 , 2 } against the ( x, y ) location in Figure 2d, where we observe clear clusteredlatent structure with respect to the direction of movement, indicating the latent dimensions learnedby AEA-svGPFA exhibit disentanglement of direction from spatial location given simple linearprojection, despite the conjunctive coding mechanism exhibited in the CA1 neurons.
By comparingto the same plot with SGP-VAE, we observe from Figure 13b that the spatial and direction variablesare more intertwined, hence qualitatively indicating the superiority of AEA-svGPFA.
Figure 2e shows the temporal evolution of the latent CCs for AEA-svGPFA and SGP-VAE, as well asthat of the distance along the track .
We observe that CCX 1 of AEA-svGPFA consistently tracks thedynamics of the distance, whereas neither CCX 1 or CCX 2 of SGP-VAE could accurately representsthe distance.
Moreover, given the additional svGP inference step over the entire trajectory (with
[Figure]
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 -9.6e
[Figure]
[Figure] [Figure]
[Figure]
CCX1CCX2posy (f) changed inducing points, see Section 3.2), CCX { 1 , 2 } of AEA-svGPFA exhibits much less temporalchunking than that of SGP-VAE.
Discussion
In this paper, we proposed structured amortised inference, which can be viewed as a strict generalisation of the SVAE framework [18], where our model allows structured amortised factor potentialsto capture the posterior correlations induced by the “explaining away” effect.
The AEA frameworkyields a strictly tighter bound on the log-likelihood than the SVAE free energy objective.
We introduce AEA-svGPFA, a combination of the proposed structured amortised inference frameworkand VAEs with GPFA latents, where scalable free-form variational Gaussian approximation can beachieved.
AEA-svGPFA supports a fully non-linear, non-conjugate extension to standard GPFA [50,1, 12], which is particularly well-suited to settings where posterior inference must be scaled toadditional data.
We show that AEA-svGPFA outperforms the baseline on all presented tasks, bothquantitatively and qualitatively.
Studying population spiking data of hippocampal neurons, we showthat AEA-svGPFA is able to identify latent signals that exhibit strong correlation with the behaviouralcovariates that are well-known to be conjunctively encoded by the recorded neurons, but withoutdirect access to these covariates.
While we have focused on temporal correlation exhibited in timeseries data, we note that the sparse amortised variational approximation presented here can also bestraightforwardly applied to other GP-latent models, such as GP-LVM [24, 49].
We note that here wehave only looked at single-cell spiking data, which could in principle be combined with the latentprocesses found by applying our model to EEG or calcium imaging datasets recorded during the samesession, possibly with different likelihoods functions.
Such joint analysis of latent structures foundby analysing different datasets under the same behaviour could potentially discover new findingswith respect to, e.g., replay, theta sweeps, etc.