Model Predictive Control of the Neural Manifold

Abstract

Neural manifolds are an attractive theoretical framework for characterizing the complex behaviors of neural populations. However, many of the tools for identifying these low-dimensional subspaces are correlational and provide limited insight into the underlying dynamics. The ability to precisely control this latent activity would allow researchers to investigate the structure and function of neural manifolds. Employing techniques from the field of optimal control, we simulate controlling the latent dynamics of a neural population using closed-loop, dynamically generated sensory inputs. Using a spiking neural network (SNN) as a model of a neural circuit, we find low-dimensional representations of both the network activity (the neural manifold) and a set of salient visual stimuli. With a data-driven latent dynamics model, we apply model predictive control (MPC) to provide anticipatory, optimal control over the trajectory of the circuit in a latent space. We are able to control the latent dynamics of the SNN to follow several reference trajectories despite observing only a subset of neurons and with a substantial amount of unknown noise injected into the network. These results provide a framework to experimentally test for causal relationships between manifold dynamics and other variables of interest such as organismal behavior and BCI performance.

1 Introduction

Neural circuits are composed of large numbers of interconnected neurons whose activity depends on other neurons in the circuit. Because of these dependencies, simultaneously recorded populations of neurons usually exhibit high levels of correlation. Equivalently, most of the variance within the high-dimensional space corresponding to the firing rates of individual neurons is confined to a lower-dimensional subspace, or neural manifold [1]. Due to their relative simplicity, manifolds have become a popular framework for understanding the complex dynamics of large neural populations. In many different systems, activity on neural manifolds has been shown to correlate with salient features of stimuli, physical position, and internal cognitive states [2, 3, 4, 5]. However, the insight that manifolds can provide into the underlying dynamics and computational principles of the circuits remains a contentious question.

Broadly speaking, neural manifolds may be seen either as descriptive tools for dealing with the inherently correlated nature of neural data arising from highly interconnected circuits, or as a method of revealing a more fundamental dynamics that exists within a latent space [6]. Using linear subspaces with autonomous dynamics as an illustrative example, the descriptive perspective can be seen as classic dimensionality reduction with

$$\mathbf{z_{t}}=\mathbf{Gx_{t}},$$ (1)

where $\mathbf{x_{t}}$ is a column vector of the activity of $n$ neurons at time $t$ and $\mathbf{z_{t}}$ is a reduced dimension representation of $\mathbf{x_{t}}$ given by the linear transformation $\mathbf{G}$ (which parameterizes the neural manifold). This model is descriptive because the latent trajectories $\mathbf{z_{t}}$ are seen as a convenient representation of high-dimensional trajectories that result from the highly coupled dynamics that exist among the individual neurons $\mathbf{x_{t}}$, and thus provide limited information for inference or mechanistic understanding [6]. In contrast, the generative perspective can be modeled as a latent factor model of the form

$$\mathbf{x_{t}}=\mathbf{Fz_{t}}+\mathbf{\epsilon_{t}},$$ (2)

where $\mathbf{x_{t}}$ is a column vector of the activity of $n$ neurons at time $t$, $\mathbf{z_{t}}$ are the latent factors that span the neural manifold with some smaller dimension $k$, $\mathbf{F}$ are the factor loadings, and $\mathbf{\epsilon_{t}}$ is a sample from some distribution (often Gaussian). This perspective views the measured neural activity as a function of the latent dynamics of $\mathbf{z_{t}}$, which emerge from but are simpler than the dynamics of $\mathbf{x_{t}}$. The descriptive and generative perspectives may also be seen as bottom-up and top-down approaches, respectively, for addressing the question of how computations are implemented by neural circuits [7, 6]. Early work on neural manifolds used linear methods such as principal components analysis (as in equation 1) and factor analysis (as in equation 2), but there is now a broad consensus that neural manifolds are often nonlinearly embedded in the full state space, requiring more sophisticated methods to identify [8].

A potential weakness in the generative approach to understanding neural manifolds is that most methods of dimensional reduction are static: they produce a time series of snapshots from an informative angle in the neural state space [9], but the dynamics have to be inferred through other means. In contrast to the bottom-up approach where there is a robust foundation of biophysics on which to build models of circuit dynamics at the level of cells and synapses, the question of how best to model dynamics in the latent space remains an active area of research [10, 11, 12, 13]. Testing these models and their causal relationship to behavior would benefit from methods for experimentally controlling the activity on the neural manifold.

In this study, we develop a framework for controlling latent dynamics in the context of a sensory system. We express activity on the neural manifold with a general state-space model

$$\mathbf{z_{t+1}} =g(\mathbf{z_{t}},\mathbf{u_{t}},\mathbf{\epsilon_{t}})$$ (3)

$$\mathbf{x_{t}} =f(\mathbf{z_{t}}),$$ (4)

where the latent dynamics on the manifold are defined by a function of the current state of the system $\mathbf{z_{t}}$, an external, time-varying stimulus $\mathbf{u_{t}}$ that will be used for control, and an intrinsic, uncontrolled source of noise $\mathbf{\epsilon_{t}}$. The high dimensional measured neural activity $\mathbf{x_{t}}$ is obtained with the observation function in equation (4). Interestingly, we show that the data-generating process (i.e., the dynamics of $\mathbf{x_{t}}$) does not need to be the same as the latent dynamics model for the framework to be useful (see Methods) and that control of a highly nonlinear neural system is possible even when agnostic to the true structure of the neural manifold.

