Title: Leveraging sparse and shared feature activations for disentangled representation learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: disentanglement

OOD generalization

multitask learning
Abstract: Recovering the latent factors of variation of high dimensional data has so far focused on simple synthetic settings. Mostly building on unsupervised and weakly-supervised objectives, prior work missed out on the positive implications for representation learning on real world data. In this work, we propose to leverage knowledge extracted from a diversified set of supervised tasks to learn a common disentangled representation. Assuming each supervised task only depends on an unknown subset of the factors of variation, we disentangle the feature space of a supervised multi-task model, with features activating sparsely across different tasks and information being shared as appropriate. Importantly, we never directly observe the factors of variations but establish that access to multiple tasks is sufficient for identifiability under sufficiency and minimality assumptions.
We validate our approach on six real world distribution shift benchmarks, and different data modalities (images, text), demonstrating how disentangled representations can be transferred to real settings.

Leveraging sparse and shared feature activations for disentangled representation learning

Abstract

Research on recovering the latent factors of variation of high dimensional data has so far focused on simple synthetic settings. Mostly building on unsupervised and weakly-supervised objectives, prior work missed out on the positive implications for representation learning on real world data. In this work, we propose to leverage knowledge extracted from a diversified set of supervised tasks to learn a common disentangled representation. Assuming that each supervised task only depends on an unknown subset of the factors of variation, we disentangle the feature space of a supervised multi-task model, with features activating sparsely across different tasks and information being shared as appropriate. Importantly, we never directly observe the factors of variations, but establish that access to multiple tasks is sufficient for identifiability under sufficiency and minimality assumptions. We validate our approach on six real world distribution shift benchmarks, and different data modalities (images, text), demonstrating how disentangled representations can be transferred to real settings.

1 Introduction

A fundamental question in deep learning is how to learn meaningful and reusable representation from high dimensional data observations [8, 75, 78, 77]. A core area of research pursuing is centered on disentangled representation learning (DRL) [56, 8, 33] where the aim is to learn a representation which recovers the factors of variations (FOVs) underlying the data distribution. Disentangled representations are expected to contain all the information present in the data in a compact and interpretable structure [46, 16] while being independent from a particular task [29]. It has been argued that separating information into interventionally independent factors [78] can enable robust downstream predictions, which was partially validated in synthetic settings [19, 58]. Unfortunately, these benefits did not materialize in real world representations learning problems, largely limited by a lack of scalability of existing approaches.

In this work we focus on leveraging knowledge from different task objectives to learn better representations of high dimensional data, and explore the link with disentanglement and out-of-distribution (OOD) generalization on real data distributions. Representations learned from a large diversity of tasks are indeed expected to be richer and generalize better to new, possibly out-of-distribution, tasks. However, this is not always the case, as different tasks can compete with each other and lead to weaker models. This phenomenon, known as negative transfer [61, 91] in the context of transfer learning or task competition [83] in multitask learning, happens when a limited capacity model is used to learn two different tasks that require expressing high feature variability and/or coverage. Aiming to use the same features for different objectives makes them noisy and often increases the sensitivity to spurious correlations [35, 27, 7], as features can be both predictive and detrimental for

different tasks. Instead, we leverage a diverse set of tasks and assume that each task only depends on an unknown subset of the factors of variation. We show that disentangled representations naturally emerge without any annotation of the factors of variations under the following two representation constraints:

Sparse sufficiency : Features should activate sparsely with respect to tasks. The representation is sparsely sufficient in the sense that any given task can be solved using few features.

Minimality : Features are maximally shared across tasks whenever possible. The representation is minimal in the sense that features are encouraged to be reused, i.e., duplicated or split features are avoided.

These properties are intuitively desirable to obtain features that (i) are disentangled w.r.t. to the factors of variations underlying the task data distribution (which we also theoretically argue in Proposition 2.1), (ii) generalize better in settings where test data undergo distribution shifts with respect to the training distributions, and (iii) suffer less from problems related to negative transfer phenomena. To learn such representations in practice, we implement a meta learning approach, enforcing feature sufficiency and sharing with a sparsity regularizer and an entropy based feature sharing regularizer, respectively, incorporated in the base learner. Experimentally, we show that our model learns meaningful disentangled representations that enable strong generalization on real world data sets. Our contributions can be summarized as follows:

We demonstrate that is possible to learn disentangled representations leveraging knowledge from a distribution of tasks. For this, we propose a meta learning approach to learn a feature space from a collection of tasks while incorporating our sparse sufficiency and minimality principles favoring task specific features to coexist with general features.

Following previous literature, we test our approach on synthetic data, validating in an idealized controlled setting that our sufficiency and minimality principles lead to disentangled features w.r.t. the ground truth factors of variation, as expected from our identifiability result in Proposition 2.1.

We extend our empirical evaluation to non-synthetic data where factors of variations are not known, and show that our approach generalizes well out-of-distribution on different domain generalization and distribution shift benchmarks.

2 Method

Given a distribution of tasks t ∼ T and data ( x t , y t ) ∼ P t for each task t , we aim to learn a disentangled representation g ( x ) = ˆ z ∈ ˆ Z ⊆ R M , which generalizes well to unseen tasks. We learn this representation g by imposing the sparse sufficiency and minimality inductive biases.

2.1 Learning sparse and shared features

Our architecture (see Figure 1) is composed of a backbone module g θ that is shared across all tasks and a separate linear classification head f ϕ t , which is specific to each task t . The backbone is responsible to compute and learn a general feature representation for all classification tasks. The linear head solves a specific classification problem for the task-specific data ( x t , y t ) ∼ P t in the feature space ˆ Z while enforcing the feature sufficiency and minimality principles. Adopting the typical meta-learning setting [34], the backbone module g θ can be viewed as the meta learner while the task-specific classification heads f ϕ t can be viewed as the base learners . In the meta-learning setting we assume to have access to samples for a new task given by a support set U , with elements ( x U , y U ) ∈ U . These samples are used to fit the linear head f ϕ ∗ leading to the optimal feature weights for the given task. For a query x Q ∈ Q , the prediction is obtained by computing the forward pass ˆ y = f ϕ ∗ ( g θ ( x Q )) .

Enforcing feature minimality and sufficiency. To solve a task in the feature space ˆ Z of the backbone module we impose the following regularizer Reg ( ϕ ) on the classification heads f ϕ with parameter ϕ ∈ R T × M × C , where T is the number of tasks, M the number of features, and C the number of classes. The regularizer is responsible for enforcing the feature minimality and sufficiency

x

U

gθ

fϕ

ˆ

z

U

gθ

ˆ

y U

L

inner

x

Q

gθ

fϕ

∗

ϕ

∗

ˆ

z

Q

gθ

ˆ

y Q

L

outer

properties. It is composed of the weighted sum of a sparsity penalty Reg L 1 and an entropy-based feature sharing penalty: Reg sharing

with scalar weights α and β . The penalty terms are defined by:

where ˜ ϕ m = 1 TC ∑ t,c | ϕ t,c,m | ∑ t,c,m | ϕ t,c,m | are the normalized classifier parameters. Sufficiency is enforced by a sparsity regularizer given by the L 1 -norm, which constrains classification head to use only a sparse subset of the features. Minimality is enforced by the feature sharing term: minimizing the entropy of the distribution of feature importances (i.e. normalized | ϕ t | ) averaged across a mini batch of T tasks, leads to a more peaked distribution of activations across tasks. This forces features to cluster across tasks and therefore be reused by different tasks, when useful.We remark that different choices for the regularizers coming from the linear multitask learning literature (e.g. [59, 39, 38]) to enforce sparse sufficiency and minimality are indeed possibile. We leave their exploration as a future direction.

2.2 Training method

We train the model in meta-learning fashion by minimizing the test error over the expectation of the task distribution t ∼ T . This can be formalized as a bi-level optimization problem . The optimal backbone model g θ ∗ is given by the outer optimization problem :

where f ϕ ∗ are the optimal classifiers obtained from solving the inner optimization problem , and ( x Q t , y Q t ) ∈ Q t are the test (or query) datum from the query set Q t for task t . Let U t be the support set with samples ( x U t , y U t ) ∈ U for task t , where typically the support set is distinct from the query set, i.e., U ∩ Q = ∅ . The optimal classifiers f ϕ ∗ are given by the inner optimization problem :

