Year: 2021

Abstract:
Understanding how neural dynamics give rise to behaviour is one of the most fundamental questions in systems neuroscience. To achieve this, a common approach is to record neural populations in behaving animals, and model these data as emanating from a latent dynamical system whose state trajectories can then be related back to behavioural observations via some form of decoding. As recordings are typically performed in localized circuits that form only a part of the wider implicated network, it is important to simultaneously learn the local dynamics and infer any unobserved external input that might drive them. Here, we introduce iLQR-VAE, a novel control-based approach to variational inference in nonlinear dynamical systems, capable of learning both latent dynamics, initial conditions, and ongoing external inputs. As in recent deep learning approaches, our method is based on an input-driven sequential variational autoencoder (VAE). The main novelty lies in the use of the powerful iterative linear quadratic regulator algorithm (iLQR) in the recognition model. Optimization of the standard evidence lower-bound requires differentiating through iLQR solutions, which is made possible by recent advances in differentiable control. Importantly, the recognition model can easily be tied to the generative model, greatly reducing the number of free parameters and ensuring high-quality inference throughout the course of learning. Moreover, iLQR is a flexible inference engine that can be used to perform amortized inference on heterogeneous trials of varying lengths. This allows for instance to evaluate the model on a single long trial after training on smaller chunks. We demonstrate the effectiveness of iLQR-VAE on a range of synthetic systems, with autonomous as well as input-driven dynamics. We further apply it to neural and behavioural recordings in non-human primates performing two different reaching tasks, and show that iLQR-VAE yields high-quality kinematic reconstructions from the neural data.

heading: 1 INTRODUCTION
text: The mammalian brain is a complex, high-dimensional system, containing billions of neurons whose coordinated dynamics ultimately drives behaviour. Identifying and interpreting these dynamics is the focus of a large body of neuroscience research, which is being facilitated by the advent of new experimental techniques that allow large-scale recordings of neural populations (Jun et al., 2017; Stosiek et al., 2003). A range of methods have been developed for learning dynamics from data (Buesing et al., 2012; Gao et al., 2016; Duncker et al., 2019; Archer et al., 2015; Hernandez et al., 2018; She and Wu, 2020; Kim et al., 2021; Nguyen et al., 2020). These methods all specify a generative model in the form of a flexible latent dynamical system driven by process noise, coupled with an appropriate observation model.
However, neural recordings are typically only made in a small selection of brain regions, leaving many areas unobserved which might provide relevant task-related input to the recorded one(s). Yet, the aforementioned methods perform Bayesian inference of state trajectories directly, and therefore do not support inference of external input (which they effectively treat as process noise and marginalize out). Indeed, simultaneous learning of latent dynamics and inference of unobserved
control inputs is a challenging problem that involves teasing apart momentary variations in the data that can be attributed to the system’s internal transition function, and those that need to be explained by unobserved inputs. This distinction can be achieved by introducing external control in the form of abrupt changes in the latent state transition function, and inferring these switching events (Ghahramani and Hinton, 2000; Linderman et al., 2017). More recently, Pandarinath et al. (2018) introduced LFADS, a sequential variational autoencoder (VAE) that performs inference at the level of external inputs as well as initial latent states. The inferred inputs were shown to be congruent with taskinduced perturbations in various reaching tasks in primates (Pandarinath et al., 2018; Keshtkaran and Pandarinath, 2019).
Here, we introduce iLQR-VAE, a new method for learning input-driven latent dynamics from data. As in LFADS, we use an input-driven sequential VAE to encode observations into a set of initial conditions and external inputs driving an RNN generator. However, while LFADS uses a separate, bidirectional RNN as the encoder, here we substitute the inference network with an optimizationbased recognition model that relies on the powerful iterative linear quadratic regulator algorithm (iLQR, Li and Todorov, 2004). iLQR solves an optimization problem that finds a mode of the exact posterior over inputs for the current setting of generative parameters. This ensures that the encoder (mean) remains optimal for every update of the decoder, thus reducing the amortization gap (Cremer et al., 2018). Moreover, having the recognition model be implicitly defined by the generative model stabilizes training, prevents posterior collapse (thus circumventing the need for tricks such as KL warmup), and greatly reduces the number of (hyper-)parameters.
While iLQR-VAE could find applications in many fields as a general approach to learning stochastic nonlinear dynamical systems, here we focus on neuroscience case studies. We first demonstrate in a series of synthetic examples that iLQR-VAE can recover the true dynamics in both autonomous and input-driven systems. Next, we show state-of-the art performance on monkey M1 population recordings during two types of reaching tasks (O’Doherty et al., 2018; Churchland et al., 2010). In particular, we show that hand kinematics can be accurately decoded from inferred latent state trajectories, and that the inferred inputs are consistent with recently proposed theories of motor preparation.

heading: 2 METHOD
text: iLQR-VAE models a set of temporal observations, such as behavioural and/or neural recordings, through a shared input-driven nonlinear latent dynamical system (Figure S1). The input encapsulates both process noise (as in traditional latent dynamics models), initial inputs that set the initial condition of the dynamics, and any meaningful task-related control input. In this section, we describe the architecture of the generative model, and the control-based variational inference strategy used for training the model and making predictions.

heading: 2.1 GENERATIVE MODEL
text: We consider the following generative model:
latent state zt+1 = fθ(zt,ut, t) (1) observations ot|zt ∼ pθ(ot|zt) (2)
where ut ∈ Rm, zt ∈ Rn and ot ∈ Rno are the input, latent state and observations at time t, respectively. Here, observations may comprise either neural activity, behavioural variables, or both – the distinction will be made later where relevant. We use the notation θ to denote the set of all parameters of the generative model. We use u0 to set the initial condition z1 = fθ(0,u0, 0) of the network1. This way, the latent state trajectory of the network z(u) = {z1, . . . ,zT } is entirely determined by the input sequence u = {u0, . . . ,uT } and the state transition function fθ(·), according to Equation 1. For fθ(·), we use either standard linear or GRU-like RNN dynamics (see Appendix B
1Note that when m < n, u0 can only reach an m-dimensional subspace of initial conditions, which could be limiting. We can circumvent this problem by spreading u0 over multiple surrogate time bins before the start of the trial, i.e. introduce {u−n/m, . . . ,u−2,u−1,u0} together with an appropriate dependence of fθ on t ≤ 0 in Equation 1, such that each of these surrogate inputs target a different latent subspace with purely integrating (“sticking”) linear dynamics before t = 1.
for details). For the likelihoods, we use Gaussian or Poisson distributions with means given by linear or nonlinear readouts of the network state of the form ōt = h(Czt + b) (Appendix C).
We place a Gaussian prior over ut≤0. We then consider two alternative choices for the prior over ut>0. The first is a Gaussian prior
pθ(ut>0) = N (0,S2) (3)
with S = diag(s1, . . . , sm). In many settings however, we expect inputs to enter the system in a sparse manner. To explicitely model this, we introduce a second prior over u in the form of a heavy-tailed distribution constructed hierarchically by assuming that the ith input at time t > 0 is
uit = si it √ ν/αt (4)
where si > 0 is a scale factor, it ∼ N (0, 1) is independent across i and t, and αt ∼ χ2ν is a shared scale factor drawn from a chi-squared distribution with ν degrees of freedom. Thus, inputs are spatially and temporally independent a priori, such that any spatio-temporal structure in the observations will have to be explained by the coupled dynamics of the latent states. Moreover, the heavy-tailed nature of this prior allows for strong inputs when they are needed. Finally, the fact that the scale factor is shared across input dimensions means that inputs are either all weak or potentially all strong at the same time for all input channels, expressing the prior belief that inputs come as shared events.
This hierarchical construction induces a multivariate Student prior at each time step:
pθ(ut) = Γ [(ν +m)/2]
Γ [ν/2] (νπ)m/2|S|
[ 1 + 1
ν uTt S −2ut
]−(ν+m)/2 (5)
where S = diag(s1, . . . , sm). Note that both S and ν are parameters of the generative model, which we will learn. We discuss the relative advantages of the Student and Gaussian priors in Appendix J.

heading: 2.2 ILQR-VAE: A NOVEL CONTROL-BASED VARIATIONAL INFERENCE STRATEGY
text: To train the model, we optimize θ to maximize the log-likelihood of observing a collection of independent observation sequences O = {o(1), . . . ,o(K)}, or “trials”, given by:
log pθ(O) = K∑ k=1 log ∫ pθ(o (k)|z(u))pθ(u) du. (6)
As the integral is in general intractable, we resort to an amortized variational inference strategy by introducing a recognition model qφ(u|o(k)) to approximate the posterior pθ(u|o(k)). Following standard practice (Kingma and Welling, 2013; Rezende et al., 2014), we thus train the model by maximizing the evidence lower-bound (ELBO):
L(O, θ, φ) = ∑ k Eqφ(u|o(k)) [ log pθ(o (k)|u) + log pθ(u)− log qφ(u|o(k)) ]
(7)
= ∑ k Eqφ(u|o(k)) [ T∑ t=1 log pθ(o (k) t |zt) + log pθ(ut)− log qφ(ut|o(k)) ] (8)
≤ log pθ(O). (9)
with respect to both θ and φ.
Here, the main novelty is the use of an optimization-based recognition model. We reason that maximizing the exact log posterior, i.e. computing
u?(o(k)) = argmax u log pθ(u|o(k)) (10)
= argmax u [ T∑ t=1 log pθ(o (k) t |u) + log pθ(ut) ] (11)
subject to the generative dynamics of Equations 1 and 2, is a standard nonlinear control problem: log pθ(o (k) t |u) acts as a running cost penalizing momentary deviations between desired outputs ot and the actual outputs caused by a set of controls u, and log pθ(ut) acts as an energetic cost on those controls. Importantly, there exists a general purpose, efficient algorithm to solve such nonlinear control problems: iLQR (Li and Todorov, 2004; Appendix D). We thus propose to use a black-box iLQR solver to parameterize the mean of the recognition density qφ(u|o) for any o, and to model uncertainty separately using a multivariate Gaussian density common to all trials. Therefore, we parametrize the recognition model as follows:
qφ(u|o) = N (u;u?(o),Σs ⊗Σt) (12) with u?(o) = iLQRsolve(o, θ). (13)
where we use a separable posterior covariance (the Kronecker product of a spatial factor Σs and a temporal factor Σt).
To optimize the ELBO, we estimate the expectation in Equation 8 by drawing samples from qφ(u|o(k)) and using the reparameterization trick (Kingma et al., 2015) to obtain gradients. A major complication that would normally preclude the use of optimization-based recognition models is the need to differentiate through the mean of the posterior. In this case, this involves differentiating through an entire optimization process. Using automatic differentiation within the iLQR solver is in general impractically expensive memory-wise. However, recent advances in differentiable model predictive control enable implicit differentiation through iLQRsolve with a memory cost that does not depend on the number of iterations (Amos et al., 2018; Blondel et al., 2021; Appendix E).

