SlowGuess/NeurIPS_2023_OCR huggingface dataset:
First, we observe a clear asymmetry in terms of performance across the  $T = 2$  datasets, even when using uniform weighing  $p_1 = 0.5$ : For instance on Taskonomy (Figure 3), several MTL models outperform SD on the depth prediction task, but none on the semantic segmentation task. Nevertheless, MTL models are very appealing when taking model efficiency into account, as they contain roughly half as many parameters. Secondly, tuning the scalarization weights is crucial in some settings: For instance in DomainNet, training with  $p_1 \in [0.65, 0.75]$  is more advantageous than uniform scalarization. Furthermore, the relative ranking of the weight  $p_1$ , with respect to model performance, does not change significantly across model capacities, for both the Taskonomy and
DomainNet examples. This suggests that optimal weights  $\mathbf{p}^*$  for a given model may be a good search starting point for another architecture of the same family (see Appendix 4).
Nevertheless, the search space for  $\mathbf{p}$  grows exponentially with  $T$ , making the search computationally prohibitive, even for one architecture. In Section 5, we propose a scalable approach to optimize scalarization weights and investigate its performance on DomainNet and CelebA.
(a) DomainNet - real + painting (higher is better)
(b) Taskonomy - normals + depth (lower is better). "shr" refers to the number of shared decoder layers.
Figure 3: MDL/MTL performance when varying tasks weight  $p_1$  for different model capacities. High opacity markers represent models that outperform their respective SDL baseline (black markers) on both tasks/domains. Dashed lines connect models with the same architecture. The marker size is proportional to the number of parameters in each model.
4.2 Conflicting gradients in practice
A widely spread explanation for task interference in the literature is that individual task gradients may point in conflicting directions, hampering training. To investigate this behavior, we measure the percentage of conflicting gradients pair encountered in each epoch, when training with uniform scalarization. Following [64], we define gradients as conflicting if and only if the cosine of their angle is negative. Finally, we provide further details and figures for this section in Appendix 3.
In Figure 4, we first illustrate an asymmetric characteristic of gradient conflicts: A high number of gradient conflicts typically translates to poorly performing models (e.g. early training), but the lower conflicts regime does not correlate well with MDL/MTL performance. In particular, it is common to encounter more gradient conflicts towards the end of training, while the loss steadily decreases. This suggests that, in practice, identifying and removing conflicts at every training iteration may be superfluous, with respect to the compute and memory cost occurred. We can also put this observation in perspective with Theorem 3 of [64], which states that in the case of two tasks, a parameter update of PCGrad leads to a lower loss than the uniform scalarization update if the tasks are conflicting enough; however, this assumption may not hold true across every training iteration.
(a) Per-epoch proportion of gradient conflicts during training
Figure 4: Proportion of conflicting gradients during uniform scalarization training with and without PCGrad on the 40 attribute classification tasks of CelebA while varying learning rate ([5e-4, 5e-3, 1e-2]), model depth ([3, 6, 9]) and width ([0.5, 0.75, 1]). While the overall number of encountered conflicts differs, the trend is consistent across both settings in that the higher number of conflicts encountered towards the end of training does not harm training or final model performance.
(b) Correlation plot with per-epoch average training loss
(c) Correlation plot with per-epoch test metric
% conflicting gradients variance (×10-4)	DomainNet	Taskonomy
Batch size	2.40	n/a
Learning rate	2.25	47.6
Model width	0.40	1.73
Encoder depth	1.09	2.63
Decoder depth	n/a	5.06
(a) We measure the variance across learning rates while keeping other parameters fixed (model capacity, batch size, etc.), then report the median variance across all settings. We repeat this process for each axis of variation.
(b) Proportion of pairwise task conflicts (for a fixed model size and learning rate) across training epochs (median), for DomainNet (6 domains) and Taskonomy (subset of 6 tasks)
Figure 5: (left) Variance of observed gradient conflicts when sweeping over different training hyperparameters (high impact) as well as model capacity (low impact), and (right) illustrating the median proportion of pairwise gradient conflicts as a measure of task affinity
Secondly, we analyze the effect of different factors of training variations on the observed proportion of gradient conflicts in Figure 5(a). On the one hand, hyperparameters directly related to weights updates, such as the learning rate of batch size, have a clear impact. In particular, since memory consumption is often a bottleneck of gradient-based MTO methods, this means that naively decreasing batch size to reduce memory usage may severely impact such methods in practice. On the other hand, model capacity has a lesser effect on the pattern of gradient conflicts, despite influencing task interference as highlighted in Figure 2. Nevertheless, we also observe that the overall magnitude of conflicts encountered for different pairs of tasks reveal interesting task affinity patterns, as illustrated in Figure 5(b): For instance, the quickdraw domain of DomainNet appears as a clear outlier (across all model capacities), which is also the case in terms of MDL performance in Figure 2.
In summary, we observe intriguing properties of gradient conflicts in practice, in particular suggesting that the extra cost of measuring, storing and correcting conflicting gradients at every training iteration can be superfluous. As an alternative, [30] shows that less costly regularization methods (e.g., early stopping) can be competitive with MTO. In this work, we consider an orthogonal direction, by tuning the scalarization/resampling weights hyperparameters to regulate the training speed of the different tasks/domains. However, the computational cost of browsing this large search space becomes a practical caveat as  $T$  grows. In the next section, we leverage a scalable hyperparameter tuning algorithm to tackle this problem and apply it on the DomainNet and CelebA datasets.
5 Population-based training for scalarization weights selection
Tuning the weights  $\mathbf{p}$  with classical parameter search methods, such as grid search or random search, becomes extremely costly when  $T$  increases. To address this high computational demand, we propose to leverage the population-based training (PBT) [24] framework which has been used for efficient hyperparameter search in reinforcement learning [24] and for data augmentation pipelines [23].
Population-based training (PBT). PBT is an evolutionary algorithm for hyperparameter search:  $N$  models are trained in parallel with different starting hyperparameters. Every  $E_{\text{ready}}$  epochs, the models synchronize: The  $Q\%$  worst models in the population are stopped, and their model weights and hyperparameters are replaced by the ones of the  $Q\%$  best models (exploit step); Then, the newly copied hyperparameters are randomly perturbed to reach a new part of the hyperparameter search space (explore step). Finally, training resumes until  $E$  epochs are reached. In other words, PBT enables dynamic exploration of the hyperparameter search space with a fixed computational cost (cost of training  $N$  models + potential overhead from synchronization). A follow-up work, Population-based Bandits (PB2) [46] proposes to leverage Bayesian optimisation [17] to better guide the explore step. In contrast to PBT, PB2 also offers theoretical guarantees on the convergence rate.
Using PBT to tune scalarization weights. PBT tuning relies on three important characteristics that may conflict with the standard scalarization MDL/MTL training pipeline:
- Models are trained with a dynamic schedule of hyperparameters. While this contrasts with standard scalarization in which  $\mathbf{p}$  is fixed, we do not expect this to be an issue as recent work has shown that scalarization with dynamic random weights performs well in practice [35].
- Secondly, models are compared against one another after a few epochs of training ( $E_{\text{ready}}$  epochs) in contrast to e.g. random search where models are usually trained until convergence, or reaching a certain stopping criterion. Consequently, tuning  $E_{\text{ready}}$  can significantly impact the search's stability and outcome, which we further discuss in Appendix 5.
- Finally, models in the population are compared using a single objective (e.g. average task/domain metric) during training. This may be an issue if the task metrics' have widely different ranges and the PBT scheduler may simply learn to favor short term improvement by giving higher weight to tasks with high metrics; While we do not observe this issue in our settings, recent work [15] also proposes a multi-objective variant of PBT which may be better fitted for MDL/MTL applications.
Nevertheless, the key advantage of PBT is its computational efficiency with regard to search space exploration: For a constant cost of training  $N$  models, and some minor overhead related to checkpointing, PBT explores up to  $N(1 + Q \times E_{\text{total}} / E_{\text{ready}})$  possible hyperparameter configurations throughout training of the population.
Table 1: Results of MDL when jointly training all 6 domains of DomainNet for scalarization (uniform and PBT-found weights) and MTO methods. PBT is run with a population size of  $N = 12$  models, such that every  $E_{\text{ready}} = 5$  epochs,  $Q = 25\%$  of the population triggers an exploit/explore step.
DomainNet (ResNet26 with 0.25 width)
average	clipart	infograph	painting	quickdraw	real	sketch
Scalarization
Uniform	46.78 ± 0.10	56.31 ± 0.04	20.46 ± 0.15	40.95 ± 0.45	60.69 ± 0.07	55.64 ± 0.01	46.64 ± 0.32
PBT	48.01 ± 0.08	58.31 ± 0.15	19.45 ± 0.08	41.32 ± 0.11	63.58 ± 0.31	60.43 ± 0.05	45.00 ± 0.26
MTO - Loss-based
Uncertainty [26]	45.12 ± 0.07	59.24 ± 0.09	17.14 ± 0.25	37.75 ± 0.16	59.85 ± 0.16	52.35 ± 0.20	44.42 ± 0.10
IMTL-L [38]	44.22 ± 0.10	58.05 ± 0.22	16.41 ± 0.16	37.53 ± 0.21	59.22 ± 0.29	51.27 ± 0.19	42.83 ± 0.38
MTO - Gradient-based
CAGrad [37]	42.82 ± 0.06	54.08 ± 0.03	18.26 ± 0.04	36.79 ± 0.14	56.46 ± 0.28	49.52 ± 0.15	41.78 ± 0.01
GradDrop [8]	42.52 ± 0.05	53.34 ± 0.03	18.16 ± 0.07	37.25 ± 0.02	55.09 ± 0.13	49.94 ± 0.26	41.35 ± 0.09
PCGrad [64]	42.78 ± 0.14	53.55 ± 0.04	18.29 ± 0.34	37.31 ± 0.38	55.60 ± 0.08	50.41 ± 0.13	41.52 ± 0.67
DomainNet (ResNet26 with original width)
average	clipart	infograph	painting	quickdraw	real	sketch
Scalarization
Uniform	48.83 ± 0.08	58.69 ± 0.05	21.58 ± 0.39	42.83 ± 0.03	62.51 ± 0.21	58.31 ± 0.06	49.05 ± 0.15
PBT	49.27 ± 0.12	58.41 ± 0.50	19.30 ± 0.19	44.53 ± 0.16	63.08 ± 0.42	60.08 ± 0.16	50.23 ± 0.01
MTO - Loss-based
Uncertainty [26]	46.96 ± 0.10	60.71 ± 0.36	18.74 ± 0.03	40.22 ± 0.40	60.92 ± 0.04	54.37 ± 0.12	46.79 ± 0.23
IMTL-L [38]	46.04 ± 0.21	59.76 ± 0.78	18.21 ± 0.15	39.12 ± 0.76	60.24 ± 0.39	53.06 ± 0.38	45.87 ± 0.29
MTO - Gradient-based
CAGrad [37]	44.91 ± 0.18	56.56 ± 0.38	19.63 ± 0.32	38.84 ± 0.47	58.06 ± 0.41	51.80 ± 0.58	44.58 ± 0.45
GradDrop [8]	45.15 ± 0.08	56.22 ± 0.43	19.89 ± 0.16	39.70 ± 0.03	57.55 ± 0.12	52.81 ± 0.04	44.76 ± 0.18
PCGrad [64]	44.96 ± 0.14	55.79 ± 0.24	19.82 ± 0.20	39.65 ± 0.29	57.30 ± 0.30	52.57 ± 0.11	44.65 ± 0.65
In Table 1 and Table 2, we report results for searching optimal scalarization weights  $\mathbf{p}^*$  when training for all 6 domains of DomainNet and for 8 tasks (attribute subsets) of CelebA. For PBT results, we first run the search algorithm using the implementation from Raytune [33]. We use  $70\%$  of the training set for training, and use the remaining  $30\%$  to rank models in the population by measuring their average accuracy on this set. Once the search is done, we retrain a model on the full training set using the scalarization weights found by PBT. For comparison, we also report results for uniform scalarization and MTO methods, using the implementation from [31]. All final models are trained for three different learning rates and the best metric is reported, averaged across 2 random seeds.
On the DomainNet example, we observe that the scalarization weights found by PBT outperforms all methods, confirming our insights that tuning weights  $\mathbf{p}$  can further enhance scalarization. Note that MTO methods were not designed or employed for multi-domain settings such as DomainNet, which may explain why gradient-based MTO methods all underperform their loss-based counterparts in the results of Table 1, while they do exhibit good performance on CelebA MTL. In fact, on CelebA, while PB2 search reaches the highest overall average metric, we find that results across tasks exhibit more variance, with PB2 and CAGrad yielding the best, comparable, performance. Overall, the very narrow differences in accuracy makes it difficult to highlight one specific method, which also raises the issue on whether CelebA is a robust enough MTL benchmark. Nevertheless, both experiments
show that scalarization can outperform more complex optimization methods when its weights are tuned properly, which can be done efficiently using scalable hyperparameter search methods like PBT or PB2. We discuss further insights and compute details in Appendix 5.
Table 2: Results of MTL when training on all 8 tasks (subset of attributes) of CelebA defined in Appendix 1. For PBT and PB2 we use slightly different parameters than DomainNet to account for the fact that CelebA contains more tasks, and hence has a larger search space: All PBT runs use a population size of  $N = 12$  models, such that every  $E_{\text{ready}} = 3$  epochs,  $Q = 40\%$  of the population triggers an exploit/explore step. For PB2 runs we use a population size of  $N = 8$  and otherwise the same  $Q$  and  $E_{\text{ready}}$  hyperparameters. For the sake of space, we omit standard deviations in CelebA in the main text (in the range  $1e^{-4}$ ), and only report results for the four best performing MTO baselines.
ViT-S/4, 6 layers, full width
average	age	clothes	face structure	facial hair	gender	hair color	hair style	mouth
Scalarization
Uniform	91.23	86.70	92.79	85.16	95.45	98.10	92.99	91.68	86.95
PBT	91.21	86.79	92.77	85.18	95.42	98.02	92.92	91.66	86.95
PB2	91.28	86.96	92.87	85.17	95.47	98.10	92.90	91.81	86.94
MTO
IMTL-L [38]	91.19	86.60	92.80	85.18	95.46	97.97	92.93	91.66	86.94
CAGrad [37]	91.27	86.92	92.89	85.11	95.49	98.17	92.93	91.74	86.89
GradDrop [8]	91.27	87.05	92.73	85.27	95.48	98.11	92.93	91.64	86.94
PCGrad [64]	91.18	86.71	92.74	85.09	95.37	98.12	92.90	91.66	86.84
ViT-S/4, 9 layers, full width
average	age	clothes	face structure	facial hair	gender	hair color	hair style	mouth
Scalarization
Uniform	91.17	87.33	92.50	85.10	95.45	97.93	92.85	91.39	86.80
PBT	91.15	87.43	92.51	85.18	95.46	97.78	92.51	91.36	87.00
PB2	91.25	86.83	92.85	85.12	95.49	98.19	92.90	91.78	86.81
MTO
IMTL-L [38]	91.16	87.40	92.44	85.09	95.44	97.92	92.87	91.39	86.75
CAGrad [37]	91.22	87.37	92.66	85.13	95.40	97.92	92.92	91.61	86.74
GradDrop [8]	91.07	87.41	92.36	85.01	95.41	97.73	92.75	91.21	86.70
PCGrad [64]	91.14	87.53	92.42	85.00	95.38	97.87	92.84	91.37	86.68
Since we are building off linear scalarization, our study suffers from the same issue: It has been shown, for instance in [5], that scalarization only finds solutions on the convex parts of the Pareto front. Same as previous studies [63, 52], we do not observe this to be an issue in practice, but there may be some cases where scalarization can simply not perform as well as more advanced MTO methods. In addition, our results and observations relies on an experimental study. While we do attempt to experiment over a diverse set of benchmarks and model capacities to get as many data points as possible, we can not guarantee the generality of our claims across all MDL/MTL settings.
This work presents a comprehensive evaluation of scalarization's effectiveness in multi-domain and multi-task learning, spanning diverse model capacities and dataset sizes. Our analysis reveals that larger-capacity models often benefit more from joint learning across diverse settings. In addition, tuning scalarization weights is key to reach optimal performance at inference and improve the MTL/MDL model generalization. Nevertheless, given a specific set of specific tasks/domains, the optimal weights are rather robust to changes in model capacities within the same architecture family. We then investigate the impact of model capacity on gradient conflicts observed during training and observe low correlation with MTL/MDL performance. Finally, to tackle the large search space of tuning scalarization weights, we propose to leverage population-based training as a scalable, efficient method for tuning scalarization weights as the number of tasks/domains increases.
[53] Sebastian Ruder, Joachim Bingel, Isabelle Augenstein, and Anders Søgaard. Latent multi-task architecture learning. In Conference on Artificial Intelligence (AAAI), 2019. 2
[54] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115:211-252, 2015. 15
[55] Ozan Sener and Vladlen Koltun. Multi-task learning as multi-objective optimization. In Conference on Neural Information Processing Systems (NeurIPS), 2018. 2, 4
[56] Changjian Shui, Mahdieh Abbasi, Louis-Émile Robitaille, Boyu Wang, and Christian Gagné. A principled approach for learning task similarity in multitask learning. In International Joint Conferences on Artificial Intelligence (IJCAI), volume abs/1903.09109, 2019. 1
[57] Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus. Indoor segmentation and support inference from rgbd images. European Conference on Computer Vision (ECCV), 7576: 746-760, 2012. 5
[58] Trevor Scott Standley, Amir Roshan Zamir, Dawn Chen, Leonidas J. Guibas, Jitendra Malik, and Silvio Savarese. Which tasks should be learned together in multi-task learning? *ArXiv*, abs/1905.07553, 2019. 1, 3, 4, 5, 16, 17
[59] Gjorgji Strezoski, Nanne van Noord, and Marcel Worrring. Many task learning with task routing. In International Conference on Computer Vision (ICCV), pages 1375–1384, 2019. 2
[60] Mihai Suteu and Yike Guo. Regularizing deep multi-task networks using orthogonal gradients. arXiv preprint arXiv:1912.06844, 2019. 2
[61] Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Marc Proesmans, Dengxin Dai, and Luc Van Gool. Multi-task learning for dense prediction tasks: A survey. IEEE transactions on pattern analysis and machine intelligence, 44(7):3614-3633, 2021. 2
[62] Zirui Wang, Yulia Tsvetkov, Orhan First, and Yuan Cao. Gradient vaccine: Investigating and improving multi-task optimization in massively multilingual models. In International Conference on Learning Representations (ICLR), 2021. 2, 18
[63] Derrick Xin, B. Ghorbani, Ankush Garg, Orhan First, and Justin Gilmer. Do current multi-task optimization methods in deep learning even help? In Conference on Neural Information Processing Systems (NeurIPS), 2022. 1, 2, 4, 10, 18
[64] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. In Conference on Neural Information Processing Systems (NeurIPS), 2020. 1, 2, 4, 7, 18
[65] Amir Roshan Zamir, Alexander Sax, Bokui (William) Shen, Leonidas J. Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. In Conference on Computer Vision and Pattern Recognition (CVPR), 2018. 1, 4, 5, 16
[66] Youshan Zhang. A survey of unsupervised domain adaptation for visual recognition. ArXiv, abs/2112.06745, 2021. 3
Supplemental Material to "Scalarization for Multi-Task and Multi-Domain Learning at Scale"
1 Experimental settings and training hyperparameters
1.1 Multi-Domain
We refer to benchmarks as "multi-domain" when they contain multiple input visual domains with a shared set of output classes (i.e.,  $\forall i \neq j, X_i \neq X_j$  and  $Y_i = Y_j$ ).
CIFAR-10 and STL-10. CIFAR-10 [29] is a classical benchmark for image classification containing 50k training samples uniformly distributed across 10 classes. STL-10 [9] is a semi-supervised dataset which was designed to resemble CIFAR-10. Specifically, we only use the 5000 annotated images in STL-10, which are also uniformly distributed across the same 10 classes as CIFAR. In STL-10, the images themselves are from the ImageNet [54] dataset, and cropped/resized to 96 pixels. We further resize them to 32 pixels to align with CIFAR. In summary, the key difficulties are (i) the input distribution shift between the two datasets and (ii) the high imbalance in data availability.
We use a vision transformer backbone (ViT-S) optimized for small-scale datasets [18] (compared to the original ViT-S, this backbone contains smaller patch sizes, fewer transformer layers and narrower embeddings, but a higher number of heads). To control model capacity, we vary the depth (number of transformer layers) in  $\{3,6,9\}$  and the width (token dimension) in  $\{48,96,144,192\}$ , Finally, we train each model from scratch on a single NvidiaV100 GPU with a batch size of 256 images for 300 epochs (including 30 epochs of linear learning rate warmup), using a learning rate of 0.001 and weight decay of 0.05 with the AdamW optimizer and cosine learning rate decay.
DomainNet. DomainNet [47] is a classification dataset of 6 visual domains annotated for 345 classes, for a total of roughly 410k training samples. DomainNet was initially introduced for the problem of multi-source domain adaption, in which one or more of the domains does not have training annotations; the key difficulty is thus to learn representations that are aligned across domains. In contrast to the CIFAR+STL example, DomainNet exhibits distribution shifts across both the input domains and output classes, as visualized in Figure A.
Figure A: Illustrating the data imbalance in DomainNet with a contour plot of the number of samples per class and domains in DomainNet. Each of the corners of the hexagon represents one of the six domains in DomainNet, and the lines (levels of the contour plot) represent the number of samples, drawn every 15 classes. For comparison, a uniformly distributed dataset would yield perfect hexagons.
Following previous literature [47, 32], we use a ResNet-101 as our main backbone. There is no domain-dependent layer in the architecture: the final classifier layer is shared across all domains. We first perform a training sweep over the largest domain (real) to select the best-performing learning rate from  $\{0.3, 0.03, 0.003, 0.0003\}$  and the number of epochs from  $\{30, 60, 90\}$ . Following these results, we use a learning rate of 0.03 and train for 30 epochs with a batch size of 512 in subsequent experiments. We train with the AdamW optimizer with a weight decay of  $1e - 4$ . We also apply linear learning rate warm-up during the first five training epochs and use cosine schedule learning rate decay for the rest of the training. Finally, to control model capacity, we vary the depth (backbone) in  $\{\mathrm{ResNet-26}, \mathrm{ResNet-50}, \mathrm{ResNet-101}\}$  and the width factor.