MUX-PLMs: Data Multiplexing for High-throughput Language Models
Year: 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 everincreasing 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 MUXPLMs that are competitive with vanilla PLMs while achieving 2x/5x inference speedup with only a 1− 4% drop on a broad suite of tasks. 1

heading: 1 Introduction
text: Language models like ChatGPT (OpenAI, 2023), PaLM (Chowdhery et al., 2022), T5 (Raffel et al., 2020), and CM3 (Aghajanyan et al., 2022), 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
1Code + Models: https://github.com/ princeton-nlp/datamux-pretraining/.
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) (Havasi et al., 2021; Ramé et al., 2021; Murahari et al., 2022) 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 (Blumhagen et al., 2011; Akam and Kullmann, 2014; Pirschel and Kretzberg, 2016; Hong et al., 2016; Friedrich et al., 2004).
Murahari et al. (2022) 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-
ar X
iv :2
30 2.
12 44
1v 2
[ cs
.L G
] 2
2 M
ay 2
02 3
196
tively. MUX-PLMs, like other pre-trained language models, provide general model initialization that can be fine-tuned for any downstream task. We demonstrate the effectiveness and generality of our MUX-PLMs class of pre-trained models by training MUX-BERT and MUX-ELECTRA models, which are trained with pre-trained objectives adapted from BERT (Devlin et al., 2019) and ELECTRA (Clark et al., 2020) respectively, although in a MIMO fashion with data multiplexing.
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.

heading: 2 Related Work
text: Efficient Inference with Transformers Recent methods in NLP rely heavily on transfer learning through Transformer-based (Vaswani et al., 2017) language models trained on large text corpora using self-supervised objectives, such as autoregressive (Radford and Narasimhan, 2018) or masked language modeling (Devlin et al., 2019). Prior analysis on pre-training language models has observed power-law scaling of model performance with respect to model size (Kaplan et al., 2020), leading the community to develop ever-larger language models. It is also generally recognized that pre-trained language models are significantly overparameterized; effectively learning a subnetwork that utilizes only a relatively small number of their total parameters (Voita et al., 2019; Michel et al., 2019; Gordon et al., 2020; Prasanna et al., 2020).
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 (Liu et al., 2019; Wang et al., 2020; Lagunas et al., 2021; Xia et al., 2022; Yang et al., 2022), knowledge distillation (Hinton et al., 2015;
Sanh et al., 2019; Sun et al., 2020; Jiao et al., 2020; Yin et al., 2021), quantization (Zafrir et al., 2019; Shen et al., 2020), and data multiplexing (Murahari et al., 2022). 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 (Chen et al., 2020; Sanh et al., 2020) and pre-trained (Zafrir et al., 2021; Jiang et al., 2022) 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 (Havasi et al., 2021; Ramé et al., 2021; Murahari et al., 2022) train models using mixed-instance representations, i.e. Zhang et al. (2018), in order to obtain predictions for multiple instances simultaneously. Unlike efficiency methods, Havasi et al. (2021) and Ramé et al. (2021) 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 (Murahari et al., 2022), 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.

heading: 3 Methodology
text: We briefly introduce readers to the data multiplexing MIMO architecture (Murahari et al., 2022), which we denote T-MUX. We then detail our novel approach to train MUX-PLMs to yield highthroughput and performant language models.

