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 minutes, capable of producing desired target EMG recruitment during in vivo testing. 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.

1. Introduction

Damage to the central nervous system can lead to the chronic impairment of motor, sensory, and autonomic function [1, 2]. However, recent advances within the field of therapeutic electrical neuromodulation have demonstrated restoration of a wide range of capabilities, including the ability to voluntarily control muscles after chronic, motor complete spinal cord injury (SCI) [3]. Specifically, epidural electrical stimulation (EES), wherein an electrode array is implanted in the epidural space below the SCI lesion site (Fig. 1a), has been shown to support walking, standing, cycling, and swimming years after severe SCI [4–12]. Evidence from computational models [13–15] and experimental studies [16–18] suggest that EES primarily recruits large-diameter, afferent fibers at their entrance to the spinal cord through dorsal root activation, which engages spinal reflex circuitry resulting in activation of interneuronal circuits and projecting motor neurons [19]. However, evoking motor activity in vivo is a complex, multi-faceted process. Factors including the anatomy of the spinal cord, the placement of electrodes, the extent of the spinal injury, EES-supported plasticity, and the current functional state of the spinal circuits orchestrate spinal neural activity and motor outputs that are highly dynamic and non-linear. According to a recent survey of SCI physicians, the lack of closed-form solutions or clear “guidelines” outlining stimulation parameters with deterministic effects has become one of the main barriers for the clinical translation of EES [20]. Therefore, there is a need to develop robust and systematic algorithms to automatically identify EES parameters that drive sensorimotor networks to aid in the restoration of motor function.

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.

The current strategy for determining patient-specific sensorimotor responses to EES, here defined as the EES evoked transynaptic activation of motor neurons via a subset of sensory fibers, is to construct a detailed associative map documenting the spatiotemporal patterns of electromyographic (EMG) activations [5, 7, 18, 21]. EES pulses are applied at varying amplitudes while the subject is at rest and the corresponding muscle responses are recorded to compute recruitment curves. This associative map is subsequently used to construct EES protocols to drive targeted muscle recruitment in more complex tasks such as standing or stepping. There are, however, two primary limitations with this approach. First, relying on associative maps implicitly assumes a deterministic relationship between EES and sensorimotor network activation. Yet, the inverse problem of identifying EES parameters that produce specific patterns of muscle activity is complex, under-constrained and stochastic (Fig. 1). Second, the recruitment-curve based approach for building an atlas of EMG activations generally overlooks covariances, or inter-dependencies, between EES parameters and muscle outputs. Deriving recruitment curves requires many repeated stimulation episodes, which further limits the number of muscle response combinations that can be tested. For this reason, the standard practice in the field is to rely on recruitment curves derived by varying stimulation amplitude alone. However, such measurements do not account for the signal variance incurred by jointly manipulating stimulation amplitudes along with other EES parameters, such as the frequency and location of the stimulation. Given the highly non-linear nature of EMG responses to EES (Fig. 2c), simple interpolation using recruitment curves may lead to erroneous or sub-optimal EES identification.

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.

To address these issues, we adopt a system identification perspective [22] and cast the inverse problem of computing optimal EES parameters, for targeted activation of sensorimotor networks, as probabilistic inference (Fig. 1). Conceptually, the idea is to treat EES parameters θ, and EMG activity x as specified by a joint probability model p(θ, x), in order to estimate the posterior density p(θ|x) via Bayes’ theorem:

$$p(\theta \mid x)=\frac{p(x\mid \theta )p(\theta )}{p(x)},$$ (1)

where p(θ) is a suitable prior over EES parameters, p(x|θ) is the “likelihood” of observing x under a certain set of EES parameters θ, and p(x), known as the evidence, signifies the probability of observing an EMG pattern x. The mathematical object p(θ|x) is traditionally a probability density function from which one can draw samples, for example via Markov Chain Monte Carlo (MCMC) [23]. This formulation naturally avoids the aforementioned caveat of determinism. Moreover, for a specific EMG observation x₀, EES parameter values sampled from the conditional distribution p(θ|x = x₀) will reflect the inherent trade-offs between the different components of θ (the frequency and amplitude of EES and the stimulation location), thus alleviating the covariance bottleneck.

Traditionally, mechanistic models are used as forward models to compute sensorimotor network activation in response to EES. One challenge with such statistical inference formulation has traditionally been the requirement that the class of mechanistic models under consideration has to be restricted to those with analytical likelihood functions because sampling algorithms such as MCMC rely on the evaluation of the likelihood p(x|θ) to sample from the posterior density. This severely restricts the class of models that can be considered and hence the generality of the approach. However, recent innovations in Approximate Bayesian Inference methods alleviate this problem making it possible to perform posterior inference from simulation models without explicit analytical access to p(x|θ) [24–26], thus permitting the inclusion of arbitrary forward models. However, despite these improvements, the practical cost of constructing, tuning, and simulating data from mechanistic models of spinal sensorimotor circuits remains prohibitively high, both from a time and human-effort standpoint.

To facilitate the rapid inference of EES parameters, we reformulate the standard inference to take advantage of recent advances in deep learning. Specifically, we propose to functionally approximate the computations of spinal sensorimotor circuits using a deep neural network f_ϕ (function f parameterized by weights ϕ) that map EES parameters onto corresponding EMG outputs. EMG signals were selected as the model output as they have traditionally been used to guide EES parameter optimization [4, 7]. We refer to f_ϕ as our “forward model” (Fig. 1b). Training f_ϕ once, effectively amortizes the cost of simulating spinal computations, which can then be efficiently reused for inference. Because of the embarrassingly parallel nature of deep neural networks coupled with hardware specializations, simulating EES-evoked sensorimotor activity data can be done exponentially faster than simulating mechanistic biophysical models. Additionally, end-to-end data-driven automation avoids the heuristics and biases that arise from manual interventions.