Given the latent dynamics specified in equation (3), the control problem is to find an external stimulus $\mathbf{u}$ such that the time evolution follows a specified trajectory $\mathbf{z^{*}}$. The field of control theory provides a rich mathematical background to find system inputs to achieve desired system outputs. Many techniques exist with feedback (closed-loop) methods being particularly attractive due to their ability to correct for unknown perturbations to the system. Broadly speaking, feedback controllers can be categorized as being either reactive or anticipatory. Reactive controllers use present and past errors in state tracking (the difference between $\mathbf{z^{*}}$ and $\mathbf{z}$) to find a control signal $\mathbf{u}$. The classic PID controller is a popular implementation of a reactive controller due to its computational simplicity and strong performance. A notable issue with reactive controllers however is that one can only correct for errors in state tracking once they have been made. This can be unacceptable behavior if certain errors in state tracking are associated with pathological states. For example, suppose that a particular region of neural state space corresponded to epileptic firing. A reactive controller would only be able to respond after the system entered this region, at which point it may be much more difficult to re-establish control. In contrast, anticipatory controllers are designed to predict future errors, which can allow them to prevent the system from ever entering undesirable regions of state space. Model predictive control (MPC) is an anticipatory controller that uses a model of the system dynamics to make predictions on how present and future inputs will affect errors in state-tracking [14]. Additionally, MPC is a type of optimal controller because it attempts to find a control input sequence $\mathbf{u_{1:T}}$ that minimizes a loss function of the form

$$J(\mathbf{x_{0}})=\sum_{i=0}^{T}\ell(\mathbf{x_{i}},\mathbf{u_{i}}),\\$$ (5)

with constraints

$$\mathbf{x_{n+1}} =f(\mathbf{x_{n}},\mathbf{u_{n}})$$

$$\mathbf{x_{LB}} \leq\mathbf{x}\leq\mathbf{x_{UB}}$$

$$\mathbf{u_{LB}} \leq\mathbf{u}\leq\mathbf{u_{UB}},$$

where $\mathbf{x_{0}}$ is the value of the state at the current time step and $\ell(\mathbf{x_{i}},\mathbf{u_{i}})$ is the loss associated with $i$th time step (relative to $\mathbf{x_{0}}$), which is a function of the state variable(s) $\mathbf{x}$ and input(s) $\mathbf{u}$. The controller uses the dynamical model to simulate $T$ time steps into the future. Many types of loss functions are possible, but typically involve the state error and energy cost of the control signal. The constraints allow one to specify the dynamics of the system and to give lower and upper bounds for the state variables and inputs. More sophisticated versions of MPC allow for additional constraints where knowledge of any measurement or process noise can be incorporated [15]. Although MPC is only guaranteed to be globally optimal for linear systems with convex loss functions, it can also be used in many nonlinear systems [14, 16], in part because the controller can correct for model errors.

Only the first value of the sequence $\mathbf{u_{1:T}}$ is used as input into the system, with the optimization performed again at the next time step. By repeatedly solving this optimization problem and only using the first value, the controller can correct for errors in system modeling and anticipate future changes in the desired state trajectory. This anticipation of dynamics can result in better controller performance compared to traditional reactive controllers such as PID [16].

There are two major issues when using MPC in practice that are relevant to control of neural dynamics. The first is that the optimization procedure is computationally expensive and can result in poor controller performance if the time steps between measurements are small [14]. Neural recordings are often of high dimension (e.g. extracellular recordings with high-density silicon probes) and evolve at fast time scales. The stimuli (corresponding to the input $\mathbf{u}$) may also be of a high dimension, which will result in even more complexity in optimizing equation (5). We can now see the utility of modeling activity on the neural manifold as latent generative process. By doing so, we can solve the optimization problem in a low-dimensional state $\mathbf{z}$ and reduce the computational complexity. If we used the descriptive approach as in equation (1), we would need to optimize in the original measurement dimension $\mathbf{x}$ to force the activity on the manifold to follow a specified reference trajectory.

The second issue with applying MPC to neural systems that it requires a dynamical model of the system to be controlled [17]. Although there are many dynamical models rooted in biology for individual neurons, the putative latent dynamics of a neural manifold are an emergent property that is difficult to model from first principles. This requires us to take a data-driven approach, where unknown parts of the system can be modeled via function approximation and used to predict the time-evolution of the system in response to various inputs.

Fitting these models is achieved by observing a temporal sequence of the state and input variables with some sampling period $\Delta t$,

$$\mathbf{Z}=[\mathbf{z_{0}},\mathbf{z_{1}},...,\mathbf{z_{N}}],\mathbf{U}=[\mathbf{u_{0}},\mathbf{u_{1}},...,\mathbf{u_{N}}],$$ (6)

where $\mathbf{z_{n}}=\mathbf{z}(n\Delta t)$. A discrete-time model can be parameterized such that

$$\mathbf{\hat{z}_{n+1}}=f_{\theta}(\mathbf{z_{n}},\mathbf{u_{n}})$$ (7)

These type of models are often referred to as forecasting models since the model predicts how the system will change across time. The goal is to find a set of parameters $\theta$ such that given some initial state value $\mathbf{z_{0}}$ and some known input $\mathbf{U}$,

$$[\mathbf{z_{0}},\mathbf{\hat{z}_{1}},...,\mathbf{\hat{z}_{N}}]\approx[\mathbf{z_{0}},\mathbf{z_{1}},...,\mathbf{z_{N}}]$$ (8)

for any general temporal data sequence produced from the true dynamical system.

Data-driven approaches have been successfully applied to MPC problems in diverse fields [18, 19, 15, 20, 21], and the field is rapidly growing. In order for data-driven models to be useful for MPC applications in neuroscience, these models must be able to accurately predict the states to be controlled based only on observable state measurements, be agnostic to the number of hidden states, and generalize to a control scheme where control signals may be outside the training set. These challenges are not unique to neuroscience, but are still important to consider when selecting a data-driven approach to model the system dynamics. MPC has already been successfully applied in multiple areas of neuroscience research. At the individual cell level, simulated Hodgkin-Huxley neurons have been shown to be controllable via MPC using both biophysical [22, 23] and data-driven dynamics models [24, 25]. Additionally, both simulated and in vivo systems of neurons have been controlled with MPC using optogenetic stimulation as the control signal [26]. However, to the best of our knowledge, there has been no explicit attempt to apply these methods to activity on the neural manifold.

