Title: Conditional Adapters: Parameter-efficient Transfer Learning with Fast Inference
Venue: NeurIPS 2023
Accepted: yes
Keywords: conditional computation

inference efficiency

parameter efficiency

large models
Abstract: 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.

Conditional Adapters: Parameter-efficient Transfer Learning with Fast Inference

Abstract

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.

1 Introduction

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

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 topk operation for computing the token selection decision. This soft topk operation, which can be seen as a generalization of softmax and a relaxation of hard topk , 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.

2 Related Work

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 Method

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 = glyph[ceilingleft] n/r glyph[ceilingright] 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 att () followed by a feed forward sub-layer F ffn () . Each layer also employs layer normalization [Ba et al., 2016], namely LN att () and LN ffn () , before applying the attention and feed forward functions. We define X ∈ R n × 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 X , we first apply layer normalization, namely X norm = LN att ( 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 adapter () denote the transformation function of the adapter branch. The output is defined as

Similar to the previous approaches, F adapter () is realized using a feed forward network with a small hidden size such as 64. As a result, computing Z 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 adapter () is applied to all input tokens X ∈ R n × d .

Conditional branch The computation of the conditional branch takes three steps. First, each CODA layer defines a router function F router () to select k tokens for the conditional branch. The router function in each layer returns two outputs

where P ∈ { 0 , 1 } k × 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 ˜ X . m ∈ [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,

where k rows in X norm are selected to construct the k -byd matrix X routed . Similar to a standard Transformer layer, the conditional branch applies attention and feed forward transformations to the selected input:

where ¯ Z routed , Z routed ∈ R k × 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 -tok attention using X 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 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 ¯ Z routed and Z routed are combined and projected back to the same shape of the original input

Layer output

Layer input

all tokens

Tunable &

Fast adapter

FFN (     )

Pretrained &

Expensive

attention (      )

Pretrained &

Expensive FFN

(      )

Finally Z cond merges with the adapter output and the original input of the current layer to produce the output of the layer:

m glyph[circledot] Z cond is an element-wise multiplication that scales the rows of Z cond using weight m . This operation can be seen as a gating operation, where the hidden state Z 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 router () that is learned to select a subset of tokens for favorable model performance. Given the token representation X norm, our router first computes dot-product score s = wX glyph[latticetop] norm , where w ∈ 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 () : R n → [0 , 1] n , and clipped to produce the selection score m :

Here Top ( λ , k ) ∈ { 0 , 1 } n is an indicator function which returns a binary mask indicating the topk values in λ . The one-hot matrix P defined in (2) can be created according to Top ( λ , k ) . In short, the highest values of λ will be selected by the router.

Function f () 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 () 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 () 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 topk operator that generalizes softmax should be used for general k > 1 . This is indeed possible by formalizing soft topk as the following optimization problem:

Here H ( λ ) = ∑ n i =1 -λ [ i ] log λ [ i ] is a generalized entropy function (applied to any positive vector λ instead of a distribution), and glyph[epsilon1] > 0 is a small coefficient.

This optimization problem is closely related to the softmax and topk operation. Specifically, when glyph[epsilon1] = 0 , it becomes a linear program which returns Top ( s , k ) as the solution. In addition, when k = 1 , it can be shown that its solution is softmax ( s /glyph[epsilon1] ) . Broadly speaking, (10) will return a soft topk mask and the smoothness is controlled by glyph[epsilon1] (and hence glyph[epsilon1] must be positive to act as a temperature).

Problem (10) does not have a closed-form solution for an arbitrary glyph[epsilon1] > 0 and k > 1 , but its solution can be obtained using an iterative algorithm. Specifically, let a ∈ R and b ∈ R n be two auxiliary variables (which can be initialized to zeros). The solution takes the form λ = exp( s + b + a glyph[epsilon1] ) . The values of a and b can be obtained using the following iterative updates:

In practice, we use T = 20 iterations and the function f ( s ) returns exp( s + b + a glyph[epsilon1] ) using a and b from the last iteration. The function f ( s ) remain 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.

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-30x 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.

4 Experimental setup

CODA is evaluated on three domains including natural language processing (NLP), computer vision and speech processing, and on a range of applications such as classification, question answering, summarization and speech recognition. The experiments are organized as follows: We first demonstrate the effectivenss of CODA conduct analyses on its design choices using the publicly available T5 models (§5). In our final results (§6), we pretrain Transformer models from scratch and extend our evaluation to vision and speech domains.

Datasets We use the C4 corpus [Raffel et al., 2020] for pretraining text models. For speech models, we use the LibriLight corpus [Kahn et al., 2020] for pretraining. Our vision Transformer models use the same data and training procedure in Pix2Struct [Lee et al., 2022]. Our finetuning datasets for text models include the MNLI [Williams et al., 2018], RTE [Dagan et al., 2005, Haim et al., 2006, Giampiccolo et al., 2007, Bentivogli et al., 2009], BoolQ [Clark et al., 2019], SQuAD [Rajpurkar et al., 2016] and XSum [Narayan et al., 2018] datasets. The speech models are evaluated on the speech recognition task using the LibriSpeech dataset [Panayotov et al., 2015]. Finally, we use the OCR-VQA [Mishra et al., 2019], DocVQA [Mathew et al., 2021], and Screen2Words [Wang et al., 2021] datasets for vision models.

Adapter for

Fast adapter

FFN

Input tensor

Expensive

5 Understanding and Analyzing CODA Expensive FFN Fast adapter FFN

attention

Transformer

Expensive

FFN

Expensive attention

Setup We present several analyses to validate the design choices of CODA in this section. We initialize CODA using the version 1.1 release of T5 checkpoints 2 , and perform CODA pretraining using the same setting as the T5 models. During pretraining, we set routing capacity to k = 192 given input sequence length n = 512 . We do not tune the value of k for pretraining, but will report the results of using different k values in finetuning. We perform 100K gradient steps, which is 10% of the total number of steps used to train the T5 dense models. The overall computational cost is over 20x less than the full training of dense models, since CODA only applies heavy computation on less than half of the tokens. Expensive attention Expensive FFN Fast adapter FFN k tokens all tokens k tokens Expensive attention Expensive FFN Fast adapter FFN k tokens all tokens k tokens

Input tensor

For simplicity, we evaluate on classification tasks for various ablation studies of CODA. Specifically, we report results on the MNLI, RTE and BoolQ datasets, and test three different model sizes including the Base, Large and XL size of T5. We will extend our evaluation to generation tasks such as question answering in the full result section (§6).

Dense

Can CODA be fast and accurate? Table 2 presents the finetuning results of CODA. As a comparison, we also report the results of Parallel Adapter, which is similar to CODA except that it applies the expensive Transformer layers to all input tokens. This constitutes an upper-bound, and is a strong baseline that has been reported as the best among a range of adapter and prompt tuning methods [He et al., 2021]. As shown in Table 2, CODA can achieve 3-5x computation reduction ( r = 3 , 5 ) in the Transformer layers at a cost of less than 1.0 point drop on average accuracy. As expected, our k -to-all attention variant achieves consistently better accuracy than the k -tok variant, since it can access the full attention context. On the other hand, the k -tok attention variant runs faster in practice, which can be beneficial for tasks with very long inputs. We select the k -to-all version in the final result section (§6). pre-training CODA pre-training CODA fine-tuning × N CODA pre-training CODA fine-tuning × N

How many pretraining steps are needed? Figure 3 plots the finetuning accuracy by varying the number of pretraining steps for CODA. Because CODA can be initialized using pretrained dense models, it requires as few as 20K steps to obtain competitive finetuning results. Of course, using more pretraining steps can improve the downstream accuracy. The fact that CODA can be quickly updated without repeating the expensive pretraining will be very beneficial in real-world applications.

Does learned routing matter? Weanalyze the impact of learned routing in Table 3 by comparing our soft topk router with other router implementations. We implement a variant that replaces soft topk with the sigmoid activation function, so the selection weight of each token activates on its own (without considering the capacity constraint). As shown in the table, this variant

Dense pre-training

MNLI

BoolQ

95

95

91-2—91-2

91-3

91.4

87-3

88

88.2

90

89:5

89-7

668

85

84

84:5

84.8

85

776.4

83.1

r=3

r=5

80

65

0

0

25000

50000

75000

100000

25000

50000

75000

100000

XL

r=1

XL

Speed v.s. Quality

CoDA

P-Adapter

XL

Large

Large

Base

Examples per second

Speed v.s. Quality

CoDA

P-Adapter r=3

Large

Base

Examples per second

BoolQ

BoolQ

CoDA

r=5

Rouge-2

XL

Small

Base

P-Adapter

Examples per second

Examples per second

Reduction r

Avg fine-tuning score

84

7.9.9

77.9

7.8.2

76

75:3

71.0

68

63.9

CoDA

P-Adapter

60

Base

Large

XL

Decoding speed-up

Base

Large

XL

Reduction r

Speed v.s. Quality

Examples per second

Rouge-2

Base

Small

Large

XL

r=3

r=5

r=1

CoDA

P-Adapter

XL

Large

Base

Examples per second

Speed v.s. Quality

26

4

3.7

Small

2.8

20

3

2.5

2.6

Base

2.4

14

2

2.2

Large

1.0

XL

2.0

006010

3

5

Large

Examples per second Examples per second Examples per second Examples per second Base Small Large Base CoDA P-Adapter achieves worse accuracy on almost all tasks and model sizes tested, getting 2.0 point worse on average. We also implement a 'no-learning' baseline that simply selects the first k tokens, which is equivalent to truncating the input sequence. 3 This baseline performs much worse, resulting in more than 10 point decrease in accuracy for small k (and equivalently large r ). This analysis confirms the importance of learning a good routing in order to retain strong model performance.

Reduction r 6 Full Results

SQuAD XSum Base Large XL Setup In this section, we apply our best training recipe to all tasks and application domains. We first pretrain dense Transformer models, followed by the CODA training procedure in §3.2. Our speech models are pretrained using a masked language modeling (MLM) objective similar to BERT [Devlin et al., 2019], and random quantized output label space [Chiu et al., 2022]. Our vision and text models use an encoder-decoder architecture similar to T5 but incorporate a few changes. Following PaLM [Chowdhery et al., 2022], we use multi-query attention [Shazeer, 2019] that shares the same key and value projection for multiple query heads. We only use 6 decoder layers and increase the feed forward hidden size (to compensate for the decrease in the number of layers). These modifications have a neutral effect on model quality, but speed up auto-regressive decoding significantly. We will show CODA is compatible with these changes and can further speed up inference by a considerably large factor. We provide more details of our experimental setup in Appendix A.

NLP results In addition to the classification datasets used in Section 5, we also evaluate our final models on the SQuAD, ReCord and XSum datasets which require generating an answer or a summary

SQuAD

CoDA Reduction r

XL

Original Image

Layer 0

Layer 8

Layer 17

Whoistheauthorofthisbook?

Whoistheauthorofthisbook?

Whoistheauthorof thisbook?

Whoistheauthorofthisbook?

FORMULA

FORMULA

FOFMULA

FORMUL

ONE

ONE

ONE

OINE

THEILLUSTRATEDHISTORY

THEILLUSTRATED HISTORY

THEILLUSTRATEDIHISTORY

THELLUSTRATEDHISTORY

BRUCEJONES

BRUCEJONES

BRUCEJONES

Fetrngechsecort

TONYBROOKS

76.2

given the input. Table 4 contains the finetuning results of XL models. Compared to the parallel adapter baseline that uses full computation, CODA achieves 3x and 5x computation reduction with only 1.0 and 1.7 point loss in average score.

Figure 4 and 5 highlight the scaling trend of CODA. CODA runs much faster with slightly worse quality than the parallel adapter baseline. This is expected because the baseline processes all tokens in every layer, whereas CODA only selects 1 /r of tokens for heavy processing. Importantly, this quality gap reduces as the model size increases (as shown in Figure 4), making CODA a computationally efficient choice for large models. Indeed, CODA can trade off quality for speed by varying the number of selected tokens. Figure 5 (left) demonstrates that CODA achieves much stronger speed-quality trade-off compared to dense models without conditional computation. The black line indicates the results of Parallel Adapter when the model size grows from Small to XL, and each blue line represents the speed-quality trade-off of CODA using r = 1 , 3 , 5 . Moreover, Figure 5 (middle) shows that larger CODA models exhibit higher inference speed-ups. These observations are consistent on other tasks. We provide additional results in Appendix §B.

Speech recognition results We further validate the performance of CODA in the speech domain. Our model uses a Transformer encoder and a 2-layer LSTM Transducer [Graves, 2012]. Similar to NLP setups, we test the performance of the speech model on 3 scales - Base, Large and XL (see Appendix A for details). Table 5 demonstrates that with sizable reduction ratios ( r = 2 , 4 ), the change on word error rate (WER) is consistently minimal on the test-clean and test-other sets of LibriSpeech across different model sizes (and on other sets in §B.2). Moreover, our results are comparable to the top-performing models, such as w2v-BERT [Chung et al., 2021] and BEST-RQ [Chiu et al., 2022], that are fully finetuned by updating all parameters. Figure 5 (right) highlight again that applying conditional computation leads to better speed-quality trade-off compared to dense models.

Vision results We extend our experiments to visual tasks that involves natural language within the image, such as documents and user interfaces. Our experiments are based on Pix2Struct [Lee et al., 2022], where an image-encoder-text-decoder is pretrained by learning to predict simplified HTML from webpage screenshots. Table 6 shows the results on three tasks that were also evaluated in the

original Pix2Struct paper. In OCRVQA and Screen2Words, we observe relatively small drops in performance when reducing the number of routed tokens (i.e. patches). When the capacity is 1/16th of the original sequence length, leading to around 13 × speedup, we only lose about 1 point. We speculate that this is due to the high-level sparsity in the inputs for these two tasks. For DocVQA, where there is comparatively more textual information, we observe a steeper performance-speed trade-off but still achieve a 8 × speedup with a 4-point drop.

To provide a more intuitive understanding why CODA works, we visualize the router behavior for the OCR-VQA model in Figure 6. We show which patches the routers prefers the most (warmest colors) and least (coolest colors), for several layers. The first, immediately obvious, observation is that router avoids low-frequency patches, i.e. patches likely to be 'whitespace', since they can be adequately handled by the cheap adapter layers. The second, more subtle, observation is that the router progressively converges on a small number of key patches that we hypothesize serve as representations for larger regions. The visualization confirms that CODA is able to select meaningful and representative patches that are useful for the prediction task.

7 Conclusion and Limitation

We present CODA, a parameter-efficient adapter method that enables fast inference. CODA relies on conditional computation to selectively activate model computation on important input units, providing a novel way to balance model expressivity and efficiency.

In this work, we focus on encoder-heavy applications such as summarization, speech recognition and visual question answering, by applying our method to the encoder. One limitation of CODA is that the current routing mechanism (i.e. token selection in a given sequence) is not directly applicable to decoder-only models for auto-regressive token generation. Enabling fast token generation using conditional activation in decoder layers is an interesting direction we plan to explore in future work.

References

Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https://www.tensorflow.org/ . Software available from tensorflow.org.

Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David Uthus, Mandy Guo, James Lee-Thorp, Yi Tay, et al. Colt5: Faster long-range transformers with conditional computation. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP) , 2023.

Anas Awadalla, Mitchell Wortsman, Gabriel Ilharco, Sewon Min, Ian Magnusson, Hannaneh Hajishirzi, and Ludwig Schmidt. Exploring the landscape of distributional robustness for question answering models. arXiv preprint arXiv:2210.12517 , 2022.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Ankur Bapna, Naveen Arivazhagan, and Orhan Firat. Controlling computation versus quality for neural sequence models. arXiv preprint arXiv:2002.07106 , 2020.

Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv:2004.05150 , 2020.

Luisa Bentivogli, Peter Clark, Ido Dagan, and Danilo Giampiccolo. The fifth pascal recognizing textual entailment challenge. In TAC , 2009. URL https://tac.nist.gov/publications/ 2009/additional.papers/RTE5_overview.proceedings.pdf .

Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your ViT but faster. In International Conference on Learning Representations , 2023.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, et al. Jax: composable transformations of python+ numpy programs. Version 0.2 , 5:14-24, 2018.

Chung-Cheng Chiu, James Qin, Yu Zhang, Jiahui Yu, and Yonghui Wu. Self-supervised learning with random-projection quantizer for speech recognition. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 3915-3924. PMLR, 17-23 Jul 2022. URL https://proceedings. mlr.press/v162/chiu22a.html .

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Yu-An Chung, Yu Zhang, Wei Han, Chung-Cheng Chiu, James Qin, Ruoming Pang, and Yonghui Wu. W2v-bert: Combining contrastive learning and masked language modeling for self-supervised speech pre-training. In 2021 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU) , pages 244-250. IEEE, 2021.

Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , 2019. URL https://aclanthology.org/N19-1300 .

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Proceedings of the PASCAL Challenges Workshop on Recognising Textual Entailment , 2005. URL http://www.cs.biu.ac.il/~glikmao/rte05/ .

Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. arXiv preprint arXiv:2305.14314 , 2023.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. ArXiv , abs/1810.04805, 2019.

Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learning , pages 5547-5569. PMLR, 2022.