where ˆ y U t = f ϕ ( g θ ( x U t ) . For both the inner loss L inner and outer loss L outer we use the cross entropy loss.

Task generation. Our method can be applied in a standard supervised classification setting where we construct the tasks on the fly as follows. We define a task t as a C -way classification problem. We first select a random subset of C classes from a training domain D train which contains K train classes. For each class we consider the corresponding data points and select a random support set U t with elements ( x U t , y U ) ∈ U and a disjoint random query set Q t with elements ( x Q t , y Q ) ∈ Q t .

Algorithm. In practice we solve the bi-level optimization problem (4) and (5) as follows. In each iteration we sample a batch of T tasks with the associated support and query set as described above. First, we use the samples from the support set S t to fit the linear heads f ϕ by solving the inner optimization problem (5) using stochastic gradient descent for a fixed number of steps. Second, we

use the samples from the query set Q t to update the backbone g θ by solving the outer optimization problem (4) using implicit differentiation [11, 31]. Since the optimal solution of the linear heads ϕ ∗ depend on the backbone g θ , a straightforward differentiation w.r.t. θ is not possible. We remedy this issue by using the approximation strategy of [28] to compute the implicit gradients. The algorithm is summarized in section B.1 of the Appendix.

2.3 Theoretical analysis

We analyze the implications of the proposed minimality and sparse sufficiency principles and show in a controlled setting that they indeed lead to identifiability. As outlined in Figure 2, we assume that there exists a set of independent latent factors z ∼ ∏ d i =1 p ( z i ) that generate the observations via an unknown mixing function x = g ∗ ( z ) . Additionally, we assume that the labels y t for a task t only depend on a subset of the factors indexed by S t ∼ P ( S ) , where S is an index set on z ∈ Z , via some unknown mixing function y t = f ∗ t ( z ) (potentially different for different tasks). We formalize the two principles that are imposed on f ∗ by:

̸

sufficiency : f ∗ t = f ∗ t | S t for S t ∼ p ( S ) 2. minimality : ̸ ∃ S ′ = S t ⊂ S s.t. f ∗ t | S ′ = f ∗ t ,

where f | S t denotes that the input to a function f is restricted to the index set given by S t (all remaining entries are set to zero). (1) states that f ∗ t only uses a subset of features, and (2) states that there are not be duplicate features.

Proposition 2.1. Assume that g ∗ is a diffeomorphism (smooth with smooth inverse), f ∗ satisfies the sufficiency and minimality properties stated above, and p ( S ) satisfies: p ( S ∩ S ′ = { i } ) > 0 or p ( { i } ∈ ( S ∪ S ′ ) -( S ′ ∩ S )) > 0 . Observing unlimited data from p ( X,Y ) , it is possible to recover a representation ˆ z that is an axis aligned, component wise transformation of z .

Remarks: Overall, we see this proposition as validation that in an idealized setting our inductive biases are sufficient to recover the factors of variation. Note that the proof is non-constructive and does not entail a specific method. In practice, we rely on the same constraints as inductive biases that lead to this theoretical identifiability and experimentally show that disentangled representations emerge in controlled synthetic settings. On real data, (1) we cannot directly measure disentanglement, (2) a notion of global ground-truth factors may even be ill-posed, and (3) the assumptions of Proposition 2.1 are likely violated. Still, sparse sufficiency and minimality yield some meaningful factorization of the representation for the considered tasks.

Relation to [47] and [58] : Our theoretical result can be reconnected with concurrent work [47] and can be seen as a corollary with a different proof technique and slightly relaxed assumptions. The main difference is that our feature minimality allows us to also cover the case where the number of factors of variations is unknown, which we found critical in real world data sets (the main focus of our paper). Instead, they only assume sparse sufficiency, which is enough for identifiability if the ground-truth number of factors is known, but is not enough to recover high disentaglement when this is not the case (see Figure 3) and does not translate well to real data, see Table 16 with the empirical comparison in Appendix D.8. Interestingly, their analysis also hints at the fact that our approach also benefits in terms of sample complexity on transfer learning downstream tasks. Our proof technique follows the general construction developed for multi-view data in [58], adapted to our different setting. Instead of observing multiple views with shared factors of variation, we observe a single task that only depend on a subset of the factors.

3 Related work

Learning from multiple tasks and domains. Our method addresses the problem of learning a general representation across multiple and possibly unseen tasks [15, 103] and environments [105, 32, 44, 97, 63, 94, 64] that may be competing with each other during training [61, 91, 83]. Prior research tackled task competition by introducing task specific modules that do not interact during training [67, 101, 80]. While successfully learning specialized modules, these approaches can not leverage synergistic information between tasks, when present. On the other hand, our approach is closer to multi-task methods that aim at learning a generalist model, leveraging multi-task interactions [106, 5]. Other approaches that leverage a meta-learning objective for multi-task learning have been

z

S

x

y

g

*

T

f

*

S t

formulated [18, 81, 50, 9]. In particular, [50] proposes to learn a generalist model in a few-shot learning setting without explicitly favoring feature sharing, nor sparsity. Instead, we rephrase the multi-task objective function encoding both feature sharing and sparsity to avoid task competition.

Similar to prior work in domain generalization, we assume the existence of stable features for a given task [64, 4, 86, 40, 90] and amortize the learning over the multiple environments. Differently than prior work, we do not aim to learn an invariant representation a priori. Instead, we learn sufficient and minimal features for each task, which are selected at test time fitting the linear head on them. In light of [32], one can interpret our approach as learning the final classifier using empirical risk minimization but over features learned with information from the multiple domains.

Disentangled representations. Disentanglement representation learning [8, 33] aims at recovering the factors of variations underlying a given data distribution. [56] proved that without any form of supervision (whether direct or indirect) on the Factors of Variation (FOV) is not possible to recover them. Much work has then focused on identifiable settings [58, 25] from non-i.i.d. data, even allowing for latent causal relations between the factors. Different approaches can be largely grouped in two categories. First, data may be non-independently sampled, for example assuming sparse interventions or a sparse latent dynamics [30, 55, 13, 100, 2, 79, 48]. Second, data may be non-identically distributed, for example being clustered in annotated groups [37, 41, 82, 95, 60]. Our method follows the latter, but we do not make assumptions on the factor distribution across tasks (only their relevance in terms of sufficiency and minimality). This is also reflected in our method, as we train for supervised classification as opposed to contrastive or unsupervised learning as common in the disentanglement literature. The only exception is the work of [47] discussed in Section 2.3.

4 Experiments

We start by highlighting here the experimental setup of this paper along with its motivation.

Synthetic experiments. We first evaluate our method on benchmarks from the disentanglement literature [62, 14, 71, 49] where we have access to ground-truth annotations and we can assess quantitatively how well we can learn disentangled representations. We further investigate how minimality and feature sharing are correlated with disentanglement measures (Section 4.1) and how well our representations, which are learned from a limited set of tasks, generalize their composition. The purpose of these experiments is to validate our theoretical statement, showing that if the assumptions of Proposition 2.1 hold, our methods quantitatively recover the factors of variation.

Domain generalization. On real data sets, we can neither quantitatively measure disentanglement nor are we guaranteed identifiability (as assumptions may be violated). Ultimately, the goal of disentangled representations is to learn features that lend themselves to be easily and robustly transferred to downstream tasks. Therefore, we first evaluate the usefulness of our representations with respect to downstream tasks subject to distribution shifts, where isolating spurious features was found to improve generalization in synthetic settings [19, 58] To assess how robust our representations are to distribution shifts, we evaluate our method on domain generalization and domain shift tasks on six different benchmarks (Section 4.2). In a domain generalization setting, we do not have access to samples coming from the testing domain, which is considered to be OOD w.r.t. to the training domains. However, in order to solve a new task, our method relies on a set labeled data at test time to fit the linear head on top of the feature space. Our strategy is to sample data points from the training

distribution, balanced by class, assuming that the label set Y does not change in the testing domain, although its distribution may undergo subpopulation shifts.

Few-shot transfer learning. Lastly, we test the adaptability of the feature space to new domains with limited labeled samples. For transfer learning tasks, we fit a linear head using the available limited supervised data. The sparsity penalty α is set to the value used in training; the feature sharing parameter β is defaulted to zero unless specified.

Experimental setting. To have a fair comparison with other methods in the literature, we adopt the standard experimental setting of prior work [32, 44]. Hyperparameters α and β are tuned performing model selection on validation set, unless specified otherwise. For comparison with baselines, we substitute our backbone with that of the baseline (e.g. for ERM models, we detach the classification head) and then fit a new linear head on the same data. The linear head module trained at test time on top of the features is the same both for our and compared methods. Despite its simplicity, we report the ERM baseline for comparison in our experiments in the main paper, since it has been shown to perform best in average on domain generalization benchmarks [32, 44]. We further compare with other consolidated approaches in the literature such as IRM [4], CORAL [85] and GroupDRO [73] and include a large and comprehensive comparison with [99, 10, 51, 53, 26, 54, 65, 102, 36, 45] in AppendixD.4. Experimental details are fully described in Appendix C.

4.1 Synthetic experiments

We start by demonstrating that our approach is able to recover the factors of variation underlying a synthetic data distribution like [62]. For these experiments, we assume to have partial information on a subset of factors of variation Z , and we aim to learn a representation ˆ z that aligns with them while ignoring any spurious factors that may be present. We sample random tasks from a distribution T (see Appendix C.3 for details) 5and focus on binary tasks, with Y = { 0 , 1 } . For the DSprites dataset an example of valid task is 'There is a big object on the left of the image' . In this case, the partially observed factors (quantized to only two values) are the x position and size . In Table 1, we show how the feature sufficiency and minimality properties enable disentanglement in the learned representations. We train two identical models on a random distribution of sparse tasks defined on FOVs, showing that, for different datasets [62, 14, 49, 71], the same model without regularizers achieves a similar in-distribution (ID) accuracy, but a much lower disentanglement.

Wethen randomly draw and fix 2 groups of tasks with supports S 1 , S 2 (18 in total), which all have support on two FOVs, | S 1 | = | S 2 | = 2 . The groups share one factor of variation and differ in the other one, i.e. S 1 ∩ S 2 = { i } for some { i } ∈ Z . The data in these tasks are subject to spurious correlations, i.e. FOVs not in the task support may be spuriously correlated with the task label. We start from an overestimate of the dimension of ˜ z of 6 , trying to recover z of size 3 . We train our network to solve these tasks, enforcing sufficiency and minimality on the representation with different regularization degrees. In Figure 3, we show how the alignment of the learned features with the ground truth factors of variations depend on the choice of α, β , going from no disentanglement ( DCI = 27 . 8 ). to good alignment as we enforce more sufficiency and minimality. The model that attains the best alignment ( DCI = 98 . 8 ) uses both sparsity and feature sharing. Sufficiency alone (akin to the method of [47]) is able to select the right support for each task, but features are split or duplicated, attaining lower disentanglement ( DCI = 71 . 9 ). The feature sharing penalty ensures clustering

[47]

100

α,β=0

Ours

80

60

DCI

40

20

0

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

β

in the feature space w.r.t. tasks, ensuring to reach high disentanglement, although it may result in the failure cases, when β is too high ( β > 0 . 25 ).

Disentanglement and minimality are corre-

lated. In the synthetic setting, we also show the role of the feature sharing penalty. Minimizing the entropy of feature activations across mini-batches of tasks results in clusters in the feature space. We investigate how the strength of this penalty correlates well with disentanglement metrics [22] training different models on Dsprites which differ by the value of β . For 15 models trained increasing β from 0 to 0 . 2 linearly, we observe a correlation coefficient with the DCI metric associated to representations compute by each model of 94 . 7 , showing that the feature sharing property strongly encourages disentanglement. This confirms again that sufficiency alone (i.e. enforcing sparsity) is not enough to attain good disentanglement.

Task compositional generalization. Finally, we evaluate the generalization capabilities of the features learned by our method by testing our model on a set of unseen tasks obtained by combining tasks seen during training. To do

100

90

80

70

2

60

Accurac

50

40

30

20

10

NoReg

Ours

0

IS|=2

IS|=3

IS|=4

IS|=5

Task support

A

this, we first train two models on the AbstractDSprites dataset using a random distribution of tasks, where we limit the support of each task to be within 2 (i.e. | S | = 2 ). The models differ in activating/deactivating the regularizers on the linear heads. Then, we test on 100 tasks drawn from a distribution with increasing support on the factors of variation ( | S | = 3 , | S | = 4 , | S | = 5) , which correspond to composition of tasks in the training distribution; see Figure 4, with the accompaning Table 9 in Appendix D.

4.2 Domain Generalization

In this section we evaluate our method on benchmarks coming from the domain generalization field [32, 93, 70] and subpopulation distribution shifts [73, 44], to show that a feature space learned with our inductive biases performs well out of real world data distribution.

Subpopulation shifts. Subpopulation shifts occur when the distribution of minority groups changes across domains. Our claim is that a feature space that satisfies sparse sufficiency and minimality is more robust to spurious correlations which may affect minority groups, and should transfer better to new distributions. To validate this, we test on two benchmarks Waterbirds [73], and CivilComments [44] (see Appendix C.1).

For both, we use the train and test split of the original dataset. In Table 4, last row, we report the results on the test set of Waterbirds for the different groups in the dataset (landbirds on land, landbirds on water, waterbirds on land, and waterbirds on water, respectively). We fit the linear head

on a random subset of the training domain, balanced by class, repeat 10 times and report accuracy and standard deviation on test. For CivilComments we report the average and worst accuracy in Figure 5, where we compare with ERM and groupDRO [73]. While performing almost on par w.r.t. ERM, our method is more robust to spurious correlation in the dataset, showing the higher worst group accuracy. Importantly, we outperform GroupDRO, which uses information on the subdomain statistics, while we do not assume any prior knowledge about them. Results per group are reported in the Appendix (Table 11).

DomainBed. We evaluate the domain generalization performance on the PACS , VLCS and OfficeHome datasets from the DomainBed [32] test suite (see Appendix C.1 for more details). On these datasets, we train on N -1 and leave one out for testing. Regularization parameters α and β are tuned according to validation sets of PACS , and used accordingly on the other dataset. For these experiments we use a ResNet50 pretrained on Imagenet [17] as a backbone, as done in [32] To fit the linear head we sample 10 times with different samples sizes from the training domains and we report the mean score and standard deviation. Results are reported in Table 4, showing how enforcing sparse sufficiency and minimality leads consistently to better OOD performance. Comparisons with 13 additional baselines is in Appendix D.4.

Camelyon17. The model is trained according to the original splits in the dataset. In Table 3 we report the accuracy of our model on in-distribution and OOD

100

90

80

70

60

50

40

ERM

30

DRO

20

Ours

10

0

splits, compared with different baselines [84, 4]. Our method shows the best performance on the OOD test domains. The intuition of why this happens is that, due to minimality, we retain more features which are shared across the three training domains, giving less importance to the ones that are domain-specific (which contain the spurious correlations with the hospital environmental informations). This can be further enforced at test time, as we show in the ablation in Appendix D.9, trading off in distribution performance for OOD accuracy.

4.3 Few-shot transfer learning.

We finally show the ability of features learned with our method to adapt to a new domain with a small number of samples in a few-shot setting. We compare the results with ERM in Table 2, averaged by domains in each benchmark dataset. The full scores for each domain are in Appendix D.5 for 1-shot, 5-shot, and 10-shot setting, reporting the mean accuracy and standard deviations over 100 draws. Our approach achieves consistently higher accuracy than ERM, showing the better adaptation capabilities of our minimal and sufficently sparse feature space.

4.4 Additional results

In Appendix D we report a large collection of additional results, including comparison with 14 baseline methods on the domain shift benchmarks (D.4), a qualitative and quantitative analysis on the minimality and sparse sufficiency properties in the real setting (D.2), a favorable additional comparison on meta learning benchmarks, with 6 other baselines including [47](D.8), an ablation study on the effect of clustering features at test time (D.9), and a demonstration on the possibility to obtain a task similarity measure as a consequence of our approach (D.7).

5 Conclusions

In this paper, we demonstrated how to learn disentangled representations from a distribution of tasks by enforcing feature sparsity and sharing. We have shown this setting is identifiable and have validated it experimentally in a synthetic and controlled setting. Additionally, we have empirically shown that these representations are beneficial for generalizing out-of-distribution in real-world settings, isolating spurious and domain specific factors that should not be used under distribution shift.

Limitations and future work : The main limitation of our work is the global assumption on the strength of the sparsity and feature sharing regularizers α and β across all tasks. In real settings these properties of the representations might need to change for different tasks. We have already observed this in the synthetic setting in Figure 3, where when β > 0 . 25 features cluster excessively and are unable to achieve clear disentanglement and do not generalize well. Future work may exploit some level of knowledge on the task distribution (e.g. some measure of distance on tasks) in order to tune α, β adaptively during training, or to train conditioning on a distribution of regularization parameters as in [21], enabling more generalization at test time. Another limitation is in the sampling procedure to fit the linear head at test time: sampling randomly from the training set (balanced by class) may not be enough to achieve the best performance under distributions shifts. Alternative sampling procedures, e.g. ones that incorporate knowledge on the distribution shift if available (as in [43]), may lead to better performance at test time.

Acknowledgments and Disclosure of Funding

Marco Fumero and Emanuele Rodolà were supported by the ERC grant no.802554 (SPECGEO), PRIN 2020 project no.2020TA3K9N (LEGO.AI), and PNRR MUR project PE0000013-FAIR. Marco Fumero and Francesco Locatello were partially at Amazon while working at this project. We thank Julius von Kügelgen, Sebastian Lachapelle and the anonymous reviewers for their feedback and suggestions.

References

Julius Adebayo, Justin Gilmer, Michael Muelly, Ian J. Goodfellow, Moritz Hardt, and Been Kim. Sanity checks for saliency maps. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada , pages 9525-9536, 2018.

Kartik Ahuja, Karthikeyan Shanmugam, Kush R. Varshney, and Amit Dhurandhar. Invariant risk minimization games. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event , volume 119 of Proceedings of Machine Learning Research , pages 145-155. PMLR, 2020.

Isabela Albuquerque, João Monteiro, Mohammad Darvishi, Tiago H Falk, and Ioannis Mitliagkas. Generalizing to unseen domains via distribution matching. ArXiv preprint , abs/1911.00804, 2019.

Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. ArXiv preprint , abs/1907.02893, 2019.

Jinze Bai, Rui Men, Hao Yang, Xuancheng Ren, Kai Dang, Yichang Zhang, Xiaohuan Zhou, Peng Wang, Sinan Tan, An Yang, et al. Ofasys: A multi-modal multi-task learning system for building generalist models. ArXiv preprint , abs/2212.04408, 2022.

Peter Bandi. Camelyon17 dataset.

Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European conference on computer vision (ECCV) , pages 456-473, 2018.

Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence , 35(8):1798-1828, 2013.

Luca Bertinetto, João F. Henriques, Philip H. S. Torr, and Andrea Vedaldi. Meta-learning with differentiable closed-form solvers. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019 . OpenReview.net, 2019.

Gilles Blanchard, Aniket Anand Deshmukh, Ürun Dogan, Gyemin Lee, and Clayton Scott. Domain generalization by marginal transfer learning. The Journal of Machine Learning Research , 22(1):46-100, 2021.

Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe LlinaresLópez, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. ArXiv preprint , abs/2105.15183, 2021.

Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. ArXiv preprint , abs/1903.04561, 2019.

Johann Brehmer, Pim De Haan, Phillip Lippe, and Taco Cohen. Weakly supervised causal representation learning. ArXiv preprint , abs/2203.16437, 2022.

Chris Burgess and Hyunjik Kim. 3d shapes dataset. https://github.com/deepmind/3dshapesdataset/, 2018.

Rich Caruana. Multitask learning. Machine learning , 28(1):41-75, 1997.

Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett, editors, Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain , pages 2172-2180, 2016.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei Li. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 2009), 20-25 June 2009, Miami, Florida, USA , pages 248-255. IEEE Computer Society, 2009.

