LiFT: Learning to Fine-Tune via Bayesian Parameter Efficient Meta Fine-Tuning

1. More meta-learning baselines.

Most methods the reviewer suggested assume small class cardinality cases (eg, learning the class prototypes or linear readout heads), hence they are not straightforwardly applicable to open-ended NLP (Crossfit) benchmarks. On the VTAB benchmark, however, we have been able to implement ProtoNet and more recent memory-efficient version ProtoNet-LITE. The latter (LITE) essentially aims to represent the prototype vectors using the entire support data (instead of a minibatch), but to remedy computational overhead they set a randomly chosen large portion of the support data as non-backpropagatable.

[ProtoNet-LITE] Memory Efficient Meta-Learning with Large Images, NeurIPS 2021

The results are shown in the following table (the full scores for individual tasks are shown in Table 2 of the revised paper).

2. The results do not include standard deviations. Are the improvements with LiFT statistically significant?

Due to the computational overhead, it was time and compute-expensive to run multiple random seed experiments, and we only managed to run a single experiment. However, we are able to perform multiple 5 random seed runs on the Crossfit-CLS45 split to collect standard deviations and $p$-values. The results are shown in the following table ($p$-values of the existing methods against our LiFT-K=5 shown in the parentheses):

As shown, our LiFT model outperforms the existing methods statistically significantly.

We have added this result in the revised paper (Appendix E.8 and Table 14).

3. Some limitations, such as the joint training on the training tasks, which can be infeasible for large-scale datasets.

We did mention this limitation in Sec. 4.5 Utilizing Task-wise Pre-fine-tuned Models. In this section, we also introduced a heuristic remedy of clustering pre-fine-tuned models. Table 5 shows the evaluation of this remedy, and the results are promising.

4. Which architectures and hyperparameters are used?

We used the same architectures and fair hyperparameters for MAML baselines (Appendix C). We agree that MAML is not scalable to large feature extractor networks. For fair comparison and for efficiency, we apply MAML in the same way that we apply LIFT. IE: We train MAML-of-LORAs, rather than MAML of the raw VIT-B/16 backbone, which would be intractable and unstable.

5. Multiple inner steps for MAML.

We have run the experiments to see the effect of the multiple inner loop updates in MAML. On the Cross-fit CLS-45 benchmark, increasing the number of inner iterations can improve the test performance at the cost of memory, but MAML still underperforms our LiFT by large margin.

More than 5 iters in MAML incurred OOM issues.

We have added this result in the revised paper (Appendix E.7 and Table 13).

1. Missing critical related works (8 papers [1-8]).

In response to the reviewer’s suggestion, we have made an empirical comparison with the recent method LITE [6]. The rest of the models suggested are all tightly coupled with specific backbone network architectures, eg, CNAPs [2], and/or assume closed set classification problems e.g. [7]. As such they are not straightforwardly applicable to our main task of LoRA PEFT framework for adapting LLMs in open-ended text generation. But as we see that they are important latest meta learning algorithms, we have cited them in our revised paper (Please see Extended related work section in Appendix D, the second paragraph).

The LITE [6] (memory efficient meta learning) is generally not tied to a specific network architecture, and as suggested therein, the ProtoNet meta learner can be applied. To this end, we ran ProtoNet-LITE on the VTAB benchmark since ProtoNet typically assumes small-way classification problems, not adequate for LLM’s large vocabulary output cardinality. LITE [6] essentially aims to represent the prototype vectors using the entire support data (instead of a minibatch), but to remedy computational overhead, they set a randomly chosen large portion of the support data as non-backpropagatable.

The results are shown in the following table (the full scores for individual tasks are shown in Table 2 of the revised paper).

2. Need to consider more recent baselines, and more effective meta-learning baselines (3 papers [1,2,3]).

[3] aims to learn a task-specific set of sparse masks for the new parameters added to the pre-trained ones while the new parameters are shared across the tasks (task-agnostic). Since they used ViT full backbone, and we are using the LoRA PEFTs, it is difficult to compare the two approaches directly. Only ViT-small was tested in [3] (sparse interpolation experts), and we also have ViT-B experiments.

