MUX-PLMs: Data Multiplexing for High-throughput Language Models
Authors: Vishvak Murahari
Ameet Deshpande
Carlos E Jimenez
Izhak Shafran
Mingqiu Wang
Yuan Cao
Karthik R Narasimhan
Venue: EMNLP 2023 Findings (2023)
Keywords: Efficient Inference
Multi-input Multi-output architectures
Data Multiplexing
TL;DR: We develop MUX-PLMs, a class of high-throughput high performance language models trained with multi-input multi-output (MIMO) architectures.

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 MUX-PLMs that are competitive with vanilla PLMs while achieving 2x/5x inference speedup with only a 1-4 % drop on a broad suite of tasks.

MUX-PLMs: Data Multiplexing for High-throughput Language Models
Vishvak Murahari Izhak Shafran²
Ameet Deshpande'
Carlos E. Jimenez'
Mingqiu Wang
Yuan Cao
Karthik Narasimhan'
Princeton University
2 Google Brain
murahariecs. princeton.edu
Abstract
that can be deployed at scale to meet the quickly growing demand.
The widespread adoption of large language models such as ChatGPT and Bard has led unprecedented demand for these technolo- gies. The burgeoning cost of inference for ever- increasing model sizes coupled with hardware shortages has limited affordable access and poses pressing need for efficiency approaches geared towards high throughput and perfor- mance. Multi-input multi-output (MIMO) al- gorithms such as data multiplexing, offer promising solution with many-fold increase in throughput by performing inference for mul- tiple inputs at the cost single imput. Yet not currently performant enough to be deployed in modern systems. We change that by developing MUX-PI Ms,,s a class of deployable high throughput pre-trained lan- guage models (PLMs) trained with data mul- tiplexing, that can fine-tuned any down- stream task. Our novel multiplexing and demul- tiplexing modules proficiently entangle and dis- entangle imputs, and enable high-performance high throughput MUX-PLAN that are competi- with vanilla PL.Ms while achieving 2x/5x inference speedup with only 1 4% perfor- mance drop on a broad suite of tasks.
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-agnosicc paradigm that perform inference for multiple inputs simultaneously at the cost of a single input. This efficiency paradigm is natively geared towards yielding high-the thghthhroughput models, in addition to being complementary in ap- proach and motivation to current efficiency meth- ods such pruning, quantization, and distilla- tion. Interestingly, MIMO approaches are partly inspired by the human brain's extraordinary abil- ity to process multiple inputs and propagate in- formation at a high bandwidth with a few neural codes (Blumhagen et al., 2011; Akam and Kull- mann, 2014; Pirschel and Kretzberg, 2016; Hong et al., 2016; Friedrich et al., 2004).
these approaches
Murahari et al. (2022) introduced data multiplex- ing, a MIMO technique that can enable a many-fold increase in throughput. The method compresses N different instances into 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-thrrrrr models.
1 Introduction
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 manu- facturing and marketing. The proficiency of these tools has led to unprecedented demand for these models, with users facing frequent outages and ca- pacity limits. Additionally, ever-increasing model sizes and hardware shortages have constrained models' ability to handle very high load of re- quests, thus limiting large-scale affordable access these models. These trends bring into focus the need for high-the tigh-throughput, perffrraar,,,,, ef- ficient, and environmentally responsible models
In this work, introduce MUX-PLMs. a class of gightthroughut pre-trained language models trained a MIMO fashion with data multiplex- ing to process multiple inputs (2-10) simultane- ously 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 clas- sification and token classification tasks, respec- tively. MUX-PLMs, like other pre-trained lan- guage models, provide general model initializa- tion that can be fine-tuned for any downstream

---