Guneet Singh Dhillon, Pratik Chaudhari, Avinash Ravichandran, and Stefano Soatto. A baseline for few-shot image classification. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 . OpenReview.net, 2020.

Andrea Dittadi, Frederik Träuble, Francesco Locatello, Manuel Wuthrich, Vaibhav Agrawal, Ole Winther, Stefan Bauer, and Bernhard Schölkopf. On the transfer of disentangled representations in realistic settings. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

Lucas Dixon, John Li, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Measuring and mitigating unintended bias in text classification. 2018.

Alexey Dosovitskiy and Josip Djolonga. You only train once: Loss-conditional training of deep networks. In International conference on learning representations , 2020.

Cian Eastwood and Christopher K. I. Williams. A framework for the quantitative evaluation of disentangled representations. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings . OpenReview.net, 2018.

M. Everingham, L. Van Gool, C. K. I. Williams, J. Winn, and A. Zisserman. The PASCAL Visual Object Classes Challenge 2007 (VOC2007) Results. http://www.pascalnetwork.org/challenges/VOC/voc2007/workshop/index.html.

Li Fei-Fei, Rob Fergus, and Pietro Perona. Learning generative visual models from few training examples: An incremental bayesian approach tested on 101 object categories. In 2004 conference on computer vision and pattern recognition workshop , pages 178-178. IEEE, 2004.