heading: 3.1 T-MUX: Data multiplexing with Transformer
text: Data multiplexing allows parallel processing of multiple sequences with a single forward or backward pass through the model (M ) and requires
two crucial components, multiplexer, and demultiplexer.
Multiplexer The multiplexer module (MUX) combines an ordered set of multiple inputs – X1:N = ( x1, . . . ,xN ) into a single superimposed representation (xMUX). If xi ∈ Rd, the multiplexer is a transformation (MUX : RN×d → Rd) such that xMUX = MUX ( X1:N ) .
To ensure MUX is an order-preserving transformation,T-MUX samples a vector (vi ∈ Rd) from a standard multivariate Gaussian and applies the Hadamard product (element-wise multiplication) with the corresponding input representation (xi) before summing up vectors for all positions.
xMUX = MUX ( X1:N ) = 1
N
N∑
i=1
xi ⊙ vi
vi ∈ Rd ∼ N (0, I) (1)
The model processes the multiplexed representation and emits a multiplexed hidden state – hMUX = M ( xMUX ) . To multiplex Transformers’ sequenced inputs ( xi = ( xi1, . . . ,x i L )) of length L, T-MUX applies the same vi to all L positions of instance i.
xMUX = MUX ( X1:N ) =
( 1
N
N∑
i=1
xi1 ⊙ vi, . . . , 1
N
N∑
i=1
xiL ⊙ vi )
(2)
Demultiplexer A prediction needs to be made for each instance in X1:N , and T-MUX’s demultiplexer module (DeMUX) achieves this by separating the superimposed output (hMUX) into N output representations corresponding to the input (h1, . . . ,hN ).
hi = DeMUX ( hMUX,pi ) hij = DeMUX ( hMUXj ,p i ) (3)
The vector pi ∈ Rd is dynamically generated for each instance (i) with the help of a prefix that is added to the input and re-used for all positions in the instance. They add a prefixi to x
i, represented by the following pattern, where ϵi is a special token, and pi is set to be the output corresponding to token
i in the prefix.
prefix1 = [ϵ1, ϵpad, . . . , ϵpad]
prefix2 = [ϵpad, ϵ2, ϵpad, . . . , ϵpad]
· · · prefixN = [ϵpad, . . . , ϵpad, ϵN ]

heading: 3.2 MUX-PLMs: Data multiplexing for high-throughput language models
text: We propose MUX-PLMs, a class of highthroughput pre-trained Transformer-based language models trained in a MIMO fashion with data multiplexing. To demonstrate the viability
and the generality of this class of models, we pretrain Transformer models with objectives based on BERT and ELECTRA, to get MUX-BERT and MUX-ELECTRA respectively. MUX-PLMs are trained with our three stage training algorithm (Figure 1). Firstly, MUX-PLMs are trained with the token retrieval task in T-MUX, which is an autoencoding setup to decode all the tokens in the multiplexed input. This simple auto-encoding task is critical to prime the model for MIMO-style data multiplexing. The MUX-PLMs are then pre-trained with standard pre-training objectives but adapted to MIMO-fashioned training with data multiplexing. MUX-PLMs show significant throughput improvement over standard pre-trained LMs while matching their downstream task accuracies. Finally, MUX-PLMs, like other pre-trained language models, provide general model initialization that can be fine-tuned for any downstream task.
Contextual multiplexer T-MUX’s multiplexer multiplexes tokens independent of 1) tokens in the same position in other instances and 2) other tokens in the instance, which could lead to suboptimal representations. We, therefore, explore a contextual multiplexing scheme that aggregates context both from tokens in the same instance and tokens in the same position of other instances (Figure 3). We first use a single transformer layer TRANSctx to get contextual representations hictx = TRANSctx ( xi1, . . . ,x i L ) ) of length L. We apply a hadamard product with a multivariate gaussian vi to all L positions.
gictx = h i ctx ⊙ vi (4)
We generate multiplexed representations, xMUX, by applying another transformer layer TRANSinst across encoded representations from N instances
at each position from 1 to L. This is done by transposing gctx and applying TRANSinst.
xMUX = TRANSinst ( g⊤ctx )
(5)
RSA demultiplexer The demultiplexer in TMUX requires a prefix whose length scales linearly with the number of instances (N ), thus reducing the effective context length during pre-training, which degrades performance (Ainslie et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020). Furthermore, it decreases throughput during inference for large N because the model must process an extra prefix of length N for each of the N instances. To address these issues, we draw inspiration from the RSA cryptosystem (Rivest et al., 1978) to randomly initialize and learn N (private) key vectors (k1, . . . ,kN , ki ∈ Rd) which are keys that can be used to demultiplex the output representation (Figure 2).
hi = DeMUX ( hMUX,ki ) hij = DeMUX ( hMUXj ,k i ) (6)
Akin to RSA, vi and ki can be treated as the keys for multiplexing (encryption) and demultiplexing (decryption) while ensuring that, unlike T-MUX, the input sequence length does not change and thereby leading to an improvement in throughput. Importantly, this architecture ensures that the keys better transfer across the different stages of training as they are no longer conditioned on the input instances.

