A GPU-native, differentiable fiber simulator reveals a species-dependent selectivity penalty in reduced-order peripheral-nerve models

Technical Summary: A GPU-native, Differentiable Fiber Simulator Reveals a Species-Dependent Selectivity Penalty

Problem Statement
Computational models are essential for designing selective peripheral nerve stimulation therapies. However, to remain computationally tractable, standard workflows universally reduce complex nerve fascicles to a single or few representative fibers (often a centroid). While this reduction enables optimization, its impact on predicted selectivity has remained untested because validating it requires simulating the entire fiber population under optimized extracellular fields—a task previously too computationally expensive. Furthermore, existing optimization pipelines rely on CPU-based, gradient-free methods (e.g., grid sweeps, evolutionary search) that scale poorly with high-dimensional design spaces (e.g., multi-contact cuffs, waveform shaping).

Methodology
The authors developed jaxon, a fully differentiable, GPU-native reimplementation of canonical peripheral nerve fiber models within the JAX framework.

• Models Implemented: jaxon implements four standard axon models: the myelinated McIntyre–Richardson–Grill (MRG) double-cable and Sweeney models, and the unmyelinated Sundt and Rattay C-fiber models.
• Differentiability & Speed: By leveraging JAX's reverse-mode automatic differentiation and GPU vectorization (vmap), jaxon computes gradients of arbitrary functions of membrane trajectories in a single forward pass. This enables gradient-based inference at scale.
• Validation: The simulator was validated against NEURON (wrapped in PyFibers) across 943 configurations involving various fiber diameters, pulse widths, and waveforms (monophasic, biphasic, sinusoidal, etc.).
• Optimization Pipeline: The authors applied jaxon to a cohort of 29 histology-segmented vagus nerves (18 swine, 11 human). They optimized stimulation patterns on a 12-contact cuff to maximize selectivity for specific target fascicles.
• Comparative Analysis: To test the "reduced-order" assumption, the authors optimized stimulation parameters using sparse sampling (1, 3, or 10 fibers per fascicle, including the standard single-centroid reduction) and then re-scored these optimized designs on the full fiber population (~1,000 fibers per nerve) to measure the "deployment penalty" (the drop in selectivity when moving from the reduced model to the full population).

Key Results

1. Accuracy and Throughput: jaxon reproduces NEURON with high fidelity, with 99.6% of 943 configurations falling within a 1% relative threshold error. In terms of speed, jaxon is approximately 820× faster than PyFibers/NEURON when simulating 100,000 fibers on a single NVIDIA A100 GPU (peaking at ~1,377× for the Sweeney model). This speedup makes full-population simulation feasible within minutes rather than hours.
2. Species-Dependent Selectivity Penalty: The study reveals that the standard single-fiber reduction systematically overestimates deliverable selectivity, but the magnitude of this error is highly species-dependent:
   • Swine: The penalty is negligible (median selectivity index drop of 0.01).
   • Human: The penalty is substantial (median drop of 0.14, reaching up to 0.45 in some cases).
3. Mechanism of Error: The penalty arises because human fascicles are significantly larger (median 151 fibers per target fascicle) compared to swine (median 30 fibers). A field optimized for a single centroid fiber fails to recruit the entire large human fascicle, leading to under-recruitment of the target. In contrast, the smaller swine fascicles are well-represented by a single sample.
4. Recoverability: Because jaxon allows for cheap re-scoring, the authors demonstrated that re-evaluating sparse-optimized designs on the full population recovers most of the lost selectivity. The residual penalty after "optimize-then-validate" is small (median ~0.02 in humans), suggesting the issue is a workflow limitation (lack of population-scale validation) rather than an insurmountable physical barrier.

Key Contributions

• jaxon Simulator: The first fully differentiable, GPU-native implementation of peripheral nerve fiber models (MRG, Sweeney, Sundt, Rattay) that supports batched simulation of large populations and end-to-end gradient computation.
• Quantification of Sampling Bias: The first rigorous quantification of the "reduced-order" sampling assumption in peripheral nerve modeling. The paper demonstrates that this assumption introduces a significant, species-dependent bias in human models that is undetectable without full-population simulation.
• Workflow Recommendation: The identification of "population-scale validation" as a critical, previously missing step in selectivity optimization pipelines. The authors show that while sparse optimization is fast, designs must be re-scored on the full population to ensure clinical relevance, particularly for human anatomy.

Significance and Claims
The paper claims that the reduced-order fiber sampling standard in GPU-accelerated selectivity modeling (e.g., reducing fascicles to a single centroid) overestimates the actual selectivity deliverable to the nerve, specifically in human anatomies with large fascicles. This "deployment penalty" is a translational bias that reduced-order models cannot detect because they never simulate the fibers on which the bias acts.

The authors emphasize that this finding is a statement about the modeling workflow, not a clinical prescription. They argue that the computational cost of full-population simulation, previously prohibitive, is now tractable thanks to tools like jaxon. Consequently, the actionable conclusion is that selectivity designs tuned on centroid fibers must be re-scored on the full fiber population before being trusted, a step that recovers the majority of the lost selectivity. The paper positions jaxon as a complementary tool to existing ecosystems (like PyFibers or AxonML), offering exact biophysical fidelity and end-to-end gradients without the overhead of Python-NEURON wrappers or the approximation errors of learned surrogates.