SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
S _ {i} ^ {(n)} = \underbrace {\left[ \sum_ {j = 1} ^ {L} \left(\sum_ {q = 1} ^ {N} v _ {j} ^ {(q)}\right) \phi \left(\sum_ {w = 1} ^ {N} k _ {j} ^ {(w)}\right) ^ {T} \right]} _ {A _ {s}} \times \phi \left(q _ {i} ^ {(n)}\right) = \sum_ {j = 1} ^ {L} \sum_ {q = 1} ^ {N} v _ {j} ^ {(q)} \left(\phi \left(\sum_ {w = 1} ^ {N} k _ {j} ^ {(w)}\right) ^ {T} \phi \left(q _ {i} ^ {(n)}\right)\right) \tag {16}
$$
$$
\begin{array}{l} \stackrel {P} {\approx} \sum_ {j = 1} ^ {L} \sum_ {q = 1} ^ {N} v _ {j} ^ {(q)} \exp \left(\langle \sum_ {w = 1} ^ {N} k _ {j} ^ {(w)}, q _ {i} ^ {(n)} \rangle / \sqrt {D}\right) \stackrel {H} {\approx} \sum_ {j = 1} ^ {L} \sum_ {q = 1} ^ {N} v _ {j} ^ {(q)} \exp \left(\langle k _ {j} ^ {(n)}, q _ {i} ^ {(n)} \rangle / \sqrt {D}\right) (17) \\ = \underbrace {\sum_ {j = 1} ^ {L} v _ {j} ^ {(n)} \exp \left(\left\langle k _ {j} ^ {(n)} , q _ {i} ^ {(n)} \right\rangle / \sqrt {D}\right)} _ {\text {u n n o r m a l i z e d} o _ {i} \text {o f c h a n n e l} n} + \underbrace {\sum_ {q \neq n} \sum_ {j = 1} ^ {L} v _ {j} ^ {(q)} \exp \left(\left\langle k _ {j} ^ {(n)} , q _ {i} ^ {(n)} \right\rangle / \sqrt {D}\right)} _ {\text {n o i s e i n s e p a r a t e p r o t e c t e d s u b s p a c e}}. (18) \\ \end{array}
$$
The output contains the  $n^{th}$  channel together with noise. The operation  $A_{s} \times \phi(q_{i}^{(n)})$ , which takes  $\mathcal{O}(LDR)$ , must be repeated  $N$  times to produce outputs for all  $N$  channels, causing a bottleneck.
Simultaneous superposition of queries and value-key tensor products using a 2D grid. Finally, we combine the two previously described approaches to encode the superposition channels in a 2D grid of size  $N \times M$ . We multiply a constant matrix  $(A_s)$  with features derived from a superposition of queries to get the superposition vector  $S_i^{(n)}$
$$
S _ {i} ^ {(n)} = \underbrace {\left[ \sum_ {j = 1} ^ {L} \sum_ {u = 1} ^ {M} \binom {N} {q = 1} v _ {j} ^ {(u , q)}\right) \phi \left(\sum_ {w = 1} ^ {N} k _ {j} ^ {(u , w)}\right) ^ {T}} _ {\text {c o n s t r u c t A} _ {s} \text {i n} \mathcal {O} (L M D (R + N))} \times \underbrace {\phi \left(\sum_ {t = 1} ^ {M} q _ {i} ^ {(t , n)}\right)} _ {\text {c o n s t r u c t} \forall i, n \text {i n} \mathcal {O} (L N D (R + M))} \tag {19}
$$
Computing the multiplication  $\times \forall i, n$  takes  $\mathcal{O}(LNDR)$ . If we set  $M = N$ , we can evaluate  $S_{i}^{(n)} \forall i, n$  using only  $\mathcal{O}(LNDR + LN^2D)$  instead of the usual  $\mathcal{O}(LDR \cdot N^2)$ . Thus, one may achieve a speedup of  $\mathcal{O}(\min(\sqrt{N^2}, R))$  compared to FAVOR+. Since  $R$  is normally in the hundreds [15], we can assume improvements of  $\mathcal{O}(\sqrt{N^2})$  for reasonably large  $N^2 = M \cdot N$ . Eq. (19) simplifies to:
$$
\begin{array}{l} S _ {i} ^ {(n)} = \sum_ {j, u, q} v _ {j} ^ {(u, q)} \left(\phi \left(\sum_ {w} k _ {j} ^ {(u, w)}\right) ^ {T} \phi \left(\sum_ {t} q _ {i} ^ {(t, n)}\right)\right) \approx \sum_ {j, u, q} v _ {j} ^ {(u, q)} \exp \left(\frac {\langle \sum_ {w} k _ {j} ^ {(u , w)} , \sum_ {t} q _ {i} ^ {(t , n)} \rangle}{\sqrt {D}}\right) (20) \\ \approx \sum_ {j = 1} ^ {L} \sum_ {u = 1} ^ {M} \sum_ {q = 1} ^ {N} v _ {j} ^ {(u, q)} \exp \left(\langle k _ {j} ^ {(u, n)}, q _ {i} ^ {(u, n)} \rangle / \sqrt {D}\right) (21) \\ = \sum_ {u = 1} ^ {M} \underbrace {\sum_ {j = 1} ^ {L} v _ {j} ^ {(u , n)} \exp \left(\frac {\langle k _ {j} ^ {(u , n)} , q _ {i} ^ {(u , n)} \rangle}{\sqrt {D}}\right)} _ {\text {u n n o r m a l i z e d} o _ {i} \text {o f c h a n n e l} (u, n)} + \underbrace {\sum_ {q \neq n} \sum_ {u = 1} ^ {M} \sum_ {j = 1} ^ {L} v _ {j} ^ {(u , q)} \exp \left(\frac {\langle k _ {j} ^ {(u , n)} , q _ {i} ^ {(u , n)} \rangle}{\sqrt {D}}\right)} _ {\text {n o i s e i n s e p a r a t e p r o t e c t e d s u b s p a c e}}. (22) \\ \end{array}
$$
4.2 Integrating FAVOR+S into MIMOFormer
As is apparent in Eq. (19), the query superposition is along a different axis  $(M)$  than the key and value superpositions  $(N)$ . The output of attention, however, is only superposed along a single axis  $(M)$ . To be able to set up superpositions along both axes  $(M$  and  $N$ ) at the next layer, we require all channels (i.e., keys, queries, and values) in separation, i.e., not superposed, at the interface between MIMOFormer layers. We present two variants of MIMOFormer with different speedups.
The first computes in superposition exclusively during the attention mechanism. The individual tokens of the channel  $(n,m)$  are directly retrieved from  $S_{i}^{(n)}$  by unbinding with the key  $\tilde{a}^{(n,m)} = a^{(n,m)}$ , and the remaining computational steps within FAVOR+S are performed separately.
The second (and faster) MIMOFormer instance additionally performs the concatenation, the linear layer, as well as the MLP in superposition (shown in Figure 3). Unlike in the first variant, the skip connection around the attention block must account for the introduced superposition. To allow a potential embedding dimension mismatch, we instantiate two different sets of randomly drawn bipolar keys: one for the skip connection and post-MLP unbinding, and one for FAVOR+S binding. All keys are frozen during training; it is up to the trainable weights in the linear layer after concatenation to find the relationship between the binding and unbinding keys.
The function  $\phi$  in the self-attention block consists of an  $R = 256$  dimensional projection and a ReLU activation [15]. Appendix C provides theoretical justification for using ReLU in  $\phi$  and its benefits for MIMOFormer. Empirically, ReLU shows better numerical stability than unbiased softmax FAVOR+.
This section evaluates the proposed MIMONets on various model architectures and benchmarks. Appendix E and Appendix F describe the experimental setup for MIMOConv and MIMONets, respectively. All experiments are repeated five times with a different random seed. We report the mean and standard deviations of accuracy to account for variability in training.
5.1 MIMOConv
CIFAR10 and CIFAR100. Our main baseline, which we adapt to MIMOConv, is a reproduced WideResNet-28-10 [11], i.e., a 28-layer CNN with the number of feature maps of each convolutional layer enlarged by  $10\times$ . Moreover, as a stronger baseline, we include the isometry regularization in the training of WideResNet-28-10, and call the resulting network WideIsoNet-28-10. Both baselines demand 5.251 GMACs (Giga multiply-accumulate operations) per sample. MIMOConv's inference complexity per sample is 5.335 GMACs for  $N = 1$ , 2.667 GMACs for  $N = 2$ , and 1.334 GMACs for  $N = 4$ ; hence, we get a  $\approx N\times$  speedup despite not accelerating the first and last layer. See Appendix E.
Table 1 shows the accuracy of the static MIMOConv, which is exclusively trained to support either 1, 2, or 4 channels. The MIMOConv with one channel ( $N = 1$ ) outperforms both baselines, which may be attributed to regularizing effects of the key-value binding. MIMOConv with  $N = 2$  channels still outperforms WideResNet-28-10, while reducing the inference complexity by  $2 \times$ . The complexity can be further reduced by increasing the number of superpositions to  $N = 4$  at a slight accuracy drop of  $\leq 3.18\%$ , compared to WideResNet-28-10.
Next, we evaluate the dynamic partitioning of the superposition channels to select a speed-accuracy operating point instantaneously, which is a main feature of our approach and sets it apart from other static approaches that opt for a fixed performance point like model downsizing, quantization aware training, and pruning (see Appendix A for a discussion). We set up a model with four channels, but evaluate its performance in different configurations: a fast (4 inputs/pass), a normal (2 inputs/pass), and a slow mode (1 input/pass). The fast mode maps each input to one channel; the medium mode distributes two inputs over pairs of channels; and the slow mode uses all channels for the same input. The models are trained on  $80\%$  of the batches in fast mode and on  $20\%$  of the batches in slow mode. Appendix E provides more details on the trade-off between fast and slow mode training. As Table 1 shows, a single dynamic model can seamlessly switch between operation points while maintaining a high accuracy compared to the static models ( $\leq 1.82\%$  drop).
Table 1: Average accuracy  $(\%)$  of WideResNet-28-10 variants and our MIMOConv. Static models are trained to process  $N$  inputs in one pass, speeding up inference by  $N\times$ . Dynamic models are trained with a variable number of inputs  $(N = 1 - 4)$ , and can process a variable number of inputs per pass. We report the average accuracy  $\pm$  the standard deviation over five runs with different seeds.
# inputs/pass	CIFAR10	CIFAR100
1	2	4	1	2	4
WideResNet-28-10	96.82±0.06	n.a.	n.a.	81.62±0.07	n.a.	n.a.
WideIsoNet-28-10	97.31±0.11	n.a.	n.a.	82.38±0.20	n.a.	n.a.
MIMOConv static (N=1)	97.49±0.08	n.a.	n.a.	83.19±0.17	n.a.	n.a.
MIMOConv static (N=2)	n.a.	96.93±0.13	n.a.	n.a.	82.30±0.19	n.a.
MIMOConv static (N=4)	n.a.	n.a.	95.58±0.23	n.a.	n.a.	78.44±0.30
MIMOConv dynamic (N=1-4)	97.13±0.11	96.41±0.14	95.43±0.07	82.52±0.09	80.48±0.08	78.19±0.10
The detailed ablation study in Appendix E makes the following findings: (1) isometry regularization improves accuracy for any number of channels; (2) training MIMOConv for more epochs closes the performance gap to the single-input baseline; (3) an appropriate number of feature maps (32 or 64) in the first layer stabilizes training; and (4) the binding keys can be frozen during training without loss (whereas unbinding keys are never frozen).
MNIST and SVHN. Figure 4 compares MIMOConv with DataMux [13] on the MNIST dataset. Even with a trivial downsizing for fair comparison from a 28-layer very-wide  $(10\times)$  ResNet to a 10-layer narrow  $(1\times)$  network, MIMOConv scales much better to high superposition channels  $(N)$  than DataMUX does. Indeed, our model shows an accuracy of  $80.4\%$  against their  $52.9\%$  in case of  $N = 16$  superposition channels (highest number of channels reported by DataMUX for vision tasks), despite being computationally cheaper (0.47 MMAC/s vs. 0.65 MMAC/s). Also, DataMux's binding overhead results in a mere  $1.35\times$  reduction in MACs compared to our  $10.9\times$
as  $N$  goes from 1 to 16. Ergo, our method scales better in accuracy and throughput as  $N$  increases.
Figure 4: Classification accuracy  $(\%)$  on MNIST for downsized model.
Finally, we tested MIMOConv on the SVHN dataset. Despite limited hyperparameter tuning, MIMOConv achieves a high accuracy of  $97.17\%$  ( $N = 1$ ), and can maintain its performance with larger superpositions (97.05% and 96.84% for  $N = 2$  and  $N = 4$ , respectively).
5.2 MIMOFormer
LRA. We evaluate MIMOFormer on five tasks from LRA [14], and compare against the vanilla Transformer [24] and the Performer [15] using FAVOR+ attention with ReLU projection. Moreover, we also consider wide Transformer variants [26], consisting of only one layer but as many heads as their deep counterparts. Task-specific architectures and training hyperparameters are kept the same for the Performer and the MIMOFormer (see Appendix F).
Table 2: Test accuracy  $(\%)$  on the long range arena (LRA). MIMOFormer uses an equal number of query superpositions  $(M)$  and value-key superpositions  $(N)$ , i.e.,  $N = M$ . Computation in superposition is performed either in attention only (att.) or in both attention and MLP (att.+MLP).  $L$  is the number of layers,  $H$  the number of heads, and * indicates curriculum learning.
ListOps	Text	Retrieval	Image	Pathfinder	Avg.
Deep models	L=6, H=8	L=6, H=8	L=4, H=4	L=3, H=4	L=4, H=8
Transformer [24]	36.37	64.27	57.46	42.44	71.40	53.39
Performer [15]	18.01	65.40	53.82	42.77	77.05	51.41
Performer (reproduced)	38.94±0.23	65.70±0.31	81.58±0.18	40.14±0.86	73.82±0.78	60.04±0.47
MIMOFormer (N=2, att.)	38.08±0.21	65.00±0.28	79.37±0.81	38.21±0.63	72.36±0.54	58.61±0.49
MIMOFormer (N=2, att.+MLP)	37.65±0.33	64.39±0.22	76.02±0.27	33.85±0.55	67.98±0.47	55.98±0.37
MIMOFormer (N=4, att.)	37.22±0.33	64.59±0.14	60.99±0.96	28.16±0.08	55.50±4.95	49.29±2.91
MIMOFormer (N=4, att.)*	37.64±0.73	64.46±0.15	74.38±0.82	30.52±0.77	67.10±0.45	54.82±0.58
MIMOFormer (N=4, att.+MLP)	17.74±0.63	60.71±5.14	72.20±0.28	24.01±0.47	50.33±0.16	45.00±1.34
Wide models	L=1, H=48	L=1, H=48	L=1, H=16	L=1, H=12	L=1, H=32
Performer (reproduced)	39.40±0.51	65.73±0.32	83.67±0.25	41.67±0.44	74.11±0.33	60.93±0.37
MIMOFormer (N=2, att.)	38.90±0.53	65.39±0.18	81.27±0.28	40.25±0.21	73.51±0.23	59.86±0.29
MIMOFormer (N=2, att.+MLP)	37.59±0.17	64.64±0.25	78.30±0.32	36.69±0.76	68.22±0.18	57.09±0.34
MIMOFormer (N=4, att.)	37.71±0.24	64.22±0.14	74.99±0.36	35.43±0.60	69.52±0.40	56.37±0.35
MIMOFormer (N=4, att.)*	37.68±0.36	64.56±0.25	76.37±0.50	35.53±0.48	73.37±0.22	57.50±0.36
MIMOFormer (N=4, att.+MLP)	18.52±0.98	63.53±0.12	74.30±0.26	26.54±0.28	56.33±0.17	47.84±0.36
Owing to an improved training setup, our replicated deep and wide Performer baselines substantially outperform the results reported in [14] (see Table 2). Moreover, MIMOFormer enables accurate computation in superposition for both deep and wide attention models. The performance drop is less pronounced in wide models (only  $1.07\%$  drop compared to Performer with  $N = 2$ , att.), which may be attributed to the larger number of heads, increasing the effective dimension  $(D_{\mathrm{tot}} = H \cdot D_{\mathrm{head}})$ .
When computing both attention and the MLP in superposition (att.+MLP), we observe better scaling (in  $N$ ) for wide models. Also, MIMOFormer reduces the gap to the baseline as the number of epochs increases (see Appendix F).
To stabilize training in the case of  $N = 4$ , we implemented a curriculum training procedure where the number of superpositions is reduced to  $N' = N / 2$  during a warmup phase (1/6th of the training steps), improving the average accuracy of MIMOFormer in both wide and deep models.
Comparing against a reproduced DataMUX [13], MIMOFormer (att.) outperforms it on ListOps  $(38.08\%$  vs.  $30.54\%$  accuracy) when using models of similar size and  $N = 2$ , see Appendix F.
Synthetic sequence modeling. Table 3 reports the accuracy on two synthetic sequence modeling tasks, which Transformer alternatives such as S4 [27] have difficulties solving [12]. On these more nuanced NLP tasks, the accuracy of DataMUX [13] drops to  $20.04\%$  and  $6.06\%$  for  $N = 2$  despite significant efforts in training, while MIMOFormer, at a score
Table 3: Accuracy (%) on synthetic sequence modelling.
Architecture	Attention	Associative recall [12]	Induction head [12]
Transformer	Softmax	98.48±1.87	100±0.0
Performer	FAVOR+	96.32±6.26	31.58±33.67
MIMOFormer (N=2, att.)	FAVOR+	96.52±2.79	99.40±0.13
MIMOFormer (N=2, att.)	DPFP [25]	93.64±12.66	98.56±0.86
DataMUX (N=2) [13]	Softmax	20.04±1.72	6.06±2.24
of  $96.52\%$  and  $99.40\%$  respectively, succeeds. We attribute this difference in performance to attention score blurring in DataMux, discussed in Appendix F. Contrastingly, our method converges to exact attention without blurring. It is versatile and can be adjusted to other linear Transformers such as DPPF [25], achieving a score of  $93.64\%$  and  $98.56\%$ .
So far, superposition principles have been applied in order to store the weights of multiple models in a single neural network [28-30], to circumvent catastrophic forgetting in continual learning [31, 32], and to render symbolic reasoning tractable [33]. To address privacy concerns when running remote inference, single inputs were bound with random channels to implement pseudo-encryption [34]. Recently in [35], HRR was used to define an unconventional version of self-attention, whose attention scores are processed to a diagonal matrix. The value vectors are scaled according to their importance in the sequence instead of being combined in a weighted sum. In contrast to us, none of these works superpose multiple inputs into a data structure to speed up computation.
In [19, 20], an ensemble of CNN models was fit into one network. However, by only broadcasting a single input over the channels and by averaging all the outputs, this approach collapses to a single-input-single-output (SISO) network. On the contrary, we explore using protected channels for different inputs at inference, resulting in an actual multiple-input-multiple-output (MIMO) network.
There has also been a line of work to accelerate Transformers using inputs in superposition [13] [36]. DataMux [13] claims to retain high performance for language understanding tasks, even when using up to 40 inputs in superposition. However, none of the reported tasks require attention layers at all [37]. In Section 5.2 we show failure of their method when actual attention is required (see also Appendix F). MUX-PLMs [36] improves on DataMux with contextual binding and replaces token prefixes with unbinding keys, but does not address the blurry attention mechanism. In contrast to DataMUX and MUX-PLMs, our work approximates true attention and our theoretical derivations show convergence to actual dot-product attention as the dimension of attention projections increases, giving us an even stronger case for applicability to large language models.
We present MIMONets that simultaneously process multiple inputs by performing computation in superposition. MIMONets bind arbitrary inputs with high-dimensional keys, which projects them to orthogonal subspaces that, together with near-isometric subfunctions, guarantee low interference through all nonlinear layers. Unbinding with (learned) keys can safely retrieve information on individual channels. We provide two MIMONets instances, MIMOConv and MIMOFormer, that show the effectiveness of computation in superposition through two dominant operations in neural network architectures: convolution and attention. Further investigations could explore the MIMO-capability of architectures that contain additional nonlinearities (e.g., max-pooling) and use different input modalities. MIMONets could be suitable candidates to accelerate dynamically and on-demand the inference of foundation models [38].
This work is supported by the Swiss National Science foundation (SNF), grant 200800. We thank Dario Bolli for conducting initial experiments and Aleksandar Terzić for helping with ablation studies.
[37] M. Hassid, H. Peng, D. Rotem, J. Kasai, I. Montero, N. A. Smith, and R. Schwartz, “How much does attention actually attend? Questioning the importance of attention in pretrained transformers,” in Findings of the Association for Computational Linguistics: EMNLP 2022, 2022, pp. 1403–1416.
[38] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill et al., "On the opportunities and risks of foundation models," arXiv preprint arXiv:2108.07258, 2021.