William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.

Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and William B Dolan. The third pascal recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on textual entailment and paraphrasing , 2007. URL https://aclanthology.org/W07-1401.pdf .

Alex Graves. Sequence transduction with recurrent neural networks. arXiv preprint arXiv:1211.3711 , 2012.

Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-augmented transformer for speech recognition. arXiv preprint arXiv:2005.08100 , 2020.

Mandy Guo, Joshua Ainslie, David Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. LongT5: Efficient text-to-text transformer for long sequences. In Findings of the Association for Computational Linguistics: NAACL 2022 . Association for Computational Linguistics, 2022. URL https://aclanthology.org/2022.findings-naacl.55 .

R Bar Haim, Ido Dagan, Bill Dolan, Lisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and Idan Szpektor. The second pascal recognising textual entailment challenge. In Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entailment , 2006. URL http://u.cs.biu.ac.il/~nlp/RTE2/Proceedings/01.pdf .

Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. International Conference on Learning Representations (ICLR) , 2016.

Junxian He, Chunting Zhou, Xuezhe Ma, Taylor Berg-Kirkpatrick, and Graham Neubig. Towards a unified view of parameter-efficient transfer learning. arXiv preprint arXiv:2110.04366 , 2021.

Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for jax. Version 0.3 , 3: 14-26, 2020.

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 , 2015.

Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning , pages 2790-2799. PMLR, 2019.