the
chef walk
the
meal
the
chef
ate
the neal
B-LDC
D
0 R-NISC
MIX DEMUX 1,M1X
DEMUX
DEMUX
Transformer
Transformer
Transformer
MIX
X77
M12
MUX
MUX
MUX
the
chef
cooks
the
meal
MASK
chef
MASK
the
neal
Benin
won
the
first
leg
Token Retrieval
Pretraining
Finctuning (NER)
Figure I: Illustrating the training process for MUX-PLMs. MUX-PLM, Ms 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). MIMO style training with data multiplexing. The resulting MUX-BERS model, similar to standard PLMs, provides 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 correct prediction while orange indicates an incorrect prediction. Red highlighted tokens in the input indicate a position that has been masked.
task. We demonstrate the effectiveness and gener- ality of our MUX-PLM, class of pre-trained mod- cls by training JUX-BERY and EUX-ELECTRA models, which are trained with pre-trained objec- tives adapted from BERT (Devlin al., 2019) and ELECTRA (Clark et al., 2020) respectively, al- though in a MIMO fashion with data multiplexing.
architectures for large language models.
Related Work
Efficient Inference with Transformers Recent methods in NLP rely heavily on transfer learning through Transformer- (Vaswani et al., 2017) language models trained on large text corpora us- ing self-supervised objectives, such as autoregres- sive (Radford and Narasimhan, 2018) masked language modeling (Devlin et al., 2019). Prior analysis on pre-training language models has ob- served power-law scaling of model performance with respect to model size (Kaplan et al., 2020). leading the community to develop ever-larger lan- guage models. It is also generally recognized that pre-trained language models are significantly over- parameterized; effectively learning subnetwork that utilizes only a relatively small number of their total parameters (Voita al., 2019; Michel et al., 2019; Gordon et al., 2020; Prasanna al., 2020).
Our work is the first to introduce MIMO archi- tectures to PLMs. To enable this, we first develop new demultiplexing module, RSA-demux (Fig- ure 2), that randomly initializes and learns private key vectors to recover the multiple outputs from a multiplexed representation. Secondly, we intro- duce a new Contextual Multiplexer module (Fig- ure 3) that uses cross- insstnces attention blaeetion-based mechanism to aggregate context across the set of multiplexed instances, which seems to be particu- larly 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-thh-arr model compression tech- niques. We hope our work validates MIMO archi- tectures as a promising complementary direction to existing efficiency techniques. Consequently, we hope future research develops MIMO architec- tures in tandem with other efficiency approaches, leveraging the best of both paradigms. We will pub- licly release our models and code to promote open- source research on the next generation of MIMO
The ubiquity of pre-trained language models, their growing size, and over-parameterrizaton has inspired extensive research on improving inference efficiency. This includes methods such as struc- tured pruning (Liu et al., 2019; Wang et al., 2020; Lagunas et al., 2021; Xia 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 (Mura-

---

hari al., 2022). These approaches assume that PLMs are highly over-parametrized and attempt to approximate a large function by learning a smaller, compressed, version 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.
Multiplexer The multiplexer module (MUX) combines an ordered set of multiple inputs X1NN (xl...... (x'....,,""" x into a single superimposed
representation (x M If E R the multiplexer
a transformation (MUX: : Nxd -> IR:") such that x MUX = MUX (x X
To ensure MUX is an order-preserving samples vector (v' E R from standard multivariate Gaussian and applies the Hadamard product (element wise multiplica- tion) with the corresponding input representation (x') before summing up vectors for all positions.
transformation
Multi-input Multi-output Models While prun- ing, quantization, and distillation seek to reduce overparameterization by reducing models' repre- sentational capacity, other lines of work seek to exploit overparameteriizatoo in other ways. Multi- input Multi-output (MIMO) architectures (Havasi al., 2021: Ramé al., 2021; Murahari al.. 2022) train models using mixed-instane represen- tations, i.c. 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 single con- volutional model perform "ensembling for free" during inference. Data multiplexing, introduced in DataMUX (Murahari al., 2022), aims improve model efficiency by training Transformer models with mixeddinstance representations to perform si- multaneous inference for language tasks, thereby improving inference throughput many-fold. Cur- rently, MIMO architectures have only been used in a limited setting, achieving middling performance. Our work training PLMs with potent MIMO architecture, data multiplexing, dramatically im- proves inference throughput while preserving high accuracy for downstream tasks.
N
X MUX = MUX (XI:N)
= N
~
The model processes the multiplexed representa- tion and emits a multiplexed hidden state h M (x-UX). To multiplex Transformers' sequenced inputs (x* = (xf..... .. xxx) of length L, T-MUX applies the same v to all L. positions of instance
x MUX
MUX (X
Σ
x
N
Demultiplexer A prediction needs to be made for each instance in X1=N, and T-MUX's s demul- tiplexer module (DeMUX) achieves this by sep- arating the superimposed output (hMUX) into N output representations corresponding to the input (h" .....,hN),
h' DeMUX (h MUX p¹) h DeMUX (h MUX D')
3 Methodology
(3)
We briefly introduce readers to the data multiplex- ing MIMO architecture (Murahari et al., 2022), which we denote T-MUX. We then detail our novel approach train MUX-PLMS yield high- throughput and performant language models.
The vector p E IRd 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 prefix; to x represented by the following pattern, where e' is special token. and p is set to be the output corresponding to token i in the prefix.
3.1 T-MUX: Data multiplexing with Transformer
Data multiplexing allows parallel processing of multiple sequences with a single forward or back- ward pass through the model (M) and requires two crucial components, multiplexer, and demulti- plexer.
prefix* = _ ,paad cPadj prefix? [eaad 2² ePad
pady
prefix = [cPad epad cN

---

kerss
MLP
DEMUX
coples
h
Illustrating our novel RSA-inspired demultiplexing module. The module is initialized with N key vectors which are used to demultiplex the transformed multiplexed representations (h (UXU ). The keys are concatenated with & MUX and processed with MI.P generate the demultiplexed output representations (h1 h.).
Figure 2:
encoding setup to decode all the tokens in the mul- tiplexed input. This simple auto-encoding task is critical to prime the model for MIMO-style data multiplexing. The LUUX-PMMs are then pre-trained with standard pre-training objectives but adapted MIMO-fashioned training with data multiplex- ing. MUX-PLMs show significant throughput im- provement over standard pre-trained LMS while matching their downstream task accuracies. Fi- nally, MUX-PLMs, like other pre-trained language models, provide general model initialization that can be fine-tuned for any downstream task.
Cross-Tnstance Instance A Average
Cross-Instance Instance Transformer
XUX
Gaussian Haramard Projection
Transformer
the
chef
cooks
the
meal
Contextual multiplexer T-MUX's multiplexer multiplexes tokens independent of I) tokens in the same position in other instances and 2) other to- kens in the instance, which could lead to subop- timal 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
Figure 3: Illustrating our attention- multiplexing module. The module generates contextual representa- tions for instances 21 24 with Transformer layer and then applies hadamard product between the con- textual representations and the corresponding multivari- ate gaussian to generate instance-conditiones represen- tations. The final multiplexed representations ated by first applying another Transformer layer, which attends across the instances for all the positions the sequence, and then averaging across the instances.
gener-
single transformer layer TRANS to get contextual representations h TRANS (x) (x)....,xx xx,)) of length L. We apply a hadamard product with multivariate gaus- sian v to all positions.
3.2 MUX-PLMs: Data multiplexing for tigh-throughput language models
We propose MUX-PLMs, a class of high- throughput pre-trained Transformer-based lan- guage models trained in a MIMO fashion with data multiplexing. To demonstrate the viability and the generality of this class of models, we pre- train Transformer models with objectives based BERT and ELECTRA. get MUX-BERT and MUX-ELECTRA respectively. MUX-PLMs are trained with our three stage training algorithm (Fig- ure 1). Firstly, MUX-PLMs are trained with the token retrieval task in T-MUX, which is an auto-
geix
h
(4)
We generate multiplexed representations, MUX by applying another transformer layer TRANS inst across encoded representations from N instances at each position from 1 to L. This is done by trans- posing getx and applying TRANS insttt-
MUX
TRANS inst
8axx
(5)
RSA demultiplexer The demultiplexer T- MUX requires a prefix whose length scales lin-

---

GLUE
Token
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
Model
Mean (std)
Max
Mean (std)
Max
BERT
85.4 (0.0)
85.4
95.8 (0.0)
95.8
1
ELECTRA
82.1 (0.0)
82.1
95.3 (0.0)
95.3
1.0x
T-MUX
60.4 (0.6)
61.8
81.4 (0.1)
81.5
19x
MUX-BERT
82.5 (0.6)
95.2 (0.1)
2.0x
MUX-ELEC*
82.5 (0.4)
95.0 (0.0)
95.1
2.0%
WNLI (Levesque et al., 2012). We also evaluate on token classification tasks such as named entity recognition (Sang and Meulder, 2003) and POS tagging (Grünewald al., 2021). We don't report average over WNLI and CoLA these are the two smallest tasks in GLUE and we observe high variance across different seeds. All numbers are reported on the dev split. We report scores for all tasks in Appendix E.
T-MUX
59.7 (0.6)
60.6
81.3 (0.2)
81.5
4.4x
MUX-BERT
80.3 (0.4)
80.9
93.6 (0.1)
93.6
4.9x
MUX-ELECT
79.8 (0.6)
93.4 (0.0)
4.9%
T-MUX
58.1 (0.5)
59.1
(9.2)
80.0
8.4x
MUX-BERT
77.6 (0.6)
(0.1)
(0.1)
91.8
9.8x
MUX-ELEC
78.2 (0.6)
91.7 (0.1)
9.7x
Table Average GLUE and token-leve classifica- tion scores for the BASE (1=12, H=768) configura- tion, across ELECTRA, BERT, and MUX-PLMs for models and / indi-
N E {1, 2,5, 5, 10]. & indicates cates throughput increase w.r.t. to vanilla BERT 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- pre-trained multiplexed models MUX-BERT and for N 2, 5 and 10. To sim- plify training, we use a random generator 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 in- stead, use the RSA demultiplexing module, Refer to Appendix B and C for implementation details.
early with the number of instances (N), thus reduc- ing the effective context length during pre-training, which degrades performance (Ainslic et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020). Further- more, it decreases throughput during inference for large N because the model must process an tra 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 ran- domly initialize and learn N (private) key vectors (kı,... ky. K; E Rd) which are keys that can be used to demultiplex the output representation (Figure 2).
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 prun- ing and distillation papers in Table 2.
h' - DeMUX (h MUX k') h2 DeMUX (h MUX j, k*)
Multi-run evaluation We evaluate all models across 5 random seeds to reduce variance for smaller datasets which is caused by the random- ized order in which we muitiplex instances in the batch. We report both the average and maximum scores across seeds in Table to understand the importance of ordering the multiplexed instances and report average across seeds for all other results.
(6)
Akin to RSA, v; and k; can be treated as the keys for multiplexing (encryption) and demulti- plexing (decryption) while ensuring that, unlike T-MUX, the input sequence length does not change and thereby lcading to improvement in through- put. 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.
Results
5.1 MUX-PLMs outperform PLMs and T-MUX
Table shows that both MUX-BERT and MUX- ELECTRA 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' effi- cient RSA-inspired demultiplexing method allows
4 Experimental Setup
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

---

Model
QNLI
QQP
SST2
Config
Model
GLUE
Token
BERT
1.0%
90.5
91.2
91.7
BERT
80.6
94,0
x
Morcover, MUX-PLMs provide a significant boost in throughput (N times faster) when com- pared to PLMs, without 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 2x 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 because a large N implies that MUX-PLMs must parallelly process more instances, thus having to share network pa- rameters and activations with a larger number of instances, thus improving throughput and degrad- ing performance. For example, the gap between ELECTRA and MUX-ELECTRA on TOKEN for MUX-BERT (N=2)
2.0%
88.2
90.4
90.6
SMALL
T-MUX
59.5
81.8
8.7x
MUX-BERT 5N-5)
4.9%
85.6
88.8
86.9
MUX-BERT
79.0
93.3
11.5x
Use additional unlabelled or task-speciflc data
BERT
85.4
95.8
1.0x
DisiiBERRT
2.0x
89.2
88.5
91.3
BASE
T-MUX
60.4
81.4
1.9x
Block Pruning
2.7x
89.7
91.2
MUX-BERT
82.5
95.2
2.0x
Prune OFA
1.0x
90.3
91.2
91.5
Hyhrid Approuches
BERT
85.8
95.6
0.3x
MUX-BERT (N=2)+ + CoFi
4.0x
88.4
90.4
90.0
LARGE
T-MUX
61.7
80.9
0.6%
BERT,
2.0x
91.1
91.1
93.0
MUX-BERT
84.1
95.2
0.6x
CoFi
2.7x
91.3
AURTTInyBERT
4.3x
89.7
89.9
91.4
Table 3: Changing the model size for MUX-BERT
MobileBERT
2.3x
91.0
92.1
(N = 2) models. Across different model sizes, MUX-
BERT outperforms T-MUX and achieve higher through- put (indicated under / column). 5 our models.
Table 2: MUX-PLMs (variants are underlined) are com- plementary to existing efficiency methods, while being competitive standalone. Contrary existing methods, MUX-PLMs do not use additional unlabelled and task- specific data and can be easily fine-tuned for any down- stream task without architectural modifications. The inference speedups (/)) are reported against BERTBASE.
without additional unlabeled and task-specific data, and can be easily fine-tuned on any downstream task without any architectural modifications (unlike pruning). For instance, when compared to Distil- BERT, BUX-BBRTTAN (N 2) does 1 point worse on QNLI and points better on QQP while be- ing equally fast and not requiring any additional unlabeled data.
it to achieve faster throughput than T-MUX, in- creasing it by over 16% for N = 100
More broadly, methods like CoFi, Auto'Tin-- BERT, and MobileBERT show that combining qualitatively different efficiency paradigms is a promising approach towards efficient models. For example, CoFi combines structured pruning and knowledge distillation, and AutoTinyBERT com- bines knowledge distillation and neural architec- ture search. This places importance on discovering new efficiency paradigms that dovetail with exist- ing ones, rather than making incremental changes to current techniques.
Given that MIMO architectures like data mul- tiplexing propose novel direction towards effi- ciency, we demonstrate its complementary nature by combining it with pruning. We perform struc- tured pruning using CoFi (sparsity of 0.6) on our MUX-PLMs. The resulting pruned model is 2x faster than the original MUX-PLM while being performant. We believe that the best combination of efficiency techniques is largely hardware and use-case dependent and we hope that MIMO archi- tectures evolve in tandem with other approaches.
N is 0.2 points and increases 3.5 points for N 10, which shows that N serves as a parame- tradc-off. We explore this further in Section 5.3 and Figure 4.
ter to control the performance-throughe
5.2 Comparing MUX-PLMs with other model compression methods
5.3 Effect of varying model size In this section, show that our multiplexing techniques work host of model sizes and report results for MUX-BERT on three models sizes, SMALL, BASE, and LARGG for N = 2 (Ta- ble 3). We report results for other values of N in the appendix. MUX-BERT's performance is
We compare our MUX-PLM models with other efficient learning methods, such pruning and dis- tillation, in Table 2. Contrary to other methods, our vanilla MUX-PLMs achieve competitive per- formance and significant throughput improvement

---

close 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).
MNLI
QQP
Model
Mux (N)
Eins
Eins
Ens
A
80.6
81.2
0.6
90.4
90.8
0.4
MUX-BERT
5
77.2
78.8
88. 8
89.7
0.9
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 re- spectively. MUX-BERT's throughput is always 2x that of BERT, which shows that spectrum MUX-PLM model sizes can be multiplexed dur- ing pre-training with competitive performance and with significantly higher throughput.
2
80.3
80.8
+.0.5
90.6
90,9
MUX-RLEC
5
77.0
78.4
89.1
89.9
0.8
Table 4: Ensembling results for MUX-BERT and MUX- ELECTRA models for E {2,5, 10). Ens denotes Ensembling. Ensembling improves performance for all the models, with the gains increasing with increasing N. This suggests that the multiplexing approach can be nat- urally adapted to load-balancing applications, where the ensembling strategy can be changed based on demand.
87
has both better accuracy and throughput. Users would only choose models on the envelope be- cause 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 close the Pareto frontier, for both TOKEN and GLUE tasks. This suggests that given an aceuracy threshold, MUX-PLM models will usually 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 con- figuration with throughput of 2815 (in/s). but for the MUX-BERI model would be the N 2 with the SMALL configuration, achieving significantly higher throughput of 5539 (in/s).
Accuracy
81
N I
77
N 2
very
5=5
74
150
400
1100
3000
8100
22000
Throughput (instances/ / sec.)
95
B
92
N
N N=2
5,4 Ensembling MUX-PLMs
N
As opposed to feeding N different instances to MUX-PLMs to improve throughput, we consider an alternate setting where we feed the same in- stance N times and build an ensemble by averag- ing the N class logits make single prediction. We randomly permute the batch, after duplicating the instance N times, prevent distribution shift. We use the BASE size models for N e {2, 5, 10) for both MUX- BERT and MUX-ELECTRA (Ta- ble 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 non- ensemble variant (A) is better for higher N, due the larger ensemble size. This result shows that non- ensemble variants are faster but perform slightly worse, while the ensemble variant performs better but is slower. A spectrum of models lie between
85
3000 22000 Throughput (instances/ sec.)
Figure 4: (Top) BERT GLUE performance and through- put and (Bottom) BERT Token task performance and throughput, for N {1,2, 10) with the SMALL, BASE, and LARGE configurations (illustrated as S/B/L.). All multiplexed models lie either on or very close to Pareto frontier (shown in grey).
Pre-trained models typically have a performance- computational efficiency trade-off, with larger mod- els having better performance but worse compu- tational efficiency. MUX-PLMs offers similar trade-off, with large N leading to better through- put but lower performance. To understand this trade-off, we plot the performance and through- put of BERT and MUX-BERY for different model sizes and draw the pareto-optimal envelope (Fig- ure 4). For any model on the envelope, no model

---

Mux (N)
Model
Mux
Demux
GLUE
Token
MUX-ELECTRA
MUX-BERT
BERT
Non- con-contextual
RSA-DeMUX
82.5
Best ticket
Worst ticket
Best ticket
Worst ticket
2
Ablation 1
Non-contextun
Prefix
83.2
95.3
Ablation 2
Contextual
RSA-DeMUX
82.3
95.3
2
83.1
82.0
1.1
83.4
81.8
1.6
80.5
78.9
1.6
80.9
79.7
1.2
MUX BERT
Non. contextual
RSA-DeMUX
80.3
93.6
77.3
77.0
5
Ablation
Non. contextual
Prefix
78.6
38.2
Ablaticn
Consextual
RSA-DeMUX
768
94.2
MUX MUX-BERT
Nos- contexetllll
DSA.DeMUX
77.8
91.6
Table 6: We consider 5 random seeds for every model variant, which can be viewed lottery tickets as the seeds control the composition of instances. We present the difference between the worst and the best- performing ticket across GLUE tasks and regularly see - point difference.
Ablation
con-contextual
25.6
Ablation 2
Consextual
RSA-DeMUX
76.0
93.3
Table 5: Ablation analysis for MUX-BERT (base con- figuration) for N E (2, 5, 10]. Across most configura- tions, the prefix demultiplexing variant performs worse than our proposed approach and fails to converge for token-level tasks for . {5, 10) (underlined numbers). The new contextual multiplexing variant (Contextual) outperforms Non-contenn conteontettull on token-level tasks.
6.2 Effect of data sampling strategies during inference
During inference, our MUX-PLMs sample N in- stances uniformly at random from the evaluation set. However, other data-sampling strategies such as clustering similar instances based on word- overlap 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, consider 5 random seeds which be viewed as lottery tickets (Frankle and Carbin, 2018). Since the random seed controls the composition of instances, we measure the differ- ence (A) between the best- pest-perrorming ticket and the worst pooss-performing ticket and average the perfor- mance for all the GLUE tasks. A 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 improvements and we leave this to future work.
these two extremes, where only fraction of the N multiplexed representations be ensembled, allowing users trade off performance and speed.
6
Analysis
6.1 Ablation study
We analyze multiplexing and demultiplexing com- ponents of MUX-PLMs and report the results in Table 5. We consider two variants, one which uses the prefix demultiplexing proposed in T-MUX in- stead of our proposed RSA-DeMUX and another which uses Contextual multiplexing instead of Non- contextual. We note that Variant 1, which uses pre- fix demultiplexing, performs worse than our MUX- BERT, other than for N 2. In fact, Variant 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.
Conclusion
We introduce MUX-PLMs, class of high- throughput pre-trained language models trained with data multiplexing, a mulli-inpu multi-output (MIM) 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 infer- ence throughput. MUX-PLMs, similar to standard PLMs, can be fine-tuned on any downstream task to yield high-throughhut, high-performance mod- els. We hope our work inspires future research in MIMO architectures for PLMs complementary efficiency paradigm to existing approaches.
Variant 2 uses Contextual multiplexing which takes into account other tokens present in the in- stance and also tokens present in the same position 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 multi- plexing's ' better performance in TOKEN is because the model needs to make prediction for every single position in the instance, which requires it 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 con-contextual multiplexing suffices.

---

8 Limitations
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. 2020. The lottery ticket hypothesis for pre-trained bert networks. In Advances in Neural Information Processing Systems, volume 33, pages 15834-15846. Curran Associates, Inc.
Our MUX-PLMs class of -hhhhhhoughhuu high- performance pre-trained models demonstrates the efficacy of the MIMO paradigm for language mod- els. However, MUX-PLMs need MIMO-style train- ing for both pre-trat vre-trainin and fine-tuning. It would be more efficient to introduce MIMO-style training only during the fine-tuning stage as would al- low stakeholders to rapidly create high-throughput models from any off-the-shelve pre-trained model. We conducted initial experiments in this setting but faced issues getting models to converge and per- form well. This setting is incredibly challenging as shell pre-trained models have been trained to only process one instance ducing our novel multiplexing and demultiplexing modules would create a large distribution shift for the model. We hope that future work focuses on innovative solutions to address these limitations.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. 2022. Palm: Scaling language modeling with pathways. arXiv preprint a1Xiv:2204.02311.
Kevin Clark, Minh-Thang Luong, Quoe Le, and Christopher D Manning. 2020. Electra: Pre-training text encoders discriminators rather than gener- ators. In International Conference on Learning Representations.
time and intro-
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language under- standing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171-4186 Minncapolis, Minnesota. Asso- ciation for Computational Linguisties.
References
ruorrr Dataset- Re Seatt-Reeeassee lasee Question-Paine Ac-
Quora. data.
William B Dolan and Chris Brockett. 2005. Automati- cally constructing a corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP2005).
cessed: 2022-10-15.
Armen Aghajanyan, Bernie Huang, Candace Ross, Vladimir Karpukhin, Hu Xu, Naman Goyal, Dmytro Okhonko, Mandar Joshi, Gargi Ghosh, Mike Lewis, et al. 2022. Cm3: A causal masked multi- modal model of the internet. arXiv preprint arXiv:2201.00 .07520.
Wikimedia Foundation. Wikipedia.
Jonathan Frankle and Michael Carbin. 2018. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations
Joshua Ainslie, Santiago Ontañón, Chris Alberti, Philip Pham, Anirudh Ravula, and Sumit Sanghai. 2020. ETC: encoding long and structured data transform- ers. CoRR, ahs/2000.08488.
Rainer W Friedrich, Christopher J Habermann, and Gilles Laurent. 2004. Multiplexing using synchrony in the zebrafish olfactory bulb. Nature neuroscience. 7(8):862-871.
Thomas Akam and Dimitri M Kullmann. 2014. Os- cillatory multiplexing of population codes for selec- tive communication in the mammalian brain. Nature Reviews Neuroscience, 11-122.
Mitchell Gordon, Kevin Duh, and Nicholas Andrews. 2020. Compressing BERT: Studying the effects of weight pruning on transfer learning. In Proceedings of the 5th Workshop on Representation Learning for NLP, pages 143-155, Online. Association for Com- putational Linguistics.
Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020. Longformer: The doccmentt transformer. arXiv preprint arXiv:2004.051500
Francisca Blumhagen, Peixin Zhu, Jennifer Shum, Yan- Ping Zhang Schärer, Emre Yaksi, Karl Deisseroth, and Rainer W Friedrich. 2011. Neuronal filter- ing of multiplexed odour representations. Nature, 479(7374):493-498
Stefan Grünewald, Prisca Piccirilli, and Annemarie Friedrich. 2021. Coordinate constructions in glish enhanced universal dependencies: Analysis and computational modeling. In Proceedings of the 16th Conference of the European Chapter of the Association Computational Linguistics: Main Volume, pages 795-809.
Daniel Cer, Mona Diab. Eneko Agirre, lñigo Lopez- Gazpio, and Liiaaa Specia. 2017. Semeval-2017 task I: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the 11th International Workshop on Semantic Evaluation (SemEval-2011)) pages 1-14.
Marton Havasi, Rodolphe Jenatton, Stanislav Fort, Jeremiah Zhe Liu, Jasper Snoek, Balaji Lakshmi- narayanan, Andrew Mingho Dai, and Dustin Tran. 2021. Training independent subnetworks for robust

---

prediction. International Conference on Learning Representations.
Alec Radford and Karthik Narasimhan. 2018. Im- proving language understanding by generative pre- training.
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015. Distilling the knowledge in a neural network. ArXiv, 150.05333..
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020. Exploring the lim- its of transfer learning with unified text-to-text transformer. Journal Machine Learning Research, 21:1-67.
Sungho Hong, Mario Negrello, Mare Junker, Aleksan- dra Smilgin, Peter Thier, and Erik De Schutter. 2016. Multiplexed coding by cerehellar purkinje neurons.
Elife, 5:e13810.
Alexandre Ramé, Rémy Sun, and Matthieu Cord. 2021. Mixmo: Mixing multiple imputs for multiple out- puts via deep subnetworks. In Proceedings of the IEFE/CVF International Conference on Computer Vision (ICCV), pages 823-833.
Ting Jiang. Deqing Wang, and Fuzhen Zhuang. 2022. Pruning pre-trained language models without fine- tuning.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang. Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2020. TinyBERT: Distilling BERT for natural language un- derstanding. pages 4173-4174.
Ronald L Rivest, Adi Shamir, and Leonard Adleman. 1978. method for obtaining digital signatures and public-key cryptosystems. Communications of the ACM, 21(2):120-126
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models.
Erik Tjong Kim Sang and Fien De Meulder. 2003. In- troduction to the conll-200 shared task: Language- independent named entity recognition. In CoNLL.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. DistilBERT, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint 1910v:1910.01 108.
François Lagunas, Ella Charlaix, Victor Sanh, and Alexander M Rush. 2021. Block pruning for faster transformers, arXiv preprint arXiv:2109004836
Hector Levesque, Frnest Davis, and Leora Morgen- ster.. 2012. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning.
Victor Sanh, Thomas Wolf, and Alexander Rush. 2020. Movement pruning: Adaptive sparsity by fine-tuning. In Advances in Neural Information Processing Systems, volume 33, pages 20389. Curran Associates, Inc.
Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang. and Trevor Darrell. 2019. Rethinking the value of network pruning. In International Conference on Learning Representations.
Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W. Mahoney, and Kurt Keutzer. 2020. Q-bert: Hessian based ultra low pre- cision quantization of bert. In AAAL
Paul Michel, Omer Levy, and Graham Neubig. 2019. Are sixteen heads really better than one? In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. 2013. Recursive deep models for semantic compositionality over sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pages 1631 1642.
Vishvak Murahari, Carlos E Jimenez, Runzhe Yang. and Karthik R Narasimhan. 2022. DataMUX: Data multiplexing for neural networks. In Thirty-Sixte Conference on Neural Information Processing Systems.
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. 2020. MoBileBERT. compact task-agnostic bert for resource-limited de- vices. pages 2158-2170.
OpenAl. 2023. Introducing chatgpt.
Friederice Pirschel and Jutta Kretzherg. 2016. Mul- tiplexed population coding of stimulus proper- ties by leech mechanosensory cells. Journal of Neuroscience, 36(13):3636 3647.
Ashish Vaswani. Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In NIPS.
Sai Prasanna, Anna Rogers, and Anna Rumshisky. 2020. When BERT Plays the Lottery, All Tickets Are Win- ning. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing Online. Association for Computational Linguistics.
Elena Voita. David Talbot, Fedor Moiseev, Rico Sen- nrich, and Ivan Titov. 2019. Analyzing multi-head self-attention: Specialized heads do the heavy lift- ing, the rest can pruned. Proceedings of the 57th Annual Meeting of the Association for
(EMNLP), pages

---

Computational Linguisties, pages 5797-5808, Flo- rence, Italy Association for Computational Linguis- tics.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhut- dinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. 2015. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In The IEEE International Conference on Computer Vision (ICCV).
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. 2018. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference Learning Representations.
Ziheng Wang, Jeremy Wohlwend, and Tao Lei. 2020. Structured pruning of large language models. In Proceedings of the 2020 Conference on Empirical Methods Natural Language Processing (EMNLP). Association for Computational Linguistics.
Alex Warstadt, Amanpreet Singh, and Samuel Bow- man. 2019. Neural network acceptability judgments. Transactions of the Association for Computational Linguistics, 7:625-641.
Adina Williams, Nikita Nangia, and Samuel R Bow- man. 2018. A broad-coverage challenge corpus for sentence understanding through inference. In NAACL-HLT.
Mengzhou Xia, Zexuan Zhong, and Danqi Chen. 2022. Structured pruming learns compact and accurate mod- In Association for Computational Linguistics (ACL).
Ziqing Yang, Yiming Cui, and Zhigang Chen. 2022. TextPruner: A model pruning toolkit for pre- trained language models. the 60th Annual Meeting of the Association for Computational Linguisties: System Demonstrations, pages 35 Dublin, Ireland. Association for Com- putational Linguistics.
Proceedings of
Yichun Yin, Cheng Chen, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. 2021. BEERTTTyEERT: Automatic hyper- parameter optimization for efficient pre-trained language models. pages 5146-5157.
Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. 2019. Q8bert: Quantized 8bit bert. 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing NeurlPS Edition (EMC2-NIPS), pages 36-39.
Ofrr Zafrir, Ariel Larey, Guy Boudoukh, Haihao Shen, and Moshe Wasserblat, 2021. Prune once for all: Sparse pre-trained language models. arXiv preprint arXiv:2111.05554
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago On- tamon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020. Big bird: Transformers for longer sequences. Advances Neural Information Processing Systems, 33: 17283-17297
Hongyi Zhang, Moustapha Cisse, Yann Dauphin, and David Lopez-Paz. 2018. mixup: Beyond empirical risk minimization. International Conference on Learning Representations.

---

Hyperparameter
MUX-BERT
MUX-ELECTRA
SMALL
BASE
LARGE
BASE
Number of layers
12
24
Hidden Size
512
768
1024
768
FFN intermediate hidden size
2048
3072
4096
3072
Attention heads
8
12
16
12
Attention head size
64
64
64
64
Mask percent
15
15
15
N/A
Learning Rate Decay
Linear
Linear
Linear
Linear
Warmup steps
10000
10000
10000
10000
Learning Rate
[le-4, 5c-5]
[le-4, Se-5]
[lc-4, 5c-5]
[1c-4, 5c-5]
Adam
le-6
le-6
1e-6
1e-6
Adam $1
0.9
0.9
0.9
0.9
Adam $2
0.999
0.999
0.999
0.999
Attention Dropout
0.1
0.1
0.1
0.1
Dropout
0.1
0.1
0.1
0.1
Batch Size
256
256
256
256
Sequence Length
512
512
512
512
Train Steps
IM
1M
1M
IM
Table 7: Pre-train hyper-parameters for MUX-BERT and MUX-ELECTRA models. We only report results for the Base configuration MUX-ELECTRA models.
A Appendices
Pre-training Details
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 small hyper-parameter search over over two learning rates. All pre-trained models are primed with the token retrieval task introduced in Murahari al. (2022). We train on the Wikipedia and Bookscorpus datasets for up 10000 training steps with learning rate of le 4, and with a sequence length of 512.
For MUX-ELECTRA models, we don't it 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.
C Fine-tuning Details
We report all the fine- fine-tuning related soper-parameters in Table 8. We run small hyper-parameet search on the learning rate, batch size and number of training steps for different tasks. All models were trained with half prefisrecision. We report numbers on the validation split. For GLUE tasks, we use the default metrics in Wang et al. (2018) and use Fl for the token-level tasks. All fine-tuning experiments were trained on 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 BERT base model. We calculate throughput by averaging across 3 different trials trial = 200 mini-batches) and use batch size of 128 and sequence length of 128 following prior work (Xia et al., 2022). We measure throughput for sequence-classinn classffcction tasks on QQP and measure throughput for token-level classification tasks on named entity recognition.

