JAXLEY: differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics

Biophysical neuron models provide insights into cellular mechanisms underlying neural computations. A central challenge has been to identify parameters of detailed biophysical models such that they match physiological measurements or perform computational tasks. Here we describe a framework for simulating biophysical models in neuroscience-JAXLEY-which addresses this challenge. By making use of automatic differentiation and GPU acceleration, JAXLEY enables optimizing large-scale biophysical models with gradient descent. JAXLEY can learn biophysical neuron models to match voltage or two-photon calcium recordings, sometimes orders of magnitude more efficiently than previous methods. JAXLEY also makes it possible to train biophysical neuron models to perform computational tasks. We train a recurrent neural network to perform working memory tasks, and a network of morphologically detailed neurons with 100,000 parameters to solve a computer vision task. JAXLEY improves the ability to build large-scale data- or task-constrained biophysical models, creating opportunities for investigating the mechanisms underlying neural computations across multiple scales.