MUX - PLM s: Pre-training Language Models with Data Multiplexing (2023)

Abstract: The widespread adoption of large language models such as ChatGPT and Bard has led to unprecedented demand for these technologies. The burgeoning cost of inference for ever-increasing model sizes coupled with hardware shortages has limited affordable access and poses a pressing need for efficiency approaches geared towards high throughput and performance. Multi-input multi-output (MIMO) algorithms such as data multiplexing, offer a promising solution with a many-fold increase in throughput by performing inference for multiple inputs at the cost of a single input. Yet these approaches are not currently performant enough to be deployed in modern systems. We change that by developing MUX-PLMs, a class of high throughput pre-trained language models (PLMs) trained with data multiplexing, that can be fine-tuned for any downstream task to yield high-throughput high-performance. Our novel multiplexing and demultiplexing modules proficiently entangle and disentangle inputs, and enable high-performance high throughput that are competitive with vanilla PLMs while achieving 2x/5x inference speedup with only a 1−4% drop on a broad suite of tasks.

## introduction language models like chatgpt (openai, 2023), palm @xcite , t5 @xcite , and cm3 @xcite , have seen unprecedented adoption in diverse sectors ranging from education and healthcare to manufacturing and marketing. the proficiency of these tools has led to unprecedented demand for these models, with users facing frequent outages and capacity limits. additionally, ever-increasing model sizes and hardware shortages have constrained models' ability to handle a very high load of requests, thus limiting large-scale affordable access to these models. these trends bring into focus the need for high-throughput, high-performance, efficient, and environmentally responsible models that can be deployed at scale to meet the quickly growing demand. multi-input multi-output architectures (mimo) @xcite @xcite are a promising hardware-agnostic and architecture-agnostic paradigm that perform inference for multiple inputs simultaneously at the cost of a single input. this efficiency paradigm is natively geared towards yielding high-throughput models, in addition to being complementary in approach and motivation to current efficiency methods such as pruning, quantization, and distillation. interestingly, mimo approaches are partly inspired by the human brain's extraordinary ability to process multiple inputs and propagate information at a high bandwidth with a few neural codes @xcite @xcite @xcite . @xcite introduced data multiplexing, a mimo technique that can enable a many-fold increase in throughput. the method compresses n different instances into a single "multiplexed" hidden representation before decompressing it into n independent predictions. while they show the plausibility of mimo training, their method leads to a significant drop in performance (20 -30% points) compared to state-of-the-art models. in this work, we introduce mux-plms, a class of high-throughput pre-trained language models trained in a mimo fashion with data multiplexing to process multiple inputs (2-10) simultaneously with a forward pass over a single instance. mux-plms offer up to 400% improvement in throughput over baseline pre-trained models while only being ∼ 4 points and ∼ 2 points worse than baseline pre-trained language models for text classification and token classification tasks, respec-< l a t e x i t s h a 1 _ b a s e 6 4 = " w v + u n z t f n y 3 c v 0 0 / z y a k d o 4 k v o 8 = " > a a a c b 3 i c b v b n s 8 n a e n 3 u r 1 q / q h 4 f w s y c p 5 j i u y 9 f l 1 6 e c q y t n k f s t t t 2 6 w 4 s d i d i c b l 5 8 a 9 4 8 a c i v / + c n / + n 2 7 q h b x 0 w z o o 9 g x b n b b h g g m z 7 2 y o s l a + s r h x x s x u b w 9 s 7 5 d 2 9 p o 4 s r z l l i x g p d k a 0 e z x k l n a q r b 0 r r m q g w c s y x u 3 8 1 j 1 t m k f h h y x j 5 k s y c h m f u w j g 6 p y p p w a p o g s a d 4 a 0 9 b q k q u a b t 5 1 l w b d c s a t 2 d r x i n b m p o b k a 3 f k x 1 4 t o i l k i v b c t o 4 4 d g 5 8 s b z w k l p w 8 r l o y 0 b e z s i 6 h i z f m + 2 l + r 4 a p j d l d / u i z c g h n 6 u + n l e i t x z i w k 5 l a u m 9 7 e / e / r 5 n a / 8 j p e r g n w e i 6 f a i f c a w r n o s c e 1 w x c m j s c k g k m 7 9 i o i s k u d d r l u w i z v z j i 6 r 5 w n x o q r x b w q v + o y u j i a 7 q e t p b d j p h d x s n g s h f f d 2 i z / s k 3 q w n 6 8 v 6 t z 6 m o w v r t r o p / s d 6 / a g k i z p @xmath0 figure 1: illustrating the training process for mux-plms. mux-plms are first primed for mimo style training with a token-retrieval auto-encoding task, where the model is trained to output the tokens in the n inputs. mux-plms are then pre-trained by adapting standard pre-training objectives (bert in this example), to mimo style training with data multiplexing. the resulting mux-bert model, similar to standard plms, provides a general model initialization that can be fine-tuned on any downstream task (ner in this example). output predictions are shown above the system head with highlighted predictions contributing to the gradient update; violet indicates a correct prediction while orange indicates an incorrect prediction. red highlighted tokens in the input indicate a position that has been masked. our work is the first to introduce mimo architectures to plms. to enable this, we first develop a new demultiplexing module, rsa-demux (figure 2 ), that randomly initializes and learns private key vectors to recover the multiple outputs from a multiplexed representation. secondly, we introduce a new contextual multiplexer module (figure 3 ) that uses a cross-instance attention-based mechanism to aggregate context across the set of multiplexed instances, which seems to be particularly effective for token-level tasks. thirdly, our three-stage training algorithm (figure 1 ) enables stable and efficient training of mux-plms. importantly, mux-plms are complementary to existing state-of-the-art model compression techniques. we hope our work validates mimo architectures as a promising complementary direction to existing efficiency techniques. consequently, we hope future research develops mimo architectures in tandem with other efficiency approaches, leveraging the best of both paradigms. we publicly release our models and code to promote opensource research on the next generation of mimo architectures for large language models. the ubiquity of pre-trained language models, their growing size, and over-parameterization has inspired extensive research on improving inference efficiency. this includes methods such as structured pruning @xcite @xcite @xcite , knowledge distillation @xcite @xcite @xcite , quantization @xcite , and data multiplexing @xcite . these approaches assume that plms are highly over-parametrized and attempt to approximate a large function by learning a smaller, compressed, version of the original model. past work has also focused on unstructured pruning for both task finetuning @xcite and pre-trained @xcite language model settings, but don't increase model throughput due to hardware limits. multi-input multi-output models while pruning, quantization, and distillation seek to reduce overparameterization by reducing models' representational capacity, other lines of work seek to exploit overparameterization in other ways. multiinput multi-output (mimo) architectures @xcite @xcite train models using mixed-instance representations, i.e. @xcite , in order to obtain predictions for multiple instances simultaneously. unlike efficiency methods, @xcite and @xcite try to obtain better performance by inducing multiple subnetworks in a single convolutional model to perform "ensembling for free" during inference. data multiplexing, introduced in datamux @xcite , aims to improve model efficiency by training transformer models with mixed-instance representations to perform simultaneous inference for language tasks, thereby improving inference throughput many-fold. currently, mimo architectures have only been used in a limited setting, achieving middling performance. our work training plms with data multiplexing, dramatically improves inference throughput while preserving high accuracy for downstream tasks. ## methodology we briefly introduce readers to the data multiplexing mimo architecture @xcite , which we denote t-mux. we then detail our novel approach to train mux-plms to yield highthroughput and performant language models. ## experimental setup datasets we pre-train all models on wikipedia (foundation) and bookscorpus @xcite . we evaluate on all datasets from the glue benchmark @xcite , which are cola @xcite , sst-2 @xcite , qqp (qqp), sts-b @xcite , mnli @xcite , qnli @xcite , rte @xcite , and wnli @xcite . we also evaluate on token classification tasks -named entity recognition @xcite and pos tagging @xcite . we don't report average over the two smallest tasks in glue, wnli and cola, as we observe high variance across seeds. all numbers are reported on the dev split. we report scores for all tasks in appendix e. ## conclusion we introduce mux-plms, a class of highthroughput pre-trained language models trained with data multiplexing, a multi-input multi-output (mimo) architecture. our mux-plms models, trained with novel mimo modules, are competitive with state-of-the-art plms on several downstream tasks while achieving a many-fold increase in inference throughput. mux-plms, similar to standard plms, can be fine-tuned on any downstream task to yield high-throughput, high-performance models. we hope our work inspires future research in mimo architectures for plms as a complementary efficiency paradigm to existing approaches.