Edward Hu, Yelong Shen, Phil Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021.

Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International Conference on Machine Learning , pages 9099-9117. PMLR, 2022.

Kahn, M. Rivière, W. Zheng, E. Kharitonov, Q. Xu, P. E. Mazaré, J. Karadayi, V. Liptchinsky, R. Collobert, C. Fuegen, T. Likhomanenko, G. Synnaeve, A. Joulin, A. Mohamed, and E. Dupoux. Libri-light: A benchmark for asr with limited or no supervision. In ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 7669-7673, 2020.

Yoon Kim and Alexander M. Rush. Sequence-level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing (EMNLP) , 2016. URL https://aclanthology.org/D16-1139 .

Harold W Kuhn and Albert W Tucker. Nonlinear programming. In Traces and emergence of nonlinear programming , pages 247-258. Springer, 2014.

Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, and Kristina Toutanova. Pix2struct: Screenshot parsing as pretraining for visual language understanding. arXiv preprint arXiv:2210.03347 , 2022.

Tao Lei. When attention meets fast recurrence: Training language models with reduced compute. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing . Association for Computational Linguistics, 2021. URL https://aclanthology.org/2021. emnlp-main.602 .

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing (EMNLP) , 2021. URL https://aclanthology.org/2021.emnlp-main.243 .

Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning , pages 6265-6274. PMLR, 2021.

Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2101.00190 , 2021.

Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei. Autoregressive knowledge distillation through imitation learning. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , 2020. URL https://aclanthology.org/2020.emnlp-main. 494 .

