N euro P rune: A Neuro-inspired Topological Sparse Training Algorithm for Large Language Models (2024)

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 10x faster in terms of training time for a given level of sparsity, simultaneously exhibiting measurable improvements in inference time in many cases.

## introduction in the past decade, transformer-based models @xcite leveraging attention mechanisms have led to state-of-the-art performance on nlp tasks and other multimodal applications, in both classification and generation settings. despite the performance improvements, the computational overhead required for training, and inference, hinders progress. the models are large and are typically parameterized by many dense matrices, which also begs the question as to whether this complexity is necessary for better performance. ## notation a transformer consists of multiple identical units. each unit in turn is comprised of a multi-head attention (mha) layer and a feed forward (ffn) or mlp layer (used interchangeably). each attention layer is partitioned into multiple heads h i composed of query q h i , key k h i and value v h i parameter matrices. if d denotes the embedding dimension of each token in an input matrix x then, a mha layer with k heads computes the attention of all heads in parallel and concatenates them. m @xmath0, . . . , h k )w o , where w o is an output dense matrix. the ffn layer in turn has two linear layers, one to expand the dimensions l in : r de × r d and the other to project it back to the original dimension l out : r d × r de . typically d @xmath1.g. in bert d @xmath2. if q, k, v are the query, key, and value matrices of all the heads in a mha layer concatenated together, then q h i , k h i v h i corresponds to the (i -1) d k + 1 to i d k columns in each such matrix respectively. let l in,h i denote the corresponding columns of the mlp and let the use of the superscript in a (l) denotes (attention) layer l of the transformer. we use -→ x to signify that x is a vector. ## related work quantization @xcite , knowledge distillation to a smaller model @xcite , and model pruning are some ways to alleviate the extensive computational cost required by llms. here we review prior work on model pruning, which is most relevant to us. ## method we propose (topological) sparsification of a transformer block at three levels: i) the two (expand and contract) multi-layer perceptron (mlp) layers, ii) the attention layers, and iii) head pruning at the level of attention layers. our method, neu-roprune, is detailed in algorithm 1 with two sub-procedures given in algorithms 2 and 3. the three sparsifications are described next. ## experiments we now test our method in two different settings: i) varying sparsity and ii) varying number of heads. in each setting, we run our method on the glue @xcite tasks, where the dev set is used for testing foot_0 . for i) we also test our method on the cnn/daily mail @xcite summarization task. for ii) we also run our method on a machine translation task for german to english on the iwslt @xcite dataset. baselines and models: when varying sparsity we compare against cofi @xcite , which is a state-of-the-art (sota) method for inducing structured sparsity in llms. since the author-shared code is for bert (encoder) we compare with cofi on bert(-base) for glue tasks. we additionally implemented neuroprune for t5 (encoderdecoder) @xcite and opt (decoder) @xcite models. since a cofi implementation was not available for t5 and opt, and would require modifying the architecture, we apply l 1 sparsity to the attention and mlp blocks of the transformer as a baseline. for summarization, we use a t5(-base) model and again l 1 as a baseline. when varying the number of heads, cofi does not provide an easy way to control for this number, and hence we compare against a specialized head pruning method called differential subset pruning (dsp) @xcite , another sota head pruning method. here too, code is available for bert, but not for t5 or opt, and hence we compare neuroprune on bert(-base) for the glue tasks. we do not show head removal results for the other models as there are no natural baselines like we had for sparsity (l 1 ). for machine translation, we use an 18-layer encoder-decoder model with 6 heads per layer as done in @xcite . metrics: for performance we report accuracy (acc) for the glue datasets except cola where mathew's correlation (mcorr) is a standard metric, rouge for summarization and bleu scores for machine translation. we also report (average) inference and train times. please note that by train time we mean the total time taken to train the model using dynamic sparsity and inference time is the time taken by the model to process a test example based on the sparse model that has been learned. hardware: all experiments were conducted on an nvidia a100 gpu with 40 gb memory. setup details: neuroprune results were obtained by varying α and β from 10 -7 to 0.1 in multiples of 10. the l 1 penalty parameter also took these values. θ took values in {0.15, 0.2, 0.25} for the head removal experiments where the default was set to 0.15 for the glue and summarization tasks. qualitatively, similar results are obtained for t5-large and opt-1.3b as seen in figure 11 in the appendix. it was sufficient to run neuroprune for a single epoch for the larger glue datasets (viz. mnli, qqp, qnli and sst2) and the summarization task, while for the other smaller glue datasets we ran it for 3 to 5 epochs. ϵ for mlp sparsification was set to 10 -4 . cofi finetunes the model before it starts pruning. for smaller glue datasets, finetuning before pruning epochs is 4 and total epochs is 100, while for larger datasets, these numbers are 2 and 20 respectively. all the other parameters were unchanged. to make the comparison fair, we turned off the distillation option. dsp runs 3 epochs of joint (finetuning and mask learning) training. for machine translation, neuroprune results were obtained with α ∈ {0.005, 0.01, 0.05, 0.1, 0.25, 0.5}, β ∈ {0.05, 0.1} and θ varying from 0.1 to 0.4 in increments of 0.02. ϵ was set as in glue. since there was no pretrained model, 15 epochs of pretraining were done followed by 15 epochs of neuroprune finetuning for a total of 30 epochs. dsp joint training was also done with a total of 30 epochs. ## discussion this work shows that neuroprune, inspired by sparsity, modularity, and preferential attachment topology in brain functional networks, is competitive and sometimes even superior to other sota dynamic sparse training methods, even though it does not solely try to optimize performance. it is also more efficient than them in train time with speed-ups also seen in inference. not to mention it is easily transferable across transformer architectures (i.e. encoder, encoder-decoder, decoder) as it does not require adding additional gating variables to the architecture or modifying the architecture in any way. also note that since dynamic sparse training and post-pruning are complimentary to each other one can potentially apply both, that is, the latter after the former. moreover, since dynamic sparse training approaches employ sparsity during training the benefits of sparsity in terms of storage and also sometimes computation are leveraged both during training and inference. in contrast, post-pruning methods apply pruning after (dense) network training, benefiting from sparsity at inference alone. there are multiple avenues for future research. first, it would be interesting to combine our redundancy-based and head-importance pruning methods to produce even more aggressive and efficient pruning of heads. second, the structured strategies we used for fine-tuning could also be tried during pre-training. third, one could test the generalizability of models trained using neu-roprune on related, but different tasks and measure if similar gains can be secured. we believe our work will spur further progress on efficient dynamic sparse training methods that are also performant.