Marco Fumero, Luca Cosmo, Simone Melzi, and Emanuele Rodolà. Learning disentangled representations via product manifold projection. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 3530-3540. PMLR, 2021.

Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The journal of machine learning research , 17(1):2096-2030, 2016.

Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence , 2(11):665-673, 2020.

Zhengyang Geng, Xin-Yu Zhang, Shaojie Bai, Yisen Wang, and Zhouchen Lin. On training implicit models. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 24247-24260, 2021.

Ian J. Goodfellow, Quoc V. Le, Andrew M. Saxe, Honglak Lee, and Andrew Y. Ng. Measuring invariances in deep networks. In Yoshua Bengio, Dale Schuurmans, John D. Lafferty, Christopher K. I. Williams, and Aron Culotta, editors, Advances in Neural Information Processing Systems 22: 23rd Annual Conference on Neural Information Processing Systems 2009. Proceedings of a meeting held 7-10 December 2009, Vancouver, British Columbia, Canada , pages 646-654. Curran Associates, Inc., 2009.

Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations , 2020.

Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation . SIAM, 2008.

Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

Irina Higgins, Loïc Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings . OpenReview.net, 2017.

Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey. Meta-learning in neural networks: A survey. arXiv preprint arXiv:2004.05439 , 2020.

Ziniu Hu, Zhe Zhao, Xinyang Yi, Tiansheng Yao, Lichan Hong, Yizhou Sun, and Ed H Chi. Improving multi-task generalization via regularizing spurious correlation. ArXiv preprint , abs/2205.09797, 2022.

Zeyi Huang, Haohan Wang, Eric P Xing, and Dong Huang. Self-challenging improves crossdomain generalization. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part II 16 , pages 124-140. Springer, 2020.

Aapo Hyvärinen, Hiroaki Sasaki, and Richard E. Turner. Nonlinear ICA using auxiliary variables and generalized contrastive learning. In Kamalika Chaudhuri and Masashi Sugiyama, editors, The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan , volume 89 of Proceedings of Machine Learning Research , pages 859-868. PMLR, 2019.

Ali Jalali, Sujay Sanghavi, Chao Ruan, and Pradeep Ravikumar. A dirty model for multi-task learning. In J. Lafferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta, editors, Advances in Neural Information Processing Systems , volume 23. Curran Associates, Inc., 2010.

Hicham Janati, Marco Cuturi, and Alexandre Gramfort. Wasserstein regularization for sparse multi-task regression. In Kamalika Chaudhuri and Masashi Sugiyama, editors, The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan , volume 89 of Proceedings of Machine Learning Research , pages 1407-1416. PMLR, 2019.

Yibo Jiang and Victor Veitch. Invariant and transportable representations for anti-causal domain shifts, 2022.

Ilyes Khemakhem, Diederik P. Kingma, Ricardo Pio Monti, and Aapo Hyvärinen. Variational autoencoders and nonlinear ICA: A unifying framework. In Silvia Chiappa and Roberto Calandra, editors, The 23rd International Conference on Artificial Intelligence and Statistics, AISTATS 2020, 26-28 August 2020, Online [Palermo, Sicily, Italy] , volume 108 of Proceedings of Machine Learning Research , pages 2207-2217. PMLR, 2020.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings , 2015.

Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. ArXiv preprint , abs/2204.02937, 2022.

Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton Earnshaw, Imran S. Haque, Sara M. Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. WILDS: A benchmark of in-the-wild distribution shifts. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 5637-5664. PMLR, 2021.

David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning , pages 58155826. PMLR, 2021.

Tejas D. Kulkarni, William F. Whitney, Pushmeet Kohli, and Joshua B. Tenenbaum. Deep convolutional inverse graphics network. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett, editors, Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada , pages 2539-2547, 2015.

Sébastien Lachapelle, Tristan Deleu, Divyat Mahajan, Ioannis Mitliagkas, Yoshua Bengio, Simon Lacoste-Julien, and Quentin Bertrand. Synergies between disentanglement and sparsity: a multi-task learning perspective. ArXiv preprint , abs/2211.14666, 2022.

Sébastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie E Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ica. In Conference on Causal Learning and Reasoning , pages 428-484. PMLR, 2022.

Yann LeCun, Fu Jie Huang, and Leon Bottou. Learning methods for generic object recognition with invariance to pose and lighting. In Proceedings of the 2004 IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2004. CVPR 2004. , volume 2, pages II-104. IEEE, 2004.

Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto. Meta-learning with differentiable convex optimization. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019 , pages 10657-10665. Computer Vision Foundation / IEEE, 2019.

Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales. Learning to generalize: Metalearning for domain generalization. In Proceedings of the AAAI conference on artificial intelligence , volume 32, 2018.

Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales. Deeper, broader and artier domain generalization. In IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29, 2017 , pages 5543-5551. IEEE Computer Society, 2017.

Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot. Domain generalization with adversarial feature learning. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 5400-5409, 2018.

Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European conference on computer vision (ECCV) , pages 624-639, 2018.

Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M. Asano, Taco Cohen, and Stratis Gavves. CITRIS: causal identifiability from temporal intervened sequences. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , volume 162 of Proceedings of Machine Learning Research , pages 13557-13603. PMLR, 2022.

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA , volume 97 of Proceedings of Machine Learning Research , pages 4114-4124. PMLR, 2019.

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. A sober look at the unsupervised learning of disentangled representations and their evaluation. J. Mach. Learn. Res. , 21:209:1-209:62, 2020.

Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event , volume 119 of Proceedings of Machine Learning Research , pages 6348-6359. PMLR, 2020.