Andre Martins and Ramon Astudillo. From softmax to sparsemax: A sparse model of attention and multi-label classification. In Proceedings of The 33rd International Conference on Machine Learning , 2016.

Minesh Mathew, Dimosthenis Karatzas, and CV Jawahar. Docvqa: A dataset for vqa on document images. In Proceedings of the IEEE/CVF winter conference on applications of computer vision , pages 2200-2209, 2021.

Anand Mishra, Shashank Shekhar, Ajeet Kumar Singh, and Anirban Chakraborty. Ocr-vqa: Visual question answering by reading text in images. In 2019 international conference on document analysis and recognition (ICDAR) , pages 947-952. IEEE, 2019.

Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. ArXiv , abs/1808.08745, 2018.

Piotr Nawrot, Jan Chorowski, Adrian Ła´ ncucki, and Edoardo M. Ponti. Efficient transformers with dynamic token pooling, 2022.

Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: an asr corpus based on public domain audio books. In 2015 IEEE international conference on acoustics, speech and signal processing (ICASSP) , pages 5206-5210. IEEE, 2015.

Ofir Press, Noah A Smith, and Omer Levy. Improving transformer models by reordering their sublayers. arXiv preprint arXiv:1911.03864 , 2019.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research , 2020. URL http://jmlr.org/papers/ v21/20-074.html .

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing , 2016. URL https://nlp.stanford.edu/pubs/ rajpurkar2016squad.pdf .

Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor, Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorne, Aitor Lewkowycz, Alex Salcianu, Marc van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bulian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee, Dan Garrette, James Lee-Thorp, Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Maitin-Shepard, Noah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, and Andrea Gesmundo. Scaling up models and data with t5x and seqio . arXiv preprint arXiv:2203.17189 , 2022. URL https://arxiv.org/ abs/2203.17189 .

Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems , 34:17555-17566, 2021.