Our goal for this study was to provide a framework for controlling the latent dynamics on a neural manifold with MPC. An artificial neural circuit was simulated using a spiking neural network (SNN) that was driven by images of handwritten digits $\mathbf{u}$. The activity of the network was measured in accordance with an extracellular recording experiment where the spike times of the SNN served as neural states $\mathbf{x}$. Modern extracellular probes are able to record from dozens to hundreds of neurons simultaneously in both anesthetized and awake animals, but this is only a small subset of the neurons in a typical local circuit. Thus, in this simulation we only used a subset of the simulated neurons to fit a latent dynamics model of the whole network. We also added a substantial amount of random synaptic noise to simulate the influence of uncontrolled spontaneous activity arising from unobserved neurons within and outside the local circuit.

Using variational autoencoders (VAEs), we found low-dimensional representations of both the neural states and stimulus images. A latent dynamics model for the controller was fit using the lower-dimensional states and inputs (referred to as $\mathbf{z}$ and $\mathbf{v}$ respectively). We then found optimal latent inputs via MPC to force the latent states on the manifold to follow specified reference trajectories $\mathbf{z^{*}}$. We first controlled the latent dynamics of the SNN to stay at specified set points in the presence of intrinsic system noise. This result provided a proof of principle that MPC can control activity of the neural manifold in a simulation of a experiment using extracellular recording. We then show that the controller is able to force the latent states to follow multiple time-varying reference trajectories, with the optimized visual stimuli showing striking differences. The ability to control the system to follow distinct trajectories would allow for experimenters to see if there is a causal relationship between activity on the manifold and macro-scale behaviors. Finally, we examined the relationship between the proportion of measurable neurons and MPC performance.

Figure 1: MPC Control Loop of Artificial Circuit. Exponentially filtered spike trains are encoded into the latent state $\mathbf{z}$ and is compared to a reference trajectory $\mathbf{z^{*}}$ to produce an error signal. The controller uses a model of the latent dynamics to find an optimal input that minimizes a loss function of the state error. This input is then projected back into the original stimulus dimension using the sVAE decoder which stimulates the artificial circuit.

2 Methods

2.1 Artificial Circuit

2.1.1 Architecture

The activity of an artificial circuit evoked by an external visual stimulus was simulated with a SNN composed of three layers: sensory, reservoir, and output. Each neuron in the SNN was modeled with a recurrent leaky integrate-and-fire (rLIF) model, with the discrete time approximation

$$V_{n+1}=\begin{cases}\beta V_{n}+w^{T}X_{n+1}+r^{T}S_{n},&\text{if }V_{n}<\Theta\\ 0,&\text{if }V_{n}\geq\Theta\end{cases}$$ (9)

where

$$V_{n}:\text{membrane voltage at the {n}th time step}$$

$$\Theta:\text{spiking threshold}$$

$$\beta:\text{decay parameter}$$

$$X_{n}:\text{feedforward input vector at the {n}th time step}$$

$$w:\text{feedforward weights}$$

$$S_{n}:\text{layer spiking vector at the {n}th time step}$$

$$r:\text{recurrent weights}$$

Whenever the $V$ variable was reset to $0$, a spike was recorded at that time step. For a given layer with $N$ neurons, this produced a binary vector $S_{n}=[s_{1},s_{2},...,s_{N}]^{T}$ where the value of each element was either a $0$ or $1$ indicating if the corresponding neuron had fired at that time step. This allowed the activity of each neuron in a layer to be affected not only by its own firing (i.e., spiking inhibition/facilitation), but also to receive inputs from the other neurons in that layer.

Each neuron in the sensory layer received feedforward input in the form of a grayscale image reshaped into a 784 dimensional vector. This input served as the external stimulus that was the primary driver of SNN activity. Gaussian noise was also added to the feedforward inputs of each neuron in every layer to produce stochasticity in activity. This noise modeled the effects of natural variability in neural firing and the effects of unknown exogenous inputs. Thus the subthreshold activity of the three layers was given by the equations

$$\text{{Sensory}}: V_{n+1}^{sen}=\beta V_{n}^{sen}+w_{sen}^{T}(I_{n}+\epsilon)+r_{sen}^{T}S_{n}^{sen}$$ (10)

$$\text{{Reservoir}}: V_{n+1}^{res}=\beta V_{n}^{res}+w_{res}^{T}(S_{n}^{sen}+\epsilon)+r_{res}^{T}S_{n}^{res}$$ (11)

$$\text{{Output}}: V_{n+1}^{out}=\beta V_{n}^{out}+w_{out}^{T}(S_{n}^{res}+\epsilon)+r_{out}^{T}S_{n}^{out}$$ (12)

where $\epsilon\sim N(0,\eta^{2})$ for each element in the feedforward input and $I_{n}$ is the stimulus image presented at the $n$th time step. Note that $\beta$ incorporates both the membrane time constant and the time step of the discretization of the continuous LIF [27]. The time units for the simulation can be arbitrarily chosen by modifying this parameter. For simplicity, we use the units of milliseconds (ms) for each time step in all simulations which implicitly imposes a spiking refractory period of 1 ms. See Supplementary Information T1 for SNN training and rLIF hyperparameters.

2.1.2 Training the Network

In principle the feedforward and recurrent weights for each neuron could be randomly distributed, such as in reservoir computing [28]. However, to give the network dynamics that implement a useful computation [29], the weights were trained to perform a classification task. Using the Python package snntorch [27], the SNN was trained to accurately classify digits from the MNIST data set [30]. Training SNNs requires additional considerations compared to traditional artificial neural networks. The resetting of a neuron’s membrane voltage when it reaches the threshold $\Theta$ produces a non-differentiable function [27] making training with gradient descent impossible. One solution to this problem is to use surrogate gradient descent [31], where the non-differentiable function is preserved in the forward pass of the network but is replaced with a sigmoid function during the backward pass. This results in a function differentiable everywhere and allows the network to be trained with valid gradients.