By training our forward model from data, the traditionally more challenging inverse problem of estimating EES parameters given data becomes tractable (Fig. 1c) for the following reason. Our amortized forward model f_ϕ serves as a close approximation of the data generating process underlying sensorimotor network activation following EES. Thus, it now suffices to compute posterior densities over the inputs to the forward model f_ϕ as opposed to parameters of mechanistic models such as muscle spindle feedback circuits (as in [27] for instance). Drawing from the recent successes of neural density estimators for simulation-based inference [24–26, 28–30], we define our “inverse” model f_ψ to be an invertible conditional flow, a special formulation of a neural network that learns the posterior over the EES parameters from data simulated using the forward model f_ϕ (Fig. 1c). The main source of biological stochasticity that we are addressing is the evaluation of EES evoked responses in the unloaded state due to factors such as differences in spinal cord anatomy and electrode placement that require individualized optimization of EES parameters. Although further optimization may be required in the contexts of active tasks that may increase stochasticity due to increased somatosensory input, the data presented here provides a methodology to evaluate EES to muscle activation mappings.

In summary, we introduce a general data-driven machine learning approach to automatically identify optimal EES parameters and produce desired target patterns of sensorimotor network and lower extremity musculature activations. To the best of our knowledge, our formulation is the first to cast the problem of spinal EES identification in a probabilistic framework. Our simulation-based inference solution to motor control is fast, only requiring a minimal compute infrastructure. We validate its utility through both extensive offline, as well as in vivo evaluations in sheep. Our approach provides a solution to the lack of systematic parameter selection that exists as a barrier to the clinical translation of EES. We make our code publicly available at https://github.com/serre-lab/ees-inference.

2. Materials and Methods

2.1. Study subjects

All study procedures were conducted with the approval of the Brown University Institutional Animal Care and Use Committee (IACUC) and in accordance with the National Institutes of Health Guidelines for Animal Research (Guide for the Care and Use of Laboratory Animals). Four sheep (two female, two male) in the age range one to four years and weighing 83.43±7.58 kg were used for this study. Animals were kept in separate cages in a controlled environment on a 12-h light/dark cycle with ad libitum access to water, and were fed twice daily. Briefly, the sheep underwent laminectomy procedures to surgically implant two 24-contact EES electrode arrays (Micro-Leads Medical, Somerville, MA) placed onto the dorsal aspect of the epidural space of the lumbosacral spinal cord, spanning approximately the L4-L6 vertebral bodies [31]. Additionally, a ground wire and a reference wire were placed into the epidural space and paraspinal musculature, respectively. The sheep were allowed at least two weeks of recovery time prior to experimentation.

2.2. Data acquisition

Muscle activity was measured via surface electromyography (EMG) from four bilateral lower extremity muscles: biceps femoris (BF), gracilis (GR), peroneus longus (PL), and gastrocnemius (GA). The EMG sensors (Trigno Avanti model, Delsys Inc., Natick, MA) were placed longitudinally along the muscle belly following preparation of the attachment site via shaving and cleaning. These muscles were selected based on our previous experience collecting surface EMG in the sheep model [31]. BF is a proximal muscle responsible for hock flexion and hip extension, and GR is a proximal muscle responsible for hip abduction, hock flexion, and hock internal rotation. PL is a distal muscle responsible for hock flexion and GA is a distal muscle responsible for hock extension. During the experimental procedure, the animals were conscious and they were raised in a support sling (Panepinto, Fort Collins, CO, USA) until clearance between the ground and their hooves was observed. Stimulation was delivered via a Grapevine Macro+Stim or Nano2+Stim front end (Ripple Neuro, Salt Lake City, UT), and the accompanying Trellis software package. Custom code was written in MATLAB R2020a (The MathWorks Inc., Natick, MA, USA) to deliver a randomized set of stimulation values that varied in both stimulation amplitude and stimulation frequency. 300 ms long trains of stimulation pulses were delivered every 1 second. The stimulation pulse waveform was a charge-balanced, cathode-leading, biphasic square wave with a 1:3 aspect ratio and 167 μs pulse width. Stimulation was delivered in a monopolar configuration, with 1 electrode serving as the cathode and a stainless steel ground electrode serving as the return path/anode. Each stimulation train had amplitude and pulse frequency drawn randomly from a set of predefined ranges (60 μA - 1500 μA amplitude, 10 – 100 Hz frequency). The maximum stimulation amplitude was determined prior to each recording by manually increasing amplitude until the sheep could not tolerate the stimulation as indicated by a visual negative response to stimulation, or the maximum set tolerance of 1500 μA was reached [31]. 18 stimulation electrodes were chosen across the electrode arrays to acquire EMG responses at various anatomical regions including multiple spinal segments along the rostral-caudal axis. Additionally the location was varied along the medial-lateral axis to target bilateral, left, and right muscle activation. EMG signals from each muscle were bandpass filtered (20 – 450 Hz, second order Butterworth) and digitized (16-bit resolution, 2148 Hz sampling rate) on the sensor before streaming wirelessly to a base station connected to a PC. The dataset used for the neural network analysis consisted of a mixture of sub- and supra-threshold EMG samples, here defined as an observable EMG signal following data processing. Representative bilateral EMG responses to EES are shown in Appendix A (Fig. 8). As shown in Fig. 8, application of disparate stimulation parameters can selectively activate different sets of hindlimb muscles, and similar sets of stimulation parameters can have significantly different effects when applied at different electrode contacts, indicating the complex interactions between EES and muscle recruitment. Additionally, synchronized video recordings of hind limb motion were obtained from three sRGB cameras (Blackfly S, Teledyne FLIR LLC, Wilsonville, OR). Hind limb kinematics were obtained offline from these videos using state-of-the-art computer vision techniques. For more details, we refer the reader to Appendix B.

