CoTBal: Comprehensive Task Balancing for Multi-Task
Visual Instruction Tuning

In current works involving visual instruction tuning, instruction-following data from various tasks are typically indiscriminately mixed for fine-tuning LMMs. The training loss is obtained by averaging the cross-entropy losses calculated across all valid tokens, as represented by the following formula:

where $N$ is the total number of tasks, $S_{i}$ is the number of samples for Task $i$, $T_{ij}$ is the number of valid tokens in the $j$th sample for Task $i$, and $t_{ijk}$ is the $k$th valid token in the $j$th sample for Task $i$. However, this approach is incompatible with the task weighting paradigm of traditional multi-task optimization algorithms, where single-task losses are individually computed and aggregated through weighted summation to get the total loss. Therefore, we introduce the GTW paradigm, specifically tailored for multi-task visual instruction tuning. The training loss of GTW is defined as:

where $\lambda_{i}$ denotes the weight of Task $i$. The losses are assigned task-specific weights and aggregated at the token level rather than at the sample or task level. GTW allows for more equitable consideration of each valid token, ensuring that the model is not biased towards certain tasks due to variations in sample sequence length or data amount across tasks. Besides, we also perform weighting in the denominator to enable a fair comparison with the indiscriminate data mixing approach (see Equation 1), where the weights are uniformly set to $1$. The GTW paradigm is employed in our CoTBal algorithm, while also laying a solid foundation for subsequent studies.

Although the focal points of distinct tasks vary in multi-task visual instruction tuning, a key shared objective exists: achieving more accurate comprehension and reasoning of visual information. As shown in Figure 1(a), the data of detailed image captioning on ShareGPT-4V (Chen et al., 2023) and visual question answering on VQAv2 (Goyal et al., 2017) both involve color information (pink and yellow dishes) in the image, which exemplifies the overlapping knowledge domains among tasks. Therefore, it is reasonable to hypothesize that different visual tasks could potentially provide mutual enhancement in their performance, which can be defined as the inter-task contribution. The extent of the overlapping knowledge domains varies, leading to differing degrees of inter-task contributions.

In practice, the inter-task contribution of Task $i$ to Task $j$ can be quantified by the validation performance for Task $j$ of the model trained on Task $i$, which is normalized by the validation performance for Task $j$ of the model trained on Task $j$ itself. However, a model trained exclusively on one task may struggle to adhere to the instruction demands of other tasks. To address this, we incorporate mini subsets from all tasks into the training set, enabling the model to understand the instruction demands of each task. Therefore, the inter-task contribution of Task $i$ to Task $j$ can be calculated as:

where $V_{j}(i+mini)$ represents the validation performance for Task $j$ of a model trained on the full dataset from Task $i$ alongside mini subsets from other tasks, and $V_{j}(mini)$ signifies the validation performance for Task $j$ of a model trained on mini subsets from all tasks. In the formula, $V_{j}(mini)$ is subtracted from both the numerator and the denominator to mitigate the impact of incorporating mini subsets from all tasks into the training set on the validation performance for Task $j$.

Furthermore, based on the accurate quantification of the inter-task contribution, we propose two task weighting strategies for inter-task contribution balancing. Firstly, we examine the average inter-task contribution of one given task to all other tasks as $C_{one2all}$, representing the extent to which this task assists all other tasks. The greater the assistance provided by one task to all other tasks, the more substantial its overall contribution to the entire training process of multi-task visual instruction tuning. Therefore, tasks that have greater $C_{one2all}$ should be assigned more weights to enhance overall performance. The specific task weights $\bm{\lambda_{one2all}}$ can be computed as:

where $C_{one2all,i}$ signifies $C_{one2all}$ for Task $i$ and $\bm{C_{one2all}}$ represents the $N$-dimensional vector of $C_{one2all}$ for all tasks. $T$ denotes the temperature hyperparameter that controls the degree of smoothness in the weight vector. Secondly, we consider the average inter-task contribution of all other tasks to one given task as $C_{all2one}$, denoting the degree to which this task receives benefits from all other tasks. If one task receives minimal benefits from other tasks, it tends to exhibit poorer performance compared to tasks that receive greater benefits. To maintain balanced overall performance, such type of tasks that have lower $C_{all2one}$ should also be assigned more weights. The specific task weights $\bm{\lambda_{all2one}}$ can be computed as:

where $C_{all2one,i}$ signifies $C_{all2one}$ for Task $i$ and $\bm{C_{all2one}}$ represents the $N$-dimensional vector of $C_{all2one}$ for all tasks. $T$ denotes the same temperature hyperparameter in Equation 5. Subsequently, we integrate the aforementioned two strategies to formulate the task weighting strategy for inter-task contribution balancing, where the task weights $\bm{\lambda_{C}}$ can be calculated as:

In addition to the inter-task contribution, another critical aspect in multi-task visual instruction tuning is the intra-task difficulty, which refers to the inherent learning difficulty within each task. Tasks that achieve near-optimal performance with a limited dataset are considered to have poor intra-task difficulties. Conversely, tasks that require the full dataset to reach optimal performance are deemed to have great intra-task difficulties. As illustrated in Figure 1(b), different tasks exhibit distinct patterns of performance improvement with increasing training data amount. Arranged by increasing intra-task difficulty, the sequence of these three tasks is as follows: visual question answering on VQAv2 (Goyal et al., 2017), detailed image captioning on ShareGPT-4V (Chen et al., 2023) and visual grounding on RefCOCO (Kazemzadeh et al., 2014; Mao et al., 2016).

Practically, the intra-task difficulty for Task $i$ is measured by the validation performance gap between a model trained on the full dataset and that trained on a mini subset from Task $i$, which is normalized by the validation performance of the former model. This metric offers a precise measure of potential performance degradation when using the mini subset of training data, thereby reflecting the inherent learning difficulty of the task. Notably, to ensure a fair measurement across each task, the ratio between the number of samples in the mini subset and the total number of samples in the full dataset should be kept consistent.

However, training extra models using both the full dataset and the mini subset from each task is necessary to obtain the intra-task difficulty, which will require additional time comparable to the training time of the final model. To alleviate this, we repurpose the models trained for computing inter-task contributions. Specifically, we substitute the model trained on the mini subset from Task $i$ with that trained on mini subsets from all tasks, and replace the model trained solely on the full dataset from Task $i$ with that trained on the full dataset from Task $i$ alongside mini subsets from other tasks. Due to the minimal inter-task contributions of others tasks to Task $i$ when compared to the contribution from Task $i$ to Task $i$ itself, the impact of mini subsets from other tasks on the validation performance for Task $i$ is negligible. Therefore, this approach significantly reduces training time with minimal error. The intra-task difficulty for Task $i$ is calculated as:

where $V_{i}(mini)$ represents the validation performance for Task $i$ of a model trained on mini subsets from all tasks, and $V_{i}(i+mini)$ denotes the validation performance for Task $i$ of a model trained on the full dataset from Task $i$ coupled with mini subsets from other tasks.

Moreover, owing to the varying intra-task difficulties across different tasks, treating each task equally during the training process may result in underfitting of the more challenging tasks, despite the simpler ones being adequately trained or even overfitted. Therefore, we propose a weighting strategy that assigns more weights to tasks with greater intra-task difficulties. The task weights $\bm{\lambda_{D}}$ can be calculated as:

where $\bm{D}$ represents the $N$-dimensional vector of intra-task difficulties for all tasks, and $T$ is the same temperature hyperparameter used in Section 3.2.

After individually establishing the strategies for inter-task contribution balancing and intra-task difficulty balancing, the final step involves integrating them to create the CoTBal algorithm. The algorithm is designed to synergistically leverage the strengths of both two balancing methods, thereby ensuring a more comprehensive and effective multi-task optimization process in visual instruction tuning. The specific task weights $\bm{\lambda_{CoTBal}}$ for comprehensive task balancing can be calculated as:

where $\alpha$ is a hyperparameter that controls the relative influence of inter-task contribution balancing and intra-task difficulty balancing. The training process of CoTBal is summarized in Algorithm 1.

Inspired by Liu et al. (2023a), we incorporate response format instructions into the data to clarify task requirements for the model and employ multiple data processing strategies to reduce training costs and ensure fairness, detailed as follows:
(1) For ShareGPT4V, the data is randomly partitioned into a validation set of 2k and a test set of 2k, with the remainder designated for training.
(2) For all VQA datasets and RefCOCO, data from the same training image are shuffled and merged into a single conversation.
(3) For RefCOCO, training conversations are segmented into parts, each with fewer than 10 turns.
(4) For OCRVQA, 80k conversations are sampled from the training set.
(5) For VQAv2, GQA and OCRVQA, 20k data are sampled from the validation set.
(6) For ShareGPT, invalid conversations are filtered out as Zheng et al. (2023), while long conversations that surpass 2048 tokens are truncated.
The training data sizes and response format instructions for each task are presented in Table 1.

Note that we have not compared traditional gradient aggregation multi-task optimization algorithms. Such algorithms require computing update gradients via backpropagation for each task separately in each iteration, followed by the aggregation of gradients across all tasks. In multi-task visual instruction tuning, the large number of tasks and the massive volume of model parameters make this process impractical and excessively time-consuming.

