SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Multi-domain, resampling and training length In the multi-domain setting, scalarization weights become resampling probabilities for each dataset, as shown in (2). Consequently, the notion of "epoch" is hard to define compare to the standard mono-dataset setting. To resolve this, we always define epochs with respect to one of the domains. For instance, in the CIFAR+STL case, we use STL as our reference. Therefore, "one epoch" translates to seeing as many samples as in the original STL dataset (5000) using the current batch size (256), i.e. roughly 20 training steps. In the DomainNet case, we define epochs relatively to the real domain. This definition has the advantage of not being impacted by the sampling weights  $\mathbf{p}$ ; In particular, this means that both the MDL models and the single dataset (SD) baselines are trained for as many training steps, and see the same amount of training samples, only sampled from different data distributions.
1.2 Multi-task
We define multi-task benchmarks as datasets where every image is fully annotated for multiple output tasks (i.e.,  $\forall i \neq j, X_i = X_j$  and  $Y_i \neq Y_j$ ). This setting is particularly popular for scene understanding problems where every scene is labelled with multiple dense predictions (e.g. depth, normals, segmentation mas, edges, etc.)
Celeba. CelebA is a binary attribute classification dataset containing 40 attributes and roughly 162k training images. To turn CelebA into a multi-task problem, it is common to consider each attribute as a binary classification task: More specifically, we use a fully shared backbone with a final linear layer of 40 outputs, outputting logits for every task. The model is then trained using 40 binary cross-entropy losses, one for each attribute. To make our comparative analysis more scalable, we define several tasks as subsets of attributes, grouped based on semantic similarity (e.g. all hair colors are in the same subgroup). The 8 resulting subsets of attributes are described in Table A. In the scalarization setting, this simply means that some of the attributes share the same importance weight.
As a backbone, we use the same ViT-S/4 based architecture as for CIFAR/STL. We train for 50 epochs with 5 epochs of learning rate warmup. We use a learning rate of 0.0005 with cosine schedule decay and train with the AdamW optimizer with a weight decay of 0.05. We use input images of size 32 (with tokens of size 4), a batch size of 256, and RandAugment data augmentations.
Table A: The eight tasks defined as subsets from CelebA attributes used in our main analysis. Attributes in the same subset share a common importance weight  $p$
Hair color	Hairstyle	Facial Hair	Mouth	Clothes	Face Structure	Gender	Age
Black Hair	Bald	5' o'Clock Shadow	Big Lips	Eyeglasses	Big Nose	Male	Young
Blond Hair	Bangs	Mustache	Mouth Slightly Open	Heavy Makeup	Chubby
Brown Hair	Receding Hairline	No Beard	Smiling	Wearing Earrings	Double Chin
Gray Hair	Sideburns	Goatee	Wearing Lipstick	Wearing Hat	High Cheekbones
Straight Hair	Wearing Necklace	Oval Face
Wavy Hair	Wearing Necktie	Pointy Noise
Taskonomy. Taskonomy [65] is a large dataset containing a variety of dense prediction tasks for indoor scenes. We use the tiny split of Taskonomy which contains roughly 275k images. Taskonomy was originally introduced for the problem of task clustering: The original work [65] proposes a task affinity metric to define a taxonomy of tasks. This taxonomy structure is then used to determine which tasks should be trained from scratch and which tasks could benefit from others via transfer learning. Closer to our setting, follow-up works [58, 16] propose to use this taxonomy to determine which tasks should be grouped or not in multi-task learning. Once the groupings are determined, a
separate backbone is trained for each group of tasks. Instead, for our analysis, we use Taskonomy-tiny in a more standard multi-task framework, where a backbone is shared across tasks.
For training, we follow the methodology of [58]. We use a ResNet-26 backbone (with varying bottleneck width) with a mirrored decoder; By default, only the encoder is shared across tasks and each task receives its own decoder. To vary model capacity, we add the option to share more or fewer layers of the decoders across tasks. We use the same learning rate of 0.1 and training for 100 epochs using a batch size of 256. We train with SGD with a momentum of 0.9 and a weight decay of  $1e - 4$ . Following [58], all output prediction maps are rescaled to have zero mean and unit variance on the training set, and all dense tasks are trained with L1 loss.
2 Additional analysis results
2.1 Complete results and methodology for Figure 2
In Section 4, we perform MDL and MTL experiments on several pairs of datasets, each time comparing to the single dataset (SD) baseline trained for the same model capacity and training length. All results are run for three random seeds on CelebA and CIFAR+STL, and two random seeds for DomainNet and Taskonomy. To present these results in a condensed form in Figure 2, we first find the scalarization weights  $\mathbf{p}^{*} = (p_{1}^{*}, 1 - p_{1}^{*})$  that yield the best average accuracy across both datasets. Then we report the difference in metrics between MDL trained with weights  $\mathbf{p}^{*}$  and the corresponding SD baseline, for each dataset. Note that for the Taskonomy case, where the tasks are evaluated via L1 loss, we measure the negative difference instead to keep the same interpretation as the other settings where a positive value means MDL improves over SD.
For completeness, we report all results for the CIFAR/STL case as trade-off plots (accuracy on dataset 1 versus accuracy on dataset 2) in Figure B (CIFAR/STL) and in in Figure C (segmentation 2D and depth tasks of Taskonomy). We observe the same trends as summarized in the main paper: First, when increasing model capacity, the MDL performance over the SD baseline increases; This is best seen when width increases (across columns). Second, the optimal weights vary across model sizes: At low width, the best performance is reached for a ratio in the range of [0.3, 0.4]. While larger models prefer  $p_{\mathrm{STL}}^* \in [0.1, 0.3]$ . Finally, it is interesting to note that these weights also differ from heuristics commonly used to set scalarization weights: such as uniform scalarization  $p_{\mathrm{STL}} = 0.5$  or for instance setting the weights to match the number of samples in each dataset  $p_{\mathrm{STL}} = 0.09$ . This further highlights the fact that tuning scalarization weights can make scalarization into a stronger baseline for MDL/MTL.
2.2 MDL helps generalization
When comparing the MDL/MTL and SD performance, we often observe that MDL/MTL improvements over SD are visible at inference but not at training time (as illustrated for instance in Figure D). This indicates that MDL/MTL training helps generalization of the model compared to training on a single dataset. In particular, in the MDL setting, this draws an interesting parallel with data augmentations: In fact, MDL training can be seen as adding additional input data from a new distribution, with a probability given by the scalarization weights  $\mathbf{p}$ , while sharing the same semantic classes. And similarly to data augmentations, adding this extra data source makes the training distribution harder to fit (hence the SD baseline outperforming MDL at train time) but can greatly benefit generalization performance (hence the inverse trend at inference).
To push the analogy further, the experimental study of [41] suggests that a good data augmentation should be one with a high affinity to the original data distribution (i.e., the distribution shift between the original data and augmented one should not be too significant) as well as a high diversity (i.e., the added data should be complex enough, which can be measure e.g. with magnitude of the training loss). This hypothesis also matches our observations: For instance adding infograph data to the real domain leads to a low affinity pairing but with high diversity and yields weaker MDL performance on the real images compared to the SD baseline (see Figure 2).
Finally, we note that the problem of finding optimal scalarization weights mirrors the one of finding data augmentation hyperparameters, which has been extensively explored for the task of image classification. It has led to now commonly-adopted augmentation strategies such as RandAugment [12], AutoAugment [11], or PBA [23], which also uses Population-based training to tackle this problem.
Figure B: Complete analysis results for the CIFAR+STL scenario. Each row corresponds to a different model depth and each column to a model width, in increasing order. In each plot, we plot the model's test accuracy on CIFAR-10 versus the test accuracy on STL-10. The single dataset baseline (SD) is drawn in black and corresponds to the accuracy obtained when training two separate networks, one on each dataset independently. The circle markers correspond to the MDL model trained for different scalarization weights  $\mathbf{p}$ . The value of  $p_{\mathrm{STL}}$  is represented as the color of the marker, while the remaining weight is always set to  $p_{\mathrm{CIFAR}} = 1 - p_{\mathrm{STL}}$ .
3 Methodology for measuring gradients conflicts
In this appendix, we briefly describe our methodology for Section 4.2. We use the same definition of gradient conflicts as in PCGrad [64]: Two task/domain specific gradients are conflicting if and only if the cosine of the angle between them is strictly negative. We train a model using standard uniform scalarization, measure the number of conflicting pairs of task/domain gradients over one epoch of training, and report it as a percentage (of all pairs), for each epoch during training.
We report these results in the main text in Figure 4 and in Figure E below. Our main observation is that the presence or absence of gradient conflicts does not correlate well with actual MDL/MTL performance throughout trained. This challenges the assumption underlying many multi-task optimization (MTO) methods [64, 38, 62] that reducing gradient conflicts leads to improved MTL performance. This also aligns well with recent results of [63, 30] showing that MTO methods that reduce gradient conflicts do not outperform simpler scalarization approaches in practice, and with the hypothesis of [30] that many MTO methods can be reinterpreted as regularization techniques.
4 Consistency of optimal scalarization weights
As noted in the analysis from Section 4, the optimal weights  $\mathbf{p}^*$  is rather consistent across model depths and widths. For instance, on the CIFAR/STL case,  $\mathbf{p}^*$  always falls in the range of [0.2, 0.4].
Figure C: Complete analysis results for the segmentation 2D and depth prediction tasks in the Taskonomy scenario. Each row corresponds to a different number of layers shared in the decoder (shr) and each column to a model width. In each plot, we plot the model's test Le loss of each task on either axis. The single dataset baseline (SD) is drawn in black and corresponds to the accuracy obtained when training two separate networks, one on each dataset independently. The circle markers correspond to the MDL model trained for different scalarization weights  $\mathbf{p}$ . The value of  $p_{\mathrm{segmentation2D}}$  is represented as the color of the marker.
This can also be seen from the qualitative results of the population-based training search of scalarization weights in Section 5.2: While the history of hyperparameter changes during the search differ, PBT tends to converge to similar distribution for the scalarization weights  $\mathbf{p}$  across different model depths and widths. This suggests that the theoretical search space for  $\mathbf{p}$  may be reduced in practice leading to a more computationally efficient search: Performing a rough initial search on a smaller model from the same architecture family can provide a promising range for  $\mathbf{p}$ , and can then be refined by searching with the larger target architecture.
5 PBT results
Because the search space for scalarization weights  $\mathbf{p}$  grows exponentially with the number of tasks, classical hyperparameter search methods such as grid search or random search would struggle to scale as the number of tasks increases. Bayesian optimization (BO) [17] allows for faster results by browsing the search space in a smart way by building and following a probabilistic model of the hyperparameters. However, BO still requires training models to convergence (or until an early stopping criterion is met) which can be computationally expensive. Instead, we experiment with the Population-based Training framework [24] for searching for the optimal scalarization weights. PBT relies on the assumption that the "goodness" of a certain hyperparameter choice can be evaluated in a few epochs during training, rather than having to finish a full run of training.
5.1 Compute resources
As mentioned in the main paper we perform all training runs on a single NVIDIA V100 machine with 32GB of memory. However, both the PBT and MTO models require higher compute resources than a single normal run of training, which we discuss below.
Figure D: Train-test Discrepancy when comparing MDL/MTL improvement over the SD baseline visualized on the CIFAR+STL example. In particular, in the MDL setting, this matches the classical interpretation of data augmentations: Adding additional semantically relevant data from an input distribution may be harder to fit at training time but leads to improved generalization performance at inference.
(a) DomainNet. Individual curves are colored by batch size.
Figure E: Additional result on measuring conflicting gradients when training with standard uniform scalarization for DomainNet (6 domains).
- PBT requires training a population of  $\mathbf{N}$  models that are regularly synchronized; then followed by a final training run with the found optimal weights. Following previous works [23], we perform the hyperparameter search on a smaller subset of the data (in our experiments  $r = 0.7$  fraction of the training set). In summary, the expected computational cost is roughly  $Nr + 1$  times higher than standard training. On the CelebA example, we also observe that using PB2 [46], which combines the benefits of Bayesian Optimisation and PBT, yields better hyperparameters using a smaller population size. In terms of memory usage, PBT is the same as a standard training run: Synchronization is handled via checkpoints saving and loading, such that only one model lives in memory. Finally, we use the publicly available PBT implementation from Raytune[33] which handles all synchronization operations across the population. The implementation would also scale well to more compute resources, as the Ray API allows for easy parallelization.
- MTO. As shown in Figure 1, the bottleneck in most gradient-based methods is memory usage. Consequently, this requires us to decrease the batch size to meet memory requirements and compensate with gradient accumulation (or parallelism if multiple devices are available).
For instance, on the DomainNet experiments with all 6 domains, we need to decrease from a batch size of 512 to a batch size of 128 with 4 steps of accumulation to still fit in memory requirements. This also raises the question of how to handle synchronization across batches: For instance, in PCGrad, the gradient conflicts (and projections) can be computed either (i) per local batch, before accumulation: this may lead to noisier updates; or (ii) after gradients are accumulated: However this is more memory-intensive as this requires to store the previously accumulated per-task gradient as well as the one being currently computed. In practice we use the implementation of [31] for all MTO methods.
When comparing different hyperparameter searches for scalarization, PBT allows for much faster exploration of the search space than classical techniques such as grid search. However, comparing PBT+scalarization with MTO is less straightforward as the computational cost depends on many factors (e.g. population size, number of tasks, and impact on memory usage, etc.), but generally, the "scalarization + hyperparameter search" approach is more favorable in case of low memory requirements as it does not change memory costs compared to standard training. However, soTA gradient-based methods are not very costly for a low number of tasks (e.g. 2-3) as shown in Table 1 which makes them appealing in settings with a few tasks. Nevertheless, one of our key takeaways is that allocating extra resources for tuning scalarization weights, to mirror the extra resources needed for MTO training, makes scalarization into a much stronger baseline, on-par or even outperforming MTO methods as shown in Section 5. Finally, another important difference is that hyperparameter search methods directly optimize for the target objective: the optimal hyperparameters are found by maximizing the average task/domain accuracy on a hold-out validation dataset. In contrast, MTO methods optimize for a proxy metric (such as reducing gradients conflicts) that may not always correlate with final performance as shown in Section 4.2.
In this section, we report some qualitative results of the hyperparameter scheduled found by PBT and PB2. At the end of PBT search, we select the model with the highest validation performance and backtrack its history to backtrack its choices of hyperparameter values during training: This yields the policy of optimal weights found by PBT which is then used to retrain a model on the full training set. We also experimented with retraining a model using the last weights of the policy, but this usually slightly underperform using the whole history of weights in the majority of cases.
In Figure F, we report examples of the policy of weights found by PBT and PB2 search for the parameters  $E_{\text{ready}} = 3$ ,  $Q = 40\%$  and  $N = 12$  for PBT (respectively  $N = 8$  for PB2).
5.3 Quantitative results
Here we report additional results for the CelebA and DomainNet experiments of Section 5, in the style of Table 1 and Table 2 in the main text: We include results using a ResNet-50 on DomainNet in Table B, and results for additional widths in CelebA in Table C and Table D.
Table B: Results of MDL when jointly training on all 6 domains of DomainNet for scalarization (uniform and PBT-found weights) and MTO methods with a ResNet50 backbone. 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.
DomainNet (ResNet50 with 0.25 width)
average	clipart	infograph	painting	quickdraw	real	sketch
Scalarization
Uniform	49.69 ± 0.05	59.90 ± 0.15	22.45 ± 0.01	43.90 ± 0.14	63.13 ± 0.15	58.95 ± 0.09	49.79 ± 0.10
PBT	50.69 ± 0.10	61.69 ± 0.45	21.27 ± 0.10	44.72 ± 0.27	63.96 ± 0.13	62.43 ± 0.13	50.06 ± 0.17
MTO - Loss-based
Uncertainty [26]	40.51 ± 0.19	53.33 ± 0.67	15.70 ± 0.02	34.44 ± 0.41	54.27 ± 0.61	47.86 ± 0.39	37.45 ± 0.37
IMTL-L [38]	37.04 ± 0.17	48.85 ± 0.59	13.93 ± 0.42	30.64 ± 0.36	51.60 ± 0.34	44.12 ± 0.21	33.12 ± 0.50
MTO - Gradient-based
CAGrad [37]	39.82 ± 0.10	50.68 ± 0.05	16.94 ± 0.04	34.37 ± 0.35	52.16 ± 0.47	46.59 ± 0.00	38.20 ± 0.07
GradDrop [8]	39.18 ± 0.15	49.80 ± 0.04	16.77 ± 0.65	33.95 ± 0.52	51.18 ± 0.06	46.04 ± 0.28	37.36 ± 0.20
PCGrad [64]	39.48 ± 0.31	50.42 ± 0.97	16.83 ± 0.31	34.63 ± 0.92	51.14 ± 0.53	46.37 ± 0.51	37.49 ± 0.98
DomainNet (ResNet50 with original width)
average	clipart	infograph	painting	quickdraw	real	sketch
Scalarization
Uniform	51.53 ± 0.06	61.89 ± 0.12	23.63 ± 0.01	45.87 ± 0.01	64.50 ± 0.08	61.46 ± 0.09	51.83 ± 0.33
PBT	51.83 ± 0.06	62.22 ± 0.06	22.61 ± 0.20	46.61 ± 0.29	64.71 ± 0.05	61.91 ± 0.05	52.93 ± 0.10
MTO - Loss-based
Uncertainty [26]	42.90 ± 0.20	56.24 ± 0.44	17.36 ± 0.12	36.91 ± 0.82	56.11 ± 0.08	50.33 ± 0.48	40.49 ± 0.51
IMTL-L [38]	39.69 ± 0.13	52.51 ± 0.59	15.51 ± 0.14	33.45 ± 0.14	53.54 ± 0.05	46.37 ± 0.48	36.75 ± 0.17
MTO - Gradient-based
CAGrad [37]	41.90 ± 0.13	53.32 ± 0.41	17.94 ± 0.42	36.72 ± 0.39	54.15 ± 0.03	48.31 ± 0.25	40.94 ± 0.14
GradDrop [8]	42.15 ± 0.14	53.38 ± 0.54	18.68 ± 0.33	37.00 ± 0.45	53.85 ± 0.11	49.04 ± 0.27	40.95 ± 0.04
PCGrad [64]	41.94 ± 0.20	53.46 ± 0.69	18.20 ± 0.28	36.95 ± 0.66	53.29 ± 0.13	48.88 ± 0.48	40.87 ± 0.49
Table C: (Table best seen zoomed in PDF) Results of MTL when training on all 8 tasks (subset of attributes) of CelebA for a depth of 6 layers 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.
ViT-S/4, 6 layers, 0.25 width
average	age	clothes	face structure	facial hair	gender	hair color	hair style	mouth
Scalarization
Uniform	90.82 ± 2.1e-04	86.96 ± 8.5e-04	92.23 ± 6.8e-04	84.64 ± 6.6e-04	95.30 ± 1.8e-04	97.61 ± 7.1e-04	92.58 ± 3.4e-04	91.03 ± 5.8e-04	86.23 ± 4.4e-04
PBT	90.93 ± 1.8e-04	86.94 ± 7.6e-04	92.37 ± 3.3e-04	84.78 ± 4.5e-04	95.18 ± 2.1e-04	97.58 ± 5.3e-04	92.82 ± 4.3e-04	91.48 ± 3.5e-04	86.25 ± 7.0e-04
PB2	90.90 ± 2.0e-04	87.10 ± 1.2e-03	92.13 ± 3.2e-04	84.82 ± 5.6e-04	95.32 ± 2.8e-04	97.41 ± 4.8e-04	92.67 ± 5.6e-04	91.23 ± 2.3e-04	86.50 ± 2.1e-04
MTO - Loss-based
Uncertainty [26]	90.82 ± 1.9e-04	86.94 ± 1.1e-03	92.22 ± 7.2e-04	84.63 ± 3.0e-04	95.32 ± 2.7e-05	97.62 ± 2.8e-04	92.57 ± 3.9e-04	91.02 ± 4.1e-04	86.22 ± 1.9e-04
IMTL-L [38]	90.82 ± 2.0e-04	86.94 ± 1.2e-03	92.22 ± 7.3e-04	84.63 ± 3.1e-04	95.32 ± 2.7e-05	97.62 ± 3.2e-04	92.57 ± 3.8e-04	91.02 ± 4.1e-04	86.22 ± 1.6e-04
MTO - Gradient-based
CAGrad [37]	90.92 ± 2.7e-04	86.96 ± 2.1e-03	92.35 ± 3.0e-05	84.92 ± 1.2e-04	95.38 ± 3.9e-04	97.56 ± 1.4e-04	92.73 ± 3.8e-04	91.30 ± 2.2e-04	86.14 ± 8.9e-05
GradDrop [8]	90.65 ± 2.9e-04	86.76 ± 1.7e-03	92.03 ± 8.1e-05	84.48 ± 1.2e-04	95.23 ± 3.9e-04	97.41 ± 5.7e-04	92.46 ± 2.1e-04	90.83 ± 1.4e-03	85.98 ± 3.5e-05
PCGrad [64]	90.86 ± 1.5e-04	87.04 ± 1.1e-04	92.32 ± 7.3e-04	84.65 ± 4.0e-04	95.27 ± 1.9e-04	97.62 ± 3.2e-04	92.64 ± 3.7e-04	91.16 ± 2.8e-04	86.22 ± 5.9e-04
ViT-S/4, 6 layers, 0.5 width
average	age	clothes	face structure	facial hair	gender	hair color	hair style	mouth
Scalarization
Uniform	91.28 ± 2.8e-04	87.33 ± 1.9e-03	92.71 ± 8.6e-05	85.15 ± 2.8e-04	95.49 ± 2.2e-04	98.03 ± 6.0e-04	92.99 ± 3.2e-04	91.63 ± 5.8e-04	86.87 ± 3.8e-04
PBT	91.29 ± 1.7e-04	87.56 ± 3.3e-04	92.77 ± 1.6e-04	85.34 ± 2.5e-04	95.40 ± 2.9e-04	97.81 ± 6.4e-04	92.99 ± 2.1e-04	91.56 ± 7.2e-04	86.89 ± 7.6e-04
PB2	91.38 ± 2.3e-04	87.71 ± 1.1e-03	92.84 ± 3.3e-04	85.20 ± 6.0e-04	95.54 ± 4.5e-04	98.00 ± 7.3e-04	92.95 ± 6.2e-04	91.73 ± 2.2e-04	87.05 ± 7.0e-04
MTO - Loss-based
Uncertainty [26]	91.30 ± 3.1e-04	87.52 ± 2.2e-03	92.72 ± 1.4e-04	85.14 ± 4.6e-04	95.48 ± 4.7e-04	98.08 ± 5.3e-04	92.97 ± 3.5e-04	91.65 ± 1.9e-04	86.86 ± 6.4e-04
IMTL-L [38]	91.30 ± 2.6e-04	87.49 ± 1.4e-03	92.74 ± 6.6e-05	85.14 ± 1.0e-04	95.51 ± 1.3e-04	98.06 ± 7.8e-04	92.98 ± 5.8e-04	91.66 ± 6.7e-04	86.85 ± 1.0e-03
MTO - Gradient-based
CGrad [37]	91.32 ± 3.2e-04	87.47 ± 2.3e-03	92.86 ± 4.9e-04	85.26 ± 3.1e-04	95.51 ± 4.3e-04	98.03 ± 0.0e+00	93.00 ± 5.9e-04	91.73 ± 1.5e-04	86.72 ± 2.3e-04
GradDrop [8]	91.19 ± 2.3e-04	87.42 ± 8.5e-04	92.65 ± 5.2e-04	85.03 ± 9.5e-04	95.43 ± 2.9e-04	97.96 ± 6.7e-04	92.80 ± 1.8e-04	91.54 ± 8.5e-04	86.70 ± 3.0e-04
PCGrad [64]	91.30 ± 4.6e-04	87.45 ± 2.8e-03	92.72 ± 6.0e-04	85.16 ± 8.9e-04	95.54 ± 5.0e-04	98.08 ± 2.1e-03	92.97 ± 1.6e-04	91.63 ± 1.2e-04	86.82 ± 3.8e-04