heading: 2.3 COMPLEXITY AND IMPLEMENTATION
text: We optimize the ELBO using Adam (Kingma and Ba, 2014) with a decaying learning rate ∝ 1/ √ i where i is the iteration number. Averaging over data samples can be easily parallelized; we do this here using the MPI library and a local CPU cluster. In each iteration and for each data sample, obtaining the approximate posterior mean through iLQR is the main computational bottleneck, with a complexity ofO(T (n3 +n2no)). To help mitigate this cost, we find it useful to re-use the previously inferred control inputs to initialize each iLQRsolve.

heading: 3 EXPERIMENTS AND RESULTS

heading: 3.1 ILQR-VAE ENABLES FAST LEARNING OF DYNAMICS
text: Before demonstrating the method on a number of synthetic and real datasets involving ongoing external inputs, we begin with a simpler example meant to illustrate some of the main advantages of iLQR-VAE (Figure 1). We generated data from an autonomous (i.e. non-input-driven) linear dynamical system (n = 8, m = 3 where n is the number of latents and m the number of input channels) seeded with a random initial condition in each of 56 trials. The state zt was linearly decoded with added Gaussian noise to produce observation data, which we used to train a model in the same class.
At the beginning of learning, iLQR-VAE originally relies on large ongoing inputs that control the generator into producing outputs very similar to the observations in the data (Figure 1, red box, left), resulting in a rapidly decreasing loss. Subsequently, the amount of input required to fit the observations gradually decreases as the system learns the internal dynamics of the ground truth system. Eventually, the inferred control inputs become confined to the first time bin, i.e. they act as initial conditions for the now autonomous dynamics of the generative model. Thus, iLQR-VAE operates in a regime where the output of the generator explains the data well at all times, and learning consists in making the inputs more parsimonious. We note that this regime is facilitated here by our choice of generator dynamics, which we initialised to be very weak (i.e with a small spectral radius) initially and therefore easily controllable.
We contrast this with learning in a modified version of iLQR-VAE where we allowed u0 to vary freely (with a Gaussian prior of adjustable variance) but effectively fixed ut>0 to be 0. In other words, we constrained the dynamics of the generator to remain (near-)autonomous throughout learning (Figure 1, grey box, top). Although this incorporates important information about the ground
forced autonomous
free adaptive prior
−0.5 0
0.5 1
1.5 2
2.5 3
3.5 4
0 2 4 6 8 10 12 14
ne ga
tiv e
E LB
O
training iteration (×103)
ū|o
ōi |o
ū|o
ōi |o
Figure 1: Fast and robust learning in iLQR-VAE. Example run of iLQR-VAE on a synthetic dataset generated by an autonomous linear dynamical. iLQR-VAE with an adaptive prior over control inputs (red) uses large inputs to fit the observations initially, but gradually pushes those back into initial conditions as it acquires the ground truth autonomous dynamics. In contrast, iLQR-VAE with a rigid input prior imposing autonomous dynamics gets stuck in plateaus and learns considerably more slowly (black; see text for details). For each setting, insets show the three inferred inputs for a given test trial (top; posterior mean ū|o), and the posterior predictions for the first two corresponding outputs (ōi|o; black dots: ground truth; blue: posterior mean with 95% c.i.). We also compare learning curves with LFADS (yellow curve). We used the same generator architecture in all scenarios, and the learning rate was handtuned for each example.
truth generator (which is itself autonomous), counter-intuitively we found that it makes learning worse. At the beginning of training iLQR is unable to find initial conditions that would explain the data well, resulting in a much higher initial loss. The model then gets stuck in plateaus that are seemingly avoided by the free iLQR-VAE, leading to much slower convergence, as shown in Figure 1 (see multiple realizations in Figure S3).
On the same toy dataset, we also compare iLQR-VAE to LFADS (Pandarinath et al., 2018), keeping the generative model in the same model class (see Appendix G for details). We find that LFADS learns in a similar manner to iLQR-VAE (Figure 1), also progressively doing away with inputs.
In summary, we find that allowing inputs to enter the system at all times helps escape plateaus in the loss landscape, resulting in fewer iterations to reach convergence. While all methods run in similar wallclock time in this example, this will ultimately be implementation and data-dependent.

heading: 3.2 ILQR-VAE FOR NONLINEAR SYSTEM IDENTIFICATION
text: Next, we illustrate the method on an autonomous nonlinear dynamical system, the chaotic Lorenz attractor (Lorenz, 1963; Appendix L). This is a standard benchmark to evaluate system identification methods on nonlinear dynamics (Nguyen et al., 2020; Hernandez et al., 2018; Champion et al., 2019), and one typically considers the dynamics to be learned if the trained model can recreate the whole attractor structure.
Here, we show that iLQR-VAE can learn these complex nonlinear dynamics. Before training, the inferred inputs were large and explained the output observations by forcing the internal state of the generator into appropriate trajectories (Figure 2A, top). At the end of learning, the inputs were confined to initial conditions, with only weak ongoing corrections to the now stronger, near-autonomous dynamics of the generator. In Figure 2B we show that, conditioned on an initial bout of test data, the model perfectly predicts the rest of the trajectory. Moreover, starting from a random initial condition, the model can recreate the whole attractor structure (Figure 2C).
To quantitatively assess how well the dynamics have been learned, we computed the k-step coefficient of determination, R2k (Appendix L), as in Hernandez et al. (2018). This metric evaluates how well the model can predict the true state k steps into the future, starting from any state inferred along a test trajectory. Hernandez et al. reported R230 ≈ 1 but did not show results for larger k. For iLQR-VAE, R230 = 0.998 and the forward interpolation was still very high at 50 time steps, with R250 = 0.996.

heading: 3.3 INFERRING SPARSE INPUTS
text: To demonstrate iLQR-VAE’s ability to infer unobserved inputs and learn the ground truth dynamics of an input-driven system, we generated synthetic data from a system with n = 3,m = 3 and no = 10, which evolves with linear dynamics for T = 1000 time steps (see Appendix K for an example with input-driven nonlinear dynamics). The system was driven by sparse inputs, and the output corrupted with Gaussian noise. Input events were drawn in each time step from a Bernoulli distribution with mean p = 0.03. Whenever an input event occurred, the magnitude of inputs in each channel was drawn from a standard mutivariate Gaussian distribution.
iLQR-VAE captured most of the variance in these sparse inputs (Figure 3A; R2 = 0.94 ± 0.02; 5 random seeds), and recovered the eigenvalue spectrum of the transition matrix almost perfectly (Figure 3B). We also evaluated LFADS on the same data, using a generator with n = m = 3 and linear dynamics, and manually optimized hyperparameters. LFADS performed poorly in this case (R2 = 0.05 ± 0.02 for input reconstruction; 3 random seeds), as well as in several other similar comparisons on datasets of different sizes and trial numbers (Appendix H). This is unsurprising, as LFADS assumes a dense (auto-regressive) Gaussian prior over the inputs, which is not overridden by the relatively small amount of data used here (see Appendix J for a detailed comparison of Student and Gaussian priors). Nonetheless, when applied to a set of 56 trials of 100 time steps driven by Gaussian autoregressive inputs, iLQR-VAE still captured the structure in the inputs more accurately than LFADS did (R2 = 0.81± 0.01 vs. 0.29± 0.06).

heading: 3.4 PREDICTING HAND KINEMATICS FROM PRIMATE M1 RECORDINGS