Table 2 presents the comparative results for multi-task instruction tuning. With the same foundational models and training data, CoTBal achieves the optimal average per-task performance improvement ($\Delta I\%$), alongside maintaining the near-lowest average per-task performance error ($\Delta E\%$). As shown in Figure 2, compared to the most commonly employed EW method, CoTBal significantly enhances the performance on ShareGPT4V, ChartQA and RefCOCO-bbox tasks while keeping competitive performance on other tasks. This validates the effectiveness of our algorithm in terms of both overall performance and the degree of performance imbalance. Specifically, as depicted in Figure 3, CoTBal effectively captures the variances in mutual contributions and inherent learning difficulties among these visual tasks, thereby providing appropriate task weights for the final model training, which maximally mitigates task conflicts.

Besides, we observe that TLA is significantly inferior to EW in both $\Delta I\%$ and $\Delta E\%$. TLA guarantees equality for each task in the final loss computation. However, variations in sample sequence length and data amount across different tasks may implicitly assign inappropriate task weights to the losses of valid tokens. The implicit weight is inversely related to the total number of valid tokens in each task, leading to poorer overall performance and a marked performance imbalance. This indicates that maintaining equality at the token level is more logical than preserving it at the task level, thereby demonstrating the viability of the GTW paradigm in multi-task visual instruction tuning.

As for the compared traditional multi-task optimization algorithms (RLW, DWA and IGBv1), directly applying them to visual instruction tuning yields suboptimal results in both $\Delta I\%$ and $\Delta E\%$. We contend that assigning task weights based on training losses is imprecise, because the fine-tuning losses in large models fails to accurately reflect training progress. This is also the reason why both the inter-task contribution and the intra-task difficulty are quantified by performance-based metrics in our CoTBal algorithm.

As shown in Table 3, we analyze the impact of different training settings on model performance from three aspects: the temperature hyperparameter configuration, the task weighting strategy selection and the calculation approach for intra-task difficulties. The complete ablation results are presented in Appendix A. The compared methods include: EW; CoTBal ($T\!=\!2.0/1.0/0.5$) where the temperature hyperparameter $T$ is set to $2.0$, $1.0$ oder $0.5$; CoTBal ($\bm{\lambda_{one2all}}/\bm{\lambda_{all2one}}/\bm{\lambda_{C}}/\bm{\lambda_{D}}/\bm{\lambda_{CoTBal}}$) where task weights are set as $\bm{\lambda_{one2all}}$, $\bm{\lambda_{all2one}}$, $\bm{\lambda_{C}}$, $\bm{\lambda_{D}}$ oder $\bm{\lambda_{CoTBal}}$; and CoTBal (precise / real Difficulty) where the precise or real calculation approach for the intra-task difficulty is employed. Specifically, the precise calculation approach trains extra models using the full dataset and the mini subset from each task, while the real calculation approach repurposes the models trained for computing inter-task contributions to reduce additional training time.

In terms of the temperature hyperparameter configuration: CoTBal consistently outperforms EW in $\Delta I\%$, maintaining its superiority across all $T$ values and enhancing its advantage as $T$ decreases. The degree of task balancing increases as $T$ decreases, leading to an improved $\Delta I\%$, which demonstrates the efficacy of comprehensive task balancing. Conversely, CoTBal exhibits a slight increase in $\Delta E\%$ as $T$ decreases. When the degree of non-smoothness in task weights becomes excessive, tasks with significantly smaller weights inevitably underperform, resulting in the slight imbalance in performance.

In terms of the task weighting strategy selection: On the one hand, compared to the EW method, CoTBal ($\bm{\lambda_{one2all}}$) enhances $\Delta I\%$ while maintaining $\Delta E\%$ constant, due to its preference for tasks that offer substantial contributions to other tasks. On the other hand, CoTBal ($\bm{\lambda_{all2one}}$) significantly reduces $\Delta E\%$, mitigating the performance imbalance issue by prioritizing tasks that receive minimal benefits from other tasks. CoTBal ($\bm{\lambda_{C}}$) integrates the aforementioned two strategies, achieving more balanced $\Delta I\%$ and $\Delta E\%$. Moreover, CoTBal ($\bm{\lambda_{D}}$) markedly enhances $\Delta I\%$ by assigning more weights to tasks that have greater learning difficulties, yet concurrently exacerbates the performance imbalance issue. Finally, CoTBal ($\bm{\lambda_{CoTBal}}$) integrates all three strategies to maximize overall performance while mitigating the performance imbalance issue.

In terms of the calculation approach for the intra-task difficulty: The precise approach and the real approach exhibit similar levels of performance, with the real one even marginally surpassing the precise one in both $\Delta I\%$ and $\Delta E\%$. When calculating the intra-task difficulty of Task $i$, training with mini subsets from any other tasks has negligible impact on performance in Task $i$, hence our CoTBal algorithm employs the real calculation approach to significantly reduce training time while ensuring performance.