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 assuming thatdistinct observed variables are derived from different latent nodes.
Here, we extend the output of amortised variational inference to incorporate structured factorsover multiple variables, able to capture the observation-induced posterior dependence between latents that results from “explaining away” and thus allow complexobservations to depend on multiple nodes of a structured graph.
We show thatappropriately parameterised factors can be combined efficiently with variationalmessage passing in elaborate graphical structures.
We demonstrate the frameworkusing synthetic data with a known generative process, We then fit a structuredmodel to high-dimensional neural spiking time-series from the hippocampus offreely moving rodents, demonstrating model identifies latent signals that correlatewith 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 [20, 41,12, 13, 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; 17] 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” [37].
Our goal here is to extend the amortised inferencesetting to capture such observation-induced dependence in the DGM setting.
zzθ pgm zzzzzzx nθ x n... x nD ψ ij ( z i ,z j )
N ψ i ( z i ) (a) 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 (b) We develop the extension in two contexts. First we consider a simple discrete latent factor model with both apriori and a posteriori coupling, evaluating the utility of structured amortised inference in this case on synthetic data from a known generativeprocess.
Second, we extend Gaussian Process Factor Analysis (GPFA)models that are widely used in neuroscience, to incorporate deep networkbased nonlinear mappings from thelatent Gaussian process values to observations.
The corresponding likelihood functions induce posterior “explaining-away” correlations between the processes.
We develop“amortised explaining-away sparse variational GPFA” (AEA-svGPFA) to achieve scalable amortisedinference of the full-covariance posterior distribution within the sparse inducing-point GP framework.
We apply this model to spike-count data collected from a population of neurons in the hippocampalcomplex, and show that the unsupervised approach automatically captures dimensions underlyingneural firing that reflect relevant behavioural correlates.
Background
2.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 a single θ for simplicity, but hyperpriors on θ can be incorporated in the contextof Variational Bayes (VB).
The posterior distribution, p ( z | y )
∝ p ( z | θ ) p ( y | z, γ ) , is often analytically 62 intractable.
Hence it is common practice to seek a tractable variational approximation to the trueposterior, q ( z ) , which is learned by minimising the KL-divergence with respect to the true posteriordistributions, and is equivalent to the maximisation of the variational free energy [46]:
[Equation]
The Variational Autoencoder (VAE) extends VI by modelling the parameters of the variationalapproximation using a recognition network that takes as input the observed values, and is able tooptimise both generative and recognition network parameters using stochastic optimisation [20, 41].
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 restrictiveto capture the complexity of the true posterior [43, 15].
Intuitively, increasing the flexibility of theclass of variational distributions leads to a tighter free energy lower bound, and often leads to betterempirical results in general [4].
A number of approaches have been proposed for improving the expressiveness of variational approximation [3, 45, 28, 26].
Here we specifically look at structured VAE [SVAE; 17].
Unlike standard VAEs, SVAE assumes the generative prior distribution, p ( z | θ ) , to be specified by a structured probabilistic graphical model (PGM).
Moreover, the amortised inference network outputs factors, r ( z | y, ϕ ) ,that approximate the generative likelihood function rather than the variational distribution itself.
These are combined with the structured prior distribution to obtain the variational posterior.
[Equation]
Johnson et al.
[17] 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. 1 is able tocarry the dependency structure in the prior distribution to the variational posterior distribution.
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 [50].
Standard GPFA assumes the following generative model.
[Equation]
) where c nk and d n are parameters defining the affine mapping from the latent space to the neural space,and g ( · ) is some smooth function.
We assume m kθ ( · ) =
unless stated otherwise.
The temporalcorrelation in the high-dimensional neural space are fully captured by the latent temporal correlationsmodelled by the latent processes.
The PGM describing the GPFA generative model is shown in 1b.
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 [39].
To this end, an extension of the sparse variational approximation for GPs is applied to the GPFAframework, leading to the svGPFA model [44, 1].
Details of svGPFA formulations are discussed inour discussion of AEA-svGPFA in Section 3.3.
Structured amortised inference and explaining away3.1 Structured amortised inference with SVAE
We assume a structured prior distribution, specified by the factor potentials, p ( z | θ ) = 1 Z (cid:81) c ∈ C ψ c ( z c ) .
Unlike the SVAE approach of enforcing a fully factorised form, we assume the recognition networkoutputs structured factor potentials over z , r ( z | y, ϕ ) 1 Z r (cid:81) c ∈ C r ξ c ( z c ) , approximating the likelihoodfunction, p ( y | z, θ ) , where C r is the set of pre-specified factors of the amortised potentials and Z r isthe normalisation constant. Then the variational distribution is proportional to the product of the priordistribution and the amortised factor potentials (given partial optimisation, see Johnson et al.
[Equation]
We assume conjugate factors unless otherwise stated.
Hence the analytical (approximate) form ofq AEA ( z | x, ϕ, θ ) can be computed with (loopy, if necessary) belief propagation.
In the most generalform, we could let r ( z | y, ϕ ) be a single joint factor potential over the complete set of latent variables.
Similar to the SVAE formulation, the free energy objective takes the following form.
[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 C).
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 how our framework can be instantiated with different latent PGM structure, theempirical evaluations of the discussed methods can be found in Section 5.
3.2 Tree-Structured Latent PGM
We instantiate our method with discrete tree-structured PGM (Figure 1a), which enables exactinference can be achieved given two iterations of message passing. A general tree-structured PGM zzzzzzzzx nx n... x nD ψ ij ( z i ,z j )
N ξ i ( z i | x n ) ξ ij ( z i ,z j | x n ) ψ i ( z i ) (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) takes the following density function. [Equation] Consider an arbitrary non-linear generative likelihood function p ( y | z ) . Then the optimal approximation to p ( y | z ) as a density function of z will in general contain non-singleton factors of z .
Weprovide a analytical example illustrating this in appendix A.1.
The optimal variational distribution should have the same PGM structure as the posterior, which oftencontains joint factors over many (possibly all) latent variables.
However, as the number of statesgrows exponentially with the number of latent variables, in practice it is usually not scalable to haveamortised inference output a joint factor that approximates the exact posterior joint factor potentials. We hence seek a scalable alternative by enforcing the amortised factor potential to be of the samestructure as the prior PGM; in this case, a tree-structured distribution with the same set of nodes andpairwise connections (Figure 2a). [Equation] Given these tree-structured amortised potentials, we can carry out exact inference with messagepassing to compute the variational distribution, then update the model parameters by maximising thefree energy objective (Eq. 4) through stochastic optimisation.
Further details of the tree-structuredamortised inference can be found in appendix A.
3.3 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 underliethe 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 popularsuch model is GPFA [50].
In order to alleviate the scalability issue, sparse variational approximationinference for GPFAs with inducing points were proposed [44, 1, 10].
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. 2) 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 .
It is common practice to use the mean-field variational approximation of the joint posterior distributionof the latent process and the inducing points [44, 1], q ( f 1: K , u 1: K ) =
(cid:81)
Kk =1 p ( f k u k ) q ( u k ) , where q ( u k ) =
N ( m k , S k ) .
The free energy objective for the svGPFA model given the mean-fieldvariational approximation takes the following expression.
[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 150 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.
Previous work has extended GPFA to deal with non-conjugate observation likelihoods and point-wisenonlinearities [1, 10, 18].
However, adding a complex non-linear likelihood function introducesadditional correlation between the latent processes given observations through “explaining away”.
Earlier work did not capture the cross-latent posterior correlations due to the mean-field variationalapproximation adopted.
Here we describe a novel inference procedure that represents the fullcovariance structure of the posterior distribution of the inducing points over all latent processes in ascalable manner.
Given the general expression of the GPFA generative model and the sparse approximation withinducing points, we have that for n = 1 , . . . , N .
[Equation]
[Equation]
where
[Equation]
 with U =
[ u 1 , . . .
, u K ] being the collection of all inducing points for all latents, and K n =diag [ κ kθ ( x n , x n )] .
Assume that the amortised inference outputs a fully factorised Gaussian potentialon h n , r ( h n | y n , ϕ )
=
N ( µ hn , Ψ hn ) . By the generative model described in Eq. 3, we can derive thevariational distribution on U given the relationship p ( h | U ) .
[Equation]
n N ( µ hn , Ψ hn ) , with
[Equation]
Hence giving
[Equation]
Then the complete (reparameterised) Monte Carlo estimate of the free energy objective is
[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.
8), 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 1b) but replacing g ( · ) withsome general non-linear mapping modelled by a neural network.
The complete inference procedurefor AEA-svGPFA described above is different from that of svGPFA (and that of other DGMs with
GP latents [2]; Figure 2b), and is graphically illustrated in Figure 2c.
Related Works
Structured Deep Generative Models.
Ample literature has studied the inclusion of latent graphicalstructures into DGM [22, 8, 17, 29, 27].
The prominent work by Johnson et al.
[17] 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.
[17], 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 firing pattern,usually lie on a low-dimensional manifold that can be characterised by the small number of dynamicalcovariates that fully explain the temporal correlation exhibited in the observations.
Extraction of thelatent covariates is particularly essential in analysing neural data, where the set of external measurablevariables and internal mental states are potentially entangled and go through a highly non-linearprocess to influence the neuronal firing.
Hence latent variable models such as GPFA are usuallyadopted for extracting the relevant latent covariates that contains non-trivial temporal correlations [50,7, 10].
Recent advances have proposed sparse variational inference (sVI) and methods for dealing withnon-conjugate likelihood function to promote the scalability and applicability of GPFA models [1, 10,19 ].
There also exists prior works of combining DGM (e.g., VAEs) with GP latents [6, 5, 2].
However, AEA-svGPFA differs from all prior works as it is able to combine svGPFA with amortised inferenceto allow scalable variational inference of a “free-form” GP posterior (i.e., with full-covariance).
[Figure]
Results
We first evaluate tree-structured amortisedinference on a synthetic dataset. We consider the “Bar” dataset of horizontal andvertical bars in D × D square-grids. [Equation]
where pa ( z ) denote the parent node of node z , σ ( · ) = 1 1+exp( · ) is the sigmoid function, and B ( p )denotes the Bernoulli distribution with rate p .
Each z i denotes the activation of one bar in thesquare-grid (e.g., z 1 = 1 indicates the presence of the first horizontal bar).
The generative parameters
W and b are defined as [Equation] where { ω i } D 2 i =1 are the temperature parameters that control the degree of randomness in the data 221 generation. Exemplary samples from the “Bar” dataset if shown in Figure 3a ( D = 8) .
We apply our approach, SVAE with tree-structured amortised potentials (TreeSVAE), standard VAE,and standard SVAE, all with 16 -dimensional latents (corresponding to the activation of 16 possible bars), on the “Bar” dataset.
The training curves of the free-energy objective is shown in Figure 3b.
We observe that TreeSVAE consistently outperforms standard SVAE and VAE through the trainingprocess in terms of the free energy objective.
Further details on the generation quality and dealingwith mismatches in the PGM structures in can found in appendix D.1.
5.2 AEA-svGPFA5.2.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.
Synthetic Dataset We generate from the GPFA generative model (Eq. 2), 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.
[40], 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 [40, 2](we include results and further discussion on partial observability in appendix D.2).synthetic dataset
[Table]
Our main baseline model is SGPVAE [2], it can be interpreted as a version of the SVAE modelwith independent GP latents. SGPVAE has the same generative model as the svGPFA generativemodel (Figure 1b), whilst having identity affine transformation between f 1: K and h (i.e., C = Iand d = 0 ), and g ( · ) is a general non-linear function parameterised by a neural network.
However,the non-trivial affine transformation in AEA-svGPFA (and svGPFA) can be easily modelled by theneural network likelihood function, hence the main difference between AEA-svGPFA and SGPVAElie in the structured amortised inference step (Section 3.3).
Given the evaluation comparison in
Table 4, we see that AEA-svGPFA achieves stronger performance on both the synthetic and EEGdataset, with respect to both the SMSE and NLL metric.
Hence, similar to the empirical findings withtree-structured amortised inference, we found that the inclusion of structured amortised inferencealways improve the performance.
Moreover, the computational complexity of AEA-svGPFA is onthe same order as SGPVAE, see appendix D.2 for further analysis and discussion.
5.2.2 Population Neuronal Firing Data
It is well-known that the hippocampal CA1 place cells and mEC grid cells are spatially modulated [34,33], and previous work has also established speed and directional modulation of CA1 and mECneurons [31, 42, 23].
Moreover, there are often non-trivial correlations between the behaviouralcovariates.
For instance, one might reasonably expect correlation between speed and spatial locationis an animal runs fastest in the central area of the environment and slower near the boundaries.
Neither svGPFA or SGPVAE can capture the posterior correlations between latent variables giventheir mean-field assumptions, whereas AEA-svGPFA, given its free-form variational approximationwith full covariance structure, is able to capture arbitrary correlation between the latents, and hence isan ideal model for unsupervised learning of the latent structure that correctly reflect the behaviouralcovariates underlying the neural firing pattern.
We use the single-cell spiking data of neurons in the hippocampal CA1 and medial entorhinal cortex(mEC) regions of rats recorded during exploration in a Z-shaped track reported by ´Olafsd´ottir et al.
The dataset contains 28 experiment sessions, each spanning 10 minutes.
Exemplary neuralfiring patterns of the population are shown in Figure 14.
Note that due to the nature of the task, wecan interpret the spatial location 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).
Otherconsidered behavioural correlates are speed, direction of movement (inbound v.s. outbound), andhead-direction.
We again use the same generative model in Eq. 2, but with Poisson observation likelihood.
[Equation]
13) 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 firstly apply AEA-svGPFA and SGPVAE to the dataset, and quantitative comparison indicatesthat AEA-svGPFA outperforms SGPVAE with respect to SMSE of prediction (Table 1).
Having trained the model, and given the firing pattern observations, y , we could immediately computethe variational posterior latent distribution over the latent processes.
[Equation]
[Equation] where q ( U ) is computed as described in Eq. 9.
Note that due to computational constraints andthe stochastic mini-batch nature of the training process, the number of inducing points is muchsmaller than the total number of observations ( 64 v.s. 6000 in this case), hence causing temporalchunking artifacts in the posterior estimation.
AEA-svGPFA can resolve the temporal chunkingwith an additional sVI step over a much larger time window (i.e., increasing the number of inducingpoints) with the weights of the recognition network fixed at the final state of the training process(see further discussion in Section D.2.2).
We note that SGPVAE is unable to remove the chunkingartifacts since it learns a recognition factor directly onto the inducing points, hence no further sVIsteps can be performed with respect to a larger set of inducing locations.
We perform the two-dimensional Canonical Correlation Analysis (CCA) on the variational posteriormeans and the behavioural correlates over the entire trajectory [14].
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 D.2.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 4a.
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 D.2.2).
In Figure 4b, 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 4b 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 4d, 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 SGPVAE, we observe from Figure 15 that the spatial and direction variables aremore intertwined, hence qualitatively indicating the superiority of AEA-svGPFA.
[Figure]
a
[Figure]
Figure 4e shows the temporal evolution of the latent CCs for AEA-svGPFA and SGPVAE 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 SGPVAE could accurately representsthe distance.
Moreover, given the additional sVI step with larger time windows, CCX { 1 , 2 } of AEA-svGPFA exhibits much less temporal chunking than that of SGPVAE.
Discussion
In this paper we proposed structured amortised inference, which can be viewed as a strict generalisation of the SVAE framework [17], where we allow structured amortised factor potentials to capture theposterior correlations induced by the “explaining away” effect.
We firstly instantiate our frameworkwith tree-structured latent PGM and demonstrate that the resulting model consistently yields higherfree energy (lower negative free, Figure 5b) and better generation quality (appendix D.1).
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.
The inclusion of non-linear likelihoods both expand the expressiveness of the model andinduces interesting posterior correlations that could inspire new findings.
We show that AEA-svGPFAoutperforms the baseline on all presented tasks, both quantitatively and qualitatively.
By applyingto population spiking data of hippocampal neurons, we show that AEA-svGPFA is able to identifylatent signals that exhibit strong correlation with the behavioural covariates that are well-known tobe conjunctively encoded by the recorded neurons.
While we have focused on temporal correlationexhibited in time-series data, we note that the sparse amortised variational approximation presentedhere can also be straightforwardly applied to other GP-latent models, such as GP-LVM and KISSGP [24, 49, 47].
We note that here we have only looked at single-cell spiking data, which could inprinciple be combined with the latent processes found by applying our model to EEG or calciumimaging datasets recorded during the same session, possibly with different likelihoods functions.
Such joint analysis of latent structures found by analysing different datasets under the same behaviourcould 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., typeof 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
or 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 amount
spent on participant compensation?
[N/A]