Instead we have implemented and run [2] bootstrapped-maml in our crossfit experimental framework with LoRA PEFT as parameters to learn. In the bootstrapped maml, the target model is first obtained by applying SGD updates on the train data then batch data sequentially. The meta learning loss is then the KL divergence between the predictive distributions of the train-data-updated model and the target model where the latter is treated as constant (no backprop).

Comparison with “bootstrapped-maml” [2] on Crossfit CLS-45:

We have added this result in the revised paper (Appendix E.6, Table 12).

3. Scalability to large-scale LLM (e.g., more than 1B param).

Our current results are already applied on VIT-B, as mentioned above, which is already larger scale than many meta-learning papers’ experiments. Nevertheless, we agree that it is interesting to explore scaling further.

We now have run our algorithm on an even larger LLM, FLAN-T5-XL (3B), and the results on the Crossfit CLS-45 task split are summarized as follows. We ran all methods on a single A100 80GB GPU where iMAML (and its mixtures) and BMAML baselines incurred OOM.

The result shows that our LiFT algorithm is scalable to an LLM with 3B parameters, where it continues to show large improvements over the baselines. Thanks for suggesting this experiment. We have added this result in our revised paper (Appendix E.1 and Table 6).

4. How about "jointly learning the source task" without a Bayesian approach?

“Jointly learning the source task without a Bayesian approach” – This is exactly what “union-traininig” does, and it underperforms ours by large margin.

Thank you for the detailed response and clarification. My concerns are well-addressed and I have changed the score accordingly.

1. Why the hierarchical Bayesian method works.

To show the importance of hierarchical modeling, we have implemented and tested two non-hierarchical Bayesian models. In the first model, the PEFT LoRA parameters are treated as random variables as ours, but they are shared across all tasks. Hence the posterior inference in this case amounts to learning task-shared (task-agnostic) information from meta training data. We used the SGLD posterior inference. This model performs way worse than our hierarchical model as shown in the table below.

As a second non-hierarchical Bayesian model, we can think of a model where each task is represented by its own PEFT LoRA parameters, but there is no governing higher-level random variables as in our LiFT model. Consequently this model aims to learn only task-specific information, and would not transfer anything to novel tasks. Hence it can be seen as a Bayesian version of "No Meta Train" method, which only runs test-time adaptation via SGLD posterior inference. This model also performs poorly compared to our hierarchical model as shown in the following table.

We have added this result in the revised paper (Appendix E.2 and Table 7).

2. Need ablation study on $\sigma^2$ and $\beta^2$.

To see the robustness to these hyperparameters, we have run the ablation study with our LiFT $K=5$ model on the Crossfit CLS-45 task split. As shown in the following table, the test performance is quite robust to the choice of ($\sigma^2$, $\beta^2$).

Please see also the plot of this ablation study in our revised paper (Appendix E.3 and Figure 7).

3. How the proposed algorithm scales with #adapted parameters (can it only be used with PEFT methods or even with FFT?). Training times for the warm-up and burn-in steps.

There is no point in adapting the full backbone parameters (FFT), which is known to overfit and underperform PEFTs in situations with modest downstream task sizes.

We report the wall clock training times (on a single V100 GPU) in the following table. Our LiFT with $K=5$ GMM is compared with the two baselines “Union Train” and “MAML” on two different PEFT (rank 4 and 64) cases. The result below signifies that our model is as efficient as the vanilla SGD update (ie, “Union Train”) for both low and high ranks, thanks to our efficient SGLD-Gibbs sampler. Asymptotically (in Big-$O$ notation), both the vanilla SGD update and our SGLD-Gibbs (Eq.9-11) $+$ online EM (Eq.13-14) take the same $O(T+d)$ time per iteration where $T$ is the time for the backbone forward/backward computation for evaluating $\nabla_{\theta_i^a} \log p(D_i|\theta_i^a)$ and $d$ is the number of PEFT parameters. This is because of the Gaussian $p(\theta_i^a|\phi)$ which allows analytic gradient $\nabla \log p(\theta_i^a|\phi)$ with respect to both $\theta_i^a$ and $\phi$. The EM iteration also takes $O(d)$ time.

