SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We then 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} \cap S_{2} = \{i\}$  for some  $\{i\} \in 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  $\tilde{\mathbf{z}}$  of 6, trying to recover  $\mathbf{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  $\alpha, \beta$ , 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
Figure 3: Role of minimality: We plot the DCI metric of a set of models (red dots) trained on fixed tasks from DSprites: Training without regularizers leads to no disentanglement (green). Enforcing sparsity alone (yellow, akin to [47]) achieves good disentanglement ( $DCI = 71.9$ ), but some features may be split or duplicated. Enforcing both minimality and sparse sufficiency (magenta) attains the best DCI (98.8). When  $\beta$  is too high ( $>0.25$ ) activated features collapses into few clusters with respect to tasks. For complete results and experiments on additional datasets see Table 8 and Figures 6, 7 in Appendix.
in the feature space w.r.t. tasks, ensuring to reach high disentanglement, although it may result in the failure cases, when  $\beta$  is too high ( $\beta > 0.25$ ).
Table 1: Enforcing disentanglement: DCI [22] disentanglement scores and ID accuracy on test samples for a model trained without enforcing sufficiency and minimality (top row), and model with the regularizers activated (bottom row). While attaining similar performance on accuracy, the model with the activated regularizer always show higher disentanglement. See Table 7 for additional scores.
Dsprites	3Dshapes	SmallNorb	Cars
No reg (DCI,Acc)	(16.6,94.4)	(44.4,96.2)	(16.5,96.1)	(60.5,99.8)
α,β (DCI,Acc)	(69.9,95.8)	(87.7, 95.8)	(55.8,95.6)	(92.3,99.8)
Disentanglement and minimality are correlated. 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  $\beta$ . For 15 models trained increasing  $\beta$  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
Figure 4: Task compositional generalization: Mean accuracy over 100 random test tasks reported for group of tasks of growing support (second, third, fourth column) for a model trained without inductive biases (blue, attaining  $DCI = 29.4$ ) and enforcing them (orange,  $DCI = 59.4$ ). The latter show better compositional generalization resulting from the properties enforced on the representation. Exact values are reported in Table 9 in Appendix.
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 accompanying Table 9 in Appendix D.
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
Table 2: Quantitative results for few-shot transfer learning, with our method consistently outperforming ERM across all sample sizes and data sets.
N-shot/Algorithm	OOD accuracy (averaged by domains)
1-shot	PACS	VLCS	OfficeHome	Waterbirds
ERM	80.5	59.7	56.4	79.8
Ours	81.5	68.2	58.4	88.4
5-shot
ERM	87.1	71.7	75.7	79.8
Ours	88.3	74.5	77.0	87.6
10-shot
ERM	87.9	74.0	81.0	84.2
Ours	90.4	77.3	82.0	89.2
Table 3: Quantitative evaluation on Camelyon17: we report accuracy both on ID and OOD splits. Our approach achieves significantly higher validation and test OOD accuracy.
Validation(ID)	Validation (OOD)	Test (OOD)
ERM	93.2	84	70.3
CORAL	95.4	86.2	59.5
IRM	91.6	86.2	64.2
Ours	93.2 ±0.3	89.9±0.6	74.1±0.2
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  $\alpha$  and  $\beta$  are tuned according to validation sets of PACS, and used accordingly on the other dataset. For these experiments we use a ResNet50 pretrained onImagenet [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
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.
Figure 5: Quantitative results on CivilComments: we report the accuracy on test averaged across all demographic groups (left group), and the worst group accuracy, on the right. Our method (green) performs similarly in terms of average accuracy and outperforms in terms of worst group accuracy, without using any knowledge on the group composition in the training data. For exact values and error estimates, see Table 10 in the Appendix.
Table 4: Results for domain generalization on DomainBed. Our approach achieves consistently higher average OOD generalization, outperforming ERM in all cases except one.
Dataset/Algorithm	OOD accuracy (by domain)
PACS	S	A	P	C	Average
ERM	77.9 ± 0.4	88.1 ± 0.1	97.8 ± 0.0	79.1 ± 0.9	85.7
Ours	83.1 ± 0.1	86.7± 0.8	97.8 ± 0.1	83.5 ± 0.1	87.5
VLCS	C	L	V	S	Average
ERM	97.6± 1.0	63.3 ± 0.9	76.4 ± 1.5	72.2 ± 0.5	77.4
Ours	98.1± 0.2	63.4± 0.5	78.2 ± 0.7	73.9± 0.8	78.4
OfficeHome	C	A	P	R	Average
ERM	53.4± 0.6	62.7 ± 1.1	76.5 ± 0.4	77.3 ± 0.	67.5
Ours	56.3± 0.1	66.7 ± 0.7	79.2± 0.5	81.3 ± 0.4	70.9
Waterbirds	LL	LW	WL	WW	Average
ERM	98.6 ± 0.3	52.05 ± 3	68.5 ± 3	93 ± 0.3	81.3
Ours	99.5 ± 0.1	73.0 ± 2.5	85.0 ± 2	95.5 ± 0.4	90.5
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 sufficiently 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).
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  $\alpha$  and  $\beta$  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  $\beta > 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  $\alpha$ ,  $\beta$  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.
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 Kugelgen, Sebastian Lachapelle and the anonymous reviewers for their feedback and suggestions.
[1] 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.
[2] 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.
[3] 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.
[4] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. *ArXiv preprint*, abs/1907.02893, 2019.
[5] 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.
[6] Peter Bandi. Camelyon17 dataset.
[7] 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.
[8] 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.
[9] 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.
[10] 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.
[11] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. ArXiv preprint, abs/2105.15183, 2021.
[12] 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.
[13] Johann Brehmer, Pim De Haan, Phillip Lippe, and Taco Cohen. Weakly supervised causal representation learning. *ArXiv preprint*, abs/2203.16437, 2022.
[14] Chris Burgess and Hyunjik Kim. 3d shapes dataset. https://github.com/deepmind/3dshapes-dataset/, 2018.
[98] 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.
[99] Shen Yan, Huan Song, Nanxiang Li, Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training. arXiv preprint arXiv:2001.00677, 2020.
[100] 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.
[101] 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.
[102] 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.
[103] Yu Zhang and Qiang Yang. An overview of multi-task learning. National Science Review, 5(1):30-43, 2018.
[104] 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.
[105] Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey. 2021.
[106] 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.