heading: 4 Experimental Setup
text: Datasets We pre-train all models on Wikipedia (Foundation) and Bookscorpus (Zhu et al., 2015). We evaluate on all datasets from the GLUE benchmark (Wang et al., 2018), which are CoLA (Warstadt et al., 2019), SST-2 (Socher et al., 2013), MRPC (Dolan and Brockett, 2005), QQP (qqp), STS-B (Cer et al., 2017), MNLI (Williams et al., 2018), QNLI (Wang et al., 2018), RTE (Wang et al., 2018), and WNLI (Levesque et al., 2012). We also evaluate on token classification tasks – named entity recognition (Sang and Meulder, 2003) and POS tagging (Grünewald et al., 2021). 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.
Model N GLUE Token ↗ Mean (std) Max Mean (std) Max
BERT 1 85.4 (0.0) 85.4 95.8 (0.0) 95.8 1.0×ELECTRA 82.1 (0.0) 82.1 95.3 (0.0) 95.3 1.0× T-MUX
2 60.4 (0.6) 61.8 81.4 (0.1) 81.5 1.9×
MUX-BERT‡ 82.5 (0.6) 83.4 95.2 (0.1) 95.4 2.0× MUX-ELEC‡ 82.5 (0.4) 83.1 95.0 (0.0) 95.1 2.0× T-MUX
5 59.7 (0.6) 60.6 81.3 (0.2) 81.5 4.4×
MUX-BERT‡ 80.3 (0.4) 80.9 93.6 (0.1) 93.6 4.9× MUX-ELEC‡ 79.8 (0.6) 80.5 93.4 (0.0) 93.5 4.9× T-MUX
10 58.1 (0.5) 59.1 79.7 (0.2) 80.0 8.4×
MUX-BERT‡ 77.8 (0.6) 78.8 91.6 (0.1) 91.8 9.8× MUX-ELEC‡ 78.2 (0.6) 79.0 91.7 (0.1) 91.8 9.7×
Table 1: Average GLUE and token-level classification scores for the BASE (L=12, H=768) configuration, across ELECTRA, BERT, and MUX-PLMs for N ∈ {1, 2, 5, 10}. ‡ indicates our models and ↗ indicates throughput increase w.r.t. to a vanilla BERTBASE model. All models are evaluated on 5 seeds with mean and max scores reported.
Models We experiment with ELECTRA and BERT pre-training objectives and present the pre-trained multiplexed models MUX-BERT and MUX-ELECTRA for N = 2, 5 and 10. To simplify training, we use a random generator to train MUX-ELECTRA models, presented as an ablation in Clark et al. (2020), instead of using a smaller masked LM. Except where otherwise noted, we do not use the contextual MUX module, but instead, use the RSA demultiplexing module. Refer to Appendix B and C for implementation details.
Baselines We run experiments to compare our models against T-MUX, from Murahari et al. (2022) and baseline PLMs - ELECTRA and BERT, across three different model configurations (SMALL, BASE, and LARGE). We also provide a comparison to results reported in recent PLM pruning and distillation papers in Table 2.
Multi-run evaluation We evaluate all models across 5 random seeds to reduce variance for smaller datasets which is caused by the randomized order in which we multiplex instances in the batch. We report both the average and maximum scores across seeds in Table 1 to understand the importance of ordering the multiplexed instances and report average across seeds for all other results.

heading: 5 Results

