U ltra S parse BERT : 99% Conditionally Sparse Language Modelling (2024)

Abstract: We present UltraSparseBERT, a BERT variant that uses 0.3% of its neurons during inference while performing on par with similar BERT models. UltraSparseBERT selectively engages just 12 out of 4095 neurons for each layer inference. This is achieved by reorganizing feedforward networks into fast feedforward networks (FFFs).To showcase but one benefit of high sparsity, we provide an Intel MKL implementation achieving 78x speedup over the optimized feedforward baseline on CPUs, and an OpenAI Triton implementation performing forward passes 4.1x faster than the corresponding native GPU implementation. The training and benchmarking code is enclosed.

## introduction feedforward layers hold the majority of the parameters of language models @xcite . however, not all of their neurons need to be engaged in the computation of the feedforward layer output at inference time for every input. a growing body of work is taking advantage of this fact in a top-down fashion, making use of a method commonly referred to as "mixture of experts" @xcite @xcite . this method consists of subdividing a large feedforward network into blocks ("experts"), designating some blocks to form a gating network, and jointly training both the experts and the gating network to produce the layer's outputs while using only a fraction of layer parameters, conditionally on the input. the covariant approach, dubbed "fast feedforward networks", is to introduce conditional exe-cution in a bottom-up fashion, utilizing individual neurons rather than blocks to perform gating and be executed conditionally @xcite . we employ this approach and produce ultrasparsebert, a variant of the bert architecture @xcite that reorganizes feedforward networks into simplified fast feedforward networks (fffs). in terms of downstream performance, ultrasparsebert performs on par with other bert-like models that are similar in size and undergo similar training procedures. the intermediate layers of ultrasparsebert are, however, effectively much sparser by design: given a feedforward (ff) and a fast feedforward (fff) network, each with n neurons, the fff uses the parameters of only o (log 2 n) neurons instead of o (n) as for ff. this is a consequence of the fact that fffs organize their neurons into a balanced binary tree, and execute only one branch of the tree conditionally on the input. in terms of output produced by the intermediate layers, such a method of execution is equivalent to treating the weights of all unused neurons as zeroes and manifests itself as conditional sparsity, since the choice of effectively non-zero neurons is conditional on the layer input. performing inference on an fff amounts to performing conditional matrix multiplication (cmm), in which the rows of the input dot with the columns of neural weights one at a time, and the weight column to proceed with is chosen depending on the output of the previous dot-product operation. in this manner, all neurons are used only by some inputs and no input needs more than just a handful of neurons to be handled by the network. this is in contrast with dense matrix multiplication (dmm), which lies at the heart of the traditional feedforward networks, and which computes the dot products of all rows with all columns. recent advances in deep learning infrastructure have made it possible to produce efficient implementations of conditional matrix multiplication based on both popular computational frameworks as well as custom kernel code. we showcase and provide three implementations of fff forward pass based on advanced pytorch compilation, the ope-nai triton framework, and the intel mkl routines. in a later section, we give a comparison of each implementation to the corresponding optimized baseline and note that while there is already clear evidence of significant acceleration, there is potential for more. reproducibility. we share our training, finetuning, and benchmarking code as well as the weights of our best model. for a quick conceptual verification, the fact that only 12 neurons are used in the inference of ultrasparsebert can be verified simply by zeroing the output of all but the chosen neurons, and we also give the code for this. ## inference fffs as a part of large language models have a considerable acceleration potential. at the center of their promise sits the operation of conditional matrix multiplication. ## limitations a limitation of our training work is that for most fff configurations, we only perform one training run. it is possible that the downstream performance of the individual configurations would vary across multiple training runs. this is partially mitigated by the use of multiple fine-tuning runs to find the downstream task score as per the precedent for bert models on the glue benchmark. a major weakness of inference speed measurements is that they depend heavily on the hardware used as well as the low-level optimization provided as the interface to the hardware. to illustrate how fast the landscape is changing: in october 2023, neither the non-copying bmm nor the triton implementation leveraging local conditionality would have been possible. our sparsity argument, however, remains intact, and is easily verifiable through the (default provided) implementation that zeroes out the contributions of all unused neurons. our work focuses on efficiency of existing models and inherits the risks of the models used, if any.