---

Hyperparameter
Value
Learning Rate
[2e-5, 5e-5]
Adam
le-8
Adam $1
0.9
Adam /32
0.999
Learning rate decay
Linear
Warmup fraction
0.1
Attention Dropout
0.1
Dropout
0.1
Weight Decay
0
Batch Size
[32, 128] for SMALL/ BASE, [16, 64] for LARGE
Train Steps
2000 for RTE and WNLI
10000 for MRPC, COLA and STSB
20000 for NER, SST2, QNLI and POS
[20000, 100000] for MNLI and QQP
Sequence Length
128
Table 8: Fine-tune hyperparameters
D Analysis details
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
To address this, randomly permute the instances in the batch after duplicating the instances N times. This ensures that the input the multiplexer lies in similar distribution what the model was trained on.
D.2 Muxology setup
To analyze the hidden states of pre-trained MUX-BERT models different layers, we take the average absolute value of hidden states and every layer for both multiplcxed and bascline 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 WikiText- 103 and average across all the samples in the evaluation split.
E Task performance breakdown for all variants

---

Model Size
N
MNII
QQP
QNLI
MRPC
WNLI
STSB
RTE
SST2
COL.A
GLUE
GLUP
WNLL COLA
1
77.86100
88.99.00
84.00400
77.20400
56,34,00
84,25500
62.45.00
88.8
43.48
73.77
80.59
2
75.09
88 18
84.31.02
19775
50.99
u2.65
55,52.14
87.04
30.64.1.7
70.54
79.03
SMALL
5
70.
70.50401
81,23
74.26.10
54,68
79.90.03
58.5621.
32.57
12.78
66.76
76.20
10
61.98,01
80.85
63,47
70.69
56.62.43
36.93
53,57
80.39
1.10
56.18
63.98
1 84.24200
90.4400 90n
27.75
26.34
89.18
63.18200
91.74.000
58.79 20.0
79.22
85.40
2
80.59:
90.36101
88.17
83.77
50.70amm
85,84.0.1
19.19
90.62:0.6
55.61
75,98
82.51
RASE
5
77. 18
88, 1
85 18
80 10a0
53,52495
24,28
59.13
N6. 88
12.35
69.75
80.28
10
73.62.p5
86.44gu
82,08
28.63
52.68
81.62,02
58.27
83.44
0.00
90
1 85.79
91,46100
92.29:00
83,82.00
56.34.00
89,53,00
66.06.00
91,40
57.79 2o.0
79,39
85,76
2
90.85 83
90.66
s4.90.0s
56.34,00
59.2 s8.22
91.38.n.4
57.89 g1.5
78.08
84.06
LARGIE
79,55
89.37a01
87.41
85.77a11
54,93100
85, 85.86
57,26220
88 88.65
46.601299
74.83
81.70
10
35.
63. 63.18,00
50,54,00
68.38100
56,90,53
82.81
52.13
50.92.00
1.87
51.35
57.63
Table 9: We show the full GLUE results for MUX-BERT, We report the mean accuracy and standard deviation over 5 seeds. Extrema and values within their standard deviation are emphasized for each model size.
Model Size
N
MNLI
QQP
QNLI
MRPC
WNLI
STSB
RTE
SS12
COLA
GLUE
GLUE WNUL COLA
77.86
88.99
84,00
77.70
56.34
84.25
62.45
88.88
43.48
73.77
80.59
2
75.21
89.01
84.61
80.64
61.97
82.97
58.12
87.84
33.08
72.61
79.77
SMALL
5
70,66
86.46
81.60
75.74
61.97
80.24
60.65
83,49
15,57
68.49
76.98
10
62.17
80.93
63.85
71.81
63.38
38.20
55.96
80.96
2.63
57.77
64.84
84.24
91.19
90.54
87.75
56.34
89.18
63.18
91.74
58.79
79.22
85.40
2
80.82
90.47
88,28
86.03
66,20
86.06
60.65
91.51
56.93
78.55
83.40
BASE
77.66
88.89
85.70
81.13
59.15
84.47
60.65
87.50
15.79
71.22
80.86
10
74.04
87.03
82,45
79.41
63.38
81,89
62.45
84,29
0.00
68.33
78.79
85.79
91.46
92.29
83.82
56.34
89,53
66.06
91.40
57.79
79.39
85.76
83.40
90,94
90.96
86.27
56.34
88.50
60.29
91.86
60.50
78.78
84.60
LARGE
5
79.69
89.43
87.81
84.80
57.75
86.49
60.65
89.45
47.56
75.96
82.62
10
35.46
63.18
50.89
68.38
61.97
83.04
55.60
50.92
7.55
53.00
58.21
Table 10: We show the full GLUE results for MUX-BERT. We report the maximum accuracy over 5 seeds. Extrema are emphasized.
N
MNLI
QQP
QNLI
MRPC
WNLI
STSB
RTE
SST2
COLA
GLUE
GLUE WNLL COLA
81.49100
90.73100
89.73_o0
75.98_00
$6.34_0.0
87.73_0.0
$7.76_..0
91.51
56.7910.0
76.45
82.13
2
90.58=0.1
88.39_o.
83.73=0.7
$7.18-2.1
86.80=o.1
58.77
88. 38.65
51.92+177
76.26
82.46
76.99103
89.08100
85.40_0.
g0.25
24.5
84,27
57.26
a5.09.1.
26.8911.2
71.35
79.76
10
74.62402
87.63
82,70_02
77.89
50,99
81.96
$9.86_2.1
82. 22.7
27.76
69.57
78.20
Table 11: We show the full GLUE results for MUX-ELECTRABASE We report the mean accuracy and standard
deviation over 5 seeds. Extrema and values within their standard deviation are emphasized for each model size.
Retreival Rate
MNLI
QQP
QNI.
MRPC
WNLI
STSB
RTF
SST2
COLA
GLUE
GLUE
WNLL
0.0
83 01.33.c
90.85
90.66.00
84.90_..
$6.34.00
88.22
59.21
91.38
$7.89_1.5 15
78.08
84.06
0.1
83.55
90.90.a1
90.58.002
85,49_1.1
56.34.00
88.28
57.76_14
90.69403
59.36.3.4
78.11
$3.89
2
0.2
83.50
90.96.01
90.69
84,95
56.34.nm
88.28 uo.3
58.34.1.
90.69105
59.17...
78.10
83.92
0.5
83.41.02
90.91
90:47 1n.
85.25 .03
56.34.nm
88.02
59.35
89.52,08
59.41:3
78.08
K3.85
0.0
79.55 usu2
10.
87.41.02
83.77...1
14.
16.8610.1
57.26
18.65107
46.66101
74,83
N1.70
0.1
79.49
89.34.m
N7.25
81 $1.81 113
53.24116
85.80
55.60
88.19
47.60
74.26
B1.07
0.2
79.37-60
89,42.01
87.23
82,40_11
$4,93
85.85
58.38
87.84403
43.58
74.00
$1.07
0.5
79 79.24-601
89.30401
87.2140.
82.06-21.
$6,34an0
85.97
52,27=40
88.58400
47.01
34.22
80.66
0.0
35.45 60.0
63.18.on
50.54.00
68.38 20
26.90
82.81
52.13
50.92400
1.87_16
51.35
$7.63
0.1
38.45 66.
63.18400
50.65p02
68.38_.0
54.93250
4.45
51,48
50.92400
1.341a
42.31
46.36
10
0.2
35.45 10.0
63.18
50.21105
68.43
54.66
0,2311.
52,35
51.72
0.29
0.29
41.83
45.94
0.5
35.45.co
63.18.00
50.43
_..0
56.06.n
82.01.a4
52.71
50.922m0
1.5
51.18
57.58
Table 12: GLUE results for MUX-BERT LARGE when using a retrieval auxiliary objective during MLM pretraining with different trade-off rates to the MLM objective. We report the average accuracy over 5 seeds. Extrema and values within their standard deviation are emphasized for each value of N.