The MNIST data set is composed of handwritten images of the digits 0 through 9. Classification of the images was performed by associating the label of the image with a corresponding neuron in the output layer of the SNN (e.g., label 4 with neuron 4). For time step $n$, the cross-entropy $\ell_{n}$ was given by

$$p_{n}^{i}=\frac{\exp(V_{n}^{i})}{\sum_{k=0}^{9}\exp(V_{n}^{k})}$$ (13)

$$\ell_{n}=-\sum_{i=0}^{9}y_{i}log(p_{n}^{i}),$$ (14)

where $V_{n}^{i}$ is the membrane voltage of the $i$th neuron which corresponds to the prediction of label $i$, and $y$ is a one-hot encoded vector of the true label. Due to the inherently temporal nature of SNNs, one must specify how many time steps a stimulus is be presented before class prediction takes place. This can be interpreted as a combination of reaction time and evidence accumulation. Thus the loss function to be minimized is given by

$$\mathcal{L}_{CE}=\sum_{n=0}^{T}\ell_{n},$$ (15)

where the cross-entropy loss at each time step is summed for some trial time length $T$. This forces the neuron of the associated predicted class to have the highest firing rate compared to the other neurons in the output layer. For the given time window that the image is presented, the feedforward and recurrent weights in the SNN are updated using backpropagation through time (BPTT).

The first half of the MNIST data set ($n=35,000$) was used for the training and validation of the SNN with an 80/20 split. Due to the stochastic nature of the noise added to every neuron in the AC, performance after training was assessed by presenting the stimuli 30 times to obtain an accuracy distribution. Accuracy on the training (n=28,000, $M=43.7$%, $s=8.4$%) and validation ($n=7,000$, $M=43.9$%, $s=8.4$%) set were largely similar, indicating there was no overfitting to the training data. Although the accuracies were far below what would be considered competitive performance on a classification task, the purpose of training the SNN was to ensure the connections between the neurons were not random.

2.2 Dimensionality Reduction of Stimuli and Neural States

The dimensionalities of the MNIST digit stimuli and SNN neural activity were reduced using variational autoencoders (VAEs). The parameters for encoders $f_{\theta}(.)$ and decoders $h_{\phi}(.)$ were found by minimizing the loss function

$$\ell_{VAE}=\ell_{recon}+\alpha\ell_{KL}$$ (16)

$$\ell_{recon}=\sum_{i=1}^{B}||\mathbf{x_{i}}-h_{\phi}(\mathbf{z_{i}})||_{2}^{2}$$ (17)

$$\ell_{KL}=\sum_{i=1}^{B}\sum_{j=1}^{k}(1+log\>\sigma^{2}_{ij}-\mu^{2}_{ij}-\sigma^{2}_{ij})$$ (18)

where $\alpha$ is a hyperparameter that scales the importance the of KL-divergence term in the loss function and $B$ is the batch size used for gradient descent. Because the latent variable $\mathbf{z}$ is parameterized by a Gaussian, minimizing this loss function is equivalent to maximizing the evidence lower bound (ELBO) [32].

The VAE framework was chosen for two reasons. First, VAEs embed high-dimensional data into nonlinear low-dimensional subspaces. This allows for a flexible approach for finding the latent dynamics of the SNN and obtaining a more generalizable compression of the data compared to linear transformations [33]. Second, the use of KL-divergence promotes nearby values of $\mathbf{z}$ in the latent space to be decoded into similar values in the original space of $\mathbf{x}$ [34]. By doing so, this allows for easier interpolation between a set of training data points in latent space when constructing the dynamics model and finding the optimal latent inputs with MPC. If a basic autoencoder (i.e., $\alpha=0$) was used instead of a VAE, nearby points in latent space would not be guaranteed to be similar in the original measurement space, which would obviously be detrimental to modeling dynamics in the latent space.

The control problem can now be expressed as given a latent embedding of the neural states $\mathbf{x}$ and stimulus states $\mathbf{u}$,

$$\mathbf{z_{n}}=\mathbb{E}(f_{\theta}^{\text{neural}}(\mathbf{x_{n}})),\mathbf{v_{n}}=\mathbb{E}(f_{\theta}^{\text{stimulus}}(\mathbf{u_{n}}))$$ (19)

we seek to find an optimal set of latent inputs

$$\mathbf{v}_{1:T}^{*}=\operatorname*{arg\,min}_{\mathbf{v}_{1:T}}\sum_{n=0}^{T}\ell(\mathbf{z_{n}},\mathbf{v_{n}})$$ (20)

with the dynamics model

$$\mathbf{z_{n+1}}=g(\mathbf{z_{n}},\mathbf{v_{n}}).$$ (21)

The SNN can then be stimulated with the decoded latent inputs

$$\mathbf{u_{n}}=h_{\phi}^{\text{stimulus}}(\mathbf{v^{*}_{n}})$$ (22)

to produce the neural states $\mathbf{x}$ at the next time step.

2.3 Stimulus VAE

A VAE was trained to find a low dimensional representation of MNIST digit stimuli (sVAE). The sVAE encoder was composed of two single-channel convolutional layers and a single feedforward layer. Each of the three layers used a ReLU activation function. The size of the latent space $\mathbf{v}$ was chosen to be 2. The sVAE decoder was symmetric to the architecture of the encoder with one key difference; a sigmoid activation function was used on the final output layer to ensure that the values of the reconstructed stimuli were between 0 and 1 (the bounds of all pixel values in the training images). See Supplementary Information T2 for layer hyperparameters.