heading: 5.1 Comparing MUX-PLMs with PLMs and T-MUX
text: Table 1 shows that both MUX-BERT and MUXELECTRA outperform T-MUX at all levels of multiplexing (N ), with improvements between 12 and 20 points on GLUE and token-classification tasks respectively. Furthermore, MUX-PLMs’ efficient RSA-inspired demultiplexing method allows it to achieve faster throughput than T-MUX, increasing it by over 16% for N = 10.
Moreover, MUX-PLMs provide a significant boost in throughput (N times faster) when compared to PLMs, without a significant loss in performance. For example, MUX-ELECTRA (N = 2) is 0.4 points better and only 0.3 points worse than ELECTRA for GLUE and TOKEN tasks respectively, while being 2× faster. Similarly, MUX-BERT (N = 2) is within 3 and 0.6 points of BERT for GLUE and TOKEN tasks respectively, while being significantly faster. We also observe that as N increases, MUX-PLMs’ throughput is significantly better, though performance compared to PLMs can degrade. This is because a large N implies that MUX-PLMs must parallelly process more instances, thus having to share network parameters and activations with a larger number of instances, thus improving throughput and degrading performance. For example, the gap between ELECTRA and MUX-ELECTRA on TOKEN for N = 2 is 0.2 points and increases to 3.5 points for N = 10, which shows that N serves as a parameter to control the performance-throughput trade-off. We explore this further in Section 5.3 and Figure 4.

heading: 5.2 Comparing MUX-PLMs with recent model compression methods
text: We compare our MUX-PLM models with other efficient learning methods, such as pruning and distillation, in Table 2. Contrary to other methods, our vanilla MUX-PLMs achieve competitive performance and significant throughput improvement without additional unlabeled and task-specific data, and can be easily fine-tuned to any downstream task without any architectural modifications. For instance, when compared to DistilBERT, MUXBERT (N = 2) does 1 point worse on QNLI and 2 points better on QQP while being equally fast and not requiring any additional unlabeled data.
More broadly, methods like CoFi, AutoTinyBERT, and MobileBERT show that combining a
wide range of paradigms (for example, CoFi combines structured pruning and knowledge distillation, AutoTinyBERT combines knowledge distillation and neural architecture search, and MobileBERT combines knowledge distillation with novel architectural innovations) is a promising approach towards efficient high-performance models.
Towards this end, MUX-PLMs are complementary in both approach and motivation to these methods, and can evolve in tandem with existing efficiency methods. MUX-PLMs demonstrate the viability of MIMO architectures for PLMs, in addition to being complementary to existing approaches, and we hope that MIMO architectures develop and evolve with other efficiency approaches while leveraging the best of all efficiency methods.

heading: 5.3 Effect of varying model size
text: In this section, we show that our multiplexing techniques work on a host of model sizes and report results for MUX-BERT on three models sizes, SMALL, BASE, and LARGE for N = 2 (Table 3). We report results for other values of N in the appendix. MUX-BERT’s performance is close to that of BERT for all model sizes while having a significantly better throughput (the gap is less than 0.7 points for TOKEN tasks and 2.9 points for GLUE for close to twice the throughput). Multiplexing works effectively on all model sizes, with the drops with respect to BERT being 1.6 and 1.7 points on GLUE for SMALL and LARGE respectively. MUX-BERT’s throughput is always ≈ 2× that of BERT, which shows that a spectrum of MUX-PLM model sizes can be multiplexed during pre-training with competitive performance and with significantly higher throughput.
Pre-trained models typically have a performancecomputational efficiency trade-off, with larger models having better performance but worse compu-
tational efficiency. MUX-PLMs offers a similar trade-off, with large N leading to better throughput but lower performance. To understand this trade-off, we plot the performance and throughput of BERT and MUX-BERT for different model sizes and draw the pareto-optimal envelope (Figure 4). For any model on the envelope, no model has both better accuracy and throughput. Users would only choose models on the envelope because for every model within the envelope, there always exists a model on the envelope which has both better performance and throughput. We note that all multiplexed models lie either on or very close to the Pareto frontier, for both TOKEN and GLUE tasks. This suggests that given an accuracy threshold, MUX-PLM models will usually be faster than PLMs. For instance, if we wanted the highest throughput model with a performance ≥ 77% on GLUE, the optimal BERT model is the SMALL configuration with a throughput of 2815 (in/s), but for the MUX-BERT model would be the N = 2 with the SMALL configuration, achieving a significantly higher throughput of 5539 (in/s).

