Muratore, P., Moradimanesh, Z., Shaikh, R.A., Zargarnezhad, N., Meek, A., Kuzovkin, I.
Neuromatch Academy, Neuromatch, Inc., USA
Published at zenodo.org/records/15126313
Keywords: bio-inspired architectures, neural networks scaling laws, efficiency-performance trade-off
Abstract
Artificial neural networks (ANNs) demonstrate remarkable performance across various domains, yet their computational demands remain a significant challenge. In contrast, biological neural networks achieve exceptional efficiency through sparse structured connectivity shaped by evolution, while maintaining high performance. This study investigates whether bio-inspired connectivity patterns can improve ANN design by balancing efficiency and performance. We systematically compare the performance of Neural Circuit Architectural Priors (NCAP), a neuro-inspired reinforcement learning model, with that of Multi-Layer Perceptrons (MLPs) compressed through pruning and knowledge distillation, in a swimming task. Our findings reveal that NCAP achieves comparable performance to larger, more complex MLPs with significantly fewer parameters. We also observe that sparse, but optimally connected small MLPs often outperform larger MLPs with suboptimal connectivity, underscoring the importance of connectivity in efficient learning. This work paves the way for the development of computationally efficient models that maintain high performance by integrating biological connectivity principles into ANN architectures.
Introduction
Despite significant advancements in artificial neural networks (ANNs), the design of computationally efficient architectures remains a major challenge. The search for an optimal architecture is inherently complex due to the vast design space created by the numerous possible configurations of network structures. This is while biological neural networks, evolved through millions of years of evolution, achieve remarkable performance, despite their small scale compared to state-of-the-art models. In these systems, structural connectivity, often characterized by hierarchical, sparse, and recurrent structures, is a fundamental property that influences function and results in high efficiency. In ANNs, however, the trade-off between efficiency and performance remains a crucial concern. While increasing model size is considered a solution to improve performance, it often comes at a substantial cost in resources. Aiming to reduce these costs, efforts have been made to compress models into smaller networks. These efforts have resulted in methods such as connection pruning, in which redundant or less important information such as weights, neurons, or layers are removed, and knowledge distillation, in which a smaller and more efficient student model is trained on the knowledge of a larger and more complex teacher model. However, these approaches themselves impose high computational demands. This raises a critical question of whether structural connectivity patterns observed in biological neural networks can effectively serve as a blueprint for designing ANNs that achieve both efficiency and high performance.
Recent research has begun to explore this question more directly by drawing inspiration from biological systems when designing ANN architectures. Rather than relying solely on compression techniques, these approaches incorporate biologically motivated structural priors to bridge the gap between the efficiency of biological networks and the performance demands of artificial ones. For example, Neural Circuit Architectural Priors (NCAP) is a Reinforcement Learning (RL) model with a specific architecture that is translated from the neural circuits of C. elegans. In a swimming task, the NCAP model utilizes only a small fraction of the parameters required by Multi-Layer Perceptron (MLP) controllers and achieves comparable performance. In addition, further studies have demonstrated the effectiveness of biologically inspired connectivity patterns in more complex architectures, such as those supporting quadruped locomotion. Importantly, such bio-inspired models demonstrate strong performance from the very outset of training, in contrast to the extended warm-up periods often required by fully connected MLPs. This phenomenon is commonly observed in nature. For instance, mammals can walk within hours of birth, gradually refining their abilities through practice.
In this study, we investigate the computational advantages of biologically inspired architectures by comparing the performance of NCAP, a neuro-inspired model, with MLPs compressed through pruning and knowledge distillation in a swimming task. Our goal is to examine the role of connectivity patterns in achieving an optimal balance between efficiency and performance, and to investigate whether compression alone can yield models with a similar efficiency-performance balance to biological neural networks. We explore this within the context of RL, which offers a relevant experimental framework for evaluating bio-inspired architectures.
Methods
NCAP provides a stunning example of how a neural network with strong priors can control six joints (five links) to achieve high performance locomotion while having only four trainable parameters. In comparison, according to our experiments, a naive MLP architecture of at least 100 parameters is required to reach comparable performance. However, as the NCAP model’s biological priors have demonstrated, the number of parameters is not the only factor that drives performance. Keeping the number of parameters fixed, but changing the connectivity pattern has a huge effect on the network’s ability to learn and its final performance. To systematically explore the role of connectivity, we conducted 250 experimental trials of a swimming task without a designated target within a MuJoCo environment, assigning each trial a randomly generated MLP architecture that varied along three primary dimensions: (1) network depth, sampled between 1 and 4 layers; (2) layer width, with each layer comprising between 1 and 64 neurons; and (3) sparsity coefficient, defined as the proportion of active connections, ranging between 0 and 1, where a value of zero indicates a fully connected layer. It is worth noting that the depth, width, and sparsity coefficients were selected uniformly at random to ensure even coverage across the chosen range of values. For each architecture, training was performed using one of the two approaches:
Reinforcement Learning from Scratch
RL via Proximal Policy Optimization (PPO), in which models were trained from scratch using policy gradient methods (black squares in Figure 1). We selected PPO for its stability and relative insensitivity to hyperparameters, which allowed us to focus on architectural differences without training instabilities confounding our results. While alternatives such as Soft Actor-Critic (SAC) and Deep Deterministic Policy Gradient (DDPG) offer greater sample efficiency, they tend to be more prone to instability in environments with sparse or noisy rewards, which was a relevant concern in our setting.
Knowledge Distillation
Knowledge distillation, where the top-performing model with the highest reward from previous runs served as the teacher, transferring knowledge to the student model via a mean squared error (MSE) loss function (grey circles in Figure 1). More specifically, for a given episode within the environment, the teacher model was executed, and the resulting tuples of observations and corresponding teacher actions were recorded. The student model was then trained in a supervised manner, using backpropagation, to predict the teacher’s actions based on the observations by minimizing the MSE loss. The student’s performance was subsequently evaluated by deploying it in the environment and assessing its effectiveness. Here, distillation provides a framework to explore the minimal model size necessary to maintain desired performance levels, effectively offering a tighter bound on the extent to which model size can be reduced without significant loss in accuracy.
Comparison with NCAP
For each trained model, we recorded Maximum Reward (the highest reward achieved during training) and reward standard deviation. Then we ranked the models based on the number of parameters. Using the Maximum Reward values and their standard deviation we bootstrapped (1000 samples per model) the distribution of Maximum Rewards for each parameter number. Finally, we calculated the Z-score of NCAP performance given these non-parametric distributions for comparison.
Results
As illustrated in Figure 1, NCAP occupies a distinct position in the performance-parameter space (the purple star), effectively balancing performance and efficiency. The performance-parameter space is defined as a two-dimensional space where the x-axis represents the number of parameters in an agent’s neural network, and the y-axis corresponds to the average reward the agent accumulates, capturing the trade-off between model complexity and performance. Statistical analysis, based on bootstrapped reward distributions across varying numbers of model parameters, further supports this observation: NCAP ranks within the top 1% among models with up to 39 parameters (p < 0.01), the top 5% among models with up to 45 parameters (p < 0.05), and the top 10% among models with up to 74 parameters (p < 0.1). Overall, NCAP remains within the top 31% of all evaluated models.
Furthermore, it is remarkable that certain smaller, yet better connected, MLPs outperform networks that are orders of magnitude larger but poorly connected, highlighting the critical role of connectivity. Moreover, it can be observed that as the number of layers increases, from a depth of one to four, the minimum number of parameters necessary to attain maximal reward also increases, indicating that at least for this task, deeper networks pose greater learning challenges. This suggests that merely scaling the network by adding more layers may introduce unnecessary complexity, thereby further emphasizing the critical significance of effective connectivity. Finally, it’s noteworthy that distillation enables faster achievement of good performance, to the point where, for a depth of four, it took too long to train networks within our time constraint of 5 and 2.5 million iterations for training and distillation, respectively. As a result, at this depth, only the networks trained using distillation achieved maximum performance, with just one single exception (a single black square that reaches a reward close to 1000).
Conclusions and Discussion
In conclusion, these findings emphasize the pivotal role of connectivity patterns in achieving an optimal balance between efficiency and performance. They demonstrate that merely increasing network size is not an effective strategy for improving performance, as we observed many instances where larger networks with suboptimal connectivity performed significantly worse than smaller well-connected networks. It is worth mentioning that factors such as learning rate schedules, initialization methods, and random seed effects may influence these outcomes, and while controlled in this study, their potential as confounders warrants further exploration to enhance the robustness of these findings. This study underscores the promise of utilizing biological connectivity patterns as a blueprint for designing more efficient and effective neural architectures. Concurrently, the distinct position of NCAP, located far from any configuration we have tested, suggests that while connectivity is essential and has facilitated progress toward the NCAP level of efficiency-performance, there is still much to be investigated in order to fully realize brain-like efficiency and advance toward brain-like ANNs. Overall, by integrating insights from neuroscience and AI, this research paves the way for the development of models that are both computationally sustainable and high performing.
Acknowledgements
We would like to extend our gratitude to the Impact Scholars Program for its invaluable support. Additionally, since this study originated as a project for the NeuroAI course at Neuromatch Academy in July 2024, we would like to acknowledge Kaitlyn Hoesterey and Ravikrishnan Jayakumar, members of the original team, for their contributions to the initial conceptualization.
References
Bhattasali, N. X., Zador, A. M., & Engel, T. A. (2022). Neural circuit architectural priors for embodied control. arXiv. https://doi.org/10.48550/arXiv.2201.05242
Bhattasali, N. X., Pattabiraman, V., Pinto, L., & Lindsay, G. W. (2024). Neural circuit architectural priors for quadruped locomotion. arXiv. https://doi.org/10.48550/arXiv.2410.07174
Elsken, T., Metzen, J. H., & Hutter, F. (2019). Neural architecture search: A survey. Journal of Machine Learning Research, 20(55), 1–21.
Friston, K. (2008). Hierarchical models in the brain. PLoS Computational Biology, 4(11), e1000211. https://doi.org/10.1371/journal.pcbi.1000211
Garwicz, M., Christensson, M., & Psouni, E. (2009). A unifying model for timing of walking onset in humans and other mammals. Proceedings of the National Academy of Sciences, 106(51), 21889–21893. https://doi.org/10.1073/pnas.0905777106
Gou, J., Yu, B., Maybank, S. J., & Tao, D. (2021). Knowledge distillation: A survey. International Journal of Computer Vision, 129(6), 1789–1819. https://doi.org/10.1007/s11263-021-01453-z
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., & Amodei, D. (2020). Scaling laws for neural language models. arXiv. https://doi.org/10.48550/arXiv.2001.08361
Marinó, G. C., Petrini, A., Malchiodi, D., & Frasca, M. (2023). Deep neural networks compression: A comparative survey and choice recommendations. Neurocomputing, 520, 152–170. https://doi.org/10.1016/j.neucom.2022.11.072
Meng, G. (2023). Leveraging Knowledge Distillation for Efficient Deep Reinforcement Learning in Resource-Constrained Environments. In 2023 International Conference on Image Processing, Computer Vision and Machine Learning (ICICML) (pp. 811-814). IEEE. http://dx.doi.org/10.1109/ICICML60161.2023.10424815
Neumann, O., & Gros, C. (2023). Scaling laws for a multi-agent reinforcement learning model. arXiv. https://doi.org/10.48550/arXiv.2210.00849
Pearce, T., Rashid, T., Bignell, D., Georgescu, R., Devlin, S., & Hofmann, K. (2024). Scaling laws for pre-training agents and world models. arXiv. https://doi.org/10.48550/arXiv.2411.04434
Schrimpf, M., Kubilius, J., Hong, H., Majaj, N. J., Rajalingham, R., Issa, E. B., Kar, K., Bashivan, P., Prescott-Roy, J., Geiger, F., Schmidt, K., Yamins, D. L. K., & DiCarlo, J. J. (2018). Brain-score: Which artificial neural network for object recognition is most brain-like? bioRxiv. https://doi.org/10.1101/407007
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017). Proximal policy optimization algorithms. arXiv. https://doi.org/10.48550/arXiv.1707.06347
Todorov, E., Erez, T., & Tassa, Y. (2012). MuJoCo: A physics engine for model-based control. 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, 5026–5033. https://doi.org/10.1109/IROS.2012.6386109