The warm-up steps in our method refers to the vanilla SGD steps for the first 1000 iterations, and the burn-in steps amounts to running the proposed SGLD recurrences for the next 1000 iterations (without collecting the posterior samples). After the burn-in steps, we collect the posterior samples while running the SGLD. The wall clock times for these two stages are shown below, and they are quite reasonable times.

We have added the above results in the revised paper (Appendix E.5, Table 10 and 11).

(Our responses continue in the thread below.)

1. A theoretical justification of the SGLD-Gibbs

The convergence of the SGLD recurrences to the target distribution has been studied recently in (Zou et al., 2021, Xu et al., 2018, Raginsky et al., 2017). Also it is well known in the MCMC literature that the Gibbs sampling, if all the variables are visited sufficiently many times, converges to the stationary distribution since it satisfies both the detailed balance equation and ergodicity. Then since our SGLD-Gibbs has the transition kernel in MCMC that is a composition of these two operators, it should converge to the target distribution.

We have added the above argument in our revised paper (Footnote 1 in p.4).

However, what is not evident in theory and needs further theoretical investigation, is whether our asynchronous update scheme in Eq.(9-11) would also converge to the stationary distribution of the chain. This is mainly because the sum of the gradients $\nabla_\phi \log p(\theta_i^a|\phi)$ over all tasks $i$ in the $\phi$ update Eq.(7) is not exact in the asynchronous scheme since we use the cavity sum (ie, the sum except the current task $i$) that is computed from the old $\phi$s.

We will be investigating further this theoretical analysis, but it may be difficult for us to come up with a decent theoretical conclusion during this rebuttal stage.

2. An alternative version to the $J$ update scheme. Eg, sample approximation of the sum in (7).

To see the impact of our $J$ update strategy (Eq. 9-11), we have run the stochastic approximate version of (7), which replaces the average of $\nabla_\phi \log p(\theta_i|\phi)$ by the current task iterate alone. As shown in the table below, it has decent performance but slightly underperforms our original proposal of the $J$ update strategy. We thank the reviewer for an interesting suggestion, and we have mentioned it in the revised paper (Please see Appendix E.4 and Table 9).

3. Related work section with the online EM literature.

We realize that there have been many prior works on the online EM methods in the literature. We have cited them along with the papers suggested by the reviewer in the extended related work section Appendix D in the revised paper. Although we find that none of them are identical to the one proposed in our paper, most are similar in nature to ours (e.g., exploiting the recursive structure of the EM). Hence we believe that these prior online EM methods can be employed in our online posterior estimation, and can be equally successful.

The discussions on the related online EM algorithms can be found in Appendix D in the revised paper.

4. Code open sourcing.

We consider releasing our code should the paper be accepted.

5. Motivation on the $J$ update scheme (Need more elaboration).

Although $J$ is actually introduced in L:179 p.4 where we stated that $J$ represents and maintains $\sum_i \nabla_\phi \log p(\theta_i|\phi)$, we now elaborate this further in the revised paper. Please see the text before Eq.(9). We reiterate it as follows:

Another computational bottleneck is the sum of the gradients $\nabla_\phi \sum_i \log p(\theta_i^a|\phi)$ over all $i=1,\dots,N$ for each $\phi$ update in (7). To remedy this issue, we introduce an auxiliary variable $J$ that maintains this sum of the gradients, and we let it asynchronously updated (11): at each iteration for task $i$, subtracting the old gradient for $i$ from $J$ to approximate the cavity sum $\sum_{i'\neq i} \nabla_\phi \log p(\theta_{i'}^a|\phi)$ and adding a new gradient $\nabla_\phi \log p(\theta_{i}^a|\phi)$ to $J$.

6. Provide results of ablation study comparing against Gaussian confirming that such enrichment is actually useful.

Fig. 3 already shows comparison between LiFT Gaussian posterior approximation (K=1) vs. LiFT GMM posterior approximation (K=3 and K=5). Clearly, GMM is better than Gaussian (K=1) for most cases.

7. Report means of the absolute deltas between the baseline and the candidate algorithm.

We report the mean of absolute deltas on the CLS-45 task split as follows: