One Network, Many Masks:
Towards More Parameter-Efficient Transfer Learning

Recently, as the pre-trained language models get larger and larger, some parameter-efficient transfer learning methods that only update a few extra parameters while freezing the PLM backbone have been proposed. Adapter-tuning Houlsby et al. (2019) fine-tuned adapter modules inserted after each attention and feed-forward layer. Prefix-tuning Li and Liang (2021) placed an additional trainable prefix to the keys and values matrix in the attention module. LoRA Hu et al. (2022) injected tunable rank decomposition matrices into each Transformer layer. Based on these parameter-efficient transfer learning methods, He et al. (2022a) gave a unified framework that allows for the transfer of design elements across various PETL approaches. However, when applied to larger PLMs and a broader range of tasks, these methods still require a large storage space because the number of extra parameters is directly proportional to the number of layers and tasks. Inspired by the parameter-sharing techniques of ALBERT Lan et al. (2020a), we propose sharing the additional parameters in PETL modules across layers and tasks. Our method can thus obtain higher parameter-efficient efficiency with a significantly smaller portion of additional storage than existing PETL methods.

Pruning is one of the most popular methods to reduce unnecessary weights from over-parameterized neural networks while maintaining comparable performance. Recently, Frankle and Carbin (2019) proposed Lottery Ticket Hypothesis and stated that in a randomly initialized dense model, a sparse sub-network exists that, when trained in isolation, can achieve performance comparable to dense models. Accompanied by this hypothesis, many pruning-before-training methods have emerged Lee et al. (2019); Bai et al. (2022); Sreenivasan et al. (2022). Xu et al. (2021) further propose a method that prunes the backward gradient of the neural network, as opposed to pruning the network parameters themselves. Based on these methods, some works He et al. (2022b); Ma et al. (2022) also proposed to combine pruning algorithms with parameter-efficient methods to further decrease the additional module size. However, those methods only focus on discarding redundant parameters. A parameter is either discarded or retained without any sharing. They fail to make full use of the additional parameters and cannot achieve highly sparse sub-networks without significantly compromising accuracy.

Multi-task learning Zhang and Yang (2022), which involves training a single model to perform well on multiple tasks, has gained popularity as a research direction in machine learning. However, this approach can be hindered by catastrophic forgetting Kirkpatrick et al. (2016) and data imbalance among tasks, which will result in overfitting on low-resource tasks and underfitting on high-resource tasks Arivazhagan et al. (2019). Houlsby et al. (2019) propose adapter tuning that only introduces and updates small additional parameters for each task while freezing the pre-trained model. Based on such a parameter-efficient method, Mahabadi et al. (2021) train a hyper-network named Hyperformer, which generates task-specific weights for the adapter modules when fed with different task embeddings.

In parameter-efficient transfer learning, we freeze the parameters $\theta_{lm}$ of the pre-trained language model and then introduce additional fine-tunable parameters denoted as $\theta_{t}$. Given dataset $\{X_{i},Y_{i}\}^{N}_{i=1}$, the goal of parameter-efficient fine-tuning is to maximize the following likelihood of the label $Y$ by only updating the additional parameters $\theta_{t}$:

Such parameter-efficient methods suggest a more effective way to adapt pre-trained language models over downstream tasks than fully fine-tuning. We give a brief introduction of the three most widely used PETL modules, namely adapter Houlsby et al. (2019), prefix-tuning Li and Liang (2021), and LoRA  Hu et al. (2022) in Appendix A. However, there is still a storage limitation when we handle a large range of tasks using these methods. In this paper, we investigate the potential for further enhancing parameter efficiency in neural network models by reducing storage requirements. While previous PETL methods have primarily focused on decreasing the number of parameters to improve efficiency, our approach posits that employing varying bit lengths (e.g., 1-bit, 8-bit, 32-bit) during storage can lead to significant improvements in parameter efficiency by reducing the overall number of bits used by the parameters. To this end, we use bits to measure the storage, which we call Bit-Level Storage (BLS), to take into account the fact that different parameters may have different bit lengths. Consider a neural model, where each parameter has a specific bit length. Then we divide these parameters into $N$ distinct groups based on their respective bit lengths. Let $\{\rho_{i}\}_{i=1}^{N}$ denote the number of parameters within the group $i$, with corresponding bit lengths $\{b_{i}\}_{i=1}^{N}$. The BLS for these parameters can subsequently be determined as follows:

Parameter-efficient methods like adapter and prefix-tuning tend to introduce an additional module in each Transformer layer. Assuming the Transformer has $L$ layers, we can split the parameters $\theta_{t}$ into $[\theta_{t,1},...,\theta_{t,L}]$ according to their layer indexes. Therefore, we can rewrite Equation 2 as:

Inspired by ALBERT Lan et al. (2020b), in our methods, we first introduce additional parameters for a single PETL module as our prototype network, denoted as $\theta_{pro}$. Then, we share the prototype network across different layers. Assuming that the number of parameters in a single PETL module is $n$, we can decrease total parameters from $nL$ to only $L$, which significantly improves the parameter efficiency. Therefore, we convert the objective function to a more concise one:

Sharing the parameters alone will reduce the model’s capacity to capture meaningful representation in different layers, leading to suboptimal results. Inspired by Zhou et al. (2019) and Ramanujan et al. (2020), we believe that parameters and network structures are both crucial contributing factors to the model’s representative capacity. To this end, we introduce different binary masks $m_{l}=\{0,1\}^{n}$ in each Transformer layer $l$ (left part of Figure 2), where $n$ denotes the number of parameters in a single PETL module. Each mask represents a corresponding subnetwork of the shared prototype network. Even though the prototype is shared among all layers, we can use different masks to create different sub-networks for each layer $l$, whose parameter will be $\theta_{sub,l}=\theta_{pro}\odot m_{l}$, where $\odot$ indicates the element-wise product. With this, we can get our final objective function as:

where $\theta_{sub}$ = [$\theta_{pro}\odot m_{1}$, $\theta_{pro}\odot m_{2}$, …, $\theta_{pro}\odot m_{L}$].

To learn such masks, we develop our training algorithm based on the edge-popup approach Ramanujan et al. (2020). Specifically, for each binary mask $m_{l}$, we introduce floating-point scores $s_{l}\in\mathbb{R}^{n}$. In the forward pass, we generate the binary mask $m_{l}$ by setting the top-$k$% with the largest absolute value in $s_{l}$ as 1 and the rest as 0. We denote such top-$k$% thresholding function as $h$, and $m_{l}=h(s_{l})$. We refer to the hyperparameter $k\%$ as the sparsity ratio in subsequent sections of this paper. During the backpropagation, we use the straight-through gradient estimator Bengio et al. (2013) to approximate the gradient of the scores, where function $h(\cdot)$ is treated as the identity function. In addition to training the masks, we also jointly optimize the prototype network.

Our approach employs learnable floating-point scores for each binary mask during fine-tuning, leading to $nL$ parameters. When integrated with the prototype network, ProPetl updates a total of $nL+n$ parameters during the training phase, which is marginally more than the conventional PEFT methods which have $nL$ extra parameters. A comprehensive overview of the ProPetl algorithm can be found in Algorithm 1. After training, we discard the floating-point scores and retain only the binary masks (1-bit) together with the shared prototype network (32-bit). Assuming that the 32-bit prototype network requires $p$ bit-level storage and the binary mask of the same dimension demands $p/32$, our ProPetl can achieve a substantial decrease in storage from around $pL$ to $p(1+L/32)$. To reconstruct the network structure during inference, we adhere to the following steps: (1) first load the unpruned, shared 32-bit prototype network, (2) then load the binary masks (1-bit) for each layer/task, and (3) extract and use the pruned subnets from the shared prototype network based on specific binary masks during the inference step.

Rather than just sharing a PETL module across layers under single-task learning, we can also allow for efficient sharing of the prototype network across multiple tasks. In our approach, we leverage layer masks, as introduced in the previous section, to support parameter sharing within the model. Additionally, we introduce task masks to support parameter sharing across multiple tasks. By performing a logical OR operation²²2We provide ablations regarding the choice of mask combining methods in Appendix D.1. on these masks, we can obtain a hybrid mask for a specific layer in a specific task, as shown in the right side of Figure 2.

With the design of the hybrid mask, given $T$ tasks and $N$ layers in the pre-trained language models, we only require one ProPetl module, $N$ layer masks, and $T$ task masks, further reducing the BLS from the single task scenario (e.g., 0.011% BLS as in Table 2). In addition, layer masks and task masks, which can be combined as hybrid masks, will potentially help infuse the knowledge into the shared prototype network from layers and tasks.

We present the results of multi-task learning and also provide baseline results from single-task learning using the T5_BASE in Table 2. Our best-performing model, ProPetl_Adapter with bottleneck dimension of 64, surpasses the fully fine-tuned T5. We also compare ProPetl with Hyperformer++  Mahabadi et al. (2021), a hyper-network that is specifically designed to transfer knowledge across tasks. The latter uses significantly more task-specific bit-level storage (26$\times$: $0.011\%$ v.s. $0.29\%$ per task), while only increasing the average score by $0.51$. Compared to the vanilla adapter, ProPetl_Adapter can be marginally better under a similar fine-tuning parameter budget but with 1/9 of the original storage. Besides, we experiment with an extreme case, where we set the bottleneck dimension to 6. Our results show that the accuracy of adapter tuning decreases from 85.48 to 83.94, while ProPetl_Adapter still maintains a comparable performance to the fully fine-tuned model (85.32 v.s. 85.47) with a remarkably small percentage ($0.001\%$) of bit-level storage per task. This demonstrates that normal adapter tuning can not make full use of the parameters and may fail to perform well with a relatively small bottleneck dimension. However, in the case of ProPetl_Adapter, even with a bottleneck dimension that is only 6, it can still achieve a reasonable result. To further validate that ProPetl is effective for larger-sized models, we also carry out experiments on the T5 3B variant and present the findings in Table 3. The outcomes align with our conclusions drawn from the T5 base model. We believe that, in ProPetl_Adapter, the structure information learned by the mask can make up for the performance drop due to the shrink of the bottleneck dimension to a certain extent. We further compare adapter tuning with ProPetl_Adapter with different percentages of bit-level storage by varying the bottleneck dimension.

The results are presented in Figure 3 (c). It shows ProPetl_Adapter is able to reach the fine-tuned performance with as few as $0.002\%$ task-specific bit-level storage. We can also see that the curve shares a similar trend with those in Figure 3 (a) and (b), which we will further discuss in the next section.