NeuroPrune: A Neuro-inspired Topological Sparse Training Algorithm for Large Language Models

TL;DR

NeuroPrune addresses the computational bottlenecks of training and inference in transformer-based LLMs by enforcing topology-aware sparsity through brain-inspired mechanisms. It combines three sparsification strategies—MLP sparsification with preferential attachment, attention group sparsity, and redundancy-based head pruning mapped to a dominating-set problem—into a model- and task-agnostic dynamic sparse training routine. Across BERT, T5, and OPT on GLUE, summarization, and machine translation, NeuroPrune achieves competitive or superior performance while delivering substantial speedups in training (up to about 10x) and favorable inference-time improvements due to structured sparsity. The work demonstrates the practicality of leveraging network topology in sparsification and points to future avenues for combining with other sparse training methods and extending to pre-training.

Abstract

Transformer-based Language Models have become ubiquitous in Natural Language Processing (NLP) due to their impressive performance on various tasks. However, expensive training as well as inference remains a significant impediment to their widespread applicability. While enforcing sparsity at various levels of the model architecture has found promise in addressing scaling and efficiency issues, there remains a disconnect between how sparsity affects network topology. Inspired by brain neuronal networks, we explore sparsity approaches through the lens of network topology. Specifically, we exploit mechanisms seen in biological networks, such as preferential attachment and redundant synapse pruning, and show that principled, model-agnostic sparsity approaches are performant and efficient across diverse NLP tasks, spanning both classification (such as natural language inference) and generation (summarization, machine translation), despite our sole objective not being optimizing performance. NeuroPrune is competitive with (or sometimes superior to) baselines on performance and can be up to $10$x faster in terms of training time for a given level of sparsity, simultaneously exhibiting measurable improvements in inference time in many cases.

Figures (11)

• Figure 1: Resulting sparsity patterns ($\approx 50\%$, $\approx 90\%$) as determined by NeuroPrune in an intermediate Transformer layer of a BERT(-base) model learned on the SST2 GLUE benchmark dataset glue. NeuroPrune sparsifies according to a preferential attachment topology as entire rows/columns of the attention and MLP matrices are zeroed out. Quantitatively, the standard deviation (sd) between the connectivity of neurons in the MLP layers increases up to two orders of magnitude ($50\%$: $12.12$, $90\%$: $4.16$) compared with standard fine-tuning ($0.13$) as seen in Figure \ref{['fig:topo_sparsity_mlp']}, This increase in sd is indicative of preferential attachment, similar to what is seen in biological neurons lynn2024heavy, where a minimal rich-get-richer mechanism is used to produce sparse and heavy-tailed networks. The pattern is qualitatively similar for other layers too, as can be seen in the Appendix.
• Figure 2: Left is the MLP degree distribution for the SST2 dataset using a BERT model indicative of preferential attachment for NeuroPrune as sparsity increases (echoing the degree distribution in brain functional networks vertes2012simple). Standard fine-tuning creates a dense network (black vertical line). Right we see the non-uniformity in connectivity at different sparsity %s across GLUE tasks using NeuroPrune, indicative of this preferential attachment across tasks.
• Figure 3: Left is the fraction of times a head is removed using NeuroPrune when fine-tuning on SST2 with a BERT model. The overall numbers (blue curve) are averaged across layers and runs ($\pm$ sd), where at least $10$ heads are removed. We also show individual layer numbers averaged across runs for the top three layers where most pruning of heads happens (Figure \ref{['fig:headimp']}(right)). We see there is a significant bias towards keeping the last head in each layer leading to a more modular structure and also showcasing preferential attachment, as neurons from the previous and next layers will connect only to these heads. The middle head (head $5$) is also maintained more than most other heads, as it replaces many of the earlier heads. Results averaged across GLUE tasks on the right are similar.
• Figure 4: Performance ($1^{\text{st}}$ column), inference time ($2^{\text{nd}}$ column) and train time ($3^{\text{rd}}$ column) for NeuroPrune and CoFI/$l_1$ on GLUE tasks at different sparsity percentages. The $1^{\text{st}}$, $2^{\text{nd}}$ and $3^{\text{rd}}$ rows correspond to BERT-base, T5-base and OPT-125m models respectively. In the $1^{\text{st}}$ row we see that NeuroPrune outperforms CoFI on the smaller GLUE datasets and is competitive on larger ones, with consistently better inference and train times. In the next two rows, we see that NeuroPrune is largely better than $l_1$ sparsity, especially at intermediate sparsities (25-80$\%$), with notable inference time gains and comparable train time. Qualitatively, similar results are obtained for T5-large and OPT-1.3b as seen in Figure \ref{['fig:sparsity2']} in the appendix.
• Figure 5: Performance (left), inference time (center) and train time (right) for NeuroPrune and DSP on GLUE tasks, where different number of heads are present in a BERT-base model are shown above. NeuroPrune is better or similar (rarely worse) in performance to DSP on most datasets, where it is notably more efficient to train. Inference time is (slightly) improved when many heads are removed, however, the DSP code (simply) masks heads rather than explicitly pruning them like ours does and hence if these masked heads are removed the inference time of DSP might also improve as shown in their paper.