DataMUX: Data Multiplexing for Neural Networks
Anonymous Author(s)
AfﬁliationAddress email
Abstract
In this paper, we introduce data multiplexing (DataMUX), a technique that enablesdeep neural networks to process multiple inputs simultaneously using a singlecompact representation.
DataMUX demonstrates that neural networks are capableof generating accurate predictions over mixtures of inputs, resulting in increasedinference throughput with minimal extra memory requirements.
Our approach usestwo key components – 1) a multiplexing layer that performs a ﬁxed linear transformation to each input before combining them to create a ‘mixed’ representationof 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 intoindependent representations before producing predictions for each input.
We showthe viability of DataMUX for different architectures (Transformers, and to a lesserextent MLPs and CNNs) across six different tasks spanning sentence classiﬁcation, named entity recognition and image classiﬁcation.
For instance, DataMUXfor Transformers can multiplex up to 20x/40x inputs, achieving 11x/18x increasein inference throughput with minimal absolute performance drops of < 2% and< 4% respectively on MNLI, a natural language inference task.
We also provide atheoretical construction for multiplexing in self-attention networks and analyze theeffect of various design elements in DataMUX.
Introduction
Deep neural networks are very effective at modeling complex functions from real-valued vectorinputs to vector outputs.
However, recent studies have hinted that modern networks are vastlyoverparameterized Kaplan et al.
Allen-Zhu et al.
(2019a)
and only require a fraction of theirparameters for capturing many functions Frankle and Carbin (2018); Frankle et al.
Thisraises a natural question: if networks contain greater processing capacity than necessary, could itbe possible for them to model a function over multiple inputs simultaneously, similar to how radiochannels share bandwidth to carry multiple information streams at once?
Such a capability couldenable signiﬁcant improvements in efﬁciency and inference throughput without requiring substantialincreases in model size or computation.
In this paper, we propose Data Multiplexing (DataMUX), a technique that enables neural networksto process multiple inputs simultaneously as a single ‘mixed’ representation of the same size as asingle input, during both training and inference (Figure 1).
This is complementary to and differentfrom minibatch processing that adds an extra batch dimension to the input, which in effect increasescomputational 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 canbe processed by a neural network without sacriﬁcing task performance.
To this end, our procedureinvolves the addition of a multiplexing layer and a demultiplexing layer to both ends of a vanillaneural network
The multiplexing layer applies a ﬁxed linear transformation to each input before
[Table]
|
[Table]
+uoNI5qETCIVROuu68ToZUQhp4JNi71Es5jQCRmxrqGShEx72fzqqX1ulKEdRMqURHuu/p7ISKh1GvqmMyQ41sveTPzP6yYY1LyMyzhBJuliUZAIGyN7FoE95IpRFKkhhCpubrXpmChC0QRVNCG4yy+vkla14l5V3PvLcv0mj6MAp3AGF+DCNdThDhrQBAoKnuEV3qwn68V6tz4WrWtWPnMCf2B9/gAMnZI3</latexit> ˆ y them to create a representation.
fed is then used to
, which allows us to train the jointly through that training time to convergence increases almost linearly training gets more challenging since the has to make predictions.
In order to train multiplexed models effectively we also a special 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 modelscan successfully process between 2-40 input instances simultaneously, with minimal losses in taskperformance.
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 bettermultiplexing methods may yet improve them in the future.
We further perform theoretical and empirical analyses to better understand our results.
First, we showthat it is possible to theoretically construct self-attention networks that can simultaneously processinputs in N independent subspaces.
Next, we perform several analyses of various multiplexingstrategies and the effect of attention heads and model size.
Finally, we also measure throughputstatistics 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 neuralnetworks and hope that future research can develop improved strategies for DataMUX and expand itsapplicability to more architectures and tasks.
Related Work
Multiplexing in deep learning While multiplexing is a standard idea in signal processing Rabinerand 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 withinthe network and only process one input at a time.
mixup Zhang et al.
(2018) proposes a trainingscheme 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 implicitregularization 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 maynot be used during inference.
Another line of work uses multiple-input multiple-output trainingto 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 improvesingle-instance performance through ensembling sub-networks’ predictions.
In contrast, we aremotivated by achieving optimal performance for prediction on multiple inputs and improving modelthroughput during inference.
Multiplexing in the brain On the biological side, several studies hint at the capability of the brainto multiplex and transmit tremendous amounts of information with a single neural code and processthem in parallel Blumhagen et al.
; Akam and Kullmann (2014); Pirschel and Kretzberg(2016);
Hong et al.
; Friedrich et al.
Recently, Lankarany et al.
(2019) conﬁrmed thatneurons 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) furtherprovide a computational model to explain multiplexing with an ensemble of homogeneous corticalneurons.
These studies increasingly demonstrate that data multiplexing is a key feature of naturalneural networks.
Over-paramaterization in neural networks
Our work also relates to the over-parameterizationof 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 traineddeep 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 atleveraging 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 alongthe 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 existingin deep neural networks.
However, to our knowledge, we are the ﬁrst to utilize data multiplexing tosimultaneously 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 consistof three components: a multiplexer module to combine the multiple input instances into a superposedrepresentation, the neural network backbone and a demultiplexing method to disentangle the processedrepresentations for individual prediction.
A schematic illustrating the ﬂow of representations is shownin Figure 2.
We detail the general requirements for each component below, as well as speciﬁcimplementation details used in our experiments.
3.1 Multiplexing
The multiplexer module, denoted  , combines a tuple of inputs, either images or sentences froma 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 distinguishingintra-sequence interactions in the case of sequenced inputs (e.g. token sequences).
Towards thisend, for each input x i with index i 2 [1 , N ] of the input tuple, the multiplexer module performsa transformation  i ( R d 7!
R d ), on the instance before ﬁnally averaging all inputs into a singlemultiplexed 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 vectorrepresentation, this operation uses the same transformation  i for each token in the sequence beforeaveraging 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 amultiplexed 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 interferencebetween their representations.
Finally, this multiplexed representation, x 1: N , is used as input to theneural 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 theinput x 1: N .
To make a prediction for each input, one can explicitly disentangle h 1: N into N individualhidden representations, h 1 , . .
, h N , respectively.
We ﬁrst obtain each h i with a demultiplexingfunction # 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 hiddenrepresentation to prevent a substantial increase in the number of parameters and improve trainingefﬁciency.
We use two alternatives for the demultiplexing function # i :1.
MLP Demuxing This strategy employs N MLPs to generate each indices’ hidden representationas
h i =
MLP i ( h 1: N ) .
We use this method for both NLP and vision tasks.
Although this method isconceptually 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 , andtransformed 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 specialtokens, 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 consistsof a index token ✏ i in it’s
i ’th position while the remaining tokens are a special pad token ✏ pad .
Thepreﬁx sequences then take on the following pattern:
[Equation]
=
[Equation]
, ✏ pad ,
[Equation] . , ✏ pad ] [Equation] We then prepend each sequence x i of the input sequence with the corresponding preﬁx i .
The tuple ofprepended sequences is then passed to the multiplexing module.
When ﬁnally generating individualhidden representations, we use the corresponding hidden representation of each index token ✏ i as theindex embedding p i .
We use the Index Embeddings demultiplexing strategy on language tasks for the Transformer architecture.
The preﬁx tokens may implicitly enable the Transformer to do instance-speciﬁc computationswhen processing the multiplexed representation and further enable demultiplexing for large N .
[Figure] [Table]
+zOPJwAqdwDh5cQxXuoQYNYIDwDK/w5jw6L86787FozTnZzDH8gfP5A3FrjLY=</latexit> ) … <latexit sha1_base64="8yUu3vCxxhATJADq3W5rx9n/Zl0=">AAAB8nicbVDLSsNAFJ3UV42vqks3g0VwVRIRdSMW3biSCvYBaSyT6aQdOsmEmZtCCf0MNy4U6db/cO9G/Bunj4W2HrhwOOe+g0RwDY7zbeWWlldW1/Lr9sbm1vZOYXevpmWqKKtSKaRqBEQzwWNWBQ6CNRLFSBQIVg96N2O/3mdKcxk/wCBhfkQ6MQ85JWAkr9knCroMyONdq1B0Ss4EeJG4M1K8+rAvk9GXXWkVPpttSdOIxUAF0dpznQT8zDTkVLCh3Uw1SwjtkQ7zDI1JxLSfTVYe4iOjtHEolYkY8ET9XZGRSOtBFJjMiEBXz3tj8T/PSyG88DMeJymwmE4HhanAIPH4ftzmilEQA0MIVdzsimmXKELBfMk2T3DnT14ktZOSe1Y6vXeK5Ws0RR4doEN0jFx0jsroFlVQFVEk0RN6Qa8WWM/WmzWapuasWc0++gPr/QfHWJSV
[Equation]
/latexit
[Equation]
N
[Equation]
latexit sha1_base
[Equation]
"ORNheOk5xc+wg1dRwU8RY7kjHYQ
[Equation]
AAAB6HicbVDLTgJBEOzFF
[Equation]
IL9ehlIjHhRHYNUY9ELx4hkUcCGzI79MLI7OxmZtaEEL7AiweN8eonefNvHGAPClbSSaWqO91dQSK4Nq777eQ2Nre2d/K7hb39g8Oj4vFJS8epYthksYhVJ6AaBZfYNNwI7CQKaRQIbAfju7nffkKleSwfzCRBP6JDyUPOqLFSo9wvltyKuwBZJ15GSpCh3i9+9QYxSyOUhgmqdddzE+NPqTKcCZwVeqnGhLIxHWLXUkkj1P50ceiMXFhlQMJY2ZKGLNTfE1MaaT2JAtsZUTPSq95c/M/rpia88adcJqlByZaLwlQQE5P512TAFTIjJpZQpri9lbARVZQZm03BhuCtvrxOWpcV76pSbVRLtdssjjycwTmUwYNrqME91KEJDBCe4RXenEfnxXl3PpatOSebOYU/cD5/AG/njLU=</latexit> (
B-ORG O
O
O
O B-ORG
O I-ORG <latexit sha1_base64="lMTUNrAmOOi8mafVYacnuW/AHlU=">AAAB+XicbVDLSsNAFL2pr1pfUZduBovgqiQiKq6KblxJBfuANpbJdNIOnUzCzKRQQv7EjQtF3Pon7vwbJ20W2npg4HDOvdwzx485U9pxvq3Syura+kZ5s7K1vbO7Z+8ftFSUSEKbJOKR7PhYUc4EbWqmOe3EkuLQ57Ttj29zvz2hUrFIPOppTL0QDwULGMHaSH3b7oVYj/wgHWVPqXt9n/XtqlNzZkDLxC1IFQo0+vZXbxCRJKRCE46V6rpOrL0US80Ip1mllygaYzLGQ9o1VOCQKi+dJc/QiVEGKIikeUKjmfp7I8WhUtPQN5N5TrXo5eJ/XjfRwZWXMhEnmgoyPxQkHOkI5TWgAZOUaD41BBPJTFZERlhiok1ZFVOCu/jlZdI6q7kXtfOH82r9pqijDEdwDKfgwiXU4Q4a0AQCE3iGV3izUuvFerc+5qMlq9g5hD+wPn8Aer+TkQ
[Figure]
</latexit
[Figure]
(NER)
) using a shared task prediction head
3.3 Retrieval warm-up for multiplexing
For Transformer models, we ﬁnd that naively adding the multiplexing and demultiplexing layersto the model fail to converge.
This is likely because the gradients for individual instances from thetask loss get mixed up in the backward pass.
To overcome this, we propose Retrieval warm-up – aself-supervised pre-training task to promote the ability of DataMUX models at distinguishing theorder and the content of individual sequences in a multiplexed representation.
This task consists ofretrieving 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 toretrieve a token from one random sentence for each token position, yielding the following objective:
[Equation]
where h Ij is a demultiplexed hidden representation of the j -th token in a randomly selected sentencewith the index I ⇠ U [1 , N ] , generated using the methods described in Section 3.2.
We use thisobjective 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 retrievaltask 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 compareour 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 Section3.3.
Though there is no multiplexing done for B meaning this operation could be solved by simplycopying input tokens to the output layer)
we ﬁnd that the retrieval pre-training produces differencesin 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:
Token-level classiﬁcation: This evaluates models’ ability to perform token-level tasks on mul-
tiplexed 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ﬁcationtasks found in the General Language Understanding Evaluation (GLUE) benchmark Wang et al.
(2019): the sentiment classiﬁcation task SST-2 Socher et al.
(2013), the sentence similarity task QQP 1 , 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 andlearn 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-upon the Wikitext-103 dataset Merity et al.
(2017) as described in Section 3.3.
In addition, we alsocontinue to use the retrieval task as an auxiliary objective during task training.
The total loss is acombination 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 showsperformance across four sentence classiﬁcation tasks (MNLI, QNLI, QQP, SST2) and a token-levelclassiﬁcation task (NER).
We observe that it is possible to multiplex up to 40 instances and maintainreasonable performance.
For easier tasks like QQP, SST2 and QNLI, we observe that the dropin 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% for40 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 doesnot 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 noticemuch variance between ﬁne-tuning runs and plot variance for the "Hadamard + Index Embedding"setting in Section A.2. (R2)
Perfect multiplexing on the retrieval warm-up task for large N Figure 4b shows the testaccuracy on the retrieval warm-up task described in Section 3.3.
We ﬁrst note that across differentmultiplexing and demultiplexing strategies, models have a retrieval accuracy of nearly 100% for up to20 instances, demonstrating the surprising ability of T-MUX to multiplex perfectly for large N .
Notethat this task does not require any aggregation of context across the sequence and thereby is mucheasier than sentence classiﬁcation or token-level classiﬁcation tasks.
Therefore, performance on thiswarm-up task indicates a soft upper bound on the number of instances we can multiplex for sentenceand 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 speedfor processing ⇠ 20,000 samples on the MNLI dataset.
We use four different batch sizes for allthe conﬁgurations and take the max throughput (details in A.6).
Figure 4c shows that multiplexingincreases throughput many folds ( 18 x for 40 instances, 11 x for 20 instances).
T-MUX enablessuperior throughput as batch size can be effectively increased by a factor of N .
We would expectthe speedup to scale linearly with N , however a large N corresponds to more preﬁx tokens whichincreases the sequence length.
Therefore, having 40 instances leads to almost a 20 x speedup asopposed to the expected 40 x. Future work can potentially improve this speedup by designing bettermultiplexing and demultiplexing strategies.
4.3 Analysis
(A1)
The number of attention heads seems invariant to multiplexing To understand the roleof the number of attention heads in multiplexing, we train a variant of T-MUX with 2 self-attentionheads per layer.
We use T-MUX with the (Hadamard + Index Embed.) conﬁguration.
We ﬁnd that
[Table]
Demultiplexer
[Table]
[Figure]
reducing self-attention heads has minimal effect on performance on the retrieval warm-up task andachieves a retrieval accuracy of ⇠ 100% up to N = 20 .
We then ﬁnd that on token and sentence-levelclassiﬁcation tasks, T-MUX with 2 self-attention heads performs comparably to T-MUX with 12self-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 Transformers 2 to multiplex: a 12 layer with hidden size of 384 (12L / 384H), a 4 layer with hiddensize of 768 (4L / 768H).
Figure 5b shows that these smaller T-MUX models can also multiplex upto 20 instances with competitive performance.
Figure 4c illustrates the speedup from the smallermodels.
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 toonly 18 x for the full-sized T-MUX with 40 instances.(A3)
Performance varies more across different indices as N increases The prediction of aninstance is conditioned on the index of the instance.
Figure 6b illustrates performance on MNLI fordifferent indices across different choices of N for the (Hadamard + Index Embed.) conﬁguration.
Weobserve that performance varies more across different indices for large N (For N= 40 , performancevaries across ⇠ 10 percentage points).
4.4 Theoretical construction for multiplexing in self-attention models
To provide some theoretical explanation for why Transformers may be amenable to multiplexing, wealso detail below a construction for self-attention based neural networks which process multiplexed token embeddings in N independent subspaces.
Our construction relies on particular structuralproperties of singular spaces of linear transformations across layers – we provide a brief sketch below,with more details in Appendix A.1.
In a model with self-attention, the multi-head attention projects the queries, keys, and values h timeswith different, learned linear projections to d K , d K , and d V dimensions respectively.
Thus, given asequence of d -dimensional multiplexed token embeddings { w 1: N t } t 2 [1 ,L ] , each head looks like:
[Equation]
Now, w 1: N t = 1 N
P
Nk =1  k ( w kt ) , and let us assume each function  k can be trained to projecteach embedding into a subspace that is least linearly-dependent with the others.
So, if we deﬁneu ( 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 V i > W V i can be grouped into N non-overlapping subsets{ r (1)1 , . . .
, r (1) m } , · · · , { r ( N ), . . .
, r ( N ) m } , where r ( k ) ` are orthonormal vectors (since the Gramian isreal symmetric), and span the same input subspaces, denoted as D 1 , . .
, D N .
In this case, the vectorafter 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 richoperation 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]
where ⌧ ( k ) i,t,t 0 is a scalar only depending on the k -th input sequence.
Thus, the self-attention operationat each position can be seen as retrieving values based on the average of query-key similarity scoresof N sequences.
[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 eachhead to only focus on one input sequence, by setting ⌧ ( k 0 )
i,t,t 0 =
for all k 0
= k , which is easilyachievable by controlling singular values of W Qi or W K i .
More details of this construction areprovided in Appendix A.1.
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 onimage classiﬁcation as techniques for robustness and data augmentation Ramé et al.
(2021); Havasiet al. (2021)
; Soﬂaei et al.
These works implicitly employed the frontend layers of aconvolutional net as multiplexing layers and suggest that a convolutional neural network can learnat most 3-4 independent subnetworks concurrently Havasi et al.
Since Transformers can bemultiplexed for up to 40 instances without a severe performance drop, we explore the DataMUXscheme 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 thatfuture work on better multiplexing strategies can make them more viable.
Multiplexing for MLPs Figure 6a (solid lines) illustrates performance of DataMUX for MLPs withvarious multiplexing strategies, paired with the MLP demultiplexing strategy.
As a baseline, we showmultiplexing using the identity transformation before combining instances (MLP baseline).
Sincethis transformation does not preserve the order of the multiplexed instances, accuracy expectedlydecreases 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 baselinewithout multiplexing.
Since random orthogonal projections cannot separate inputs perfectly, wealso 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 6a (dashed lines).
Like in the MLP case, all methods shown use the MLP demultiplexing strategy.
We notice that thebaselines (CNN baseline) has a roughly N -fold performance drop because of the unidentiﬁability ofinput 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 relyon.
We therefore also explore N two-layer convolutional networks with a tanh activation as outmultiplexing transformations  i (CNN + Nonlinear).
We note that even with this transformation,the dimensionality of the multiplexed representation still remains equal to the dimensionality of asingle 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 .
Overall, multiplexing for MLPs and CNNs seems to be more challenging than Transformers asevidenced by the sharper performance drops with increasing N .
However, we believe these numbersare still non-trivial and show the potential for multiplexing in CNNs and MLPs, with multiplexingand 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 multipleinputs using a single ‘mixed’ representation.
We introduced DataMUX, a novel setting for trainingand inference with multiplexing, and demonstrate its viability using a variety of network architecturesand tasks.
Our results reveal the surprising ability of neural networks to predict on up to 40 inputinstances simultaneously using data multiplexing.
Our efforts have sought to show the potentialfor data multiplexing using methods that require limited computational overhead or added learnedparameters; achieving results that incur minimal degradation in performance while simultaneouslyincreasing system throughput dramatically.
While the scope of this paper is to demonstrate the general ability and potential for neural networksto use data multiplexing, there are several avenues for future research to investigate the limits andtheoretical implications of DataMUX.
Directions of particular interest include: large-scale pretraining, multi-modal processing, multi-lingual models, different mechanisms of multiplexing anddemultiplexing, and a more rigorous understanding of the architectural, data, and training conditionsenabling data multiplexing.