Aurelie C. Lozano and Grzegorz Swirszcz. Multi-level lasso for sparse multi-task regression. In Proceedings of the 29th International Conference on Machine Learning, ICML 2012, Edinburgh, Scotland, UK, June 26 - July 1, 2012 . icml.cc / Omnipress, 2012.

Chaochao Lu, Yuhuai Wu, José Miguel Hernández-Lobato, and Bernhard Schölkopf. Invariant causal representation learning for out-of-distribution generalization. In International Conference on Learning Representations , 2022.

Zvika Marx, Michael T Rosenstein, Leslie Pack Kaelbling, and Thomas G Dietterich. Transfer learning with an ensemble of background tasks. Inductive Transfer , 10, 2005.

Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement testing sprites dataset. https://github.com/deepmind/dsprites-dataset/, 2017.

John Miller, Rohan Taori, Aditi Raghunathan, Shiori Sagawa, Pang Wei Koh, Vaishaal Shankar, Percy Liang, Yair Carmon, and Ludwig Schmidt. Accuracy on the line: on the strong correlation between out-of-distribution and in-distribution generalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 7721-7735. PMLR, 2021.

Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In Proceedings of the 30th International Conference on Machine Learning, ICML 2013, Atlanta, GA, USA, 16-21 June 2013 , volume 28 of JMLR Workshop and Conference Proceedings , pages 10-18. JMLR.org, 2013.

Hyeonseob Nam, HyunJae Lee, Jongchan Park, Wonjun Yoon, and Donggeun Yoo. Reducing domain gap by reducing style bias. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 8690-8699, 2021.

Boris N. Oreshkin, Pau Rodríguez López, and Alexandre Lacoste. TADAM: task dependent adaptive metric for improved few-shot learning. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada , pages 719-729, 2018.

G. Parascandolo, N. Kilbertus, M. Rojas-Carulla, and B. Schölkopf. Learning independent causal mechanisms. In Proceedings of the 35th International Conference on Machine Learning, PMLR 80:4036-4044 , 2018.

Ji Ho Park, Jamin Shin, and Pascale Fung. Reducing gender bias in abusive language detection. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing , pages 2799-2804, Brussels, Belgium, 2018. Association for Computational Linguistics.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, highperformance deep learning library. In Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019.

Jielin Qiu, Yi Zhu, Xingjian Shi, Florian Wenzel, Zhiqiang Tang, Ding Zhao, Bo Li, and Mu Li. Are multimodal models robust to image and text perturbations? ArXiv preprint , abs/2212.08044, 2022.

Scott E. Reed, Yi Zhang, Yuting Zhang, and Honglak Lee. Deep visual analogy-making. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett, editors, Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada , pages 1252-1260, 2015.

Bryan C Russell, Antonio Torralba, Kevin P Murphy, and William T Freeman. Labelme: a database and web-based tool for image annotation. International journal of computer vision , 77(1):157-173, 2008.

Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. ArXiv preprint , abs/1911.08731, 2019.

Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event , volume 119 of Proceedings of Machine Learning Research , pages 8346-8356. PMLR, 2020.

Ruslan Salakhutdinov. Deep learning. In Sofus A. Macskassy, Claudia Perlich, Jure Leskovec, Wei Wang, and Rayid Ghani, editors, The 20th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD '14, New York, NY, USA - August 24 - 27, 2014 , page 1973. ACM, 2014.

Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. ArXiv preprint , abs/1910.01108, 2019.

Jürgen Schmidhuber. Learning factorial codes by predictability minimization. Neural computation , 4(6):863-879, 1992.

Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE , 109(5):612-634, 2021.

Anna Seigal, Chandler Squires, and Caroline Uhler. Linear causal disentanglement via interventions. ArXiv preprint , abs/2211.16467, 2022.

Amanpreet Singh, Ronghang Hu, Vedanuj Goswami, Guillaume Couairon, Wojciech Galuba, Marcus Rohrbach, and Douwe Kiela. FLAVA: A foundational language and vision alignment model. ArXiv preprint , abs/2112.04482, 2021.

Jake Snell, Kevin Swersky, and Richard S. Zemel. Prototypical networks for few-shot learning. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA , pages 4077-4087, 2017.

Peter Sorrenson, Carsten Rother, and Ullrich Köthe. Disentanglement by nonlinear ICA with general incompressible-flow networks (GIN). In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 . OpenReview.net, 2020.

Trevor Standley, Amir Roshan Zamir, Dawn Chen, Leonidas J. Guibas, Jitendra Malik, and Silvio Savarese. Which tasks should be learned together in multi-task learning? In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event , volume 119 of Proceedings of Machine Learning Research , pages 9120-9132. PMLR, 2020.

Baochen Sun, Jiashi Feng, and Kate Saenko. Correlation alignment for unsupervised domain adaptation. In Domain Adaptation in Computer Vision Applications , pages 153-171. Springer, 2017.

Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision , pages 443-450. Springer, 2016.

Victor Veitch, Alexander D'Amour, Steve Yadlowsky, and Jacob Eisenstein. Counterfactual invariance to spurious correlations: Why and how to pass stress tests, 2021.

Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In 2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017 , pages 5385-5394. IEEE Computer Society, 2017.

Oriol Vinyals, Charles Blundell, Tim Lillicrap, Koray Kavukcuoglu, and Daan Wierstra. Matching networks for one shot learning. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett, editors, Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain , pages 3630-3638, 2016.

Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The caltech-ucsd birds-200-2011 dataset. 2011.

Zihao Wang and Victor Veitch. A unified causal view of domain invariant representation learning. ArXiv preprint , abs/2208.06987, 2022.

Zirui Wang, Zihang Dai, Barnabás Póczos, and Jaime G. Carbonell. Characterizing and avoiding negative transfer. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019 , pages 11293-11302. Computer Vision Foundation / IEEE, 2019.

Martin Wattenberg, Fernanda Viégas, and Ian Johnson. How to use t-sne effectively. Distill , 1(10):e2, 2016.

Florian Wenzel, Andrea Dittadi, Peter V. Gehler, Carl-Johann Simon-Gabriel, Max Horn, Dominik Zietlow, David Kernert, Chris Russell, Thomas Brox, Bernt Schiele, Bernhard Schölkopf, and Francesco Locatello. Assaying out-of-distribution generalization in transfer learning. In Neural Information Processing Systems , 2022.

Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre-Alvise Rebuffi, Ira Ktena, Krishnamurthy Dvijotham, and Ali Taylan Cemgil. A fine-grained analysis on distribution shift. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 . OpenReview.net, 2022.

Matthew Willetts and Brooks Paige. I don't need u: Identifiable non-linear ica without side information. ArXiv preprint , abs/2106.05238, 2021.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. ArXiv preprint , abs/1910.03771, 2019.

Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gontijo Lopes, Ari S. Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , volume 162 of Proceedings of Machine Learning Research , pages 23965-23998. PMLR, 2022.

