SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We assume that we have access to multiple instances or views of the same event, where only a subset of generative factors changes between views. The knowledge about the data collection process provides a form of weak supervision. For example, we have two images of a robot arm as depicted here on the left side (see (Gondal et al., 2019)), which we would describe using high-level concepts such as color, position or rotation degree. From the data collection process, we know that a subset of these generative factors is shared between the two views. We do not know how many generative factors there are in total nor how many of them are shared. More precisely, looking at the robot arm, we do not know that the views share two latent factors, depicted in red, out of a total of four factors. Please note that we chose four generative in Figure 10 only for illustrative reason as there are seven generative factors in the  $mpi3d$  toy dataset. Hence, the goal of learning under weak supervision is not
Figure 10: Motivation for the Partitioning of Generative Factors under weak supervision. The knowledge about the data collection process provides a weak supervision signal. We have access to a dataset of pairs of images of the same robot arm with a subset of shared generative factors (in red). We want to learn the shared and independent generative factors in addition to learning from the data. The images of the robot arms are taken from Locatello et al. (2020) but originate from the mpi3d toy dataset (see https://github.com/rr-learning/disentanglement_dataset). The image is from Sutter et al. (2023) and their ICLR 2023 presentation video (see https://iclr.cc/virtual/2023/poster/10707).
only to infer good representations, but also inferring the number of shared and independent generative factors. Learning what is shared and what is independent lets us reason about the group structure without requiring explicit knowledge in the form of expensive labeling. Additionally, leveraging weak supervision and, hence, the underlying group structure holds promise for learning more generalizable and disentangled representations (see (e.g., Locatello et al., 2020)).
C.3.1 Generative Model
We assume the following generative model for DRPM-VAE
$$
\begin{array}{l} p (\boldsymbol {X}) = \int_ {\boldsymbol {z}} p (\boldsymbol {X}, \boldsymbol {z}) d \boldsymbol {z} (51) \\ = \int_ {z} p (\boldsymbol {X} \mid \boldsymbol {z}) p (\boldsymbol {z}) d \boldsymbol {z} (52) \\ \end{array}
$$
where  $\mathbf{z} = \{\mathbf{z}_s, \mathbf{z}_1, \mathbf{z}_2\}$ . The two frames share an unknown number  $n_s$  of generative latent factors  $\mathbf{z}_s$ , and an unknown number,  $n_1$  and  $n_2$ , of independent factors  $\mathbf{z}_1$  and  $\mathbf{z}_2$ . The RPM infers  $n_k$  and  $\mathbf{z}_k$  using  $Y$ . Hence, the generative model extends to
$$
\begin{array}{l} p (\boldsymbol {X}) = \int_ {\boldsymbol {z}} p (\boldsymbol {X} \mid \boldsymbol {z}) \sum_ {Y} p (\boldsymbol {z} \mid Y) p (Y) d \boldsymbol {z} \\ = \int_ {\mathbf {z}} p \left(\mathbf {x} _ {1}, \mathbf {x} _ {2} \mid \mathbf {z} _ {s}, \mathbf {z} _ {1}, \mathbf {z} _ {2}\right) \sum_ {Y} p (\mathbf {z} \mid Y) p (Y) d \mathbf {z} \\ = \int_ {\boldsymbol {z} _ {s}, \boldsymbol {z} _ {1}, \boldsymbol {z} _ {2}} p \left(\boldsymbol {x} _ {1} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {1}\right) p \left(\boldsymbol {x} _ {2} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {2}\right) \sum_ {Y} p \left(\boldsymbol {z} _ {s}, \boldsymbol {z} _ {1}, \boldsymbol {z} _ {2} \mid Y\right) p (Y) d \boldsymbol {z} _ {s} d \boldsymbol {z} _ {1} d \boldsymbol {z} _ {2} \tag {53} \\ \end{array}
$$
Figure 11 shows the generative and inference models assumptions in a graphical model.
(a) Generative Model
Figure 11: Graphical Models for DRPM-VAE models in the weakly-supervised experiment.
(b) Inference Model
C.3.2 DRPM ELBO
We derive the following ELBO using the posterior approximation  $q(\mathbf{z}, Y \mid \mathbf{X})$
$$
\begin{array}{l} \mathcal {L} _ {E L B O} (\boldsymbol {X}) = \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p (\boldsymbol {X} \mid \boldsymbol {z}, Y) - \log \frac {q (\boldsymbol {z} , Y \mid \boldsymbol {X})}{p (\boldsymbol {z} , Y)} \right] (54) \\ = \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p (\boldsymbol {X} \mid \boldsymbol {z}) - \log \frac {q (\boldsymbol {z} \mid Y , \boldsymbol {X}) q (Y \mid \boldsymbol {X})}{p (\boldsymbol {z}) p (Y)} \right] (55) \\ = \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p \left(\boldsymbol {x} _ {1}, \boldsymbol {x} _ {2} \mid \boldsymbol {z}\right) - \log \frac {q (\boldsymbol {z} \mid Y , \boldsymbol {X})}{p (\boldsymbol {z})} - \log \frac {q (Y \mid \boldsymbol {X})}{p (Y)} \right] (56) \\ = \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p (\boldsymbol {x} _ {1} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {1}) \right] - \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p (\boldsymbol {x} _ {2} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {2}) \right] \\ - \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log \frac {q \left(\boldsymbol {z} _ {s} , \boldsymbol {z} _ {1} , \boldsymbol {z} _ {2} \mid Y , \boldsymbol {X}\right)}{p \left(\boldsymbol {z} _ {s} , \boldsymbol {z} _ {1} , \boldsymbol {z} _ {2}\right)} \right] - \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log \frac {q (Y \mid \boldsymbol {X})}{p (Y)} \right] (57) \\ \end{array}
$$
Following Lemma 4.2, we are able to optimize DRPM-VAE using the following ELBO  $\mathcal{L}_{ELBO}(\pmb {X})$  ..
$$
\begin{array}{l} \mathcal {L} _ {E L B O} \geq \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p \left(\boldsymbol {x} _ {1} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {1}\right) \right] - \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p \left(\boldsymbol {x} _ {2} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {2}\right) \right] (58) \\ - \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log \frac {q \left(\boldsymbol {z} _ {s} , \boldsymbol {z} _ {1} , \boldsymbol {z} _ {2} \mid Y , \boldsymbol {X}\right)}{p \left(\boldsymbol {z} _ {s} , \boldsymbol {z} _ {1} , \boldsymbol {z} _ {2}\right)} \right] (59) \\ - \mathbb {E} _ {q (Y | \boldsymbol {X})} \left[ \log \left(\frac {\left| \Pi_ {Y} \right| \cdot q (\boldsymbol {n} \mid \boldsymbol {X} ; \boldsymbol {\omega})}{p (\boldsymbol {n} ; \boldsymbol {\omega} _ {p}) p \left(\pi_ {Y} ; \boldsymbol {s} _ {p}\right)}\right) \right] (60) \\ - \log \left(\max  _ {\tilde {\pi}} q \left(\tilde {\pi} \mid \boldsymbol {X}; \boldsymbol {s}\right)\right), (61) \\ \end{array}
$$
where  $\pi_{Y}$  is the permutation that leads to  $Y$  during the two-stage resampling process. Further, we want to control the regularization strength of the KL divergences similar to the  $\beta$ -VAE (Higgins et al., 2016). The ELBO  $\mathcal{L}(\mathbf{X})$  to be optimized can be written as
$$
\begin{array}{l} \mathcal {L} _ {E L B O} = \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p \left(\boldsymbol {x} _ {1} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {1}\right) \right] + \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log p \left(\boldsymbol {x} _ {2} \mid \boldsymbol {z} _ {s}, \boldsymbol {z} _ {2}\right) \right] (62) \\ - \beta \cdot \mathbb {E} _ {q (\boldsymbol {z}, Y | \boldsymbol {X})} \left[ \log \frac {q \left(\boldsymbol {z} _ {s} , \boldsymbol {z} _ {1} , \boldsymbol {z} _ {2} \mid Y , \boldsymbol {X}\right)}{p \left(\boldsymbol {z} _ {s} , \boldsymbol {z} _ {1} , \boldsymbol {z} _ {2}\right)} \right] (63) \\ - \gamma \cdot \mathbb {E} _ {q (Y | \boldsymbol {X})} \left[ \log \left(\frac {\left| \Pi_ {Y} \right| \cdot q (\boldsymbol {n} ; \boldsymbol {\omega} (\boldsymbol {X}))}{p (\boldsymbol {n} ; \boldsymbol {\omega} _ {p})}\right) \right] (64) \\ - \delta \cdot \mathbb {E} _ {q (Y | \boldsymbol {X})} \left[ \log \left(\frac {\operatorname* {m a x} _ {\tilde {\pi}} q (\tilde {\pi} ; \boldsymbol {s} (\boldsymbol {X}))}{p \left(\pi_ {Y} ; \boldsymbol {s} _ {p}\right)}\right) \right] (65) \\ \end{array}
$$
where  $s(\mathbf{X})$  and  $\omega(\mathbf{X})$  denote distribution parameters, which are inferred from  $\mathbf{X}$  (similar to the Gaussian parameters in the vanilla VAE).
Figure 12: Setup for the weakly-supervised experiment. The three methods differ only in the View Aggregation module.
As in vanilla VAEs, we can estimate the reconstruction term in Equation (58) with MCMC by applying the reparametrization trick (Kingma and Welling, 2014) to  $q(z \mid Y, X)$  to sample  $L$  samples  $z^{(l)} \sim q(z \mid Y, X)$  and compute their reconstruction error to estimate Equation (58). Similarly, we can sample from  $q(Y \mid X)$  times. We use  $L = 1$  to estimate all expectations in  $\mathcal{L}_{ELBO}$ .
C.3.3 Implementation and Hyperparameters
In this experiment, we use the disentanglement_lib from Locatello et al. (2020). We use the same architectures proposed in the original paper for all methods we compare to. The baseline algorithms, LabelVAE (Bouchacourt et al., 2018; Hosoya, 2018) and AdaVAE (Locatello et al., 2020) are already implemented in disentanglement_lib. For details on the implementation of these methods we refer to the original paper from Locatello et al. (2020). HGVAE is implemented in Sutter et al. (2023). We did not change any hyperparameters or network details. All experiments were performed using  $\beta = 1$  as this is the best performing  $\beta$  (according to Locatello et al. (2020). For DRPMVAE we chose  $\gamma = 0.25$  for all runs. All models are trained on 5 different random seeds and the reported results are averaged over the 5 seeds. We report mean performance with standard deviations.
We adapted Figure 12 from Sutter et al. (2023). It shows the baseline architecture, which is used for all methods. As already stated in the main part of the paper, the methods only differ in the View Aggregation module, which determines the shared and independent latent factors. Given a subset  $S$  of shared latent factors, we have
$$
q _ {\phi} \left(z _ {i} \mid \boldsymbol {x} _ {j}\right) = a v g \left(q _ {\phi} \left(z _ {i} \mid \boldsymbol {x} _ {1}\right), q _ {\phi} \left(z _ {i} \mid \boldsymbol {x} _ {2}\right)\right) \quad \forall i \in S \tag {66}
$$
$$
q _ {\phi} \left(z _ {i} \mid \boldsymbol {x} _ {j}\right) = q _ {\phi} \left(z _ {i} \mid \boldsymbol {x} _ {j}\right) \quad \text {e l s e} \tag {67}
$$
where avg is the averaging function of choice (Locatello et al., 2020; Sutter et al., 2023) and  $j \in \{1, 2\}$ . The methods used (i.e. Label-VAE, Ada-VAE, HG-VAE, DRPM-VAE) differ in how to select the subset S.
For DRPM-VAE, we infer  $\omega$  from the pairwise KL-divergences  $KL_{pw}$  between the latent vectors of the two views.
$$
K L _ {p w} \left(\boldsymbol {x} _ {1}, \boldsymbol {x} _ {2}\right) = \frac {1}{2} K L \left[ q \left(\boldsymbol {z} _ {1} \mid \boldsymbol {x} _ {1}\right) \right\| q \left(\boldsymbol {z} _ {2} \mid \boldsymbol {x} _ {2}\right) \bigg ] + \frac {1}{2} K L \left[ q \left(\boldsymbol {z} _ {2} \mid \boldsymbol {x} _ {2}\right) \right\| q \left(\boldsymbol {z} _ {1} \mid \boldsymbol {x} _ {1}\right) \bigg ] \tag {68}
$$
where  $q(\boldsymbol{z}_j \mid \boldsymbol{x}_j)$  are the encoder outputs of the respective images. We do not average or sum across dimensions in the computation of  $KL_{pw}(\cdot)$  such that the  $KL_{pw}(\cdot)$  is  $d$ -dimensional, where  $d$  is the latent space size. The encoder  $E$  in Figure 12 maps to  $\mu(\boldsymbol{x}_j)$  and  $\sigma(\boldsymbol{x}_j)$  of a Gaussian distribution. Hence, we can compute the KL divergences above in closed form. Afterwards, we feed the pairwise KL divergence  $KL_{pw}$  to a single fully-connected layer, which maps from  $d$  to  $K$  values
$$
\log \omega = F C \left(K L _ {p w} \left(\boldsymbol {x} _ {1}, \boldsymbol {x} _ {2}\right)\right) \tag {69}
$$
where  $d = 10$  and  $K = 2$  in this experiment.  $d$  is the total number of latent dimensions and  $K$  is the number of groups in the latent space. To infer the scores  $s(\mathbf{X})$  we again rely on the pairwise KL divergence  $KL_{pw}$ . Instead of using another fully-connected layer, we directly use the log-values of the pairwise KL divergence
$$
\log \boldsymbol {s} = \log K L _ {p w} \left(\boldsymbol {x} _ {1}, \boldsymbol {x} _ {2}\right) \tag {70}
$$
Similar to the original works, we also anneal the temperature parameter for  $p(n; \omega)$  and  $p(\pi; s)$  (Grover et al., 2019; Sutter et al., 2023). We use the same annealing function as in the clustering experiment (see Appendix C.2). We anneal the temperature  $\tau$  from 1.0 to 0.5 over the complete training time.
Figure 13: Samples from the noisyMultiMNIST dataset with increasing noise ratio in the right task.
C.4 Multitask Learning
C.4.1 MultiMNIST Dataset
The different tasks in multitask learning often vary in difficulty. To measure the effect of discrepancies in task difficulties on DRPM-MTL, we introduce the noisyMultiMNIST dataset.
The noisyMultiMNIST dataset modifies the MultiMNIST dataset (Sabour et al., 2017) as follows. In the right image, we set each pixel value to zero with probability  $\alpha \in [0,1]$ . This is done before merging the left and right image in order to only affect the difficulty of the right task. Note that for  $\alpha = 0$  noisyMultiMNIST is equivalent to MultiMNIST and for  $\alpha = 1$  the right task can no longer be solved. This allows us to control the difficulty of the right task, without changing the difficulty of the left. A few examples are shown in Figure 13.
C.4.2 Implementation & Architecture
The multitask loss function for the MultiMNIST dataset is
$$
\mathbb {L} = w _ {L} \mathbb {L} _ {L} + w _ {R} \mathbb {L} _ {R} \tag {71}
$$
where  $w_{L}$  and  $w_{L}$  are the loss weights, and  $\mathbb{L}_L$  and  $\mathbb{L}_R$  are the individual loss terms for the respective tasks  $L$  and  $R$ . In our experiments, we set the task weights to be equal for all dataset versions, i.e.  $w_{L} = w_{R} = 0.5$ . We use these loss weights for the DRPM-MTL and ULS method. For the ULS method, it is by definition and to see the influence of a mismatch in loss weights. The DRPM-MTL method on the other hand does not need additional weighting of loss terms. The task losses are defined as cross-entropy losses
$$
\mathbb {L} _ {t} = - \sum_ {c = 1} ^ {C _ {t}} \boldsymbol {g} t _ {c} \log \boldsymbol {p} _ {c} = - \boldsymbol {g} \boldsymbol {t} ^ {T} \log \boldsymbol {p} \tag {72}
$$
where  $C_L = C_R = 10$  for MultiMNIST,  $\pmb{g}t$  is a one-hot encoded label vector and  $\pmb{p}$  is a categorical vector of estimated class assignments probabilities, i.e.  $\sum_{c}\pmb{p}_{c} = 1$ .
The predictions for the individual tasks  $\pmb{p}_t$  are given as
$$
\boldsymbol {p} _ {t} = h _ {\theta_ {t}} (\boldsymbol {z}), \quad \text {w h e r e} \tag {73}
$$
$$
\boldsymbol {z} = \operatorname {e n c} _ {\theta} (\boldsymbol {x}) \tag {74}
$$
for a sample  $\pmb{x} \in \pmb{X}$  (see also Figure 14). We use an adaptation of the LeNet-5 architecture LeCun et al. (1998) to the multitask learning problem (Sener and Koltun, 2018). Both DRPM-MTL and ULS use the same network  $\mathrm{enc}_{\theta}(\cdot)$  with shared architecture up to some layer for both tasks, after which the network branches into two task-specific sub-networks that perform the classifications. Different to the ULS method, the task-specific networks in the DRPM-MTL pipeline predict the digit using only a subset of  $z$ . DRPM-MTL uses the following prediction scheme
$$
\boldsymbol {p} _ {t} = h _ {\theta_ {t}} \left(\boldsymbol {z} _ {t}\right), \quad \text {w h e r e} \tag {75}
$$
$$
\boldsymbol {z} _ {t} = \boldsymbol {z} \odot \boldsymbol {y} _ {t} \tag {76}
$$
$$
\boldsymbol {y} _ {t} = \operatorname {D R P M} (\boldsymbol {\omega}, \boldsymbol {s}) _ {t} = \operatorname {D R P M} \left(\operatorname {e n c} _ {\varphi} (\boldsymbol {x})\right) _ {t} \tag {77}
$$
The DRPM-MTL encoder first predicts a latent representation  $\mathbf{z} \gets \mathrm{enc}_{\theta}(\mathbf{x})$ , where  $\mathbf{x}$  is the input image. Using the same encoder architecture but different parameters  $\varphi$ , we predict a partitioning encoding  $z' \gets \mathrm{enc}_{\varphi}(\mathbf{x})$ . With a single linear layer per DRPM log-parameter  $\log \omega$  and  $\log s$  are computed. Next we infer the partition masks  $\mathbf{y}_{\mathrm{L}}, \mathbf{y}_{\mathrm{R}} \sim p(\mathbf{y}_{\mathrm{L}}, \mathbf{y}_{\mathrm{R}}; \boldsymbol{\omega}, s)$ . We then feed the masked latent representations  $\mathbf{z}_{\mathrm{L}} \gets \mathbf{z} \odot \mathbf{y}_{\mathrm{L}}$  and  $\mathbf{z}_{\mathrm{R}} \gets \mathbf{z} \odot \mathbf{y}_{\mathrm{R}}$  into the task specific classification networks  $h_{\theta_{\mathrm{L}}}(\mathbf{z}_{\mathrm{L}})$  and  $h_{\theta_{\mathrm{R}}}(\mathbf{z}_{\mathrm{R}})$  respectively to obtain the task specific predictions. Since the two tasks in the MultiMNIST dataset are of similar nature, the task-specific networks  $h_{\theta_{\mathrm{L}}}^{}$  and  $h_{\theta_{\mathrm{R}}}^{}$  share the same architecture, but have different parameters.
Figure 14: Overview of the multitask learning pipeline of the DRPM-MTL method.
C.4.3 Training
For both the ULS and the DRPM-MTL model, we use the Adam optimizer with learning rate 0.0005 and train them for 200 epochs with a batch size of 256. We again choose an exponential schedule for the temperature  $\tau$  and anneal it over the training time, as is explained in Appendix C.2.3.
In our ablation we use  $\alpha \in \{0, 0.1, 0.2, \ldots, 0.9\}$  and train each model with five different seeds. The reported accuracies and partition sizes are then means over the five seeds with the error bands indicating the variance and standard deviation respectively. We evaluate each model after the epoch with the best average test accuracy.
C.4.4 CelebA for MTL
In addition to the experiment shown in Section 5.3, we show additional results for DRPM-MTL on the CelebA dataset (Liu et al., 2015). In MTL, each of the 40 attributes of the CelebA dataset serves as an individual task. Hence, using CelebA for MTL results is a 40 task learning problem making the scaling of different task losses more difficult compared to MultiMNIST (see Section 5.3) where we only need to scale two different tasks.
We again use the newly introduced DRPM-MTL method and compare it to the ULS model. We use the same pipeline as for MultiMNIST dataset but with different encoders and hyperparameters (see Appendices C.4.2 and C.4.3). We use the pipeline of Sener and Koltun (2018) with a ResNet-based encoder to map an image to a representation of  $d = 64$  dimensions. For architectural details, we refer to Sener and Koltun (2018) and https://github.com/is1-org/MultiObjectiveOptimization.
Again, ULS inputs all  $d = 64$  dimensions to the task-specific sub-networks whereas DRPM-MTL partitions the intermediate representations into  $n_T$  different subsets, which are then fed to the respective task networks.  $n_T$  is the number of tasks.
Compared to the MultiMNIST experiment (see Appendix C.4.2), we introduce an additional regularization for the DRPM-MTL method. The additional regularization is based on the upper bound in Lemma 4.2 and is penalizing size of  $|\Pi_Y|$  for a given  $n$ . Hence, the loss function changes to
$$
\mathbb {L} = \frac {1}{n _ {T}} \sum_ {t = 1} ^ {n _ {T}} \mathbb {L} _ {t} + \lambda \cdot \mathbb {L} _ {\text {r e g}} \tag {78}
$$
$$
\text {w h e r e} \mathbb {L} _ {\text {r e g}} = \log \left(\prod_ {t = 1} ^ {n _ {T}} n _ {t}!\right) = \sum_ {t = 1} ^ {n _ {T}} \log \Gamma \left(n _ {t} + 1\right) \tag {79}
$$
For all versions of the experiment (i.e.  $n_T \in 10, 20, 40$ ), we set  $\lambda = 0.015 \approx \frac{1}{64}$ , which is the number of elements we want to partition. The task losses  $\mathbb{L}_t$  are simple BCE losses similar to the MultiMNIST experiments but with two classes per task only.
We perform two different experiments based on the CelebA experiment. First, we use form a MTL experiments using the first 10 attributes out of the 40 attributes. Second, we increase the number of different tasks to 20. Because we sort the attributes alphabetically in both cases, the first 10 tasks are
Table 7: Results for the MTL experiment on the CelebA dataset. We compare the DRPM-MTL again to the ULS method. We assess the performance of both methods on two sub-experiment of the CelebA experiment. In Table 7a, we form a MTL experiment with 10 different tasks. In Table 7b, we form a MTL experiment with 20 different tasks where the first 10 tasks are the same as in the 10 tasks experiment. And in Table 6c, a MTL experiment with all 40 tasks from the dataset. We train both methods for 50 epochs using a learning rate of 0.0001 and a batch size of 128. The temperature annealing schedule remains the same as in the MultiMNIST experiment. We report the per task classification accuracy in percentages (\%) as well as the average task accuracy in the bottom row of both subtables.
(a) 10 Tasks
ULS	DRPM
T0	92.0±0.5	92.4±0.5
T1	83.8±0.4	83.7±0.2
T2	80.2±0.5	80.2±0.4
T3	81.9±0.8	82.2±0.6
T4	98.5±0.2	98.5±0.1
T5	95.2±0.2	95.3±0.2
T6	80.0±1.4	82.4±0.4
T7	82.0±0.3	82.2±0.2
T8	89.7±0.7	90.7±0.2
T9	94.6±0.5	95.0±0.2
avg(Tasks)	87.8±0.3	88.3±0.1
(b) 20 Tasks
ULS	DRPM
T0	92.4±0.7	93.0±0.2
T1	83.7±0.6	83.9±0.7
T2	79.9±0.6	80.1±0.4
T3	82.4±0.5	83.0±0.7
T4	98.6±0.1	98.6±0.1
T5	95.2±0.1	95.5±0.0
T6	82.0±1.3	84.4±0.4
T7	82.5±0.1	82.8±0.2
T8	90.1±0.9	91.0±0.4
T9	94.7±0.2	95.1±0.1
T10	95.9±0.1	95.9±0.1
T11	84.9±0.1	84.6±0.3
T12	91.0±0.4	91.6±0.2
T13	94.7±0.1	94.9±0.1
T14	95.4±0.3	96.0±0.1
T15	99.2±0.0	99.2±0.1
T16	95.8±0.3	96.0±0.1
T17	97.3±0.3	97.5±0.2
T18	91.2±0.3	91.2±0.1
T19	87.0±0.3	87.3±0.2
avg(Tasks)	90.7±0.2	91.1±0.1
shared between the two experiment versions. And third, we set  $n_T = 40$ , where the first 20 tasks are the shared with the previous experiment.
Table 7 shows the results of all CelebA experiments for both methods, ULS and DRPM-MTL. We see that the DRPM-MTL scales better to a larger number of tasks compared to the ULS method, highlighting the importance of finding new ways of automatic scaling between tasks. Interestingly, the DRPM-MTL outperforms the ULS method on most tasks for the 20-tasks experiment even though it has only access to  $d / n_{T} = 64 / 20 = 3.2$  dimensions on average. And even more extrem for  $n_{T} = 40$ , DRPM-MTL on average has only access to  $d / n_{T} = 64 / 40 = 1.6$  dimensions. On the other hand, the ULS method can access the full set of 64 dimensions for every single task.
C.5 Supervised Learning
Given the true partition  $Y$ , we can also adapt the DRPM to learn partitions in a supervised fashion. One instance where we know the true partition of elements is in the case of classification. There, for a given batch  $X$  containing  $B$  samples, we have  $Y := (\pmb{y}_1, \dots, \pmb{y}_B)$  where  $\pmb{y}_i \in \mathbb{R}^K$  is the one-hot encoding of the label of the  $i$ -th sample in the batch. In this ablation, we infer  $\hat{Y} := DRPM(\omega_{\theta_1}(X), s_{\theta_2}(X))$ , where we compute  $\omega(X)$  and  $s(X)$  as in the variational clustering experiment (Appendix C.2.2) and use the DRPM without resampling as in the multitask experiment (Section 5.3). We optimize