SlowGuess/NAACL_2022_OCR huggingface dataset:
Empowering Parameter-Efficient Transfer Learning by Recognizing the Kernel Structure in Attention
Yifan Chen $^{1*}$  Devamanyu Hazarika $^{2*}$  Mahdi Namazifar $^{2}$  Yang Liu $^{2}$  Di Jin $^{2\dagger}$  Dilek Hakkani-Tur $^{2}$
<sup>1</sup>University of Illinois Urbana-Champaign <sup>2</sup>Amazon Alexa AI
The massive amount of trainable parameters in the pre-trained language models (PLMs) makes them hard to be deployed to multiple downstream tasks. To address this issue, parameter-efficient transfer learning methods have been proposed to tune only a few parameters during fine-tuning while freezing the rest. This paper looks at existing methods along this line through the kernel lens. Motivated by the connection between self-attention in transformer-based PLMs and kernel learning, we propose kernel-wise adapters, namely Kernel-mix, that utilize the kernel structure in self-attention to guide the assignment of the tunable parameters. These adapters use guidelines found in classical kernel learning and enable separate parameter tuning for each attention head. Our empirical results, over a diverse set of natural language generation and understanding tasks, show that our proposed adapters can attain or improve the strong performance of existing baselines.
Transfer learning using large-scale transformer-based pre-trained language models (PLMs) (Radford et al., 2019) has become the standard scheme for various natural language processing (NLP) tasks. Among many strategies, fine-tuning these PLMs emerges as the predominant strategy to adapt the generic models to a specific task (Howard and Ruder, 2018). However, deploying these models is a challenge as curating customized models across a wide variety of tasks would lead to scalability issues. It requires one to store (and sometimes move) multiple copies of the PLM parameters for different tasks, which is inefficient.
A popular approach to tackling such scalability issues is to make the PLM-based transfer learning more parameter-efficient. This can be done
by freezing most of the PLM parameters and inserting small trainable modules into the PLM. Adapters (Houlsby et al., 2019; Pfeiffer et al., 2020; Mahabadi et al., 2021; Hu et al., 2021) and Prefix-Prompt-tuning (Shin et al., 2020; Li and Liang, 2021; Lester et al., 2021; Liu et al., 2021c,b) have emerged as the prominent approaches under this paradigm. These methods are incredibly parameter-efficient and have comparable performance to full fine-tuning models on many common NLP tasks (mainly in Natural Language Understanding) by tuning only  $0.1 - 3\%$  task-specific parameters of the original PLMs.
However, most of these studies take the PLMs as a black box, i.e., these methods are not customized to transformers. This raises whether parameter-efficient fine-tuning has fully utilized the transformer structure in PLMs. Therefore in this work, we propose kernel-wise adaptation, which recognizes and utilizes the kernel structure within self-attention—the core component in a transformer. Specifically, we take inspiration from recent work that connects self-attention to kernel learning (Choromanski et al., 2020; Chen et al., 2021; Tsai et al., 2019) to treat the different attention heads in a transformer's attention sub-layer as separate kernel estimators. We hypothesize that parameter-efficient tuning can benefit from some useful guidelines in classical kernel learning literature and incorporate them into our proposed methods. These include:
1. By interpreting attention heads as kernel estimators, we design the adaptation to be head-specific;
2. We assign more budgets of tunable parameters to learn the value components in the attention mechanism, which correspond to coefficients in kernel methods.
We discuss these guidelines in detail in  $\S 4.1$ . We also evaluate our hypotheses through rigorous
*Equal contribution. This work was performed while the first author was interning at Amazon Alexa AI.
† Correspondence to: Di Jin <djinamzn@amazon.com>
Findings of the Association for Computational Linguistics: NAACL 2022, pages 1375 - 1388 July 10-15, 2022 ©2022 Association for Computational Linguistics
empirical evaluation. First, we test the effectiveness of the two guidelines above by comparing the default LoRA (Hu et al., 2021)—a state-of-the-art approach for efficient adaptation—and two of its variants that implement the two guidelines, respectively. Next, we evaluate our variant of kernel-wise adaptation on three Natural Language Generation (NLG) benchmarks and two Natural Language Understanding (NLU) tasks using the GPT-2 architecture. While parameter-efficient work has extensively covered NLU tasks, it is unknown how well the results transfer to NLG tasks. As language generation typically requires more expressive models, we put more emphasis on multiple NLG tasks that include data-to-text, free-form question answering (QA), and summarization. The empirical results in §6 demonstrate that with the same parameter budgets, our proposed method can attain better generation quality and classification accuracy than previous techniques, and in many settings, it is close to or even outperforms the full parameter fine-tuning.
The literature on parameter-efficient adaptation can be broadly categorized as follows:
Adapters. Originally proposed by Houlsby et al. (2019); Pfeiffer et al. (2020), adapters modulate the output of a transformer layer by inserting small Multi-layer Perception (MLP) bottleneck layers. Recent work has proposed many variants of the original adapters, including dropping adapters across several layers (Rückle et al., 2020) or constraining adapters to be low-rank operators (Mahabadi et al., 2021).
A recent line of work focused on identifying the important subset of parameters within the PLMs. Ben Zaken et al. (2021) proposed to only tune the bias terms in the PLMs. MPOP (Liu et al., 2021a) suggested decomposing the weight matrices in PLMs through matrix product operators (MPO) and only trained the matrices of small size (freezing the large matrices) obtained from the decomposition, which implicitly recognizes the small matrices as the important subset.
Low-rank adaptation (LoRA) (Hu et al., 2021) directly assumes that the update of the weight matrices during training can be approximate low-rank, and accordingly proposed to re-parameterize the original weight matrix  $\mathbf{W}$  by  $\mathbf{W} + \mathbf{BA}$ , where  $\mathbf{W}$  is frozen to its pre-trained weights whereas  $\mathbf{A} \in \mathbb{R}^{r \times N_h p}$ ,  $\mathbf{B} \in \mathbb{R}^{N_h p \times r}$  are updated in training
1. We note that LoRA too introduces new weights  $A$  and  $B$  similar to adapters, but they are used only to re-parameterize the existing weights and do not add extra sandwiched layers that modify the original model architecture.
Prefix-tuning. Originally shown in GPT-3 (Brown et al., 2020), prompts are extra tokens that help in the task adaptation of PLMs. Transitioning from the manual design of prompts, Shin et al. (2020) searched for the prompts over the discrete space of tokens based on the task-specific training data; Li and Liang (2021); Lester et al. (2021); Liu et al. (2021c,b) further extended the search space to continuous prompts and tuned the prompts through back-propagating the error in training. Prompt-based methods have been shown to be similar to adapters by (He et al., 2021).
We start by providing a brief introduction to the transformer architecture (§ 3.1) and then revisit the connection between attention and kernel estimators (§ 3.2), building on which we propose the kernel-wise adapter in § 4.
3.1 Transformer Architecture
Transformers (Vaswani et al., 2017) are composed of  $L$  stacked layers, where each layer comprises of a multi-headed attention and a fully connected feedforward network (FFN) sub-layer. The attention sub-layer, assuming  $N_{h}$  heads and dimension size  $p$  for each head, first maps an input  $\mathbf{X} \in \mathbb{R}^{n \times N_{h}p}$  into the query  $(Q)$ , key  $(K)$ , and value  $(V)$  matrices through the following affine transformations:
$$
\boldsymbol {Q} / \boldsymbol {K} / \boldsymbol {V} = \boldsymbol {X} \boldsymbol {W} _ {[ q / k / v ]} + \boldsymbol {1 b} _ {[ q / k / v ]} ^ {T}, \tag {1}
$$
where  $Q, K, V \in \mathbb{R}^{n \times N_{h}p}$ ,  $W_{q}, W_{k}, W_{v}$  are  $N_{h}p$ -by- $N_{h}p$  weight matrices, and  $\boldsymbol{b}_{q}, \boldsymbol{b}_{k}, \boldsymbol{b}_{v} \in \mathbb{R}^{N_{h}p}$  are the bias terms<sup>3</sup>. After the transformation, the three components  $Q, K, V$  are split into  $N_{h}$  blocks corresponding to different heads. For example,  $Q$  is re-written as  $Q = (Q^{(1)}, \dots, Q^{(N_{h})})$ , where each block  $Q^{(h)} = X W_{q}^{(h)} + 1(b_{q}^{(h)})^{T}$  is an
In implementation, LoRA also trains the bias terms in the linear transform besides the matrices  $A, B$ , while for brevity, the bias terms are omitted throughout the paper.
2For simplicity we omit the cross-attention module in transformer-based encoder-decoder models.
3To ease the notations we adopt the setting where  $X,Q,K,V$  have the same shape.
$n$ -by- $p$  matrix, and  $\mathbf{W}_q^{(h)}, \mathbf{b}_q^{(h)}$  are the corresponding parts in  $\mathbf{W}_q, \mathbf{b}_q$ . The attention output for the  $h^{th}$  head is then computed as:
$$
\begin{array}{l} \boldsymbol {L} ^ {(h)} \boldsymbol {V} ^ {(h)} := \operatorname {s o f t m a x} (\boldsymbol {Q} ^ {(h)} (\boldsymbol {K} ^ {(h)}) ^ {T} / \sqrt {p}) \boldsymbol {V} ^ {(h)} \\ = \left(\boldsymbol {D} ^ {(h)}\right) ^ {- 1} \boldsymbol {M} ^ {(h)} \boldsymbol {V} ^ {(h)}, \tag {2} \\ \end{array}
$$
where  $M^{(h)}\coloneqq \exp \left(Q^{(h)}(K^{(h)})^T /\sqrt{p}\right)$  and  $D^{(h)}$  is a diagonal matrix in which  $D_{ii}^{(h)}$  is the sum of the  $i$ -th row in  $M^{(h)}$ , corresponding to the normalization part in softmax.
After we obtain the outputs in each head, they are concatenated as,
$$
\boldsymbol {L} := (\boldsymbol {L} ^ {(1)} \boldsymbol {V} ^ {(1)}, \dots , \boldsymbol {L} ^ {(N _ {h})} \boldsymbol {V} ^ {(N _ {h})}),
$$
followed by the overall output,
$$
\boldsymbol {L} \boldsymbol {W} _ {o} + \boldsymbol {1} \boldsymbol {b} _ {o} ^ {T}, \tag {3}
$$
where  $\mathbf{W}_o$  and  $\mathbf{b}_o$  are similarly sized as the other matrices in Equation (1).
3.2 Attention as Kernel Estimators
For each head in the attention module, we have given the expression of attention output in Equation 2. In this subsection, we will re-write attention as a kernel estimator to show the connection.
In computing the attention output (of a single head), we have a length- $n$  input sequence  $\{x_{i}\}_{i=1}^{n}$  (the rows in  $X$ ) and accordingly we can obtain  $N^4$  key vectors  $\{k_{j}\}_{j=1}^{N} \subset \mathbb{R}^{p}$  (from the key matrix  $K$ ) and query vectors  $\{q_{i}\}_{i=1}^{n} \subset \mathbb{R}^{p}$  (from  $Q$ ). The original goal of self-attention is to obtain the representation of each input token  $x_{i}$ :  $g(x_{i})$ . By denotation exchange:  $q_{i} := x_{i}$  and  $f(q_{i}) := g(x_{i})$ , we can also understand the aforementioned self-attention module as returning the representation  $f(q_{i})$  of the input query vector  $q_{i}$  through  $\{k_{j}\}_{j=1}^{n}$ , which behaves as a kernel estimator (Choromanski et al., 2020; Peng et al., 2020; Chen et al., 2021). Specifically, for a single query vector  $q_{i}$ , a Nadaraya-Watson kernel estimator (Wasserman, 2006, Definition 5.39) models its representation as,
$$
f \left(q _ {i}\right) = \sum_ {j = 1} ^ {n} \ell_ {j} \left(q _ {i}\right) c _ {j}, \tag {4}
$$
where  $\ell_{j}(q_{i})\coloneqq \frac{\kappa(q_{i},k_{j})}{\sum_{k = 1}^{N}\kappa(q_{i},k_{j})}.$
Here,  $\kappa (\cdot ,\cdot)$  is a kernel function, and  $c_{j}$ 's are the coefficients ( $c_{j}$  can either be a scalar or a vector in different applications) that are learned during training. In this estimator,  $\{k_j\}_{j = 1}^n$  serve as the supporting points which help construct the representation for an input  $q_{i}$ .
For kernel function  $\kappa (x,y) = \exp \left(\langle x,y\rangle /\sqrt{p}\right)$ , we slightly abuse the notation  $\kappa (\boldsymbol {Q},\boldsymbol {K})$  to represent an  $n$ -by- $N$  empirical kernel matrix, whose element in the  $i$ -th row and the  $j$ -th column is  $\kappa (q_i,k_j),\forall i\in [n],j\in [N]$ . With these notations, the representation of the whole sequence  $\mathbf{Q}$  will be,
$$
D ^ {- 1} \kappa (Q, K) C, \tag {5}
$$
where  $D$  is a diagonal matrix for row normalization in Eq. (4), and  $C$  is an  $N$ -by- $p$  matrix whose  $j$ -th row is  $c_{j}$ . Considering the correspondence between Equation (5) and the standard softmax attention in Equation (2), we can have a finer division of the attention module: the empirical kernel matrix  $\kappa(Q, K)$  ( $D$  is decided by  $\kappa(Q, K)$ ) and the coefficient part  $C$ , which includes but is not limited to value matrices in attention (see § 4). In what follows, we will discuss how we build adapters for these two parts.
We introduce our Kernel-mix method in this section, which builds upon the proposed Kernel-wise adaptation. To explain the principle behind Kernel-wise, we first illustrate the guidelines we aim to adopt in our adapter design and show that existing methods fail to satisfy them (§ 4.1). With these details, we finally propose our method in § 4.2 and § 4.3.
4.1 Guidelines Motivated by Kernel Learning
Given the analogy between attention in PLMs and kernel estimators, we hypothesize that parameter-efficient adaptation should be aware of this connection in transformer-based PLMs and utilize desirable guidelines emerging from the literature on kernel learning. Here we discuss the guidelines introduced in § 1 in further detail.
Guideline-1 suggests that the adaptation should be head-specific. Conceptually, different heads correspond to different empirical kernel matrices (distinct distribution of attention scores), and it will be beneficial to adapt the attention module in a head-specific manner. The effect of head-specific
Note that  $N$  may not always equal  $n$ , such as in cross attention  $(N \neq n)$  or in prefix-tuning  $(N > n$  due to the prefix pretended to the key matrix) (Li and Liang, 2021). In this subsection we omit the superscript  $(h)$  for simplicity since the discussion is limited within a single head
adaptation is also observed by other work, e.g., (He et al., 2021) that mentioned multi-head influence can make methods such as prefix-tuning more expressive.
Guideline-2 is that we should assign more parameter budgets to the coefficient (or value) part of attention compared to the empirical kernel matrix part (query and key). This guideline comes from the classical optimization procedure in kernel learning (Wasserman, 2006, Definition 5.29) where we fix the kernel in use and only perform the unconstrained optimization for the coefficients  $(c_{j}$ 's in Equation (4)). This practice in kernel learning can be justified by Representative Theorem (Scholkopf et al., 2001) that the minimizer  $f^{*}$  of some certain empirical risks admits a representation of the form:
$$
f ^ {*} (\cdot) = \sum_ {j = 1} ^ {N} \alpha_ {j} \kappa (\cdot , k _ {j}),
$$
where  $\alpha_{j}$ 's are the free parameters to optimize. Representer Theorem indicates that the target estimator  $f^{*}(\cdot)$  is simply a linear combination of  $\kappa (\cdot ,k_{j})$ 's, and therefore many kernel methods focus on optimizing the coefficients  $\alpha_{j}$ 's. Revisiting Equation (4), we find  $\alpha_{j} = \frac{c_{j}}{\sum_{j^{\prime} = 1}^{N}\kappa(q_{i},k_{j^{\prime}})}$  under the setting of transformers, which motivates us to apply Guideline-2 to better model the tunable coefficient part  $c_{j}$ 's in attention.
In addition, kernel learning theory concludes that the sample efficiency of a Nadaraya-Watson kernel estimator is mainly influenced by its bandwidth (the scaling factor, corresponding to the factor  $\frac{1}{\sqrt{p}}$  in Equation (2)), rather than the concrete form of the empirical kernel matrix  $\kappa(Q, K)$  (Wasserman, 2006, Section 5.4). This implies that the adaptation to the empirical kernel matrix can be conservative. This is also similar to the conjecture by (He et al., 2021), which mentions "attention learns pairwise positional interactions which do not require large capacity for adapting to new tasks."
Do Existing Adapters Satisfy the Guidelines?
Adapters are designed to modify the hidden states in a certain step in a transformer, and their mechanism can be stated as,
$$
\boldsymbol {H} \leftarrow \boldsymbol {H} + \Delta \boldsymbol {H},
$$
where  $H$  is the "hidden state" in a certain step, and  $\Delta H$  is the update given by the adapter. As
shown in (He et al., 2021), this definition embodies most of the recent proposals for efficient adaptation, such as adapters, prefix-tuning, LoRA, and similar variants.
The original MLP-based adapters, which only adjust the output of a particular layer (Houlsby et al., 2019), do not modify the empirical kernel matrix in the attention sub-layer.
Prefix-tuning satisfies the first guideline as it is head-specific by nature (it prepends trainable continuous prefixes to both key and value matrices in each head). However, it fails to satisfy the second guideline as it enforces an equal assignment of tunable parameters to both the kernel and the coefficient parts (since the prefixes for key matrices and value matrices correspond to each other).
As for weight-updating adapters, such as LoRA (Hu et al., 2021), its proposed setup disobeys both guidelines. The original LoRA updates the whole weight matrix, which is not head-specific. To explain this, consider the weight matrix  $W_{q}$  as an example<sup>6</sup>. In each training step, LoRA updates  $W_{q}$  with a low-rank matrix  $BA$  of the same size as  $W_{q}$ . However, if we denote the  $r$ -by-  $N_{h}p$  matrix  $A$  as  $(A^{(1)}, A^{(2)}, \ldots, A^{(N_{h})})$ ,  $A^{(h)} \in \mathbb{R}^{r \times p}, \forall i \in [N_{h}]$ , the modification to the weight matrix  $W_{q}^{(h)}$  for each head would be  $BA^{(h)}$ 's,  $\forall i \in [N_{h}]$ . We observe that the updates for all the heads share the same column space spanned by  $B$ . In the extreme case of rank-1  $B$  (for example), the updates for each column in the weight matrix will be in the same direction, which is not ideal for adapting all the heads. Further, as suggested by Hu et al. (2021), LoRA evenly assigns the parameter budgets to the weight matrices for  $Q$  and  $V$ , which deviates from the second guideline.
4.2 Kernel-wise Adaptation
We choose LoRA as our primary base to develop our method because of its flexibility in assigning parameters to different weight matrices—both in empirical kernel matrix and coefficient components.
4.2.1 Guideline-1: Head-specific Adaptation
To incorporate Guideline-1, we extend the framework of LoRA and propose Kernel-wise  $^{7}$  that sat
The case of  $W_{v}$  is similar to  $W_{q}$ , while for  $W_{o}$ , the role of  $\mathbf{A}$  and  $\pmb{B}$  would be exchanged. For simplicity, we will always assume we are modifying  $W_{q}$  and discuss how to make  $\pmb{B}$  head-specific throughout the paper.
$^{7}$ Kernel-wise is a concrete scheme to adjust a specific weight matrix. It will have multiple variants modifying different weight matrices.
(b) Kernel-wise
(c) Kernel-wise-lite
Figure 1: (a) Original LoRa with two heads. (b) Kernelwise variant of LoRa that satisfies Guideline-1 (c) lightweight version of (b) that uses  $\pmb{W}_k^{(h)}$  to reduce the trainable parameters in  $\pmb{B}^{(h)}$ . Note that colored matrices are tunable, whereas gray-scale matrices are fixed (Best viewed in color).
isfies the first guideline.
Here, we allow the low-rank weight matrix updates for each head to have customized column spaces, by training distinct  $B^{(h)} \in \mathbb{R}^{N_h p \times r}$  for head-  $h$ ,  $\forall h \in [N_h]$ . In this case, the weight matrix  $W^{(h)}$  in head-  $h$  would be updated by
$$
\boldsymbol {W} ^ {(h)} \leftarrow \boldsymbol {W} ^ {(h)} + \boldsymbol {B} ^ {(h)} \boldsymbol {A} ^ {(h)},
$$
and is expected to be more expressive due to the non-shared column spaces (see Figure 1(b)).
On the downside, this design suffers from inflexibility with a small parameter budget. For all the  $B^{(h)}$ 's, to provide rank- $r$  updates in each head, the new adaptation takes around  $N_{h}^{2}pr$  parameters. However, if, for instance, only  $4N_{h}p$  parameters are assigned to modulate a weight matrix, we can still implement the original LoRA by using rank-2  $A$ ,  $B$ , while the construction of Kernel-wise would be prohibited since even rank-1 updates in each head will require more parameters than the budget.
To resolve the issue, we provide a lightweight alternative to the head-specific adaptation above, which we call Kernel-wise-lite. In this version, we propose to use the frozen  $\mathbf{W}_k^{(h)} \in \mathbb{R}^{N_h p \times p}$  as the head-specific basis for head-  $h$  ( $\mathbf{W}_k^{(h)}$  is the  $h$ -th block in the weight matrix  $\mathbf{W}_k$ , c.f. § 3.1). Therefore any target weight matrix  $\mathbf{W}^{(h)}$  would be updated by
$$
\boldsymbol {W} ^ {(h)} \leftarrow \boldsymbol {W} ^ {(h)} + \boldsymbol {W} _ {k} ^ {(h)} \boldsymbol {B} _ {k} ^ {(h)} \boldsymbol {A} ^ {(h)},
$$
where  $B_{k}^{(h)}$  is a  $p$ -by-  $r$  matrix (see Figure 1(c)).
Utilizing this  $\pmb{W}_{k}^{(h)}$  with a smaller  $B_{k}^{(h)}$  allows the adaptation to be head-specific while containing the same number of trainable parameters as the original LoRA, which is  $2N_{h}pr$ . This comes at the cost of restricting the basis spaces of updates for each head from the unconstrained  $B^{(h)}$  to  $W_{k}^{(h)}B_{k}^{(h)}$ , where  $W_{k}^{(h)}$  is fixed.
But why should we choose  $\mathbf{W}_k$  for the lightweight updates? The motivation comes from results in kernel learning that encourage adapting of the coefficient part using the basis spaces spanned by key matrices. In a kernel estimator, the coefficient  $\mathbf{C}$  is independent of the query sequence  $\mathbf{Q}$  as it is trained solely by the supporting points  $\mathbf{K}$  asymptotically,  $c_j$ , the  $j$ -th row in  $\mathbf{C}$ , is only decided by  $k_j$ ,  $j \in [N]$  (Yang et al., 2017). Concretely, given the loss function and the kernel function,  $c_j$  is in general influenced by  $k_j$  and  $\mathbf{K}_{-j}$  (all the points except  $k_j$ ), while when  $N \to \infty$ ,  $\mathbf{K}_{-j}$  can be fully specified by a fixed distribution. This implies that in attention, compared to queries, keys are more related to values, and following which, we turn to  $\mathbf{W}_k$  to form the basis for the low-rank updates for the conceptual motivation.
Combining LoRA with Kernel-wise. Our proposed adaptation can make fine-grained adjustments to each attention head and thus improve the representations. However, increased representation power might at times trade-off with lower-rank updates. As such, we propose our main variant—Kernel-mix, to combine the original LoRA and Kernel-wise, to attain the best of both worlds—larger basis (therefore higher rank updates) and specific adaptation to each head. Its update expression is as follows,
$$
\boldsymbol {W} ^ {(h)} \leftarrow \boldsymbol {W} ^ {(h)} + (\boldsymbol {B} _ {L o R A}, \boldsymbol {B} ^ {(h)}) \left( \begin{array}{c} \boldsymbol {A} _ {L o R A} ^ {(h)} \\ \boldsymbol {A} ^ {(h)} \end{array} \right),
$$
where  $B_{\text{LoRA}}$  is shared among all the heads, while  $B^{(h)}$ 's are head-specific. We remark Kernel-mix also has a lightweight alternative, Kernel-mixlite, which is the combination of the original LoRA and Kernel-wise-lite. We compare different variants through experiments in  $\S 6$ .
4.2.2 Guideline-2: Making Coefficients more Expressive
To incorporate Guideline-2, we propose to make coefficients more expressive by allowing the modification of both  $W_{v}$  and  $W_{o}$  for the coefficient part, compared to only updating  $W_{v}$  in the original LoRA. We achieve this by re-writing the attention sub-layer under the kernel estimator perspective, which extends the scope of attention by including  $W_{o}$  in its head-specific computation as well. If we represent  $W_{o}$  as,