Jianxiong Xiao, James Hays, Krista A. Ehinger, Aude Oliva, and Antonio Torralba. SUN database: Large-scale scene recognition from abbey to zoo. In The Twenty-Third IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2010, San Francisco, CA, USA, 13-18 June 2010 , pages 3485-3492. IEEE Computer Society, 2010.

Shen Yan, Huan Song, Nanxiang Li, Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training. arXiv preprint arXiv:2001.00677 , 2020.

Weiran Yao, Yuewen Sun, Alex Ho, Changyin Sun, and Kun Zhang. Learning temporally causal latent processes from general temporal data. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 . OpenReview.net, 2022.

Lu Yuan, Dongdong Chen, Yi-Ling Chen, Noel Codella, Xiyang Dai, Jianfeng Gao, Houdong Hu, Xuedong Huang, Boxin Li, Chunyuan Li, Ce Liu, Mengchen Liu, Zicheng Liu, Yumao Lu, Yu Shi, Lijuan Wang, Jianfeng Wang, Bin Xiao, Zhen Xiao, Jianwei Yang, Michael Zeng, Luowei Zhou, and Pengchuan Zhang. Florence: A new foundation model for computer vision. ArXiv preprint , abs/2111.11432, 2021.

Marvin Zhang, Henrik Marklund, Nikita Dhawan, Abhishek Gupta, Sergey Levine, and Chelsea Finn. Adaptive risk minimization: Learning to adapt to domain shift. Advances in Neural Information Processing Systems , 34:23664-23678, 2021.

Yu Zhang and Qiang Yang. An overview of multi-task learning. National Science Review , 5(1):30-43, 2018.

Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 million image database for scene recognition. IEEE transactions on pattern analysis and machine intelligence , 40(6):1452-1464, 2017.

Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey. 2021.

