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/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 frameworkwith one element of such inference, combining a structured prior distribution with fully factoredmean-field (and functionally conjugate) amortised likelihood potentials.
However many generativemodels of interest involve observations that depend on more than one latent variable.
This inducesjoint potentials in the likelihood, coupling latents 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 various synthetic and real dataset, including a spikecount dataset collected from a population of neurons in the hippocampal complex, and show thatthe unsupervised approach automatically captures dimensions underlying neural firing that reflectrelevant behavioural correlates.
Background2.1 Variational inference and structured variational autoencoders
We consider the variational inference (VI) problem of approximating the posterior distribution of the latent 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 in thecontext of Variational Bayes (VB).
The posterior distribution, p ( z | y )
∝ p ( z | θ ) p ( y | z, γ ) , is often 55 analytically intractable.
Hence it is common practice to seek a tractable variational approximation tothe true posterior, q ( z ) , which is learned by minimising the KL-divergence with respect to the trueposterior distributions, and is equivalent to the maximisation of the variational free energy [47]:
[Equation]
The Variational Autoencoder (VAE) extends VI by modelling the variational parameters using arecognition network that takes as input the observed values, and is able to jointly train both generativeand recognition network parameters with stochastic optimisation of the free energy objective [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 factor potentials(that are conjugate with p ( z | θ ) ).
Comparing to the mean-field VI in 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 widely used model in neural data analysis for inferringthe dynamic latent structure underlying high-dimensional population spike trains [51].
Standard
GPFA assumes the following generative model.
[Equation]
where c jk and d n are parameters defining the affine mapping from the latent space to the neuralspace (or a subspace of the neural space in more general likelihood), 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, 51].
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
Non-linear generative likelihood induces additional latent dependency structure that does not exist apriori (we provide a simple example in Appendix B illustrating such effect).
Despite introducing thestructured prior distribution, the mean-field recognition potential assumption of SVAE only injectssingleton factor potential to the prior distribution, prohibiting accurate inference over the latents withchanged graphical structure (relative to the prior PGM). Instead of enforcing a fully factorised form, we assume the recognition network outputs structuredfactor potentials over z , providing amortised estimates of the additional dependencies induced 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 antuencoding 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) the population neuronal spiking trains, and relate these to behaviour or environmental variables so asto gain understanding of the functional and cognitive aspects of neural computations.
One popularmodel is GPFA [51].
In order to alleviate the scalability issue, sparse variational approximationinference for GPFAs with inducing points were proposed [45, 1, 12].
For k = 1 , . . .
, K , the inducingpoints u k represents the function values of the latent process f k at the M k inducing locations, z k .
We then augment 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 137 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. However, naive application ofthe amortised mean-field VI, q ( U ) = N ( U | m ( y ; ϕ ) , Ψ ( y ; ϕ ) , is inappropriate for GP models sincethe strong a priori temporal correlation will be carried into the posterior distribution. In order to preserve the within-process latent correlations, Ashman et al.
[2] introduced the SGP-VAE, wherethe variational distribution is computed by combining the prior GP distribution with an amortisedapproximation of the likelihood function given a set of observation pair, { ( x b , y b ) } Bb =1 . [Equation] where [Equation]
[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, despite that 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]
[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 [51, 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 [52].
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 potentials tocapture the posterior correlations induced by the “explaining away” effect.
We theoretically provethat the AEA framework yields an upper bound on 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[51, 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.
By applying to population spiking data of hippocampal neurons,we show that AEA-svGPFA is able to identify latent signals that exhibit strong correlation with thebehavioural covariates that are well-known to be conjunctively encoded by the recorded neurons.
While we have focused on temporal correlation exhibited in time-series data, we note that the sparseamortised variational approximation presented here can also be straightforwardly applied to other GP-latent models, such as GP-LVM and KISS-GP [24, 50, 48].
We note that here we have only lookedat single-cell spiking data, which could in principle be combined with the latent processes found byapplying our model to EEG or calcium imaging datasets recorded during the same session, possiblywith different likelihoods functions.
Such joint analysis of latent structures found by analysingdifferent datasets under the same behaviour could potentially discover new findings with respect to,e.g., replay, theta sweeps, etc.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A] If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]