💻 bioinformatics

Surrogate Gradients for Gradient-Based Parameter Estimation in Simplified Neuron Models

This paper introduces a differentiable Adaptive Exponential Integrate-and-Fire model using surrogate gradients within the Jaxley framework, demonstrating that while this approach enables efficient gradient-based parameter estimation for subthreshold dynamics, it currently fails to outperform gradient-free methods in recovering full spike trains due to the non-differentiable nature of the spike mechanism.

Original authors: Mayer, P., Kozlov, A.

Published 2026-09-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mayer, P., Kozlov, A.

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

The brain is a vast, intricate network of billions of cells, each firing electrical signals to carry thoughts, sensations, and commands. To understand how this machinery works, scientists build computer models of these cells. For decades, the most accurate models were like detailed blueprints, simulating every tiny chemical channel and physical structure inside a neuron. While precise, these models are so computationally heavy that simulating even a small piece of brain tissue takes immense time and power. To study the brain as a whole, researchers turned to simplified models. These are like rough sketches that capture the essential firing behavior of a neuron without the heavy baggage of every microscopic detail. They are fast and efficient, making them the only practical way to simulate large networks. However, there is a catch: these simplified models rely on a sudden, sharp jump when a neuron fires, a mechanism that breaks the smooth mathematical rules computers use to learn and improve. This has forced scientists to use slow, trial-and-error methods to tune these models to match real brain data, rather than the faster, more direct methods used in modern artificial intelligence.

A team of researchers at KTH Royal Institute of Technology in Stockholm set out to see if they could bridge this gap. They wanted to know if they could teach these simplified models to fit real brain data using the same fast, gradient-based methods that power modern machine learning. To do this, they built a new version of a popular simplified model called the Adaptive Exponential Integrate-and-Fire model. They implemented it in a powerful, modern software framework designed for high-speed computing. Crucially, they added a mathematical trick known as a surrogate gradient. In simple terms, this trick allows the computer to pretend the sudden jump of a neuron firing is smooth and continuous just for the purpose of calculating how to improve the model, even though the actual simulation still fires in sharp, discrete jumps. This allowed them to run the model on powerful graphics processors, making it hundreds of times faster than the standard software used by neuroscientists.

With this new tool in hand, the researchers ran a massive series of tests to see if the fast, gradient-based method could actually find the correct settings for the neuron model. They created thousands of synthetic brain recordings and asked the computer to adjust the model's parameters to match them perfectly. They compared their fast method against the traditional, slower trial-and-error approach. The results were surprising and clear. While the fast method worked perfectly when the model was already close to the right answer, it failed to find the correct settings when starting from a distance. In every scenario where the starting point was slightly off, the fast method got stuck or found the wrong solution, while the slower, traditional method consistently succeeded. The researchers discovered that the problem was not the speed of the method, but the landscape of the problem itself. The sudden jump of the neuron firing creates a jagged, uneven terrain for the computer to navigate. The fast method, which relies on smooth slopes to guide it, gets confused by these jagged edges and falls into dead ends, whereas the slower method is robust enough to climb over them.

The study also revealed that the type of error the computer tries to minimize matters deeply. When the researchers asked the computer to simply match the voltage levels of the recorded signal at every single moment, the fast method would often decide that the easiest way to reduce error was to stop the model from firing at all. It found a quiet, non-firing state to be a better solution than a slightly misaligned firing pattern. This happened because the math of the fast method could not distinguish between a slightly wrong spike and no spike at all in a way that encouraged the model to keep firing. The researchers found that using more complex ways to measure the difference between the model and the data, such as counting spikes or measuring the timing between them, helped a little, but not enough to overcome the fundamental difficulty. Even with these improvements, the fast method could only recover the correct settings if the starting guess was already extremely close to the truth.

Ultimately, the paper concludes that for these specific simplified neuron models, the promise of using fast, gradient-based fitting to replace slow trial-and-error methods has not yet been realized. The mathematical shortcut used to make the models differentiable introduces a bias that misleads the optimizer, and the jagged nature of the firing mechanism creates a landscape that is too difficult for the fast method to navigate reliably. The researchers showed that while their new software implementation is incredibly fast and efficient for running simulations, the optimization strategy itself hits a wall. The slow, derivative-free methods, though computationally expensive, remain the most reliable way to fit these models to data. The work highlights a specific limitation in applying modern machine learning techniques to certain types of biological models, suggesting that for now, the most effective path forward is to stick with the proven, albeit slower, methods or to look for entirely different types of neuron models that do not have these jagged mathematical barriers.

Technical Summary: Surrogate Gradients for Gradient-Based Parameter Estimation in Simplified Neuron Models

Problem Statement
Simplified spiking neuron models, such as the Adaptive Exponential Integrate-and-Fire (AdEx) model, are essential for large-scale brain simulations due to their computational efficiency compared to biophysically detailed Hodgkin-Huxley (HH) models. However, fitting the parameters of these simplified models to experimental recordings has historically relied on derivative-free optimization methods (e.g., grid search, evolutionary algorithms, Nelder–Mead). This reliance stems from the discrete spike-and-reset mechanism inherent in these models, which renders them non-differentiable and prevents the use of standard automatic differentiation (AD) and gradient-based optimization. While surrogate gradients (SGs) have successfully enabled gradient-based training of synaptic weights in Spiking Neural Networks (SNNs) for classification tasks, their efficacy for estimating biophysical parameters in single-neuron models remains untested.