heading: 3.4.1 TRIAL-STRUCTURED MAZE TASK
text: To highlight the utility of iLQR-VAE for analyzing experimental neuroscience data, we next apply it to recordings of monkey motor (M1) and dorsal premotor (PMd) cortices during a delayed reaching task (‘Maze’ dataset of Kaufman et al., 2016; DANDI 000128). This dataset contains 108 different reach configurations over nearly 3000 trials, and has recently been proposed as a ‘neuroscience MNIST’ to benchmark neural data analysis methods (Pei et al., 2021). Several latent variable models were evaluated on this dataset in (Pei et al., 2021), allowing us to directly compare the performance of iLQR-VAE to those reported results.
For this experiment, we fitted iLQR-VAE to the neural activity, using a model with GRU dynamics (n = 60), a Student prior over inputs (m = 15), and a Poisson likelihood (no = 182 neurons). To allow for direct comparison with baselines reported in Pei et al. (2021), we trained models on 1340
o1 ō1|o1
400 500 600 700 800
−0.2
0
0.2
0.9 1
true
learned
true ‖u‖
inferred ‖u‖
time [a.u.]
Im (λ
)
Re(λ)
A B
Figure 3: Inferring sparse inputs to a linear system. (A) Top: example observations (black dots) and inferred posterior mean (blue line). Bottom: true and inferred inputs. iLQR-VAE is able to infer the timing and magnitude of the inputs almost perfectly, despite being trained on only a single timeseries of 1000 time steps. Note that iLQR-VAE fails to infer one of the smallest inputs, whose effect on the observations is largely masked by observation noise. (B) Comparison of the true (black) and learned (blue) eigenvalue spectra. This shows that iLQR-VAE recovers the ground-truth dynamical system up to a similarity transformation.
trials spanning all reach conditions and restricting data to a time window starting 250 ms before and ending 450 ms after movement onset, binned at 5ms.
iLQR-VAE inferred inputs that were confined to initial conditions, from which smooth single-trial dynamics evolved near-autonomously (Figure 4A). As a first measure of performance, we evaluated the model on a “co-smoothing” metric defined as the normalized log-likelihood score of held-out neurons’ activity (see Appendix M) conditioned on held-in neurons. Conditioning of 137 neurons (i.e using 45 held-out neurons), we obtained a co-smoothing of 0.331 ± 0.001 (over 5 random seeds). For comparison, Pei et al. (2021) reports 0.187 for GPFA (Yu et al., 2009), 0.225 for SLDS (Linderman et al., 2017), 0.329 for Neural Data Transformers (Ye and Pandarinath, 2021) and R2 = 0.346 for AutoLFADS (LFADS with large scale hyperparameter optimization Keshtkaran et al., 2021) on the same dataset.
Next, we assessed how well hand velocity could be decoded from neural activity (another highly relevant metric for neural data). We applied ridge regression to predict the monkey’s hand velocity (with a 100 ms lag) from momentary neuronal firing rates (mean of the posterior predictive distribution) on test data. This reconstruction could be performed with very high accuracy R2 = 0.896 ± 0.002 (over 5 random seeds), compared to 0.640 for GPFA, 0.775 for SLDS, 0.897 for Neural Data Transformers and 0.907 for AutoLFADS (Pei et al., 2021).
These experiments place iLQR-VAE on par with state-of-the-art methods, without any extensive hyperparameter optimization.

heading: 3.4.2 CONTINUOUS REACHING TASK
text: While a large number of neuroscience studies perform neural and behavioural recordings during trial-structured tasks, much can be learned by analyzing the dynamics of more naturalistic, less constrained behaviours. iLQR-VAE’s flexible recognition model is well-suited to the analysis of such less structured tasks, as it can easily be trained and tested on trials of heterogeneous lengths. To illustrate this, we applied iLQR-VAE to a self-paced reaching task during which a monkey had to reach to consecutive targets randomly sampled from a 17x8 grid on a screen (O’Doherty et al., 2018; Makin et al., 2018). This dataset consists of both neural recordings from primary motor cortex (M1) together with continuous behavioural recordings in the form of x- and y-velocities of the fingertip.
In this experiment, we used GRU dynamics (n = 50), a Gaussian prior over inputs (m = 10), and a combination of a Poisson likelihood (130 neurons) and a Gaussian likelihood (2 kinematic variables). Here, we also experimented with fitting the spike trains and hand velocities jointly, as we found that it allowed iLQR-VAE to reach a similar kinematic decoding performance as when fitting neural activity alone, but using a smaller network. More generally, we reason that a more
natural approach to making behavioural predictions from neural data using a probabilistic generator is to fit it to both jointly, and then use the posterior predictive distribution over behavioural variables (conditioning on spike trains only) as a nonlinear decoder. In future work, this could provide more accurate predictions in those motor tasks where linear regression struggles (see e.g. Schroeder et al., 2021).
For our analyses, we used the first ∼ 22 minutes of a single recording session (indy 20160426), excluded neurons with overall firing rates below 2 Hz, and binned data at 25 ms resolution. Although it is not a formal requirement of our method, we chunked the data into 336 non-overlapping pseudotrials of 4 s each, in order to enable parallelization of the ELBO computation during training. We only trained the model on a random subset of 168 trials.
To highlight the flexibility of iLQR as a recognition model, we then evaluated the model by performing inference on the first 9 minutes of the data, as a single long chunk of observations. Note that this is not generally possible in LFADS or any sequential VAE where an encoder RNN has been trained exclusively on trials of the same fixed length. Despite the lack of trial structure, we found that neurons display a stereotyped firing pattern across multiple instances of each reach. This is revealed by binning the angular space into 8 reach directions, temporally segmenting and grouping the inferred firing rates according to the momentary reach direction, and aligning these segments to the time of target onset (Figure 4B). Moreover, hand kinematics could be linearly decoded from the inferred firing rates with high accuracy (Figure 4C; R2 = 0.75± 0.01 over 5 random seeds), on-par with AutoLFADS (R2 = 0.76; Keshtkaran et al., 2021), and considerably higher than GPFA and related approaches (Jensen et al., 2021).
We next wondered if we could use iLQR-VAE to address an open question in motor neuroscience, namely the extent to which the peri-movement dynamics of the motor cortex rely on external inputs (possibly from other brain areas). Such inputs could arise during movement preparation, execution, neither, or both. We thus examined the relationship between the inputs inferred by iLQR-VAE and the concurrent evolution of the neuronal firing rates and hand kinematics. Overall, neuronal activity
tends to rise rapidly starting 150 ms before movement onset (Figure 4D, red), consistent with the literature (Shenoy et al., 2013; Churchland et al., 2012). Interestingly, we found that inputs tend to arise much earlier (around the time of target onset), and start decaying well before the mean neural activity has finished rising (Figure 4D top), about 150 ms before the hand started to move (Figure 4D, bottom). While these results must be interpreted cautiously, as inference was performed using information from the whole duration of the trial (i.e. using iLQR as a smoother), they show that the data is best explained by large inputs prior to movement onset, rather than during movement itself. Interestingly, the timing of these inputs is globally consistent with target-induced visual inputs driving preparatory activity in M1, whose dynamics then evolve in a more autonomous manner to drive subsequent motion.

heading: 4 DISCUSSION

heading: LIMITATIONS AND FUTURE WORK
text: While we have demonstrated that iLQR-VAE performs well on various toy and real datasets, the method has a number of limitations that could be addressed in future work. Firstly, using iLQR as a way of solving maximum a posteriori inference in state space models comes at a high computational cost, and with the risk that iLQR may converge to a local minimum. Note that both these issues could potentially be tackled at once if process noise in the generator was modelled separately from control inputs, as the MAP estimation problem could then be solved using some of the highly efficient algorithms available in the framework of linearly solvable stochastic control (Todorov, 2009; Dvijotham and Todorov, 2013; Kappen, 2005). Moreover, for simplicity we modelled posterior input uncertainty using a common covariance across all data samples. This might be limiting, for example when modelling neural populations that exhibit coordinated global firing fluctuations giving rise to data samples with highly variable information content. A better solution would be to amortize the computation of the posterior uncertainty, ideally by reusing some of the computations performed in iLQR. We plan to investigate this in future work. Finally, the problem of decoupling ongoing inputs from dynamics is degenerate in general, and there is no guarantee that iLQR-VAE will always successfully identify the ground-truth. This problem will be exacerbated in the low data regime, or if there is a large mismatch between our prior over inputs and the true input distribution. Further generalization tests such as extrapolations can however be used to assess post-hoc how well the dynamics have been learned.

heading: CONCLUSION
text: The rise of new tools and software now makes it possible to record from thousands of neurons while monitoring behaviour in great detail (Jun et al., 2017; Mathis et al., 2018; Musk et al., 2019). These datasets create unique opportunities for understanding the brain dynamics that underlie neural and behavioural observations. However, identifying complex dynamical systems is a hard nonlinear filtering and learning problem that calls for new computational techniques (Kutschireiter et al., 2020). Here, we exploited the duality between control and inference (Toussaint, 2009; Kappen and Ruiz, 2016; Levine, 2018) to bring efficient algorithms for nonlinear control to bear on learning and inference in nonlinear state space models.
The method we proposed uses iLQR, a powerful general purpose nonlinear controller, to perform amortized inference over inputs in an RNN-based generative model. Using an optimization-based recognition model such as iLQR has two advantages. First, it brings important flexibility at test time, enabling predictions on arbitrary, heterogeneous sequences of observations as well as seamless handling of missing data. Second, owing to parameter sharing between the generative and recognition models, the ELBO gap is reduced (Appendix I), making learning more robust (in particular, to initialization) and reducing the number of hyperparameters to tune. With the advent of automatically differentiable optimizers (Blondel et al., 2021), we therefore hope that optimization-based recognition models will open up new avenues for VAEs.

heading: Appendix
text: A GRAPHICAL SUMMARY OF THE MODEL
input u
RNN
latent state z
observations o
ne ur
on s
time
inference (iLQR)
generation
Figure S1: Illustration of the model iLQR-VAE is trained to model spiking data. Every training iteration involves inference of the input u (green) given observations and the current parameters of the generator. Intuitively, the optimal inputs are ones that produce latent trajectories in the RNN that are most compatible with the data, without overfitting. After training, conditioning on neural activity and inferring the input generates predictions.
B IMPLEMENTATION OF THE DYNAMICS
We considered different functional forms for the discrete-time dynamics of the latent state. In the following, zt and ut denote the latent state and an external input at time t, respectively.

