Fast inference of spinal neuromodulation for motor control using amortized neural networks

Abstract

Objective.Epidural electrical stimulation (EES) has emerged as an approach to restore motor function following spinal cord injury (SCI). However, identifying optimal EES parameters presents a significant challenge due to the complex and stochastic nature of muscle control and the combinatorial explosion of possible parameter configurations. Here, we describe a machine-learning approach that leverages modern deep neural networks to learn bidirectional mappings between the space of permissible EES parameters and target motor outputs.Approach.We collected data from four sheep implanted with two 24-contact EES electrode arrays on the lumbosacral spinal cord. Muscle activity was recorded from four bilateral hindlimb electromyography (EMG) sensors. We introduce a general learning framework to identify EES parameters capable of generating desired patterns of EMG activity. Specifically, we first amortize spinal sensorimotor computations in a forward neural network model that learns to predict motor outputs based on EES parameters. Then, we employ a second neural network as an inverse model, which reuses the amortized knowledge learned by the forward model to guide the selection of EES parameters.Main results.We found that neural networks can functionally approximate spinal sensorimotor computations by accurately predicting EMG outputs based on EES parameters. The generalization capability of the forward model critically benefited our inverse model. We successfully identified novel EES parameters, in under 20 min, capable of producing desired target EMG recruitment duringin vivotesting. Furthermore, we discovered potential functional redundancies within the spinal sensorimotor networks by identifying unique EES parameters that result in similar motor outcomes. Together, these results suggest that our framework is well-suited to probe spinal circuitry and control muscle recruitment in a completely data-driven manner.Significance.We successfully identify novel EES parameters within minutes, capable of producing desired EMG recruitment. Our approach is data-driven, subject-agnostic, automated, and orders of magnitude faster than manual approaches.

Keywords: approximate Bayesian inference; artificial neural networks; epidural electrical stimulation; machine learning; neuromodulation; spinal cord stimulation.

Figure 8:. Representative data samples.

We randomly selected EES trials from a single animal to visualize bilateral rectified EMG responses to a single EES pulse train. Data shown here includes a 100ms pre-stim and a 400ms post-stim period. The onset of EES response is marked by the vertical gray dotted line at 0ms.

Figure 9:. Video acquisition setup.

We recorded sRGB videos from three cameras simultaneously. The cameras were positioned to maximally cover the range of our subject’s hind limb motion.

Figure 10:. EMG serves as a proxy for physical limb motion.

The summarized EMG activity recorded at each muscle contact highly correlated with the EES-evoked angular displacement of the ipsilateral hind limb. Each point in these plots refer to a unique trial. Data presented here only includes trials where EES was applied on a midline electrode contact.

Figure 11:. Partitioning the EES parameter space

We “hold out” trials that used EES parameters in the marked regions from the forward neural network model training. These trials are exclusively used to evaluate the goodness of fit of the forward model and test its generalization capabilities.

Figure 12:. Training and validation curves

We demonstrate the ability of f_ϕ (trained on each animal individually) to effectively learn the forward mapping from EES to EMG. Here we show loss curves over the course of training as given by L1 errors, as well as correlations between the predicted and ground truth validation EMG.

Figure 1:. Goal:

To automatically determine optimal Epidural Electrical Stimulation (EES) parameters (θ) that yield target sensorimotor network activation (x). (a) EES is thought to primarily recruit large-diameter, afferent fibers at their entrance to the spinal cord through dorsal roots when stimulation is delivered at motor threshold. (b) Using mechanistic models to study the transformation of EES into motor outputs such as those measured by Electromyography (EMG), referred to here as the forward problem, impedes the rapid discovery of solutions to the more general inverse problem (c). We propose to leverage recent progress in deep learning to amortize the cost of simulating spinal sensorimotor computations via a universal function approximator f_ϕ, which can subsequently be reused by another neural network f_ψ for solving our problem of interest.

Figure 2:. Measuring motor outputs evoked by epidural spinal stimulation.

(a) Intra-operative X-Ray image of the implanted electrode arrays in L3-6. (b) Experimental setup for data acquisition in sheep. Spinal stimulation (μ ∈ [10, 100]Hz, a ∈ [60, 1500]μA) was delivered for 300ms at the start of each trial. Here, we show a sample trial stimulating a Caudal electrode with μ = 100Hz, a = 420μA. Motor output is measured by both electromyography (EMG) and kinematics as given by the Hip, Hock and Hoof positions. (c-d) Evoked EMG responses measured upon systematic changes to EES frequency and amplitude respectively reveal the nonlinearities of spinal sensorimotor computations. EES onset is indicated by gray dashed lines.

Figure 3:. Forward model f_ϕ accurately predicts EES-evoked motor outputs in silico

(a) f_ϕ learns to transform parameterized EES into a summary EMG measure by backpropagating error gradients. (b) Numerical evaluation of the forward model predictions on data from four different sheep shows that it captures nearly all of the explained variance (L1 → 0) in the motor outputs. The red crosses denote a “random” baseline performance (c) Marginal parameter recovery of parameters μ (left) and a (right) for each of the four sheep using the correspondingly trained f_ϕ confirms successful amortization of spinal sensorimotor computations.

Figure 4:. Neural network-guided optimal recruitment of muscles.

(a) Electrode-conditioned inverse models use the trained forward model f_ϕ to estimate the posterior density over EES parameters given a target pattern of EMG activity. (b) Sampling from this posterior density yields optimal EES parameters which we test in silico and in vivo. We successfully demonstrate the ability to control sensorimotor network activation through our fully automated procedure.

Figure 5:. Discovering functionally equivalent EES protocols.

(a-c) Our automated search through the EES parameter space optimally identifies pulse train parameters for different spinal electrode contacts to evoke similarly correlated sensorimotor responses. (d) Our approach is also well-calibrated and correctly assigns uniformly low posterior probabilities to sites where no EES configuration can elicit desired EMG response.

Figure 6:. Neural network guided EES evokes desired motor outputs in vivo.

We measure the efficacy of our model-guided EES trials by computing the correlation between the evoked EMG response and the target EMG response. We note that this correlation profile closely mimics the soft upper-bound given by the correlation between EMG signals evoked by the same EES parameters recorded at different times.

Figure 7:. Parameter trade-offs in EES emerge naturally from our system.

Our forward model f_ϕ is able to identify discernably different EES parameters that evoke identical motor outcomes. This system as a whole allows for systematic investigation of compensatory spinal mechanisms involved in the sensorimotor pathways.