iLQR-VAE : control-based learning of input-driven dynamics with applications to neural data
Authors: Marine Schimel
Ta-Chu Kao
Kristopher T Jensen
Guillaume Hennequin
Venue: ICLR 2022 Oral (2022)
Keywords: neuroscience
latent variable models
RNN
VAE
motor control
control theory
dynamical systems

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, having the recognition model be implicitly defined by the generative model greatly reduces the number of free parameters and allows for flexible, high-quality inference. This makes it possible 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. 

Published as a conference paper at ICLR 2022
ILQR-VAE CONTROL-BASBI LEARNING OF INPUT- DRIVEN DYNAMICS WITH APPLICATIONS TO NEURAL DATA
Marine Schimel
Ta-Chu Kao
Department of Engineering University of Cambridge Cambridge, UK mmcs 3 cam ac.
Gatsby Computational Neuroscience Unit
University College London
London, UK
C. ikaoeucl.ac. uk
Kristopher T. Jensen Department of Engineering University of Cambridge Cambridge, UK kt ktjj1ee ccam.cc.uk
Guillaume Hennequin
Department of Engineering University of Cambridge Cambridge, UK g wennequ ngeng cam. ac uk
ABSTRACT
Understanding how neural dynamics give rise to behaviour is one of the most fundamental questions systems neuroscience. To achieve this, common ap- proach is to record neural populations in behaving animals, and model these data as emanating from latent dynamical system whose state trajectories can then be related back to behavioural observations via some form of decoding. As record- ings are typically performed in localized circuits that form only a part of the wider implicated network, it is important to simultancously learn the local dynamics and infer any unobserved external input that might drive them. Here, we intro- duce iLQR-VAE 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 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 iL.QR solutions, which is made possible by recent advances in differentiable control. Importantly, the recognition model is naturally tied the generative model, greatly reducing the number of free pa- rameters and ensuring high-quality inference throughout the course of learning. Moreover, iLQR can be used to perform inference flexibly on heterogeneous trials of varying lengths. This allows for instance to evaluate the model on single long trial after training on smaller chunks. We demonstrate the effectiveness of iLQR- VAE range of synthetic systems, with autonomous well input-driven dynamics. We further apply to neural and behavioural recordings in non-human primates performing two different reaching tasks, and show that LLQR-VAE yields high-quality kinematic reconstructions from the neural data.
control-based approach to variational inference in nonlinear
INTRODUCTION
The mammalian brain is a complex, high-dimensional system, containing billions of neurons whose coordinated dynamics ultimately drive behaviour. Identifying and interpreting these dynamics is the focus of 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

---

Published as a conference paper at ICLR 2022
generative model in the form of flexible latent dynamical system driven by process noise, coupled with an appropriate observation model.
Importantly. neural recordings are typically only made in 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 there- fore 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, generally degenerate problem that involves teasing apart momentary variations in the data that can be attributed the system's internal transition function, and those that need to be explained by forcing inputs. This distinction can be achieved by introducing external con- trol 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,, sequential variational autoencoder (VAE) that performs inference at the level of external inputs well as initial latent states. The inferred inputs were shown to be congruent with task-induced perturbations in various reaching tasks in primates (Pandarinath et al., 2018; Keshtkaran and Pandarinath, 2019). Further related work is discussed in Appendix A.
Here, we introduce iLQR- VAE, a new method for learning input-driven latent dynamics from data. As in LFADS, conditions and external inputs driving an RNN generator. However, while LFADS uses a separate, bidirectional RNN as the encoder, here substitute the inference network with optimization- based 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 KI. warmup), and greatly reduces the number of (hraatr)parrmmeters
an input-driven sequential VAE to encode observations into set of initial
While iLQ VAE could find applications in many fields as a general approach to learning stochastic nonlinear dynamical systems, here we focus neuroscience case studies. We first demonstrate in series of synthetic examples that iLQR-VA VAE can learn the dynamics of both autonomous and input- driven systems. Next, we show state-of- the art performance on monkey MI population recordings during two types of reaching tasks (O' Doherty et al., 2018; Churchland et al., 2010). In particular, show that hand kinematics be accurately decoded from inferred latent state trajectories, and that the inferred inputs are consistent with recently proposed theories of motor preparation.
2 METHOD
il HLQR-VAE models set of temporal observations, such as behavioural and/or neural recordings, through a shared input-driven nonlinear latent dynamical system (Figure S1). The input encapsu- lates 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 de- scribe the architecture of the generative model, and the control-based variational inference strategy used for training the model and making predictions. A graphical summary of the model can be found in Appendix B.
2.1 GENERATIVE MODEL
We consider the following generative model:
latent state Z+11 = fe(z,, Ue,
(I)
observations 01/21 po (2))
where Ui E RT, 2e E R. and o, E R are the input, latent state and observations at time t, respectively. Here, observations may comprise either neural activity, behavioural variables, both the distinction will be made later where relevant. We use the notation 0 to denote the set of all parameters of the generative model. We use un to set the initial condition ž fo ug 0) of the
2

---