The second half of the MNIST data set ($n=35,000$) was used for training and validation of the sVAE. Eighty percent of this data ($n=28,000$) was used for training and the remaining 20% for validation ($n=7,000$). The sVAE was trained via gradient descent in Pytorch using the Adam optimizer.

A sequence of latent inputs for latent dynamics model identification was constructed using the validation set of images. Discrete points in the latent sVAE space were obtained by running k-means clustering (n=100) on the low dimensional embedding of the validation images. Half of the centers were used for the dynamics model training ($V_{train}$) and half for testing ($V_{test}$). These discrete points were converted to a continuous time series by one of three methods: step function, fast interpolation, slow interpolation. The step function method took a sequence of centers and held each center constant for 500 time steps (ms). The fast and slow methods took the sequence and linearly interpolated values between each element in the sequence but at different time scales (200 ms for the fast and 1000 ms for the slow). Both the $V_{train}$ and $V_{test}$ sequences were 46.3 seconds long (46,300 time steps). The use of these three methods was to have a input sequence that could show how the latent states of the SNN responded to inputs changing at different time scales and frequencies. Figure 2 shows the latent input sequences and their sVAE decoded values.

Figure 2: Latent Embedding of MNIST Digits. A) Each colored dot is one of the MNIST digits embedded in the 2-dimensional nonlinear subspace of the sVAE encoder. Notice the clustering of the digits by label (color), indicating that digits with identical labels were often embedded in nearby latent space. In order to generate a latent sequence of inputs used to stimulate the artificial circuit, points from this latent space where sampled using k-means clustering (100 centers). B) Training and testing latent sequences where each generated using half of the centers from A. C) Three points from the training inputs decoded into the original stimulus dimension.

2.4 Neural VAE

Figure 3: Latent Embedding of Neural Activity. A) Results of the VAE on training SNN neural activity (nVAE). The exponentially filtered spikes are embedded in a latent 2D space through the nVAE encoder. Activity in this latent space can be projected back into the original dimension with the nVAE decoder. B) Results of the nVAE on the testing SNN neural activity.

The SNN was stimulated using the sVAE-decoded $V_{train}$ and $V_{test}$ input sequences with the resulting spiking activity used to fit a VAE for the neural states (nVAE). A random sample of 20% of the neurons ($n=122$) in the SNN was used to build the model. These neurons were the only units in the SNN that were measured for the entirety of experiments I and II. Recall that in a typical extracellular recording, only a subset of the neurons are observable. The purpose of this random sampling was to mimic the incomplete information that would be obtained in a real recording.

The measured binary spiking states were converted to continuous states with an exponential filter, where the smoothed state $\mathbf{x_{n}}$ of spiking state $\mathbf{y_{n}}$ is given by

$$\mathbf{x_{n+1}}=\omega\mathbf{y_{n}}+(1-\omega)\mathbf{x_{n}}$$ (23)

where $\omega$ was chosen to be $0.5$ and $\mathbf{x_{0}}=\mathbf{y_{0}}$.

The nVAE encoder and decoder were symmetric, with each having five feedforward layers. The smoothed state $\mathbf{x_{n}}$ was z-scored when entering the first layer of the encoder. A small value $\epsilon=1\times 10^{-5}$ was added to the denominator of the standardization since some time steps had near identical values. Each hidden layer used a ReLU activation function followed by a batch normalization layer. The size of the latent dimension $\mathbf{z}$ was chosen to be 2. Because the addition of noise to the rLIF neurons resulted in temporal spiking jitter to the same input stimulus, the variance of the latent representation $\mathbf{z}$ was constrained to be above 1. This was achieved by having the encoder learn the log of the variance instead of the variance directly and then applying the softplus function to the estimate. By using this minimum variance, it acted as a regularizing parameter to the noise present in the training data. See Supplementary information T3 for layer and training hyperparameters.

2.5 Latent Dynamics Model

Figure 4: Forecasting Performance of Latent Dynaimcs Model in Experiment I. (Left) Latent state $\mathbf{z}$ (red) forecasting on the training data. On top is the forecast for $\mathbf{z_{1}}$ and on bottom $\mathbf{z_{2}}$. In black is the actual latent trajectory of the training data. The product-moment correlation between the actual and predicted latent states is shown above each figure. (Right) Same as on the left, but with the testing data. The fits between the training and testing forecast are largely identical.

A linear latent dynamics model of the form

$$\mathbf{z_{n+1}}=A\mathbf{z_{n}}+B\mathbf{v_{n}}$$ (24)

was estimated using ridge regression with the L2-hyperparameter chosen via leave-one-out cross-validation. All model fitting was performed using the sklearn Python package. The dynamics model was fit with the data produced by the $V_{train}$ input sequence. Model performance was assessed by using an initial value of $\mathbf{z}$ and forecasting for the entire training sequence length. At every time step, the known value of $V_{train}$ and the model’s previous prediction of $\mathbf{z}$ was used to forecast the next value. An inadequate model would produce a time-series that was a poor approximation of the actual latent state trajectory produced by $V_{train}$ and could possibly even diverge due to compounding errors.

As an example, the forecasted values of the training data in Experiment I were a close fit to the actual $\mathbf{z}$ training trajectories as measured with the product-moment correlation ($R_{z_{1}}=0.87$, $R_{z_{2}}=0.73$). The possibility of overfitting was assessed by forecasting with the latent states elicited from $V_{test}$. The resulting predicted trajectory was also similar to the actual latent trajectory ($R_{z_{1}}=0.87$, $R_{z_{2}}=0.72$) indicating that the dynamics model would be useful for control with MPC. See Figure 4 for the forecasted latent trajectories of the training and testing data in Experiment I.

3 Results

3.1 Experiment I: MPC Can Force the Latent Dynamics on the Neural Manifold to Stay at set points

As a simple first experiment, the latent dynamics were controlled to follow a step function. This reference trajectory was composed of two set points in latent space, each held constant for 500 ms. The values were chosen by running k-means clustering on the latent state training data and using the resulting centroids as the set points. The controller had a predictive time horizon $T$ of 30 time steps and optimized the loss function