2.3. Data handling

EMG preprocessing

EMG waveforms were epoched from −100 ms prior to stimulation to 300 ms after each stimulation train. Raw values were downsampled to 2000 Hz and digitally high-pass filtered with 100 Hz cutoff (two-pole acausal Butterworth filter) to reduce any noise outside the frequency spectrum of interest that may have been caused by issues such as movement artifacts. EMG envelopes were calculated by taking the absolute value of the raw waveforms and applying a 2 pole acausal Butterworth low-pass filter with a corner frequency of 40 Hz.

Outlier removal

EMG data recorded from awake and alert animals included incidental factors, such as movement artifacts (especially during unprompted voluntary movements by the animal), unrelated to the electrical stimulation. We computed an internal consistency metric to aid in the removal of gross outliers. EMG envelopes, grouped by trial condition, were epoched into 100 ms windows. We calculated the sample covariance for each window and compared the maximum observed Mahalanobis distance to the 99.999% confidence interval of the corresponding chi-squared distribution. Trials with at least one window violating this confidence interval were flagged as outliers and excluded from subsequent analysis. Generally, this heuristic was very strict and led to the rejection of 5% – 10% of the data (confirmed as anomalous upon visual inspection).

Unreliable samples removal

To promote the discovery of repeatable associations between EES parameters and EMG, we calculated a “reliability” score for each EMG response, quantifying the degree to which the signal remains stable with repeated presentation of identical stimulation parameters. Sources of trial-to-trial variability included sensor noise, animal motion, and biological stochasticity. For each stimulation configuration θ_j, we computed its reliability coefficient ${\xi }_{{\theta }_j}$ using a bootstrapping procedure as:

$${\xi }_{{\theta }_j}=\frac{1}{N_b}\sum _{b=1}^{N_b}\frac{{\sum }_i\left({\widehat{x}}_{b,i}-\overline{\widehat{x_b}}\right)\left({\widehat{y}}_{b,i}-\overline{\widehat{y_b}}\right)}{\sqrt{{\sum }_i{\left({\widehat{x}}_{b,i}-\overline{{\widehat{x}}_b}\right)}^2{\sum }_i{\left({\widehat{y}}_{b,i}-\overline{{\widehat{y}}_b}\right)}^2}},$$ (2)

where for configuration θ_j and bootstrap round b we split EMG responses from repeats of this stimulation into two random halves, whose mean responses are ${\widehat{x}}_b$ and ${\widehat{y}}_b$. We deemed trials for which ${\xi }_{{\theta }_j}\le {\xi }_{thresh}$ as unreliable and discarded them from subsequent analysis. In our experiments, we set N_b = 100 and ξ_thresh = 0.4, which was approximately the median data reliability measure. Furthermore, if the majority of the trials for a given EMG sensor were deemed unreliable or the EMG sensor became dislodged due to animal movement, that sensor was not used in subsequent analysis. If more than one EMG sensor on each leg was deemed unreliable, data from that experiment was not used for analysis.

Subthreshold EMG removal

We defined sub-threshold EMG samples as trials where the stimulation was insufficient to evoke a sensorimotor response discernible from background noise. Previous publications have simply indicated threshold values of 20 μV to determine what was considered a supra-threshold EMG response to spinal stimulation [32, 33]. We describe a more robust methodology that takes into account any noise in the signal that may rise above the simple thresholds used in previous studies. To quantify the sub-threshold samples, we followed a similar bootstrap procedure as described above, computed the mean covariance between the split-half EMG means ${\widehat{x}}_b$ and ${\widehat{y}}_b$, and scaled these values to the range [0, 1]. Trials with this covariance metric less than 0.1 were rejected as sub-threshold samples. The covariance threshold is a hyperparameter that was empirically determined and was motivated by our average-case reliability metric.

EMG summarization

Prior research, via a combination of computational modeling, patient studies and experimenter intuition, laid out guidelines for the effective summarization of raw temporal EMG waveforms into compact representations [5, 7]. We adhered to said guidelines for constructing lower-dimensional EMG representations as:

$$x_{r,k}=\mathrm{Clip}\left[\frac{1}{{\eta }_q(k)}\underset{z_{r,k}}{\underbrace{{\int }_0^T\left|{\tilde{x}}_{r,k}(t)\right|dt}},\{0,1\}\right],$$ (3)

Here, ${\tilde{x}}_{r,k}(t)$ denotes the raw EMG waveform obtained during trial r from EMG contact k. η_q(k) denotes the q-th percentile activity value of z_r,k, over all trials r ∈ {1..R}. |.| is the absolute value function, and Clip[x, {a, b}] simply denotes the hard clip operation confining the range of x to [a, b]. Unless specified otherwise, we set q = 95 and T = 300ms for our analyses. The EMG summary for trial r is thus x_r = [x_r,1, x_r,2, …, x_r,K], where K is the total number of used EMG contacts.

EES parameterization

In our setup, we configured EES for each trial using three degrees of freedom. For the neural network analyses, we represented EES for trial r as the vector θ_r = [μ_r, a_r, E_r,e], where μ_r and a_r denote the frequency and amplitude, respectively, of the pulse train used for trial r. E_r,e is a binary vector with a value of 1 at index e (the stimulated electrode contact during trial r) and 0 elsewhere. Frequency and amplitude values were normalized to the range [0, 1), for neural network training purposes.

2.4. Forward model (f_ϕ) specification