heading: 5.4 Ensembling MUX-PLMs
text: As opposed to feeding N different instances to MUX-PLMs to improve throughput, we consider an alternate setting where we feed the same instance N times and build an ensemble by averaging the N class logits to make a single prediction. We randomly permute the batch, after duplicating the instance N times, to prevent distribution shift. We use the BASE size models for N ∈ {2, 5, 10} for both MUX-BERT and MUX-ELECTRA (Table 4). The ensemble model does significantly better than the non-ensemble variant on both MNLI and QQP for all values of N (e.g., 1.6 and
0.9 points on N = 5 MUX-BERT for the two tasks). We note that the improvement over the nonensemble variant (∆) is better for higher N , due to the larger ensemble size. This result shows that nonensemble variants are faster but perform slightly worse, while the ensemble variant performs better but is slower. A spectrum of models lie between these two extremes, where only a fraction of the N multiplexed representations can be ensembled, allowing users to trade off performance and speed.

heading: 6 Analysis

heading: 6.1 Ablation study
text: We analyze multiplexing and demultiplexing components of MUX-PLMs and report the results in Table 5. We consider two variants, one which uses the prefix demultiplexing proposed in T-MUX instead of our proposed RSA-DeMUX and another which uses Contextual multiplexing instead of Noncontextual. We note that Variant 1, which uses prefix demultiplexing, performs worse than our MUXBERT, other than for N = 2. In fact, Variant 1 does not converge for TOKEN tasks for N = 5 and N = 10 and performs 1.7 and 1.2 points worse on GLUE when compared to MUX-BERT.
Variant 2 uses Contextual multiplexing which takes into account other tokens present in the instance and also tokens present in the same position of other instances. This variant performs better than Non-contextual for TOKEN tasks (almost over 1.7 points on TOKEN for N = 10) but performs worse for GLUE tasks. We believe that Contextual multiplexing’s better performance in TOKEN is because the model needs to make a prediction for every single position in the instance, which requires it to efficiently multiplex all token positions in the
output. However, for GLUE tasks, the model needs to make a prediction only for the [CLS] token, for which Non-contextual multiplexing suffices.

heading: 6.2 Muxology: Analyzing hidden representations of multiplexed models
text: To understand the nature of representations being learned by MUX-BERT models, we analyze the absolute value of activations and entropy of the attention distribution across all the layers of the Transformer encoder, averaged over the evaluation split of WikiText-103 (Merity et al., 2016) (Figure 5). We report this analysis for different values of N and for different model sizes.

heading: 1. Activation norms spike for MUX-BERT in the
text: last layer. Figure 5 (top) shows that activation norms spike in the last layer for multiplexed models. We believe this is because the model is preparing for demultiplexing and is packing information from all N instances, which makes the activations denser. We believe MUX-BERT has learned to efficiently encode multiple instances until the last layer where it needs to make independent predictions for them.

heading: 2. Entropy of the attention weights of MUXBERT is lower than BERT for higher layers.
text: Figure 5 (bottom) suggests that MUX-BERT tends to have lower entropy attention distributions on average as opposed to BERT for higher layers. This could be related to Deshpande and Narasimhan (2020)’s observation of pre-trained models having peaky attention distributions in the higher layers, with small irregularities. Since the model implicitly has to use the same attention distribution for all the multiplexed instances, the peaky distribution gets reinforced and is further corroborated by higher N having lower entropy in the final layer. We, therefore, believe that MUX-BERT has learned to create shared representations for multiple instances to effectively use the instance-independent attention distribution.