$$J(\mathbf{z_{0}})=\mathbf{z}^{\intercal}_{\mathbf{T}}\mathbf{S}\mathbf{z_{T}}+\sum_{i=0}^{T-1}\mathbf{z}^{\intercal}_{\mathbf{i}}\mathbf{Q}\mathbf{z_{i}}+\Delta\mathbf{v}^{\intercal}_{\mathbf{i}}\mathbf{R}\Delta\mathbf{v_{i}}$$ (25)

where

$$\mathbf{Q,S}=\begin{bmatrix}150&0\\ 0&150\end{bmatrix},\mathbf{R}=\begin{bmatrix}10,000&0\\ 0&10,000\end{bmatrix}.$$ (26)

No constraints other than the dynamics model were included. Due to the stochastic nature of the SNN responses, 50 independent trials of MPC were performed. The normalized mean square error (nMSE) was used to quantify the controller performance of the latent state, which normalizes mean square error by the difference between the maximum and minimum $\mathbf{z^{*}}$ values. This was done for ease of comparing controller performance across dimensions with different scales. All MPC optimizations and implementations were performed using the do-mpc python package [35]. This package utilizes CasADi [36] and IPOPT [37] for interior-point optimization and automatic differentiation methods.

The controller was able to achieve good performance, especially when considering the noise in the system and that only 20% of the neurons of the SNN were observable. Even though the noise that was added to all neuron inputs was Gaussian, the nonlinearities in the rLIF models propagate highly complex noise structures throughout the network. At the beginning of control and after the set point was switched, the system converged to the desired location in the latent space within 20–50 time steps and then remained there despite continual noise inputs that produced large perturbations (Figure 5A). When the set point was changed, there was a large excursion of the control stimulus within its latent space, and there were large fluctuations in the input throughout the experiment as the controller corrected for noise-induced perturbations.

Because the latent space has only two dimensions, we can visualize the dynamics of the forecasting model as a vector field and the current and desired states of the system as points within this field. The dynamics can be decomposed into an autonomous component $Az_{n}$ and the forcing from the stimulus $Bv_{n}$, and as the combined system is linear, there is at most a single equilibrium. As seen in Figure 5B, at time points when the set point was constant, the controller chose stimuli such that the dynamics had a stable equilibrium close to the set point. During the transition to a new set point, the controller initially moved the equilibrium to a location beyond the set point but then settled to an input that created a fixed point at the desired location. This overshoot may have allowed the controller to more rapidly move the system to its desired location.

As seen in Figure 5C, there was a qualitative change in the activity of the population when the reference point changed at 500 ms. Interestingly, while there was an initial drop in the average firing rate, it returned to a similar value as the system converged on the new set point. This indicates that the latent states are not just functions of the population firing rate, but are functions of particular firing patterns and activity from specific neurons. We can also examine the stimuli that were actually presented to the SNN during control. As seen in Figure 5D, the control stimuli all had digit-like characteristics, which undoubtedly reflects the fact that the sVAE was trained only on digits. Notice, however, that the image presented during the transition to the second reference point is a larger and more intense ‘2’ than the image presented after the system had converged.

Figure 5: Results of Experiment I. A) (Left) The results of MPC of the SNN across 50 trials. The reference trajectories (black) were composed of two set points that changed at 500 ms. In light red are the controlled latent trajectories across the 50 trials and the average of these trajectories are shown in dark red. (Right) The latent inputs produced by the optimizer. In light blue are the inputs used across the 50 trials and the average input shown in dark blue. B) At each time step, the latent dynamics model predicts a vector field on the latent states. Three snapshots of this vector field are shown at 200, 500, and 800 ms. The desired reference points are shown by a black X and the controlled latent states are indicated by red dots. Notice the magnitude of the vectors when the reference point changes at 500 ms. The blue dots are the fixed points of the dynamics model given the value of the latent inputs $\mathbf{v}$. C) (Above) The smoothed spike trains of the measured SNN neurons across the 50 trials are shown in light blue with the average shown in black. The smoothed spikes of a sample trial is shown in dark blue. (Below) The measured spike trains of the sampled trial corresponding to the dark blue curve above. D) Decoded latent inputs from the optimizer across three time points of the sampled trial. These are the visual stimuli that drove activity in the network. An animated visualization of the latent state and the stimuli used for control can be found in Supplementary Information V1. Link: https://doi.org/10.6084/m9.figshare.26072803.v1

3.2 Experiment II: Distinct Reference Trajectories Produce Different Levels of Controller Performance

The previous experiment revealed how the controller successfully used the stimulus-dependent dynamics of the forecasting model to drive the network to specific locations in the latent space. However, the forecasting model is only a linear approximation of the simulated network’s dynamics, which are nonlinear and of a much higher dimension. To examine how control could be used to probe the underlying dynamics of the network, it is not sufficient to characterize the start and end points of a trajectory, but instead what specific path was taken. For our next experiment, the reference trajectory was replaced with two time-varying functions (reference trajectory 1 and 2). Each of these reference trajectories had the same initial ($\mathbf{z_{0}^{*}}$) and final values ($\mathbf{z_{f}^{*}}$), but took different paths through latent state space. Using a parameterized function of a circle passing through points $\mathbf{z_{0}^{*}}$ and $\mathbf{z_{f}^{*}}$, trajectories 1 and 2 were the opposite arcs of the resulting circle. This ensured that both trajectories were of equal length through the latent state space. The observation model from Experiment I was used for this task (20% observability) and the values of $\mathbf{z_{0}^{*}}$ and $\mathbf{z_{f}^{*}}$ were the set points from the same experiment. Fifty control trials were performed for each of the reference trajectories with each trial having the same MPC hyperparameters from Experiment I.