In this study, we used a densely connected, feedforward deep neural network model, also known as a Multi-Layer Perceptron (MLP), as our forward model (Fig. 3a). We trained this model end-to-end using Stochastic Gradient Descent, implemented via the Adam algorithm [34], to learn a mapping between EES parameters and corresponding evoked EMG. Specifically, we factorized our forward model into embedding, core, and readout modules. The embedding module received as input the parameterized EES, and projected this into a feature space that is fed to the core module. The core module was responsible for transforming the embedded features such that the readout module is able to use these to predict the summarized EMG response. We adopted this factorization to explicitly provide a framework for fairly evaluating diverse model architectures, in place of the core module, in the future. In our experiments, each module consisted of a linear layer followed by a non-linear rectifying (ReLU) activation function. Our embedding and core feature spaces were 16 and 256-dimensional respectively. We trained the model by minimizing the L1 loss between the predicted and groundtruth EMG summaries:

$$ℒ(\mathbf{y},\mathbf{x})=\frac{1}{N}\sum _{r=1}^N\frac{1}{K}\sum _{k=1}^K\left|y_{r,k}-x_{r,k}\right|$$ (4)

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.

In Eq. 4, y denotes the predicted EMG summary while N, K refer to the total number of training samples and number of EMG contacts used respectively. Model parameters were updated using the Adam optimization algorithm [35] with an initial learning rate lr = 1e − 2 for 1000 epochs. Additionally, we used a L2 weight decay of 5e − 4 and 50% variational dropout before the readout module to regularize the network and prevent over-fitting. Additional details regarding the train vs. test data splits are specified in Appendix C.

2.5. Inference model $\left(\left\{f_{{\psi }_{(i)}}\right\}\right)$ specification

To tackle the inverse problem (Fig. 1c), we adopted and augmented the Sequential Neural Posterior Estimation algorithm for our use case of controlling sensorimotor activation via EES [24]. We outline our procedure in Algorithm 1. Our primary contribution was in the reformulation of the typical approximate Bayesian inference setting to instead use amortized neural networks as “simulators”, and estimate posterior densities over their input space. Additionally, to circumvent the problem of dealing with discrete vectors in the EES parameter space, we instead trained electrode-conditional models. For training an electrode-conditional inverse model, we treated EES parameters μ and a as random variables, and clamp E to reflect the binary vector E_i with a value of 1 at index i (the conditioned electrode contact) and 0 elsewhere. We trained all such electrode-conditioned inverse models in parallel, and thus at no additional time cost.

For each electrode-conditioned inverse model, we started by defining a Uniform prior over the range of frequencies and amplitudes, and sampled a number of parameter sets from this prior. We treated our amortized forward model f_ϕ as a “simulator”, and used it to construct a simulated dataset {(θ_j, x_j)} from these sampled parameter sets. The benefits of using f_ϕ as a simulator, as opposed to mechanistic models in typical simulation-based inference settings, were two-fold. First, simulations using f_ϕ only involved forward passes through a neural network, which occurred in a massively parallel manner and was rapid. In contrast, simulations using mechanistic models are typically orders of magnitudes slower. Second, the “learnability” of f_ϕ imparted flexibility both in the fundamental nature of the mapping being approximated, as well as input/output parameterizations.

Algorithm 1.

Controlling EMG activity by identifying optimal EES protocols

Require: π(θ)	▷ prior over EES parameters
Require: fϕ : Concat([θ, E_]) → x	▷ trained forward model mapping EES to EMG
Require: x*	▷ target EMG
Require: Ne	▷ number of EES electrode contacts
for i ϵ {0..Ne} in parallel do
r ← 1
pi(θ) := π(θ)	▷ initialize target density with prior
while r ≤ 2 do
${\widehat{\mathrm{\Theta }}}_{1\dots K}\sim p_i(\theta )$	▷ sample a batch of parameters
${\mathbf{x}}_{1\dots K}\leftarrow f_{\varphi }( \text{Concat} [{\widehat{\text{Θ}}}_{1\dots K},{\mathbf{E}}_i])$	▷ simulate via an electrode-conditioned forward pass
${\psi }_{(i)}\leftarrow {\psi }_{(i)}-\lambda \frac{\partial ℒ\left({\psi }_{(i)},{\mathbf{x}}_{1\dots K}\right)}{\partial {\psi }_{(i)}}$	▷ gradient descent
$p_i(\theta ):=q_{f_{{\psi }_{(i)}}}{\left(x^{*}\right)}^{(\theta )}$	▷ update target density
r ← r + 1
${\tilde{p}}_i\left(\theta \mid x^{*}\right)\leftarrow p_i(\theta )$	▷ get the final posterior
$\mathbf{\Theta }:=\{(\mathrm{Concat}([{\theta }_i^{(j)},{\mathbf{E}}_i]),{\tilde{p}}_i({\theta }_i^{(j)}))\mid {\theta }_i^{(j)}\sim {\tilde{p}}_i(\theta ),\forall i\in \left\{\mathrm{0..}{N}_e\right\}\}$	▷ sample EES
${\theta }^{*}=\arg {\min }_{\left(\theta ,q\left(\theta \mid x^{*}\right)\right)\in \mathbf{\Theta }}‖f_{\varphi }(\theta )-x^{*}‖-q\left(\theta \mid x^{*}\right)$	▷ determine optimal EES

We instantiated each inverse model $f_{{\psi }_{(i)}}$ as a Masked Autoregressive Flow (MAF) [36]. $f_{{\psi }_{(i)}}$ comprised multiple stacks of invertible transformations and could be interpreted as a deep normalizing flow. The overall objective of the inverse model $f_{{\psi }_{(i)}}$ was to learn to approximate the posterior density p(Concat[μ, a, E_i] | x). Training proceeded by using the simulated data (EMG) {x_j} (as context) and parameters {θ_j} as inputs to $f_{{\psi }_{(i)}}$, which transformed random values from a preset, standard distribution (such as a multivariate normal) to our target distribution. Network parameters ψ_(i) were optimized by minimizing the negative log probability loss (Eq. 5).

$$ℒ\left({\psi }_{(i)}\right)=-\sum _{j=1}^{N_s}\log \left(q_{f_{{\psi }_{(i)}}}\left({\text{x}}_j\right)\left({\theta }_j\right)\right),$$ (5)

