SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tao Lei* Junwen Bai Siddhartha Brahma Joshua Ainslie Kenton Lee Yanqi Zhou Nan Du Vincent Y. Zhao Yuexin Wu Bo Li Yu Zhang Ming-Wei Chang
We propose Conditional Adapter (CoDA), a parameter-efficient transfer learning method that also improves inference efficiency. CoDA generalizes beyond standard adapter approaches to enable a new way of balancing speed and accuracy using conditional computation. Starting with an existing dense pretrained model, CoDA adds sparse activation together with a small number of new parameters and a light-weight training phase. Our experiments demonstrate that the CoDA approach provides an unexpectedly efficient way to transfer knowledge. Across a variety of language, vision, and speech tasks, CoDA achieves a 2x to 8x inference speed-up compared to the state-of-the-art Adapter approaches with moderate to no accuracy loss and the same parameter efficiency.
Large pretrained models have achieved groundbreaking results but the main impediment to deploy them has been the cost of adaptation and inference. Due to the ever growing size of the pretrained models, for example, finetuning has become increasingly expensive as it requires a separate copy of the full model and updates to all parameters for every downstream task. Parameter-efficient transfer learning such as Adapter [Houlsby et al., 2019] and Prompt Tuning [Lester et al., 2021] have been proposed to address this issue. These methods only update a small subset of parameters for each downstream task, allowing the model to retain knowledge and avoid catastrophic forgetting [Vu et al., 2022]. Noticeably, these methods can match the accuracy of a fully finetuned model, while achieving better accuracy on out-of-domain data distributions [Lester et al., 2021, Awadalla et al., 2022].
Unfortunately, standard parameter-efficient transfer learning methods only bring parameter efficiency, not inference efficiency. For example, while only a few small projection matrices are added into the pretrained model in the Adapter approach, all the model inputs (such as tokens) still use all parameters during inference. Therefore, the inference speed is the same (or slightly lower) with respect to the full finetuning method. Moreover, prior studies have shown that these parameter-efficient learning methods are most effective when the size of the pretrained model is large [Lester et al., 2021], making many advantages of these methods difficult to realize in practice.
In this paper, we propose Conditional Adapter (CoDA), a parameter-efficient transfer learning method that offers both parameter and inference efficiency. CoDA is a generalization of the adapter approach, built with the following intuition - we can treat the pretrained model as a universal source of knowledge but only query against it for necessary inputs. Figure 1 compares CoDA with finetuning and standard adapter approaches. Similar to standard adapter approaches, our model adds and updates a small adapter in each layer, while fixing the pretrained Transformer blocks for downstream adaptation. Unlike previous approaches, however, CoDA assumes that many of input
*Correspondence: taoleics@gmail.com, junwen@google.com
Figure 1: Comparison between different ways to use pretrained Transformer models, including (1) standard finetuning (left) where all parameters are tunable and computation is dense, (2) standard adapters (center) where a small set of new tunable parameters are added while the computation remains dense, and (3) CoDA (right) where the computation is sparsely activated.
Table 1: CoDA significantly reduces the inference time compared to the Parallel Adapter approach [He et al., 2021], while still maintaining parameter efficiency.
New param	MNLI (text)
Acc ↑	Speedup
P-Adapter	0.4%	91.5	1.0x
CoDA	0.4%	90.7	3.2x
New param	OCR-VQA (vision)
EM ↑	Speedup
P-Adapter	2.8%	67.5	1.0x
CoDA	2.8%	67.6	8.0x
New param	Librispeech (speech)
WER ↓	Speedup
P-Adapter	2.5%	1.4/2.7	1.0x
CoDA	2.5%	1.4/2.8	2.2x
token representations (of each layer) are not important for the prediction task and therefore do not require heavy computation. In such cases, the pretrained Transformer block can be skipped. Given that many tokens are not processed by the Transformer block, CODA runs significantly faster than previous methods.
While conditional activation has clear speed benefits, CoDA must learn to select important tokens for heavy computation in order to maintain model accuracy. To this end, we introduce a soft top- $k$  operation for computing the token selection decision. This soft top- $k$  operation, which can be seen as a generalization of softmax and a relaxation of hard top- $k$ , utilizes entropy-regularized optimization techniques similar to computational optimal transport [Cuturi, 2013]. As a result, its output can be computed using fast and differentiable iterations, allowing token selection to be directly optimized for model performance.
We apply CoDA on encoder-heavy tasks and evaluate its effectiveness on three different domains – natural language processing, computer vision and speech processing. Overall, CoDA achieves 2 to 8 times inference speed-up over standard adapter approach with moderate to no accuracy loss. Table 1 showcases our results by selecting one of the best performing tasks in each domain. We also conduct comprehensive ablation studies to analyze the effectiveness, efficiency and scalability of CoDA. For example, we found that with just a little to no router pretraining, existing dense pretrained models such as T5 [Raffel et al., 2020] can be efficiently converted into CoDA models to gain both parameter efficiency and speed advantages.
Parameter-efficient transfer learning methods Due to the ever-growing number of parameters in the pretrained Transformer models, various methods have been proposed for transfer learning with minimal parameter updates. Prompt tuning [Lester et al., 2021] and prefix tuning [Li and Liang, 2021] introduce new virtual token embeddings that can be finetuned as model parameters. Adapter approaches [Houlsby et al., 2019, He et al., 2021] add a small number of new, learnable parameters to each layer while keeping the pretrained parameters fixed. Another popular method, Low-Rank Adaptation [LoRA; Hu et al., 2021], injects learnable low-rank decomposition matrices into pretrained model parameters. In addition to requiring less storage cost, parameter-efficient methods have been shown to be more sample-efficient and achieve better out-of-domain generalization than standard finetuning. CoDA is an adapter approach but can be easily combined with other parameter-efficient methods such as LoRA to accelerate their inference.
Conditional computation The development of sparsely and conditionally activated models has been a very active research area. For example, Mixture-of-Experts (MoE) models [Shazeer et al.,
2017] and many recent advances [Du et al., 2022, Fedus et al., 2021] have been proposed to scale up the size of language models without increasing the computation cost. Many recent works have explored better token routing methods for MoE models, for example using random hashing [Roller et al., 2021], balanced assignment [Lewis et al., 2021] and expert-choosing router [Zhou et al., 2022]. CoDA applies conditional computation to both attention and feed-forward blocks of the model, whereas MoE models only focus on sparse activation in the feed-forward blocks.
Similar to our approach, various recent methods have achieved computation efficiency by skipping computation on a subset of input tokens. However, the selection mechanism can be very different, such as using pooling [Nawrot et al., 2022], token merging [Bolya et al., 2023], token pruning [Rao et al., 2021, Yin et al., 2022], learned sigmoid gates [Bapna et al., 2020] and early exiting [Schuster et al., 2022]. While most of the token merging and pruning methods have been proposed for vision tasks, we show that CoDA is applicable to multiple domains including text, vision and speech. In addition, token merging and our token selection method are built with different inductive biases and intuition. Token merging leverages redundancies in visual tokens, while token selection assumes a spike of token relevance. That is, only a few tokens are necessary for the prediction task. Another major difference is that CoDA dynamically routes and updates token representations in each layer, whereas if a token is pruned (or merged), it will never be re-used by subsequent layers. We believe our token routing mechanism is more suited for text and speech applications, such as question answering, where different tokens might play important roles in different layers, or given different input queries.
Finally, CoDA is closely related to a concurrent work, CoLT5 [Ainslie et al., 2023], which also utilizes conditional activation (token selection) for inference efficiency. The focus of CoLT5 and CoDA are very different. CoLT5 specifically tailors its model architecture for long text (e.g. over 16k tokens), for example, by combining local attention with routed attention. The CoLT5 models are pre-trained from scratch and all parameters are finetuned for downstream tasks. In comparison, CoDA is directly initialized and adapted from an already pretrained dense model, and we optimize its performance on parameter-efficient transfer learning. The strengths of CoDA and CoLT5 can be combined for long text applications.
Efficient Transformer models Many efficient Transformer variants have been proposed to accelerate model computation. Examples include creating fast attention variants [Wang et al., 2020a, Beltagy et al., 2020, Guo et al., 2022, Hua et al., 2022], searching network architectures [Press et al., 2019, So et al., 2021, Su et al., 2021] and utilizing non-attention neural modules for efficiency [Gulati et al., 2020, Lei, 2021]. CoDA utilizes conditional computation as an orthogonal approach for efficiency.
Model compression Apart from building efficient model architectures, model compression methods such as pruning [Han et al., 2016, Zhu and Gupta, 2017, Wang et al., 2020b, Xia et al., 2022] and distillation [Hinton et al., 2015, Kim and Rush, 2016, Turc et al., 2019, Lin et al., 2020] can be adopted to speed up model inference. Compared to these methods, CoDA retains all model parameters of the pretrained large model, and therefore avoids retraining a new model from scratch or knowledge forgetting caused by parameter removal. In addition, CoDA can be seen as a dynamic version of layer pruning because it can activate different Transformer layers for each token, and can be further combined with distillation to reduce the loss of accuracy caused by conditional computation.
3.1 Architecture
Throughout this and the experiment section, we build CoDA on top of parallel adapters [He et al., 2021]. However, note that our method can be generalized to other types of adapters such as sequential adapters [Houlsby et al., 2019] and LoRA [Hu et al., 2021]. We present additional experimental results using LoRA in Appendix B.3. Figure 2 illustrates our architecture and shows how CoDA computes its output by selecting only a small subset of input tokens to query against the pretrained model. When parallel adapters are used, CoDA introduces a small number of learnable parameters in the parallel branches, while the vast majority of model parameters (associated with the pretrained Transformer layers) remain fixed. In addition, CoDA only sends  $k = \lceil n / r \rceil$  tokens for heavy processing. We define  $r > 1$  as the reduction factor, a constant (such as 4) to control the computation saving.
Next, we briefly introduce our notations and describe the computation of CoDA in detail. We use  $F()$  to denote a parameterized neural network and the corresponding function defined by the network. For instance, a Transformer layer [Vaswani et al., 2017] consists of an attention sub-layer  $F_{\mathrm{att}}()$  followed by a feed forward sub-layer  $F_{\mathrm{ffn}}()$ . Each layer also employs layer normalization [Ba et al., 2016], namely  $LN_{\mathrm{att}}()$  and  $LN_{\mathrm{ffn}}()$ , before applying the attention and feed forward functions. We define  $X \in \mathbb{R}^{n \times d}$  as the input of a Transformer encoder layer, where  $n$  is the number of input tokens and  $d$  is the hidden size of the model.
Given layer input  $\mathbf{X}$ , we first apply layer normalization, namely  $\mathbf{X}_{\mathrm{norm}} = LN_{\mathrm{att}}(\mathbf{X})$ . The normalized input will be processed by the adapter branch and the conditional Transformer branch. Their outputs are then added and combined as the final output of the layer.
Adapter branch Let  $F_{\mathrm{adapter}}(\cdot)$  denote the transformation function of the adapter branch. The output is defined as
$$
Z _ {\text {a d a p t e r}} = F _ {\text {a d a p t e r}} \left(X _ {\text {n o r m}}\right) \tag {1}
$$
Similar to the previous approaches,  $F_{\mathrm{adapter}}()$  is realized using a feed forward network with a small hidden size such as 64. As a result, computing  $Z_{\mathrm{adapter}}$  only incurs a small number of floating point operations and its cost is often negligible compared to the cost of the heavy Transformer branch. The adapter branch does not conditionally select tokens. In other words,  $F_{\mathrm{adapter}}()$  is applied to all input tokens  $X \in \mathbb{R}^{n \times d}$ .
Conditional branch The computation of the conditional branch takes three steps. First, each CoDA layer defines a router function  $F_{\text{router}}()$  to select  $k$  tokens for the conditional branch. The router function in each layer returns two outputs
$$
\boldsymbol {m}, \boldsymbol {P} = F _ {\text {r o u t e r}} \left(\boldsymbol {X} _ {\text {n o r m}}\right) \tag {2}
$$
Figure 2: Illustration of a single CoDA layer with parallel adapter.  $k$  tokens are selected and processed by the frozen pretrained Transformer layer, and all tokens are processed by the fast adapter layer.
where  $P \in \{0,1\}^{k \times n}$  is a matrix consisting of  $k$  one-hot vectors indicating the selection of tokens. Here  $P[i,j] = 1$  if and only if the  $i$ -th selected token is the  $j$ -th input token from  $\tilde{X}$ .  $m \in [0,1]^n$  is a weight mask in which  $m[j]$  is the selection weight for the  $j$ -th input token.  $m[j] = 0$  if the token is not selected. We will describe how the router learns the selection in more details later in this section.
After the routing decision is made, the input representations of the selected tokens can be collected using a matrix multiplication,
$$
\boldsymbol {X} _ {\text {r o u t e d}} = \boldsymbol {P} \boldsymbol {X} _ {\text {n o r m}} \quad \in \mathbb {R} ^ {k \times d} \tag {3}
$$
where  $k$  rows in  $\mathbf{X}_{\mathrm{norm}}$  are selected to construct the  $k$ -by- $d$  matrix  $\mathbf{X}_{\mathrm{routed}}$ . Similar to a standard Transformer layer, the conditional branch applies attention and feed forward transformations to the selected input:
$$
\bar {Z} _ {\text {r o u t e d}} = F _ {\text {a t t}} \left(\boldsymbol {X} _ {\text {r o u t e d}}\right) \tag {4}
$$
$$
\boldsymbol {Z} _ {\text {r o u t e d}} = F _ {\text {f f n}} \left(L N _ {\text {f f n}} \left(\boldsymbol {X} _ {\text {r o u t e d}} + \bar {\boldsymbol {Z}} _ {\text {r o u t e d}}\right)\right) \tag {5}
$$
where  $\bar{Z}_{\mathrm{routed}}$ ,  $Z_{\mathrm{routed}} \in \mathbb{R}^{k \times d}$  denote the output of the attention network and the feed forward network respectively.
We consider two attention variants which differ in how they compute key-value vectors. One variant applies a  $k$ -to-  $k$  attention using  $X_{\mathrm{routed}}$  as both the query vectors and key-value vectors. The other variant applies a  $k$ -to-all attention using the entire input vectors  $X_{\mathrm{norm}}$  as the attention keys and values. The  $k$ -to-all variant runs slower but obtains higher quality close to the full model. We compare the performance of the two variants in Section 5.
The attention and feed-forward output  $\bar{Z}_{\mathrm{routed}}$  and  $Z_{\mathrm{routed}}$  are combined and projected back to the same shape of the original input
$$
\boldsymbol {Z} _ {\text {c o n d}} = \boldsymbol {P} ^ {\top} \left(\bar {\boldsymbol {Z}} _ {\text {r o u t e d}} + \boldsymbol {Z} _ {\text {r o u t e d}}\right) \quad \in \mathbb {R} ^ {n \times d} \tag {6}
$$
Finally  $Z_{\mathrm{cond}}$  merges with the adapter output and the original input of the current layer to produce the output of the layer:
$$
\boldsymbol {Y} = \boldsymbol {X} + \boldsymbol {Z} _ {\text {a d a p t e r}} + \boldsymbol {m} \odot \boldsymbol {Z} _ {\text {c o n d}} \tag {7}
$$
$m \odot Z_{\mathrm{cond}}$  is an element-wise multiplication that scales the rows of  $Z_{\mathrm{cond}}$  using weight  $m$ . This operation can be seen as a gating operation, where the hidden state  $Z_{\mathrm{cond}}[i]$  of the  $i$ -th token is weighted by the token selection score  $m[i]$  assigned by the router. This enables gradient propagation from  $m$  to the router parameters, such that the token selection can be jointly optimized with other model components during training.
Learned router An important ingredient of CoDA is the router function  $F_{\mathrm{router}}(\cdot)$  that is learned to select a subset of tokens for favorable model performance. Given the token representation  $X_{\mathrm{norm}}$ , our router first computes dot-product score  $s = wX_{\mathrm{norm}}^{\top}$ , where  $w \in \mathbb{R}^{d}$  is a parameter vector associated with the router in this layer. The dot-product score  $s$  is further normalized by a function  $f(\cdot): \mathbb{R}^n \to [0,1]^n$ , and clipped to produce the selection score  $m$ :
$$
\boldsymbol {\lambda} = f (\boldsymbol {s}) \tag {8}
$$
$$
\boldsymbol {m} = \boldsymbol {\lambda} \odot \operatorname {T o p} (\boldsymbol {\lambda}, k) \quad \in \mathbb {R} ^ {n} \tag {9}
$$
Here  $\operatorname{Top}(\lambda, k) \in \{0, 1\}^n$  is an indicator function which returns a binary mask indicating the top- $k$  values in  $\lambda$ . The one-hot matrix  $P$  defined in (2) can be created according to  $\operatorname{Top}(\lambda, k)$ . In short, the highest values of  $\lambda$  will be selected by the router.
Function  $f(\cdot)$  must remain differentiable with respect to its input ( $s$  in this case) such that we can update the router parameters  $w$  during training. One possible choice for  $f(\cdot)$  is the sigmoid activation function which normalizes the values in  $s$  independently. However, this does not explicitly model the constraint that we need to select  $k$  tokens from  $n$  available tokens. Consider a simple case where  $k = 1$ , a natural choice for  $f(\cdot)$  would be the softmax function. Since softmax provides global normalization over the input scores, a gradient update to increase one of the scores would also decrease the other scores, a desirable effect for learning top-1 selection.
We hypothesize that a soft top- $k$  operator that generalizes softmax should be used for general  $k > 1$ . This is indeed possible by formalizing soft top- $k$  as the following optimization problem:
$$
f (\boldsymbol {s}) := \underset {\boldsymbol {\lambda}} {\arg \max } \boldsymbol {s} ^ {\top} \boldsymbol {\lambda} + \epsilon H (\boldsymbol {\lambda})
$$
$$
\text {s . t .} \quad \mathbf {1} ^ {\top} \boldsymbol {\lambda} = k, \quad \boldsymbol {\lambda} [ i ] \in [ 0, 1 ] \forall i = 1, \dots , n \tag {10}
$$
Here  $H(\lambda) = \sum_{i=1}^{n} -\lambda[i] \log \lambda[i]$  is a generalized entropy function (applied to any positive vector  $\lambda$  instead of a distribution), and  $\epsilon > 0$  is a small coefficient.
This optimization problem is closely related to the softmax and top- $k$  operation. Specifically, when  $\epsilon = 0$ , it becomes a linear program which returns  $\mathrm{Top}(s,k)$  as the solution. In addition, when  $k = 1$ , it can be shown that its solution is  $\mathrm{softmax}(s / \epsilon)$ . Broadly speaking, (10) will return a soft top- $k$  mask and the smoothness is controlled by  $\epsilon$  (and hence  $\epsilon$  must be positive to act as a temperature).
Problem (10) does not have a closed-form solution for an arbitrary  $\epsilon > 0$  and  $k > 1$ , but its solution can be obtained using an iterative algorithm. Specifically, let  $a \in \mathbb{R}$  and  $\pmb{b} \in \mathbb{R}^n$  be two auxiliary variables (which can be initialized to zeros). The solution takes the form  $\lambda = \exp \left(\frac{s + b + a}{\epsilon}\right)$ . The values of  $a$  and  $b$  can be obtained using the following iterative updates:
$$
a ^ {\prime} = \epsilon \ln (k) - \epsilon \ln \left(\sum_ {i = 1} ^ {n} \exp \left(\frac {\boldsymbol {s} [ i ] + \boldsymbol {b} [ i ]}{\epsilon}\right)\right)
$$
$$
\boldsymbol {b} ^ {\prime} = \min  \left(- s - a ^ {\prime}, 0\right) \tag {11}
$$
In practice, we use  $T = 20$  iterations and the function  $f(s)$  returns  $\exp \left( \frac{s + b + a}{\epsilon} \right)$  using  $a$  and  $b$  from the last iteration. The function  $f(s)$  remains differentiable with respect to  $s$  using these iterative updates, so we can train the router jointly with other model parameters. We provide additional discussion and the derivation of the updates in Appendix §C.
Model	Reduction r	Base	Large	XL	Δ Avg
MNLI	RTE	BoolQ	MNLI	RTE	BoolQ	MNLI	RTE	BoolQ
Parallel Adapter (w/o conditional computation)	-	87.1	71.5	77.9	90.3	84.8	85.8	91.5	89.9	88.4	±0.0
CoDA, k-to-all attention	3	86.6	72.6	76.6	90.2	85.9	85.1	91.4	91.3	89.4	+0.2
CoDA, k-to-k attention	86.3	72.2	76.2	89.8	87.0	83.7	91.4	89.5	88.2	-0.3
CoDA, k-to-all attention	5	86.0	70.8	76.0	89.7	85.2	84.3	91.0	91.3	87.2	-0.6
CoDA, k-to-k attention	82.5	70.8	75.4	88.1	87.0	81.8	89.9	87.7	84.8	-2.1
Table 2: Results of applying CoDA to T5 v1.1 models. CoDA achieves significant computation savings while retaining accuracy close to the dense baseline. We compare CoDA to a corresponding parallel adapter method that processes all tokens without conditional computation. We report accuracy on the development set on 3 tasks  $\times$  3 model sizes, and set the number of selected tokens  $k = \lceil n / r\rceil$ . The last column shows the change on average accuracy with respect to the parallel adapter method. We select the  $k$ -to-all version as our default (shown in bold).
3.2 Training
CoDA can be directly initialized from an existing Transformer model. Given a pretrained model such as T5 [Raffel et al., 2020], the Transformer layers are directly re-used and copied in the conditional branches of CoDA, and only the adapter and router parameters are randomly initialized. Because pretraining a large dense model can be expensive, our method reduces the overall training cost.
The routers and neural network components in CoDA must co-operate and be optimized for accurate model predictions. When the available finetuning data is limited, a random initialization for the router (and adapter) parameters can be sub-optimal. We demonstrate that CoDA can be further pretrained using the same pretraining objective as the dense model, in order to enhance downstream performance. Importantly, CoDA requires significantly fewer training steps during pretraining, since most of its parameters are taken from an already pretrained model. We show that the cost of CoDA pretraining can be  $10 - 30\mathrm{x}$  lower than the pretraining of its original dense model. We present this analysis in Section 5.
Finally, we train CoDA on downstream tasks by only updating the adapter, router and layer normalization parameters. The size of the adapters is small (e.g. 5M parameters), and each router and layer normalization block only introduces  $d$  parameters, where  $d$  is the model dimension. As a result, CoDA remains parameter-efficient similar to previous adapter and prompt-tuning methods.