Bernhard Schmitzer. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM Journal on Scientific Computing , 41(3):A1443-A1481, 2019.

Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems , 2022.

Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150 , 2019.

Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 , 2020.

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538 , 2017.

Jonathan Shen, Patrick Nguyen, Yonghui Wu, Zhifeng Chen, et al. Lingvo: a modular and scalable framework for sequence-to-sequence modeling, 2019.

David So, Wojciech Ma´ nke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc V Le. Searching for efficient transformers for language modeling. Advances in Neural Information Processing Systems , 34:6010-6022, 2021.

Xiu Su, Shan You, Jiyang Xie, Mingkai Zheng, Fei Wang, Chen Qian, Changshui Zhang, Xiaogang Wang, and Chang Xu. Vision transformer architecture search. arXiv e-prints , pages arXiv-2106, 2021.

Kai Sheng Tai, Taipeng Tian, and Ser-Nam Lim. Spartan: Differentiable Sparsity via Regularized Transportation. In Advances in Neural Information Processing Systems , 2022.

Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv:1908.08962 , 2019.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Tu Vu, Aditya Barua, Brian Lester, Daniel Cer, Mohit Iyyer, and Noah Constant. Overcoming catastrophic forgetting in zero-shot cross-lingual generation. arXiv preprint arXiv:2205.12647 , 2022.

Bryan Wang, Gang Li, Xin Zhou, Zhourong Chen, Tovi Grossman, and Yang Li. Screen2words: Automatic mobile ui summarization with multimodal learning. In The 34th Annual ACM Symposium on User Interface Software and Technology , UIST '21, page 498-510, New York, NY, USA, 2021. Association for Computing Machinery. ISBN 9781450386357. doi: 10.1145/3472749.3474765. URL https://doi.org/10.1145/3472749.3474765 .

Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020a.

Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) . Association for Computational Linguistics, 2020b. URL https://aclanthology. org/2020.emnlp-main.496 .

Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers) , 2018. URL https://aclanthology.org/N18-1101 .

Stephen J Wright. Coordinate descent algorithms. Mathematical Programming , 151(1):3-34, 2015.

Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. In Association for Computational Linguistics (ACL) , 2022.

Qizhe Xie, Minh-Thang Luong, Eduard Hovy, and Quoc V Le. Self-training with noisy student improves imagenet classification. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10687-10698, 2020.

Hongxu Yin, Arash Vahdat, Jose M. Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov. A-ViT: Adaptive tokens for efficient vision transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 10809-10818, June 2022.

Yu Zhang, Daniel S Park, Wei Han, James Qin, Anmol Gulati, Joel Shor, Aren Jansen, Yuanzhong Xu, Yanping Huang, Shibo Wang, et al. Bigssl: Exploring the frontier of large-scale semi-supervised learning for automatic speech recognition. IEEE Journal of Selected Topics in Signal Processing , 16(6):1519-1532, 2022.

Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, and James Laudon. Mixture-of-experts with expert choice routing, 2022. URL https://arxiv.org/abs/2202.09368 .

Michael Zhu and Suyog Gupta. To prune, or not to prune: exploring the efficacy of pruning for model compression. arXiv preprint arXiv:1710.01878 , 2017.

A Experimental details

Model implementation For our text and vision experiments, we implement our models using JAX [Bradbury et al., 2018]. Specifically, our training and model modules are built on top of the T5X, Flax and Flaxformer framework [Roberts et al., 2022, Heek et al., 2020]. Following the T5 v1.1 implementation and PaLM [Chowdhery et al., 2022], our Transformer models use the GLU variant [Shazeer, 2020] as the feed forward network and multi-query-attention [Shazeer, 2019] as the attention block. These modifications are shown to enhance modeling capacity and speed up decoding respectively.

For the speech experiments, we use TensorFlow [Abadi et al., 2015] and the Lingvo framework [Shen et al., 2019]. The state-of-the-art Transformer variant for speech recognition is the Conformer architecture [Gulati et al., 2020] which additionally uses depth-wise convolution in each layer. Since the convolution operation is applied to consecutive inputs and does not immediately support routing, we use the standard Transformer architecture [Vaswani et al., 2017] instead. Swish activation is used in the feed forward blocks, following Gulati et al. [2020]. We provide the model configuration details in Table 7.

Model training We use the same data and procedure described in T5 [Raffel et al., 2020], BESTRQ [Chiu et al., 2022] and Pix2struct [Lee et al., 2022] for pre-training the respective text, speech and vision models. We use the same training hyper-parameters, such as batch size, input sequence length, the number of pre-training steps and the choice of optimizer and learning rate scheduling. All models have been pre-trained using 128 or 256 TPUv3/TPUv4 chips.

We run CODA pre-training for text and vision models, using an additional 100K steps and 200K steps respectively. For text models, the input sequence length is n = 512 and we set the number of selected tokens k = 192 . For vision models, the input sequence contains n = 4096 image patches and we set k = 1024 . CODA pre-training is not used for our speech models because there are