heading: B.1 LINEAR DYNAMICS
text: The simplest case considered is that of linear dynamics:
zt+1 = Azt +But (S1)
One issue with linear dynamics is that they may become unstable, such that repeated application of the operator A will lead to a divergence of ‖z‖ and the associated gradients. This can become problematic, especially when modelling long sequences of observations. To circumvent this issue, we used a parametrization of the propagator A that ensured it remained stable at all times. To find a stable linear parametrization, we considered the Lyapunov stability condition (Bhatia and Szegö, 2002). The discrete time dynamics of Equation S1 are asymptotically stable if and only ifA satisfies
P −APAT = I (S2) for some positive definite matrix P with eigenvalues ≥ 1. It is easy to verify that the following parameterization of the state matrixA satisfies this criterion:
A = UD1/2Q(D+ I)−1/2UT (S3)
with U and Q arbitrary unitary matrices, and D an arbitrary non-negative diagonal matrix. Conversely, any stable matrix can be reached by this parameterization. Note that the matrix P that satisfies Equation S2 is then given by P = U(D + I)UT . Finally, as we are also learning the B and C matrices in Equation S1, we can without loss of generality set U = I .

heading: B.2 GRU DYNAMICS
text: To fit the monkey reaching data as well as the Lorenz attractor, we chose the dynamical system to be Minimal Gated Unit (MGU). More specifically, we used the MGU2 variant of the MGU proposed in Heck and Salem (2017):
ft = σ(Ufzt−1) (S4) ẑt = g(Uh(ft zt−1) +Wxt + bh) (S5) zt = (1− ft) zt−1 + ft ẑt−1 (S6)
where xt = But denotes the input entering the dynamical system. Note that the latent state z is often denoted by h is the literature. We found that the MGU2 gave better and more stable performance than the MGU. We hypothesize that this is due to the input entering the system in the update gate only (as opposed to entering it through both forget and update gates), thus making the system more easily controllable. We chose σ(·) to be a sigmoid function, and g(·) to be a soft ReLu-like nonlinearity,
g(x) = x+ √ x2 + 4
2 . (S7)

heading: C LIKELIHOOD FUNCTIONS
text: The likelihood of the observations appear both in the ELBO and in the iLQR cost. Minimization of the latter via iLQR requires computing the momentary Jacobians and Hessians of the likelihood function w.r.t. the internal state zt. Although these quantities can be obtained generically via automatic differentiation, iLQR is always faster when they are provided directly (Appendix D), which we did here using the analytical expressions given below.

heading: C.1 GAUSSIAN LIKELIHOOD
text: For the Gaussian likelihood, we assumed observations o were linearly decoded from latents z, and corrupted with Gaussian noise, such that o ∼ N (Cz+ b,Σ), withC the readout matrix, b a vector of biases, and Σ a diagonal matrix of variances. This gives the following log-likelihood function:
logP (ot|z) = − 1
2
[ (Czt + b− ot)TΣ−1(Czt + b− ot) + no log(2π) +
∑ i log Σii
] (S8)
The Jacobian of this expression is as follows :
∂ logP (ot|zt) ∂zt = −CTΣ−1(Czt + b− o) (S9)
Finally, the Hessian is given by :
∂2 logP (ot|zt) ∂zt∂zTt = −CTΣ−1C (S10)

heading: C.2 POISSON LIKELIHOOD
text: To model spike trains, we assume that they are generated by a Poisson process with an underlying positive rate function for neuron i given by:
µi = βif((Cz)i + bi) ∆ (S11)
where ∆ denotes the time bin size and βi is a neuron-specific gain parameter. This yields the following log-likelihood :
logP (ot|zt) = no∑ i=1 (oi logµi − µi + log oi!) (S12)
where the sum is over neurons. Using the shorthand notations h(x) = log f(x) and at = Czt + b, the Jacobian and Hessian of this expression are given by :
∂ logP (ot|zt) ∂zt = CT [ot h′(at)−∆β f ′(at)] (S13)
∂2 logP (ot|zt) ∂zt∂zTt = CT [diag(ot h′′(at))−∆ diag(β f ′′(at))]C (S14)
D ILQR ALGORITHM
Our recognition model makes use of the iterative Linear Quadratic Regulator algorithm (iLQR; Li and Todorov, 2004; Tassa et al., 2014) to find the mean of the posterior distribution qφ(u|o). Iterative LQR is used to solve finite-horizon optimal control problems with non-linear dynamics and non-quadratic costs by (i) linearizing the dynamics locally around some initial trajectory, (ii) performing a quadratic approximation to the control cost around that same trajectory, (iii) solving the linear-quadratic problem generated by the local approximation to obtain better control inputs, and (iv) repeat until convergence, each time linearizing around the trajectory induced by the new inputs. Below, we first introduce the linear-quadratic regulator (LQR), and detail the approximation used in iLQR to turn any non-linear non-quadratic problem into one that can be solved with LQR. Moreover, we provide pseudo-code for our implementation of iLQR (see Algorithm 1).
The Linear Quadratic Regulator is concerned with finding the set of controls u ∈ Rm that minimize a quadratic cost function C(u) under deterministic linear dynamics, given by:
C(u) = T∑ t=0 1 2 ( zTt C zz t zt + u T t C uu t ut + z T t C zu t ut + u T t C uz t zt ) + zTt c z t + u T t c u t (S15)
s.t. zt+1 = Atzt +Btut + ht. (S16)
Here,At ∈ Rn×n is a (possibly time-dependent) transition matrix,Bt ∈ Rn×m represents the input channels at time t, and ht is a state and input-independent term. Note that while the state z ∈ Rn appears in Equation S15, it is a deterministic function of the initial condition z0 and the sequence of inputs u. LQR finds the inputs minimizing Equation S15 using a dynamic programming approach, by recursively finding the feedback rule (Kt,kt) which gives the optimal inputs to minimize the cost-to-go at each time t as ut = Ktzt + kt. Details can be found in function Backward in Algorithm 1.
iLQR is an extension of LQR to general dynamics and cost functions. Specifically, iLQR minimizes
Cθ(u) = T−1∑ t=0 rθ(zt,ut, t) subject to zt+1 = fθ(zt,ut, t) (S17)
where θ denotes a set of parameters. At iteration i, iLQR approximates both the dynamics and the cost around the current trajectory τ i = (zi,ui) as:
f̃θ(δz i t, δu i t, t) ≈ (∇zfθ)T δzt + (∇ufθ)T δut (S18)
and
r̃θ(δz i t, δu i t, t) ≈
1
2
[ δzTt (∇2zzrθ)δzt + 2δzTt (∇2zurθ)δut + δuTt (∇2uurθ)δut ] (S19)
+ δzT (∇zrθ) + δuT (∇urθ) (S20)
Here, δz and δu refer to perturbations around the current nominal trajectory, and all ∇ operators correspond to partial differentiation evaluated at the current nominal trajectory (uit, z t i) and corresponding time t. The above equations are readily identified as a local LQR problem of the form of Equation S15, which can thus be solved using standard dynamic programming tools. Once δu? minimizing Equation S19 has been computed, the inputs are updated as ui+1 = ui + δu?,i, and the new state trajectory follows from simulating the dynamics forward with these new inputs. After each LQR update, we thus obtain a new trajectory τ i+1 and the process repeats until convergence to some locally optimal trajectory τ ?.
Implementation details can be found in Algorithm 1. Note that the backward LQR pass involves inversion of the matrix Quu (defined in Algorithm 1 function Backward). Depending on the specific form of the iLQR cost function, this might not always be positive-definite. Therefore, we include an adaptive Levenberg-Marquard-type regularizer (not described in the pseudo-code) Quu ← Quu + λI to maintain positive definiteness. Thus, iLQR effectively reverts to first-order gradient descent, as opposed to second-order optimization, whenever the locally quadratic approximation is a bad one.
Algorithm 1 iLQRsolve(Cθ(u),uinit)), with u ∈ Rm and Cθ defined in Equation S17. Parameters: θ, γ

heading: I iLQR
text: τ 0 = Rollout(uinit) for i = 1 to converged do
for t = 0 to T do F zt = ∇zfθ, F ut = ∇ufθ czt = ∇zrθ, cut = ∇urθ, Czzt = ∇2zrθ, Cuut = ∇2urθ, Cuzt = ∇2uzrθ end for k[0:T−1],K[0:T−1] = Backward( F z[0:T ],F u [0:T ], c z [0:T ], c u [0:T ],C zz [0:T ],C uz [0:T ],C uu [0:T ]))
τ i = Forward(K[0:T−1],k[0:T−1], τ i−1) end for
function Rollout(u) z0 = 0 for t = 1 to T do zt+1 = fθ(zt,ut) end for return τ = {z,u}
function Backward( F x[0:T ],F u [0:T ], c z [0:T ], c u [0:T ],C zz [0:T ],C uz [0:T ],C uu [0:T ]) vT = c z T ,VT = C zz T for t = T − 1 to 0 do Qzzt = C zz t + F z t >Vt+1F z t
Quzt = C uz t + F u t >Vt+1F z t Quut = C uu t + F u t >Vt+1F u t qxt = c z t + F
zvt+1 qut = c u t + F
uvt+1 Kt = − (Quut ) −1 Quzt kt = − (Quut ) −1 qut Vt = Q zz t +Q zu t Kt vt = q z t +K T t q u t
end for return k[0:T−1],K[0:T−1]
function Forward(K[0:T−1],k[0:T−1], τ = {u, z}) α = 1 repeat ẑ0 = 0 û0 = αk0 for t = 1 to T do ẑt = fθ(ẑt−1, ût−1) ût = Kt(ẑt − zt) + αkt
end for α = γα
until Cθ(û) < Cθ(u) return τ = {ẑ, û}