Published as a conference paper at ICLR 2022
network'. This way, the latent state trajectory of the network z(u) (.1,..., _2r is entirely de- termined by the imput sequence u = (Mo..... ur) and the state transition function fo(-), according to Equation 1. For fo(-), use either standard linear GRU-like RNN dynamics (see Appendix C 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 Or h(Cz, b) (Appendix D).
We place a Gaussian prior over 4,50- We then consider two alternative choices for the prior over 44>0- The first Gaussian prior
pe(ur>0) = S0)
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 heavy-tailed distribution constructed hierarchically by assuming that the 7th input at time / > 0 is
Uett S.cuvv/oa
where 8; is scale factor, 6ür (0, 1) is independent across and i, and 0z x3 is shared scale factor drawn from a chi-squared distribution with y degrees of freedom. Thus, inputs are spatially and temporally independent 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 imputs when they are needed. Finally, the fact that the scale factor is shared across input dimensions means that inputs are either all weak 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:
T [ + m
1 s
(u-m)/2
po(us) [[/)
1 1
2
(5)
//
where S = diag(si .1,...,,m.) Note that both S and v are parameters of the generative model, which
we will learn.
2.2 ILQR- VAE: A NOVEL CONTROL-BASED VARIATIONAL INFERENCE STRATEGY
To train the model, we optimize 0 to maximize the log- log-likel of observing a collection of inde- pendent observation sequences O = {o(1),... (k)), }. or "trials", given by:
Σ log / po(( | ((u))po(u) du.
K
log (O)
As the integral is in general intractable, we resort to variational inference strategy by introducing a recognition model qo (u o (k)) to approximate the posterior pa (u)) Following standard prac- tice (Kingma and Welling, 2013; Rezende et al.. 2014), we thus train the model by maximizing the evidence lower-bound bound (ELBO):
L(O,0, Φ) Σ
T Ry (ucloitt [oo po (+) + (oog - log (u [())]]
(7)
k
Σ
E (alloat)
log po (0x (2)) + logpo(us) log (u (k)) Σ
(8)
< log po (O).
with respect to both and Φb
NNte that when m < n, ut can only reach an m-dimensional subspace of initial conditions, which could be limiting. We can circumvent this problem by spreading un over multiple surrogate time bins before the start of the trial, ie. introduce {u- (u-a/my.... u-2, fo together with an appropriate dependence of fo on f < 0 in Equation such that each of these surrogate inputs target different latent subspace with purely integrating ("sticking"" linear dynamics before t
3

---

Published as a conference paper at ICLR 2022
Here, the main novelty is the of optimization-ban based recognition model. We reason that maximizing the exact log posterior, i.e. computing
u (o! argmax log pa (o((((()))
(10)
argmax Σ log pa (0, (u) /0)
log po (ui)
(11)
subject to the generative dynamies of Equations 1 and 2, is standard nonlinear control problem: log po(;
(h11/u (u) acts as running cost penalizing momentary deviations between desired outputs o and the actual outputs caused by set of controls u, and log (p() acts an energetic cost those controls. Importantly, there exists a general purpose, efficient algorithm to solve such nonlinear control problems: iLQR (Li and Todorov, 2004; Appendix E). We thus propose to use a black-box iLQQR solver to parameterize the mean of the recognition density Is (o(u)o) (u) for any o, and to model uncertainty separately using a multivariate Gaussian density common to all trials. Therefore, parametrize the recognition model as follows:
gs q o)) = N(u; u (0), Σ, + Σ with u² (o) iL.QRsolveo,, (o,
(12)
(13)
where we use separable posterior covariance (the Kronecker product of spatial factor Σ, and temporal factor Σ,).
To optimize the ELBO, we estimate the expectation in Equation 8 by drawing samples from (s(u ({))) (=) and using the reparameterization trick (Kingma et al., 2015) to obtain gradients. A major complication that would normally preclude the use of optimization-ooe based recogmition 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 isole with a memory cost that does not depend on the number of iterations (Amos et al., 2018; Blondel et al., 2021; Appendix F).
2.3 COMPLEXITY AND IMPLEMENTATION
We optimize the ELBO using Adam (Kingma and Ba, 2014) with decaying learning rate 1 1/vī where i is the iteration number. Averaging over data samples can be easily parallelized; we do this here using the MPI library and local CPU cluster. In each iteration and for each data sample, ob- taining the approximate posterior mean through iLQR is the main computational bottleneck, with complexity of (T +n)n To help mitigate this cost, we find useful to re-use the previously inferred control inputs to initialize each Qsoolvv..
3
EXPERIMENTS AND RESULTS
3.1 ILQR- VAE ENABLES FAST LEARNING OF DYNAMICS
Before demonstrating the method on a number of synthetic and real datasets involving ongoing external inputs, we begin with simpler example meant illustrate some of iLQR-VAE's main properties (Figure 1). We generated data from an autonomous (i.e. non-input-driven) linear dynam- ical system (n latents, m input channels) seeded with a random initial condition in each of 56 trials. The state 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-VAR 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 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.c. they act as initial conditions for the now autonomous dynamics of the generative model. Thus, iLQR- VAE operates in regime where the output of the generator explains the data well all times, and learning
4

---

Published as a conference paper at ICLR 2022
Fast and robust learning in iLQR- VAE. Example run of iLQR- CR-VAE on synthetic dataset generated by an autonomous linear dynamical. ILQR-VAE with an adap- tive prior over control inputs (red) initially uses large inputs to fit the observations, but gradually pushes those imputs back into initial conditions as it acquires the ground truth au- tonomous dynamics. In contrast, iL.QR- VAE with rigid input prior imposing autonomous dynamics gets stuck in plateaus and learns considerably more slowly (black; text for details). For each setting, insets show the three inferred inputs for given test trial (top: posterior mean ū|o, rescaled by the maxi- mum input in the sequence), and the poste- rior predictions for the first two correspond- ing outputs (ōg 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 architec- ture in all scenarios, and the learning rate was hand-tuned for each example.
Figure 1:
forced autonomous
G
free adaplive prior
3.5
ELBO
2.5
1.5
0
0.5
10
12
14
training iteratio (xc10)
consists 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 small spectral radius) initially and therefore easily controllable.
We contrast this with learning in modified version of ILQR-VAE where allowed u4 vary freely (with a Gaussian prior of adjustable variance) but effectively fixed ut>o to be 0. In other words, constrained the dynamics of the generator remain (near-)autonomuus throughout learn- ing (Figure 1, grey box, top). Although this incorporates important information about the ground truth generator (which itself autonomous), counter-intuitively At the beginning of training, iLQR is unable to find initial conditions that would explain the data well, resulting in much higher initial loss. The model then gets stuck in plateaus that seemingly avoided by the free version of iLQR- VAE (see Figure S2 for independent repeats of this experiment).
found that impairs learning.
On the same toy dataset, we also compared iLQR- VAE to LFADS (Pandarinath et al., 2018), keeping the generative model in the same model class (see Appendix G for details). We found that LFADS learns in a similar manner to iLQR- VAE (Figure 1), also progressively doing away with inputs.
3.2 ILQR- VAE FOR NONLINEAR SYSTEM IDENTIFICATION
Next, we illustrate the method on an autonomous nonlinear dynamical system, the chaotic Lorenz attractor (Lorenz, 1963; Appendix H). This is standard benchmark 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 be learned if the trained model can recreate the whole attractor structure.
Here, we show that iLQR-VAR can learn these complex nonlinear dynamics. Before training, the inferred inputs are large throughout the trial, and explain the output observations by forcing the in- ternal state of the generator into appropriate trajectories (Figure 2A, top). At the end of learning, the inputs remain confined to the first time bin, setting the initial condition of the trajectories which are now driven by the stronger, eear-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 trajec- tory. Moreover, starting from a random initial condition, the model can recreate the whole attractor structure (Figure 2C).
5

---

Published as a conference paper at ICLR 2022
A
B
extrapolations
prediction
C
example autonomous trajectory after learning
true
01 j
0
-10
10
negaiive ELBO
4
(sleps)
2
10 100 1000 10000
training iteration
0
Figure 2: Learning nonlinear autonomous dynamics. (A) Evolution of the negative ELBO dur-
ing training of an iL-RR-VE model with nonlinear RNN and a Gaussian likelihood (n 20, m = 5, no = 3). The inputs inferred by iLQR are originally strong (blue, before learning), but are progressively pushed to initial conditions (red, after learning) as the autonomous dynamics of the Lorenz attractor are approximated with increasing accuracy. (B) Four example test trajectories (columns), conditioning on the noisy data within the first half (green shading), and predicting in the second half. (C) Single long autonomous trajectory after training (setting ueso 0), starting from a random initial condition and running the dynamics for 10000 steps. The model displays the but- terfly topology characteristic of the Lorenz system, and completes multiple cycles without deviating from the attractor, suggesting that the ground-truth dynamics have been learned.
To quantitatively assess how well the dynamics have been learned, we computed the k-step coeffi- cient of determination, RE in Hernandez et al. (2018). This metric evaluates how well the model can predict the true state & steps into the future, starting from any state inferred along a test trajectory (see Appendix H for details). Hernandez et al. reported R2g R 1 but did not show results for larger k. For iLQR- VAE, R23 0.998 and the forward interpolation was still very high at 50 time steps, with RS 0.996.
3.3 INFERRING SPARSE INPUTS
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. = 10, which evolves with linear dynamics for T = 1000 time steps (see Appendix I for an example with input-drive 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 Bernoulli distribution with mean p 0.03. Whenever an input event occurred, the magnitude of inputs in each channel was drawn from standard mutivariate Gaussian distribution.
3, m 3 and
We fit both iLQR- VAE and LEADS models to these data, choosing the generator to be within the ground-truth model class for both. ILQR-VAB captured most of the variance in the inputs (Figure 3A: R² = 0.94 + 0.02; 5 random seeds), and recovered the eigenvalue spectrum of the transition matrix almost perfectly (Figure 3B). LFADS however performed poorly this example (R² = 0.05 = 0.02 for input reconstruction; 3 random seeds), as well as in several other similar comparisons datasets of different sizes and trial numbers (Appendix J). 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. Note however that when applied to set of 56 trials of 100 time steps driven by Gaussian autoregressive inputs, iLQR-VAL VAE still captured the structure in the inputs more accurately than LFADS did (It² 0.81 = 0.01 vs. 0.29 + 0.06). We hypothesize that this reflects the difficulty of learning a good recognition model from a small amount of data. We evaluate the effect of the choice of prior more extensively Appendix K.
3.4 PREDICTING HAND KINEMATICS FROM PRIMATE M1 RECORDINGS
3.4.1 TRIAL-STRUCTURED MAZE TASK
To highlight the utility of iLQR- VLLL--VAEAAAAE for analyzing experimental neuroscience data, next applied it to recordings of monkey motor (MI) and dorsal premotor (PMd) cortices during a delayed reach- ing task ('Maze' dataset of Kaufman et al., 2016; DANDI 000128). This dataset contains 108 differ-

---

Published as a conference paper at ICLR 2022
A
true
learned
0
true [u]]
Inferred [u]]
-0.2
400
500
600
700
800
0.9
time [a.u.]
Re(A)
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. iL.QR-VAE is able to infer the timing and magnitude of the inputs almost perfectly, despite being trained on only single timeseries of 1000 time steps. Note that -LQR-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 similarity transformation.
ent reach configurations over nearly 3000 trials, and has recently been proposed as a neuroscience benchmark for neural data analysis methods (Pei et al., 2021). We compared the performance of iLQR-VAE to several other latent variable models, evaluated on this dataset in Pei et al. (2021).
Consistent with previous findings (Pandarinath et al., 2018), iLQR- VAE inferred inputs that were confined initial conditions, from which smooth single-trial dynamies evolved near-autonomously (Figure 4A). As a first measure of performance, we evaluated the models on "co-smoothing", i.e the ability to predict the activity of held-out neurons conditioned on a set of held-in neurons (see Appendix L for details). Conditioning of 137 neurons (i.c 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), 7), 0.329 for Neural Data Transformers (Ye and Pandarinath, 2021) and R² 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 metric of interest to neuroscientists. We applied ridge regression to predict the monkey's 's hand velocity (with 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 Rt2 = 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- BR-VEE par with state- of-the-art methods, without any extensive hyperparameter optimization.
3.4.2 CONTINUOUS REACHING TASK
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 reach to consecutive targets randomly sampled from 17x8 grid on 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 and y-velocities of the fingertip.
VAE to self- self-p reaching task during which a monkey had to
In this example, we experimented with fitting the spike trains and hand velocities jointly (combining a Poisson likelihood for the 130 neurons and Gaussian likelihood for the 2 kinematic variables, see Appendix M for further details). We found that it allowed iLQR-VAL VAE to reach a similar kinematic decoding performance when fitting neural activity alone, but using a smaller network. More gen- erally, we reason that a natural approach to making behavioural predictions from neural data using probabilistic generator is to fit it to both jointly, and then the posterior predictive distribution
7

---

Published as a conference paper at ICLR 2022
maze task
continuous reaching
A
B neuron 16
D
120
inferred inout
neurons
hand speed
rate (H2)
40
v. (a.u.)
neuron 44
-300 target
+300 ms
true interred
onset
(a.u.)
(a.u)
-250
+450
target
+750ms
-300
move
+300 ms onset
time (s)
Figure 4: iLQR-VAR can be used to decode kinematics from neural data, and learns dynamics relying strongly on preparatory inputs. (A) 50 example hand trajectories (top) from the monkey reaching 'Maze' dataset, corresponding single-trial firing rate timecourse (middle; one example neu- ron), and inferred input (bottom; one example input dimension). (B) Mean (thick) and single- 'trial' (thin) firing rates of two example neurons during reaches of various directions (colours), aligned to target onset. Note that three single- "tglle'trrill' firing rates are shown for only three of the 8 reach directions for which averages are shown. Interestingly, single 'trial' activities evolve tightly around their trial averages, and resembles the firing rate responses shown in A (middle). (C) Example spike raster (top) and hand kinematics (bottom) for a second-long chunk of test data in the continuous monkey reaching task. Un and by refer to hand x- and y-velocities respectively. (D) Overall magni- tude of the inferred inputs ||u, (bluc), average population spiking activity (red), and hand velocity llyell ll (black), each z-scored, averaged across movement episodes, and aligned to target onset (top) or movement onset (bottom).
over behavioural variables (conditioning spike trains only) as 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., 2022).
For our analyses, we used the first 22 minutes of a single recording session (indy 201 60426), excluded neurons with overall firing rates below 2 Hz, and binned data at 25 ms resolution. Although it is not formal requirement of our method, we chunked the data into 336 non-overlapping pseudo- trials of 4 s each, in order to enable parallelization of the ELBO computation during training. We only trained the model on a random subset 168 trials.
To highlight the flexibility of iLQR forming inference on the first 9 minutes of the data, as a single long chunk of observations. Note that this is generally possible in LFADS any sequential VAE where encoder RNN has been trained exclusively on trials of the same fixed length. Despite the lack of trial structure, we found that neurons display stereotyped firing pattern across multiple instances of each reach. This was revealed by binning the angular space into 8 reach directions, temporally segmenting and grouping the inferred firing rates according 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; R² 0.75 0.01 over 5 random seeds), on-par with Autol (R² 0.76; Keshtkaran et al., 2021) 1), and considerably higher than GPFA and related approaches (R2 0.6; Jensen et al., 2021).
recognition model, we then evaluated the model by per-
We next wondered whether we could use ILQR-VAE to address an open question in motor neu- roscience, 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 prepara- tion, 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), con- sistent with the literature (Shenoy et al., 2013; Churchland et al., 2012). Interestingly, we observed
8

---

Published as a conference paper at ICLR 2022
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, performed using information from the whole duration of the trial (i.e. using iI.QR 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 MI, whose dynamics then evolve in a more autonomous manner to drive subsequent motion.
inference was
4 DISCUSSION
LIMITATIONS AND FUTURE WORK
While we have demonstrated that iLQR-VAE performs well on various toy and real datasets, the method has number of limitations, some of which could be addressed in future work. Firstly, the problem of decoupling ongoing imputs from dynamics is degenerate in general, and there is no guarantee that ILQRR-AE 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. While further generalization tests such as extrapolations can be used to assess post-hoc how well the dynamics have been learned, the lack of identifiability will often make interpretation of the model parameters difficult. Secondly, using iLQR as way of solving maximum a posteriori inference in state space models comes at high computational cost, and with the risk that iLQR may converge to a local minimum. We 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). Finally, for simplicity we modelled posterior input uncertainty using 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, left to future work, would be amortize the computation of the posterior uncertainty by reusing some of the computations performed in iLQR.
CONCLUSION
The rise of new tools and software now makes 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 hard nonlin- ear 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, 6016; Levine, 2018; Appendix N) 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 inference over inputs in an RNN-based generative model. Using an optimization- 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 reduced (Appendix O), making learning more robust (in particular, to initialization) and reducing the number of hyperparameters optimizers (Blondel et al., 202 1), we therefore hope that optimization recognition models will open up new avenues for VAEs.
recognition
tune. With the advent of automatically differentiable
ACKNOWLEDGEMENTT
We thank Jasmine Stone and Javier Antorán for helpful comments on the manuscript, and Anil Mad- havapeddy for providing computing resources at early stages. M.S. was funded by an EPSRC DTP studentship and K. K.T.J. was funded by Gates Cambridge scholarship. This work performed

---

Published as a conference paper at ICLR 2022
using resources provided by the Cambridge Tier-2 system operated by the University of Cambridge Research Computing Service (htp://www.hne..amm..uuk funded by EPSRC Tier-2 capital grant EP/P020259/1.
REFERENCES
Amos, B., Rodriguez, L. D. J., Sacks, J., Boots, B., and Kolter, J. Z. (2018). Differentiable MPC for end-to-end planning and control. arXiv preprint arXi::1810.1344
Archer, E., Park, M., Buesing, L., Cunningham, J., and Paninski, L. (2015). Black box variational inference for state space models. arXiv preprint arXiv:1511.0755
Blondel, M., Berthet, Q. Cuturi, M.. Frostig, R., Hoyer, S., Llinares-López. F., Pedregosa, F., and Vert, J.-P. (2021). Efficient and modular implicit differentiation. arXiv preprint arXiv:2/05.15183
Buesing, I., Macke, J. H., and Sahani, M. (2012). Learning stable, regularised latent models of neural population dynamics. Network: Computation in Neural Systems, 23:24
Champion, K., Lusch, B., Kutz, N., and Brunton, S. L. (2019). Data-driven discovery of coordinates and governing equations. Proceedings of the National Academy of Sciences, 116(45):24445222511
Churchland, M. M., Cunningham, J. P., Kaufman, M. T., Foster, J. D., Nuyujukian, P., Ryu, S. I., and Shenoy, K. V. (2012). Neural population dynamics during reaching. Nature, 487:51-56.
Churchland, M. M., Cunningham, J. P., Kaufman, M. T., Ryu, S. I., and Shenoy, K. (2010). Cortical preparatory activity: representation of movement or first cog in a dynamical machine? Neuron, 68(3):887-400.
Cremer, C., Li, X., and Duvenaud. D. (2018). Inference suboptimality in variational autoencoders. In International Conference on Machine Learning, pages 1078-1086. PMLR.
Duncker, L., Bohner, G., Boussard, J., and Sahani, M. (2019). Learning interpretable continuous- time models of latent stochastic dynamical systems. In International Conference on Machine Learning, pages 1726-1734.
Dvijotham, K. and Todorov, E. (2013). Linearly solvable optimal control. In Reinforcement learning and approximate dynamic programming for feedback control, volume 17, pages 119-141. Wiley Online Library.
Gao, Y., Archer, E., Paninski, L., and Cunningham, J. (2016). Linear dynamical neural population models through nonlinear embeddings. arXiv preprint 1605v:6005.8...
Ghahramani, and Hinton, G. E. (2000). Variational learning for switching state-space models. Neural comput., 12:831-864.
Hernandez, D., Moretti, A. K., Wei, Z., Saxena, S., Cunningham, J., and Paninski, L. (2018). Non- linear evolution via spatially-dependent linear dynamics for electrophysiology and calcium data. arXiv preprint arXiv: arXhv:1811.0244
Jensen, K. T., Kao, T.-C., Stone, J. T., and Hennequin, G. (2021). Scalable Bayesian GPFA with automatic relevance determination and discrete noise models. Advances in Neural Information Processing Systems, 34.
Jun, J. J., Steinmetz, N. A., Siegle, J. H., Denman, D. J., Bauza, M., Barbarits, B., Lee, A. K., Anastassiou, C. A., Andrei, A., Aydmn, Ç., et al. (2017). Fully integrated silicon probes for high- density recording of neural activity. Nature, 551(7679):232-233
Kappen, H. J. (2005). Linear theory for control of nonlinear stochastic systems. Phys. Rev. Lett., 95:200201.
Kappen, H. and Ruiz, H. C. (2016). Adaptive importance sampling for control and inference. Journal of Statistical Physics, 162:1244-1266.
10

---

Published as a conference paper at ICLR 2022
Kaufman, M. T., Seely, J. S., Sussillo, D., Ryu, S. L., Shenoy, K. V., and Churchland, M. M. (2016).
The largest response component in the motor cortex reflects movement timing but not movement type. eNeuro, 3(4).
Keshtkaran, M.
R. and Pandarinath, C. (2019). Enabling hyperparameter optimization in sequential autoencoders for spiking neural data, arXiv preprint arXiv: 1908. 07896.
Keshtkaran, M. R., Sedler, R., Chowdhury, R. H., Tandon, R., Basrai, D., Nguyen, S. I.,, Sohn, H., Jazayeri, M., Miller, I.. E., and Pandarinath, C. (2021). A large-scale neural network training framework for generalized estimation of single-trial population dynamics. bioRxiv.
Kim, T. D., Luo, T. Z., Pillow, W., and Brody, C. (2021). Inferring latent dynamies underlying neural population activity via neural differential equations. In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceed- ings of Machine Learning Research, pages 5551-5561. PMLR.
Kingma, D. P. and Ba, J. (2014), Adam: A method for stochastic optimization. arXiv preprint arXiv: 4412.6980.
Kingma, D. P., Salimans, T., and Welling, M. (2015). Variational dropout and the local reparame- terization trick. arXiv preprint arXiv:1506.02555
Kingma, D. P. and Welling, M. (2013). Auto-encoding variational Bayes. arXiv preprint arXiv:1312.6114.
Kutschireiter, A., Surace, S. C., and Pfister, J.-P. (2020). The hitchhiker's guide to nonlinear filtering. Journal of Mathematical Psychology, 94: 102007.
Levine, S. (2018). Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint 188iv:/805.00009.
Li, W. and Todorov, E. (2004). Iterative linear quadratic regulator design for nonlinear biological movement systems. ICINCO (I), pages 222-229.
Linderman, S., Johnson, M., Miller, A., Adams, R., Blei, D., and Paninski, L. (2017). Bayesian learning and inference in recurrent switching linear dynamical systems. In Artificial Intelligence and Statistics, pages 914-922. PMLR.
Lorenz, N. (1963). Deterministic nonperiodic flow. Journal of atmospheric sciences, 20(2):130- 141.
Makin, J. G., O'Doherty, J. E., Cardoso, M. M., and Sabes, N. (2018). Superior arm-movement decoding from cortex with a new, unsupervised-leumiin algorithm. Journal of neural engineer- ing. 15(2):026010.
Mathis, A., Mamidanna, P., Cury, K. M., Abe, T., Murthy, V. N., Mathis, M. W., and Bethge, M. (2018). DeepLabCut: markerless pose estimation of user-defined body parts with deep learning. Nature neuroscience, 12(9):2281-12)
Musk, E. et al. (2019). An integrated mrainnmachine interface platform with thousands of channels. Journal of medical Internet research, 21:e16194.
Nguyen, D., Ouala, S., Drumetz, L., and Fablet, R. (2020). Variational deep learning for the iden- tification and reconstruction of chaotic and stochastic dynamical systems from noisy and partial observations. arXiv preprint 2=Xiv:2009.022900
O'Doherty, J. E.,, Cardoso, M. M. B., Makin, J. G., and Sabes, P. N. (2018). Nonhuman Primate Reaching with Multichannel Sensorimotor Cortex Electrophysiology: broadband for indy 20160630 01. This research supported by the Congressionally Directed Medical Re- search Program lyzed Veterans of America. JGM was supported by fellowship from the Swartz Foundation.
JEO was supported by fellowship #2978 from the Para-
Pandarinath, C., O' Shea. D. J., Collins, J., Jozefowicz, R., Stavisky, S. D., Kao, J. C., Trautmann, E. M., Kaufman, M. T.. Ryu, S. I., Hochberg, L. R., et al. (2018). Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10))805-815.
11

---

Published as a conference paper at ICLR 2022
Pei, F., Ye, J., Zoltowski, D., Wu, A., Chowdhury, R. H., Sohn, H., O'Doherty. E.,, Shenoy, K. V., Kaufman, M. T., Churchland, M., et al. (2021). Neural latents benchmark 221: Evaluating latent variable models of neural population activity. arXiv preprint arXtv:2109.04663
Rezende, D. J., Mohamed, S., and Wierstra, D. (2014). Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pages 1278-1286 PMLR.
Schroeder, K. E., Perkins, S. M., Wang, Q., and Churchland, M. M. (2022). Cortical control of virtual self-motion using task-specific subspaces. Journal of Neuroscience, 42(2):220-239,
She, Q and Wu, A. (2020). Neural dynamics discovery via Gaussian process recurrent neural networks. In Uncertainty in Artificial Intelligence, pages 454-664 PMLR.
Shenoy, K. V., Sahani, M., and Churchland, M. M. (2013). Cortical control of arm movements: dynamical systems perspective. Annu. Rev. Neurosci., 36:337-359,
Stosiek, C., Garaschuk, O., Holthoff, K., and Konnerth, (2003). In vivo two-photon calcium imaging of neuronal networks. Proceedings of the National Academy of Sciences, 100(12):7319- 7324.
Todorov, E. (2009). Efficient computation of optimal actions. Proc. Natl. Acad. Sci., 106:1 1478- 11483.
Toussaint, M. (2009). Robot trajectory optimization using approximate inference. In Proceedings of the 26th annual international conference on machine learning, pages 1049-1056.
Ye, J. and Pandarinath, C. (2021). Representation learning for neural population activity with neural data transformers. bioRxiv.
Yu, Byron, M., Cunningham, P., Santhanam, G., Ryu, S. I., Shenoy, K. V., and Sahani, M. (2009). Gaussian-proces: factor analysis for low-dimensional single-trial analysis of neural population activity. In Advances in neural information processing systems, pages 1881-1888.
12

---

Published as a conference paper at ICLR 2022
Appendix
A ADDITIONAL RELATED WORK
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 devel- oped 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 sys- tem (LDS), which parameters can be learned using an Expectation ceteraation-Maxmizazainn (EM) approach Ghahramani and Hinton (1996). While linear models are typically very efficient as they allow es- timates 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 sys- tems (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 tran- sitions between those. Alternatively, Costa al. (2019) proposes 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 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- EL-R-VAE, where the recognition parameters are kept tied to the generative parameters.
locally linear; interestingly, the proposed
Another way to keep the problem solvable while allowing for richer dynamics to approximate those using& 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 similar manner, popular approach involves modelling the dynamics as linear in the space of observables (which can include lincar or nonlinear mappings from the state of the system), as is done in dynamic mode decomposition Schmid (2010); Kutz et al. (2016) (see Brunton al. (2016a) for applications to neural data). This approach is closely related to Koopman operatory theory, which finds set of dynamic modes and uses those to approximate the data as 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; Hermandez 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 (DMDe; Proctor et al., 2016), which can be generalized into Koopman operators with inputs and control (KIC; Proctor et al., 2018).
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 (Pandarinath et al., 2018) 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 models the dynamics as (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 limited subset of inputs driving

---

Published as a conference paper at ICLR 2022
transitions at given time points, and thus only learned the magnitude of those inputs and not their timing.
Finally, approach related to the modelling of unobserved inputs (which give rise 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 state transitions, although those "inputs" are restricted to live a discrete subspace.
B GRAPHICAL SUMMARY OF THE MODEL
generation
RNN
observations o
neurons
input
II
time
latent state
inference (iLQR)
Figure S1: Illustration of the model iLQR-VAE trained to model a of noisy observations (here, spike trains). During each training iteration, iLQR is used to infer the input (green), given observations and the current parameters the generator. Intuitively, the optimal inputs are ones that produce latent trajectories in the RNN that are most compatible with the data, without overfitting. Predictions are generated by running the dynamics forward, conditioned on a given input and set of parameters.
C IMPLEMENTATION OF THE DYNAMICS
We considered different functional forms for the discrete-time time dynamics of the latent state. In the following, 2, and u, denote the latent state and external input at time t, respectively.
C.1 LINEAR DYNAMICS
The simplest case considered is that of linear dynamics:
#t+1 Azt Bu,
(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 ||=|| and the associated gradients. This can become problematic, especially when modelling long sequences of observations. To circumvent this issue, used parametrization of the propagator A that ensured it remained stable at all times. To find a stable linear parametrization, we considered the Lyapumov stability condition (Bhatia and Szegö, 2002). The discrete time dynamics of Equation S1 are asymptotically stable if and only if A satisfies
P APAT =I
(S2)
2

---

Published as a conference paper at ICLR 2022
for some positive definite matrix P with eigenvalues 1. It is easy to verify that the following parameterization of the state matrix satisfies this criterion:
Q(( 1)
(S3)
with U and Q arbitrary unitary matrices, and D arbitrary non-negative diagonal matrix. Con- versely, 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)U Finally, are also learning the B and C matrices Equation S1, we can without loss of generality set U
C.2 GRU DYNAMICS
To fit the monkey reaching data well the Lorenz attractor, chose the dynamical system be Minimal Gated Unit (MGU). More specifically, we used the MGU2 variant of the MGU proposed Heck and Salem (2017):
oU=oUU 3; = -U 2 2:-1) Wa, + br)
(S4)
(S5)
2. - (1- 21-1 2:-1
(S6)
where æ; Bu, denotes the input entering the dynamical system. Note that the latent state is often denoted by the literature. We found that the MGU2 gave better and more stable perfor- mance than the MGU. We hypothesize that this is due 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 o-) to be sigmoid function, and g(-) to be a soft Rel.u-ikk nonlinearity,
I+ Vx2 + 1.
g(x)
(S7)
D LIKELIHOOD FUNCTIONS
The likelihood of the observations appears 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 zy. Although these quantities can be obtained generically via auto- matic differentiation, iLQR is always faster when they are provided directly (Appendix E), which we did here using the analytical expressions given below.
D.1 GAUSSIAN LIKELIHOOD
For the Gaussian likelihood, we assume observations linearly decoded from latents and corrupted with Gaussian noise, such that o ~ N(Cz + b, E), with the readout matrix, a vector of biases, and Σ diagonal matrix of variances. This yields the following log-lielihood function:
log P(o, = - (Cz, + -0) (C2, b o,) + no log((n)) + Σ log Eni
(S8)
The Jacobian of this expression is as follows
i log P(o, zx) Oz,
-C"'-'(Zz++bo
(S9)
Finally, the Hessian is given by:
8 log P(or|z)) = - 02,02?
(S10)
D.2 POISSON LIKELIHOOD
To model spike trains, we assume that they are generated by a Poisson process with an underlying positive rate function for neuron given by:
14 3.f((Cz) A
(S11)
3

---

Published as a conference paper at ICLR 2022
where R" - Rt is nonlinear function (chosen to be exponential when modelling the monkey recordings, and a soft ReLU-like nonlinearity elsewhere), A denotes the time bin size, and 3 is neuron-specific gain parameter. This yields the following log-likelihood:
log P(o.)z))
Σ (o, log μ4 - p4 + log
(S12)
where the sum is performed over neurons. Using the shorthand notations h(z) log fx) and a, = Ozz + b, the Jacobian and Hessian this expression are given by
i log P(o)lz))
CT [ox o - 3 f(a)))
(S13)
Oz,
8 log P(or
CT [diag(or or o ha))) - A diag(3 @ ff (a)))
(S14)
0z,0zi
E ILQR ALGORITHM
Our recognition model makes of the iterative Linear Quadratic Regulator algorithm (iLQR; Li and Todorov, 2004; Tassa et al., 2014) to find the mean of the posterior distribution ga(u) Iterative LQQ is used 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 quadratic approximation 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 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 iLQQR (see Algorithm 1).
The Linear Quadratic Regulator concerned with finding the set of controls u E R.th that minimize quadratic cost function C(u) under deterministic linear dynamics, given by:
T C(u)
Σ5
C, 21 + u. C. un + IT ".', + u. C 21) + 2 c + u
(S15)
2
1=
s.t. 2++1 = A,2, + B + h4-
(S16)
Here, A, E Brxa a (possibly time-dependent) transition matrix, B, Raxm represents the input channels at time I, and h. is state and input-independent term. Note that E R" is a deterministic function the initial condition žg and the sequence of inputs u. LQR finds the inputs minimizing Equation S15 using dynamic programming approach, by recursively finding the feedback rule (K,, k2) which gives the optimal inputs to minimize the cost-to-go at each time a u, K, + kt- Details can be found in function on Backward in Algorithm 1.
iLQR an extension of LQR to general dynamics and cost functions. Specifically, iLQR minimizes
T-1
(u() = Σ ro(zi, t) subject to 21+1 fo u4
(S17)
where 0 denotes set of parameters. At iteration i, iLQR approximates both the dynamics and the cost around the current trajectory T' = (z', ;2,,00 u²) as:
(6 Su, t) R fo( + (V.fo) 8z + (Vufo) Tou
(S18)
and
(8zz, ou t) = ro(T))+
2 12 (V 0oz 282 (V re ro) - + dut (V2 (enrre)8u
(S19)
+ +zz (V.ro) + Su (V (V.ro)
(S20)
Here, 8z and su refer to perturbations around the current nominal trajectory, and all V operators correspond to partial differentiation evaluated at the current nominal trajectory (u', x²) and corre- sponding time
4

---

Published as a conference paper at ICLR 2022
The above equations readily identified local LQR problem of the form of Equation S15, which can thus be solved using standard dynamic Programming tools. Once dua minimizing Equa- tion S19 has been computed, the imputs are updated as u u + Óu", and the new state trajectory follows from simulating the dynamics forward with these new inputs. After each LQR update, we thus obtain trajectory - and the process repeats until convergence some locally optimal trajectory T"
Implementation details can be found in Algorithm 1. Note that the backward LQR pass involves inversion of the matrix Que. (defined in Algorithm funct ion Backward). Depending on the specific form of the iLQR cost function, this might not always be positive-defnnite Therefore, Sarterrg-Mruurrd--yyy -ype regularizer (not described in the pseudo-code) Qun < Quv + XI to maintain positive definiteness. Thus, iI.QR effectively reverts to first-order gradient descent, opposed to second-order optimization, whenever the locally quadratic approxi- mation is a bad one.
we include an adaptive
F DIFFERENTIATING THROUGH ILQR
Here we discuss how to efficiently differentiate through the iLQR algorithm. This becomes nec- essary when one wishes to differentiate through a function involving an iLQRsolve, such as the posterior mean of our recognition model (Equation 11). While naive but simple strategy to achieve this would be unroll the algorithm and gather gradients for every step, this is expensive both com- putationally and memory-wise. Amos et al. (2018) derived a way to analytically obtain gradients with respect to parameters of iLQR, at the cost single LQR pass. Specifically, differentiat- ing through iLolve is achieved by running iLQR to convergence, forming linear-quadrate approximation around the converged trajectory, following the steps described Appendix E and dif- ferentiating through the corresponding LQR problem. Below, provide alternative derivation to Amos et al.' 's of the gradients of an LQR solution.
F.1 LQR OPTIMALITY CONDITIONS
We now introduce use the more compact notation T; which will be used in the rest of this section.
Zt U4
F,
[A,
B.). C.
=
C Csu
Cuu
C
As described in Appendix E, the horizzho discrete-time LQR problem involves minimizing:
= Σ T C.T. + T
(S21)
subject to constraints on its dynamics
T++1 F.T, + fcr
(S22)
following the notation from Appendix E. To solve this problem, write down the Lagrangian:
T C= Σ 2 C..++ c)T) + Σ (T fc T(+1).
(S23)
L=
where A1, A2. XT are adjoint (dual) variables that enforce the dynamic constraint. Differenti- ating with respect to A, and T, enables 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 (C + c.) + F. A - A =
(S24)
FIT. + ft [I 0) T =
(S25)
CTTT + cT
[I 0 XT 0
(S26)
Zn
[
0
(S27)
5

---

Published as a conference paper at ICLR 2022
Algorithm il. il.QRsolve(Co(u)) )ii) with a IR' and Co defined in Equation S17.
Parameters: 7 7
iLQR
T = Rollout(u for 2 converged do
for f = 0 to I" do
F V fo, Ft Vafe
C vrro, C = vuro, C²²² = Viro. C.Lb = Viro. C.¹² = V Nz ro
end for
11]: K 10:T-1] Backward( Fio:T 10:T7 ()r c00:z C [0:2]
C
[0:7]
Cuu [0:T]
T' Forward (o:r- 110:1-1);;7 7
end for
function Rollout(u)
Zo = 0 for / 1 to T do
Zr+1 = fo(Zt, Ut)
end for return (z,
function Backward( To:TI 0.r]. 9o:T1. VT C., VT
Fu
(0orr)
C
(0:1]:
Cuz
[0:2]
[0:T]
CF
for | = T 1 to do
Q =C²² F
V F
Q CH Fu V
Q C + Fu V4+ FR
q. c₁ Fo0411
q. C + F'4411
K; = (Q!") Q kE
A (QU") L
V. Q + 2 K₁
U: qi + K. g
end for
return k 1), K [0:T-1]
function Forward 10: k(o:-11111
{u, x})
=
repeat
20 = 0
to aki
for / = to T do
2. f0 ú;
ú, = Ki(2 zx. ak;
end for
until (ū) < Cul
return (2,

---

Published as a conference paper at ICLR 2022
Rearranging, rewrite the KKT conditions in matrix form
C
F
Ca
F
0
[-I 0)
A+t1
ft
(S28)
[-I
0J²
C.+1
F 0
T++1
C++1
F+11
A++2
fi+1
These optimality conditions are satisfied for the solution to the optimization problem p" (Ti. Tz. .t... Xp). Equation S28 implies that the solution of the LQR problem p will sat- isfy:
p²=-K- q.
(S29)
Computing this quickly becomes infeasible as K grows with long-time horizons, and Equation $28 typically solved linear time using a dynamic programming approach, as described in Ap- pendix E.
F.2 BACKPROPAGATING THROUGH THE LQR SOLVER
Differentiating through LQR solve boils down differentiating through the backsolve in Equa- tion S29. In the following, we denote the adjoint of parameter 0 as 0. From Giles (2008), we know that the adjoint of the backsolve operation given by:
P. K- pp qp"
q = - K T
(S30)
(S31)
We note that Equation S30 has the same form Equation S29, which means can compute q = (-.. Cr. )"' by solving another LQR problem. After solving for T, we can then compute K as outer-product of with y to get:
Kc. K K F.
(S32)
Koa K Futs
KFy
K
:
C₁
ft
Tt T
Tê-
Xisz T ...)
(S33)
C++1
fe+l
C
CAtTT
(S34)
C4117H11 T c411711"
C411A111 T
T

---

Published as a conference paper at ICLR 2022
Collecting all the gradients of C₁ and Fú, arrive
C ; (Kc. *K²)) ( + T/cc)
(S35)
2 F. = Kr. + RT F? SITIT!" Xi+1
(S36)
Note that we have symmetrized the adjoint of C₁ which ensures that C₁ remains symmetric after each gradient update. The antisymmetric part of C, does contribute the LQR cost.
Finally, one subtlety arises from the fact that Equation S22 and Equation S21 are written function of in the general LQR setting. In the iLQR case however, the LQR problem is local at each iteration, and år vanishes at convergence. If we denote by the last iteration before declaring convergence, one can however write the problem function of the variable of interest r", using:
J
( (T
)C((T
- +8) + - -())
(S37)
t=
CiTi + (cf C())T; + cst Σ 2
(S38)
subject to constraints on its dynamics
1 + FiT
(S39)
This implies that the values for C and need to be adjusted accordingly, such as reflect the switch of variable from åt during the optimization to the fixed point r to compute gradients. Note that at convergence we can use T¹ Ae T", giving access to all the necessary variables to compute gradients with respect to 0.
G DETAILS OF EXPERIMENT
Figure S2: Illustration of the learning curves for 5 different runs of the "forced au- tonomous" iLQRR-VE model. The model consistently gets stuck in plateaus during the optimization, leading to a slower convergence than their "free" counterparts (see Figure 1).
nogative ELBO
5
10
15
20
iteration number x10
The data in Section 3.1 was generated from an autonomous linear dynamical system with n 8, m = 3, and n. 8 where n, is the dimension of the observation space. All the models were fit using the dynamies within the ground-truie -trth model class, i.e with linear dynamics, n 8, m 3, and n. = 8. We optimized the model parameters with Adam, using (manually optimized) learning rates of 0.04/(1+ /k/1) for the free iL.QR-VAE model, 0.04/(1+ Vk/I for autonomous il iL.QR-VAL RAE 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). Note that while all methods run in similar wallclock time this example, this will ultimately be implementation and data-dependent.
In Figure $2, we show additional learning curves for the "forced autonomous" models; these show that, even for different initializations and trajectories through the loss landscape, the model consis- tently gets stuck plateaus. This can be contrasted with the free-form iLQR- VAE models.
H FURTHER DETAILS OF LORENZ ATTRACTOR
The chaotic Lorenz attractor consists of a three-dimensional state (E1, E2, Ea) evolving according
10( (1)
i = (28 (x) - E2
Eg 6162 8lg/3
(S40)
8

---

Published as a conference paper at ICLR 2022
For example, generated data by integrating Equation S40 long time period using a Runge-Kutte 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 iL.RR-VAR VAE. on this dataset (Figure 2B, bottom). We then fitted these data using GRU dynamics with n. 20 and 5.
The normalized ssstep mean equa-equared error was defined as follows:
MSE
Ct
(S41)
MSE
R
(S42)
-t=
|| a||²
where Ee4. the prediction time k. and the mean for this trial.
I LEARNING INPUT-DRIVEN NONLINEAR DYNAMICS
To bridge the gap between autonomous nonlinear dynamics (see Section 3.2) and real data, evaluated iLQR-VAL VAE on an input-driven nonlinear system, the Duffing oscillator. We generated Duffing trajectories that included a perturbation of the Duffing state way through. We then embedded those 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.
A
X2
B
200
neurons
100
X1
50
PC
negative ELBO
0.17
-5
PC2
50
0.168
50
|||||
10
0.166
200
400
600
800
1000
100
200
training iterations
time [a.u.]
Figure S3: (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 ground- truth 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 shown in gray. Bottom: norm of the inferred input.
PCs (after the perturbation)
The dynamics of the Duffing oscillator are given by
(S43)
To generate each training sample, we integrated Equation $43 from two different random initial conditions for 100 time steps each using Runge-Kutte solver (RK4) with dt = 0.03. Example such trajectories are shown in Figure S3A (top); note that each trajectory can be understood as the evolution of the system in state-space for given energy level of the oscillator. We then concatenated those two trajectories to yield a single trajectory of 200 steps with a perturbation the middle. We then linearly mapped the now-dimensional oscillator onto 300-dimensional state, before passing through the nonlinearity of Equation S7 to obtain a set of firing rates, which then gave rise to observations via Poisson process (Figure S3B, top).

---

Published as a conference paper at ICLR 2022
We generated 112 training and 1 12 testing trials in this way. We fit these data using iLQR- with n = 20, m = 4, and found that could successfully infer the latent trajectories (see Figure S3B, middle). Importantly. iLQR- VAE learned to fit most of the trajectories 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 S3B, 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.
J COMPARISON OF LFADS AND ILQR- VAE ON A TOY INPUT INFERENCE TASK
A
true learned
true |[u]
inferred ||u|
0.2
time [a.u.]
Re(A)
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 http:////gttthb.ccoooooogeeeeeee wtaptter//dadd_ttttorreesee which we modified
computation-tie coon-thru-dyynn dynami to include linear dynamics and Gaussian likelihoods. We then evaluated the quality of the input reconstruction by measuring how much input variance captured by the models. We report this as the R² from inferred to true inputs.
We used generative model within the ground truth model class. For each dataset, we performed a LFADS.
search to choose the pest-performing encoder architecture and learning rate for
Results of this experiment are summarized in Table S1. iLQR-VAC 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 VAE
1x1000
0.94 0.01
10x100
0.04 0.06
32x
0.29
0.83
56x100
0.80 0.08
10x200
0.20
0.93 0.01
AR 56x100
0.28 0.02
0.81 + 0.02
Table S1: Comparison of iLQR- VAE and LFADS on input inference tasks. Results are reported R² (mean + sem) over 3 random seeds for each.
Our results suggest that LFADS' performance improves with larger amounts of data. More surpris- ingly, LFADS also seems to perform better when the data is split into shorter trials. In particular, we found 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, il --QR-VAE inferred
10

---

Published as a conference paper at ICLR 2022
eutoregressive
sparse
A
D
E
aba
time ims
time ims
1000
Gaussian Student
Gluessaa
obss
1009
abs
log
10
time Ims
10
20
30
40
50
sime /im
0
10
20
30
50
Beeaaiio
Itertion
Figure S5: Comparison of the Student (pink) and Gaussian (green) prior for learning a linear dy- namical system driven by autoregressive inputs (AR) or sparse inputs. (A) Fit of the observations using the Gaussian (top) and Student (bottom) priors. Both allow to fit the observations highly ac- curately. (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 of v. (C) Evolution of the loss for both choices of prior. The loss curves closely aligned, and both models converge to similar ELBO value. (D) Fit of the observations using the Gaussian (top) and Student (bottom) priors. Both models allow fit the observations, but the Student prior allows to obtain 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 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 higher ELBO value.
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' lower R², is the prior over inputs (auto-regressive prior for LFADS and Student for (LQR-VAE) In Figure $4 we show example of LFADS, one of the test examples of the S 56x 100 dataset. In this example, LFADS infers its largest input concurrently to the ground truth input, but also infers small inputs when there are however that this is not the only effect at play here, as emphasized by the lower performance on the AR dataset. The impact of the choice of prior in iLQR-VAE. is discussed further in Appendix K.
in the ground truth. This has a significant impact
the R² metrie. Note
K COMPARISON OF THE STUDENT AND GAUSSIAN PRIORS
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, n. 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 first example, both priors yield extremely similar results (see Figure S5(A-C)). Indeed, the Student prior learns 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 see in Figure S5(D-F), the Gaussian prior learns large variance fit the sparse inputs, leading to higher baseline noise than in the true system.
As shown here, the Student prior offers more flexible model, the Gaussian case is recovered for large и values. Note however that using the Gaussian prior ensures that the input term in the iLQR cost function always convex in u, which can facilitate the optimization and allow il IIQQR to
11

---

Published as a conference paper at ICLR 2022
converge faster in cases. Moreover, in the case of autonomous dynamics (e.g Lorenz attractor and Maze dataset) both priors will converge to the same solution.
L FURTHER DETAILS OF SINGLE TRIAL ANALYSES
L.1 BENCHMARKING AGAINST EXISTING METHODS
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://gutt dandia archirrchhiii ve. /
We used 1720 training trials and 510 validation trials, which were drawn randomly for each instanti- ation of the model to avoid overfitting to test data. The risk of overfitting to the dataset was lowered by the fact that iLQR- VAE requires very little hyperparameter optimization. For this experiment, fitted iLQR- VAE to the neural activity using model with MGU dynamics (n = 60), a Student prior over inputs (m 15), and a Poisson likelihood (n, = 182 neurons). We trained models on trials spanning all reach conditions and restricting data time window starting 250 ms before and ending 450 ms after movement onset, binned at 5ms. For regression to hand velocity, we introduced lag of 100ms between neural activity and hand velocity. As the data used in the NLB chal- lenge 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 mini-batch size of 168 trials.
The co co-smoothing metric used to assess how well the model fit the data is defined as log ligeliidd score
(C(A;
Dm.e) Dm.e) - L(Ans (()n:)n.))
(S44)
Thpp log
where the overall log- ligelikll for all neurons, À denotes the vector inferred time-varying firing rates, Å, is the mean firing rate for neuron n and nap is the total number of spikes.
the sum of all the log- ligeliikss evaluated at all points and
I.2 FURTHER ANALYSES
A key feature of monkey MI motor cortical recordings is the prevalence of rotational dynamics in the data (Churchland et al., 2012). These can be captured using jPCA, method developed to find the subspace in which the dynamics are most rotational, which was recently generalized by Rutten al. (2020). Here, we found that we could uncover clean rotational dynamics from the single-trial firing rates, similarly to Pandarinath et al. (2018).
Figure S6: 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 fir- ing rates. Projection of the neural activity yields very clean rotational trajectories.
JPCA2
JPCA1
M FURTHER DETAILS OF THE CONTINUOUS REACHING TASK ANALYSIS
M. DETAILS OF THE ANALYSES
For our analyses of the primate data in Section 3.4, we considered the first 22 minutes of the record- ing session i ndy 201 60426 from O' O'Doherty et al. (2018). We binned spikes at 25 ms resolution
12

---

Published as a conference paper at ICLR 2022
and considered all neurons with firing rate of at least 2Hz. 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 z-scored the hand velocity and shifted it by 120ms, following on Jensen et al. (2021).
To fit MiiL---AARE 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 held-out test dataset. We fitted a model with " 50, m 10 using the non-linear dynamics described in Equation S4. The latent state was then mapped onto both the kinematics and neural observations. We used a linear readout from latents 2D kinematics variables, and a linear readout following by nonlinearity from latents to the firing rates of 130 neurons.
After fitting the -LQR-VAE to neural activity and behavior jointly, we then proceeded infer u from neural activity alone. Next, we computed the kinematic reconstruction error on the test dataset as the fraction of variance captured in both 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 ms- We aligned the z-scored input 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.
M.2 COMPARISON OF ILQR-VAE AND BGPFA
Figure S7: Comparison of the firing rates inferred by bGPFA and iLQR-VAL VAE 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 better decoding of the kinematics.
neuron
neuron 12
neuron 110
60 _
60
60
50
50
50
40
30
30
30
20
20
10
0
0
ILOR-VAE (Hz)
60
60
60
50
50
50
40
40
30
30
20
20
20
10
25
50
75
100
25
50
75
100
25
50
75
100
time (ms)
time (ms)
time (ms)
As further way of understanding the relative benefits and disadvantages of iL.QR-VAE VAF, we com- pared its performance with BGPFA, a fully Bayesian extension of GPFA YYu et al., 2009) that enables the use of non- Gan-Gaussi likelihoods, scales to very large datasets, and was recently shown to out- perform standard GPFA on this same continuous reaching dataset (Jensen et al., 2021). Importantly, places smooth prior directly on the la- tents with no explicit notion of dynamics. We fit both methods using 10 minutes of data (chunked into pseudo-triae -truls for VLQR-VAE VAE and as continuous trial for PGPFA). 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 R² when regressing from firing rates to hand velocity). Note that the firing rate estimates here 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.
bGPFA makes different assumptions iLQR-
From Figure S7, 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
13

---

Published as a conference paper at ICLR 2022
well as ILQR-VAE (R² 0.6 for bGPFA and 0.76 for iLQR--AE.) This is indeed what
in Figure S8.
True
0.6
ILOR- PRAVAE
bGPFA
0.4
Figure S8: Ground truth hand velocity (red) and decoded kinematics using iLQR- VAE (blue) and bGPFA (dotted black). A lin- ear decoder was trained on 9 minutes data and then evaluated on the first minute of the recording (shown here). bGPFA struggles to capture the biggest peaks in the velocity, con- sistent with the smoother firing rates and the lower R²a
V, (a.u.)
0.2
m
0
-0.2
0.44
1200
1300
1400 time (ms)
1500
1600
N LINK TO KALMAN FILTERING
The Linear Quadratic Regulator and the Kalman filter (Kalman, 1964) algorithms designed for systems with linear dynamics and Gaussian noise. LQR finds the optimal feedback control law to minimize 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 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 Kalman filter.
LINEAR QUADRATIC CONTROL AS FILTERING
The Kalman smoother assumes dynamics of the form
1 Az, + B ww
(S45)
0, = Cz, +
(S46)
with w I). NO E.). and the initial condition is assumed to be generated by Gaussian distribution with known parameters 21 ~ (N, IT).
On the other hand, LQR assumes the following fully-determinnstic dynamics:
Ze+1 Az BB'
(S47)
0: = Cz
(S48)
with #1 known exactly.
Note than in the iLQR-VAE framework we have thus far only considered cases where observed external inputs were given. However these can be straightforwardly included as an additional Bü term in Equation S47 and Equation S45.
The Kalman smoother's objective is minimize the expected mean squared error between the in- ferred latent state and the true state, E ||| - & [ As deseribed in Aravkin et al. (2017), with linear dynamics and Gaussian noise, this becomes equivalent to minimizing the following objective w.r.t z:
((*) II (21 21]2 +
Σ 1=1
B (2141 Az))
Σ (/2
Cz
(S49)
14

---

Published as a conference paper at ICLR 2022
where the first two terms correspond to the prior 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 minimize to find the posterior mean in iLQR-VA Q--VEE (Equation 11):
T-
T
((u) = 1|2 uoll?
D uu
1/2
D (2 - Cz
(S50)
u
1=1
1=
7
|| || uol1
1 1/2
3141
A2;)11 +
(S51)
t=1
The right-hand sides of Equation S49 and Equation S50 become identical when Σ. = B* and Eg II. Note that the introduction of the B' matrix in Equation S45 unties the two formulations slightly by allowing for further mixing between the input channels that isn't 't accounted for by the prior. In the examples we consider next, we therefore set B'=
The above equations show how LQR can be used to solve the standard Kalman filtering problem, with the key difference being that the optimization performed over inputs u (u....... ur-11 rather than latent trajectories = = ......., _T} directly. This is illustrated in Figure S9(A), where Rauch- Kung-Striebel (RKS) smoother and LQR were ran on the same set of 8-dimensional observa- tions arising from 8-dimensional linear dynamical system, and inferred the same latent trajectory given the ground-trute parameters. As we only use LQR to parametrize the mean of the posterior distribution, we trained the recognition model for 100 steps get the uncertainty the latents, which was very similar to the output of the RKS smoother.
LEARNING A KALMAN FILTER
We then proceeded to learn the parameters of the models using either ILQR-VAE Expectation- Maximization (EM) procedure, or direct minimization of the negative log likelihood of the data (Figure S9B-C).
Interestingly, the EM algorithm is closely related to iLQR- VAE, since the E-step finds the latent tra- jectories minimizing Equation S49, when iLQR- RRE solves Equation S50 in an inner optimization loop. While there exists an analytical solution for the M-step in the case of the Kalman filter, this does generalize to nonlinear dynamics and non-Gaussian noise. Therefore, we used gradient descent procedure for the maximization 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 S9C that iLQR-VAE reaches a smaller NLL in considerably fewer iterations than gradient descent, which we hypothesize is due the good preconditioning given by iLQR (discussed in Figure 1). Note however that the cost of one iLQR- VAE iteration is higher than the direct computation Equation $49.
In this section, we have shown in a simple linear-quadratic example how il IL.R--VAR VAE performs filter- ing by inferring the process noise as inputs. While this 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 particle filtering unscented Kalman filters (Doucet and Johansen, 2009; Wan et al., 2001). iLQR- VAE offers another way to solve this problem, with close links to the aforementioned approaches.
O ANALYSIS OF THE INFERENCE GAP
In order to evaluate the benefits of defining the recognition model implicitly through the generative parameters, we compared iLQR-VAR VAE to a more standard sequential variational auto-encoder, using bidirectional recurrent neural network 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 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 al.,
3, 3, n. 10, Student prior), such that the
15

---

Published as a conference paper at ICLR 2022
A
C
ILOR VAE
talento
Figure S9: Comparison of the posterior LQR and a Kalman smoother on noisy observations generated by a latent system (see text for de- tails). (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 uncer- tainty is shown for both cases on half of the data. The iLQR-VA. 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 the posterior distribution using iLQR v.s direct minimization of the NLL. On unseen data, both were able to converge to 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 and axes logarithmic scale.
of Kalman smoother and LQR. We both
2015), which is computed as
k p(o, 24)
LIWAE
E
(u o)
log
T:
(u,lo)
(S52)
where we used Monte-Carlo sampling with 1000 samples to evaluate the expectation. This then allowed us to compute the inference gap (Cremer et al., 2018) of both models as AI ELBO. As shown in Figure S10, iLQR-VAE has 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.
Figure S10: Comparison of iLQQ with biRNN recognition model. (A- B) Loss and inference gap as a function of iteration number (starting from iteration 20) for iLQR-VAE (blue) and the biRNN model (pink). (C) In- ferred input norm the end of training for iLQR-VAE and the biRNN model. Ground truth input is shown in dotted black lines. (D) Inferred out- put at the end of training for iLQR-VAL and the biRNN model. Both models explain the observations (grey dots) well, but iLQR- VAE captures the sharp transitions better.
LOR VAE BIANN
effrence
1.5
LOR VAE LL LOR--AE ELBO
BIANN biRNN ELBO
0.5
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
C
D
3.5
3
2.5
observations
00.5
-1.
10
30
40
70
80
90
100
20
30
60
100
time (ms)
time
(ms)
REFERENCES
Amos, B., Rodriguez, I. D. J., Sacks, J., Boots, B., and Kolter, J. (2018). Differentiable MPC for end-to-end planning and control. arXiv preprint 1rXiv:1810. 13400.
16

---

Published as a conference paper at ICLR 2022
Aravkin, A., Burke, J. V., Ljung, L., Lozano, A., and Pillonetto, G. (2017). Generalized Kalman smoothing: Modeling and algorithms. Automatica, 86:63-86.
Bhatia, N. P. and Szegö, G. (2002). Stability theory of dynamical systems. Springer Science & Business Media.
Boyd, S., Boyd, S. P., and Vandenberghe, L. (2004). Convex optimization. Cambridge university press.
Brunton, B. W., Johnson, L. A., Ojemann, J. G., and Kutz, J. N. (2016a). Extracting spatial-temporal coherent patterns in large-scale neural recordings using dynamic mode decomposition. Journal of neuroscience methods, 1-15
Brunton, S. L., Proctor, J. L., and Kutz, J. N. (2016b). Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 1 113(15):3932-33)
Burda, Y., Grosse, R., and Salakhutdinov, R. (2015). Importance weighted autoencoders. arXiv preprint orXiv:1509.005/9
Champion, K., Lusch, B., Kutz, N., and Brunton, S. (2019). Data-driven discovery of coordinates and governing equations. Proceedings of the National Academy of Sciences, 116(45):2245-22255.
Chen, T., Amos, B., and Nickel, M. (2020). Learning neural event functions for ordinary differ- ential equations. arXiv preprint arXiv:2011.03900
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. (2018), Neural ordinary differential equations. arXiv preprint arXiv:1806.07366
Churchland, M. M., Cunningham, J. P., Kaufman, M. T., Foster, D., Nuyujukian, P., Ryu, S. I., and Shenoy, K. V. (2012). Neural population dynamics during reaching. Nature, 487:51-56.
Costa, A. C.,, Ahamed, T., and Stephens, G. J. (2019). Adaptive, locally linear models of complex dynamics. Proceedings of the National Academy of Sciences, 116(5)1501-1510.
Cremer, C., Li, X., and Duvenaud, D. (2018). Inference suboptimality in variational autoencoders. In International Conference on Machine Learning, pages 1078-1086. PMLR.
Doucet, A. and Johansen, A. M. (2009). A tutorial on particle filtering and smoothing: Fifteen years later. Handbook of nonlinear filterin 12(656-704):3.
Fieseler, C., Zimmer, M., and Kutz, J. N. (2020). Unsupervised learning of control signals and their encodings in Caenorhabditis elegans whole-brain recordings. Journal of the Royal Society Interface, 17(13))22000449.
Ghahramani, Z. and Hinton, G. E. (1996). Parameter estimation for linear dynamical systems.
Ghahramani, Z. and Hinton, G. E. (2000). Variational learning for switching state-space models. Neural comput., 12:831-864.
Giles, M. (2008). An extended collection of matrix derivative results for forward and reverse mode automatic differentiation.
Heck, J. C. and Salem, F. M. (2017). Simplified minimal gated unit variations for recurrent neural networks. In 20/7 IEEE 60th International Midwest Symposium on Circuits and Systems (MWS- CAS), pages 1593-1596. IEEE.
Heraanezz, D., Moretti, A. K., Wei, Z.. Saxena, S., Cunningham, J., and Paninski, L. (2018). Non- linear evolution via spatially-depender linear dynamics for electrophysioloyy and calcium data. arXiv preprint 4r11v:1811.0245,
Jensen, K. T., Kao, T.-C., Stone, J. T., and Hennequin, G. (2021). Scalable Bayesian GPFA with automatic relevance determination and discrete noise models. Advances in Neural Information Processing Systems, 34.
17

---

Published as a conference paper at ICLR 2022
Kalman, R. E. (1964). When is linear control system optimal?
Karush, W. (2014). Minima of Functions of Several Variables with Inequalities as Side Conditions. Springer Basel.
Koppe, G., Toutounji, H., Kirsch, P., Lis, S., and Durstewitz, D. (2019). Identifying nonlinear dynamical systems via generative recurrent neural networks with applications to fmri. PLoS computational biology, 15(8)e10072633.
Kuhn, H. W. and Tucker, A. W. (2014). Nonlinear programming. In Traces and emergence of nonlinear programming, pages 247-258. Springer.
Kutz, J. N., Brunton, S. L.,, Brunton, B. W., and Proctor, J. L. (2016). Dynamic mode decomposition: data-driven modeling of complex systems. SIAM.
Li, W. and Todorov, E. (2004). Iterative linear quadratic regulator design for nonlinear biological movement systems. In ICINCO (1), pages 222-222 Citeseer.
Linderman, S., Johnson, M., Miller, A., Adams, R., Blei, D., and Paninski, L. (2017), Bayesian learning and inference recurrent switching linear dynamical systems. Artificial Intelligence and Statistics, pages 112-922. PMLR.
Morrison, M. J., Fieseler, C., and Kutz, J. N. (2020). Nonlinear control in the nematode C. elegans. Frontiers in Computational Neuroscience, 14:123.
O'Doherty, J. E., Cardoso, M. M. B., Makin, J. G., and Sabes, P. N. (2018). Nonhuman Primate Reaching with Multichannel Sensorimotor Cortex Electrophysiology: broadband for This research was supported by the Congressionally Directed Medical Re-
indy_201 (ily_20000300000 search Program (W8IXWW-114--00100) JEO was supported by fellowship #2978 from the Para- lyzed Veterans of America. JGM was supported by fellowship from the Swartz Foundation.
Pandarinath, C.,, S'ea, D. J., Collins, J., Jozefowicz, R., Stavisky, S. D., Kao, J. C., Trautmann, E. M., Kaufman, M. T., Ryu, S. I., Hochberg, L. R., et al. (2018). Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10))805-815.
Pei, F., Ye, Zoltowski, D., Wu, A., Chowdhury, R. H., Sohn, H., O'Doherty, E., Shenoy, K. V., Kaufman, M. T., Churchland, M., et al. (2021). Neural latents benchmark 21: Evaluating latent variable models of neural population activity. arXiv preprint arXiv:2 04463.
Proctor, L.., Brunton, L.,, and Kutz, J. (2016). Dynamic mode decomposition with control. SIAM Journal on Applied Dynamical Systems, 15(1)1142-161
Proctor, J. L., Brunton, S. L., and Kutz, J. N. (2018). Generalizing Koopman theory to allow for inputs and control. SIAM Journal on Applied D Dynamical Systems, 17(1 ((11):909-30.
Rutten, V., Bernacchia, A., Sahani, M., and Hennequin, G. (2020). Non-reversible Gaussian pro- cesses for identifying latent dynamical structure in neural data. Advances in Neural Information Processing Systems.
Schmid, P. J. (2010). Dynamic mode decomposition of numerical and experimental data. Journal of fluid mechanics, 656:5-28.
Tassa, Y., Mansard, N., and Todorov, E. (2014). Control-limited differential dynamic programming. In 2014 IEEE International Conference Robotics and Automation (ICRA), pages 1168-1175. IEEE.
Wan, E. A., Van Der Merwe, R., and Haykin, S. (2001). The unscented Kalman filter. Kalman filtering and neural networks, 5(2007):211288)
Yu, Byron, M., Cunningham, P., Santhanam, G., Ryu, S. Shenoy, K. V., and Sahani, M. (2009). Gaussian-proces factor analysis for low-dimensional single-trial analysis of neural population activity. In Advances in neural information processing systems, pages 1881-1888.
18


Reviews:
[{"review_id": "bZSyAe4_hug", "reviewer": "Reviewer_YcP3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "strengths\n---------\nThe problem of simultaneously learning dynamics and ongoing external inputs is a crucial one, especially in fields such as neuroscience where complex dynamical systems are often analyzed through incomplete measurements of the system (i.e. not every neuron is monitored). This paper presents a sound approach to tackling this problem, though I admit my understanding if the control field is limited.\n\nThe paper is clearly written and easy to follow.\n\nI appreciate building up the complexity of the toy examples.\n\nComaprison to state-of-the-art methods is nice and clear.\n\n\nweaknesses\n----------\nFig. 1: \"Indeed, iLQR is unable to find initial conditions that would explain the data well, resulting in a much higher initial loss.\" While the paper implicitly points to this outcome as an advantage of the more flexible prior, it is concerning that iLQR cannot find the initial condition of this very simple example. Is this an issue with the iLQR algorithm itself? The chosen prior? It would be very useful to understand this behavior in the simple example before moving on to more complex examples.\n\nThe leap from example 1 (nonlinear autonomous dynamics) and example 2 (linear dynamics + sparse inputs) to the real data (presumably nonlinear dynamics + ongoing inputs) begs the question: how well does the model perform with (nonlinear dynamics + sparse inputs)? Showing good performance in this regime would increase my confidence in the model's ability to handle more complex, real-world data.\n\nThere is very little reference to related work, especially with regards to inferring control inputs in dynamical systems (linear or nonlinear). The paper would feel much more complete with this information (which could be provided as an additional appendix if space is short).\n\nThe general problem of decoupling ongoing dynamics from external inputs is underdetermined, and results rely heavily on the structure imposed on these aspects of the model. In this paper the prior on the inputs is a sparse one, but what happens if there is mismatch between this prior and the real data? How will that affect the results (and their interpretation)? The paper comments on this several times (and shows how an AR prior in LFADS performs when inputs are sparse, but not the reverse), but I'd appreciate a more thorough discussion of this issue in the Discussion.\n\nLearning in the proposed model seems to have some strong parallels with the EM algorithm, except that maxima from iLQR replace expectations. Is this true? Could be interesting to briefly draw this connection when describing the inference strategy.\n\n\nminor\n-----\nFig. 1 caption: first sentence, add \"system\" to end?\n\nparagraph that starts \"At the beginning...\", second sentence, typo: \"learning consists in _making_\"\n\nsame paragraph: \"We note that this regime is facilitated here by our choice of generator dynamics, which we initialised to be very weak initially and therefore easily controllable.\" What does it mean for the generator dynamics to be \"initialised to be very weak\"? This is potentially an interesting point that is unclear to me as written. What does \"weak initialisation\" mean in this context?", "summary_of_the_paper": "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.", "summary_of_the_review": "I'd lean towards accepting this paper, due to its novelty and thoroughness. I'd be willing to argue even more strongly for it if the authors could address some of the points above (especially the more complex toy example).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Xfk41Zsi_Po", "reviewer": "Reviewer_gF5H", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Post-rebuttal: Thanks for addressing my concerns. The authors have updated the manuscript and added experiments. I will not change my recommendation.\n\n##########\n\nPros: \n\n+ Implicit recognition model implied by the generative model. It incorporates the inferred forward dynamics.\n\n+ Competitive performance to STOA method\n \n+ Small parameter space for easy hyperparameter tuning\n\nConcerns: \n\n- The method is extended and compared to LFADS. How about with other methods?\n\n- Unlike the bidirectional RNN in LFADS, the implicit recognition model should incorporate the dynamics from the trained generative model. This bias could be good or bad for the training. Can the authors address more?\n\n- There are other methods like VIND consider to incorporate the dynamics from the trained generative model. This better be compared or discussed.\n\n- Is this method fast or slow comparing to LFADS or other methods?\n\n- How well can the proposed method learn other typical types of dynamical systems such as fixed attractors, continuous attractors and etc.\n\nMinor:\n\n* Typo: \"LFDADS\" in paragrah 2, page 2", "summary_of_the_paper": "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.\n", "summary_of_the_review": "Overall, I vote for accepting. The paper is well written. I like the idea of transforming to control problem. The implicit recognition model also interesting. This would be useful to learn neural dynamics from large population recordings. Hopefully the authors can address my concern in the rebuttal period. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "0GSrgb07970", "reviewer": "Reviewer_Ve7m", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Authors present a novel method for simultaneously learning latent dynamics and inferring unobserved control inputs. The method performs very well on several toy datasets (autonomous and input-driven linear dynamical system + Lorenz attractor), is on par with state of the art methods on benchmark datasets for neural data analysis methods, and allows for better reconstruction of hand kinematics for primate recordings in a continuous reaching task. The main novelty of the method lies in the utilization of iLQR with implicit differentiation. The paper also extensively discusses LFADS and provides very interesting insights on LFADS.\n\nAn important component  of the method is the prior on input distribution. Authors show how to use a Gaussian prior or a multivariate Student prior (it allows for strong inputs when needed, represents the fact that inputs come as shared events and are spatially and temporally independent, and authors mention it might be one advantage over LFADS's autoregressive prior). This is a key idea of the paper and I would like to see more experiments highlighting the importance of the prior. I would recommend that the authors show (possibly in the appendix), the differences between the inferred inputs and performances when using different priors. \n\nThe derivation of the ELBO does not seem straightforward to me. To go from equation 7 to 8, don't you need to condition p(o|u) on z_0 as well? Maybe I am not correct, but I would appreciate if authors clarify this point (and correct the manuscript if needed)  \n\nAbout the computational complexity. The model is linear in T and cubic in n, which seems fast for low-dimensional latents only. Authors mention that iQLR-VAE enables fast learning of dynamics, but they also mention it as a limitation in the discussion. Could authors clarify this? Moreover, I believe the computational complexity comparison with LFADS in Figure 1 bottom left is not fair. Authors compare the number of training iteration without comparing the complexity of each iteration. I would like to see comparison for other datasets as well. \n\nThere are many performance metrics used in this paper to show that iQLR-VAE performs as well as other state of the art methods. However, there are no comparisons of the processes learnt by the different methods. Showing superimposed processes inferred by different methods  (not only LFADS) for several experiments would allow to gain insight into the methods and how they compare (in term of uncertainty, smoothness, etc...). Similarly, for Figure 4, it would be interesting to see a low dimensional representation of neural activity (and comparisons of different methods). For this experiment, is it necessary to use a latent processes of dimension 50? Wouldn't it be sufficient and more efficient to use lower dimensional latent for this task? \n\nMinor point : \n - In figure 4 caption, there is a typo and the last part should be panel D not C.\n", "summary_of_the_paper": "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. \nAuthors show comparisons to other models (the closest one being LFADS) on toy datasets and benchmark datasets for neural data analysis methods. \niLQR-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. ", "summary_of_the_review": "I recommend acceptance of the paper. This is a very interesting and novel method, that is shown to perform very well on several different datasets. \nThere are many quantitative experiments but showing qualitative comparisons would be very beneficial for gaining insights into how different methods compare. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces a novel control-based variational inference approach that learns latent dynamics in an *input-driven* state-space model. An optimal control solution (iLQR) is implicitly used as the recognition model which is fast and compact. Reviewers unanimously agree on the high quality writing and high significance of the work. This paper advances the horizon of nonlinear dynamical system models with unobserved input, an impactful contribution to the neuroscience and time series communities.
Decision: Accept (Oral)