BoolQ

Examples per second

Examples per second

Reduction r

BoolQ

SQuAD

XSum

Base

Large

XL

(a)

Examples per second

Examples per second

Examples per second

F1 score

F1 score

Rouge-2

(b)

4

3.7

3.1

3

3

3

2.5

2.3

2.3

2

2

2

1.0

1.0

1.0

3

5

3

5

3

5

BoolQ

SQuAD

XSum

96

98

26

CoDA

P-Adapter

XL

85

92

20

Large

75

86

14

Base

Small

65

2°

80

602

2000000

8

0000

sufficient fine-tuning data. Following standard practice in speech, we use the 1K hour data from the LibriSpeech dataset [Panayotov et al., 2015] and another 30K hour data generated using the noisy student self-training method [Xie et al., 2020, Zhang et al., 2022].

Table 8 lists the hyper-parameters used for fine-tuning, including the sequence length, learning rate, batch size and the number of fine-tuning steps used. For NLP datasets, we set the maximum input length and decoding length to the 98th percentile of lengths in the training set. For vision datasets, we set the input length following the suggested values in Pix2struct. We also find that annealing the number of routed tokens k can achieve better finetuning results. Specifically, we decrease k linearly from the sequence length n down to the target value n/r using the first 10% to 20% of the finetuning steps.

CoDA Reduction r

B Additional results

B.1 NLP

Table 9 contains the complete fine-tuning results on the 6 language datasets. As discussed in §6, the gap between CODA and its counterpart without conditional computation is large at Base size. As the model size increases, CODA retains almost the same level of quality given 3x computation reduction ( r = 3 ). The reduction leads to decoding speed-ups, as shown in Figure 7. More importantly, we see that larger model benefits more from CODA, achieving a speed-up factor close to the reduction factor

r . These results highlight the potential of CODA for large-scale models, which we plan to investigate in future work.

B.2 Speech

Table 10 extends Table 5 by including WER results on dev-clean and dev-other splits. From the table, one can observe that XL with CoDA ( r = 2 , 4 ) are consistently better than the Large parallel adapter model on each split, and the Large model with CoDA ( r = 2 , 4 ) are also consistently better than the Base PA on each split. Given the inference speeds for CoDA models shown in Table 8, larger CoDA models are generally faster and better than smaller dense ones (even with PA) with regard to either time cost or computation GFLOPs. Therefore, it is likely for CoDA to help scale up ASR models with decent computation resources and time cost.

B.3 Combining CODA and LoRA

CODA can be easily combined with other types of adapter methods. To see this, we additionally implemented a variant that combines with Low-Rank Adapter [LoRA; Hu et al., 2021], which is another parameter-efficient transfer learning method that recently became the most popular choice for LLMs. We incorporate the latest development suggested in the QLoRA paper [Dettmers et al., 2023], which adds low-rank adapters to every linear projection matrix in the Transformer layers. This is found to obtain better fine-tuning performance than the original implementation. Our CODA variant with LoRA simply removes the parallel adapter branches and instead adds low-rank adapters to the projection matrices of the pretrained layers.

Table 11 shows the finetuning results. The new LoRA baseline achieves stronger accuracy than the Parallel Adapter baseline (84.0 v.s. 82.9 on average), highlighting the effectiveness of recent development on LoRA. In addition, our CODA variant using LoRA still achieves very close accuracy compared to its dense counterpart (84.0 v.s. 84.0 or 83.7 on average). We believe the additional

Small

Base

Large

XL

Examples per second

Speedv.s.Quality

GFLOPs

1200

2.8

2.6

800

2.4

400

2.2

2.0

0

Base

Large

XL

Dense

CODA(r=2)

■

CODA (r=4)

results strengthen our claims - that CODA enables a strong trade-off between accuracy and efficiency using conditional activation, and this technique can be combined with other developments in PETL.

C Soft topk algorithm

C.1 Derivation of the iterative updates

We present the derivation of iterative updates (11) for solving the soft topk problem (10) in Section 3. The soft topk operation is defined as a maximization problem (10). For the derivation, we rewrite it as an equivalent minimization problem:

Note the term glyph[epsilon1] 1 glyph[latticetop] λ will be a constant glyph[epsilon1] × k , but we include it in the minimization object to make our derivation simpler later.

Now, let a ∈ R and b ∈ R m be the Lagrangian variables corresponding to the linear constraints 1 glyph[latticetop] λ = k and λ [ i ] ≤ 1 ∀ i . 4 The minimization problem is equivalent to its Lagrangian expression:

The objective function (12) is strongly convex and the solution space of λ is a convex set. As a result, strong duality holds so we can instead solve the dual problem. The dual problem exchanges the min and max operators in (13):

The optimal solution ( a, b , λ ) must have the Karush-Kuhn-Tucker (KKT) conditions hold [Kuhn and Tucker, 2014], namely

Substituting λ using the above equation in (14), the dual problem now has a simple form:

We can solve this problem using coordinate descent [Wright, 2015] by successively maximizing the function with either a or b fixed. That is, we find the optimal a that maximizes the dual objective given a fixed b , and vice versa. This leads to the iterative updates (11) described in Section 3.

In short, we obtain the iterative updates of the soft topk problem (10) by solving its dual problem and by performing coordinate decent of the dual variables a and b . The iterative updates are in fact the coordinate decent steps.

C.2 The glyph[epsilon1] -scaling trick

