SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nicolas Menet $^{1,2*}$
menetn@ethz.ch
Michael Hersch<sup>1,2</sup>
her@zurich.ibm.com
Geethan Karunaratne
kar@zurich.ibm.com
Luca Benini²
lbenini@iis.ee.ethz.com
Abu Sebastian<sup>1</sup>
ase@zurich.ibm.com
Abbas Rahimi
abr@zurich.ibm.com
<sup>1</sup>IBM Research - Zurich, <sup>2</sup>ETH Zurich
With the advent of deep learning, progressively larger neural networks have been designed to solve complex tasks. We take advantage of these capacity-rich models to lower the cost of inference by exploiting computation in superposition. To reduce the computational burden per input, we propose Multiple-Input-Multiple-Output Neural Networks (MIMONets) capable of handling many inputs at once. MIMONets augment various deep neural network architectures with variable binding mechanisms to represent an arbitrary number of inputs in a compositional data structure via fixed-width distributed representations. Accordingly, MIMONets adapt nonlinear neural transformations to process the data structure holistically, leading to a speedup nearly proportional to the number of superposed input items in the data structure. After processing in superposition, an unbinding mechanism recovers each transformed input of interest. MIMONets also provide a dynamic trade-off between accuracy and throughput by an instantaneous on-demand switching between a set of accuracy-throughput operating points, yet within a single set of fixed parameters. We apply the concept of MIMONets to both CNN and Transformer architectures resulting in MIMOConv and MIMOFormer, respectively. Empirical evaluations show that MIMOConv achieves  $\approx 2 - 4\times$  speedup at an accuracy delta within  $[+0.68, -3.18]\%$  compared to WideResNet CNNs on CIFAR10 and CIFAR100. Similarly, MIMOFormer can handle 2-4 inputs at once while maintaining a high average accuracy within a  $[-1.07, -3.43]\%$  delta on the long range arena benchmark. Finally, we provide mathematical bounds on the interference between superposition channels in MIMOFormer. Our code is available at https://github.com/IBM/multiple-input-multiple-output-nets.
Driven by the successes of deep learning in image and natural language processing tasks, increasingly large neural network models have been developed to reach state-of-the-art performance [1-4]. These large models, however, increase computational complexity in terms of operation count for every event of input processing. One viable option to reduce the computational cost of processing per input is to create a compositional data structure where a variable number of input items (i.e., values) can be bound to corresponding protection keys, creating key-value pairs that can coexist and be processed
*Research conducted at IBM Research – Zurich.
Figure 1: MIMONets simultaneously pass multiple inputs through a nonlinear function, e.g., a deep convolutional network (on top) or a Transformer (on bottom). Input samples are bound with high-dimensional keys to project the samples into quasi-orthogonal subspaces. The results of the individual samples are retrieved at the end of the network by unbinding with corresponding keys.
concurrently. This variable-sized data structure can be represented by fixed-width distributed representations in vector-symbolic architectures (VSAs) [5-7]. In VSAs, the composition of different items in the data structure is based on functional compositionality (i.e., key-value binding), which yields a dimensionality-preserving distributed representation, rather than concatenative compositionality. Interestingly, the resulting fixed-width distributed data structure can be transformed by a one-time application of a function, whereby all input items are jointly transformed, leading to holistic transformation or computation in superposition [8-10]. This concept of computation in superposition can reduce the effective number of operations per input by a factor of the number of input items in the data structure, because the function is applied to the data structure holistically without decomposing the constituent items for individual transformations. However, processing the VSA data structure via computation in superposition has so far been limited to linear maps [8-10].
Motivated by these observations, we make the following contributions:
(1) We introduce a principled and transparent approach to Multiple-Input-Multiple-Output Neural Networks (MIMONets) based on VSA, enabling computation in superposition for highly nonlinear transformations in neural networks (Section 2). The MIMONets concept can be applied to various architectures, embracing the rich capacity provided by increasingly large models. The resulting network can handle many inputs at once, thus reducing the computational cost per input. We describe and overcome the challenges of computation in superposition, which originates from nonlinear interference of inputs in separate superposition channels.
(2) We propose MIMOConv, a realization of MIMONets for deep convolutional neural network (CNN) architectures (Section 3). We provide several strategies for mitigating interference between different superposition channels, including a novel locality-preserving binding operation (PWHRR) and isometry-inducing regularization. Empirical evaluations on CIFAR10 and CIFAR100 show that a MIMOConv built on a WideResNet-28-10 [11] can process concurrently two inputs in superposition  $(\approx 2\times$  speedup) even with a slightly higher accuracy  $(0.11 - 0.68\%$  gain), and four inputs  $(\approx 4\times$  speedup) with a marginal drop  $(1.24 - 3.18\%)$  (see Section 5.1).
(3) We further extend the concept of MIMONets to Transformer architectures, where the calculation of attention scores poses additional challenges for the computation in superposition paradigm. To this end, we propose MIMOFormer, which relies on a 2D grid binding scheme for computing attention in superposition (Section 4). We derive probabilistic tail bounds on the distortion caused by inter-channel interference and show that our method converges to noise-free attention in the limit of high dimension. Our method succeeds empirically ( $\geq 96.52\%$  accuracy) at synthetic sequence modelling tasks [12], while previous work [13] fails ( $\leq 20.04\%$ ). We also provide evaluations on the long range arena (LRA) dataset [14] in Section 5.2 using a MIMOFormer that is based on the Performer architecture [15]. Compared to the Performer, MIMOFormer maintains a high average accuracy with a marginal drop of  $1.07\%$  and  $3.43\%$  when handling two and four inputs at once, respectively.
(4) MIMONets allow a dynamic trade-off at inference time between accuracy and speed, i.e., they offer an instantaneous on-demand switching between accuracy-throughput operating points using  $a$
single set of fixed parameters. Experimental results in Section 5.1 show that a dynamic MIMOConv can seamlessly operate at various modes ( $\approx 1 - 4 \times$  speedup) while maintaining a high accuracy compared to the best static models ( $\leq 1.82\%$  drop).
2 MIMONets enabling Computation in Superposition
The central idea behind MIMONets (see Figure 1) is to simultaneously pass multiple inputs as a superposition through a nonlinear function  $f_{\theta}$  parameterized by neural network weights  $\theta$ . We isolate the individual inputs into separate protected channels by binding them with protection keys resulting in a key-value data structure [5-7].
Concept. Assuming two inputs  $(x^{(1)}$  and  $x^{(2)})$ , which can be generic embeddings either from images or natural language, we define a unique high-dimensional key  $(a^{(1)}$  and  $a^{(2)})$  for each protected channel, drawn randomly at initialization. Owing to the Blessing of Dimensionality, randomly drawn high-dimensional vectors are quasi-orthogonal with high probability (see Appendix A). Consequently, binding  $(\odot)$  the inputs with these keys yields quasi-orthogonal key-value pairs  $(x^{(1)}\odot a^{(1)}$  and  $x^{(2)}\odot a^{(2)})$ , which enables one to superpose the pairs with low interference:
$$
s = a ^ {(1)} \odot x ^ {(1)} + a ^ {(2)} \odot x ^ {(2)}. \tag {1}
$$
As discussed in Appendix A,  $s$  admits a noisy retrieval of  $x^{(1)}$  and  $x^{(2)}$  through unbinding:
$$
\hat {x} ^ {(1)} = a ^ {(1)} \circledast s = a ^ {(1)} \circledast a ^ {(1)} \odot x ^ {(1)} + a ^ {(1)} \circledast a ^ {(2)} \odot x ^ {(2)} = x ^ {(1)} + n o i s e \tag {2}
$$
To accelerate computing, inspired by the above-mentioned noisy retrieval, we pass the superposition  $s$  through a nonlinear function  $f_{\theta}$  with parameters  $\theta$ , such as a neural network, before retrieval. The quasi-orthogonality of the bound inputs allows processing each in a separate protected subspace—all with a single function call. To be able to recover the first processed sample  $f_{\theta}(x^{(1)})$  from  $f_{\theta}(s)$ , we aim to find an unbinding key  $\tilde{a}^{(1)}$  for which
$$
\begin{array}{l} \tilde {a} ^ {(1)} \circledast f _ {\theta} (s) \approx \tilde {a} ^ {(1)} \circledast f _ {\theta} \left(a ^ {(1)} \odot x ^ {(1)}\right) + \tilde {a} ^ {(1)} \circledast f _ {\theta} \left(a ^ {(2)} \odot x ^ {(2)}\right) (3) \\ \approx f _ {\theta} \left(x ^ {(1)}\right) + \tilde {a} ^ {(1)} \circledast f _ {\theta} \left(a ^ {(2)} \odot x ^ {(2)}\right). (4) \\ \end{array}
$$
The first approximation holds exact for linear  $f_{\theta}$ . As discussed in Section 3, a nonlinear  $f_{\theta}$  can still be encouraged to allow such an approximation through appropriate weight regularization techniques and well-suited activation functions. Further, by optimizing over unbinding keys  $(\tilde{a}^{(i)})$ , the second estimation (Eq. (4)) can be achieved. Consequently, matching binding and unbinding keys  $(a^{(i)}$  and  $\tilde{a}^{(i)})$  that confirm the approximation (Eq. (3) and (4)) set up a protected channel through the nonlinear function  $f_{\theta}(s)$ . Appendix A lists the design choices of the adopted VSA, which define the operations of key-value binding and unbinding, for all MIMONet variants presented in this work. In the case of image embeddings, we use circular convolution [6] for binding and Matrix Binding of Additive Terms (MBAT) [16] for unbinding. In the case of sequence tokens, we bind and unbind using the Hadamard product [17]. Binding and unbinding keys are always data-independent, i.e., they depend only on the index of the protected channel. See [18] for alternative binding and unbinding options.
Dynamic Inference. Setting up  $N$  protected channels through a neural network  $f_{\theta}$  gives almost a speedup of  $N \times$  due to most computations taking place in superposition. However, as is explored empirically, increasing  $N$  adds inter-channel noise leading to a decrease in predictive accuracy. If a fixed trade-off is unsatisfactory, one can build a dynamic model capable of running a superposition of one up to  $N$  different inputs. By inserting the same input into multiple channels and averaging the output, one effectively forms an in-network ensemble, similar to [19,20]. Using all protected channels for different inputs leads to a fast but less accurate model, whereas using all protected channels for the same input yields a slower but accurate ensemble model. By partitioning the superposition channels on demand, arbitrary configurations in between may be reached. Note that our method can instantaneously adapt to the current computational demand, without loading different model weights from the main memory. To perform across slow and fast configurations, the model should randomly switch between them during training. See Appendix A for a more detailed explanation.
Figure 2: MIMOConv configured with  $N = 3$  channels. The input images are passed individually through the first convolutional layer before binding each feature value with a channel-specific high-dimensional key. The key-value pairs are superposed yielding a dimensionality-preserving composition and passed through the rest of the CNN layers. The output is unbound with corresponding keys, and the unbound representations are classified separately with a shared fully-connected layer.
3 MIMOConv
This section presents how the MIMONets concept, introduced in Section 2, can be applied to construct a multiple-input-multiple-output CNN (MIMOConv) capable of simultaneously processing multiple images in superposition. The MIMOConv architecture is shown in Figure 2. Multiple input samples  $(N)$  are passed through the first convolutional layer, bound with a unique high-dimensional key based on Holographic Reduced Representations (HRR) [6], and superposed by element-wise addition. After passing the superposed tensors through the network's main CNN layers, we obtain a combined feature vector with information on all inputs. By unbinding with separately learned keys based on MBAT [16], which amounts to a matrix multiplication, we extract the individual processed information, which is then passed through a final fully-connected layer to produce logits for classification. In the following, we introduce our three main contributions that lead to a highly accurate MIMOConv.
Augmenting CNNs with locality-preserving variable bindings. We embrace a principled and transparent binding mechanism from HRR. Accordingly, binding is performed using circular convolution with a binding key of dimension  $D$ , drawn from an i.i.d. Gaussian distribution with zero mean and  $1 / D$  variance. Instead of convolving the flattened image tensor, we repeatedly apply circular convolution between the binding key and each pixel volume spanning across the feature maps  $(D \times 1 \times 1)$ . This binding operation, which we call position-wise HRR (PWHRR), is translation equivariant and maintains locality, an essential property for subsequent layers with limited receptive fields. More concretely, binding and unbinding are performed as
$$
\left(a ^ {(k)} \odot x ^ {(k)}\right) _ {:, w, h} = a ^ {(k)} * x _ {:, w, h} ^ {(k)} \tag {5}
$$
$$
\left(\tilde {a} ^ {(k)} \circledast h\right) _ {:, w, h} = \tilde {a} ^ {(k)} \cdot h _ {:, w, h}, \tag {6}
$$
with image tensors  $x^{(k)} \in \mathbb{R}^{D \times W \times H}$ , hidden representation  $h \in \mathbb{R}^{D' \times W \times H}$ , binding key  $a^{(k)} \in \mathbb{R}^D$  and unbinding key  $\tilde{a}^{(k)} \in \mathbb{R}^{D' \times D'}$ . Here,  $D, D', W$ , and  $H$  denote the hidden dimension at binding, the hidden dimension at unbinding (generally differs from  $D$ ), image width, and image height, respectively. * is the circular convolution,  $\cdot$  the matrix multiplication, and  $k$  indexes the superposition channel. Unbinding is applied after the global (average) pooling to reduce computational costs. The binding keys can be either learned or fixed during training (see ablation study in Appendix E).
Embracing high dimensional embeddings. According to the Blessing of Dimensionality (see Appendix A), random vectors quickly become quasi-orthogonal as their dimension increases. To reduce interference between protected channels, we increase the number of feature maps by adopting Wide Residual Networks [11], the most commonly used CNN architecture to achieve state-of-the-art accuracy on CIFAR100 [21]. The input tensors are passed individually through the first convolutional layer before being superposed in a suitably high dimensional space. We set the number of feature maps after the first convolutional layer to  $D = 64$ . This is  $4 \times$  more than the standard Wide-ResNet-28 [11], which results in improved training stability at a marginally higher compute cost (see Section 5.1).
Encouraging isometric layers. We aim to preserve the quasi-orthogonality of our protected channels as the superposition is passed through many network layers. To that end, residual connections are used and each subfunction  $g_{\theta}$  of type (strided) spatial convolution or activation function is made approximately inner-product preserving, i.e.,
$$
\langle g _ {\theta} (x), g _ {\theta} (y) \rangle \approx \langle x, y \rangle . \tag {7}
$$
Inspired by [22], we deploy a regularization to the CNN weights and use a parametric ReLU [23] activation function, a learnable affine combination between identity and ReLU. Those adjustments lead to a near-isometric network. The regularization for the individual CNN layers is determined by
$$
L (W) = \frac {\gamma}{2} \| C o n v (W, W) - \delta^ {C _ {o}} \| _ {F} ^ {2} \quad \text {w h e r e} \quad \delta_ {:; j, l} ^ {C _ {o}} = I _ {C _ {o} \times C _ {o}} \cdot \mathbb {1} _ {j, l = \lfloor \frac {k}{2} \rfloor}, \tag {8}
$$
$$
L \left(W ^ {T}\right) = \frac {\gamma}{2} \left\| C o n v \left(W ^ {T}, W ^ {T}\right) - \delta^ {C _ {i}} \right\| _ {F} ^ {2} \quad \text {w h e r e} \quad \delta_ {:,: j, l} ^ {C _ {i}} = I _ {C _ {i} \times C _ {i}} \cdot \mathbb {1} _ {j, l = \lfloor \frac {k}{2} \rfloor}, \tag {9}
$$
with  $\gamma$  as hyperparameter. We use  $L(W)$  if  $C_i > C_o$ , else  $L(W^T)$ . See Appendix B for more details.
4 MIMOFormer
This section presents MIMOFormer, which applies the principles of computation in superposition to dot-product self-attention [24]. Figure 3 shows a MIMOFormer layer with four protected channels, consisting of a single-head $^2$  attention block, a concatenation, a linear layer, an MLP, and a skip connection.
Merely superposing protected attention keys<sup>3</sup> and queries does not yield the desired result. As discussed in Appendix F, with scalar attention scores between pairs of tokens, vanilla dot-product attention irreversibly combines attention scores of separate protected channels, effectively blurring the attention weights. By building instead on linear Transformers [15] [25], attention scores are not collapsed to scalars, thus enabling computation in superposition.
Despite being compatible with other linear transformers (such as DPFP [25]), for concreteness we discuss changes to the Performer's FAVOR+ attention block [15]. Enabling computation in superposition, we label the block as FAVOR+S.
Given attention keys  $(k_{j})_{j = 1}^{L}$  queries  $(q_{i})_{i = 1}^{L}$  and values  $(v_{j})_{j = 1}^{L}$  FAVOR+ estimates dotproduct attention at sequence index  $i$  through
Figure 3: MIMOFormer layer applying computation in superposition to single-head FAVOR+S attention and to the MLP. This example passes four channels  $(N\cdot M = 2\cdot 2 = 4)$  to the FAVOR+S attention (see Eq. (19)). The individual outputs are retrieved by unbinding after the MLP. The skip connection superposes the individual inputs for alignment, using the same protection keys as in unbinding.
$$
o _ {i} = \sum_ {j = 1} ^ {L} v _ {j} \frac {\exp \left(\langle k _ {j} , q _ {i} \rangle / \sqrt {D}\right)}{\sum_ {l = 1} ^ {L} \exp \left(\langle k _ {l} , q _ {i} \rangle / \sqrt {D}\right)} \approx \frac {1}{B _ {i}} \underbrace {\left[ \sum_ {j = 1} ^ {L} v _ {j} \phi \left(k _ {j}\right) ^ {T} \right]} _ {A} \times \phi \left(q _ {i}\right), \tag {10}
$$
where  $\phi : \mathbb{R}^D \to \mathbb{R}_+^R$  approximates the softmax kernel  $\exp(\langle k_j, q_i \rangle / \sqrt{D})$  as an explicit inner product. Since computing  $\phi$  has a computational complexity of  $\mathcal{O}(DR)$ , the construction of  $A \in \mathbb{R}^{D \times R}$  takes  $\mathcal{O}(LDR)$ . Equally, multiplying  $A \times \phi(q_i) \forall i$  takes  $\mathcal{O}(LDR)$ . Thus, FAVOR+ breaks the quadratic dependence on sequence length  $L$  of attention. The denominator  $B_i$  is discussed in Appendix C.
2The application to multi-head attention is straightforward; empirical results are shown in Section 5.2.
3Protection keys are denoted by the letter  $a$  and attention keys by the letter  $k$  to distinguish between them.
Our extension, FAVOR+S, computes attention in superposition, yielding a square-root speedup in the number of protected channels. It encodes them in an  $M \times N$  grid, distributing the computational burden among the setup of the value-key matrix  $A$  and its product with  $\phi(q_i)$ . Importantly, in the limit of high dimensional projections  $D$  our mechanism converges to exact self-attention, completely separating the protected channels from one another. In the following, we assume the token values  $\overline{v}_j$ , keys  $\overline{k}_j$ , and queries  $\overline{q}_j$  to be in protected subspaces:
$$
v _ {j} ^ {(m, n)} := \bar {v} _ {j} ^ {(m, n)} \odot a ^ {(m, n)}, \quad k _ {j} ^ {(m, n)} := \bar {k} _ {j} ^ {(m, n)} \odot a ^ {(m, n)}, \quad q _ {j} ^ {(m, n)} := \bar {q} _ {j} ^ {(m, n)} \odot a ^ {(m, n)}, \tag {11}
$$
where  $a^{(m,n)}$  are i.i.d. bipolar vectors of Rademachers [17] and  $(m,n)$  denotes a channel.
MIMOFormer benefits from the low time complexity  $(\mathcal{O}(D))$  of the Hadamard product, especially since binding and unbinding are performed in every MIMOFormer layer. Our derivations rely on two estimates:
$$
\phi (k) ^ {T} \phi (q) \stackrel {{P}} {{\approx}} \exp \left(\langle k, q \rangle / \sqrt {D}\right) \quad \langle \sum_ {w = 1} ^ {N} k _ {j} ^ {(u, w)}, \sum_ {t = 1} ^ {M} q _ {i} ^ {(t, n)} \rangle \stackrel {{H}} {{\approx}} \underbrace {\langle k _ {j} ^ {(u , n)} , q _ {i} ^ {(u , n)} \rangle} _ {\text {i n t e n d e d s i g n a l}} \tag {12}
$$
The approximation  $P$ , which improves with increasing  $R = \dim(\phi(\overline{q}_i))$ , is due to FAVOR+ and is quantified in [15]. On the other hand, the approximation  $H$  follows from:
Inter-channel distortion. The probability that inter-channel attention distorts the intended signal of the dot-product by a factor outside  $[1 - \alpha, 1 + \alpha]$  has various upper bounds, most notably decaying exponentially w.r.t.  $D\alpha^2\cos^2(\angle (\overline{k}_j^{(u,n)}, \overline{q}_i^{(u,n)})) / (NM - 1)^2$ . See Appendix D for the full theorem.
4.1 FAVOR+S: Computing self-attention in superposition
We first discuss separately how to use a one-dimensional grid to carry out either the multiplication  $(\times)$  or the construction of  $A$  in superposition. Finally, the integration into a 2D grid will be shown.
Placing multiple queries in superposition. We set up channels  $1, \ldots, M$  by simultaneously generating a superposition in the construction of  $A_{s}$  and of the queries to be applied. To avoid inter-channel attention we superpose value-key tensor products, i.e., we do not construct tensor products between superposed values and superposed keys:
$$
\begin{array}{l} S _ {i} = \underbrace {\left[ \sum_ {j = 1} ^ {L} \sum_ {u = 1} ^ {M} v _ {j} ^ {(u)} \phi \left(k _ {j} ^ {(u)}\right) ^ {T} \right]} _ {A _ {s}} \times \phi \left(\sum_ {t = 1} ^ {M} q _ {i} ^ {(t)}\right) = \sum_ {j = 1} ^ {L} \sum_ {u = 1} ^ {M} v _ {j} ^ {(u)} \left(\phi \left(k _ {j} ^ {(u)}\right) ^ {T} \phi \left(\sum_ {t = 1} ^ {M} q _ {i} ^ {(t)}\right)\right) (13) \\ \approx \sum_ {j = 1} ^ {L} \sum_ {u = 1} ^ {M} v _ {j} ^ {(u)} \exp \left(\langle k _ {j} ^ {(u)}, \sum_ {t = 1} ^ {M} q _ {i} ^ {(t)} \rangle / \sqrt {D}\right) \approx \sum_ {j = 1} ^ {L} \sum_ {u = 1} ^ {M} v _ {j} ^ {(u)} \exp \left(\langle k _ {j} ^ {(u)}, q _ {i} ^ {(u)} \rangle / \sqrt {D}\right) (14) \\ = \sum_ {u = 1} ^ {M} \underbrace {\sum_ {j = 1} ^ {L} v _ {j} ^ {(u)} \exp \left(\left\langle k _ {j} ^ {(u)}, q _ {i} ^ {(u)} \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} u}. (15) \\ \end{array}
$$
We obtain a superposition of bound output values; hence, the cost of computing  $A \times \phi(q_i)$  for all  $i$  is amortized across channels. However, the construction complexity of  $A_s$  is increased  $M$ -fold to  $\mathcal{O}(LDR \cdot M)$ , hence the complexity per protected channel remains at  $\mathcal{O}(LDR)$ .
Constructing value-key tensor products in superposition. Next, we demonstrate a value-key tensor product  $(A_{s})$  shared across all channels, but with a mere  $\mathcal{O}(LD(R + N))$  setup complexity. In  $\mathcal{O}(LND)$  we compute  $\sum_{q = 1}^{N}v_{j}^{(q)}$  and  $\sum_{w = 1}^{N}k_{j}^{(w)}$  for all  $j$ . These are then (re)used to build  $A_{s}$ .