DataMUX: Data Multiplexing for Neural Networks
Abstract
In this paper, we introduce data multiplexing (DataMUX), a technique that enables deep neural networks to process multiple inputs simultaneously using a single compact representation.
DataMUX demonstrates that neural networks are capable of generating accurate predictions over mixtures of inputs, resulting in increased inference throughput with minimal extra memory requirements.
Our approach uses two key components – 1) a multiplexing layer that performs a ﬁxed linear transformation to each input before combining them to create a ‘mixed’ representation of the same size as a single input, which is then processed by the base network, and 2) a demultiplexing layer that converts the base network’s output back into independent representations before producing predictions for each input.
We show the viability of DataMUX for different architectures (Transformers, and to a much lesser extent MLPs and CNNs) across six different tasks spanning sentence classiﬁcation, named entity recognition and image classiﬁcation.
For instance,
DataMUX for Transformers can multiplex up to 20x/40x inputs, achieving up to 11x/18x increase in inference throughput with absolute performance drops of < 2% and < 4% respectively compared to a vanilla Transformer on MNLI, a natural language inference task.
We also provide a theoretical construction for multiplexing in self-attention networks and analyze the effect of various design elements in DataMUX.
Introduction
Deep neural networks are very effective at modeling complex functions from real-valued vector inputs to vector outputs.
However, recent studies have hinted that modern networks are vastly overparameterized Kaplan et al.
Allen-Zhu et al.
(2019a)
and only require a fraction of their parameters for capturing many functions Frankle and Carbin (2018); Frankle et al.
This raises a natural question: if networks contain greater processing capacity than necessary, could it be possible for them to model a function over multiple inputs simultaneously, similar to how radio channels share bandwidth to carry multiple information streams at once?
Such a capability could enable signiﬁcant improvements in efﬁciency and inference throughput without requiring substantial increases in model size or computation.
In this paper, we propose Data Multiplexing (DataMUX), a technique that enables neural networks to process multiple inputs simultaneously as a single “mixed” of the same size as a
[Table]
|
[Table]
+uoNI5qETCIVROuu68ToZUQhp4JNi71Es5jQCRmxrqGShEx72fzqqX1ulKEdRMqURHuu/p7ISKh1GvqmMyQ41sveTPzP6yYY1LyMyzhBJuliUZAIGyN7FoE95IpRFKkhhCpubrXpmChC0QRVNCG4yy+vkla14l5V3PvLcv0mj6MAp3AGF+DCNdThDhrQBAoKnuEV3qwn68V6tz4WrWtWPnMCf2B9/gAMnZI3</latexit> ˆ y single input, during both training and inference (Figure 1).
This is complementary to and different from minibatch processing that adds an extra batch dimension to the input, which in effect increases computational and memory requirements.
Due to the use of a single combined representation,
DataMUX can increase network inference throughput with minimal time and memory overhead.
The key challenge for DataMUX lies in effectively compressing a mixture of signals so that it can be processed by a neural network without sacriﬁcing task performance.
To this end, our procedure involves the addition of a multiplexing layer and a demultiplexing layer to both ends of a vanilla neural network.
The multiplexing layer applies a ﬁxed linear transformation to each input before combining them to create a single “mixed” representation.
This is then fed into the base network, which processes it to obtain a “mixed” vector representation.
The demultiplexing layer then converts this vector back into a set of vector representations corresponding to each original input, which is then used to make ﬁnal predictions for each instance.
Our multiplexing and demultiplexing layers are end-to-end differentiable, which allows us to train the entire model jointly through standard gradient descent methods.
We ﬁnd that training time to convergence increases almost linearly with the number of multiplexed in stances as training gets more challenging since the model has to make multiple simultaneous predictions.
In order to train multiplexed models effectively, we also introduce a special token retrieval task as a pre-training objective.
We demonstrate the effectiveness of our DataMUX approach using three different architectures (Transformers Vaswani et al.
(2017), MLPs, and CNNs) on six tasks spanning sentence classiﬁcation, named entity recognition and image classiﬁcation.
Our results demonstrate that DataMUX models can successfully process between 2-40 input instances simultaneously, with minimal losses in task performance.
For instance, on the sentiment analysis dataset SST-2, our multiplexed Transformer (T-MUX) can handle 40x inputs with only a 2% drop in accuracy compared to a vanilla Transformer.
To our best knowledge, this is the ﬁrst demonstration of large-scale data multiplexing in Transformers.
While our results on CNNs and MLPs for image classiﬁcation are not as strong as those on
Transformers, we present them to demonstrate the generality of our method and believe that better multiplexing methods may yet improve them in the future.
We further perform theoretical and empirical analyses to better understand our results.
First, we show that it is possible to theoretically construct self-attention networks that can simultaneously process inputs in N independent subspaces.
Next, we perform several analyses of various multiplexing strategies and the effect of attention heads and model size.
Finally, we also measure throughput statistics to demonstrate the efﬁciency gains provided by DataMUX (up to 18x inference speedup).
We believe multiplexing holds great promise in improving the efﬁciency and versatility of neural networks and hope that future research can develop improved strategies for DataMUX and expand its applicability to more architectures and tasks.
Related Work
Multiplexing in deep learning While multiplexing is a standard idea in signal processing Rabiner and Gold (1975), its use in deep learning has been limited.
Lu et al.
(2020) develop MUXConv, a layer that multiplexes both spatial and channel information for convolutional neural networks (CNNs).
Instead of performing data multiplexing , these models multiplex different features within the network and only process one input at a time.
mixup Zhang et al.
(2018) proposes a training scheme where networks are trained with a convex combination of instances to predict the instances’ label distributions from this combined representation.
They ﬁnd that this scheme provides implicit regularization of models and improves generalization during inference.
Unlike our work however, mixup does not preserve the order of the combined instances and therefore mixed instances may not be used during inference.
Another line of work uses multiple-input multiple-output training to obtain multiple sub-networks within a neural network for a relatively small number of inputs
Havasi et al.
; Ramé et al.
; Soﬂaei et al.
These approaches aim to improve single-instance performance through ensembling sub-networks’ predictions.
In contrast, we are motivated by achieving optimal performance for prediction on multiple inputs and improving model throughput during inference.
Multiplexing in the brain On the biological side, several studies hint at the capability of the brain to multiplex and transmit tremendous amounts of information with a single neural code and process them in parallel Blumhagen et al.
; Akam and Kullmann (2014); Pirschel and Kretzberg (2016);
Hong et al.
; Friedrich et al.
Recently, Lankarany et al.
conﬁrmed that neurons in the somatosensory cortex multiplex different types of features (low-contrast, high-contrast) using their timing and rate of synchronous and asynchronous spiking.
Rezaei et al.
(2021) further provide a computational model to explain multiplexing with an ensemble of homogeneous cortical neurons.
These studies increasingly demonstrate that data multiplexing is a key feature of natural neural networks.
Over-paramaterization in neural networks Our work also relates to the over-parameterization of deep neural networks that has been investigated in several pieces of prior work Allen-Zhu et al.
(2019b,a); Radhakrishnan et al.
The lottery ticket hypothesis Frankle and Carbin (2018);
Frankle et al.
(2020); Malach et al.
(2020) demonstrated that only a fraction of neurons in a trained deep model are sufﬁcient to capture a required function.
Neural architecture search Zoph and Le (2016); Liu et al.
Elsken et al.
(2019) aims to balance performance with parameter efﬁciency, by discovering the most optimal architectures for each task.
There have also been several efforts at leveraging this over-parameterization effect for multi-task learning Hu and Singh (2021).
Cheung et al.
(2019) propose a method to superpose many models for different tasks into a single neural network, and show that it helps mitigate catastrophic forgetting.
Wortsman et al.
(2020) is another work along the same lines which uses a base network, along with specialized sub-networks for different tasks.
Similar to these paradigms, our work possibly takes advantage of the over-parameterization existing in deep neural networks.
However, to our knowledge, we are the ﬁrst to utilize data multiplexing to simultaneously process multiple inputs during both training and inference, and are able to show that a ﬁxed set of parameters can process a mixture of inputs (even up to 40) with minimal loss in accuracy.
Method
Our primary goal is to process a set of input instances simultaneously over a shared neural network (multiplexing) with minimal overhead during inference.
To this end, we design DataMUX to consist of three components: a multiplexer module to combine the multiple input instances into a superposed representation, a neural network backbone to process mixed representations, and a demultiplexing module to disentangle the processed representations for individual prediction.
A schematic illustrating the ﬂow of representations is shown in Figure 2.
We detail the general requirements for each component below, as well as speciﬁc implementation details used in our experiments.
3.1 Multiplexing
The multiplexer module, denoted  , combines a tuple of inputs, either images or sentences from a batch, ( x 1 , . .
, x N ) , for x i 2 R d , into a more compact representation x 1: N 2 R d in an orderdependent way, which enables effective demultiplexing after processing, as well as distinguishing intra-sequence interactions in the case of sequenced inputs (e.g. token sequences).
Towards this end, for each input x i with index i 2 [1 , N ] of the input tuple, the multiplexer module performs a transformation  i
( R d 7! R d ), on the instance before ﬁnally averaging all inputs into a single multiplexed representation as:
[Equation]
For sequenced inputs (e.g. token sequences), we combine N sequences by multiplexing tokenwise.
That is for inputs of the form x i = { w ij } j 2 [1 ,L ] , where w ij 2 R d is a token’s input vector representation, this operation uses the same transformation  i for each token in the sequence before averaging over each position across indices, such that x 1: N = { w 1: N j } j 2 [1 ,L ] where w 1: N j 2 R d is a multiplexed representation of N tokens at position j , i.e., x 1: N = {  ( w 1 j , . . . , w Nj ) }
j 2 [1 ,L ] . For  i , we experiment with using either (1) a linear projection with a random ﬁxed orthogonal matrix (denoted “Ortho”) or (2) the Hadamard product with a ﬁxed Gaussian random vector (denoted “Hadmard”, equivalent to a linear map using a diagonal matrix in our case).
We hope these transformations map instances at different indices into distinguishable regions and consequently reduce interference between their representations.
Finally, this multiplexed representation, x 1: N , is used as input to the neural network backbone, which is architecturally unchanged.
3.2 Demultiplexing
The output of the neural network backbone will be a multiplexed hidden representation h 1: N of the input x 1: N .
To make a prediction for each input, one can explicitly disentangle h 1: N into N individual hidden representations, h 1 , . .
, h N , respectively.
We ﬁrst obtain each h i with a demultiplexing function # i , i.e.,
[Equation]
For sequenced input, demultiplexing is done position-wise, i.e., h ij =
# i ( h 1: N j ) for each position j .
Finally, predictions are made using a shared task head on each inputs’ respective individual hidden representation to prevent a substantial increase in the number of parameters and improve training efﬁciency.
We use two alternatives for the demultiplexing function # i :
MLP Demuxing
This strategy employs N MLPs to generate each indices’ hidden representation as h i =
MLP i ( h 1: N ) .
We use this method for both NLP and vision tasks.
Although this method is conceptually simple, it adds learnable parameters proportional to N . 2.
Index Embeddings We
generate index embeddings p i , which are then concatenated to h 1: N j , and transformed by a shared multi-layer network to generate each individual hidden representation, i.e., h ij =
MLP shared ( h 1: N j , p i ) .
To generate the index embeddings p i , we add a sequence of N special tokens, called the preﬁx, to the beginning of each sequence of the input tuple.
For multiplexing with
N sequences, we add N corresponding preﬁxes, denoted preﬁx i for i 2 [1 , N ] .
Each preﬁx i consists of a index token ✏ i in it’s
i ’th position while the remaining tokens are a special pad token ✏ pad .
The preﬁx sequences then take on the following pattern:
[Equation]
We then prepend each sequence x i of the input sequence with the corresponding preﬁx i .
The tuple of prepended sequences is then passed to the multiplexing module.
When ﬁnally generating individual hidden representations, we use the corresponding hidden representation of each index token ✏ i as the index embedding p i .
We use the Index Embeddings ( 3.2) demultiplexing strategy on language tasks for the Transformer architecture.
The preﬁx tokens may implicitly enable the Transformer to do instance-speciﬁc computations when processing the multiplexed representation and further enable demultiplexing for large N .
3.3 Retrieval warm-up for multiplexing
For Transformer models, we ﬁnd that naively adding the multiplexing and demultiplexing layers to the model fail to converge.
This is likely because the gradients for individual instances from the task loss get mixed up in the backward pass.
To overcome this, we propose Retrieval warm-up – a self-supervised pre-training task to promote the ability of DataMUX models at distinguishing the order and the content of individual sequences in a multiplexed representation.
This task consists of retrieving the correct tokens and order for each position and sequence of the input tuple (Figure 4a).
Although we could add this loss for every sentence in each position, memory constraints force us to retrieve a token from one random sentence for each token position, yielding the following objective:
[Equation]
[Table]
Transformer
[Table]
nato said
[Table]
Demultiplexer
[Table] [Table]
where h I j is a demultiplexed hidden representation of the j -th token in a randomly selected sentence with the index I ⇠ U [1 , N ] , generated using the methods described in Section 3.2.
We use this objective to optimize all the parameters of the end-to-end multiplexed model.
For sequenced inputs, we ﬁnd the retrieval auxiliary objective crucial to learning and report performance on the retrieval task in Section 4.2.
Multiplexing for Transformers (T-MUX) 4.1 Experimental setup
Models We ﬁrst evaluate the capabilities and limits of data multiplexing speciﬁcally for the Transformer architecture on a range of text classiﬁcation tasks.
We apply DataMUX on a 12-layer
Transformer based model with a hidden dimension size of 768 and 12 self-attention heads built on the
Huggingface Wolf et al.
(2019) framework, and refer to the resulting model as T-MUX.
We compare our T-MUX models to 2 baselines: (B1)
A 12-layer 768 hidden dimension vanilla Transformer.
(B2)
A 12-layer 768 hidden dimension Transformer pretrained using the retrieval task described in Section 3.3.
Though there is no multiplexing done for B2 (meaning this operation could be solved by simply copying input tokens to the output layer)
we ﬁnd that the retrieval pre-training produces differences in performance and we show this baseline for completeness.
We also apply DataMUX to smaller models (see A2 ) and compare with similar baselines.
Tasks We evaluate our models and the baselines on two types of text classiﬁcation tasks: 1.
Token-level classiﬁcation: This evaluates models’ ability to perform token-level tasks on multiplexed inputs.
This poses a particular challenge for data-multiplexing models since they must maintain a high level of individual token disentanglement while also producing representations capable of solving the task.
We evaluate token-level classiﬁcation on the CoNLL-2003 Named Entity Recognition (NER) task Sang and Meulder (2003).
Sentence-level classiﬁcation: We evaluate models on a subset of the sentence-level classiﬁcation tasks found in the General Language Understanding Evaluation (GLUE) benchmark Wang et al.
(2019): the sentiment classiﬁcation task [-SST-
-]SST-2 Socher et al.
(2013), the sentence similarity task
QQP 2 , and the natural language inference tasks MNLI Williams et al.
and QNLI Wang et al.
; Rajpurkar et al.
By evaluating on a variety of sentence-level tasks, we can gain a better sense of the capabilities of data multiplexing neural networks on tasks that require aggregating contextual information.
Similar to previous works, we prepend a [CLS] token to all sequences and learn a task head on top of the demultiplexed [CLS] token representation.
Auxiliary retrieval objective The T-MUX models are all pre-trained using the retrieval warm-up on the Wikitext-103 dataset Merity et al.
(2017) as described in Section 3.3.
In addition, we also continue to use the retrieval task as an auxiliary objective during task training.
The total loss is a combination of the task loss and retrieval loss (we use ↵ = 0 . 1 in our experiments):
[Equation]
4.2 Main results
(R1)
Multiplexing leads to minimal drop in performance even for large N Figure 3 shows performance across four sentence classiﬁcation tasks (MNLI, QNLI, QQP, SST2) and a token-level classiﬁcation task (NER).
We observe that it is possible to multiplex up to 40 instances and maintain reasonable performance.
For easier tasks like QQP, SST2 and QNLI, we observe that the drop in performance with increasing N is insigniﬁcant while for more difﬁcult tasks like MNLI and
NER, there is a trade-off between performance and N , with performance dropping 10%-15% for 40 instances.
Since we’ve encountered unstable optimization when using MLP Demuxing, we only provide results using Index Embeddings demultiplexing.
We ﬁnd that the multiplexing strategy does not impact performance across different tasks, even slightly increasing for small values of N (e.g. 2, 5).
This increase may be attributable to implicit regularization à la mixup .
We also did not notice much variance between ﬁne-tuning runs and plot variance for the "Hadamard + Index Embedding" setting in Section A.4.
(R2)
Perfect multiplexing on the retrieval warm-up task for large N Figure 4b shows the test accuracy on the retrieval warm-up task described in Section 3.3.
We ﬁrst note that across different multiplexing and demultiplexing strategies, models have a retrieval accuracy of nearly 100% for up to 20 instances, demonstrating the surprising ability of T-MUX to multiplex perfectly for large N .
Note that this task does not require any aggregation of context across the sequence and thereby is much easier than sentence classiﬁcation or token-level classiﬁcation tasks.
Therefore, performance on this warm-up task indicates a soft upper bound on the number of instances we can multiplex for sentence and token-level classiﬁcation tasks given a particular multiplexing and demultiplexing method.
(R3)
Throughput can be increased multi-fold
We measure throughput of our multiplexed model (Hadamard + Index Embed) across different number of instances by calculating inference speed for processing ⇠ 20,000 samples on the MNLI dataset.
We use four different batch sizes for all the conﬁgurations and take the max throughput (details in A.8).
Figure 4c shows that multiplexing increases throughput many folds ( 18 x for 40 instances, 11 x for 20 instances).
T-MUX enables superior throughput as batch size can be effectively increased by a factor of N .
We would expect the speedup to scale linearly with N , however a large N corresponds to more preﬁx tokens which increases the sequence length.
Therefore, having 40 instances leads to almost a 20 x speedup as opposed to the expected 40 x. Future work can potentially improve this speedup by designing better multiplexing and demultiplexing strategies.
[Table]
Demultiplexer
[Table]
[Figure]
DmCOfFeXc
[Figure]
5q05J5s5hD9wPn8ADRmNqA
[Figure]
latexit
[Figure]
4.3 Analysis
(A1)
The number of attention heads seems invariant to multiplexing To understand the role of the number of attention heads in multiplexing, we train a variant of T-MUX with 2 self-attention heads per layer.
We use T-MUX with the (Hadamard + Index Embed.) conﬁguration.
We ﬁnd that reducing self-attention heads has minimal effect on performance on the retrieval warm-up task and achieves a retrieval accuracy of ⇠ 100% up to N = 20 .
We then ﬁnd that on token and sentence-level classiﬁcation tasks, T-MUX with 2 self-attention heads performs comparably to T-MUX withself-attention heads (Figure 5a).
A2) DataMUX also provides throughput boost with smaller Transformers We further investigate the applicability of DataMUX to smaller Transformer models.
We choose two smaller Transformersto multiplex: a 12 layer with hidden size of 384 (12L / 384H), a 4 layer with hidden size of 768 (4L / 768H).
Figure 5b shows that these smaller T-MUX models can also multiplex up to 20 instances with competitive performance.
Figure 4c illustrates the speedup from the smaller models.
As the smaller models can only multiplex up to 20 instances with reasonable performance, we see that multiplexing with 20 instances provides an even higher throughput of 25 x, compared to only 18 x for the full-sized
T-MUX with 40 instances.
A3)
Performance varies more across different indices as N increases The prediction of an instance is conditioned on the index of the instance.
Figure 7b illustrates performance on MNLI for different indices across different choices of N for the (Hadamard + Index Embed.) conﬁguration.
We observe that performance varies more across different indices for large N (For N = 40 , performance varies across ⇠ 10 percentage points).
(A4)
The demultiplexed representation of an instance is robust to the other instances it is multiplexed
with To analyze whether the representation of an instance changes with respect to the other samples it is multiplexed with, we randomly select 10 instances x 1 , x 2 , · · · x 10 from the
MNLI dataset and multiplex each of them (separately) with 30 different sets of other instances
y i 1 , y i 2 , · · · , y iN  1 for i 2 [1 · · · 30] .
This generates 30 different demultiplexed representations for each of the 10 selected samples.
We then visualize the resulting points by reducing the 768 dimensional demultiplexed representation to a 200 dimensional vector with PCA ( F.R.S. (1901)), followed by t-SNE ( van der Maaten and Hinton (2008)) visualization.
Figure 6 shows these clusters for different values of N .
We ﬁnd that across different values of N , all the points corresponding to the same x i are very close to each other, which suggests that the representation of an instance is not signiﬁcantly inﬂuenced by the set of instances it is multiplexed with.
4.4 Theoretical construction for multiplexing in self-attention models
To provide some theoretical explanation for why Transformers may be amenable to multiplexing, we also detail below a construction for self-attention based neural networks which process multiplexed token embeddings in N independent subspaces.
Our construction relies on particular structural properties of singular spaces of linear transformations across layers – we provide a brief sketch below, with more details in Appendix A.3.
In a model with self-attention, the multi-head attention projects the queries, keys, and values h times with different, learned linear projections to d K , d K , and d V dimensions respectively.
Thus, given a sequence of d -dimensional multiplexed token embeddings { w 1: N t } t 2 [1 ,L ] , each head looks like:
[Equation]
Now, w 1: N t
[Equation]
N k =1  k ( w kt ) , and let us assume each function  k can be trained to project each embedding into a subspace that is least linearly-dependent with the others.
So, if we deﬁne u ( k ) t := 1 N  k ( w kt ) , we assume h u ( k ) t , u ( k 0 )
t 0 i ⇡ 0 for all pairs of indices k 6 = k 0 and all positions t .
To preserve this independent subspace structure after self-attention, we will ﬁrst need
[Equation]
This is achievable if the eigenvectors of W Vi > W Vi can be grouped into N non-overlapping subsets { r
[Equation] , r ( [Equation]
[Equation] , { r ( N ), . . . , r ( N ) m } , where r ( k ) ` are orthonormal vectors (since the Gramian is real symmetric), and span the same input subspaces, denoted as D 1 , . .
, D N .
In this case, the vector after transformed by W Vi can be expressed as a sum of N vectors v (1) i,t , . .
, v ( N ) i,t in dual subspaces
D 1 V , . . .
, D NV which are independent of each other.
The linear maps from D k to D kV still allows rich operation on each component of a multiplex input without interference by other components.
In addition to decompress-able value vectors, we can set the query and key matrices, W Qi and W Ki , to have some subsets of right and left singular vectors such that span D 1 , . . .
, D N and D 1 V , . . .
, D NV .
Then, we can show that the inner product of the query and keys of the i -th head can be rewritten as:
[Equation]
[Equation]
where ⌧ ( k ) i,t,t 0 is a scalar only depending on the k -th input sequence.
Thus, the self-attention operation at each position can be seen as retrieving values based on the average of query-key similarity scores of N sequences.
[Equation] [Equation]
This average retrieval using soft-max could be a desired property as implicit regularization.
However, if we want perfect non-interference in retrieval, the network always has an option to specialize each head to only focus on one input sequence, by setting ⌧ ( k 0 )
i,t,t 0 =
for all k {+0 
+}
= k , which is easily achievable by controlling singular values of W Qi or W Ki .
More details of this construction are provided in Appendix A.3.
Figure 4.3 suggests that number of attention heads do not signiﬁcantly impact our empirical multiplexing results.
Therefore our theoretical construction does not fully explain what models learn empirically but nonetheless provides an existence proof for multiplexing under certain assumptions.
Multiplexing for MLPs and CNNs
We also investigate multiplexing for multilayer perceptrons (MLPs) and convolutional neural networks (CNNs).
Certain cases of data multiplexing have been explored for convolutional architectures on image classiﬁcation as techniques for robustness and data augmentation Ramé et al.
; Havasi et al.
; Soﬂaei et al.
These works implicitly employed the frontend layers of a convolutional net as multiplexing layers and suggest that a convolutional neural network can learn at most 3-4 independent subnetworks concurrently Havasi et al.
Since Transformers can be multiplexed for up to 40 instances without a severe performance drop, we explore the DataMUX scheme for MLPs and CNNs on the MNIST image classiﬁcation task LeCun and Cortes (2005).
While our results for these architectures are not as strongly positive as for T-MUX, we believe that future work on better multiplexing strategies can make them more viable.
Multiplexing for MLPs Figure 7a (solid lines) illustrates performance of DataMUX for MLPs with various multiplexing strategies, paired with the MLP demultiplexing strategy.
As a baseline, we show multiplexing using the identity transformation before combining instances (MLP baseline).
Since this transformation does not preserve the order of the multiplexed instances, accuracy expectedly decreases on the order of 1 /N .
Multiplexing with random orthogonal matrices (MLP + Ortho) works for up to 8 instances with an accuracy of ⇠ 78% , compared to ⇠ 95% for the baseline without multiplexing.
Since random orthogonal projections cannot separate inputs perfectly, we also experiment with a set of N low-rank independent transformations for multiplexing (MLP +
LowRank) and observe that this improves performance slightly for larger N (by 5% for N =8).
Multiplexing for CNNs Results for CNN multiplexing are shown in Figure 7a (dashed lines).
Like in the MLP case, all methods shown use the MLP demultiplexing strategy.
We notice that the baselines (CNN baseline) has a roughly N -fold performance drop because of the unidentiﬁability of input order.
The “Ortho” transformation performs quite poorly for CNNs (only ⇠ 56% for N = 8 ).
This is likely because this transformation destroys the property of spatial locality that CNNs rely on.
We therefore also explore N two-layer convolutional networks with a tanh activation as out multiplexing transformations
i (CNN + Nonlinear).
We note that even with this transformation, the dimensionality of the multiplexed representation still remains equal to the dimensionality of a single input.
We ﬁnd that for N  4 , performance is above 80% , which is signiﬁcantly better than
CNN+Ortho, but performance drops rapidly for N > 4 .
We show results of other multiplexing strategies for CNN in Appendix A.11.
Overall, multiplexing for MLPs and CNNs seems to be more challenging than Transformers as evidenced by the sharper performance drops with increasing N .
However, we believe these numbers are still non-trivial and show the potential for multiplexing in CNNs and MLPs, with multiplexing and demultiplexing strategies better suited for these architectures.
Discussion
In this work, we have shown that neural networks can be trained to multiplex, i.e. process multiple inputs using a single “mixed” representation.
We introduced DataMUX, a novel setting for training and inference with multiplexing, and demonstrate its viability using a variety of network architectures and tasks.
Our results reveal the surprising ability of neural networks to predict on up to 40 input instances simultaneously using data multiplexing.
Our efforts have sought to show the potential for data multiplexing using methods that require limited computational overhead or added learned parameters; achieving results that incur minimal degradation in performance while simultaneously increasing system throughput dramatically.
While the scope of this paper is to demonstrate the general ability and potential for neural networks to use data multiplexing, there are several avenues for future research to investigate the limits and theoretical implications of DataMUX.
Directions of particular interest include: large-scale pretraining, multi-modal processing, multi-lingual models, different mechanisms of multiplexing and demultiplexing, and a more rigorous understanding of the architectural, data, and training conditions enabling data multiplexing.
Acknowledgements
We thank Ameet Deshpande, Shunyu Yao, Jens Tuyls, Tianyu Gao and Mengzhou Xia for their valuable feedback on early drafts and encouragement throughout the course of this project, and the anonymous reviewers for their suggestions on improving the paper.
[-i,t,t ′ = 
-]
[-NER (F1) 768 384  24 Hidden Dimension
-]
[-768 384  24 Hidden Dimension 
-]