The iterations of a and b will converge but the number of iterations needed can be very large for small glyph[epsilon1] . In practice, we only perform a small number of iterations and return the corresponding λ , which may be close but not the exact solution to (12). In order to improve the convergence given a small number of iterations, we apply an empirical trick called the glyph[epsilon1] -scaling heuristic [Schmitzer, 2019]. Let glyph[epsilon1] t denote the value of glyph[epsilon1] at the t -th iteration. We initialize glyph[epsilon1] 0 to a larger value and gradually reduce glyph[epsilon1] t to the target glyph[epsilon1] . Specifically, we set glyph[epsilon1] t = max( βglyph[epsilon1] t -1 , glyph[epsilon1] ) at the t -th iteration, using a scaling constant β ∈ (0 , 1) . We use glyph[epsilon1] 0 = 4 throughout our experiments, glyph[epsilon1] = 0 . 03 and β = 0 . 7 for text and vision models and glyph[epsilon1] = 1 . 0 and β = 0 . 85 for speech models. Using a larger number of iterations leads to better convergence but we found T = 20 sufficient for our experiments.

C.3 Overhead of soft topk iterations

The soft topk iterations are performed for every routed Transformer layer. Although this seems computationally expensive, the actual overhead is very small compared to the overall decoding latency. The complexity only scales linearly with the number of layers and the sequence length, and does not depend on the model dimension d . Table 12 showcases the latency numbers on the BoolQ and XSum datasets, when performing batched decoding using a single TPUv4 chip. We observe that the cost of iterations is less than 2% of the total decoding latency. Moreover, the relative cost decreases dramatically as the model size increases, since it does not depend on the model dimension.

C.4 Additional discussion

This iterative algorithm is closely related to the Sinkhorn algorithm of Optimal Transport (OT). Specifically, the Sinkhorn algorithm solves the entropy-regularized version of Optimal Transport [Cuturi, 2013]. However, our work concerns an different optimization instance. While OT solves a transportation problem where the solution space is defined with the marginal constraints over the rows and columns of a transportation matrix, our optimization problem is constrained with a total budget ( ∑ i λ i = k ) and upper bounds ( λ i ≤ 1 ∀ i ). This leads to different iterative updates.

Concurrent to our work, Tai et al. [2022] have used a similar linear program (LP) formulation for soft topk operation, and have applied the operator for learning sparse neural networks (i.e. model pruning). Compared to our formulation (12), they first reduce the LP to an equivalent instance of optimal transport problem, before introducing the entropy term. As a result, the derived updates are different. In addition, Tai et al. [2022] have introduced an initialization for the dual variables to improve the convergence of their algorithm, whereas we use glyph[epsilon1] scaling instead. Their implementation can be explored for CODA as well.

Besides formulating soft topk using entropy-regularized optimization, there are other possible variants for trainable sparsity. One example is sparsemax [Martins and Astudillo, 2016] that can learn sparse multi-label probabilities. We believe that the sparsemax formulation can generalize from the top-1 to topk case, but it is beyond the scope of this work. We use the current soft topk implementation because it is a natural extension of softmax (see discussions in §3), and because it can be solved using simple iterative updates.

D Author Contributions