heading: E DIFFERENTIATING THROUGH ILQR
text: Here we discuss how to efficiently differentiate through the iLQR algorithm. This becomes necessary when one wishes to differentiate through a function involving an iLQRsolve, such as the posterior mean of our recognition model (Equation 11). While a naive but simple strategy to achieve this would be to unroll the algorithm and gather gradients for every step, this is expensive both computationally and memory-wise. Amos et al. (2018) derived a way to analytically obtain gradients with respect to the parameters of iLQR, at the cost of a single LQR pass. Specifically, differentiating through an iLQRsolve is achieved by running iLQR to convergence, forming a linear-quadratic approximation around the converged trajectory, following the steps described in Appendix D and differentiating through the corresponding LQR problem. Below, we provide an alternative derivation to Amos et al.’s of the gradients of an LQR solution.

heading: E.1 LQR OPTIMALITY CONDITIONS
text: We now introduce use the more compact notation τt = [ zt ut ] , Ft = [At Bt],Ct = [ Czzt C uz t Czut C uu t ] , which will be used in the rest of this section.
As described in Appendix D, the finite-horizon, discrete-time LQR problem involves minimizing:
J = T∑ t=0 (1 2 τTt Ctτt + c T t τt )
(S21)
subject to constraints on its dynamics
τt+1 = Ftτt + ft, (S22)
following the notation from Appendix D. To solve this problem, we write down the Lagrangian:
L = T∑ t=0 (1 2 τTt Ctτt + c T t τt ) + T−1∑ t=0 λTt+1(Ftτt + ft − xt+1), (S23)
where λ1,λ2, · · · ,λT are adjoint (dual) variables that enforce the dynamic constraint. Differentiating with respect to λt and τt enables us to obtain the set of equations satified by λ and τ , also known as the KKT conditions (Kuhn and Tucker, 2014; Karush, 2014; Boyd et al., 2004):
[I 0] (Ctτt + ct) + F T t λt+1 − λt = 0 (S24)
Ftτt + ft − [I 0] τt+1 = 0 (S25)
CT τT + cT − [I 0] T λT = 0 (S26)
z0 − [I 0] τ0 = 0 (S27) Rearranging, we can rewrite the KKT conditions in matrix form as: . . . Ct F T t Ft 0 [−I 0] [−I 0]T Ct+1 F Tt+1
Ft+1 0 . . .  ︸ ︷︷ ︸
K
 ... τt λt+1 τt+1 λt+2
...  ︸ ︷︷ ︸
p
= −  ... ct ft ct+1 ft+1
...  ︸ ︷︷ ︸
q
(S28)
These optimality conditions are satisfied for the solution to the optimization problem p? = (τ ?0 , · · · , τ ?T ,λ?1, · · ·λ?T ). Equation S28 implies that the solution of the LQR problem p? will satisfy:
p? = −K−1q. (S29) Computing this quickly becomes infeasible asK grows with long-time horizons, and Equation S28 is typically solved in linear time using a dynamic programming approach, as described in Appendix D.

heading: E.2 BACKPROPAGATING THROUGH THE LQR SOLVER
text: Differentiating through an LQR solve boils down to differentiating through the backsolve in Equation S29. In the following, we denote the adjoint of parameter θ as θ̄. From Giles (2008), we know that the adjoint of the backsolve operation is given by:
q̄ = −K−T p̄, (S30) K̄ = −K−T p̄pT = q̄pT . (S31)
We note that Equation S30 has the same form as Equation S29, which means we can compute q̄ = (· · · , c̄t, f̄t, · · · )T by solving another LQR problem. After solving for τ̄ , we can then compute K̄ as an outer-product of z̄ with y to get: . . . K̄Ct K̄FTt K̄Ft K̄Ct+1 K̄FTt+1 K̄Ft+1
. . .  ︸ ︷︷ ︸
K̄
(S32)
=  ... c̄t f̄t c̄t+1 f̄t+1
...
 ( · · · τ ?t T λ?t+1 T τ ?t+1 T λ?t+2 T · · · ) (S33)
=  . . . c̄tτ ? t T c̄tλ ? t T f̄t+1τ ? t T c̄t+1τ ? t+1 T c̄t+1λ ? t+1 T c̄t+1λ ? t+1 T
. . .
 . (S34)
Collecting all the gradients of C̄t and F̄t, we arrive at
C̄t = 1
2 (K̄Ct + K̄
T Ct) =
1 2 (c̄tτ ? t T + τ ?t c̄ T t ) (S35)
F̄t = K̄Ft + K̄ T FTt = f̄t+1τ ? t T + λ?t+1c̄ T t . (S36)
Note that we have symmetrized the adjoint of Ct, which ensures that Ct remains symmetric after each gradient update. The antisymmetric part of Ct does not contribute to the LQR cost.
Finally, one should keep in mind that Equation S22 describes dynamics as a function of τ , as opposed to the locally linearized iLQR dynamics given by Equation S18 which are instead a function of δτ . This induces an additional constant in Equation S22, which must be included when computing gradients.

heading: F LINK TO KALMAN FILTERING
text: The Linear Quadratic Regulator and the Kalman filter (Kalman, 1964) are algorithms designed for systems with linear dynamics and Gaussian noise. LQR finds the optimal feedback control law to minimize a cost C in deterministic systems, while the Kalman filter yields an estimate of the state from observations corrupted with process and observation noise. It is well-known that Kalman filtering and LQR are dual of one another, and they can both be combined into LQG to yield an
optimal control law from noisy observations. Here, we explore another link between LQR and Kalman smoothing, by showing how LQR can be used as a Kalman smoother. Moreover, in order to gain insights into the learning process of iLQR-VAE, we explore different procedures for learning the parameters of a Kalman filter.

heading: LINEAR QUADRATIC CONTROL AS FILTERING
text: The Kalman smoother assumes dynamics of the form
zt+1 = Azt +B Kwt (S37)
ot = Czt + vt (S38)
with w ∼ N (0, I), v ∼ N (0,Σv), and the initial condition is assumed to be generated by a Gaussian distribution with known parameters z1 ∼ N (µ,Π). On the other hand, LQR assumes the following fully-deterministic dynamics :
zt+1 = Azt +B Iut (S39)
ot = Czt (S40)
with z1 known exactly.
Note than in the iLQR-VAE framework we have thus far only considered cases where no observed external inputs were given. However these can be straightforwardly included as an additional B̂û term in Equation S39 and Equation S37.
The Kalman smoother’s objective is to minimize the expected mean squared error between the inferred latent state and the true state, E [ ‖x− x̂‖2 ] . As described in Aravkin et al. (2017), with linear dynamics and Gaussian noise, this becomes equivalent to minimizing the following objective w.r.t z:
L(z) = ‖Π−1/2(z1 − µ)‖2 + T−1∑ t=1 ‖B(zt+1 −Azt)‖2 + T∑ t=1 ‖Σ−1/2v (ot −Czt)‖2 (S41)
where the first two terms correspond to the prior over the initial condition and smoothness of the trajectory, and the last term represents the likelihood of the observations. Interestingly, this can be related to the objective we minimize to find the posterior mean in iLQR-VAE (Equation 11):
L(u) = ‖Σ−1/20 u0‖2 + T−1∑ t=1 ‖Σ−1/2u ut‖2 + T∑ t=1 ‖Σ−1/2v (ot −Czt)‖2 (S42)
= ‖Σ−1/20 u0‖2 + T−1∑ t=1 ‖Σ−1/2u (zt+1 −Azt)‖2 + T∑ t=1 ‖Σ−1/2v (ot −Czt)‖2. (S43)
The right-hand sides of Equation S41 and Equation S42 become identical when Σu = BK and Σ0 = Π. Note that in practice the introduction of the BI matrix in Equation S37 unties the two formulations slightly by allowing for further mixing between the input channels that isn’t accounted for by the prior. In the examples we consider next, we therefore setBI = I.
The above equations show how LQR can be used to solve the standard Kalman filtering problem, with the key difference being that the optimization is performed over inputs u = {u0, . . . ,uT−1} rather than latent trajectories u = {z1, . . . ,zT } directly. This is illustrated in Figure S2(A), where a Rauch-Kung-Striebel smoother and LQR were ran on the same set of 8-dimensional observations arising from an 8-dimensional linear dynamical system, and inferred the same latent trajectory given the ground-truth parameters. As we only use LQR to parametrize the mean of the posterior distribution, we trained the recognition model for 100 steps to get the uncertainty over the latents, which was very similar to the output of the RKS smoother.

