HelioX: A GPU-Native Framework for Simulation and Training of Biophysically Detailed Networks

Biophysically detailed neural networks represent a promising frontier for brain-inspired AI, offering intrinsic spatio-temporal dynamics to enhance the expressivity and computational density of deep learning systems. However, general-purpose deep learning frameworks suffer from a fundamental mismatch between their dense parallel optimizations and the irregular, tree-structured complexity of biological mechanisms. In this work, we propose HelioX, a \textbf{GPU-native} framework designed to unify high-performance simulation with scalable training. Unlike approaches that adapt biology to existing deep learning tools, HelioX adopts a “GPU-to-Biophysics'' paradigm. We tailor the underlying GPU parallelism to biological structures by implementing custom-fused CUDA kernels for both the Dendritic Hierarchical Scheduling (DHS) algorithm and its gradient propagation. This design eliminates the runtime overhead of generic automatic differentiation and enables multi-stream concurrency for spike generation and equation assembly. Experimental results demonstrate that HelioX outperforms standard simulators (NEURON) by orders of magnitude and surpasses prior GPU-based solvers in both speed and scalability. We successfully train deep biophysical MLPs and whole-brain-scale biophysical circuits (e.g., the BAAIWorm \textit{C. elegans} model) on a single consumer-grade GPU. HelioX establishes a new standard for computational efficiency, enabling the training of biophysically detailed models at scales previously unattainable.

Modern AI usually models each neuron as a very simple unit, which makes large systems easy to train but leaves out much of how real neurons compute. Real neurons have branching structures and rich electrical dynamics, and these details may help both neuroscience and brain-inspired AI, but they are extremely expensive to simulate and train. We built HelioX to make these more realistic neuron models practical on GPUs, the chips that power modern AI. Rather than squeezing biological models into software designed for ordinary neural networks, HelioX organizes GPU computation around the tree-like shape and time-varying behavior of biological neurons. This lets it simulate detailed neurons quickly and compute learning updates without the heavy overhead of general-purpose training tools. In our experiments, HelioX matches the numerical results of standard neuroscience simulators while running much faster and using less memory. It can train deep biophysical neural networks and adjust a whole-brain-scale circuit model of the tiny worm C. elegans on a single consumer GPU. By lowering the cost of working with realistic neuron models, HelioX gives researchers a faster way to study biological computation and explore new forms of AI inspired by the brain.