Figure 6: Latent States and Inputs for Different Reference Trajectories. A) The controlled paths (green and indigo) for the two reference trajectories (black) in Experiment II. Each of the reference trajectories had the same initial and final values ($\mathbf{z_{0}^{*}}$ and $\mathbf{z_{f}^{*}}$ respectively). The paths for 50 trials with trajectory 1 as the reference are shown in green and in indigo for trajectory 2. For comparison, the paths from Experiment I 200 ms before and after the change in set point are shown in red. The set points were identical to the $\mathbf{z_{0}^{*}}$ and $\mathbf{z_{f}^{*}}$ used here, but the path between them was not constrained to follow certain values. The dark colors show the average of the respective path types. B) Latent states $\mathbf{Z_{1}}$ (teal) and $\mathbf{Z_{2}}$ (magenta) show differing regions of poor control. While the average paths are near the reference values (dark color), each of the fifty trials (light color) may oscillate around the reference trajectory. Below are the latent inputs used to control the system across the fifty trials (light color) with the average values (dark color). C) Differences in nMSE of latent states and RMS of latent inputs across trajectories. Notice the difference in relative errors between the latent states across conditions. This indicates that errors in control along a particular dimension are not independent of errors in the other dimension. The RMS of the latent states are higher in trajectory 2 for both latent inputs. Since the same sVAE was used, the numerical values of the error are directly comparable. See Supplementary Information V2, V3 for animated visualization of the controlled trajectories.

As seen in Figure 6A, the controller successfully forced the network to take two different trajectories through the latent space, both of which were different from the trajectories the system used in Experiment I. However, the paths the system took often deviated substantially from the reference trajectory. For both trajectory 1 and 2, it appeared as though the system was attracted to specific regions of the state space and that there were other regions that the controller was unable to force the system to enter. These regions likely correspond to attractors and separatrices in the true dynamics of the SNN. Interestingly, individual trajectories appeared to show large oscillations in some regions of the state space, which could reflect rapid switching between stable equilibria or a limit cycle. It remains to be seen if a nonlinear forecasting model would enable the controller to force the system through repulsive regions of the state space better (see Discussion). There were differences between the two trajectories in the average state error (Figure 6B) and the variance of the latent inputs produced by the controller (Figure 6C). This indicated that more “power” in the latent space was needed to control the system to follow trajectory 2, even though the lengths of the two trajectories were the same.

Examining the activity in the measurement space, we see obvious differences in the spiking behavior of the SNN (Figure 7A). The visual stimuli produced from decoding the latent inputs were initially similar between the two reference trajectories but then diverged as the trajectories separated (Figure 7B). This is consistent with the behavior of the SNN, which was trained to produce distinct responses to different digits; thus, different regions of the state space are likely to correspond to specific digits. One implication of this is that MPC of the latent dynamics could reveal if specific kinds of stimuli correspond to particular latent trajectories or if the activity in the latent space is driven by the differences in the stimulus at every time step (e.g. prediction error [38]).

Figure 7: Recorded Spikes and Control Inputs for Experiment II.A) Spike trains from a single trial for the reference trajectory 1 and 2 conditions. Note that the high densities of spiking in the Trajectory 2 condition correspond to the same time interval as the high oscillations in the latent space ($\sim 500-800$ ms). B) The decoded latent inputs that were used to stimulate the SNN during specific time points of the single trials in A.

3.3 Experiment III: The Proportion of Observable Neurons Affects Controller Performance

To investigate how robust the control strategy was to partial observation of the neural population, the procedure from Experiment I was performed on different subsets of neurons randomly sampled from the SNN. For each level of neuron observability, 10 independent ensembles were sampled from the SNN ($n=80$). An nVAE was trained for each ensemble, but in order to ensure a fair comparison, the architecture was kept constant (except for the size of the input layer). As shown in Figure 8A, there was a tradeoff between the performance of the nVAE and the performance of the forecasting model. With a small number of neurons, (1% observability, 6 neurons), the nVAEs achieved almost perfect reconstruction for both training and testing data. However, the forecasting models performed very poorly, indicating that the stimulus-dependent dynamics of the full SNN could not be inferred from the behavior of just a few neurons. With larger proportions of the population observed, the performance of the nVAEs decreased, indicating that a latent space with only two dimensions was not sufficient to capture all the variance in the neural activity. There was a floor to this degradation of performance, and the reconstruction errors plateaued around 30% observability. This is consistent with a high level of correlation within the SNN. The performance of the forecasting models increased with the number of observed neurons. With 20% or more of the neurons included in the model, the dynamics of the SNN within the latent space could be predicted with a high level of accuracy.

The latent dynamics inferred at each level of observability were controlled to follow a step function of two set points obtained from k-means clustering, as in Experiment I. Ten trials of MPC were run for each of the models using the same previous loss function and hyperparameters. In practice, the hyperparameters of the controller would be tuned after fitting the forecasting model, but kee** them the same provides a fairer basis for comparison. Overall, controller performance was better with larger proportions of neurons observed (Figure 8B,C). Additionally, the variability of average nMSE across ensembles was typically higher when fewer neurons from the population were observed. This is unsurprising since there are many more possible ensembles to sample in the lower observation conditions which may differ in their controllability. However, it was still possible to sample bad ensembles when observing a large proportion of the neurons. The three highest levels of neuron observability all had an ensemble that produced higher average nMSE for a single latent state than the average of the nMSE distribution of the three lowest levels (Figure 8B).

Figure 8: Effects of Percentage of Observable Neurons on Control.A). The performance of the nVAE and forecasting models plotted as a function of the proportion of the neurons observed. For each observation condition, ten independent ensembles were sampled from the SNN. The average performance across ensembles is plotted for each of the observation percentages. One ensemble in the 1% observation condition produced divergent control trajectories and was removed from analysis. B) The average nMSE of the latent states across ensembles for each level of observation percentage. As the percentage goes up, the nMSE trends toward smaller values. However, as seen clearly in the 40% observation condition, bad samples are still possible.C) Reference (black) and controlled (red) trajectories for examplar ensembles for each observation percentage. Light traces show trajectories for individual trials ($n=10$) and dark red shows the average.

