Sparse Mixers: Combining M o E and Mixing to build a more efficient BERT (2022)

Abstract: We combine the capacity of sparsely gated Mixture-of-Experts (MoE) with the speed and stability of linear, mixing transformations to design the Sparse Mixer encoder model. Sparse Mixer slightly outperforms BERT on GLUE and SuperGLUE, but more importantly trains 65% faster and runs inference 61% faster. We also present a faster variant, prosaically named Fast Sparse Mixer, that marginally underperforms BERT on SuperGLUE, but trains and runs nearly twice as fast. We justify the design of these two models by carefully ablating through various mixing mechanisms, MoE configurations, and hyperparameters. Sparse Mixer overcomes many of the latency and stability concerns of MoE models and offers the prospect of serving sparse student models, without resorting to distilling them to dense variants.

## introduction sparsely gated mixture-of-experts (moe) models have seen a surge of interest in recent years @xcite @xcite @xcite @xcite . moe models offer the promise of sublinear compute costs with respect to the number of model parameters. by training "experts" that can independently process different slices of input data, moe layers increase model capacity with limited increases in flops. perhaps because of the favorable capacity-tocompute trade-off, most recent moe studies, including the aforementioned, have focused on using moe to scale up large models. using moe layers to scale to larger models offers quality and total train efficiency gains over dense models, but not train or inference step latency benefits. indeed, the task of serving these models in practice is either ignored or relegated to distilling the sparse teacher model to a dense student model @xcite , often with a significant quality loss relative to the sparse teacher model. for example, @xcite are only able to distill roughly 30% of the switch transformer's quality gains to a dense model. orthogonal to moe, efficient mixing models @xcite @xcite replace attention in transformer-like models with simpler linear transformations or mlp blocks that "mix" input representations. linear transformations are particularly attractive because they are faster than the combined projection and dot product operations in an attention layer. in this work, we pull on both moe and mixing threads to build low latency, sparse encoder models that we hope can used in production settings. we focus on encoder models, and bert-like models in particular, because they are widely used in practice -for example, in dual encoders for retrieval @xcite . relative to the vanilla transformer model @xcite , we speed up our model in two ways. (1) we use the increased capacity from moe sublayers to offset parameter reductions in other parts of the model. (2) we use mixing transformations to replace a large fraction of self-attention sublayers with faster, linear transformations. the resulting model, which we name sparse mixer, slightly (< 1%) outperforms bert on glue @xcite and superglue @xcite , but most importantly trains 65% faster and runs inference 61% faster. we also introduce a simple variant of sparse mixer, prosaically named fast sparse mixer, that marginally (< 0.2%) under-performs bert on superglue, but runs nearly twice as fast: training 89% faster and running inference 98% faster. an interesting finding of our work is a training stability synergy between the sparse and mixing model components. as a point of comparison, we find that simply replacing dense feed-forward sublayers in bert with moe variants yields highly unstable models; see section 5. however, these instabilities dissipate as we replace self-attention sublayers with mixing sublayers. we hypothesize that the (token-dependent) relevance weighted selfattention basis is the source of the instability, and hence that replacing the majority of self-attention sublayers with mixing sublayers renders sparse mixer models highly stable. in summary, we introduce two models: we justify the design of these models by ablating through model mixing, moe, and hyperparameter configurations. with sparse mixers, we demonstrate that the speed and stability regressions of moe models may be overcome using mixing mechanisms. this offers the promise of directly serving sparse models, rather than resorting to distilling them to dense variants. ## related work sparsely gated mixture-of-experts. mixture-of-experts (moe) models were introduced by @xcite ; jordan and @xcite and more recently popularized by @xcite . recent work, such as @xcite , has played out the promise of moe models by achieving state of the art results on a number of nlp benchmarks. as with contemporary moe studies @xcite , these models are large and primarily focus on model quality. when efficiency is studied, it is typically at the level of a total train time efficiency metric. for example, although the per training step speed of the switch transformer @xcite is slower than the vanilla transformer, because the switch transformer surpasses the vanilla model's top accuracy in a fraction of the steps, the switch transformer can be correctly described as a more efficient model. however, the slower step speed is an achilles heel for serving such models; one generally cannot ask a user to wait longer for a more accurate model response. a notable exception is @xcite , who sparsify multiple components of the transformer, primarily by replacing softmaxes with argmaxes, to achieve an over 2x speed-up in unbatched inference speed on cpus for base/large model sizes in contrast to our work, their speedups to not carry over to accelerator hardware or to training. memory mechanisms are another popular sparse technique for adding capacity to models with limited increases in compute; see, for example, @xcite @xcite . while intuitively appealing and empirically promising, suboptimal implementations (look-ups in particular) for accelerator hardware often yield memory models that have favorable theoretical compute properties, but are slow in practice. mixing. several recent works have explored mixing mechanisms, such as matrix multiplications @xcite , mlp blocks @xcite , and spectral transforms @xcite , as an efficient replacement of attention in transformer-like models. @xcite ; @xcite ; @xcite find that hybrid attention-mixing models, wherein partial or a limited number of attention sublayers are retained, were faster than transformers with only very limited accuracy degradation. building off these works, we use sparse mlp sublayers to compensate for the remaining accuracy gap. ## coordinate descent we train in a typical transfer learning setting @xcite : masked language modelling (mlm) and next sentence prediction (nsp) pretraining, followed by fine-tuning on glue @xcite and superglue @xcite . when comparing models, we always use the exact same setup for all models and baselines. in particular, we follow the pre-training setup in @xcite with a few updates: (1) we pretrain on the much larger c4 dataset @xcite ; (2) we use a 32000 sentencepiece vocabulary model @xcite in this section, we follow a "coordinate descent" through our model configurations until we arrive at the final sparse mixer design. given the large number of model hyperparameters to explore, we be computed using the fft @xcite . this requires three operations: one fft to diagonalize the computation, one to apply the diagonalized matrix multiplcation and one ifft to transform back to real space. a toeplitz matrix may be embedded in a circulant matrix to take advantage of the same fft computation. in practice, we find that for standard sequence lengths (512), using the fft is slower than direct matrix multiplications on both gpu and tpu. for our coordinate descent study, we only pretrain for 500k steps, which we found to be reasonably indicative of model performance. models are fine-tuned with the same batch size (64) on the validation split of each respective glue task for 5 epochs and the best result for each task is selected from across three default base learning rates, adapted from devlin et al. (2019): {10 -5 , 5 • 10 -5 , 10 -4 }. our final model is pretrained for longer and is evaluated on both glue and superglue for a broader set of training configurations in section 5. we prioritize efficiency -speed and accuracy. we use pre-training step speed as a proxy for model latency. we rely on downstream average glue scores as our primary accuracy metric, but fallback to upstream mlm and nsp accuracies when glue scores between model variants are similar. additional coordinate descent experiments are summarized in appendix a.2, and full glue results for all coordinate descent experiments are provided in appendix a.3. ## conclusions mixing transformations and moe play well together. utilizing moe for capacity and mixing for speed and stability, we introduced the sparse mixer -a model that slightly outperforms bert on glue and superglue, but more importantly trains 65% faster and runs inference 61% faster. we also presented a faster variant, fast sparse mixer, that marginally under-performs bert on super-glue, but trains and runs nearly twice as fast: 89% faster training and 98% faster inference. sparse mixer overcomes many of the speed and stability concerns of moe models and offers the prospect of serving sparse student models.