Methodology
The authors investigate whether surrogate gradients can facilitate efficient gradient-based parameter estimation for the AdEx model. The study involves three primary methodological components:

Implementation: A differentiable AdEx model was implemented within the Jaxley framework, a JAX-based simulator designed for biophysical models. To handle the non-differentiable reset condition, the authors employed a surrogate gradient approach using jax.custom_vjp. In the forward pass, the standard hard Heaviside function is used to maintain correct spike dynamics. In the backward pass, the derivative is replaced by a smooth surrogate function (sigmoid, exponential, or SuperSpike). The implementation is verified against the Brian2 simulator and benchmarked for runtime performance.
Loss Function Design: Three categories of differentiable loss functions were evaluated:
- Mean Squared Error (MSE): Applied directly to voltage traces.
- Feature-Based Loss: A soft approximation of the Guarino et al. approach, using differentiable surrogates (sigmoid gates and soft-argmax) to extract features like spike counts, inter-spike intervals, and spike times.
- Distance-Based Metrics: The Van Rossum distance, which convolves spike trains with an exponential kernel, and Soft-Dynamic Time Warping (Soft-DTW).
Evaluation Protocol: The effectiveness of the approach was tested via a "recovery-radius benchmark." This involved generating synthetic target traces from ground-truth AdEx parameters and then attempting to recover these parameters from perturbed initializations. The study compared three gradient-based methods (using the different loss functions) against two derivative-free baselines (Nelder–Mead on soft and hard feature losses). Success was measured by the coincidence factor ( $\Gamma$ ), with a threshold of $\Gamma > 0.5$ indicating a successful recovery. A companion subthreshold benchmark was also conducted to isolate the performance of gradient descent from the spike mechanism.

Key Contributions

Differentiable AdEx in Jaxley: The authors provide the first implementation of a surrogate-gradient-enabled AdEx model within Jaxley. This implementation is end-to-end differentiable, batchable (vmap), and portable to GPUs/TPUs. It achieves a runtime speedup of approximately 200x compared to Brian2's default runtime mode while maintaining accuracy comparable to target-compiled C++ code.
Systematic Benchmarking: The paper presents a rigorous benchmark of 2,700 paired runs comparing gradient-based fitting against derivative-free baselines across three distinct firing regimes (tonic, adaptation, initial bursting) and varying perturbation scales.
Loss Landscape Characterization: The study analyzes the geometry of loss surfaces for different loss functions, revealing that MSE leads to non-spiking attractors, while feature-based and distance-based losses create narrow, curved valleys separated by large non-spiking plateaus.

Results

Performance of Gradient-Based Methods: Across the recovery-radius benchmark, gradient-based fitting never matched the performance of the derivative-free Nelder–Mead baseline. While gradient methods could recover parameters for very small perturbations (recovery radius $\delta \approx 0.05$ ), their success rate collapsed rapidly as perturbation increased. In contrast, Nelder–Mead maintained high success rates up to $\delta = 0.8$ .
Loss Function Sensitivity:
- MSE: Performed the worst, consistently driving the optimizer toward non-spiking solutions because eliminating misaligned spikes reduced the pointwise voltage error more effectively than aligning them.
- Feature-Based (Guarino) & Van Rossum: These performed better than MSE but still failed to recover parameters for moderate-to-large perturbations. The Van Rossum distance showed slightly better robustness than the feature-based loss in some regimes, likely due to fewer softening hyperparameters.
Subthreshold Benchmark: In a control experiment where the spike mechanism was disabled (subthreshold regime), gradient descent (Adam) outperformed Nelder–Mead in terms of sample efficiency and convergence speed. This indicates that gradient descent itself is not the failure point; rather, the failure arises from the interaction between the surrogate gradient, the spike mechanism, and the loss landscape geometry.
Runtime: The Jaxley implementation is significantly faster than Brian2 in its default Python/Cython runtime mode (geometric mean speedup of 188x), though it is comparable to Brian2's compiled C++ standalone mode.

Significance and Claims
The paper concludes that surrogate gradients, while successful for training synaptic weights in SNNs, do not currently enable efficient gradient-based parameter estimation for simplified neuron models like AdEx when compared to state-of-the-art derivative-free methods. The primary obstacles identified are:

Surrogate Bias: The surrogate gradient approximates the gradient of a smoothed loss function, not the true loss. For single-neuron fitting, this bias displaces the optimal parameters, preventing exact recovery.
Loss Landscape Pathologies: The loss surfaces contain large non-spiking plateaus and narrow valleys. Gradient descent struggles to navigate these geometries, often getting stuck in local minima or being pulled toward non-spiking solutions (especially with MSE).
Vanishing Gradients: Gradients decay exponentially through the leak dynamics, particularly for parameters that only influence the system during spike events (e.g., adaptation parameters).

The authors assert that the failure is not intrinsic to gradient descent itself (as evidenced by the subthreshold results) but is a consequence of the specific interaction between the discrete spike mechanism, the surrogate approximation, and the loss function design. They suggest that future progress may require different neuron models with closed-form solutions (e.g., Quadratic Integrate-and-Fire with pseudospikes) or more complex optimization strategies involving multi-trace constraints and hybrid objectives, rather than simply tuning surrogate functions or loss weights.

Technical Summary: Surrogate Gradients for Gradient-Based Parameter Estimation in Simplified Neuron Models

More like this