4 Discussion

In this study, we demonstrated that MPC can be used to control the latent states of an artificial spiking neural network using low-dimensional, data-driven models of the dynamics. Despite the highly nonlinear dynamics of the SNN, the latent states could be fixed to certain set points or forced to follow arbitrary, time-varying reference trajectories. Control was possible even when only a limited proportion of the neurons was observed and there when there were unknown sources of noise in the network. Reducing the dimensionality of the neural activity and the stimuli used to control the network can therefore make MPC computationally tractable for high-dimensional neural systems. Based on this approach and this proof of principle, it should be possible to use MPC for real-time control of latent dynamics in sensory-driven biological networks using extracellular recordings. Achieving this level of control over neural networks would enable experimenters to probe the dynamics of neural circuits in ways that we have begun to outline here.

In Experiment II, we found that the state of the network could be forced to take different trajectories by providing distinct sequences of latent inputs, but there were regions of the latent space where control was better or worse. One possibility is the latent dynamics model made better predictions for certain regions of latent state space, where the dynamics of the full network were more linear. The use of a nonlinear dynamics model may result in a better approximation of the latent vector field, but at the cost of complexity in both model fitting and MPC optimization. Another possibility is that the VAEs may have learned a dimensional transformation that preserved the distributional information in the high-dimensional space but that distorted the temporal relationships needed to infer the latent dynamics in certain regions [11]. By fitting the VAEs and latent dynamics models in two distinct phases, this may have produced dynamics models that were not optimal for the subspaces found by the VAEs. An alternative approach would be to simultaneously learn the dimensionality reduction and latent dynamics model in a single step, for example by extending a recurrent switching linear dynamical systems approach [13] to nonlinear dimensional reduction or by using structured variational autoencoders [39]. Although this would likely improve the performance of the controller, it would also introduce greater complexity when fitting the model, as additional hyperparameters would be needed to scale the influences of the VAE reconstruction, latent forecasting, and measurement space forecasting. Extensive work may be needed to identify the best approaches for balancing forecasting accuracy with computational efficiency for a given system and set of scientific questions.

Recent years have seen the development of high-density silicon electrodes that can record extracellular spikes from hundreds to thousands of neurons simultaneously [40, 41], but this is still only a tiny fraction of the number of cells that participate in local neural circuits. In Experiment III, we found that the average performance of the forecasting model and the controller decreased as the proportion of neurons observed in the SNN decreased. However, these lower observation conditions also had higher variance in the distribution of average nMSE, indicating that good control over small ensembles is possible. It is also well-known that when observing a subset of states of a dynamical system, time-delay embedding the measurements gives information on the full dynamics [42]. In each of the experiments here, time-delay embedding was not used to fit the VAEs or latent dynamics models for simplicity. Although this did not appear to impact the controller performance for the higher observation percentage models, it may have resulted in poorer performance as the percentage decreased in value. It would be interesting to examine if using time delays would result in increased performance even when the number of observed neurons is very low. Other work has shown the success of using time-lagged autoencoders to find latent dynamics models [43], and methods exist to find optimal time delays and the number of embedding dimensions [44]. This would introduce additional complexity in practice however, because both the dimensionality reduction and dynamics model need to be fit quickly when recording from living neurons, otherwise the experiment may not be feasible in a laboratory setting due to cell death or electrode drift.

The divide between the descriptive and generative perspectives on neural manifolds is related to the question of whether there is a low-dimensional dynamical system that emerges from the much larger and more complex dynamical system defined by the biophysics of intrinsic and synaptic currents in large populations of neurons. In this study, the true dynamics of the SNN were defined by a large, nonlinear dynamical system composed of hundreds of interconnected artificial neurons. The activity of this network could be represented in a low-dimensional neural manifold (Figure 3) with a clustered structure that is suggestive of attractor basins. The MPC framework developed in this study provided a method for experimentally probing the dynamics on this manifold to better understand its structure and function. Our results are consistent with the generative perspective in that a simple linear approximation of the dynamics in the latent space was sufficient to achieve a high level of control over the network. However, we also found that that there were regions of the state space where the latent dynamical model was not a good enough model of the underlying system to provide strong control, a result that can be interpreted as support of the descriptive perspective or as a source of insight for how to improve the latent model. If MPC can be applied in biological systems, it could provide a strong test of whether manifold activity is causal by enabling experimenters to see if specific organismal behaviors can be produced by controlling the latent neural dynamics. This would be an powerful tool in furthering the understanding of how complex behaviors and computations emerge from the structure of neural circuits and the dynamics of their activity.

It may be important that there was a topological alignment between the latent spaces for neural activity and the stimulus set. The sVAE successfully discovered a dimensional map** that separated the digits and their variants into 10 distinct and well-separated clusters (Figure 2A). There was also clear evidence of clustering in the neural latent space that mapped in an orderly way to which digit the stimulus was (Figure 3). Though it is beyond the scope of the present study, it would be interesting to explore how the stimulus latent spaces discovered by other dimensionality reduction methods impacts how the latent dynamics model performs in forecasting and control. It is also interesting to consider that the dynamics of the SNN and the latent map** of the sVAE were both learned from the distribution of the stimulus, but through radically different methods. If topological alignment is required for control in this classification task and in other computational problems such as the ones explored by Susillo and Barak [29], it may speak to a simple but profound theory that derives from the ideas of James [45] and Hebb [46]: that learning is a process of aligning the latent dynamics of neural circuits to the latent dynamics of the physical world. Optimal control, both as theory and as a method for more precise experimental manipulation, may be of benefit in testing this theory in biological systems.