BoRA: Bayesian Hierarchical Low-Rank Adaption for Multi-task Large Language Models

Authors: simen  eide
arnoldo  frigessi
Venue: NLDL_org_2025_Conference
Published: 2024-07-08 08:22:24

Abstract: This paper introduces Bayesian Hierarchical Low-Rank Adaption (BoRA), a novel method for finetuning multi-task Large Language Models (LLMs). Current finetuning approaches, such as Low-Rank Adaption (LoRA), perform exeptionally well in reducing training parameters and memory usage but face limitations when applied to multiple similar tasks. Practitioners usually have to choose between training separate models for each task or a single model for all tasks, both of which come with trade-offs in specialization and data utilization.

BoRA addresses these trade-offs by leveraging a Bayesian hierarchical model that allows tasks to share information through global hierarchical priors. This enables tasks with limited data to benefit from the overall structure derived from related tasks while allowing tasks with more data to specialize. Our experimental results show that BoRA outperforms both individual and unified model approaches, achieving lower perplexity and better generalization across tasks. This method provides a scalable and efficient solution for multi-task LLM finetuning, with significant practical implications for diverse applications.

Main paper: 
 introduction: 
 Large Language Models (LLMs) have become popular models for solving a wide range of generative text problems. LLMs are typically pretrained on large and diverse datasets of texts, and then fine-tuned on new datasets dedicated to specific tasks. The finetuning process is often done by maximizing the next token log likelihood of the task-specific dataset, and the resulting model is then used to make predictions and generations for new data. A popular fine-tuning method is Low-Rank Adaptation (LoRA), which reduces the number of trainable parameters by a factor of up to 10,000 and the GPU memory requirement by a factor of 3, while still outperforming other fine-tuning techniques ( [1] ). When performing fine-tuning of LLMs, the practitioner often wants to solve multiple similar tasks at the same time. For example: A media organization might want to use an LLM to generate various textual elements like headlines, keywords, and summaries for a given news article. A website might have multiple chatbot tasks answering user inquiries across various topics. And finally, a scriptwriting AI might want to generate dialogues for different characters in a movie, each with its unique style and vocabulary.In these multi-task cases, the practitioner can either (1) train one model for each task and dataset separately, or (2) train one model for all tasks and condition the specific output wanted (e.g. [2] ). There is a trade-off between these two approaches: Train each task separately and the model will be able to specialize on each task, but it may suffer when limited data are available for some tasks. Train one model for all tasks, and the model will be able to share information between tasks, but it may not be able to specialize on each task due to limited model capacity.In this paper, we suggest a way to circumvent these trade-offs and present a generalization of the two LoRA fine-tuning methods for multi-task problems, which we call the Bayesian Hierarchical Low-Rank Adaptation method (BoRA). The BoRA method allows tasks to share information between each other through a set of global hierarchical prior parameters, in a Bayesian hierarchical setting. This means that tasks with little data can benefit from the global structure across related tasks, while tasks with more data can specialize on their specific tasks. We show in Section 3.4 that BoRA can be seen as a generalization of the two approaches above.To test our method, we evaluate BoRA on the next token generation problem of parliamentary speeches, letting each parliamentary representative be a task. In our setting, BoRA is superior to both conventional approaches described above and suggests a way to mitigate the trade-off between the two. The paper is organized as follows: Section 2 presents related works. Section 3 presents BoRA, the Hierarchical LoRA method; Section 4 describes the experimental setup, Section 5 presents the results, and Section 6 discusses the results and the implications of the new method. Finally, we make the code openly available at https://github.com/simeneide/bora.We group related work into two areas: multi-task learning and fine-tuning methods of large language models. Autoregressive LLMs and Bayesian modeling are described in Section 3. 
multi-task learning: 
 Multi-task Learning [3] is a popular machine learning approach for training one model on multiple tasks. A common approach is to share the lower layers of a neural network between tasks and have separate parameters for the last (or last few) layers for each task. Our paper suggests a different approach, where we share information between tasks through a hierarchical prior over the task-specific parameters.Modern training of language models can itself be considered a form of multi-task learning where the model is trained on a diverse set of tasks and datasets [4, 5] . However, with the trend of even bigger models, pretraining of LLMs is only available to a few large organizations due to both computational and data requirements. Therefore, there is a need for methods that can leverage the pretraining of these models for smaller tasks and datasets such as the one we present in this paper.Model merging of LLMs [6] [7] [8] is another popular approach, often showing promising performance. It takes multiple task-specific models and merges them into one model. The standard way is by doing a weighted average of the model parameters ( [6] ), but more sophisticated methods exist ( [7, 8] ). Model merging differs from our approach in that it merges the models after training, while the BoRA method shares information between tasks during training through a Bayesian hierarchical model."A popular approach to fine-tuning neural networks involves adjusting only the top layer (or the n top layers) of a pre-trained neural network (e.g., [9] ). These techniques exploit the fact that features learned in the lower layers of the network are general and transferable, while the top layers are more task-specific. This approach has been widely adopted in neural networks. However, in this paper, we focus on the LoRA method, which instead allows us to train a low-rank decomposition of the entire network.There are several variations of the LoRA algorithm. The Sparse Mixture of Low Rank Adapters (SIRA) [10] , for instance, creates a Mixture of Expert on the low-rank parameters. Another variation, the Weight Decomposed Low-Rank Adaptation (DoRA) [11] , decomposes the low-rank parameters into scale and unit norm direction parameters. There is also the Efficient Low Rank Adaptation (LoRA+) [11] , which improves the existing LoRA optimization algorithm by adjusting the learning rate of the lowrank parameters. Lastly, MoRA [12] decomposes the low-rank parameters using learnable square matrices with non-parameter based compressions and decompressions. All these methods are designed to improve the performance of language models similar to LoRA, and they can be used in conjunction with the proposed Bayesian Hierarchical LoRA method.3 Method: Hierarchical LLM This section describes the Hierarchical LoRA method for multi-task problems. We begin by defining the task and data structure, followed by an introduction to pretrained autoregressive language models. Next, we describe the LoRA method, and finally, we present the proposed Bayesian Hierarchical LoRA model. 
task and data structure: 
 We define our study across D distinct tasks, denoted by d = 1, 2, ..., D. For each task d, we have a dataset D d containing N d documents, with each document n = 1, ..., N d consisting of a sequence of W n tokens. This dataset structure can be formally represented as:EQUATIONIn this notation, w d,n,1:Wn represents the token sequence in the n th document of the d th task. For simplification in subsequent discussions, we will omit the indices d and n (and write w i ) when their reference is evident from the context.Our foundational model is a pretrained autoregressive language model, characterized by its parameters θ f ull and its architecture, referred to as a Large Language Model (LLM). The parameters θ f ull were acquired through pretraining on a large and diverse dataset, distinct from our current dataset D. The model's probability distribution is expressed as:EQUATIONwhere LLM (w i+1 |w 1:i ) denotes the probability of the next token w i+1 given the preceding tokens w 1:i in the document. For simplicity, we omit θ f ull explicitly in the LLM, since the autoregressive language model inherently depends on these parameters.A Large Language Model consists of numerous linear layers. The LoRA method reduces the number of trainable parameters by reparameterizing the weight parameters (W ∈ R n1×n2 , W ⊂ θ) of these layers into the original pretrained weights W f ull and a low-rank decomposition A and B ( [11] ):EQUATIONwhere W f ull ∈ R n1×n2 is the original pretrained weight, B ∈ R n1×r and A ∈ R r×n2 are the lowrank factors, n 1 and n 2 are the input and output dimensions of the original weight matrix in the linear layer, r min(n 1 , n 2 ) is the low-rank dimension, and α is a scaling hyperparameter.During LoRA fine-tuning, the low-rank factors A and B are learned (estimated) from the data, while the original pretrained weights W f ull remain fixed to their pretrained values.In Tasks share some common knowledge, modeled by introducing a Gaussian hierarchical prior over the task parameters θ 1:D , described by:EQUATION)where Θ is a set of hierarchical mean parameters, I is a unit diagonal matrix, and τ ≥ 0 is a scalar precision hyperparameter controlling how similar the task parameters θ d should be to the hierarchical mean parameters Θ. We let Θ have an improper uniform hyperprior (i.e. P (Θ) = 1 for all values of Θ). The prior assumes that the task parameters share a common origin because the tasks are related: input data come from similar distributions and the output tasks share similarities. The precision hyperparameter τ dictates how much the task parameters θ d are a priori expected to deviate from the hierarchical mean parameters Θ. It indicates how much structure and information is shared among tasks and should be set by the practitioner based on prior knowledge. The posterior distribution of the hierarchical model, obtained by combining equations 4 and 5, is:EQUATIONThe resulting BoRA model utilizes each task based on its dataset size and similarity to other tasks: If a task has few data points, the posterior in equation 6 will be dominated by the prior term, and the task parameters will remain near the global hierarchical parameters, borrowing information from these.Conversely, if a task has many data points, the likelihood will dominate the posterior, allowing the task parameters to focus more on the individual task.The hierarchical LoRA model can be seen as a generalization of the two conventional approaches for multi-task problems: If we let τ approach zero, the probability distribution of the prior (eq. 5) will become a flat uniform prior and the situation will be equivalent to training each task independently. On the other hand, if we let τ approach infinity, the prior will become a strong constraint on the task parameters, preventing them from deviating from the global parameters, and the situation will be equivalent to training one model for all tasks.Given the precision hyperparameter τ , we seek the maximum a posteriori (MAP) estimate of the hierarchical model. We achieve this using AdamW ( [13] ), a gradient-based optimization algorithm, by optimizing over the task parameters θ 1:D and the hierarchical mean parameters Θ using the posterior in equation 6. This implies that the task parameters will have gradients both towards minimizing the conventional LLM loss (represented by the likelihood term) and the hierarchical mean parameters (originating from the prior). The hierarchical mean parameters Θ will have gradients towards the average of the task parameters, learning to represent the global structure of the tasks.An important observation is that the gradient of the prior term in equation 6 is proportional to the precision hyperparameter τ . This means that the optimization will require more steps to converge for larger values of τ , as the task parameters will be more constrained towards the hierarchical mean parameters and a posteriori dependent on each other. Therefore, we adjust the learning rate of the AdamW optimizer to be proportional to the precision hyperparameter τ , and we find in our experiments that this causes the optimizations to converge at a similar rate for different values of τ .  
experimental setup: 
 To test our method, we optimize the model with a range of different precision hyperparameters τ , and evaluate the perplexity on the test set for each task. Perplexity is defined as the exponentiated average negative log-likelihood per token, and is a common metric for evaluating language models. A lower perplexity indicates a higher likelihood and a better model. We use the Talk of Norway Dataset [14] , which consists of speeches by Norwegian parliament politicians. Speakers come from different political parties and geographical areas, but the tasks will still share a common domain represented by parliamentary speeches in Norway over a short period of time. Therefore, we can consider different speakers of this dataset as different tasks. Specifically, we select the first 25 speakers chronologically who have more than 100 speeches in the dataset. This results in 25 tasks with samples ranging from 110 to 477 documents, and an average of 202 documents per task. For some speakers, we have significantly less data compared to others. The distribution of samples per task can be seen in Figure 1 . We reserve 33% of the data for the test set.We select a model that has been pretrained on a large and diverse dataset, but has not seen the specific tasks that we want to evaluate it on. We use the 'opt-350m' model [15] , which is predominantly trained on English, and has only seen a small amount of Norwegian data from the CommonCrawls dataset. 
results: 
 The overall results of the hierarchical LoRA model on the Talk of Norway Dataset are shown in Table 1 and Figure 2 . We see that the model achieves the best test perplexity of 12.82 when using a precision parameter of 100. The improvement is largest compared to training the models independently (τ = 0), and smaller when comparing to the limiting case where all task-specific model parameters are constrained to be equal (τ = 10000). When considering each task separately (Table A .1), we see that all tasks benefit from a hierarchical model compared to both the case when all models are trained independently (τ = 0) and the limiting case where all task-specific model parameters are constrained to be equal (τ = 10000). Figure 3 shows the relative improvement of each task when using the hierarchical model compared to the two cases versus their dataset sizes. There is an indication that the improvement from training separately is larger for tasks with less data, while this is not observed for the "one model" case. A precision parameter of 100 gives the best perplexity for all our tasks. 
discussion: 
 Our results in Figure 2 demonstrate that the hierarchical LoRA model can improve test perplexity for all tasks compared to both training the models independently and training one model for all tasks. This confirms our hypothesis that BoRA can be a useful method for multi-task problems, mitigating the trade-off between training each task separately and training one model for all tasks. Notably, BoRA also benefits tasks with relatively more data. We observe in Figure 4 that the distance between the task parameters and the global prior increases with increasing task dataset size. This is expected in a hierarchical model, as tasks with more data will have a larger likelihood term in the posterior, allowing them to specialize more on their specific task during training.We observe no consistent relationship between dataset size and the final perplexity (Figure 5 ). Although one might expect that tasks with more data would learn more about their specific task and achieve lower perplexity, this pattern is not evident in our experiments. Other task-specific factors, such as under-represented dialects and political views, might influence the final perplexity. Controlling for these factors is complex and beyond the scope of this paper, but they are likely contributors to this outcome. However, when considering the relative improvement of our best-performing model, the results in Figure 3 indicate that tasks with less data may benefit more than those with more data. Although the signal is subtle, this is an expected result, as the hierarchical model allows tasks with limited data to leverage information and structure from the global model, while tasks with more data can specialize in their specific task.In principle, one could adopt a fully Bayesian approach, instead of Maximum a Posteriori (MAP) maximization, to assess the reliability of the hierarchical model. Such an approach would necessitate the use of Markov Chain Monte Carlo (MCMC), Variational Inference, or other methods.This approach would involve sampling from the posterior distribution of the hierarchical model, al-lowing for the estimation of uncertainty in the task parameters. However, MCMC methods are computationally intensive, and it remains unclear to us how to efficiently and scalably sample from the posterior distribution of the language model. A middle ground between a full Bayesian approach and the MAP estimate could involve being Bayesian for select parameters. For example, instead of our current approach where we maximize the likelihood of the validation set concerning the precision parameter τ , one could be Bayesian and integrate over its empirical posterior distribution obtained from MCMC. This would provide insights into how well tasks share information in a specific problem context. We leave these explorations for future research.  
conclusion: 
 This paper presents a novel generalization of the LoRA fine-tuning method for multi-task problems, which we call BoRA, Bayesian Hierarchical Low-Rank Adaptation for Multi-task Large Language Models. The hierarchical model allows tasks to share information between each other through a set of global hierarchical prior parameters and can be seen as a generalization of the two conventional approaches for multi-task problems. We demonstrate that BoRA can improve test perplexity for all tasks compared to both training the models independently and training a single model for all tasks, effectively mitigating the trade-off between these two approaches. We propose adjusting the learning rate according to the design parameter that controls the trade-off between likelihood and prior. Future work includes investigating the effect of BoRA on more tasks and datasets, examining how the relative capacity of the global model affects performance, and exploring the full posterior distribution of the model. 
 
 Appendix: Table A .1 shows the test perplexities for each task when using BoRA with different precision hyperparameters τ . Proceedings of the 6th Northern Lights Deep Learning Conference (NLDL), PMLR 265, 2025. 2025 Simen Eide & Arnoldo Frigessi. This is an open access article distributed under the terms and conditions of the Creative Commons Attribution license (http://creativecommons.org/licenses/by/4.0/). 


Decision: The paper proposes to use a Bayesian model for fine-tuning optimization of a multi-task LLM model. 

The contribution could be considered incremental. 

Experiments are weak (no comparison) and the plots are low-quality. 

The paper is evaluated as clear and well-presented. 

Extending the training to a multi-task scenario is reasonable, and using the Bayesian model for the training is a sound approach. 

pros:
1. fine-tuning of LLM in a multi-task scenario could be a realistic scenario
2. The Bayesian framework is reasonable 

cons:
1. weak evaluation, with no comparison
2. low-quality plots
3. incremental contribution

Reviews: review: title: title: low-rank adaptation methods for LLM in the context of fine-tuning for multiple tasks at onc
paper_summary: paper_summary: **Summary**: The authors propose a low-rank adaptation methods for LLM in the context of fine-tuning for multiple tasks at once. In contrast to standard `LoRA` where one module would be trained **independetly** for each task, `BoRA` aims to leverage possible transferability across tasks. More specifically, BoRA also introduces separate module parameters $\theta_d$ for each task; but, to model the relationship between tasks, the join distribution $\theta_{1\dots D}$ is modelled with a Gaussian hierarchical prior, with a uniform hyperprior $\Theta$.
main_review: main_review: Taking into account the rebuttal and other reviews, I am inclined to keep my original rating: While I think the experimental section could include more bechmarks and/or baselines, the proposed idea is interesting and well explained, and experiments include an interesting ablations on $\tau$.
strength_weakness: strength_weakness: **Weaknesses:**
  * By design, BoRA seems to assume the tasks are related enough to be modelled by a joint prior. However, task interference / negative transfer is a key issue in the multi-task literature: While the model does allow for the BoRA module parameters to deviate from the common prior, it would be interesting to test this In practice. In fact, the results of Table 1 seem to imply that the benchmark considered in the paper does not suffer much from task interference, since a single LoRA module trained on all task performs better than training  independent LoRA modules for each task (13.91 vs 16.70 perplexity)
  * In terms of related work, the paper could also mention the task merging literature as it seems closely related to multi-task PEFT in general.

**Minor notes:**
  * Some notations are a bit cumbersome (e.g. introducing the intermediate notion of documents in Section 3.1 seems superfluous; for instance, most recent LLMs tend to directly refer to the number of tokens they have been trained on)
  * In Figure 2, it would be useful to plot a constant line for the perplexity value of $\tau = 0$ ($\sim$ LoRA baseline)