SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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
Figure 8: The scaling of CoDA on Librispeech. Left: Larger models with CoDA can achieve much better performance with the same or faster speed compared to smaller models. Right: GFLOPS reduction is significant as  $r$  increases.
Model	Reduction r	Base	Large	Avg on 6
MNLI	RTE	BoolQ	MNLI	RTE	BoolQ
Parallel Adapter (PA)	-	87.1	71.5	77.9	90.3	84.8	85.8	82.9
CoDA w/ PA	3	86.6	72.6	76.6	90.2	85.9	85.1	82.8
CoDA w/ PA	5	86.0	70.8	76.0	89.7	85.2	84.3	82.0
Low-rank Adapter (LoRA)	-	88.0	73.7	80.3	90.7	85.2	86.3	84.0
CoDA w/ LoRA	3	86.2	76.9	78.4	90.3	86.3	85.8	84.0
CoDA w/ LoRA	5	86.0	76.9	78.3	89.8	86.3	84.7	83.7
Table 11: Results of applying CoDA with Parallel Adapter v.s. Low-Rank Adapter. We report results on the development sets and on the Base and Large model size.
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 top-  $k$  algorithm
C.1 Derivation of the iterative updates
We present the derivation of iterative updates (11) for solving the soft top- $k$  problem (10) in Section 3. The soft top- $k$  operation is defined as a maximization problem (10). For the derivation, we rewrite it as an equivalent minimization problem:
$$
\max  _ {\boldsymbol {\lambda}} \boldsymbol {s} ^ {\top} \boldsymbol {\lambda} + \epsilon H (\boldsymbol {\lambda})
$$
$$
\iff \quad \min  _ {\boldsymbol {\lambda}} - s ^ {\top} \boldsymbol {\lambda} - \epsilon H (\boldsymbol {\lambda})
$$
$$
\Leftrightarrow \quad \min  _ {\boldsymbol {\lambda}} - s ^ {\top} \boldsymbol {\lambda} - \epsilon H (\boldsymbol {\lambda}) - \epsilon \mathbf {1} ^ {\top} \boldsymbol {\lambda} \tag {12}
$$
$$
\begin{array}{l} \text {s . t .} \mathbf {1} ^ {\top} \boldsymbol {\lambda} = k, \quad \boldsymbol {\lambda} [ i ] \in [ 0, 1 ], \quad i = 1, \dots , n. \end{array}
$$
Note the term  $\epsilon \mathbf{1}^{\top}\lambda$  will be a constant  $\epsilon \times k$ , but we include it in the minimization object to make our derivation simpler later.
Now, let  $a \in \mathbb{R}$  and  $\pmb{b} \in \mathbb{R}^{m}$  be the Lagrangian variables corresponding to the linear constraints  $\mathbf{1}^{\top} \pmb{\lambda} = k$  and  $\pmb{\lambda}[i] \leq 1 \forall i$ . The minimization problem is equivalent to its Lagrangian expression:
$$
\min  _ {\lambda \in \mathbb {R} ^ {m}} \max  _ {a \in \mathbb {R}, b \leq 0} - s ^ {\top} \boldsymbol {\lambda} - \epsilon H (\boldsymbol {\lambda}) - \epsilon \mathbf {1} ^ {\top} \boldsymbol {\lambda} + a (k - \mathbf {1} ^ {\top} \boldsymbol {\lambda}) + \boldsymbol {b} ^ {\top} (\mathbf {1} - \boldsymbol {\lambda}) \tag {13}
$$
${}^4\lambda [i]\geq 0\forall i$  is already implied by the term  $H(\lambda) = \sum_{i} - \lambda [i]\log \lambda [i]$  in the objective, due to the use of  $\log \pmb {\lambda}[i]$
The objective function (12) is strongly convex and the solution space of  $\lambda$  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):
$$
\max  _ {a \in \mathbb {R}, b \leq 0} \min  _ {\lambda \in \mathbb {R} ^ {m}} - s ^ {\top} \boldsymbol {\lambda} - \epsilon H (\boldsymbol {\lambda}) - \epsilon \mathbf {1} ^ {\top} \boldsymbol {\lambda} + a (k - \mathbf {1} ^ {\top} \boldsymbol {\lambda}) + b ^ {\top} (\mathbf {1} - \boldsymbol {\lambda}) \tag {14}
$$
The optimal solution  $(a, b, \lambda)$  must have the Karush-Kuhn-Tucker (KKT) conditions hold [Kuhn and Tucker, 2014], namely
$$
\begin{array}{l} \frac {\partial \left(- s ^ {\top} \boldsymbol {\lambda} - \epsilon H (\boldsymbol {\lambda}) + \epsilon \mathbf {1} ^ {\top} \boldsymbol {\lambda} + a (k - \mathbf {1} ^ {\top} \boldsymbol {\lambda}) + \boldsymbol {b} ^ {\top} (\mathbf {1} - \boldsymbol {\lambda})\right)}{\partial \boldsymbol {\lambda}} = 0 \\ \iff \quad \boldsymbol {\lambda} = \exp \left(\frac {\boldsymbol {s} + \boldsymbol {a} + \boldsymbol {b}}{\epsilon}\right) \quad \Longleftrightarrow \quad \boldsymbol {\lambda} [ i ] = \exp \left(\frac {\boldsymbol {s} [ i ] + \boldsymbol {a} + \boldsymbol {b} [ i ]}{\epsilon}\right) \forall i = 1, \ldots , n \\ \end{array}
$$
Substituting  $\lambda$  using the above equation in (14), the dual problem now has a simple form:
$$
\max  _ {a \in \mathbb {R}, b \leq \mathbf {0}} k \cdot a + \mathbf {1} ^ {\top} \boldsymbol {b} - \mathbf {1} ^ {\top} \exp \left(\frac {\boldsymbol {s} + a + \boldsymbol {b}}{\epsilon}\right)
$$
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.
$$
a ^ {\prime} = \epsilon \ln (k) - \epsilon \ln \left(\sum_ {i = 1} ^ {n} \exp \left(\frac {\pmb {s} [ i ] + \pmb {b} [ i ]}{\epsilon}\right)\right), \qquad \pmb {b} ^ {\prime} = \min (- \pmb {s} - a ^ {\prime}, 0)
$$
In short, we obtain the iterative updates of the soft top- $k$  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  $\epsilon$ -scaling trick
The iterations of  $a$  and  $b$  will converge but the number of iterations needed can be very large for small  $\epsilon$ . In practice, we only perform a small number of iterations and return the corresponding  $\lambda$ , 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  $\epsilon$ -scaling heuristic [Schmitzer, 2019]. Let  $\epsilon_t$  denote the value of  $\epsilon$  at the  $t$ -th iteration. We initialize  $\epsilon_0$  to a larger value and gradually reduce  $\epsilon_t$  to the target  $\epsilon$ . Specifically, we set  $\epsilon_t = \max(\beta \epsilon_{t-1}, \epsilon)$  at the  $t$ -th iteration, using a scaling constant  $\beta \in (0,1)$ . We use  $\epsilon_0 = 4$  throughout our experiments,  $\epsilon = 0.03$  and  $\beta = 0.7$  for text and vision models and  $\epsilon = 1.0$  and  $\beta = 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 top- $k$  iterations
The soft top- $k$  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 [Caturi, 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  $(\sum_{i}\lambda_{i} = k)$  and upper bounds  $(\lambda_{i}\leq 1\forall i)$ . This leads to different iterative updates.
Task	Model size	Iteration latency (ms)	Total latency (ms)
BoolQ	Base	0.9	48.0
Large	1.7	105.3
XL	1.7	297.6
XSum	Base	1.0	513.9
Large	2.0	1076.2
XL	2.0	2426.0
Table 12: Latency (in milliseconds) of the soft top-  $k$  iteration and the total decoding time per batch. We use a single TPUv4 chip and 128 sequences per batch. The maximum iteration overhead is less than  $2\%$  of the total latency.
Concurrent to our work, Tai et al. [2022] have used a similar linear program (LP) formulation for soft top- $k$  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  $\epsilon$  scaling instead. Their implementation can be explored for CoDA as well.
Besides formulating soft top- $k$  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 top- $k$  case, but it is beyond the scope of this work. We use the current soft top- $k$  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.