heading: 6.3 Effect of data sampling strategies during inference
text: During inference, our MUX-PLMs sample N instances uniformly at random from the evaluation set. However, other data-sampling strategies such as clustering similar instances based on wordoverlap could improve performance. We explore the effect of composition of N instances on the performance of MUX-PLMs in Table 6. For each model variant, we consider 5 random seeds

heading: 2 83.1 82.0 1.1 83.4 81.8 1.6

heading: 5 80.5 78.9 1.6 80.9 79.7 1.2

heading: 10 79.0 77.3 1.7 78.8 77.0 1.8
text: which can be viewed as lottery tickets (Frankle and Carbin, 2018). Since the random seed controls the composition of N instances, we measure the difference (∆) between the best-performing ticket and the worst-performing ticket and average the performance for all the GLUE tasks. ∆ is consistently greater than 1 point for all values of N for both MUX-ELECTRA and MUX-BERT, and illustrates the importance of the composition of N instances. An improved data sampling strategy could lead to improvements and we leave this to future work.

heading: 7 Conclusion
text: 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.

heading: Acknowledgements
text: We gratefully acknowledge support from Google AI Princeton, where Vishvak Murahari was a student researcher for part of this work, and the Samsung GRO program. We thank Mengzhou Xia, Jens Tuyls, and Tianyu Gao, with special thanks to our espresso machine.

heading: A Appendices

heading: B Pre-training Details
text: We report all pre-training related hyper-parameters in Table 7. We primarily use the HuggingFace Transformers implementations for BERT and ELECTRA based models. All pre-training experiments were run on 8 A100 GPUs with distributed training. We run a small hyper-parameter search over over two learning rates. All pre-trained models are primed with the token retrieval task introduced in Murahari et al. (2022). We train on the Wikipedia and Bookscorpus datasets for up to 10000 training steps with a learning rate of 1e− 4, and with a sequence length of 512.
For MUX-ELECTRA models, we don’t train a generator as in the original ELECTRA work, but only use uniform-random token replacement. This is similar to what was used in ablations in ELECTRA (Clark et al., 2020). The generator randomly replaces 15% of tokens in the input with other tokens in the vocabulary.

heading: C Fine-tuning Details
text: We report all the fine-tuning related hyper-parameters in Table 8. We run a small hyper-parameter search on the learning rate, batch size and number of training steps for different tasks. All models were trained with half-precision. We report numbers on the validation split. For GLUE tasks, we use the default metrics in Wang et al. (2018) and use F1 for the token-level tasks. All fine-tuning experiments were trained on 1 V100 GPU.
Speedup calculation For all models, we calculate throughput (samples/second) on a single V100 GPU and report throughput gains with respect to the BERTBASE model. We calculate throughput by averaging across 3 different trials (1 trial = 200 mini-batches) and use a batch size of 128 and a sequence length of 128 following prior work (Xia et al., 2022). We measure throughput for sequence-classification tasks on QQP and measure throughput for token-level classification tasks on named entity recognition.

heading: D Analysis details
text: D.1 Ensembling results setup We find that multiplexing the same instance by duplicating the instance N times leads to worse performance. This is likely because this input configuration is very out of distribution from what the multiplexed models are trained on. To address this, we randomly permute the instances in the batch after duplicating the instances N times. This ensures that the input to the multiplexer lies in a similar distribution to what the model was trained on.
D.2 Muxology setup To analyze the hidden states of pre-trained MUX-BERT models at different layers, we take the average absolute value of hidden states and every layer for both multiplexed and baseline models, across different configurations. To analyze the entropies of the attention distributions at different layers, we calculate the attention distribution across different attention heads for each position in the sequence length. To measure how peaky the attention distribution is likely to be, we calculate the entropies of the attention distributions at all positions and average across all the positions and across all the attention heads to get the average entropy for all layers. We conduct this analysis on WikiText-103 and average across all the samples in the evaluation split.
E Task performance breakdown for all variants

heading: Future Work
text: We hope our work inspires future research in MIMO architectures for PLMs as a complementary efficiency paradigm to existing approaches