iLQR-VAE : control-based learning of input-driven dynamics with applications to neural data

1 INTRODUCTION . The mammalian brain is a complex , high-dimensional system , containing billions of neurons whose coordinated dynamics ultimately drives behaviour . Identifying and interpreting these dynamics is the focus of a large body of neuroscience research , which is being facilitated by the advent of new experimental techniques that allow large-scale recordings of neural populations ( Jun et al. , 2017 ; Stosiek et al. , 2003 ) . A range of methods have been developed for learning dynamics from data ( Buesing et al. , 2012 ; Gao et al. , 2016 ; Duncker et al. , 2019 ; Archer et al. , 2015 ; Hernandez et al. , 2018 ; She and Wu , 2020 ; Kim et al. , 2021 ; Nguyen et al. , 2020 ) . These methods all specify a generative model in the form of a flexible latent dynamical system driven by process noise , coupled with an appropriate observation model . However , neural recordings are typically only made in a small selection of brain regions , leaving many areas unobserved which might provide relevant task-related input to the recorded one ( s ) . Yet , the aforementioned methods perform Bayesian inference of state trajectories directly , and therefore do not support inference of external input ( which they effectively treat as process noise and marginalize out ) . Indeed , simultaneous learning of latent dynamics and inference of unobserved control inputs is a challenging problem that involves teasing apart momentary variations in the data that can be attributed to the system ’ s internal transition function , and those that need to be explained by unobserved inputs . This distinction can be achieved by introducing external control in the form of abrupt changes in the latent state transition function , and inferring these switching events ( Ghahramani and Hinton , 2000 ; Linderman et al. , 2017 ) . More recently , Pandarinath et al . ( 2018 ) introduced LFADS , a sequential variational autoencoder ( VAE ) that performs inference at the level of external inputs as well as initial latent states . The inferred inputs were shown to be congruent with taskinduced perturbations in various reaching tasks in primates ( Pandarinath et al. , 2018 ; Keshtkaran and Pandarinath , 2019 ) . Here , we introduce iLQR-VAE , a new method for learning input-driven latent dynamics from data . As in LFADS , we use an input-driven sequential VAE to encode observations into a set of initial conditions and external inputs driving an RNN generator . However , while LFADS uses a separate , bidirectional RNN as the encoder , here we substitute the inference network with an optimizationbased recognition model that relies on the powerful iterative linear quadratic regulator algorithm ( iLQR , Li and Todorov , 2004 ) . iLQR solves an optimization problem that finds a mode of the exact posterior over inputs for the current setting of generative parameters . This ensures that the encoder ( mean ) remains optimal for every update of the decoder , thus reducing the amortization gap ( Cremer et al. , 2018 ) . Moreover , having the recognition model be implicitly defined by the generative model stabilizes training , prevents posterior collapse ( thus circumventing the need for tricks such as KL warmup ) , and greatly reduces the number of ( hyper- ) parameters . While iLQR-VAE could find applications in many fields as a general approach to learning stochastic nonlinear dynamical systems , here we focus on neuroscience case studies . We first demonstrate in a series of synthetic examples that iLQR-VAE can recover the true dynamics in both autonomous and input-driven systems . Next , we show state-of-the art performance on monkey M1 population recordings during two types of reaching tasks ( O ’ Doherty et al. , 2018 ; Churchland et al. , 2010 ) . In particular , we show that hand kinematics can be accurately decoded from inferred latent state trajectories , and that the inferred inputs are consistent with recently proposed theories of motor preparation . 2 METHOD . iLQR-VAE models a set of temporal observations , such as behavioural and/or neural recordings , through a shared input-driven nonlinear latent dynamical system ( Figure S1 ) . The input encapsulates both process noise ( as in traditional latent dynamics models ) , initial inputs that set the initial condition of the dynamics , and any meaningful task-related control input . In this section , we describe the architecture of the generative model , and the control-based variational inference strategy used for training the model and making predictions . 2.1 GENERATIVE MODEL . We consider the following generative model : latent state zt+1 = fθ ( zt , ut , t ) ( 1 ) observations ot|zt ∼ pθ ( ot|zt ) ( 2 ) where ut ∈ Rm , zt ∈ Rn and ot ∈ Rno are the input , latent state and observations at time t , respectively . Here , observations may comprise either neural activity , behavioural variables , or both – the distinction will be made later where relevant . We use the notation θ to denote the set of all parameters of the generative model . We use u0 to set the initial condition z1 = fθ ( 0 , u0 , 0 ) of the network1 . This way , the latent state trajectory of the network z ( u ) = { z1 , . . . , zT } is entirely determined by the input sequence u = { u0 , . . . , uT } and the state transition function fθ ( · ) , according to Equation 1 . For fθ ( · ) , we use either standard linear or GRU-like RNN dynamics ( see Appendix B 1Note that when m < n , u0 can only reach an m-dimensional subspace of initial conditions , which could be limiting . We can circumvent this problem by spreading u0 over multiple surrogate time bins before the start of the trial , i.e . introduce { u−n/m , . . . , u−2 , u−1 , u0 } together with an appropriate dependence of fθ on t ≤ 0 in Equation 1 , such that each of these surrogate inputs target a different latent subspace with purely integrating ( “ sticking ” ) linear dynamics before t = 1. for details ) . For the likelihoods , we use Gaussian or Poisson distributions with means given by linear or nonlinear readouts of the network state of the form ōt = h ( Czt + b ) ( Appendix C ) . We place a Gaussian prior over ut≤0 . We then consider two alternative choices for the prior over ut > 0 . The first is a Gaussian prior pθ ( ut > 0 ) = N ( 0 , S2 ) ( 3 ) with S = diag ( s1 , . . . , sm ) . In many settings however , we expect inputs to enter the system in a sparse manner . To explicitely model this , we introduce a second prior over u in the form of a heavy-tailed distribution constructed hierarchically by assuming that the ith input at time t > 0 is uit = si it √ ν/αt ( 4 ) where si > 0 is a scale factor , it ∼ N ( 0 , 1 ) is independent across i and t , and αt ∼ χ2ν is a shared scale factor drawn from a chi-squared distribution with ν degrees of freedom . Thus , inputs are spatially and temporally independent a priori , such that any spatio-temporal structure in the observations will have to be explained by the coupled dynamics of the latent states . Moreover , the heavy-tailed nature of this prior allows for strong inputs when they are needed . Finally , the fact that the scale factor is shared across input dimensions means that inputs are either all weak or potentially all strong at the same time for all input channels , expressing the prior belief that inputs come as shared events . This hierarchical construction induces a multivariate Student prior at each time step : pθ ( ut ) = Γ [ ( ν +m ) /2 ] Γ [ ν/2 ] ( νπ ) m/2|S| [ 1 + 1 ν uTt S −2ut ] − ( ν+m ) /2 ( 5 ) where S = diag ( s1 , . . . , sm ) . Note that both S and ν are parameters of the generative model , which we will learn . We discuss the relative advantages of the Student and Gaussian priors in Appendix J . 2.2 ILQR-VAE : A NOVEL CONTROL-BASED VARIATIONAL INFERENCE STRATEGY . To train the model , we optimize θ to maximize the log-likelihood of observing a collection of independent observation sequences O = { o ( 1 ) , . . . , o ( K ) } , or “ trials ” , given by : log pθ ( O ) = K∑ k=1 log ∫ pθ ( o ( k ) |z ( u ) ) pθ ( u ) du . ( 6 ) As the integral is in general intractable , we resort to an amortized variational inference strategy by introducing a recognition model qφ ( u|o ( k ) ) to approximate the posterior pθ ( u|o ( k ) ) . Following standard practice ( Kingma and Welling , 2013 ; Rezende et al. , 2014 ) , we thus train the model by maximizing the evidence lower-bound ( ELBO ) : L ( O , θ , φ ) = ∑ k Eqφ ( u|o ( k ) ) [ log pθ ( o ( k ) |u ) + log pθ ( u ) − log qφ ( u|o ( k ) ) ] ( 7 ) = ∑ k Eqφ ( u|o ( k ) ) [ T∑ t=1 log pθ ( o ( k ) t |zt ) + log pθ ( ut ) − log qφ ( ut|o ( k ) ) ] ( 8 ) ≤ log pθ ( O ) . ( 9 ) with respect to both θ and φ . Here , the main novelty is the use of an optimization-based recognition model . We reason that maximizing the exact log posterior , i.e . computing u ? ( o ( k ) ) = argmax u log pθ ( u|o ( k ) ) ( 10 ) = argmax u [ T∑ t=1 log pθ ( o ( k ) t |u ) + log pθ ( ut ) ] ( 11 ) subject to the generative dynamics of Equations 1 and 2 , is a standard nonlinear control problem : log pθ ( o ( k ) t |u ) acts as a running cost penalizing momentary deviations between desired outputs ot and the actual outputs caused by a set of controls u , and log pθ ( ut ) acts as an energetic cost on those controls . Importantly , there exists a general purpose , efficient algorithm to solve such nonlinear control problems : iLQR ( Li and Todorov , 2004 ; Appendix D ) . We thus propose to use a black-box iLQR solver to parameterize the mean of the recognition density qφ ( u|o ) for any o , and to model uncertainty separately using a multivariate Gaussian density common to all trials . Therefore , we parametrize the recognition model as follows : qφ ( u|o ) = N ( u ; u ? ( o ) , Σs ⊗Σt ) ( 12 ) with u ? ( o ) = iLQRsolve ( o , θ ) . ( 13 ) where we use a separable posterior covariance ( the Kronecker product of a spatial factor Σs and a temporal factor Σt ) . To optimize the ELBO , we estimate the expectation in Equation 8 by drawing samples from qφ ( u|o ( k ) ) and using the reparameterization trick ( Kingma et al. , 2015 ) to obtain gradients . A major complication that would normally preclude the use of optimization-based recognition models is the need to differentiate through the mean of the posterior . In this case , this involves differentiating through an entire optimization process . Using automatic differentiation within the iLQR solver is in general impractically expensive memory-wise . However , recent advances in differentiable model predictive control enable implicit differentiation through iLQRsolve with a memory cost that does not depend on the number of iterations ( Amos et al. , 2018 ; Blondel et al. , 2021 ; Appendix E ) .

Summary: 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.