heading: LEARNING A KALMAN FILTER
text: We then proceeded to learn the parameters of the models using either iLQR-VAE, an ExpectationMaximization (EM) procedure, or direct minimization of the negative log likelihood of the data (Figure S2B-C).
1 10 100 1000
la te
nt s
ou tp
ut
nl l
i
iLQR-VAE Kalman
A B C
Figure S2: Comparison of the posterior mean of a Kalman smoother and LQR. We ran both LQR and a Kalman smoother on noisy observations generated by a latent system (see text for details). (A) The Kalman smoother (blue) and LQR (dotted red) both inferred the same posterior mean for the latent trajectories, matching the true latents (black) almost perfectly. The posterior uncertainty is shown for both cases on half of the data. The iLQR-VAE uncertainty was obtained by optimizing the variance of the recognition model for 200 iterations, and then drawing 1000 samples from the recognition model. (B) We compared learning the parameters of the posterior distribution using iLQR v.s direct minimization of the NLL. On unseen data, both were able to converge to a solution close to the smoothed output trajectory. Note that we stopped the optimization after 8000 iterations. (C) Learning curves of the direct optimization and iLQR-VAE on the same 168 training trials. Note that we used Adam with the same learning rate of 0.02 in both cases, in order to directly compare the effect of the gradient steps. Both x and y axes use a logarithmic scale.
Interestingly, the EM algorithm is closely related to iLQR-VAE, since the E-step finds the latent trajectories minimizing Equation S41, when iLQR-VAE solves Equation S42 in an inner optimization loop. While there exists an analytical solution for the M-step in the case of the Kalman filter, this does not generalize to nonlinear dynamics and non-Gaussian noise. Therefore, we used a gradient descent procedure in the M-step.
Both of these were performed using Adam with a learning rate of 0.02, and with initial parameters drawn from the same distributions. We see in Figure S2C that iLQR-VAE reaches a smaller NLL in considerably fewer iterations than gradient descent, which we hypothesize is due to the good preconditioning given by iLQR (discussed in Figure 1). Note however that the cost of one iLQRVAE iteration is higher than the direct computation of Equation S41.
In this section, we have shown in a simple linear-quadratic example how iLQR-VAE performs filtering by inferring the process noise as inputs. While this is undoubtedly an unconventional approach, it becomes particularly valuable in cases where dynamics are non-linear and the noise non-Gaussian. Indeed, in such cases the problem of learning an estimator for the latent state is a very difficult one, typically solved using methods such as particle filtering or unscented Kalman filters (Doucet and Johansen, 2009; Wan et al., 2001). iLQR-VAE therefore offers another way to solve this problem, with close links to the aforementioned approaches.
G DETAILS OF EXPERIMENT 1
0
1
2
3
4
5
6
7
8
0 5 10 15 20
ne ga
tiv e
E LB
O
iteration number ×103
Figure S3: Illustration of the learning curves for 5 different runs of the “pinned autonomous” iLQR-VAE model. The models consistently get stuck in plateaus during the optimization, unlike their “free” counterparts (see Figure 1).
The data in Section 3.1 was generated from an autonomous linear dynamical system with n = 8, m = 3, and no = 8 where no is the dimension of the observation space. All the models were fit using the dynamics within the ground-truth model class, i.e with linear dynamics, n = 8, m = 3, and no = 8. We optimized the model parameters with Adam, using (manually optimized) learning
rates of 0.04/(1+ √ k/1) for the free iLQR-VAE model, 0.04/(1+ √ k/1 for autonomous iLQR-VAE
and 0.02/(1+ √ k/30 for LFADS, where k is the iteration number. We used GRU networks with 32 units to parametrize the LFADS encoders (one encoder for the initial condition and one for the inputs).
In Figure S3, we show additional learning curves for the “pinned autonomous” models; these show that, even for different initializations and trajectories through the loss landscape, the model consistently gets stuck in plateaus. This can be constrasted with the “free-form” iLQR-VAE models.
H COMPARISON OF LFADS AND ILQR-VAE ON A TOY INPUT INFERENCE TASK
o1 ō1|o1
0 100 −0.2
0
0.2
0.9 1
true
learned
true ‖u‖
inferred ‖u‖
time [a.u.]
Im (λ
)
Re(λ)
A B
Figure S4: Details of the sparse input inference in LFADS and iLQR-VAE. (A) Top: example observations (black dots) and inferred posterior mean (blue line). Bottom: true and inferred inputs. LFADS can infer the timing of the largest input, but also uses non-zero inputs during the rest of the trial. (B) Comparison of the true (black) and learned (blue) eigenvalue spectra.
We used the LFADS implementation from https://github.com/google-research/ computation-thru-dynamics/tree/master/lfads_tutorial, which we modified to include linear dynamics and Gaussian likelihoods. We then evaluated the quality of the input reconstruction by measuring how much input variance was captured by the models. We report this as the R2 from inferred to true inputs.
We used a generative model within the ground truth model class. For each dataset, we performed a hyper-parameter search to choose the best-performing encoder architecture and learning rate for LFADS.
Results of this experiment are summarized in Table S1. iLQR-VAE – which did not require any hyperparameter tuning for these examples – inferred inputs more accurately for all dataset sizes and trial lengths.
LFADS iLQR-VAE S 1x1000 0.05± 0.02 0.94± 0.01 S 10x100 0.15± 0.06 0.84± 0.01 S 32x100 0.29± 0.01 0.83± 0.05 S 56x100 0.31± 0.08 0.80± 0.02 S 10x200 0.27± 0.01 0.93± 0.01 AR 56x100 0.28± 0.02 0.81± 0.02
Table S1: Comparison of iLQR-VAE and LFADS on 6 input inference tasks. Results are reported as R2 (mean ± sem) over 3 random seeds for each.
Our results suggest that LFADS’ performance improves with larger amounts of data. More surprisingly, LFADS also seems to perform better when the data is split into shorter trials. In particular, we found it difficult to fit LFADS on the single long trial, but the dynamics could be learned more accurately if this data was split into 10 trials of 100 steps. On the other hand, iLQR-VAE inferred
inputs more accurately for longer trials. This is what we would expect if the model is well learnt, as longer trials contain more information to fit the inputs accurately.
One important distinction between the two methods, which partly explains LFADS’ lowerR2, is the prior it over inputs (auto-regressive prior for LFADS and Student for iLQR-VAE). In Figure S4 we show an example of LFADS, on one of the test examples of the S 56x100 dataset. In this example, LFADS infers its largest input concurrently to the ground truth input, but also infers small inputs when there are none in the ground truth. This has a significant impact on the R2 metric. Note however that this is not the only effect at play here, as emphasized by the lower performance on the AR dataset.

heading: I ANALYSIS OF THE INFERENCE GAP
text: In order to evaluate the benefits of defining the recognition model implicitly through the generative parameters, we compared iLQR-VAE to a more standard sequential variational auto-encoder, using a bidirectional recurrent neural network as the recognition model. We generated data from the same system as in Figure 3, in the form of 76 trials of 100 time steps. We used the same generative model in both cases (linear dynamics with n = 3, m = 3, no = 10, Student prior), such that the only difference lay in the choice of recognition model. We compared the ELBO to a more accurate estimate of the log-likelihood, the Importance Weighted Autoencoder (IWAE) bound (Burda et al., 2015), which is computed as
LIWAE = Eu1,..un∼q(u|o)
[ log ( 1
k k∑ i=1 p(o, zi) q(ui|o)
)] . (S44)
This then allowed us to compute the inference gap Cremer et al. (2018) of both models as LIWAE− ELBO. As shown in Figure S5, iLQR-VAE has a smaller inference gap throughout training, leading to faster and more robust convergence. This confirms the intuition that keeping the recognition and generative models in sync throughout training reduces the inference gap.
−6 −5 −4 −3 −2 −1 0
1
2
500 1000 1500 2000 2500 3000 0
0.5
1
1.5
2
2.5
500 1000 1500 2000 2500 3000
0
0.5
1
1.5
2
2.5
3
3.5
0 10 20 30 40 50 60 70 80 90 100 −2 −1.5 −1 −0.5
0
0.5
1
0 10 20 30 40 50 60 70 80 90 100
iteration
iLQR-VAE LL iLQR-VAE ELBO
biRNN LL biRNN ELBO in
fe re
nc e
ga p
iteration
iLQR-VAE biRNN
‖u ‖
time (ms)
ob se
rv at
io ns
time (ms)
A B
C D
Figure S5: Comparison of iLQR with a biRNN recognition model. (AB) Loss and inference gap as a function of iteration number (starting from iteration 20) for iLQR-VAE (blue) and the biRNN model (pink). (C) Inferred input norm at the end of training for iLQR-VAE and the biRNN model. Ground truth input is shown in dotted black lines. (D) Inferred output at the end of training for iLQR-VAE and the biRNN model. Both models explain the observations (grey dots) well, but iLQR-VAE captures the sharp transitions better.

heading: J COMPARISON OF THE STUDENT AND GAUSSIAN PRIORS
text: In this section, we compare the performance of the Gaussian and Student priors on two toy examples. The first consists of data generated by a linear dynamical system (n = 3,m = 3, no = 10, Gaussian likelihood) driven by autoregressive Gaussian inputs (close to the Gaussian prior). The second one uses the same system, but driven by sparse inputs (closer to the Student prior). We find that in the
10
100
1000
0 10 20 30 40 50 1
10
100
0 10 20 30 40 50
ob s
ob s
time /ms
‖u ‖
time /ms
lo ss
iteration
Gaussian Student
ob s
ob s
time /ms
‖u ‖
time /ms
lo g
lo ss
iteration
Gaussian Student
autoregressive sparse
A B
C
D E
F
Figure S6: Comparison of the Student and Gaussian prior on a dynamical system driven by autoregressive inputs and sparse inputs. (A) Fit of the observations using the Gaussian (top) and Student (bottom) priors. Both allow to fit the observations highly accurately. (B) Inferred input norm for both priors. The temporal structure of the signal is very similar in both cases, as the Student prior essentially becomes Gaussian for large ν values. (C) Evolution of the loss for both prior choices. The loss curves are closely aligned, and both models converge to a similar ELBO value. (D) Fit of the observations using the Gaussian (top) and Student (bottom) priors. Both models fit the observations, but the Student prior recovers smoother trajectories. (E) Inferred input norm for both priors. The Student prior is close to 0 at all times, except when it requires a sharp input to explain the data. On the other hand, the Gaussian prior requires a large variance to be able to fit sparse inputs, leading to non-zero inferred inputs at all time points. (F) Evolution of the loss for both prior choices. The loss curve for the Student prior is lower than the Gaussian one throughout training, and converges to a higher ELBO value.
first example, both priors yield extremely similar results (see Figure S6(A-C)). Indeed, the Student prior learns a very high value of ν ∼ 20, thus becoming nearly Gaussian. In the sparse input case however, the Student prior allows to fit the data considerably better. As we can see in Figure S6(D-F), the Gaussian prior learns a large variance to fit the sparse inputs, leading to higher baseline noise than in the true system.
As we can see here, the Student prior offers more flexibility, as the Gaussian case is recovered for large ν values. Note however that the Gaussian prior leads to an input term in the iLQR cost function that is always convex in u, which can facilitate the optimization and allow iLQR to converge faster in some cases.Moreover, in cases where the learnt dynamics are autonomous (here, Lorenz attractor and Maze dataset), both priors will find the same solution.

heading: K LEARNING INPUT-DRIVEN NONLINEAR DYNAMICS
text: To bridge the gap between autonomous nonlinear dynamics (see Section 3.2) and real data, we evaluated iLQR-VAE on an input-driven nonlinear system, the Duffing oscillator. We embedded Duffing trajectories that included a perturbation of the Duffing state half-way through, into the spiking activity of 200 neurons (see below). We found that iLQR-VAE could successfully learn the dynamics and infer the timing of the perturbations.
The dynamics of the Duffing oscillator are given by
ẋ1 = x2 ẋ2 = x1 − x31. (S45) To generate each training sample, we integrated Equation S45 from two different random initial conditions for 100 time steps each using a Runge-Kutta solver (RK4) with dt = 0.03. Example such trajectories are shown in Figure S7A. We then concatenated those two trajectories into a single trajectory of 200 steps, which then had a perturbation in the middle. We then linearly mapped the
low-dimensional oscillator onto a 200-dimensional state, before passing it through the nonlinearity of Equation S7 to obtain a set of firing rates, which then gave rise to observations via a Poisson process (Figure S7B, top).
We generated 112 training and 112 testing trials in this way. We fit these data using iLQR-VAE with n = 20, m = 4, and found that it could successfully infer the latent trajectories (see Figure S7B, middle). Importantly, iLQR-VAE learned to fit most of the trajectories as an autonomously evolving dynamical system, and only used inputs to explain the sudden change in the oscillator’s energy level triggered by the perturbation (see Figure S7B, bottom). This shows that the model can successfully disentangle ongoing dynamics from external inputs, suggesting that it is well-suited for identifying input-driven dynamics in real data.
x1
x2
0.166
0.168
0.17
0 200 400 600 800 1000
0
100
200
−50 0 50
−50 0 50
0 10
0 100 200
ne ga
tiv e
E LB
O
training iterations
ne ur
on s
P C 1 P C 2 ‖u ‖
time [a.u.]
A B
Figure S7: (A) Top: example Duffing trajectories (100 time steps, dt = 0.03) starting from random initial conditions. Bottom: negative ELBO during the course of training. (B) Top: spike raster corresponding to an example test sample. Middle: first two principal components of the posterior mean over firing rates (red), given the spiking data shown at the top. For comparison, the groundtruth PCs in the first half of the trial (before the perturbation) are shown as black dots, with their hypothetical unperturbed continuation shown as a dashed line. The second half of the ground-truth PCs (after the perturbation) are shown in gray. Bottom: norm of the inferred input.

heading: L FURTHER DETAILS OF LORENZ ATTRACTOR
text: The chaotic Lorenz attractor consists of a three-dimensional state (`1, `2, `3) evolving according to
˙̀ 1 = 10(`2 − `1) ˙̀2 = `1(28− `3)− `2 ˙̀3 = `1`2 − 8`3/3 (S46)
For our example, we generated data by integrating Equation S46 over a long time period using a Runge-Kutta solver (RK4) followed by z-scoring and splitting the resulting state trajectory into 112 non-overlapping bouts (Figure 2A). We added Gaussian noise with a standard deviation of 0.1, and trained iLQR-VAE on this dataset (Figure 2B, bottom). We then fitted these data using GRU dynamics with n = 20 and m = 5.
The normalized k-step mean-squared error was defined as follows:
MSEk = T−k∑ t=0 ‖xt+k − x̂t+k‖2 (S47) R2k = 1− MSEk∑T−k
t=0 ‖xt+k − x̄‖2 (S48)
where x̂t+k is the prediction at time t+ k, and x̄ the mean for this trial.

heading: M FURTHER DETAILS OF SINGLE TRIAL ANALYSES

heading: M.1 BENCHMARKING AGAINST EXISTING METHODS
text: To allow for direct comparison with benchmarks reported Pei et al. (2021), we first used data provided by the Neural Latents Benchmark (NLB) challenge, available at https://gui. dandiarchive.org/#/dandiset/000128. This data included recordings from 250 ms before movement onset to 450 ms after, binned at 5ms. For regression to hand velocity, we introduced a lag of 100ms between neural activity and hand velocity.
We used 1720 training trials and 510 validation trials, which were drawn randomly for each instantiation of the model, to avoid overfitting to test data. Moreover, iLQR-VAE requires very little hyperparameter optimization, lowering the risk of overfitting to the dataset. As the test data used in the NLB challenge is not publicly available, the results we reported were not computed on the exact same data split. However, the model performed highly consistently across random seeds, such that we expect iLQR-VAE’s performance to be directly comparable to the results from Pei et al. (2021). To fit these data, we ran iLQR-VAE on 168 CPUs for ∼ 6h, using a mini-batch size of 168 trials. The co-smoothing metric used to assess how well the model fit the data is defined as log-likelihood score :
1
nsp log 2
( L(λ; ŷn,t)− L(λ̂n, ŷn,t) ) (S49)
where the overall log-likelihood L is the sum of all the log-likelihoods evaluated at all points and for all neurons, λ denotes the vector inferred time-varying firing rates, λ̂n is the mean firing rate for neuron n and nsp is the total number of spikes.

heading: M.2 FURTHER ANALYSES
text: A key feature of monkey M1 motor cortical recordings is the prevalence of rotational dynamics in the data (Churchland et al., 2012). These can be captured using jPCA, a method developed to find the subspace in which the dynamics are most rotational (note that these rotations can also be recovered in the more general framework of Rutten et al. (2020). Here, we found that we could uncover clean rotational dynamics from the single-trial firing rates, similarly to Pandarinath et al. (2018).
jP C
A 2
jPCA1
Figure S8: Projection of the neural activity of 200 movements in the subspace defined by the top two jPCA axes. jPCA finds the subspace capturing most rotations in the data, while spanning the same space as the top 2 principal components. Here, the jPCA subspace was found using the single-trial firing rates. Projection of the neural activity yields very clean rotational trajectories.

heading: M.3 LOOKING FOR PREPARATORY INPUTS
text: A decade of research on the preparation and execution of arm reaches in monkeys has shown that in tasks with explicit delays between target and movement onset, cortical activity is driven into a preparatory subspace during this delay (Churchland et al., 2012; Shenoy et al., 2013; Lara et al., 2018). It has recently been suggested that this could be mediated by inputs during preparation (Kao et al., 2021). To evaluate whether iLQR-VAE could potentially be used to discover systems relying on preparatory inputs, we generated a toy example, modelling a reaching task in which the motor cortex is driven mainly by preparatory inputs and evolves autonomously during movement. More concretely, we generated a network of 40 units with non-normal connectivity, and linearly decoded the neural activity (firing rates) onto a set of torques driving a two-link arm (Todorov et al., 2005).
We then defined a task consisting of a 300 ms delay period in which the arm could not move, and a 300 ms movement phase in which it had to perform a reach to a target. We optimized the inputs to the network such that it was driven by inputs during the delay period only (see Figure S9B), yet performed a perfect reach to target. We generated data for 96 different reaches, added Gaussian noise to the neural activity traces (see Figure S9A), and we then trained an iLQR-VAE model with generator dynamics within the ground-truth model class on half of those trials and tested it on the other half. Interestingly, iLQR-VAE learned to model the data using a system mainly driven by preparatory inputs (see Figure S9C). Over all test trials, the ratio of preparatory to movement input magnitudes was 5, suggesting that the learnt dynamical system relied strongly on preparatory inputs.
-300 mov +300 -300 mov +300 -300 mov +300
u (tr
ue )
u (le
ar ne
d)
A B C
Figure S9: iLQR-VAE can infer preparatory inputs in a toy linear model. (A) Example firing rates for a given reach (10 neurons shown). (B) Example preparatory inputs in the model. (C) Example inputs inferred by iLQR-VAE when fitting the data from (A).
The above example suggests that iLQR-VAE is ideally placed to exploit population recordings to investigate hypotheses concerning movement preparation and execution. Thus, we wondered whether preparatory processes could be identified in the neural data; to answer this, we used neural recordings from the Maze dataset, from 500ms before movement onset to 500ms after. We fit these using the same model as described in Section 3.4.1, and found that iLQR-VAE could fit the entire trial by evolving autonomously from the initial condition (Figure S10;B). While unsurprising (Pandarinath et al. (2018) can fit a similar split of data using a fully autonomous GRU), this highlights the difficulty of disentangling parts of the observations that can be explained solely using the model dynamics from those requiring external inputs. We suggest that questions relating to preparatory inputs could nevertheless be answered using methods such as iLQR-VAE, for instance by evaluating the likelihood of the data under different priors.
0
40
-500 mov +500 0
50
-500 mov +500
−1
0
1
-500 mov +500
ra te
(H z)
time /ms
neuron 1
ra te
(H z)
time /ms
neuron 90
u (a
.u )
time /ms
A B
Figure S10: iLQR-VAE learns to fit both pre-movement and movement activity using a single autonomously evolving dynamical system. (A) Firing rates for neuron 1 (left) and neuron 90 (right) across reaches. Some neurons show more modulation during preparation, while others appear to encode mainly movement-related activity. (B) Input inferred by the model.

heading: N FURTHER DETAILS OF THE CONTINUOUS REACHING TASK ANALYSIS

heading: N.1 DETAILS OF THE ANALYSES
text: For our analyses of the primate data in Section 3.4, we considered the first 22 minutes of the recording session ‘indy 20160426’ from O’Doherty et al. (2018). We binned spikes at 25 ms resolution and considered all neurons with a firing rate of at least 2Hz. The resulting spike counts were squarerooted prior to model fitting. Behavioural data took the form of the velocity of the hand of the
monkey in the xy-plane and were extracted as the first derivative of a cubic spline fitted to the position over time. We shifted the hand velocity by 120ms, following on Jensen et al. (2021). Both neural activity and behavioral data was subsequently z-scored, and the hand velocities were passed through a function f(x) = sign(x) log(1 + |x|) to reduce its sparsity. To fit iLQR-VAE, the resulting dataset was divided into 336 non-overlapping pseudo-trials of which a random half were used to fit the generative model and the other half of the trials were used as a held-out test dataset. We fitted a model with n = 30,m = 10 using the non-linear dynamics described in Equation S4. The latent state was linearly read out onto the 130 neurons.
After fitting the iLQR-VAE to neural activity and behavior jointly, we then proceeded to infer u from neural activity alone. We then computed the kinematic reconstruction error on the test dataset as the fraction of variance captured in both x- and y- hand velocities.
Finally, we analyzed the inputs to the model after fitting, in relation to specific events in the task and the behaviour. We defined ‘movement onset’ after each target onset as the time at which the hand speed first exceeded 0.03 m s−1. Finally we aligned the z-scored input u on each trial to target onset and movement onset separately for visualization purposes. We performed a similar analysis for hand speed and mean z-scored neural activity which were also z-scored and aligned to target and movement onset for comparison with the control input.
N.2 COMPARISON OF ILQR-VAE AND BGPFA
0 10 20 30 40 50 60
0 10 20 30 40 50 60
25 50 75 100
0 10 20 30 40 50 60
0 10 20 30 40 50 60
25 50 75 100
0 10 20 30 40 50 60
0 10 20 30 40 50 60
25 50 75 100
bG P
FA ra
te (H
z)
neuron 1
iL Q
R -V
A E
ra te
(H z)
time (ms)
neuron 12
time (ms)
neuron 110
time (ms)
Figure S11: Comparison of the firing rates inferred by bGPFA and iLQR-VAE on the first 100ms of the continuous reaching task data, for three different neurons. bGPFA learns smoother trajectories. On the other hand, fitting iLQR-VAE with a Gaussian prior with no temporal structure allows to capture more variance in the firing rates, which in turn leads to a better decoding of the kinematics.
As a further way of understanding the relative benefits and disadvantages of iLQR-VAE, we compared its performance with bGPFA, a fully Bayesian extension of GPFA (Yu et al., 2009) that enables the use of non-Gaussian likelihoods, scales to very large datasets, and was recently shown to outperform standard GPFA on this same continuous reaching dataset (Jensen et al., 2021). Importantly, bGPFA makes different assumptions to iLQR-VAE, as it places a smooth prior directly on the latents with no explicit notion of dynamics. We fit both methods using 10 minutes of data (chunked into pseudo-trials for iLQR-VAE and as a continuous trial for bGPFA). For iLQR-VAE we then performed inference and retrained the posterior covariance on the first minute of data whilst fixing the generative parameters. We found that while both methods captured similar trends in the firing rates, bGPFA yielded smoother estimates, but iLQR-VAE captured larger modulations (consistent with the higher R2 when regressing from firing rates to hand velocity). Note that the firing rate estimates here are not as smooth as for the Maze dataset (c.f. Figure 4A), because iLQR-VAE was fit using a Gaussian prior over inputs with non-zero variance at all times, effectively implying an autoregressive prior on the latent trajectories and firing rates.
From Figure S11, one can notice that bGPFA struggles to capture larger variations in the firing rate. This suggests oversmoothing, and might explain why the method does not capture hand kinematics as well as iLQR-VAE (R2 = 0.6 for bGPFA and 0.76 for iLQR-VAE.) This is indeed what we see in Figure S12.
−0.4
−0.2
0
0.2
0.4
0.6
1200 1300 1400 1500 1600
v x (a
.u .)
time (ms)
True iLQR-VAE
bGPFA
Figure S12: Ground truth hand velocity (red) and decoded kinematics using iLQRVAE (blue) and bGPFA (dotted black). A linear decoder was trained on 9 minutes of data and then evaluated on the first minute of the recording (shown here). bGPFA struggles to capture the biggest peaks in the velocity, consistent with the smoother firing rates and the lower R2.

heading: O ADDITIONAL RELATED WORK
text: In this section, we first discuss (non-exhaustively) several methods used for identifying dynamical systems from data, before presenting the few approaches we are aware of that explictly tackle the problem of inferring unobserved control inputs to those systems.
The problem of identifying the dynamics giving rise to a set of observations is one that spans many fields, from climate modelling to neuroscience, and a variety of methods have therefore been developed to tackle it. Most existing approaches assume non-driven dynamics, as this greatly facilitates systems identification.
One common modelling paradigm is to assume the data arises from a latent linear dynamical system (LDS), which parameters can be learned using an Expectation-Maximization (EM) approach Ghahramani and Hinton (1996). While linear models are typically very efficient as they allow estimates to be computed in closed-form, they severely restrict the range of dynamics that can be approximated. Various extensions have been proposed, such as switching linear dynamical systems (Linderman et al., 2017; Ghahramani and Hinton, 2000), which assume that the data can be modelled using several latent dynamical systems with a Hidden Markov Model controlling the transitions between those. Alternatively, Costa et al. (2019) proposes to use adaptative locally linear dynamics and uses an iterative procedure to find the most likely switching points. In a similar vein, Hernandez et al. (2018) approximates the dynamics as locally linear; interestingly, the proposed method (VIND) incorporates the generative dynamics in the approximate posterior distribution over latent trajectories given data. This is reminiscent of the approach taken in iLQR-VAE, where the recognition parameters are kept tied to the generative parameters.
Another way to keep the problem solvable while allowing for richer dynamics is to approximate those using a linear combination of nonlinear basis functions. This then turns the optimization into the more simple problem of learning the weights of the expansion (with the caveat that one needs to choose the set of basis functions). This is the method used in Brunton et al. (2016b), with an additional constraint that the coefficients are sparse in the space of basis functions, yielding a more interpretable model. This was later extended in Champion et al. (2019) to allow for automatic discovery of a set coordinates in which the dynamics can be approximated as sparse.
In a similar manner, a popular approach involves modelling the dynamics as linear in the space of observables (which can include linear or nonlinear mappings from the state of the system), as is done in dynamic mode decomposition Schmid (2010); Kutz et al. (2016) (see Brunton et al. (2016a) for applications to neural data). This approach is closely related to Koopman operatory theory, which finds a set of dynamic modes and uses those to approximate the data as a single linear dynamical system.
Finally, the dynamics can be modelled using nonlinear neural networks, and the parameters learned using variational methods (see e.g (Nguyen et al., 2020; Hernandez et al., 2018; Koppe et al., 2019)).
Most of the aforementioned models can be extended to incorporate known external inputs coming into the system. This is for instance done in dynamic mode decomposition with control inputs (DMDc; Proctor et al., 2016). On the other hand, the range of methods modelling dynamics driven by unobserved inputs (which must thus be inferred) is a lot more limited. Indeed, LFADS is the first method we are aware of which explicitely models the set of control inputs driving the system. As described in the main text, LFADS Pandarinath et al. (2018) models the dynamics as a (potentially) input-driven nonlinear dynamical system, and learns both the parameters and the inputs. More recently, Fieseler et al. (2020) proposed an extension of DMDc to handle unsupervised learning of unobserved signals as well as estimation of the dynamics. This was then used this to successfully model neural recordings made in C. elegans. Crucially however, the dynamics were modelled as linear, thus restricting the range of dynamics that the learnt system could generate. Morrison et al. (2020) modelled the same data using input-driven nonlinear dynamics, but assumed that there is a limited subset of inputs driving transitions at given time points, and thus only learned the magnitude of those inputs and not their timing.
Finally, an approach related to the modelling of unobserved inputs (which give rise to changes that cannot be explained by the dynamics alone) is the explicit modelling of events which lead to discontinuities in the dynamics. This is done in Chen et al. (2020) within the framework of neural ordinary differential equations Chen et al. (2018). To some extent, one can also view switching dynamical models as inferring unobserved inputs giving rise to state transitions, although hose “inputs” are restricted to live in a discrete subspace.

Review summaries:
This paper presents a new approach for inference in a model that simultaneously provides latent dynamics, initial conditions, and - importantly - external inputs. This approach is enabled by using the outcome of an optimization algorithm (iLQR) in the recognition model, recently enabled by other work in the field. The paper demonstrates the use of this approach on simulated data and real neural data, and compares to other classic and contemporary models of nonlinear dynamical systems.

The paper proposes a control-based variational inference approach that learns latent neural dynamics in input-driven SSM. It utilizes iLQR in the recognition model that transforms it into an optimal-control problem. The recognition model in the proposed method is implicitly implied by the generative model and thus reduces the number of free parameters comparing to existing methods. The proposed methods are evaluated on sythetic chaotic attractor and real-world neural recordings. 

This paper proposes ILQR-VAE, a novel method that allows to simultaneously learn latent dynamics and infer unobserved control inputs. The method relies on IQLR solver and recent advances allowing for implicit differentiation to maximize an Evidence Lower Bound on log-likelihood of observation to infer a conditional distribution over inputs as well as latent states.  Authors show comparisons to other models (the closest one being LFADS) on toy datasets and benchmark datasets for neural data analysis methods.  iLQR-VAE is on par with state-of-the-art methods on many datasets, does not require any extensive hyperparameter optimization, and allows for fast inference when dimension of the latent processes is not too high.