Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Differentiable Approach to Multi-scale Brain Modeling

Chaoming Wang, Muyang Lyu, Tianqiu Zhang, Sichao He, Si Wu

TL;DR

This work presents a multi-scale differentiable brain modeling workflow utilizing BrainPy, a unique differentiable brain simulator that combines accurate brain simulation with powerful gradient-based optimization and offers a promising tool to bridge neuroscience data across electrophysiological, anatomical, and behavioral scales.

Abstract

We present a multi-scale differentiable brain modeling workflow utilizing BrainPy, a unique differentiable brain simulator that combines accurate brain simulation with powerful gradient-based optimization. We leverage this capability of BrainPy across different brain scales. At the single-neuron level, we implement differentiable neuron models and employ gradient methods to optimize their fit to electrophysiological data. On the network level, we incorporate connectomic data to construct biologically constrained network models. Finally, to replicate animal behavior, we train these models on cognitive tasks using gradient-based learning rules. Experiments demonstrate that our approach achieves superior performance and speed in fitting generalized leaky integrate-and-fire and Hodgkin-Huxley single neuron models. Additionally, training a biologically-informed network of excitatory and inhibitory spiking neurons on working memory tasks successfully replicates observed neural activity and synaptic weight distributions. Overall, our differentiable multi-scale simulation approach offers a promising tool to bridge neuroscience data across electrophysiological, anatomical, and behavioral scales.

A Differentiable Approach to Multi-scale Brain Modeling

TL;DR

This work presents a multi-scale differentiable brain modeling workflow utilizing BrainPy, a unique differentiable brain simulator that combines accurate brain simulation with powerful gradient-based optimization and offers a promising tool to bridge neuroscience data across electrophysiological, anatomical, and behavioral scales.

Abstract

We present a multi-scale differentiable brain modeling workflow utilizing BrainPy, a unique differentiable brain simulator that combines accurate brain simulation with powerful gradient-based optimization. We leverage this capability of BrainPy across different brain scales. At the single-neuron level, we implement differentiable neuron models and employ gradient methods to optimize their fit to electrophysiological data. On the network level, we incorporate connectomic data to construct biologically constrained network models. Finally, to replicate animal behavior, we train these models on cognitive tasks using gradient-based learning rules. Experiments demonstrate that our approach achieves superior performance and speed in fitting generalized leaky integrate-and-fire and Hodgkin-Huxley single neuron models. Additionally, training a biologically-informed network of excitatory and inhibitory spiking neurons on working memory tasks successfully replicates observed neural activity and synaptic weight distributions. Overall, our differentiable multi-scale simulation approach offers a promising tool to bridge neuroscience data across electrophysiological, anatomical, and behavioral scales.
Paper Structure (29 sections, 14 equations, 8 figures, 3 tables)

This paper contains 29 sections, 14 equations, 8 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Methods
  3. Differentiable neuron models with surrogate gradients
  4. Event-driven differentiable synaptic operators
  5. Workflow for multi-scale differentiable brain modeling
  6. Training biologically-informed spiking networks on cognitive tasks
  7. Neuron fitting
  8. Task training
  9. Conclusion and discussion
  10. Software and Data
  11. Acknowledgements
  12. Surrogate gradient functions
  13. Event-driven synaptic operators
  14. Loss functions for neuron fitting
  15. Mean square error for fitting membrane potentials
  16. ...and 14 more sections

Figures (8)

  • Figure 1: Multi-scale differentiable brain modeling workflow. The entire workflow is executed using the differentiable brain simulator BrainPy wang2023brainpywang2024brainpy. (A) At the microscale level, the single neuron and synapse model are fitted based on electrophysiological recording data and gradient-based optimizations. (B) At the mesoscopic level, connectome constraints are incorporated into the network construction, facilitating the integration of structural connectivity information. (C) At the macroscale behavior level, gradient-based optimization methods are applied to train the above data-constrained model networks to reproduce the cognitive behaviors as seen in humans or animals.
  • Figure 2: Overview of the neuron fitting procedure. (A) Experimental data: Step currents are injected into the neuron, and the resultant membrane potential responses are recorded. (B) Illustration of the optimization procedure: Parameter values are initialized from a distribution (initialization). Neurons with these parameters are simulated in parallel, and their outputs are compared with the ground truth data (simulation). The prediction error is utilized to estimate gradients (gradient), which are then used to update the initialized parameters for the subsequent iteration (update). (C) Fitting results of the HH model on a cortical pyramidal cell using five different optimization methods.
  • Figure 3: Training the biological-informed excitatory and inhibitory spiking networks using the evidence accumulation task. (A) The input spike train. (B) The recurrent spiking dynamics. (C, D) The membrane potentials of five excitatory (C) and inhibitory (D) neurons. (E, F) The synaptic weight distribution before (E) and after (F) training.
  • Figure S4: The collection of surrogate gradient functions $g'(x)$ in BrainPy wang2023brainpywang2024brainpy, where $x\ge 0$ represents the neuronal membrane potential exceeding the spiking threshold.
  • Figure S5: Architecture of the recurrent spiking EI network. The network consists of excitatory (E) and inhibitory (I) spiking units, denoted by $\mathbf{r}(t)$. These units are trained using an online gradient-based learning framework BrainScale brainscale. Time-varying inputs $\mathbf{u}(t)$ are received by the network, and the recurrent activity is encoded through time-varying outputs $\mathbf{z}(t)$. The inputs represent task-relevant sensory information or internal rules, while the outputs encode a decision in the form of an abstract decision variable, probability distribution, or direct motor output. Each spiking unit exhibits its own dynamics, and the firing rate of each unit is adjusted through our differentiable fitting method (Section \ref{['sec:neuron:fitting']}). The connectivity between the spiking units is determined based on connectomic measurements theodoni2022structural.
  • ...and 3 more figures