---

Mux Strategy
MNLI
QQP
QNLI
MRPC
WNLI
STSB
RTE
SST2
COLA
GLUE
GLUE_
MUX-BERT
80.59an
90.36
88.17-01
83.77e14
50,7047o
85.84
58. 1921
90.62=34
55.61=10
75.98
82,51
DataMUX
$1.64n22
90.67
Lo.s
102
14.17.04
56.34.n0
86.36.82
60.87.07
53.74.1
Attention
81,32402
90.65
200
88.77
80888
56.34200
86.25
$6.90 21
91.06
47.15
75.48
82.26
MUX-BERT
77.18in2
88 88.79
Loas
85.58 m1
80.10.oo
53.52
84.28 .oa
59.13
86.88 .ot
12.33 124
69.75
80.28
DataMIIX
36.32.a
89.13g0.1
84.22_1
78. 18.38
59,44,a5
81.78
54,15
86.17-2.4
28.32
70.88
78.59
Atlention
77.16.n
88.71 184
84.33
10.49
54.08
80.37.o.
$4,4122
81,95
34.67.12
69.58
76.78
73.62.0a
86.94
28
52.68.a0
58.27
83.44
0.00a00
10
DataMUX
72.74.m1
87.88
82.28
77. 30
56.34.nm
78.07.044
55,31
82.36..0
13,56150
67.32
76.56
Attention
71.83,n2
88.00 Lo.0
81,46_02
73.53
53.24
82.95.aa
52.71
81.28.o4
32.84.n6
68.63
75.97
Table 13: GLUE results for MUX-BERT BASE using alternative muliiplexing-demallileexxgg strategies. We report the average accuracy over 5 seeds. Extrema and values within their standard deviation are emphasized for each value N.
Model
MNLI
QQP
QNLI
MRPC
WNLI
STSB
S812
COLA
GLUE
GLUE GLUE_WNLL.OLL COLA
61.48102
80.33g00
60.05.02
68.43108
56.34.00
15.02:04
51.12
79.75
8.22
53.42
$9.45
SMALL
5
58.35
77.50
57.17
68.38
56.34.0n
11.31,03
51.70
77.78
6.02
51.62
$7.46
53.63402
77 20.
51,222aa
68.38.00
57.4649a
12.40
52.35
50,92,0.0
0.00g00
63,29403
81.42,01
60.38
68, 68,38402
56,9045a
17.65210
51.19417
80.78
9.62
54.40
60.44
BASE
5
60.67.12
19142.001
59,77
69.61.nx
53.80173
14.92
52.71
81.15.08
10.35
53.60
59.75
59.07402
1
57.99
68.38400
60.28
11,83
53.07
78.35
7.40g17
64.64402
82.10ga1
60.21
69.95
56.34.00
21.62.0.4
$2,71
80.34
8.72
55.18
61.65
LARGE
5
60.78103
78.561n1
60.19
69,51
56.34100
17.33
52,71
78.28.0.4
10.63
53.81
39.62
48.79406
68.41.01
55.76448
68.58406
58.59
8.38
$4.95
64.82.10
3.48,39
Table 14: GLUE results for T-MUX with the original training recipe and implementation from Murahari al. (2022). We report the average accuracy and standard deviation over 5 seeds. Extrema and values within their standard deviation are emphasized for each model size.