All authors have contributed to running experiments and discussing research ideas. Tao leads the project, developed the conditional architecture, designed the experiments and analyses. Kenton, Yu and Ming-Wei proposed the idea of applying conditional computation for large model adaptation. Joshua demonstrated the conditional architecture is applicable to attention, and implemented the initial version of conditional attention block. Tao, Yanqi, Nan, Vincent, Yuexin, Ming-Wei and Yu conducted the NLP experiments including model pre-training, fine-tuning and various ablation analyses. Siddhartha conducted the majority of the vision experiments. Kenton conducted the vision analysis and advised on the vision experiments. Junwen conducted the majority of the speech experiments. Bo and Yu assisted in trouble-shooting the speech models, ran the model pre-training and provided guidance on the speech experiments. Finally, Tao, Ming-Wei, Junwen and Kenton made the primary contributions to the writing of the paper.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes Conditional Adapter (CODA), which not only enables parameter-efficient fine-tuning but is also beneficial for the inference speedup. The main idea of CODA is to pass all tokens to a lightweight routing adapter path while selectively passing only certain tokens to the original path. The selection process, named the “Soft Top-K” operation, is conducted by solving an optimization problem that utilizes a generalized entropy function and an iterative searching algorithm. Experimental results on Encoder-Decoder Transformers in NLP, Vision, and Speech applications demonstrate significant inference acceleration with only minimal performance degradation.\n\nSTRENGTHS:\n* The idea of CODA is simple yet powerful. The authors support the effectiveness of the proposed method through sufficient experiments. Although the baseline (Parallel Adapter) achieves higher performance than CODA, I believe the inference speedup can compensate for the performance loss.\n* The derivation and implementation of the “Soft Top-K” operation are particularly unique. This operation outperforms simple sigmoid gating or Top-K truncation. Furthermore, this differentiable function has the potential for a broader range of neural network research.\n\nWEAKNESSES:\n* Although the authors briefly mentioned CoLT5, it would be better to discuss more the differences between CoLT5 and CODA. The current explanation (L78-79) is somewhat insufficient.\n\nQUESTIONS:\n* How much is the burden of the iterative searching process (T=20) during training? In other words, how does the training time/performance change as iteration steps T are reduced?\n* Regarding Figure 5(b), how could the decoding be accelerated? As I understand, CODA does not remove the number of tokens, so the input #tokens to the cross-attention would also not be reduced. Additionally, the readability could be improved if Figure 5 is enlarged in the manuscript (minor issue).\n* (Suggestion) I expect CODA can also reduce memory consumption for both training and inference. It would strengthen the paper’s value if there were a comparison regarding memory usage.\n* (Suggestion) CODA currently fixes the reduction ratio r for all layers. I understand that it is a simple and structured way, but how about setting different r for different layers? For example, small/large r for low/high layers. This could provide additional insights.\n\nLIMITATIONS:\nThe authors have addressed the limitations in Section 7.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a unique modification over parameter-efficient transfer learning by employing a conditional adapter parallel to the pre-trained model. In doing so, a dynamic token selection mechanism is proposed for each block. The selected tokens undergo heavy pre-trained operations, while all the tokens undergo light-weight adapter operations. This greatly improves the inference efficiency of the target dataset. Experiments from language, vision and speech recognition are shown to test the efficacy of the proposed method.\n\nSTRENGTHS:\n1. Wide variety of experiments provided show the superiority of CoDA as compared to parallel-adapter among others.\n2. The overall idea of improving the inference speed of large models during the transfer learning stage has huge practical value.\n3. A large chuck of questions have been addressed in the main text including - iteration wise performance, the importance of routing, etc.\n\nWEAKNESSES:\n1. The presentation of the work can be improved; after section 3.2, the paper seems unorganised, and the scale and location of tables and plots also require some work.\n2. Thorough comparison with existing works is missing. There are two streams of comparison possible - [1] Parameter-efficient fine-tuning methods [2] Dynamic/conditional inference methods. Even though experiments are shown on a variety of tasks, it is still unclear how this work compares with existing methods.\n3. Final latency with respect to the baseline model is not thoroughly presented, also additional parameters due to parallel branches are not given.\n\nQUESTIONS:\n1. Is there a simple softmax baseline instead of the top-soft-k operation?\n2. Is the iterative process for soft top-k done for each forward pass and each layer? If yes, how much overhead does it add?\n3. What is the parameter overhead of CoDA and absolute latency numbers for different tasks?\n4. There are a variety of standard benchmarks for parameter efficient fine-tuning - VTAB-1k (vision), GLUE (language), among others; why did the authors choose specific benchmarks presented in the paper? If possible, can the authors present results on these standard benchmarks?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed Conditional Adapter (CODA), which can achieve good trade-offs between speed and accuracy. CODA achieves 2x to 8x inference speed-up compared to the state-of-the-art Adapter approach while maintaining moderate to no accuracy loss and the same parameter efficiency across language, vision, and speech tasks.\n\nSTRENGTHS:\n1. The proposed method can greatly speed up inference, while keeping competitive performance.\n2. The writing of the paper is clear and easy to understand.\n3. The experiments are extensive, which can also well support the arguments of this paper.\n\nWEAKNESSES:\n1.  This paper proposes a new adapter and claims to accelerate model inference. However,  there is no direct relationship between the adapter and the inference acceleration. The token pruning approach appears more like an independent module. Therefore,  this paper is more like a combination of token pruning and traditional adapters, and both approaches are not new in the community.  \n2. The token pruning  is not new, as there have been many methods applied to ViT. Some methods can even achieve inference acceleration without fine-tuning. The author should disscuss the differences with these methods.\n\nToken Merging: Your ViT But Faster, ICLR, 2023\nDynamicvit: Efficient vision transformers with dynamic token sparsification. NeurIPS, 2021\nA-ViT: Adaptive tokens for efficient vision transformer, CVPR, 2022\n\n3. Based on 1, I'm not quite clear whether this article proposes an adapter method or a pruning method. From the author's descriptions, i lean towards the former. However, in the experimental section, the author lacks a comparison with newly proposed  parameter-efficient transferring learning (PETL) methods, such as   AdaMix, MAM adapter, and so on.\n\nAdamix: Mixture-of-adapter for parameter-efficient tuning of large language models, EMNLP, 2022\nTOWARDS A UNIFIED VIEW OF PARAMETER-EFFICIENT TRANSFER LEARNING, ICLR, 2022\n\n4. It is unclear whether this method can be applied to the currently popular autoregressive models.\n\nQUESTIONS:\nSee weakness\n\nLIMITATIONS:\nSee weakness", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
The paper introduces a novel technique called Conditional Adapters (CODA), which aims to enhance transfer learning and improve inference efficiency. CODA surpasses the state-of-the-art Adapter approach by achieving a remarkable 2x to 8x acceleration in inference speed while maintaining moderate accuracy loss and identical parameter efficiency. The paper provides examples of language, vision, and speech tasks where CODA has been extensively evaluated, showcasing the promising results obtained through its application.

STRENGTHS:
1. This paper addresses an intriguing problem of accelerating inference speed through efficient parameter tuning.
2. It is commendable that the proposed method is evaluated across three different modalities, indicating its versatility.
3. The model's simplicity and effectiveness are notable. The overall design is straightforward, and the method is technically sound.
4. The paper is well-written and easy to follow.

WEAKNESSES:
1. The availability of the source code is unclear, and it would be beneficial to know if it will be made publicly available for further exploration and reproducibility.
3. Can not be applied to current main-stream auto-regressive models.
2. While the paper covers text and speech modalities, for the visual modality, it would be valuable to benchmark the proposed method on tasks such as image recognition and object detection, in addition to OCR-VQA and Doc-VQA.

QUESTIONS:
See weakness.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-12