Jinguo Zhu, Xizhou Zhu, Wenhai Wang, Xiaohua Wang, Hongsheng Li, Xiaogang Wang, and Jifeng Dai. Uni-perceiver-moe: Learning sparse generalist models with conditional moes. ArXiv preprint , abs/2206.04674, 2022.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a novel regularization technique for disentangled representation learning. This regularizer is based on two principles formalized in the paper: 1) sparsity and 2) minimality. It assumes access to multiple tasks during training where each task uses only a subset of features, and different tasks share some of the features. The paper also provides theoretical results showing that these two principles are enough for the identification of disentangled factors. The proposed approach is validated, showing the significance of both regularization terms. Empirical results show the benefits of the proposed feature learning technique over standard ERM training.\n\nSTRENGTHS:\n-  The proposed two formalized principles (sparse sufficiency and minimality) and their sufficiency for disentangled representation learning are valuable contributions to the field.\n-  The empirical study assesses the quality of the learned representations beyond disentanglement, showing the practical benefits of the proposed approach.\n-  The overall idea and motivation are sound.\n\nWEAKNESSES:\n1. The paper lacks clarity; certain details and symbols' introductions are missing for better understanding (more in the question section). The writing can be improved.\n    - For example, proposition 2.1 does not clearly state what S' and i are. Also, the symbol $f_t^*$ is not introduced. The proposition also misses subsequent discussion for better understanding. \n    - The text has a number of mistakes. For example, L50: \"a**n** entropy\", L57: \"co~~h~~exist\", L77: \"give**n**\", L88: the formula has an extra parenthesis, L110: subscript $t$ is missing for $y$s.\n2. The empirical evaluation is insufficient and lacks comparisons with other disentangled representation learning methods as well as few-shot/meta-learning methods.\n    - Tab. 1 lacks other repr. learning methods.\n    - Tab. 2 and 4 lack other few-shot/meta-learning methods. The main paper refers to Appendix for the latter mentioning a 10% improvement on MiniImageNet. However, Tab.15 in Appendix does not support the claim, with MetaOptNet being the best.\n    - Fig. 5 lacks other domain adaptation methods.\n3. It is not entirely clear how the sharing regularizer in (3) corresponds to the minimality principle introduced in Sec. 2.3. Also, it is unclear how exactly it favors clustering. Further clarifications are needed. (see questions for a more concrete question)\n\nQUESTIONS:\n1. How the DCI metric is defined exactly? The paper refers to [22], which has three separate metrics. Is DCI a combination of those? If yes, how are they combined?\n2. The Waterbirds dataset has only two classes; how is task distribution defined?\n3. Is regularization used or not when learning a novel test task? If yes, where do multiple tasks need for it come from?\n4. For the experiments on DomainBed, it is stated that pretrained RN50 is used. How the proposed method influences the features learned in this case? Do you finetune it?\n5. In the definition of the two principles, what is the difference between $f_t^*/f^*/f^*_{S_t}$? They are not properly introduced.\n6. See 3 from weaknesses.\n    - How does the sharing term \"forces features to cluster across tasks and therefore be reused by different tasks, when useful\" (L92-93)? For example, consider 3 FoVs $f_i$ and 3 binary classification tasks $t_i$ (with logits for simplification, i.e., M=1) corresponding to each combination of only 2 features with weights +1/-1, i.e., $y = \\sigma(f_{S_t[1]} - f_{S_t[2]})$, $S_t \\in \\{(1, 2), (1, 3), (2, 3)\\}$ In this case, $\\tilde{\\phi}_m = 2/6 = 1/3$ (removing the extra $1/TC$ term, see minor comments), resulting in a uniform distribution that results in high entropy (maximum possible). On the other hand, the sparsity and minimality are satisfied. It seems that the distribution should be over $t,c$ instead of $m$.\n    - How does it correspond to the minimality principle introduced in Sec 2.3? The latter seems to suggest that no (smaller? see 9 below) subset of features can solve the task. Why having an L1 loss is not enough to satisfy this? If using fewer features is possible, the corresponding weights will be set to 0 following the L1 regularizer.\n7. While the overall idea is sound, it is unclear why the meta-learning framework is necessary. Can the same technique be applied to the standard multi-task setting where different linear heads are learned simultaneously with a shared feature extractor?\n\n\nMinor:\n8. Why scaling by $1/TC$ in L88? $\\tilde{\\phi}_m$ is not a distribution in this case.\n9. in L129 do you mean $S' \\subset S_t$? Otherwise can include $S_t$ and be sufficient for classification.\n\nLIMITATIONS:\nLimitations are well-addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the problem of learning disentangled representations in a setting with access to multiple supervised tasks, each of which is assumed to depend only on a subset of the underlying factors of variation. While the factors of variation are never directly observed, two principles are incorporated to allow their identifiability; i) sparse sufficiency (“sparsity”) ii) minimality - encouraging features to be shared / reused across tasks. The proposed approach is demonstrated to achieve a higher degree of disentanglement on synthetic data compared to previous methods, as well as better out-of-distribution generalization in real image domains.\n\nSTRENGTHS:\n1. The idea of studying disentanglement in the context of generalization to downstream tasks is well motivated, especially in real image domains.\n\n2. While the “sparsity” regularization already introduced for disentanglement in previous works, including the concurrent work of [47], the second “minimality” term seems novel. Including this term is well justified for real world problems when the number of FoV is unknown, and the results in the paper highlight its effectiveness.\n\n3. The experimental section contains both quantitative disentanglement evaluation on synthetic datasets, as well as appropriate assessment for domain generalization in real image data, where the ground truth for the factors of variation is unknown.\n\nWEAKNESSES:\n1. The idea of multitask meta-learning still looks a bit “over” complicated to me. I think the paper would benefit from including the following simple baseline for comparison: combining all the supervision from all tasks into a single multilabel classification task. The feature layer before the classification head should still include the sparsity term. Can the authors elaborate on the ERM baseline? Is this one the “simple” baseline I just mentioned?\n\n2. The dataset considered in [47] is 3Dshapes. However, I do not see a comparison on this benchmark against [47]. Could the authors report the DCI evaluation of the proposed method and [47] on 3Dshapes?\n\nQUESTIONS:\n1. There are several generative approaches to learning disentangled representations which are worth mentioning e.g. [1, 2]. However, these methods rely on weak supervision e.g. observing some of the factors of variation for some of the images. Given the binary labels constructed in this paper for some of the dataset, I wonder if there is a way to “quantify” the amount of supervision provided to the proposed model. This is not trivial as no direct labels are used but as long as many of “binary” labels are constructed the effective level of supervision is increased. Could the authors elaborate on this perspective?\n\n[1] Locatello et al. “Disentangling factors of variations using few labels”. In ICLR, 2020.\n\n[2] Gabbay et al. “An Image is Worth More Than a Thousand Words: Towards Disentanglement in the Wild”. In NeurIPS 2021.\n\n2. DCI metric is a combination of three different metrics: Disentanglement, Completeness and Informativeness. However, the results in this paper are presented as a single number. Is this the average of the three? or just the Disentanglement? I request the authors to present the metrics separately in the quantitative evaluation as they may shed more light on the performance of each of the considered methods. \n\n3. Minor typos:\n* Line 45 - w.r.t to\n* Line 57 - “cohexist” → “coexist”\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe foundation of the proposed work is a form of regularization motivated by the problem of disentanglement.  A supervised task is deconstructed into randomly generated sub-tasks, and the same trained representations are used for all sub-tasks (followed by a linear head specific to each sub-task).  Two regularizers are then introduced on the feature activations: one which encourages sparsity (L1) and one which encourages minimality by penalizing the entropy of the distribution of feature activations.  \n\nThe disentanglement motivation is such: If we assume each sub-task only requires information from a subset of some ground truth disentangled factors of variation of the data, it is argued that the training procedure and regularization will eventually disentangle the factors of variation completely.  Disentangled factors confer various advantages in representation learning, one of which being increased robustness to shifts in the data distribution because spurious factors are able to be separated from the rest. This advantage forms a large part of the proposed work.  After establishing the intended effect of the regularization on synthetic data where generative factors are known, the focus shifts to real data where domain shift is a challenge.\n\nSTRENGTHS:\nThe method is sensible while being nontrivial.  It is presented well, and the experimental results, tackling either idealized disentanglement or practical domain shift, cover a lot of ground.  It is interesting to leverage the ability to decompose label information as a way to condition representations and encourage structural properties of extracted factors of variation.\n\nWEAKNESSES:\nThere are two perspectives with which to view the introduced regularization and training procedure.  The first is in terms of what it accomplishes with respect to disentanglement in synthetic scenarios, and the second in terms of its utility to address domain shift in real datasets (disregarding disentanglement as unrealistic).  The support for both perspectives could be strengthened.  \n\nFrom the disentanglement perspective, it’s not clear how much everything depends on Y containing full information about the generative factors.  If Y depends on fewer factors than X (as is more often the case), how can Proposition 2.1 hold?  For example, if the original labels (before any task deconstruction) have no “angle” information from DSprites, no sub-task should extract angle information, and even with unlimited samples from p(x,y), the “angle” generative factor will be ignored as spurious.  If Y depends on the factors in an entangled way, e.g. I(xpos;Y)=I(ypos;Y)=0, but I(xpos,ypos;Y)>0, then there is no way to create sub-tasks out of Y that depend on xpos but not ypos.  Even with unlimited samples there should be no disentanglement of the xpos and ypos factors.  Thus it seems that the disentanglement is predicated on having I(Y;Z)=H(Z), ie all factor information in Y; if true, this is quite a severe constraint.  More clarity around this point would be helpful, either through a constructed example on the synthetic dataset or elaboration about what’s expected of Y when discussing Proposition 2.1.\n\nFrom the real world, pragmatism perspective, the domain shift results are numerous yet somewhat insubstantial.  The comparisons in the main text are predominantly to ERM, which does not seem particularly convincing.  In the research cited to support using ERM as a baseline (Gulrajani and Lopez-Paz 2020), the competitiveness seemed to only come from re-tuning for ERM anew everything pertaining to training and model selection -- but I gather all conditions match for the ERM as the proposed method.  Given two methods, identical except that one has two additional hyperparameters to tune, the one with additional hyperparameters should prevail.  There are many more domain shift results in the Supp, though they are thrown together without much in the way of a unifying perspective.  A significant number of baselines are compared to in App. D4,5, yet the results are presented without exposition.  Curiously, the proposed method is a significant outlier, performing several standard deviations above all fourteen baselines, on a single category of each experiment of D4 (S in CLSV, C in PACS, C in ACPR), yet the performance on other categories is always middling or even the worst.  Can the authors speak to this?\n\nQUESTIONS:\nCan the authors please elaborate on the sub-task generation process for the real datasets?  Take Camelyon17 as an example -- how were sub-tasks generated out of this binary classification problem?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to learn disentangled representations by leveraging knowledge from a diverse set of supervised task objectives. Specifically, the paper shows that disentangled representations can emerge without ground-truth annotations of the factors of variation under two assumptions: 1) sparse sufficiency, and 2) minimality. Practically, the paper implements a meta-learning objective for representation learning with the sparse sufficiency and minimality constraints. The authors also evaluate the learnt representations for out-of-distribution (OOD) generalization.\n\nSTRENGTHS:\n- The paper is very well written and easy to follow.\n- The paper appropriately situates the proposed method among all the related works, including highly related concurrent work.  \n- The paper has multiple significant contributions that would benefit the research community — it presents the theory as well shows experimentally that appropriate inductive biases can sufficiently recover the factors of variation from the data.\n- The disentanglement results are promising. The paper contains ablation studies to show the importance of individual components such as the feature sharing property.\n- The proposed method is evaluated for domain generalization and domain shift tasks on six different benchmarks, achieving convincing results.\n\nWEAKNESSES:\nIn my opinion, the paper does not have any major weaknesses. To me, barring a few grammatical mistakes, it already reads better than a lot of accepted papers. The only thing I could suggest is to make the caption of Figure 1 a bit more descriptive so that it conveys a full picture by itself and doesn’t rely on the main text.\n\nQUESTIONS:\nI think the paper is already in a good and sufficient state to be accepted as it is. I do not have any further questions for the authors.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}