C o B a: Convergence Balancer for Multitask Finetuning of Large Language Models (2024)

Abstract: Multi-task learning (MTL) benefits the fine-tuning of large language models (LLMs) by providing a single model with improved performance and generalization ability across tasks, presenting a resource-efficient alternative to developing separate models for each task. Yet, existing MTL strategies for LLMs often fall short by either being computationally intensive or failing to ensure simultaneous task convergence. This paper presents CoBa, a new MTL approach designed to effectively manage task convergence balance with minimal computational overhead. Utilizing Relative Convergence Scores (RCS), Absolute Convergence Scores (ACS), and a Divergence Factor (DF), CoBa dynamically adjusts task weights during the training process, ensuring that the validation loss of all tasks progress towards convergence at an even pace while mitigating the issue of individual task divergence. The results of our experiments involving three disparate datasets underscore that this approach not only fosters equilibrium in task improvement but enhances the LLMs’ performance by up to 13% relative to the second-best baselines. Code is open-sourced at https://github.com/codefuse-ai/MFTCoder.

## introduction in recent years, large language models (llms) have emerged as a focal point of research within both academia and industry, owing to their superior performance. these models are initially pretrained, designed to ensure they possess broad applicability across a variety of downstream tasks. this is followed by a finetuning stage, which meticulously adapts the models for specific tasks or scenarios. however, this phase requires individual, task-specific finetuning, leading to a complex deployment scenario in production environments. ## related work all multi-task learning (mtl) approaches (crawshaw, 2020; @xcite are designed to foster positive knowledge transfer across tasks through parameter sharing, while simultaneously minimizing any potential negative transfer, often referred to as task interference. our discussion here primarily revolves around optimization techniques within mtl, given their relevance to llms. we categorize the existing work into two distinct groups: classical mtl methods and mtl methods tailored for llms. classical methods traditional mtl strategies aimed at addressing task imbalance from an optimization standpoint fall into two categories: gradient manipulation and loss balance. gradient manipulation techniques @xcite @xcite ) create a composite update vector at each optimization step by amalgamating task gradients. this approach ensures local improvements across all tasks but can be computationally intensive, particularly for models with numerous parameters. conversely, loss balancing methods dynamically adjust each task's weight during training based on predefined factors such as task uncertainty @xcite , task difficulty prioritization @xcite , and random loss weighting @xcite . these methods are computationally more efficient but do not guarantee simultaneous convergence of all tasks. to overcome this issue, advanced solutions aiming at convergence balance include dwa @xcite , lbtw @xcite , grad-norm @xcite , metaweighting @xcite , and famo @xcite . the latter three adjust task weights based on gradients, with metaweighting additionally focusing on the validation instead of the training loss to enhance generalization performance. unfortunately, gradient-based weight adjustment can be computationally demanding, as shown in methods for llms mtl research specific to llms is still in its infancy, with a handful of notable contributions from models such as t5 @xcite , muppet @xcite , ext5 @xcite , and mftcoder @xcite . the initial three primarily aggregate data from various tasks without considering task equilibrium, often overlooking tasks with smaller datasets and favoring those with larger ones. mftcoder advances this by calculating individual loss for each task, yet assigns equal weights across the board. mftcoder acknowledges the inability of such approaches to ensure uniform validation loss convergence across tasks and suggests leveraging famo as a potential solution. the proposed method, coba, embodies the strengths of both classical and llm-specific mtl approaches, achieving convergence balance among tasks with minimal additional computational demands. crucially, it focuses on validation loss, thereby promising to maintain or improve the generalization capabilities of the model. ## convergence balancer (coba) multi-task learning (mtl) is engineered to optimize a single model, parameterized by θ ∈ r m , enabling it to adeptly perform @xmath0, potentially even in tandem. the loss function for task i at the t-th iteration is denoted by ℓ i (θ; t) : r @xmath1. this forms the foundation for the optimization challenge of mtl, expressed as: these two criteria are quantified respectively through relative convergence scores (rcs) and absolute convergence scores (acs). rcs is employed to assess the convergence pace relative among tasks, while acs measures the current convergence rate against historical rates for each task individually. note that these scores are derived from the slopes of validation losses-not the gradients-to minimize computational demands. we then integrate a divergence factor (df), highlighting the overall convergence trajectory across tasks. this factor ensures rcs impacts weights predominantly when all tasks are converging, while acs takes precedence as an arbitrary task commences diverging. in the sequel, we detail the computation of slopes, the formulation of rcs and acs, and their amalgamation with the df. ## experiments in this section, we assess the performance of the coba across four diverse datasets: the code completion (cc) dataset, encompassing five programming languages; the code-related task (crt) dataset, featuring five unique programming tasks; the xtreme-up, which delves into question-answering across nine natural languages; and multi-domain qa dataset, including question answering data in the fields of coding, mathematics, and natural language. due to the space limit, the results of the last dataset are shown in appendix d. the tasks within these datasets are inherently related and generative, making them ideal candidates for mtl experiments on llm. for further insights into the datasets, readers are directed to the appendix a. our evaluation benchmarks the coba against 8 state-of-the-art (sota) baselines foot_0 : single-task learning (stl), which finetunes each task in isolation; uniform @xcite , applying equal weights to all tasks in an mtl framework; gradnorm @xcite , a method that optimizes the task weights iteratively such that task-specific gradients are of similar magnitude; lbtw @xcite , which dynamically adjusts task weights according to the ratio of current to initial loss ω i (t) = (ℓ i (t)/ℓ 0 (t)) b , parameterized by a hyperparameter b; and famo @xcite , aimed at optimizing weights to enhance the minimal improvement rate across tasks. notably, the last three methods were originally designed based on the training loss. in pursuit of enhanced generalization for the fine-tuned models, we have adapted these methods to focus on validation loss, denoted as lbtw * , gradnorm * , and famo * . except for stl and uniform, all methods strive to balance convergence across tasks, demonstrating their potential to compete with coba. the detailed experiment setup is described in appendix b. ## conclusion in this paper, we propose coba, a novel mtl method for llms that simultaneously achieves convergence balance with low computational coma performance gap. in contrast, famo * optimizes the original training loss without performing re-normalization. indeed, we observe that famo exhibits a higher loss scale (both training and validation) in nearly all experiments compared to famo * , except in the multi-domain qa dataset where a higher learning rate is applied. this disparity is particularly pronounced in the xtreme-up experiments, where the loss for famo is approximately double that of famo * . 3. the f1 score metric employed for xtreme-up has a strong correlation with the loss, which elucidates why famo's performance on this dataset is significantly inferior to that of famo * . conversely, for code-related datasets, the pass@1 metric shows a relatively weak correlation with loss; thus, a decrease in loss does not necessarily translate to an increase in pass@1. this may account for the comparable performance of famo and famo * in these code-related datasets.