Reviews:
[{"review_id": "WcgO7Ka5N3", "reviewer": "Reviewer_dyvA", "summary": "", "questions": "- What is the primary role of the Hadamard product in the MUX module? Can it be understood as a type of identifier for each sample, similar to how position encoding is used in the transformer?\n- The pre-training of MUX-PLM involves two stages. Does this require a longer training period compared to the standard PLM training? Could you provide a comparison of the training computational times?\n- Are the poor results on WNLI and CoLA due solely to the small size of the data sets, or is it because MUX-PLM is not suitable for these types of tasks? Could you sample a subset from data sets like SST-2 to confirm whether MUX-PLM performs poorly on all small data sets?\n- What is the model size of T-MUX in Table 1? Is it the same size as MUX-BERT and MUX-ELECTRA?\n- How is the DeMUX function calculated in Line 288?\n- Section 5.4 analyzes the results of ensemble learning. Although the performance has improved compared to without ensemble learning, it is still not as good as the BERT model with equivalent computational cost. What could be the reason for this?\n", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your insightful comments and we thank your support in bringing the MIMO paradigm to our community. \n> Additionally, MUX-PLM may result in higher training costs. \n\n> The pre-training of MUX-PLM involves two stages. Does this require a longer training period compared to the standard PLM training? Could you provide a comparison of the training computational times?\n\n**Both the baseline PLMs and MUX-PLMs are pre-trained with the same number of optimization steps** (Please refer to Table 7 for other details). We also use the same hyper-parameter sweep for fine-tuning MUX-PLMs and PLMs (Please refer to Table 8).\n\n\nThe first stage (“Token Retrieval” in Figure 1) is trained for 10K optimization steps. We pre-train all models for 1M optimization steps. The first stage is a couple of orders of magnitude cheaper than pre-training and hardly leads to any overhead. Appendix B contains more details.\n\n\n> Achieving different acceleration ratios requires pre-training different MUX-PLMs while existing compression methods can obtain models with different acceleration ratios through fine-tuning a certain PLM.\n\n\nAs noted by all the reviewers, MUX-PLMs and MIMO architectures more generally are complementary to existing model compression methods. \n\n\nWe demonstrate this in Table 2 by combining MUX-BERT (N = 2) with CoFi, a SOTA structured pruning approach. The resulting model is 2x faster than the MUX-BERT model and 4x faster than the baseline while being equally performant. This demonstrates that the **acceleration ratios for both methods multiply without a significant drop in performance**.\n\n\nWe believe MIMO architectures will continue to evolve in tandem with the current and future efficiency paradigms to leverage the multiplication of these acceleration ratios to get faster and more performant models.\n\n\nFinally, MUX-PLMs are the first pre-trained MIMO style models in the community and we believe the community will make training of MUX-PLMs more efficient, where potentially existing PLMs are quickly adapted to yield MUX-PLMs.\n\n> The article introduces the Contextual MUX module, but it does not perform as well as the original MUX module on the GLUE benchmark.\n\nWe introduce the Contextual MUX module not only to improve performance but also to understand how design choices in the multiplexing module affect performance. Our experiments indicate that cleverly aggregating context does improve performance for token-level tasks, which require the model to make predictions for every single position in the instance, which consequently requires it to multiplex all token positions in the output efficiently. This might not be as critical for sequence classification tasks which require a single prediction for an entire sequence.\n\n\nWe will improve the writing to reflect this.\n\n\n> What is the primary role of the Hadamard product in the MUX module? Can it be understood as a type of identifier for each sample, similar to how position encoding is used in the transformer?\n\n\nYour understanding is very accurate, it is indeed a type of identifier for each sample that helps us generate a compressed and multiplexed representation for the set of N-independent samples.\n\n\n> Are the poor results on WNLI and CoLA due solely to the small size of the data sets, or is it because MUX-PLM is not suitable for these types of tasks? Could you sample a subset from data sets like SST-2 to confirm whether MUX-PLM performs poorly on all small data sets?\n\n\nWe would like to clarify that the performance on WNLI and CoLA are *not poor but instead have high variance*. \n\n\nTables 9 - 14 in Appendix E contain performance metrics and error bars for each individual task for various model configurations. Consider the WNLI column in Table 9, Our MUX-PLM models perform very well compared to the non-multiplexed baseline. However, the variance across the 5 seeds is very high.\n\n\n> What is the model size of T-MUX in Table 1? Is it the same size as MUX-BERT and MUX-ELECTRA?\n\n\nYes, it is the exact same size as MUX-BERT and MUX-ELECTRA. We have listed the size and the other hyper-parameter details in Table 7.\n\n\n> Section 5.4 analyzes the results of ensemble learning. Although the performance has improved compared to without ensemble learning, it is still not as good as the BERT model with equivalent computational cost. What could be the reason for this?\n\n\nA good ensemble involves a diverse set of functions that are accurate on different parts of the input space. Here, the function being ensembled is the model output for a particular index of the set of multiplexing functions. Therefore, we can think of it as an ensemble of N functions,  f_1, .. f_N, where f_i is the output of MUX-BERT for the ith instance. These functions would be incentivized to be similar given the training objectives and this consequently would constitute a weaker ensemble.\n\n\nWe introduced the ensemble primarily to showcase that one could dynamically interpolate between accuracy and throughput by controlling which and how many instances are ensembled for a prediction. We discuss this further at the end of section 5.4.\n\n\n> How is the DeMUX function calculated in Line 288?\n\n\nThe DEMUX function is instantiated with a standard feedforward network containing linear layers, layer norm layers, gelu activation, and some skip connections. We will improve the writing to clarify this.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper is well-written and easy to read.\n- Research in the field of MIMO holds significant potential and could serve as an important supplement to existing acceleration methods.\n- The experiments have demonstrated that the impact on performance is minimal when throughput is increased.", "reasons_to_reject": "- Some of the contributions made by the paper are not fully substantiated by the experiments. The article introduces the Contextual MUX module, but it does not perform as well as the original MUX module on the GLUE benchmark. Furthermore, the Contextual MUX was not employed in the main experiment.\n- Compared to other acceleration methods, MUX-PLM does not demonstrate strong competitiveness in terms of performance. It may not yet be adopted by users from a practical perspective, indicating room for further optimization. Additionally, MUX-PLM may result in higher training costs. Achieving different acceleration ratios requires **pre-training** different MUX-PLMs, while existing compression methods can obtain models with different acceleration ratios through **fine-tuning** a certain PLM.", "questions_to_authors": "- What is the primary role of the Hadamard product in the MUX module? Can it be understood as a type of identifier for each sample, similar to how position encoding is used in the transformer?\n- The pre-training of MUX-PLM involves two stages. Does this require a longer training period compared to the standard PLM training? Could you provide a comparison of the training computational times?\n- Are the poor results on WNLI and CoLA due solely to the small size of the data sets, or is it because MUX-PLM is not suitable for these types of tasks? Could you sample a subset from data sets like SST-2 to confirm whether MUX-PLM performs poorly on all small data sets?\n- What is the model size of T-MUX in Table 1? Is it the same size as MUX-BERT and MUX-ELECTRA?\n- How is the DeMUX function calculated in Line 288?\n- Section 5.4 analyzes the results of ensemble learning. Although the performance has improved compared to without ensemble learning, it is still not as good as the BERT model with equivalent computational cost. What could be the reason for this?\n", "paper_topic_and_main_contributions": "The paper explores the application of Data Multiplexing technology in pre-trained language models to enhance model throughput. The authors propose a three-stage MUX-PLM training method, along with the design of novel MUX and DeMUX modules. Experimental validation demonstrates that MUX-PLM can effectively boost model throughput with minimal impact on performance. Furthermore, this technique can be integrated with existing compression methods to achieve superior results.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "rznauWNsNI", "reviewer": "Reviewer_R2qS", "summary": "", "questions": "1)\tWhen using the contextual multiplexer, this should add some complexity overhead on the top of total inference time, and this is not considered in your speedup results. Correct?\n2)\tWhat is the Demux architecture? Like one in the TMUX paper?\n3)\tAre Trans_ins and Trans_ctx full transformers? Encoder-based transformers ? Decoder-based ?\n4)\tIf I understand well, you claim that your method does not involve architectural change during the finetuning step. So do most of the baselines you listed: for example, it is true mobile-bert has a modified architecture, but during finetuning, it is finetuned like a regular bert model. Not sure I understand what you mean. \n5)\tJust an observation: I would not call your method RSA-like because RSA works in a very different way…", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed comments. We really appreciate it.\nWe would like to clarify a few misunderstandings:\n\n\n>  Method is not novel as it was introduced in a previous paper and even tested with Bert models (TMUX).\n\n\nMUX-PLMs, to the best of our knowledge, is the **first pre-trained MIMO-style (Multi-input Multi-output) class of models**. Pre-training MIMO-style models is not trivial as it requires significant changes to both the pre-trained objective and the pre-training procedure to ensure training stability and competitive performance. \n\n\n**The difference between T-MUX and MUX-PLMs is akin to the difference between the Transformer and BERT paper, which were both novel and impactful contributions in their own right.** T-MUX has a transformer backbone but is not pre-trained and **Murahari. et al. do not propose either pre-training strategies or pre-trained models.** T-MUX models perform very poorly compared to MUX-PLMs and other SOTA models.\n\n\nThe goal of our paper is to show that MIMO-style architectures can achieve competitive performance compared to SOTA models and to that end, we explored different design choices such as the demultiplexing and multiplexing modules and the pre-training recipe/ algorithm.\n\n\n>  What is the Demux architecture? Like one in the TMUX paper?\n\n\nOur novel RSA-demux demultiplexing module (Figure 2, Section 3.2) enables large-scale pre-training of MUX-PLMs and is significantly faster and more performant than the prefix-based demultiplexer introduced by T-MUX (Section 3.1).\n\n\nPlease note that there are several major differences between the two demultiplexing modules. Firstly, the prefix-based demultiplexer adds a set of prefix tokens (Section 3.1) and uses the output of the prefix tokens as conditioning for the demultiplexer to disentangle the representations. Our novel RSA-demux module instead learns a set of N demultiplexing vectors (Figure 2, Section 3.2) along with the other layers in the module to disentangle the multiplexed representations. This implies the following: \n\na) The conditioning for the demultiplexing module does not vary for every input unlike the prefix-based demultiplexer and this leads to better training stability. This is because the N keys transfer well across the three different training stages. For example, the ablations in Table 5 indicate that the prefix-based demultiplexer from T-MUX does not even converge for token-level tasks and performs worse than our RSA-demux in almost all settings. \n\nb) This demultiplexing module is significantly faster as observed in the speedup column in Table 1.\n\n\n**Differences between T-MUX and MUX-PLMs**\n\n> From the ablation study, it seems like this method is only effective (compared to the original TMUX) when used on token-level tasks. \n\nOur method consists of all the three components listed below, and each component is in itself novel and performant. Therefore, the ablation in Table 2 does not suggest that our method only does well on token-level tasks but suggests that using the contextual multiplexer together with the RSA-demultiplexer module improves performance on token-level tasks.\n\n\n1. RSA-demux module:  Our novel RSA-demux demultiplexing module enables large-scale pre-training of MUX-PLMs and is significantly faster and more performant than prior work.\n\n2. Pre-training MIMO style models: We have addressed this in our response above. T-MUX is not a pre-trained model and also does not provide any pre-training objectives or strategies. MUX-PLMs represent the first pre-trained MIMO style models competitive with SOTA models.\n\n3. Contextual multiplexer:  We introduce our novel contextual multiplexer module to explore the importance of various design choices for effective multiplexing. Our contextual multiplexer module highlights the importance of aggregating context for token-level tasks but does not improve performance for sequence classification tasks.\n\n\n> Authors method performs worse than SOTA model compression methods such as MobileBERT and TinyBERT: while the former only increases models’ speed inference, the latter also reduce models’ sizes. On a size, we have a method that accelerates models’ inference and on the other, we have methods that compress models’ size and accelerate them even more. Which one would you pick?\n\n\nAs noted by all the reviewers, MUX-PLMs and MIMO architectures more generally are complementary to existing model compression methods.\n \nWe demonstrate this in Table 2 by combining MUX-BERT (N = 2) with CoFi, a SOTA structured pruning approach. The resulting model is 2x faster than the MUX-BERT model and 4x faster than the baseline, while being equally performant, and demonstrates that the **acceleration ratios for both methods multiply without a significant drop in performance**.\n\nTherefore, future work could easily extend MUX-PLMs to yield MUX-TinyBERT or MUX-MobileBERT to get the best of both paradigms i.e. highly compressed models with very high throughput.\n\n\n> I think this framework could have been tested on decoder-based transformers as well. I don’t see any of that on the paper.\n\nThanks for the suggestion. We agree that the framework could be extended to decoder-based transformers. Given the compute budget constraints in an academic setting, we focussed our attention on thoroughly exploring MIMO-style pre-trained models for transformer-encoder models but we expect future work to trivially extend MUX-PLMs to decoder-based transformers.\n \n> Are Trans_ins and Trans_ctx full transformers? Encoder-based transformers ? Decoder-based ?\n\n> When using the contextual multiplexer, this should add some complexity overhead on the top of total inference time, and this is not considered in your speedup results. Correct?\n\n\nTrans_ins and Trans_ctx are **single** transformer **layers**. Therefore, the contextual multiplexer module in total has two transformer layers.\n\n\nAdding the two additional layers does increase overhead but not significantly.  The overhead increases with increasing N as the two transformer layers have to process a very large batch size for a large N. \n\n\nWe include the speedup numbers for Table 5 below:\n\n\n| MUX(N) | Model      | MUX            | DEMUX     | Speedup |\n|--------|------------|----------------|-----------|---------|\n| 2      | MUX-BERT   | Non-contextual | RSA-demux | 2       |\n| 2      | Ablation 1 | Non-contextual | Prefix    | 1.9     |\n| 2      | Ablation 2 | Contextual     | RSA-demux | 1.7     |\n| 5      | MUX-BERT   | Non-contextual | RSA-demux | 4.9     |\n| 5      | Ablation 1 | Non-contextual | Prefix    | 4.4     |\n| 5      | Ablation 2 | Contextual     | RSA-demux | 4.1     |\n| 10     | MUX-BERT   | Non-contextual | RSA-demux | 9.8     |\n| 10     | Ablation 1 | Non-contextual | Prefix    | 7.9     |\n| 10     | Ablation 2 | Contextual     | RSA-demux | 5.9     |\n\n> Just an observation: I would not call your method RSA-like because RSA works in a very different way…\n\nWe were inspired by RSA when coming up with the RSA-demux demultiplexing module. If the reviewers find it confusing, we will be happy to change the name.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1)\tAuthors extend previous work done on data multiplexing by proposing a framework that outperforms its predecessor. \n2)\tTheir contextual multiplexer seems to add some gains to the model’s final performance.\n3)\tMethod can be applied to different architectures (BERT, ELECTRA).\n4)\tMethod can be combined with other model compression methods.\n5)\tPaper written in a good English.", "reasons_to_reject": "1)\tThere are some ambiguities in the paper: check “questions for authors” section.\n2)\tI don’t think there are many changes compared to the original TMUX version, a part from the contextual multiplexer (which is barely used by authors in the experiments) and the reduction of N de-multiplexing keys to one. Please correct me if I am wrong.\n3)\tAuthors method performs worse than SOTA model compression methods such as MobileBERT and TinyBERT: while the former only increases models’ speed inference, the latter also reduce models’ sizes. On a size, we have a method that accelerates models’ inference and on the other, we have methods that compress models’ size and accelerate them even more. Which one would you pick?\n4)\tI think this framework could have been tested on decoder-based transformers as well. I don’t see any of that on the paper.\n5)\tFrom the ablation study, it seems like this method is only effective (compared to the original TMUX) when used on token-level tasks.", "questions_to_authors": "1)\tWhen using the contextual multiplexer, this should add some complexity overhead on the top of total inference time, and this is not considered in your speedup results. Correct?\n2)\tWhat is the Demux architecture? Like one in the TMUX paper?\n3)\tAre Trans_ins and Trans_ctx full transformers? Encoder-based transformers ? Decoder-based ?\n4)\tIf I understand well, you claim that your method does not involve architectural change during the finetuning step. So do most of the baselines you listed: for example, it is true mobile-bert has a modified architecture, but during finetuning, it is finetuned like a regular bert model. Not sure I understand what you mean. \n5)\tJust an observation: I would not call your method RSA-like because RSA works in a very different way…", "paper_topic_and_main_contributions": "This paper presents a method to augment language models’ throughputs using data multiplexing/demultiplexing. Like signals multiplexing in signal processing, authors’ method consists in superimposing entries, feed them to the model and disentangle the output of the model to recover each entry’s output. This way, the model can process multiple sentences at the same time. Method is not novel as it was introduced in a previous paper and even tested with Bert models (TMUX). Authors also propose a new framework to multiplex the signal, called contextual multiplexer, and claim their overall framework gets its inspiration from the RSA protocol, as it uses keys (vectors) to encode and decode signals. On paper their method seems to outperform the regular TMUX (the multiplexing framework that was previously published) but performs poorly compared to SOTA model compression methods (knowledge distillation). Their method requires a 3-stage-training: 1) autoencoder setting, to learn multiplex/demultiplex inputs/outputs. 2) Regular pretraining (with self-supervised objective). 3) Finetuning on downstream tasks. To my understanding, their method only differs from the TMUX by the length of the demultiplexing key.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "zntV3bAHgU", "reviewer": "Reviewer_29YX", "summary": "", "questions": "1. Why is Table 2 not showing the performance of the Block Pruning, CoFi, and MobileBERT methods on the QQP task?\n2. Section 5.4 demonstrates the performance of MUX-PLMs in an ensemble setting, where the computational overhead remains consistent with the original PLM. How does the performance of the ensemble MUX-PLMs compare to the original PLM?\n3. In Table 5, when using the prefix as the demultiplexing layer with N = 5 or 10, the model fails to converge on the token classification task. Could you please provide an explanation for this phenomenon?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your valuable comments and we thank your support in bringing the MIMO paradigm to our community. \n\n\n> In Table 5, when using the prefix as the demultiplexing layer with N = 5 or 10, the model fails to converge on the token classification task. Could you please provide an explanation for this phenomenon?\n\nThis is a great question. The prefix demultiplexer uses the output of the prefix tokens as conditioning in the demultiplexer. This makes the conditioning sensitive to the input and also causes stability issues when training the model across three different stages with three different objectives. This problem does not arise in the RSA-demux module as the N independent keys in the module transfer very well between the three different stages.\n\n The token classification task requires the model to make predictions for every single position in the instance, which consequently requires it to multiplex and demultiplex all token positions in the output efficiently. This problem is magnified for higher values of N (5 or 10). Therefore, the prefix-based demultiplexer struggles to effectively demultiplex in the harder setting.\n\n> According to the results in Table 5, the proposed contextual multiplex layer fails to effectively enhance the model's performance on downstream tasks.\n\nWe introduce the Contextual MUX module not only to improve performance but also to understand how design choices in the multiplexing module affect performance. Our experiments indicate that cleverly aggregating context does improve performance for token-level tasks, which require the model to make predictions for every single position in the instance, which consequently requires it to multiplex all token positions in the output efficiently. This might not be as critical for sequence classification tasks which require a single prediction for an entire sequence.\n\n> Section 5.4 demonstrates the performance of MUX-PLMs in an ensemble setting, where the computational overhead remains consistent with the original PLM. How does the performance of the ensemble MUX-PLMs compare to the original PLM?\n\nEnsembling MUX-PLMs improves performance over single MUX-PLMs but it is still lower compared to the original PLM.\n\nThis is likely because a good ensemble involves a diverse set of functions that are accurate on different parts of the input space. Here, the function being ensembled is the model output for a particular index of the set of multiplexing functions. Therefore, we can think of it as an ensemble of N functions,  $f_1, \\cdots f_N$, where $f_i$ is the output of MUX-BERT for the $i^{th}$ instance. These functions would be incentivized to be similar given the training objectives and this consequently would constitute a weaker ensemble.\n\nImportantly, we introduced the ensemble primarily to showcase that one could dynamically interpolate between accuracy and throughput by controlling which and how many instances are ensembled for a prediction. We discuss this further at the end of section 5.4.\n\n> Why is Table 2 not showing the performance of the Block Pruning, CoFi, and MobileBERT methods on the QQP task?\n\nWe were not able to find official numbers for QQP for the validation split for these baselines.\n\n> The model's pre-training process involves two steps: token retrieval training followed by traditional language model objectives. Conducting an ablation study for both pre-training steps would make the experimental results more promising.\n\nWe did run experiments without the token retrieval training but we could not get the model to converge. We therefore decided to not include that setting.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The authors present an interesting multi-input multi-output algorithm for PLMs, effectively reducing the computational costs of large language models.\n2. Experimental results show that MUX-BERT and MUX-ELECTRA achieve significant reductions in computational costs while maintaining comparable performance to the original models, resulting in improved throughput.", "reasons_to_reject": "1. The authors only validate the model's performance on GLUE, NER, and POS tagging tasks. Including more challenging tasks (e.g., SQuAD for question-answering) would strengthen the experimental results' persuasiveness.\n2. According to the results in Table 5, the proposed contextual multiplex layer fails to effectively enhance the model's performance on downstream tasks.\n3. The model's pre-training process involves two steps: token retrieval training followed by traditional language model objectives. Conducting an ablation study for both pre-training steps would make the experimental results more promising.\n4. In the multi-input multi-output setting, different data in one batch can influence each other. Therefore, the authors should analyze and discuss the impact between different data within a batch.", "questions_to_authors": "1. Why is Table 2 not showing the performance of the Block Pruning, CoFi, and MobileBERT methods on the QQP task?\n2. Section 5.4 demonstrates the performance of MUX-PLMs in an ensemble setting, where the computational overhead remains consistent with the original PLM. How does the performance of the ensemble MUX-PLMs compare to the original PLM?\n3. In Table 5, when using the prefix as the demultiplexing layer with N = 5 or 10, the model fails to converge on the token classification task. Could you please provide an explanation for this phenomenon?", "paper_topic_and_main_contributions": "The paper proposes a novel pre-trained language model called MUX-PLMs, designed for multi-input multi-output scenarios. This method introduces a multiplexing layer at the input layer to compress multiple data streams and a demultiplexing layer at the output layer to decompress the hidden vectors. The experimental results demonstrate that MUX-PLMs, under both BERT and ELECTRA architectures, achieve comparable downstream performance to the original PLMs while significantly reducing the computational overhead.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces MUX-PLMs, a pre-trained language model designed to enhance throughput in multi-input multi-output scenarios. The approach uses data multiplexing to process multiple inputs simultaneously, promising improved inference speeds. Experiments focus primarily on the GLUE benchmark, with the paper arguing for both the novelty and effectiveness of MUX-PLMs. Three reviewers have provided their assessments.

Strengths:

1. Novel Approach: The authors' introduction of a MUX and DeMUX architecture in the realm of PLMs is noteworthy. This innovation aims to address the growing demand and computational cost challenges associated with large language models.
2. Positive Results: Experiments indicate that MUX-PLMs, particularly under BERT and ELECTRA architectures, can maintain performance levels comparable to traditional PLMs while offering significant reductions in computational overhead.
3. Clarity: The paper is well-structured, making the content easily accessible to readers. Several reviewers highlighted its readability.
4. Broad Applicability: The proposed method is shown to work across different architectures like BERT and ELECTRA and can also be combined with other compression methods.

Weaknesses:

1. Insufficient Validation: The model's performance validation seems restricted, with emphasis on GLUE, NER, and POS tagging tasks. An exploration into more diverse and challenging tasks would have strengthened the study's credibility.
2. Performance of Contextual MUX: The paper introduces the Contextual MUX module, but it reportedly does not outperform the original MUX module. Its absence in primary experiments and lackluster results question its utility.
3. Comparison with SOTA: The method's performance, when juxtaposed with state-of-the-art model compression techniques (like MobileBERT and TinyBERT), is reportedly not competitive. Reviewers questioned the practical adoption of MUX-PLMs given these findings.
4. Ambiguities & Lack of Depth: Several concerns were raised regarding ambiguities in the paper and a perceived lack of depth in specific sections. For instance, the method's uniqueness compared to the earlier TMUX version was questioned. A more profound exploration of the impact of different data within a batch in a multi-input setting would also have been beneficial.
5. Training Cost & Flexibility: The potential increased training costs associated with MUX-PLMs and the need to pre-train different MUX-PLMs for different acceleration ratios could be a limitation compared to existing methods that fine-tune a single PLM.

The paper, with its focus on MUX-PLMs, offers a unique lens into improving language model throughput via data multiplexing. While the premise is exciting and there are promising results, there are notable gaps in validation, comparative analysis, and practical implications that should be addressed. Future revisions should consider expanding the experimental range, addressing ambiguities, and elucidating the tangible benefits of adopting MUX-PLMs over existing methods.|This paper presents MUX-PLMs, a pre-trained language model strategically crafted to optimize throughput in multi-input multi-output scenarios. By harnessing data multiplexing, this approach concurrently handles multiple inputs, offering the prospect of significantly enhanced inference speeds. The experimental evaluation primarily centers around the GLUE benchmark, demonstrating the novelty and efficacy of MUX-PLMs. The insights of three reviewers further enrich the paper's assessment.
Decision: Accept-Findings