where q(.) is an estimate of the posterior density, whose parameters are specified by the output of $f_{{\psi }_{(i)}}(.)$. To refine this posterior to a specific target EMG activity pattern x₀, we repeated the aforementioned process albeit with minor modifications to initially sample parameters from the amortized posterior (instead of the prior) and modifying the loss function to reflect this change [30]. EES samples drawn from the refined posterior were then used in vivo to observe their effect on sensorimotor network activation. Further details regarding hyperparameter choices in this process are described in Appendix D.

3. Results

3.1. Neural networks precisely capture the function of spinal sensorimotor computations

Converging evidence for low-dimensional neural dynamics governing the non-linear transformation of sensory stimuli to motor outputs indicates that artificial neural networks might be well suited to approximate spinal sensorimotor computations. We adopt a class of densely connected feedforward DNNs to learn a mapping from EES inputs to motor outputs. Additionally, by efficiently parameterizing EES settings and summarizing EMG data, we further tailor the problem formulation to maximally benefit from DNNs’ function approximation capabilities.

To demonstrate the utility of our forward model f_ϕ (Section. 2.4), we train it on data collected from four different sheep. In each animal, we recorded EMG from four bilateral lower extremity muscles, and discard signal from noisy contacts. As per standard machine learning practice, we hold out 10% of each dataset as a dedicated test set. In particular, the test set exclusively consisted of trials with at least one of the EES parameters, frequency (μ) or amplitude (a), unseen during training. Setting up the test samples this way allowed us to evaluate the true degree to which f_ϕ learns non-linear sensorimotor computations, since this would involve both soft (interpolate along a single EES parameter) and hard (interpolate over multiple EES parameters) generalization (Appendix C).

We first evaluate the goodness of fit of f_ϕ in silico by computing the per-EMG contact L1 metric between the network predictions and the observed EMG activity for various EES inputs in our test data. Numerical results for the four animals are reported in Fig. 3b and Table. 1. An L1 value of 0 signifies “perfect” model performance, defined as capturing 100% of the explainable variance of evoked motor outputs for any given EES input. Additionally, to contextualize the performance of f_ϕ, we compute a “random” baseline prediction error (red crosses; Fig. 3b). The distribution of model prediction errors for each animal differed significantly from the baseline predictions with our model prediction errors being significantly lesser (Sheep A (Mann–Whitney U = 4.0, n₁ = n₂ = 52, p < .001); Sheep B (Mann–Whitney U = 0.0, n₁ = n₂ = 26, p < .001); Sheep C (Mann–Whitney U = 0.0, n₁ = n₂ = 136, p < .001); Sheep D (Mann–Whitney U = 1.0, n₁ = n₂ = 144, p < .001)). We report the average training time of f_ϕ to be 33.25 ± 16.9s on a standard 16-core Intel Xeon Gold 6242 CPU workstation.

Table 1: Numerical results accompanying Fig. 3c.

We report the L1 errors per-EMG contact, on a dedicated held-out test dataset, for models trained on each of the four sheep. Errors are reported as mean ± standard deviation (in arbitrary units as given by the normalized EMG scale).

EMG Contact	Sheep A	Sheep B	Sheep C	Sheep D
LPeroneus	0.041 ± 0.047	0.007 ± 0.009	0.012 ± 0.024	0.029 ± 0.06
LGastrocnemius	0.046 ± 0.059	0.016 ± 0.02	0.01 ± 0.023	0.041 ± 0.04
LBicepsFemoris	-	-	0.011 ± 0.017	-
LGracilis	0.058 ± 0.071	0.027 ± 0.029	0.023 ± 0.043	0.04 ± 0.068
RPeroneus	0.04 ± 0.044	0.03 ± 0.033	0.019 ± 0.043	0.028 ± 0.046
RGastrocnemius	0.037 ± 0.035	0.032 ± 0.048	0.024 ± 0.047	0.045 ± 0.046
RBicepsFemoris	-	-	0.03 ± 0.052	0.013 ± 0.036
RGracilis	0.043 ± 0.046	0.019 ± 0.018	0.038 ± 0.054	0.029 ± 0.07

As a precursor to performing complete probabilistic inference in vivo, we test the ability of the correspondingly trained f_ϕ to “recover” EES parameters used in a trial of interest using Maximum Likelihood Estimation (MLE). We define “recovery” here as the ability to identify the value of a specific EES parameter θ_j, which when used along with other given EES parameters θ_−j makes the observed EMG response x most probable. The results of our evaluation for θ_j ∈ {μ, a} are presented in Fig. 3c. We found that the MLE estimates of parameters μ and a were significantly correlated with the respective ground truth values on the held-out test trials. Specifically, we computed the Pearson correlation coefficient, r, between the MLE estimates and ground truth values and test for the slope of the regression line being significantly different from 0 using the Wald test with t-distribution of the test statistic (Sheep A (μ(r = .974, p < .001); a(r = 0.968, p < .001)), Sheep B (μ(r = .958, p < .001); a(r = 0.965, p < .001)), Sheep C (μ(r = .601, p < .001); a(r = 0.969, p < .001)), Sheep D (μ(r = .767, p < .001); a(r = 0.645, p < .001))). We thus verify that f_ϕ is indeed a good functional approximation of spinal sensorimotor computations.

3.2. Neural network guided EES optimally shapes muscle recruitment

Leveraging the successful amortization of spinal sensorimotor computations by our forward model, we trained and deployed our inverse models $\left\{f_{{\psi }_i}\right\}$ in vivo to drive motor outputs into targeted regimes. Figure 4 provides an overview of the process. Our electrode-conditioned inverse model uses the forward model f_ϕ, and a desired, novel target EMG pattern to estimate the posterior probability density over the full space of EES parameters (Algorithm 1). We sampled from this posterior distribution to identify candidate EES parameters, which are then tested in silico first then in vivo. The whole process was fully automated. To quantify the success of our approach we measured how well the evoked EMG in vivo aligned with the “target” EMG pattern, as well as how well it aligned with the forward model output.

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.

