SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Minyoung  $\mathbf{K}\mathbf{m}^{1}$
$^{1}$ Samsung AI Center Cambridge, UK
mikim21@gmail.com
Timothy Hospedales<sup>1,2</sup>
$^{2}$ University of Edinburgh, UK
t.hospedales@ed.ac.uk
Deep learning practice is increasingly driven by powerful foundation models (FM), pre-trained at scale and then fine-tuned for specific tasks of interest. A key property of this workflow is the efficacy of performing sparse or parameter-efficient fine-tuning, meaning that by updating only a tiny fraction of the whole FM parameters on a downstream task can lead to surprisingly good performance, often even superior to a full model update. However, it is not clear what is the optimal and principled way to select which parameters to update. Although a growing number of sparse fine-tuning ideas have been proposed, they are mostly not satisfactory, relying on hand-crafted heuristics or heavy approximation. In this paper we propose a novel Bayesian sparse fine-tuning algorithm: we place a (sparse) Laplace prior for each parameter of the FM, with the mean equal to the initial value and the scale parameter having a hyper-prior that encourages small scale. Roughly speaking, the posterior means of the scale parameters indicate how important it is to update the corresponding parameter away from its initial value when solving the downstream task. Given the sparse prior, most scale parameters are small a posteriori, and the few large-valued scale parameters identify those FM parameters that crucially need to be updated away from their initial values. Based on this, we can threshold the scale parameters to decide which parameters to update or freeze, leading to a principled sparse fine-tuning strategy. To efficiently infer the posterior distribution of the scale parameters, we adopt the Langevin MCMC sampler, requiring only two times the complexity of the vanilla SGD. Tested on popular NLP benchmarks as well as the VTAB vision tasks, our approach shows significant improvement over the state-of-the-arts (e.g.,  $1\%$  point higher than the best SOTA when fine-tuning RoBERTa for GLUE and SuperGLUE benchmarks).
Practical deep learning increasingly relies on a workflow of fine-tuning pre-trained foundation models for the specific task of interest to be solved. The foundation models (FM)s are often pre-trained at scale using large architectures and self-supervised objectives on large scale unlabelled data, allowing them to encode a great deal of general prior knowledge. Fine-tuning then specialises these general purpose FMs to the specific task of interest. The key challenge then becomes how to conduct fine-tuning so as to balance the conflicting goals of adapting to the task at hand; while retaining the knowledge obtained during foundational pre-training and avoiding overfitting on the (typically smaller) downstream training set. This is addressed from a variety of perspectives from careful learning rate scheduling to regularised fine-tuning [26, 12]. However, the mainstream approach in practice is largely based on some instantiation of sparse fine-tuning.
Sparse fine-tuning, also known as parameter efficient fine-tuning (PEFT) [23, 43, 14, 24, 22, 15, 16, 18, 45, 13, 10], manages the adaptation/forgetting trade-off by selectively updating a specific subset of parameters in the FM while keeping others frozen. These approaches also benefit from good parameter scalability with respect to the number of downstream tasks to solve since each downstream
task only requires storing a small number of parameters compared to the full FM. A large number of sparse/PEFT approaches have now been proposed, based on either additive or mask-based principles. Additive approaches [16, 15] inject a hand-chosen set of additional small modules to be learned, while keeping the original FM entirely frozen; while mask based approaches [43, 18] manually select a sparse subset of original FM parameters to consider as learnable – such as the prompt parameters in transformers [18]. Despite substantial engineering effort, prior sparse fine-tuning work mainly relies on human intuition and ad-hoc heuristics to specify additive modules and sparse fine-tuning masks. This raises the question how can we automatically select which parameters to insert or update in a principled and effective way?
In this paper, we address this issue by proposing an automated sparse fine-tuning method. Our BayesTune framework automatically discovers the parameters to update, with minimal human intervention. In contrast to some early concurrent attempts that rely on heuristics and/or heavy approximations [45, 13, 10], BayesTune is highly principled, underpinned theoretically by Bayesian marginalization of parameter uncertainty. Specifically, we place a sparse Laplace prior on the FM parameters, and the posterior inference of the sparse scale parameters reveals which parameters are most important to update from their initial values. BayesTune is efficiently implemented with Langevin MCMC, which only requires twice the cost of vanilla SGD. We demonstrate BayesTune with both RoBERTa transformer on (Super)GLUE benchmark suite and Vision Transformer on VTAB benchmark suite. BayesTune can achieve a specified sparsity level, and provides a reasonable way to identify a good sparsity/performance trade-off if a sparsity constraint is not pre-defined. Finally, we show that BayesTune can work with both mask-based search spaces to identify existing FM parameters to update by fine-tuning, as well as additive search spaces to identify which new modules to add where for additive adaptation.
2 Problem Setup and Motivation
First we introduce some key notations as follows:
$$
\begin{array}{l} \theta = (\text {V a r i a b l e}) \text {P a r a m e t e r s} (\text {i . e .}, \text {w e i g h t s a n d b i a s e s}) \text {o f t h e F M} \\ \bar {\theta} = (\text {C o n s t a n t}) \text {P a r a m e t e r s (i . e . , w e i g h t s a n d b i a s e s)} \text {o f t h e p r e - t r a i n e d F M} \\ D = \text {T r a i n i n g} \\ \end{array}
$$
We denote by  $\theta_{i}$ $(\in \mathbb{R})$  the  $i$ -th element of  $\theta$  by considering  $\theta$  as a whole vector of concatenated parameters, i.e.,  $\theta = [\theta_1,\dots ,\theta_d]^\top$  where  $d$  is the number of parameters (so,  $\theta_{i}$  is scalar). Before we start, we clarify the goal of the sparse fine-tuning formally as follows.
Desiderata. We aim to find  $\theta$  which performs well on the downstream task, and at the same time remains close to  $\overline{\theta}$  sparsely in  $L_{0}$  sense (i.e.,  $||\theta -\overline{\theta} ||_{0} = \sum_{i = 1}^{d}\mathbb{I}(\theta_{i}\neq \overline{\theta}_{i})$  is small). Roughly speaking,  $\theta_{i} = \overline{\theta}_{i}$  for many  $i$ 's. Oftentimes we define and set the sparsity level  $p\in [0,1]$ , meaning that the proportion of the FM parameters updatable from the pre-trained values is no greater than  $(100\times p)\%$  (i.e.,  $||\theta -\overline{\theta} ||_0\leq \lfloor d\cdot p\rfloor$ ).
About the sparsity level  $p$ . We remark here that the sparsity level  $p$  is often given by the users, for instance, in case we have a strict pre-specified memory constraint (e.g., embedded platforms), where the allowable extra space for saving the updated new parameters from  $\overline{\theta}$  for the downstream task is only  $(100 \times p)\%$  of the original FM storage. In other situations, it would be ideal if the sparse fine-tuning algorithm can estimate optimal  $p$  that trades off the downstream accuracy against memory overhead (e.g., find the smallest  $p$  that enables a certain level of accuracy guarantee). Unfortunately most existing sparse fine-tuning methods are not capable of this feature, and rather resort to time-consuming cross-validation-type search. Although this latter case opens up a new interesting research problem to pursue further, our proposed approach offers reasonable and principled criteria for estimating the optimal  $p$  (See Sec. 3 for details).
How about layer-wise selection? We also note that our approach is not tied to the layer-wise parameter treatment, namely selecting a few layers to be updated with the rest layers frozen, which is a popular strategy adopted in many sparse fine-tuning methods. The layer-wise treatment is based on the conjecture that whether the parameters important to the downstream task are deviated from the pre-trained ones or not, is determined in a layer-wise all-or-nothing manner. Although this is a reasonable assumption, our approach does not rely on the assumption, and is able to discover optimal sparse updates automatically by inspecting individual parameters over all layers.
Back to our notations, the downstream task must be associated with the relevant loss function  $l(\theta; z)$  with  $z \in D$  for fine-tuning. Here the task could be either supervised (i.e.,  $z = (x, y)$  for a pair of input  $x$  and its target label  $y$ ) or unsupervised ( $z = x$  with input  $x$  alone). In our model the loss function is turned into a likelihood model by the conventional trick,  $p(z|\theta) \propto \exp(-l(\theta; z))$ . Although our approach can deal with both supervised and unsupervised cases seamlessly, we predominantly focus on the supervised classification case, that is, we define the data likelihood  $p(y|x, \theta) \propto \exp(-l(\theta; x, y))$  for  $(x, y) \in D$ , where  $l(\theta; x, y)$  is typically the cross-entropy loss.
3 (Proposed) Bayesian Sparse Fine-tuning Model
We propose a (hierarchical) Bayesian model to tackle the downstream prediction task of interest. Specifically, we treat the FM parameters  $\theta$  as random variables and the downstream training data  $D$  as evidence. To encourage parameter sparsity, namely most parameters  $\theta_{i}$  remain at the pre-trained  $\overline{\theta}_i$ , we impose the Laplace prior  $p(\theta|\lambda)$ ,
$$
p (\theta | \lambda) = \prod_ {i = 1} ^ {d} p \left(\theta_ {i} \mid \lambda_ {i}\right) = \prod_ {i = 1} ^ {d} \mathcal {L} \left(\theta_ {i}; \bar {\theta} _ {i}, \lambda_ {i}\right) = \prod_ {i = 1} ^ {d} \frac {1}{2 \lambda_ {i}} \exp \left(- \left| \theta_ {i} - \bar {\theta} _ {i} \right| / \lambda_ {i}\right) \tag {1}
$$
where  $\mathcal{L}(x;\mu ,b) = \frac{1}{2b}\exp (-|x - \mu | / b)$  is the (univariate) Laplace distribution with mean  $\mu$  and scale  $b$ . Thus we fix the prior means as the pre-trained values  $\overline{\theta}$ , and only model the scale parameters which are newly introduced random variates denoted by  $\lambda (>0)$ . Note that every single parameter  $\theta_{i}$  (scalar) is associated with its own prior scale  $\lambda_{i}$  (thus  $d$  scale variables in total). Another thing to note is that the final layer (also known as the readout head) of the FM is often completely replaced by a random one for the downstream task. Since this final layer needs to be learned from the scratch with the downstream data, we do not place the Laplace prior on the parameters of the readout head.
Through the scale variables  $\lambda$  of the Laplace distributions, we can express our (prior) preference to the degree of parameter deviation from the pre-trained values: small  $\lambda_{i}$  leads to a peaky Laplace around its mean, penalizing even small deviation of  $\theta_{i}$  from  $\overline{\theta}_i$ ; on the other hand, large  $\lambda_{i}$  makes it flat, being less sensitive to deviation, allowing  $\theta_{i}$  to take values freely away from  $\overline{\theta}_i$ . By carefully choosing  $\lambda_{i}$  values, we can balance effectively between overfitting (too much deviation) and underfitting (too little deviation). Instead of choosing  $\lambda_{i}$ 's manually, our idea is to learn them automatically from data, through the principled Bayesian inference: given the evidence (i.e., the downstream data  $D$ ) we infer the most probable values of  $\lambda_{i}$ 's that best explain the evidence, while being small enough to lead to sparse updates overall.
To this end, we regard  $\lambda$  as random variables and impose hyper-prior on  $\lambda$  (hierarchical Bayes). As we prefer sparse updates in the end, the hyper-prior needs to put higher mass/density on small  $\lambda$  values. We adopt the Gamma distribution for this purpose, more specifically,
$$
p (\lambda | \alpha , \beta) = \prod_ {i = 1} ^ {d} p \left(\lambda_ {i} | \alpha , \beta\right) = \prod_ {i = 1} ^ {d} \mathcal {G} \left(\lambda_ {i}; \alpha , \beta\right) \tag {2}
$$
where  $\mathcal{G}(x;\alpha ,\beta)\propto x^{\alpha -1}\exp (-\beta x)$  for  $\alpha ,\beta >0$  . Note that although we can introduce individual  $\lambda_{i}$  -specific parameters (i.e.,  $(\alpha_{i},\beta_{i}))$  , we instead use a single  $(\alpha ,\beta)$  shared over all  $\lambda_{i}$  's for simplicity. Since the mode of Gamma is 0 if  $\alpha < 1$  and the variance is  $\alpha /\beta^2$  , we use small  $\alpha < 1$  and large  $\beta$  to enforce a priori small  $\lambda$  values. In our experiments we choose  $^\dagger$ $(\alpha = 0.01,\beta = 100)$  for both the NLP and VTAB vision benchmarks.
Our full model can be written as the following joint distribution:
$$
p (D, \theta , \lambda | \alpha , \beta) = \overbrace {p (\lambda | \alpha , \beta)} ^ {= \prod_ {i} \mathcal {G} (\lambda_ {i}; \alpha , \beta)} \times \underbrace {p (\theta | \lambda)} _ {= \prod_ {i} \mathcal {L} (\theta_ {i} ; \bar {\theta} _ {i} , \lambda_ {i})} \times \overbrace {p (D | \theta)} ^ {= \prod_ {(x, y) \in D} p (y | x, \theta)} \tag {3}
$$
With this model, our ultimate goal is to infer the posterior distribution of the network weights and scales  $p(\theta, \lambda|D)$ , namely
$$
p (\theta , \lambda | D, \alpha , \beta) = \frac {p (D , \theta , \lambda | \alpha , \beta)}{\iint p (D , \theta , \lambda | \alpha , \beta) d \theta d \lambda} \tag {4}
$$
<sup>1</sup>We also tested with models with further hierarchy by placing priors on  $\alpha$  and  $\beta$ , however, there was no significant advantage over the manually chosen ones.
We adopt the stochastic-gradient MCMC approach [40, 7, 4] to obtain samples from the posterior, especially the Langevin dynamic method (SGLD) [40], which amounts to running the following recurrence to collect posterior samples (after some burn-in steps):
$$
[ \theta , \lambda ] \leftarrow [ \theta , \lambda ] + \frac {\eta}{2} \nabla \left(\log p (\lambda | \alpha , \beta) + \log p (\theta | \lambda) + \frac {\hat {N}}{| B |} \log p (B | \theta)\right) + \gamma \epsilon \sqrt {\eta} \tag {5}
$$
where  $B(\subset D)$  is a minibatch,  $\eta$  is small step size, and  $\epsilon \sim \mathcal{N}(0,I)$ . In standard SGLD,  $\hat{N} = |D|$  (training data size) and  $\gamma = 1$ . However, we allow them to take values different from these default ones:  $\hat{N}$  is the effective training data size to account for data augmentation $^2$  (e.g., if there are 5 different augmentation strategies and each allows 10 different variations, the effective data size can be  $\hat{N} = 10^{5} \times |D|$ ); the other hyperparameter  $\gamma$  can be used to discount the noise effect $^3$ , which is validated in the range  $10^{-4:0}$  in our experiments. See Sec. 5.1 for more details. Note that in the parentheses subject to derivatives, the first two terms admit closed-form gradients while the gradient of the last term can be computed by the conventional SGD backprop. Thus each step in (5) is efficient, requiring at most only two times the complexity of the vanilla SGD step.
After a burn-in period, we can maintain those samples  $(\theta ,\lambda)$  to approximate  $p(\theta ,\lambda |D)$ . For instance, the running average of the  $\lambda$  samples, denoted by  $\hat{\lambda}$ , is a good estimate of the mean of the marginal posterior  $p(\lambda |D)$ . Since we imposed the hyper-prior that encourages small  $\lambda$ , ideally the majority of the scale parameters tend to remain small a posteriori, whereas there will be only a few large-valued  $\hat{\lambda}_i$ 's corresponding to those FM parameters that crucially need to be updated away from their pre-trained values to explain the data  $D$ . Roughly speaking,  $\hat{\lambda}_i$ s indicates how important it is for the corresponding FM parameters  $\theta_{i}$ s to be updated away from their pre-trained values  $\bar{\theta}_i$  on the downstream data. Importantly, this  $\hat{\lambda}$  can be used in the next stage (described in the next paragraph) to decide which network weights  $\theta_{i}$  need to be frozen (those with small  $\hat{\lambda}_i$ s) and which need to be updated from the pre-trained (those with large  $\hat{\lambda}_i$ s).
Thresholding scales and the second stage. Although we adopted the sparsity-inducing Laplace prior, it is not necessarily the case that the majority of posterior  $\theta$  values are sharply staying at the pre-trained values  $\bar{\theta}$ . We may need some thresholding: If  $p$  is given, then we can directly take the top  $(100 \times p)\%$  of them. Otherwise we examine for cut-off point - Those  $i$ 's with small posterior means  $\hat{\lambda}_i$  are considered as frozen weights (i.e.,  $\theta_i = \overline{\theta}_i$ ), while those  $i$ 's with large  $\hat{\lambda}_i$ s can be treated as updatable parameters. In practice, we sort and plot the  $\hat{\lambda}_i$ s to eyeball and find a reasonable cut-off point. One can also use this  $\hat{\lambda}$ -plot to decide a reasonable sparsity level  $p$  for good accuracy-memory trade-off (See Fig. 1 for illustration). Once we have decided which parameters are updated and frozen, we can run vanilla SGD to train the updatable parameters of the FM, which forms our second stage.
Our overall algorithmmm dubbed BayesTune, is summarized as pseudocodes in Alg. 1.
4.1 Comparison to Existing Sparse Fine-tuning Methods
The recent sparse/PEFT fine-tuning approaches broadly fall into the following three categories: 1) Heuristic search criteria, 2) Attaching small extra modules to the FM (e.g., adapter-based [15]), and 3) Directly formulating the sparse fine-tuning problem as an optimization problem.
1) Heuristic search criteria.
- Random selection - Randomly select  $(100 \times p)\%$  parameters to update.
Similar discussions of how to quantify dataset size when using Data Augmentation as well as data size inflation, were made in several prior Bayesian deep learning studies, for instance, [33, 34, 17].
3The rationale for noise discounting is as follows. Purely finding the posterior mean (i.e., SGLD without noise discount) risks performing poorly if the posterior is truly multi-modal, because it may converge to a low probability parameter. Also, purely searching for the posterior mode (i.e., MAP instead of SGLD) may be sensitive to data noise, because no stochasticity is taken into account properly. So for combining principle and practice it is reasonable to prefer a discounted noise procedure that balances between identifying a particular mode, but gets a mean estimate in the vicinity of that mode.
Figure 1: (Left) The  $\hat{\lambda}$ -plot in our BayesTune after stage-1 training. This is with EfficientNet-B0 on the Oxford-Pets dataset. After sorting the posterior mean  $\hat{\lambda}_i$  values, we plot several sparsity levels  $p$  and the corresponding  $\hat{\lambda}$  values. To decide a reasonable sparsity level  $p$  for good accuracy-memory trade-off, we illustrate two candidates: Cand-1 with more sparse updates ( $p = 0.05$ ) but many parameters with large  $\hat{\lambda}_i$ 's (e.g.,  $0.1 \sim 0.3$ ) are forced frozen; Cand-2 with less sparse updates ( $p = 0.1$ ) but ensures that those fixed parameters have indeed small  $\hat{\lambda}_i$ 's (e.g.,  $< 0.1$ ). (Right) Test errors in stage-2 at different cutoff points in our BayesTune (red curve). Those two candidate points (Cand-1 and Cand-2) are highlighted and superimposed, consistent with the left figure. The more sparse Cand-1 has slightly higher error than the full-update, while less sparse Cand-2 achieves the best test performance among all other choices. It exhibits nice accuracy-memory trade-off. Further increasing  $p$  (being far more dense) incurs overfitting.
Algorithm 1 BayesTune (two-stage) sparse fine-tuning algorithm.
Input: Pre-trained model  $\overline{\theta}$  , downstream data  $D$    
Initialise:  $\theta = \overline{\theta}$ $\lambda_{i} = 0.0001$  for all i.   
I. Stage-1 (Select  $(100\times p)\%$  parameters to update) For  $t = 0,1,2,\dots,T_{\mathrm{max}}$  . 0. Sample a minibatch  $B = \{(x,y)\}$  from  $D$  1. Run the SGLD step (5) to get a sample  $(\theta ,\lambda)$  2. If  $t >$  burn-in, accumulate  $\hat{\lambda}\gets (1 - \zeta)\hat{\lambda} +\zeta \lambda$  where  $\zeta = 1 / (\#$  accumulated-samples+1). End-For (If  $p$  not given) Decide the cut-off point  $\lambda_{th}$  and the corresponding sparsity level  $p$  using the  $\hat{\lambda}$  -plot. Select those  $i$  's with  $\hat{\lambda}_i > \lambda_{th}$  as updatable parameters.   
II. Stage-2 (Update the selected parameters) Do vanilla SGD with data  $D$  while updating only those updatable parameters from Stage-1.
- MixOut [23] - After training the FM model with full update, a random  $(100 \times (1 - p))\%$  parameters are reset to the pre-trained values.
- BitFit [43] - Update only the bias parameters with the rest frozen.
- MagPruning [14, 24, 22] - Take the largest  $(100 \times p)\%$  in magnitude as updatable parameters.
Compared to our proposed "BayesTune", these heuristic search methods are less principled, unable to explain why the resulting models should perform well on the downstream tasks
2) Approaches that attach small extra modules to FM. They do not search for a subset of FM parameters to update, but attach some small extra modules to the FM (e.g., insert a small MLP module after each layer). All of the FM parameters are frozen, and they only train those attached modules. So it can be seen as sparse fine-tuning considering that the size of the new modules is relatively small compared to that of the FM.
- Adapter [15] - A small bottleneck MLP is inserted at the end of each layer of the FM.
- LoRA [16] - A low-rank projection weight matrix is inserted before the query/key/value embeddings in each layer.
- Prompt learning (VPT) [18] - A set of learnable tokens is inserted in the input of a Transformer attention block in each layer.
- Neural Prompt Search [45] - A NAS-based approach (e.g., learning a supernet followed by evolutionary search for a downstream tasks specific subnet) that aims to search for the optimal combination of the above three methods.
Compared to our "BayesTune", these adapter-based models involve the difficult problem of where to place the extra modules and what would be the optimal module dimensions/sizes, which is essentially a very difficult combinatorial search problem, often requiring heuristic search methods like the evolutionary search.
3) Directly optimizing the sparse optimization problem. One can formulate the sparse fine-tuning as an optimization problem, for instance,
$$
\min  _ {\theta} \mathbb {E} _ {(x, y) \sim D} [ l (\theta ; x, y) ] \text {s . t .} | | \theta - \bar {\theta} | | _ {0} \leq \left\lfloor d \cdot p \right\rfloor . \tag {6}
$$
Unfortunately this problem becomes a discrete optimization due to the  $L_0$  cardinality constraints, thus difficult to solve in general. Some recent approaches aimed to relax the problem to continuous optimization, but relying on heavy approximation.
- DiffPruning [13] - They relax the discrete problem into a continuous one by introducing Bernoulli variables indicating update or freeze, while the reparametrized sampling is further approximated by the Gumbel-softmax treatment.
- ChildPruning [41, 31] - They iteratively train the full model parameters and calculate the projected mask to find the child network.
- Second-order approximation method (SAM) [10] - They adopt several steps of second-order Taylor approximation of the loss function, leading to a very succinct parameter selection strategy: after computing the gradients of  $\theta$  on the data  $D$ , take the  $(100 \times p)\%$  parameters with the largest gradient magnitudes.
Compared to our "BayesTune", these direct optimization strategies rely on strong assumptions and heavy approximation to relax the difficult discrete optimization.