Amortized B ayesian Meta-Learning for Low-Rank Adaptation of Large Language Models (2025)

Abstract: Fine-tuning large language models (LLMs) with low-rank adaptaion (LoRA) is a cost-effective way to incorporate information from a specific dataset. However, it is often unclear how well the fine-tuned LLM will generalize, i.e., how well it will perform on unseen datasets. Methods have been proposed to improve generalization by optimizing with in-context prompts, or by using meta-learning to fine-tune LLMs. However, these methods are expensive in memory and computation, requiring either long-context prompts or saving copies of parameters and using second-order gradient updates. To address these challenges, we propose Amortized Bayesian Meta-Learning for LoRA (ABMLL). This method builds on amortized Bayesian meta-learning for smaller models, adapting this approach to LLMs while maintaining its computational efficiency. We reframe task-specific and global parameters in the context of LoRA and use a set of new hyperparameters to balance reconstruction accuracy and the fidelity of task-specific parameters to the global ones. ABMLL provides effective generalization and scales to large models such as Llama3-8B. Furthermore, as a result of using a Bayesian framework, ABMLL provides improved uncertainty quantification. We test ABMLL on Unified-QA and Crossfit datasets and find that it outperforms existing methods on these benchmarks in terms of both accuracy and expected calibration error.

## introduction large language models (llms) handle a variety of tasks reasonably well @xcite . however, to tailor llms to specific domains, fine-tuning on specific datasets is often necessary. while methods such as low-rank adaptation (lora; @xcite ) fine-tune a pretrained llm costeffectively, a fine-tuned llm is limited to the domain it is trained on. its performance may not improve in other domains and sometimes worsens as it suffers from catastrophic forgetting. such catastrophic forgetting may result in overfitting and erasing existing capabilities of the pretrained llm @xcite . meta-learning is a strategy for solving this problem, training models on a variety of tasks in a way that supports generalization across tasks @xcite . however, meta-learning typically requires a large amount of computation and memory, making it challenging to apply to llms. one form of metalearning that has been applied to llms involves fine-tuning models on in-context prompt-response examples @xcite ). another more traditional approach, maml-en-llm @xcite , adapts the model-agnostic meta-learning (maml) @xcite framework to llms. however, both methods are limited in the size of the language models that can be used: the former requires long-context prompts, whereas the latter uses second-order gradient updates and saves a model for each task. recent work on amortized bayesian meta-learning (abml; @xcite ) addresses some of the computation and memory requirements of meta-learning. this approach posits a generative model over parameters where taskspecific parameters are generated from global parameters, and inference over task-specific parameters is amortized. in other words, the conditional distribution over task-specific parameters is shared across tasks, implying that computation and memory costs stay constant with respect to the number of tasks. this approach thus offers a path towards efficient meta-learning for llms. however, several challenges exist. first, we need to specify the generative model over weight space in the context of llms. second, the prior term used in abml no longer adapts to the setting of fine-tuning a pretrained model because the spread of its weights mismatches that of an arbitrary prior used to train a model from scratch. third, the enormous size of llms makes training difficult, as the scale of probabilities assigned to the model variables can overwhelm the influence of the data likelihood. in this paper, we present a solution to these problems, taking a bayesian approach to fine-tuning llms using abml. to define the underlying generative model and efficiently characterize the distributions involved, we use lora to express both the model weights and their uncertainty. we introduce a new prior over global variables that accounts for the spread of the parameters learned in the pretrained model. we also introduce two adjustable hyperparameters that balance reconstruction accuracy and the fidelity of task-specific parameters to the global ones. using amortized bayesian meta-learning for llm fine-tuning, we achieve both higher accuracy and better uncertainty estimation over unseen tasks compared with regular fine-tuning and other scalable methods in the meta-learning literature. figure 1 illustrates an example where incorporating uncertainty estimation in fine-tuning leads to a more calibrated model response. our method is scalable and avoids the computation and memory overhead of other meta-learning approaches, making it adaptable to larger models such as llama 3 8b. we show that amortized bayesian meta-learning provides fine-tuned llms that are accurate on domainspecific tasks, more generalizable to new tasks, and provide better uncertainty estimation. as a different approach, @xcite and @xcite explored meta-learning for llms using in-context learning. these works show that it is possible to fine-tune llms on in-context examples and achieve generalization. however, our approach does not require curation of such examples, does not place constraints on the size of the context window of a model, and is more scalable. uncertainty representation for llms. approaches to capturing uncertainty for llms can rely on the intrinsic representation of uncertainty in the model or focus on capturing extrinsic uncertainty about model parameters. intrinsic approaches produce better uncertainty calibration via prompt engineering and sampling @xcite or learning an external model @xcite . extrinsic approaches include using finetuning methods to incorporate uncertainty, such as training lora with ensembles @xcite , laplace approximation @xcite , and variational inference @xcite . our work takes the extrinsic approach but differs from existing approaches by using the metalearning setting to achieve strong uncertainty calibration through generalization across datasets. ## method our method extends amortized bayesian meta-learning, making it possible to apply to llms. this approach combines the advantages of metalearning for adapting to new tasks with bayesian inference for uncertainty representation. we use the the objective of eq. 1 from abml. in our setting, θ and ϕ i are the global and taskspecific model parameters produced as the output of lora adapters. on a high level, the generative process is algorithm 1 one epoch in the abmll algorithm. the "test section" does not need to be performed every epoch. input: likelihood model p(d i |ϕ i ), prior p(θ) and p(ϕ|θ), variational posterior q θ (ϕ i |d i ), with trainable parameters b, a; constant c, β; number of tasks m and inner-loop size k. training section for task @xmath0, 2, ..., m } do draw batch d i from task i dataset. inner-loop: for iter @xmath1, 2, ..., k} do run a step gradient descent to minimize w.r.t. ϕ i : -e q θ (ϕ i |d i ) [log p(di|ϕi)] + βkl q θ (ϕi|di) p(ϕi|θ) . end for outer-loop: run a step gradient descent to minimize w.r.t. θ: -e q θ (ϕ i |d i ) [log p(di|ϕi)] + βkl q θ (ϕi|di) p(ϕi|θ) + βkl(q(θ)||p(θ)). end for test section take unseen task i. create a copy of the above weights, and on the new weights: for iter @xmath2, 2, ..., k} do draw batch d i from task i dataset. run a step gradient descent to minimize w.r.t ϕ i : -e q θ (ϕ i |d i ) [log p(di|ϕi)] + βkl q θ (ϕi|di) p(ϕi|θ) . ## empirical evaluations model and datasets. we fine-tune llama3-8b on crossfit @xcite and unifiedqa @xcite , textual datasets commonly used to train meta-learning models. because a key aim of our paper is to study uncertainty quantification, we filter for multiple choice datasets, leading to a subset of crossfit and unifiedqa with 34 datasets with 68k training datapoints in total. they feature problems such as sentiment analysis, natural language inference, and identifying particular traits or topics in a given text. for evaluation on an unseen task, we use winogrande @xcite , a multiple choice dataset evaluating common sense reasoning. metrics. we use accuracy to evaluate general performance and expected calibration error (ece) to evaluate uncertainty estimation. implementation details. all methods use batchsize of 2 and inner-loops with 5 gradient steps. lora adapters follow standard practice with @xmath3, and learning rate is tuned in [10 -6 , 5 • 10 -5 ]. for abmll, β = 5 • 10 -10 , γ = 10 -6 , @xmath4. for the gamma prior, a 0 = 1, b 0 = 0.01, following @xcite . during validation on the unseen dataset, all models train 10 gradient steps on 10 batches from this dataset and evaluate on the rest. baselines. we use four baseline methods that can viably scale to llama3-8b. pretrained is the offthe-shelf llm. regular lora is the default lora method trained on the whole randomly shuffled training dataset. structured lora also uses the default lora, but the training dataset follows the same "structure" as our method: it is iteratively trained 5 gradient steps on one task at a time. thus, it tests the effect of our generative model on performance. the reptile @xcite algorithm primarily uses a weighted average between new weights and previous weights to achieve metalearning.