We performed in vivo testing on Sheep C and D. Sheep A damaged their electrode midway into the study. Sheep B had to be removed from the study due to non-study-related health issues, which prevented awake behaving tests. In each sheep, we apply Algorithm 1 for N_T novel EMG targets $\{{\mathbf{x}}_{n=1\dots N_T}^{*}\}$. By “novel”, we allude to the fact that neither our forward models nor inverse models have seen this pattern of EMG during their respective training phases. Our procedure returns N_P distinct EES proposals for each EMG target, which are then applied for N_R repeats. To obtain the EES proposals we first sample EES parameters from the respective electrode-conditioned posterior densities and then jointly maximize the likelihood $p\left({\theta }_n\mid x={\mathbf{x}}_n^{*}\right)$ and minimize the in silico loss $\Vert f_{\varphi }\left({\theta }_n\right)-{\mathbf{x}}_n^{*}\Vert$. In this set of experiments we set N_T = N_P = 5 and N_R = 10. The specificity of evoked EMG adjudicated the optimality of sampled EES in our context. To confirm that EMG responses are not invariant to EES parameters, we also tested “non-optimal” parameters in silico to verify that the resulting EMG patterns differ significantly from the target (Fig. 5d). Non-optimal EES were picked as having low posterior density. A complete numerical analysis of the neural network guided in vivo experiments is presented in Fig. 6. The efficacy of our model-guided EES trials parallels the achievable upper-bound as specified by the inter-trial reliability of EMG responses. A nonparameteric Mann-Whitney U test confirmed that the distribution underlying the ideal upper-bound correlations was statistically identical to the distribution underlying the correlations between targets and our model-guided EMG activities for both Sheep C (Mann–Whitney U = 635.0, n₁ = n₂ = 39, p = .105) and Sheep D (Mann–Whitney U = 1650.0, n₁ = n₂ = 60, p = .216). Over a randomly selected set of n = 50 EMG targets, we report the training time for our electrode-conditioned inverse models $t_{train}^{(\psi )}\in [42.1,600.2]s$, with a median of 69.8s. On this same set, EES proposal sampling times $t_{sample}^{(\psi )}\in [0.28,216.42]s$ with a median of 1.53s.

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.

As we note earlier in the manuscript, the transformation of EES to EMG is a complex, non-linear process which moreover can be many-to-one (a single EMG pattern can result from many EES parameters). Lending credence to this argument, we are able to automatically identify a variety of EES that evoke similar sensorimotor responses in vivo (Fig. 5a–c). Data presented here reveal that stimulating different epidural electrode contacts with appropriately chosen pulse trains can be functionally equivalent, despite these EES configurations being spatially distinct. In this spatial EES search, we also demonstrate that our system is well calibrated and correctly assigns low posterior probabilities in cases where there exists no possible way to achieve the target EMG by stimulating a specific electrode contact (Fig. 5d).

The end-to-end time for executing the training of our inverse model, sampling EES parameters for given EMG targets, and testing these EES samples in vivo was ~ 20 min per animal. The target EMG patterns used were arbitary, and novel. Despite no explicit knowledge of either spinal physiology nor specifics of the electrode anatomy made available to the system, our algorithm was functionally well-calibrated and effectively controlled spinal sensorimotor network activation.

3.3. Non-identifiable EES protocols pave the way to study parameter trade-offs

Both the Maximum Likelihood Estimation analysis (Fig. 3c) and the spatial EES search (Fig. 5) indicate the non-identifiability of the mapping from EES parameters to EMG activation. By non-identifiability, we refer to the scenario of two or more unique EES parameter configurations producing observationally equivalent EMG outputs. The presence of complex posterior geometries (for example, Fig. 5b) suggest that spinal sensorimotor circuits may embody this inherent redundancy as well. Distinct sensorimotor pathways may potentially compensate (or trade-off) for the reduction or gain in individual EES parameter values while conserving the overall function of the system. A systematic way to understand these trade-offs would help in identifying their neural substrates. Though so far we have seen the non-identifiability of EES in a varying electrode-context, here we present evidence for parameter trade-offs within the same electrode (Fig. 7).

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.

We compute level sets, in the posterior, defined as contours along which the posterior probabilities remain constant. We draw EES samples that lie along a contour and use the forward model f_ϕ to predict sensorimotor network activation in silico. EES samples on a level set in the posterior evoke near-identical sensorimotor responses thus confirming our parameter trade-off hypothesis. Fig. 7 presents a visual example wherein unique EES parameters are shown to elicit identical motor outputs. Prior work from other groups on computational models of spinal circuitry have helped delineate mechanisms of endogenous interactions between neurons from the effect of exogenous inputs such as EES [27, 37]. Such models also aid in identifying circuit modulations, either via altering spinal connections or EES parameters, that change motor outputs. However, these models were explicitly constructed to include these model features. For the first time, we show that facets of neural interactions emerge naturally from optimizing for functional outcomes. The present study lays the groundwork for a powerful framework that allows for unbiased and systematic theoretical investigation of spinal neuromodulation through in silico electrophysiology.

4. Discussion

In this study, we demonstrate the first fully-automated, data-driven, and subject-agnostic method to optimally control spinal sensorimotor network activity via EES. Herein, we present data collected from four sheep, each implanted with two 24-contact EES electrode arrays on the lumbosacral spinal cord. We first showcase the ability to functionally approximate spinal sensorimotor computations with a feedforward deep neural network. Our forward model approximation of spinal sensorimotor computations was precise, and was able to accurately predict EMG outputs based on EES parameters (Fig. 3). After this amortization step, we leveraged our inverse model family and learn to control motor outputs from the forward model by optimally manipulating EES parameters. Our method identified, within minutes of computational effort on typical hardware, novel EES parameters that produced desired target EMG recruitment in vivo (Figs. 4,5). We numerically verified the efficacy of these EES proposals by computing the correlation between the evoked and target EMG activity which approached the performance ceiling (Fig. 6). In addition to being fully automated, our approach is also orders of magnitude faster than manual optimization of EES protocols per-subject.

The ability to control neural population activity has been of long-standing interest to the broader neuroscience community [38–40]. Several model-based and model-free solutions to this problem have been proposed, each having their own merits and demerits. On the one hand, model-free approaches do not rely on explicit neural circuit-specific knowledge and are fast. However, they do not support drawing inferences across subjects. On the other hand, model-based approaches rely on domain knowledge and large amounts of training data but can generate testable hypotheses and be used for system identification. Specifically in the context of controlling motor neuron activity, multiple efforts have demonstrated that neuromodulation of the spinal cord with EES can restore voluntary motor function, including standing and stepping in individuals diagnosed with severe chronic incomplete SCI [4–12]. The most successful of approaches thus far have relied on engaging proprioceptive feedback circuits to activate central pattern-generating networks, which in-turn allowed proprioceptive signals to drive motor control [13, 27]. These feedback circuit models were also extended to methodologies for optimally selecting electrode locations for recruiting specific subsets of sensory nerve fibers to restore locomotion in rodents [41], non-human primates [21], and human patients [7] with SCI.

These systems, however, suffer from three primary drawbacks: (1) The EES stimulation protocols themselves were empirically determined via systematic search; (2) Only highly prototypical movement patterns were targeted; and (3) They relied on complex peri-surgical imaging, testing, and device personalization. Overall, this process was laborious and time-consuming with widely varying configurations between subjects due to differences in the pathophysiology of their SCI and differences in spinal cord anatomy and electrode placement. The detailed biophysical models of the proprioceptive feedback circuits, accompanying these systems, formalized the modulatory effect of exogenous EES. Model parameters were either tuned manually or selected by performing a brute-force grid search, a process both computationally expensive and limiting the capacity for generalization.

Our system can be viewed as a hybrid between model-based and model-free approaches. By defining our forward model as a universal function approximator and reformulating the problem to utilize low-dimensional output summary statistics and parameter representations, we can “learn” directly from relatively few samples without specific knowledge about the underlying neural circuitry. We also retain the ability to amortize information across subjects, which, given recent developments in the transfer learning literature can potentially help in the reduction of data budgets and thereby experimental time. Moreover, our models are parameter-efficient and the end-to-end run time of the system makes it practically feasible and convenient. The notion of using deep neural networks to amortize data generation processes has been explored in the past [25]. However, this was only in the context of synthetic data generated from simple stochastic models. In fact, machine learning approaches for surrogate modeling is now emerging as a research frontier in and of itself [42, 43]. We present the first evidence that this framework can be adopted for sophisticated experimental data for which mechanistic explanations are lacking or infeasible to generate.

In a broader context, the development of deep neural network solutions within the space of neuromodulation holds great promise for the future. Intelligent control of neuroprostheses could enable more seamless interaction between the user and the neural interface, and lead to more efficient movement after neural injury [44, 45]. Prior work has utilized Bayesian optimization of peripheral stimulation to evoke lower extremity movements [46], and machine learning has been used to predict EES outcomes [47, 48]. However, our system presents the first demonstration of the applicability of deep learning systems for spinal neuromodulation in vivo. From a practical perspective, the data-driven approach employed here only cost minimal experimental time for obtaining EES trials used for training the neural networks. Post-amortization of the mapping between EES and EMG outputs, we were able to successfully reuse these networks to derive stimulation proposals for any number of EMG targets at no additional cost. Our modeling approach did not require any specific knowledge about the anatomy of the spinal tracts and hence did not necessitate any intra-operative imaging. From a theoretical perspective, the degeneracy of the mapping between EMG outputs and EES highlighted by our system opens up the space for addressing several questions in the future. Biological principles dictate that, for effective activation of motor neurons, EES targets large-diameter afferent fibres through the dorsal roots at motor threshold [18]. Without any explicit knowledge about the location of such dorsal roots, we were able to identify multiple EES configurations that produced desired outputs. These findings suggest the existence of redundant mechanisms of complex motor output production implemented by spinal circuitry that can be probed systematically with our system in-the-loop.

From an end user’s perspective, our framework takes as input any arbitrary pattern of muscle activations and outputs relevant EES parameter choices. Since our models were trained with EMG responses that effectively captured both motor activation via transynaptic activation of sensory neurons as well as direct activation of motor neurons at high stimulation amplitudes [13, 49], target EMG patterns can be chosen to reflect sensorimotor computations of neuromuscular excitation and inhibition that may be necessary for functional behavior. Additionally, previous clinical EES studies directly translated EMG data recorded at rest without ascending or descending input to use during spatiotemporal EES where stimulation is applied at specific time-locked portions of complex tasks such as locomotion [7, 18]. Therefore, the models described here could potentially be applicable to re-enabling complex motor function in individuals with spinal cord injury.

Although the work described here provides an algorithmic framework for quickly determining effective stimulation parameters, there are limitations in this study that should be noted. Long-term use of EES in clinical populations has resulted in improved motor performance [5–7], and in select cases has led to study participants regaining movement without stimulation [50, 51]. Currently our algorithm uses static models trained once and does not account for changes with long-term use of EES in a clinical population. Development of continual learning algorithms that are capable of adapting their behavior over time are currently of great interest to the machine learning community [52]. Therefore, future work in this field should focus on the continued development of these algorithms to non-stationary conditions and task improvements over time. Additionally, our stimulation was applied while the animal was in an unloaded state. Previous literature has indicated the importance of somatosensory input from the periphery during gait and standing activities that increases biological stochasticity and is necessary to provide proprioceptive information regarding limb position during active tasks [49, 53]. Future experimentation using this model should account for changes due to increased somatosensory input, and incorporate machine learning algorithms to precisely control stimulation timing in complex tasks such as locomotion as has been performed previously using intraspinal microstimulation (ISMS) experiments in cats [54, 55]. Furthermore, although our algorithm produced comparable results across all animals included within this study, the methodology requires retraining to produce similar results across a wide range of subjects. Differences in spinal anatomy, electrode design, electrode positioning, and the presence of an SCI may alter the EES to EMG mappings in new subjects or implantations, and further work in a larger population is necessary to account for these differences. Finally, we restricted our inferred parameters to electrode location, stimulation amplitude, and stimulation frequency, while leaving pulse width constant. Computational studies have shown that pulse width can change the recruitment of neural structures during SCS [56]. Therefore, future work should include this as another variable that can be inferred using machine learning approaches to the neuromodulation field.

Spinal neuromodulation is an emerging field, and the work described here is a necessary addition to aid in the development and discovery of EES protocols in a fast, data-efficient, and subject-agnostic manner. Rapid discovery of these protocols for neuromodulation devices currently represents one of the largest barriers to the clinical translation of EES. The development of artificially intelligent tools, such as the one presented in this manuscript, will aid in scaling up EES to the clinical space and realize meaningful benefits for spinal rehabilitation in the future.

A. Representative bilateral EMG activity upon EES

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.

B. Kinematics acquisition and analysis

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.

We performed simultaneous sRGB video recording from three cameras placed at a distance of 150cm from the subject as shown in Fig. 9. The left and right cameras were placed orthogonal to each other, with the rear camera placed directly behind the subject. Camera placement was dictated by three factors: (1) The coverage of the full range of the subject’s hind limb motion; (2) The ability to disambiguate joint locations during self-occlusions; and (3) The overlap in the field of views of the three cameras. We used DeepLabcut, a markerless pose estimation toolkit, to extract hind limb kinematics [57]. Specifically, we used a standard 50 layer Residual Network (ResNet-50) as the convolutional backbone for DeepLabCut [58]. The ResNet-50 layers were initialized with ImageNet pre-trained weights, and fine-tuned for our purposes. To opportunistically maximize the precision of joint detection, we trained a separate model for each viewpoint. We randomly selected 50 ideal frames per-view, and manually annotate three key points (Hip, Hock, and Hoof) per limb. Ideal frames are defined as those with no self-occlusions.

B.1. Functionally correlating EMG and Kinematics

All the analysis presented in the manuscript relied on motor output as measured by Electromyography (EMG). Here, we show that summarized EMG activity is also highly correlated with the physical motion of the ipsilateral limb as measured by kinematics. The intent of this analysis was to validate that the EMG responses we obtained during testing also resulted in kinematic outcomes that would be useful in locomotion. To so do, we perform midline EES (3 electrode contacts, 588 trials) with frequencies μ ∈ {10, 25, 50, 100}Hz and amplitudes a ∈ [60, 1500]μA. Kinematics during these trials, computed using the computer vision system described in B, was used to compute the EES-evoked angular displacement of the hind limbs. Limb displacement was calculated as the change in hip angle in the sagittal plane from baseline as a result of stimulation. Movements caused by stimulation were primarily hip flexion/knee extension. Figure 10 presents the correlation analysis between each summarized EMG channel activation and the angular displacement of the corresponding hind limb. Angular displacement of the hip was chosen to represent the gross overall movement of the hindlimb during application of EES. As shown in Figure 10, all muscles demonstrate relatively high levels of correlation between EMG activity and kinematic movement. Although activation of each muscle has different biomechanical effects, the correlations between angular displacement of the hip and each recorded EMG were simply used to demonstrate that the mappings between EES and EMG could be useful in predicting kinematic outputs in addition to EMG.

C. Forward model specifics

As per standard machine learning practice, we partitioned EES trials from each dataset into dedicated train and test splits. To do so, we followed a scheme as presented in Fig. 11. We observed that the EMG responses to systematically varying EES are highly non-linear (Fig. 2c–d), and thus wanted to test the generalization ability of our forward model f_ϕ. Specifically, we wanted to test two types of generalization: (1) Soft generalization: the model’s ability to interpolate along a single unseen parameter dimension, and (2) Hard generalization: the model’s ability to interpolate within a completely novel parameter regime. f_ϕ is trained on trials that used EES parameters from the unshaded regions (Fig. 11). The goodness of fit of f_ϕ is computed by evaluating the network on trials with EES parameters from the shaded regions. Note that these are completely non-overlapping. The region enclosed by the blue boundary tests the “soft” generalization capacity of f_ϕ, while the region enclosed by the red boundary tests “hard” generalization.

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.

D. Inverse model hyperparameter specifications

Hyperparameter	Description	Value
num_rounds	The total number of rounds of refinement from a globally amortized posterior to a datapoint-specific posterior	2
num_simulations	The total number of EES trials to be “simulated” (and used) per round of inverse neural network model training	1024
simulation_batch_size	Batch size for simulating data from the amortized forward neural network model	1000
training_batch_size	The training batch size of the inverse neural network models	50
num_samples	The number of EES parameter sets to be drawn from the approximated posterior density to construct the pairwise density plots	10000
filteringx_ratio	The fraction of EES proposals (sampled from the posterior distribution) selected for filtering based on target correlation after sorting all EES proposals based on their log probability	0.1
num_proposals	The number of unique EES parameters (per EMG target) to be selected from the posterior samples for in vivo testing	5
timeout	Upper bound on the time to convergence (in seconds) to wait for the training of the inverse models	600