SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Amélie Royer, Tijmen Blankevoort, Babak Ehteshami Bejnordi
{aroyer, tijmen, behtesha}@qti=qualcomm.com
Training a single model on multiple input domains and/or output tasks allows for compressing information from multiple sources into a unified backbone hence improves model efficiency. It also enables potential positive knowledge transfer across tasks/domains, leading to improved accuracy and data-efficient training. However, optimizing such networks is a challenge, in particular due to discrepancies between the different tasks or domains: Despite several hypotheses and solutions proposed over the years, recent work has shown that uniform scalarization training, i.e., simply minimizing the average of the task losses, yields on-par performance with more costly SotA optimization methods. This raises the issue of how well we understand the training dynamics of multi-task and multi-domain networks. In this work, we first devise a large-scale unified analysis of multi-domain and multi-task learning to better understand the dynamics of scalarization across varied task/domain combinations and model sizes. Following these insights, we then propose to leverage population-based training to efficiently search for the optimal scalarization weights when dealing with a large number of tasks or domains.
Learning a unified architecture that can handle multiple domains and/or tasks offers the potential for improved computational efficiency, better generalization, and reduced training data requirements. However, training such models proves a challenge: In particular, when the tasks or domains are dissimilar, practitioners often observe that the learning of one task or domain interferes with the learning of others. This phenomenon is commonly known as interference or negative transfer. How to best resolve this interference is an open problem spanning numerous bodies of literature, such as multi-task learning, domain adaptation and generalization, or multi-modal learning.
In fact, there is currently no unanimous understanding of what causes, or how to predict, task interference in practice: For instance, classical generalization bounds on training from multiple data sources [3, 48, 56, 44] involve a measure of distance between the respective tasks/domains distributions, which can be intuitively thought of as a measure of interference. In more practical applications, task clustering methods [65, 58, 16] propose various metrics to measure task affinity, and train tasks with low affinity separately, with separate backbones, thereby circumventing the interference issue at the architecture level. Finally, in the multi-task optimization (MTO) literature, the prevailing view is that gradients from different tasks may point in conflicting directions and cause interference, which has led to a line of work on reducing gradient conflicts by normalizing gradients or losses statistics [37, 7, 25, 64]. While MTO have established a new state-of-the-art for training multi-task models in the past few years, recent work [63, 30] shows that, surprisingly, minimizing
a simple average of the tasks/domains losses, also known as scalarization, can yield performance trade-off points on the same Pareto front as more costly MTO methods.
Nevertheless, there still remains some unexplored areas which we aim to shed lights on in this paper. First, most previous experiments have only been conducted on multi-task settings, on benchmarks with either few tasks or few training samples (e.g. multi-MNIST, NYU-v2), and for a fixed architecture. In particular, the link between model capacity and the observed MTL performance is often overlooked. Secondly, it is not clear how the choice of scalarization weights links to and/or impacts the hypothesis of gradient conflicts as a key underlying signal of task interference. Finally, scalarization can become prohibitively expensive for a large number of tasks as the search space for the optimal loss weights grows exponentially; This raises the issue of how to efficiently browse this search space. In this work, we investigate these questions to further motivate scalarization as a simple and scalable training scheme for multi-task and multi-domain problems.
Our key contributions. We perform a large-scale analysis of scalarization for both multi-task (MTL) and multi-domain learning (MDL). We cover a wide range of model capacities, datasets with varying sizes, and different task/domain combinations. Our key conclusions are as follows:
- (C1) When compared to a model trained on each task/domain individually, MDL/MTL performance tends to improve for larger model sizes, showing that the benefits of MTL/MDL frameworks should be put into perspective with respect to the backbone architecture capacity.
- (C2) Tuning the scalarization weights for a specific tasks/domains combination is crucial to obtaining the optimal MTL/MDL performance in settings with high imbalance; Nevertheless, the relative performance of different scalarization weights is often consistent across model capacities inside a family architecture. This suggests that searching for optimal scalarization weights for a lower model depth/width is also relevant for the full model size, while using less compute.
- (C3) Gradients conflicts between tasks/domains naturally occur during MTL training: They behave differently across layers and learning rates, but are scarcely impacted by model capacity and scalarization weights choice. These observations give new insights into the practical implications of avoiding conflicting gradients throughout training.
- (C4) We leverage fast hyperparameter search methods such as population-based training [24] to efficiently browse the search space of scalarization weights as the number of tasks grows.
Multi-Task Optimization (MTO) and scalarization. MTO methods aim to improve MTL training by balancing the training dynamics of different tasks. Prior research can be split into two categories: On the one hand, loss-based methods propose to align the task losses magnitudes by rescaling them through various criteria, e.g., task uncertainty [26], task difficulty [21], random loss weighting [34], or gradients statistics [7, 14, 60]. On the other hand, gradient-based methods directly act on the per-task gradients rather than the losses. For instance, [55, 13] tackle MTL as a multi-objective optimization problem using the Multiple Gradient Descent Algorithm to directly locate Pareto-optimal solutions. Another major line of work considers conflicting gradient direction to be the main cause of task interference. Consequently, these works [64, 37, 38, 62, 8, 25] aim to mitigate conflicts across per-task gradients. For instance, PCGrad [64] suggests projecting each task's gradient onto the normal plane of other gradients to suppress conflicting directions, while GradDrop [8] ensures that all gradient updates are pure in sign by randomly masking all positive or negative gradient values during training. Unfortunately, gradient-based techniques require per-task gradients, leading to substantial memory usage and increased runtime. Furthermore, recent work [63, 30] has shown that these methods often perform on-par with the less costly approach of directly optimizing the average of the task losses, also known as uniform scalarization. Building off these insights, we further investigate the benefits of scalarization and the practical dynamics of gradient conflicts in this work.
Architectures for MTL. Orthogonal to our work, another line of research focuses on designing multi-task architectures that mitigate task interference by optimizing the allocation of shared versus task-specific parameters. Hard parameter sharing methods [27, 26, 6, 59, 61] build off a shared backbone and carefully design task-specific decoder heads optimized for the tasks of interest. In contrast, soft parameter sharing methods [45, 53, 20, 39] jointly train a shared backbone while learning sparse binary masks specific to each task/domain which capture the parameter sharing pattern across tasks. To select which tasks should share parameters or not, a prominent line of work
focuses on defining measures of task affinity [2, 58, 22, 16]: Tasks with high affinity are jointly trained, while low affinity ones use different backbones, in the hope of reducing potential interference by fully separating conflicting tasks. This results in multi-branch architectures where each branch handles a specific subset of tasks with high affinities. The scope of our study is complementary to this line of work, as we focus on how to best train a given set of tasks, without changing the architecture.
Multi-Domain Learning (MDL). MDL refers to the design and training of a single unified architecture that can operate on data sources from different domains. Domain adaptation and generalization methods often aim to align the learned representations of the different domains [19, 4, 66]; These are in particular targeting the setting where one of the domain lacks supervisory signals. In the fully-supervised scenario, several methods have been proposed to use a unified backbone, which captures common features across all tasks, while allowing the model to learn domain-specific information via lightweight adapter modules [49, 51, 50]. Similar to soft parameter sharing method, these methods bypass task interference by keeping the shared backbone frozen and only training domain-specific modules. Alternatively, a common practice when training with multiple input datasets is to use over- or undersampling techniques [1, 28, 43], in particular to handle class imbalance. As we describe later, resampling can be seen as the MDL counterpart of scalarization in MTL. Both enable a simple and lightweight training scheme, yet finding the optimal resampling/scalarizatoin weights is nontrivial. This motivates us to empirically verify our findings in the MDL as well as MTL setting.
3 Motivation and experimental setting for MTL/MDL
We start by describing the MTL/MDL setup: Given  $T$  supervised datasets with respective inputs  $X_{t}$  and outputs  $Y_{t}$ , our goal is to find model parameters  $\theta^{*}$  that minimize the datasets' respective losses:  $(\mathcal{L}_1(X_1,Y_1),\dots,\mathcal{L}_T(X_T,Y_T))$ . In practice, this can be achieved either by solving a multi-objective problem on the task losses [13, 42] or, more commonly, by minimizing a (possibly weighted) average of the losses [27]. In both cases, most training methods for MTL/MDL can be phrased as updating the model parameters  $\theta$  using a weighted average of the individual tasks' gradients:
$$
\theta_ {i + 1} = \theta_ {i} + \eta \sum_ {t = 1} ^ {T} p _ {t} ^ {i} \mathbb {E} _ {\left(x _ {t}, y _ {t}\right) \sim q \left(X _ {t}, Y _ {t}\right)} \nabla_ {\theta_ {i}} \mathcal {L} _ {t} \left(x _ {t}, y _ {t}\right), \text {w h e r e} \eta \text {i s t h e l e a n i n g r a t e} \tag {1}
$$
where  $\mathbf{p}(\mathbf{i}) = (p_1^i,\dots ,p_T^i)$  captures the dynamic importance scaling weights for each dataset at timestep  $i$ , and  $q(X_{t},Y_{t})$  is the underlying distribution the  $t$ -th dataset is drawn from. We also follow the common assumption that  $\mathbf{p}$  is a probability distribution, i.e.,  $\sum_{t}p_{t} = 1$  and  $\forall t$ ,  $p_t\in \{0,1\}$ .
We further distinguish between (i) multi-task learning (MTL) where every input is annotated for every task (i.e.,  $\forall i, j, X_{i} = X_{j}$ ), and (ii) multi-domain learning (MDL) where the goal is to solve a common output task across multiple input domains or modalities ( $\forall i, j, Y_{i} = Y_{j}$ ). In the MDL setting, we can also rephrase (1) as resampling the data distributions according to  $\mathbf{p}$ . Both formalisms are equivalent, but resampling often performs better in practice when using stochastic gradient methods [1]:
$$
\underbrace {\sum_ {t = 1} ^ {T} p _ {t} \mathbb {E} _ {q \left(X _ {t} , Y _ {t}\right)} f \left(x _ {t} , y _ {t}\right)} _ {\text {R e w e i g h n f o r m a l i s m o f (1)}} = \underbrace {\mathbb {E} _ {q ^ {\prime} (X , Y)} f (x , y)} _ {\text {R e s a m p l i n g f o r m a l i s m}} \text {w h e r e} \underbrace {q ^ {\prime} (x , y) = \sum_ {t = 1} ^ {T} \mathbb {1} _ {x \in X _ {t}} q (x , y) p _ {t}} _ {\text {R e s a m p l i n g d i s t r i b u t i o n}} \tag {2}
$$
Current state-of-the-art methods for training objectives such as (1) can be roughly organized into three categories, based on how the datasets' importance weights  $\mathbf{p}$  are defined:
- Scalarization defines the weights  $p_t$  as constant hyperparameters: The update rule of (1) reduces to computing the gradient of the weighted average loss  $\nabla_{\theta}(\sum_t p_t \mathcal{L}_t)$ . In the MDL formalism of (2), scalarization is similar to classical oversampling/undersampling, where  $p_t$  is defined either as a scalar (per-domain) or a vector (to handle both per-class and per-domain biases). In both settings, the key difficulty lies in tuning the weights  $\mathbf{p}$ , as the search space grows exponentially with  $T$ .
- Loss-based adaptive methods [38, 26] dynamically compute  $p_t$  for every batch, aiming to uniformize training dynamics across tasks by rescaling the losses. For instance, IMTL-L [38] dynamically reweighs the losses such that they all have the same magnitude.
- Gradient-based adaptive methods [7, 64, 38, 25, 8, 37, 13, 55, 36, 42] also dynamically compute weights for every batch. While this line of work usually outperforms loss-based methods, they also incur a higher compute and memory training cost, as they require  $T$  backward passes to obtain each individual dataset gradient, which in turn need to be stored in memory.
While gradient-based approaches are generally considered SotA across MTO methods, they present certain practical challenges: They come with a higher computational and memory cost, which increases with  $T$ , as illustrated in Figure 1(b). The implicit weights  $\mathbf{p}_t$  computed by these methods are also hard to extract and transfer to other training runs since they are tied to other hyperparameters (training length, batch size, gradient clipping, etc.). In addition, recent work [63, 30, 52] has shown that simple scalarization with uniform weights actually often performs on-par with both loss- and gradient-based methods. Furthermore, scalarization is highly practical: It does not incur additional costs during training, and scalarization weights are easily interpretable as a measure of "task importance" in the optimization problem. However, some experiments in [58, 52] also suggest that the benefits of uniform scalarization MTL are impacted by model capacity in some settings.
Motivated by these insights, we aim to better harness scalarization for scalable and practical MTL/MDL training. In particular, we focus on the following two remaining gray areas: First, we perform a large-scale analysis investigating the underlying effects of model capacity on MTL and MDL  $(C1)$  in Section 4. We then discuss the impact of model capacity on the choice of optimal scalarization weights  $(C2)$  in Section 4.1 and on the presence of conflicting gradients  $(C3)$  in Section 4.2. Secondly, browsing the search space of  $\mathbf{p}$  to tune the scalarization weights becomes increasingly costly as the number of tasks  $T$  grows. As an alternative to the usual grid- or random-search approaches, we propose to leverage population-based training as a cost-efficient parameter search for scalarization, and report our conclusions  $(C4)$  in Section 5.
MDL	MTL
Dataset	CIFAR+STL [29, 9]	DomainNet [47]	CelebA [40]	Taskonomy [65] (tiny split)
# tasks or domains, T	2	6	8 or 40	7
Training size, |X|	55k	~ 410k	~ 162k	~ 275k
Backbone	ViT-S/4	ResNet-101	ViT-S/4	ResNet-26 encoder/decoder
Depth sweep	{3, 6, 9}	{r26, r50, r101}	{3, 6, 9}	{4, 2, 0} shared decoder layers
Width sweep	{0.25, 0.5, 0.75, 1.0}	{0.25, 0.5, 1.0}	{0.25, 0.5, 0.75, 1.0}	{0.5, 1.0}
(a) Summary of our experimental setting
(b) Memory usage of various multi-task optimization (MTO) methods
Figure 1: Table (a) summarizes the experimental setups used throughout the paper. Figure (b) reports profiling results for popular multi-task optimization (MTO) methods in a small-scale (ResNet18, batch size 16) and large-scale (ResNet50, size 128) setting, for 224px inputs: In practice, we find that the main bottleneck for gradient-based methods is their high memory usage. Consequently, training with reasonable batch sizes requires either high parallelism (compute demand) or gradient accumulation (slower runs). In contrast, loss-based methods have much better-constrained costs. However, they usually underperform their gradient-based counterpart in the literature.
3.3 Experimental setting
Here we briefly describe the experimental setup that is used throughout the paper. We aim to cover a wide range of model sizes, datasets with varying sizes, and different task/domain combinations for MDL and MTL. A summary is given in Figure 1(a) and further details can be found in Appendix 1.
For MDL, we first consider a two-domain example composed of CIFAR10 [29] and STL10 [9] with a ViT-S backbone optimized for small datasets [18]. We then expand the results to the DomainNet
benchmark [47] containing 345 output classes and 6 input domains. We use ResNet as our main backbone, following previous works [47, 32]. In both cases, the backbone parameters are fully shared across all domains; We use cross-entropy as training loss and top-1 accuracy as our main metric.
For MTL, we use CelebA [40] with the same ViT-S backbone as for CIFAR/STL; We split the 40 attributes into semantically-coherent groups to form 8 output classification tasks, as described in the appendix. The transformer backbone is fully shared across tasks, while the last linear classification layer is task-specific. Then, we experiment on a larger MTL setting for dense prediction: Most traditional benchmarks, such as Cityscapes [10] and NYU-v2 [57], are rather small, as they only contain 2-3 dense tasks (segmentation, normals, depth) and  $\sim 5\mathrm{k}$  fully annotated images. Instead, we perform our analysis on the challenging Taskonomy dataset [65]. To keep experiments scalable, we filter the 26 tasks in Taskonomy down to 7 (e.g. by removing self-supervised tasks or clearly overlapping ones such as 2D edges and 3D edges). For training, we follow the setup described in [58]: We normalize every dense annotation to have zero mean and unit variance across the training set, and use the  $L_{1}$  loss as our training loss and common metric for all tasks. The backbone is composed of a (shared) ResNet encoder, followed by task-specific decoders built with upsampling operations and 1x1 convolutions. To control model capacity, we vary  $(i)$  the width of the model,  $(ii)$  the depth of the encoder, as well as  $(iii)$  the number of layers in the decoder(s) which are shared across tasks.
Finally, in all settings, we sweep over the domains/tasks' weights under the common assumption that  $p$  is a probability vector (i.e.  $\forall t, p_t \in \{0,1\}$  and  $\sum_{t}p_{t} = 1$ ). All reported results are averaged across 2 random seeds for DomainNet and Taskonomy, and 3 for the smaller benchmarks. Unless stated, every experiment is conducted on a single NVIDIA V100 GPU.
4 Benefits of MDL/MTL under the lens of model capacity
We first investigate the behavior of scalarization for MDL/MTL training with two tasks, while varying model capacity and for different tasks/domains combinations. To quantify the benefit or harm of joint training with multiple datasets, we compare the model performance on each dataset with that of a same-sized model, trained on a single dataset at once. We refer to this baseline as "SD" (single dataset). Since we use SD as a reference point, we ensure that its training pipeline is well tuned: We sweep over the hyperparameters (learning rate, weight decay, number of training steps, etc) of the baseline SD and keep the same values for training the MDL/MTL models, while varying the task weights  $(p_1,p_2 = 1 - p_1)$  from  $p_1 = 0$  (SD baseline of the first dataset) to  $p_1 = 1$  (SD baseline of the second dataset). Further details on the training hyperparameters can be found in Appendix 1. Finally, we summarize our analysis results in Figure 2: For each task/domain pair, we first plot the weight  $p_1^*$  yielding the best accuracy averaged across both tasks/domains, then we report the accuracy difference between the corresponding MDL/MTL model and the associated SD baseline.
Impact of model capacity. We first observe that MDL/MTL performance greatly varies across model capacities, and tends to increase with model size (C1): In some cases, the trend even inverts from MDL/MTL underperforming to outperforming SD (e.g. DomainNet's real + sketch pairing).
Selecting optimal importance weights  $\mathbf{p}$ . Secondly, the best-performing  $\mathbf{p}$  at inference are rarely the uniform  $p_1 = p_2 = 0.5$ , even in the Taskonomy scenario where the uniform training objective is identical to the average of task metrics which we aim to maximize at inference. We further discuss the effects of tuning scalarization weights for the accuracy/efficiency trade-off (C2) in Section 4.1.
MTL/MDL primarily improves generalization. As we show in Appendix 2), even when MTL/MDL outperforms the SD baseline at inference, it is generally not the case at training time. In other words, tuning the respective tasks/domains weights in MTL/MDL can be seen as a regularization technique that improves generalization over the SD baselines by balancing the two datasets' training objectives. In particular, in the MDL setting, this insight is also closely connected to data augmentations: Each new domain can be seen as a non-parametric and non-trivial augmentation function, while the associated resampling weight  $p_t$  is the probability of applying the augmentation on each sample.
Interestingly, recent work [30] suggests that many MTO methods which aim to avoid gradient conflicts can also be interpreted as regularization techniques, and that simpler tricks such as early stopping are often competitive. To better understand the link between these two insights, we analyze the natural emergence of conflicting gradients during MDL/MTL training, and how they are impacted by model capacity (C3) in Section 4.2.
(a) Analysis of MDL when pairing natural images (real domain) with each of the five remaining DomainNet domains (left), and on the CIFAR-10 + STL-10 toy example (right)
(b) Analysis of MTL on CelebA (left) and Taskonomy-tiny (right) for different task pairs and model sizes.
Figure 2: Performance of scalarization for MDL/MTL relative to SD under different model capacities; Each column corresponds to a different task/domain pair  $(T = 2)$ . The first row of each plot contains a heatmap of the best performing scalarization weights  $p^*$  wrt. to the average test accuracy on both tasks. Each of the following rows contains the difference in metrics between the MTL/MDL model and its counterpart SD baseline, where green indicates positive changes and purple, a negative one. Note that the colormaps' ranges are defined per row, and visualized as color bars at the beginning of each row. We observe the general trend that the performance improvement of MDL, relative to the corresponding SD baseline, tends to increase with model capacity.
4.1 Selecting optimal scalarization weights  $\mathbf{p}^*$
So far, we have discussed MDL/MTL improvement over the SD baseline when selecting the weights  $p^*$  that maximize average accuracy. However, in real-world settings, many criteria come into play when evaluating MTL/MDL models: For instance, one task may be more critical than others hence the model should be evaluated via weighted accuracy. Second, even if the MTL/MDL model underperforms the corresponding SD baseline at the same model capacity, it may still improve the efficiency-to-accuracy trade-off when both tasks need to be solved at once. To give a clearer overview of how these considerations affect the choice of weights  $\mathbf{p}$ , we report results for all scalarization weights in Figure 3; We also highlight points where MTL/MDL outperforms SD on both domains/tasks, and represent each model's number of parameters as the marker size.