Enhancing Compositional Generalization via Compositional Feature Alignment
Authors: Haoxiang Wang
Haozhe Si
Huajie Shao
Han Zhao
Venue: Accepted by TMLR ()

Abstract: Real-world applications of machine learning models often confront data distribution shifts, wherein discrepancies exist between the training and test data distributions. In the common multi-domain multi-class setup, as the number of classes and domains scales up, it becomes infeasible to gather training data for every domain-class combination. This challenge naturally leads the quest for models with Compositional Generalization (CG) ability, where models can generalize to unseen domain-class combinations. To delve into the CG challenge, we develop CG-Bench, a suite of CG benchmarks derived from existing real-world image datasets, and observe that the prevalent pretraining-finetuning paradigm on foundational models, such as CLIP and DINOv2, struggles with the challenge. To address this challenge, we propose Compositional Feature Alignment (CFA), a simple two-stage finetuning technique that i) learns two orthogonal linear heads on a pretrained encoder with respect to class and domain labels, and ii) fine-tunes the encoder with the newly learned head frozen. We theoretically and empirically justify that CFA encourages compositional feature learning of pretrained models. We further conduct extensive experiments on CG-Bench for CLIP and DINOv2, two powerful pretrained vision foundation models. Experiment results show that CFA outperforms common finetuning techniques in compositional generalization, corroborating CFA's efficacy in compositional feature learning.

Published in Transactions on Machine Learning Research (05/2024)
Enhancing Compositional Generalization via Compositional Feature Alignment
Haoxiang Wang
kwang264 @illinois. eduu
Deportment of Electrical and Computer Engineering University of Mllmois Urbona- Champaign
Haozhe
haozhes3@hllinois.cco edu
Department Electrical and Computer Engineering University of Mlinois Urbena- Chann-ng
Huajie Shao
hshwa @wm. edu
Department of Computer Science Willimm and Mary
Han Zhao
hanshao hanshao@iliinois.coo edu
Department Computer Science
Universily of Ilimois Cathnn--Camppigh
Reviewed on OpenReview: https: // wptp//veee...eee.eemm/// net/ forum? id= k3d5C0YufK
Abstract
Real-world applications of machine learning models often confront data distribution shifts, wherein discrepancies exist between the training and test data distributions. In the com- mon multi-domain multi-class setup, as the number of classes and domains scales up, it becomes infeasible to gather training data for every domain-class combination. This chal- lenge naturally leads the quest for models with Compositional Generalization (CG) ability, where models can generalize to unseen domain-class combinations. To delve into the CG challenge, we develop CG-Bench, a suite of CG benchmarks derived from existing real-world image datasets, and observe that the prevalent pretraining finetuning paradigm on founda- tional models, such as CL.P and DINOv2, struggles with the challenge. To address this challenge, we propose Compositional Feature Alignment (CFA), a simple two-stage fine- tuning technique that i) learns two orthogonal linear heads respect to class and domain labels, and ii) fine-tunes the encoder with the newly learned head frozen. We theoretically and empirically justify that CFA encourages compositional feature learning of pretrained models. We further conduct extensive experiments on CG-Bench for CLIP and DINOv2, two powerful pretrained vision foundation models. Experiment results show that CFA outperforms common finetuning techniques in compositional generalization, corroborating CFA's efficacy in compositional feature learning. The code is released at Aliireennlg
pretrained encoder with
1
Introduction
Over the past decade, machine learning has emerged as transformative technology, driving advancements across various domains such as computer vision (He et al.) 2016a b), natural language processing (Devlin) et al.) 2019 Brown et al. 2020), biology (Jumper et al.] 2021), etc. These innovations have been fueled by the development of increasingly sophisticated models, the availability of largo-scale datasots, and the growth of computational power. However, crucial obstacle persists in applying machine learning models to real-world scenarios: their performance tends to degrade significantly when confronted with data

---

Published in Transactions on Machine Learning Research (05/2024)
Compositional Generalization
ID
COD
Domain Generalization
Quickdraw Fanting Infogroon clpart
Domains
NeeESS0iio 14600- VOR
PRESSONAA MAGIC YUR
Rea
Sketch
Beach
Bush
Eraser
Hurricnc
Sock
Stichhes
Trumeet
Zebra
Beach
Bush
Eraser
Hurricane
Sock
Sttchee
Trumpet
Zebra
Classes
Classes
Figure 1: Compositional generalization (CG) domain generalization (DG). Masked entries are unseen domain-class combinations, while unmasked ones exist in the training dataset.
distribution shifts (Koh al.) 2021 Gulrajani & Lopez-Paz 2021 Santurkar et al. 2021), where the data distribution during testing differs from that used in training.
In an effort to overcome this problem, the machine learning research community has turned its attention to (OOD) generalization, with the goal of developing models that are robust under data distribution shifts. Existing research primarily investigates various types of data distribution shifts, such as domain generalization (Gulrajani & Lopez-Pazy Paz, 2021 Koh et al. 2021h, subpopulation shift. (Santurkar
al.
2021 Yang et al. 2023), input corruption (Deng et al.) 2009), and spurious correlation (Sagawa et al. [2020]. While generalizing to these different type of distribution shifts has garnered significant attention, there exists another realistic yet understudied challenge in OOD generalization: compositional generalization (CG).
Within the multi-domain, multi-class context, assumes we have E domains (i.c., environments) and K classes, leading to E x K pairs of domain and class combinations, which could be formulated as elements in a E x K matrix as shown in Figure [ In domain generalization (DG), the learner has access to data from all the classes and all the domains and aims to make predictions on data from a new, unseen domain. Nonetheless, in real-world scenarios, given the large number of categories, e.g., 1000 classes in ImageNet, one cannot always collect the complete data from all the domains. Put in another word, the training data might not cover all the possible combinations of domains and classes, as represented by each cell in the matrix in Figure [ This sparsity pattern becomes especially pronounced when the number of classes or environments is large because collecting comprehensive training data for each combination becomes a formidable task. In such cases, a key challenge arises: can the model generalize to unseen domain-class combinations? This is the compositional generalization (CG) challenge we aim to tackle this work.
The CG challenge manifests ubiquitously across various real-world applications. distribution shifts in certain existing DG datasets, such iWildCam (Beery et al. 2020 Koh al. 2021) are more accurately characterized by CG than DG. Moreover, finetuning pretrained (foundation) models struggles to tackle the CG challenge. This emphasizes the need for the machine learning community to recognize and address this emerging distribution shift challenge with innovative solutions.
For example, data
find that the widely-used method of
Our Contributions In our attempt to tackle this challenge, we draw inspiration from existing lines of rescarch such as invariant risk minimization (IRM) (Arjovsky et al. 2019) and invariant-feature subspace recovery (ISR) (Wang et al. 2022). In particular, Wang et al. (2022) showed that under certain structural conditions in the data generative process, post- post-processing methods via subspace projection can effectively learn invariant features that can generalize across unseen domains from the same data generative process but under different interventions on non-causal factors. Empirically, we find that the learned features (i.c., the outputs of the last hidden layer) conform to compositional structure where the subspace of domain-
2

---

Published in Transactions on Machine Learning Research (05/2024)
Domain Domain 2
Domain Domain Cless2 [Unseen]
Class
Figure 2: Illustration of desired compositional feature structure for compositional generalization.
related urat is orghoglnal to that of class-related features, the corresponding model can generalize across unknown domain-class pairs. Motivated by this observation, to induce features that match this compositional
structure, we introduce a two-stage finetuning approach termed Compositional Feature Aligmment (CFA),
which is also inspired by recent progress in the literature of neural collapse (Papyan et al. 2020 Zhu et al. 2021 Yang et al. 2022).
More specifically, upon the features given by the encoder, we construct two heads, one for predicting the target label of interest and the other for predicting the domain index. Note that the two-head architecture is not new, and has been widely used in domain adversarial neural networks (Ganin & Lempitsky 2015 Zhao et al. 2018). However, different from domain adversarial neural networks where adversarial training through minimax optimization needed, our proposed method is computationally lightweight and can be divided into two stages.
CFA first identifies proper compositional feature structure via a two-head regularized linear probing (i.e., training linear heads with the encoder frozen). Subsequently, the encoder undergoes finetuning with the heads being frozen. Leveraging tools from the neural collapse literature, we theoretically prove that CFA can effectively align features with the compositional feature structure under mild assumptions. Furthermore, we construct a synthelic Color-CIFAR datasel to examine CFA empirically and observe that CFA can indeed align features with the desired compositional feature structure.
To facilitate the studies of compositional generalization, we curate a suite of benchmarks for the composi- tional generalization challenge, building on four real-world image datasets: OfficeHome (Venkateswara al., 2017). DomainNet (Peng et al.) 2019), WILDD--WildCam (Beery et al. 2020 andd WIL.DS-FMon (Christic
et al. 2018). We consider two powerful pretrained vision encoders, CLIP [Radford et al. 2021 and DI- NOv2 1Qqub et al.] 2023), with ViT-B (Dosovitskiy et al. (202)) architecture, and apply different finetuning methods to them, including linear probing, full finetuning, LP-FT (Kumar et al. 2022), reweighting and our proposed CFA. Extensive experimental results on CG-Suite show that CFA-finctuned models can indood generalize unseen domain-class combinations better than other finetuning moth We hope that the curated CG-Suite can facilitate future research on compositional generalization.
2 Compositional Feature Alignment
The key to the CG challenge is to identify and encode the compositional relationship between classes and domains by learning the features. Hence, it is inportant to first understand what kind of feature struc- tures are desired for compositional generalization. To this end, we first provide a formal definition of the compositional feature structure, and then explain our motivations behind the definition.
Definition 1 (Compositional Feature Structure). For any input x from class y E {1...., (1...., F), its fealure e Rdd salisfies the compositional feature structure. long as can be decomposed as:
., K} and domain e
3

---

Published in Transactions on Machine Learning Research (05/2024)
Stage Multi-Head I. Linear Probing
Stage 2
Fine-tuning with Frozen Orthogonal Heads
Class
Domain
Class
Domain
Two-Head Linear Classifier
Wi
Orthagonal
W
A
W
W
Trainable Module
Encoder (Pretrained)
S
Frozen Module
Figure 3: Illustration of our proposed method, Compositional Feature Alignment (CFA).
Class Feature: MM4, e R Domain Fealure: 22 NW2 mg) R2
Total Feature: R
Zuoise
where Zuoise E Rd-da-da represent noise features irrelevant to classes and domains, and R. Iddd is full-rank orthonormal matris. Note that p4 Eq have dependence on y; while 16 En rely e.
Fig. 2 provides a visualization of this compositional feature structure for simple setup of 2 classes and 3 domains, where one domain-class combination is absent in the training set. It is evident that class features and domain fealures exist in orthogonal subspaces, as required by Definition [ In this case, a linear classifier that exclusively utilizes class features and disregards domain and noise features can effectively generalize the unscen domain-class combination.
It is noteworthy that even with the perfect alignment of learned features to this compositional structure all training data, there is no guarantee that features from unseen domain-class combinations will still conform to this structure.
2.1 Method: Training with Frozen Orthogonal Heads under Normalization
Though we have defined ideal feature structure for compositional generalization, the neural features produced by pretrained models may not align with this structure. In this section we introduce our method to encourage the learned features to follow the above structure. At a high level, the proposed method contains two stages of finetuning from pretrained models. We include an auxiliary linear head for domain prediction, complementing the pre-existing class prediction head. The first stage involves multi-label linear probing, which adjusts the two heads to achieve an optimal compositional feature structure. In the second stage, we fine-tune the encoder, while keeping the two heads frozen. The final product is finetuned encoder that generates features in alignment with the predetermined compositional feature structure from stage one. The two stages, diagrammatically represented in Fig. 3 are detailed below, followed by a discussion on our rationale behind the amgorithm design.
Stage 1 (Multi-Label Linear Probing). We begin with a pretrained encoder, 2(-) that maps inputs to dimensional features of unit norm, where > K + E.). We construct two linear heads without bias terms, denoted by W1 Rxdd and W₂ RExd Keeping Φ(-) frozen, we train these heads with two cross-entropy loss terms, which take into account both class and domain labels. An orthogonalily constraint ensures Wi and Wa spon orthogonal subspaces.
Mathematically, the optimization objective of the first stage con be written as
min W.
Σ
K LOE(BI Wld(z), v) X - W24(x), e)
(1)
(s.y,e)EDerain
31,82, and Wi E u(d)* W2 U(d) WTW. -0
(2)
where Dursin represents the training set, leE is the cross-entropy loss 31, 32 are inverse temperature param- eters (also called as logit scale in CLIP [Radford al.) 2021)) ) U(d) denotes the set of d-dimensional unit vectors, and O slands for the zero malria.
4

---

Published in Transactions on Machine Learning Research (05/2024)
Stage 2 (finctuning with Frozen Heads). We then freeze the trained W1 and W2, and proceed to fine-ture the encoder (()) end-lo-end, using the same mulli-label cross-enlropy loss function.
The optimization objective of this finetuning stage can be expressed as
min N
1
Σ
} / W1D()). + X - W20(W2 K
(3)
(a.Meeee
The following discussion explains our motivations and reasoning underlying the algorithm design:
Recent work on neural collapse indicate that during the train- ing of a multi-class classifier using cross-entropy loss, freezing the linear head according to simplex struc- ture can guide the features to align with the frozen head (Zhu et al. 2021 Yang ct al. 2022). This observa- lion implies that the features of data in class y to collapse in the direction of the row vector of the classifier corresponding to class y. In addition, we empirically observe that the head-freezing technique does not compromise the model's performance compared to end-to-end finetuning. We include the details regarding this experiment in an ablation study in Appendix C Inspired by these, we devise the two-stage strategy where the Stage determines the optimal head weights for our compositional feature structure, and the Stage 2 finetunes the encoder with frozen head weights to align the features with the feature structure. Linear Probing Two Orthogonal Heads: Unlike research work on neural collapse that focus exclusively on class prediction (Papyan et al.) 2020 Zhu et al., 2021 Yang al.j 2022). our work also accounts for the effects of domains, outlined in Definition [ We therefore introduce an auxiliary head, W2, for domain prediction, alongside the original class prediction head denoted as W1- Thus, for a sample I, the encoder, W1, and W2 predict the class and domain labels based on the feature ((x). Definition [ implicitly poscs an orthogonality requirement on domain-related and class-related features since R is an orthonormal matrix. To meet. this feature orthogonality requirement, we impose an orthogonality constraint on the two heads (i.e., W,W, = 0).
Freezing Head for Feature Alignment:
Normalizing Features and Weights to Address Data Imbalance: While Zhu et al 2021 | Yang et al. (0022) provide a theoretical justification for head freezing, their theory assumes class-balance training data. In the case of data imbalance, Thrampoulddis ct al.) (2022) shows that the head and features may become misaligned. Upon reviewing the technical details of Thrampoulidis el al. (2022 we find that this misalignment can be rectified by normalizing features and head weights to a hyper-sphere. This normalization ensures constant norms for features and head weights, thereby ensuring alignment (cf. Theorem [] Consequently, we assume that the features produced by the encoder & are also normalized to unit norm, which is a common practice in modern vision model pretraining such CLIP (Radford et al. (021),, SimCLR (Chen et al. 2020), MoCo (He et al., 2020), DINO [Caron et al. 2021 Additionally, we impose the head normalization constraint W1 E U(d) K W E U(d)* in (2) a technique already employed in CLIP (Radford et al. 2021
2.2 Theoretical Guarantee
In the algorithm above, Stage I is relatively simple, comprising joint minimization problem over two lincar heads. In contrast, Stage 2 is more complex, it optimizes neural encoder using two heads under two cross-entropy loss terms. We offer theoretical justification for Stage 2 below, demonstrating that the finetuned encoder can indeed align features with the two frozen orthogonal heads produced by Stage
thereby creating a feature structure that meets the requirement in Definition [
In line with recent research on neural collapse (Mixon et al., 2020 Fang et al. 2021 Zhu et al. 2021 Thrampoulidis et al., 2022), we adopt the unconstrained fcature model (UFM) or layer- aayer-pceled model, where ((x) is treated as free optimization variable in Rdd for every input. x.
We denote Z [o(z1)...., class labels, and environment labels, respectively. In the context of the unconstrained feature model, the
E Y [yi....,unl. and E= [e....., eN] as the stack of features,
Recennt studies show that weight. normalization for the linear head (without blas) can enhance the performance of fine-tuned CLII Goyal al. 2022 Wang et 2023
5

---

Published in Transactions on Machine Learning Research (05/2024)
optimization objective of Stage 2 is transformed to:
1 min KN (CE(BI Wiz, + A ANN ( W2Z, E) s.t. Z ZEU(d)
(4)
Theorem 1 (Feature Alignment). Assuming the feature dimension is no smaller than K + E, and training dala exists for each cluss and domain (though not necessarily for each domain-class combination), and Wi and W2 are normalized and span orthogonal subspaces such that W1 e U(d)* W2 E e((d) and W1W = 0.
Additionally, we assume $1, B2 are sufficiently large. The global minimam of (4 results in the following: for any i E [N], denole 2, as the i-th column vector of Z, we have
2 = WT az + +WJbe.
(5)
where ag E RRK is vector depending on the class label ys, and be, E RE is vector relying on the domain label 62.
This theorem, intuitively, demonstrates that upon optimizing (4) to a global minimum, for any training sample from class y in environment e, its corresponding feature =* can be decomposed as linear combination of two vectors depending on y and e, respectively, and the two vectors live in orthogonal feature subspaces. This indicates that the learned features conform to a compositional feature structure satisfying Definition [ The complete proof is found in Appendix B where we leverage theoretical tools from Thrampoulidis et al.
h022 in the proof.
3 Empirical Studies
3.1 Benchmark Development of
We create CG-Bench, compositional generalization benchmark built on four datascts previously designed for DG research: Office-Home (Venkaleswara et al. 2017), DomainNet (Peng et al.] 2019), and iWildCam (Beery et al.] 2020) & FMoW (Christie et al. 2018)) from the WILDS benchmark (Koh et al. 2021 These datasots span a wide range of application scenarios from common object recognition in web-crawled images to wildlife recognition from camera traps and building recognition from satellite images. Due the page limit, we eoraorate on the motivation for creating the CG-bench and the curation procedure by taking DomainNet example. Additional details regarding benchmark curation can be found in Appendix A
DomainNet (Peng et al., 2019) consists of objects in different art styles. In the DomainNet dataset, there are K = 345 classes and F. 6 domains: {Clipart, Infograph, Painting, Quickdraw, Real Image, Sketch). In addressing the DG challenge, prior research using DomainNet typically employed a leave-one-out vross-validation strategy. This involves training the model on data from five of the domains and subsequently evaluating its performance on the sixth, omitted domain. In addition to the DG task, it's worth noting that the CG challenge is intrinsically present within DomainNet. To underscore this point, we carried out preliminary experiment using CLIP.
CG Challenge in DomainNet We randomly divide DomianNet into training and evaluation sets, with an 80:20 split. A CLIP model is fully fine-tuned on this training data, and evaluated on validation data from all domain-class combination. We also gathered zero-shot accuracy of the CLIP model for comparison. As a final step, we examined the test accuracies for each domain-class combination, correlating them with the count of their respective training data samples. We only focus on hard domain-class combinations that the zero-shot accuracy is below 30%, and visualize evaluation results over these combinations inFigure 4a Firstly, we notice that certain domain-class combinations possess minimal or even no training samples (e.g., some combinations have only 2 images and neither of them is sampled into the training sct). This observation aligns with our considered CG scenario. Within this CG context, both the fine-tuned and zero-shol models encounter difficulties achieving high test accuracy when the training data for a specific domain-class combination is insufficient. This leads us conclude that the CG challenge inherently present in DomainNet, and current zero-shot and fine-tuned models fail address. Consequently, we are motivated to establish benchmark for methodical investigation of this challenge.
6

---

Published in Transactions on Machine Learning Research (05/2024)
Test Accuracy Training Data Size (DomainNet)
Number Domains Combinations
Zero-Shot Test Accuracy Distribution (DomainNet)
0.6
Model Type
200
Median Tast Accuracy
0.5
finetuned
zero-shot
150
100
0.2
0.1
50
0.0
[8. 16)
[16, 32)
[32, 64) [64. 128) [128. 256)
0.0
0.2
0.4
0.6
0.8
1.0
Training Size domain-class combination)
Accuracy
(a) Test. Accuracy va. Training Data Size
(b) Zero-shot Test Accuracy Distribution
Figure Test accuracy statistics for each domain-class combination on DomainNet dataset. Left: The test accuracy compared with the number of training data. Points median accuracy while the shading area is bounded by 25% and 75% quantiles. Right: The number of domain-class combinations at different zero-shot test accuracies.
Benchmark Curation Setup Consider data belonging to E domains (i.e., environments) and classes, resulting in an Ex K matrix of domain-class combinations (such as the demo shown in Fig. [ A binary mask Mía E {0, 1) BxX is applied to this matrix indicate in-distribution domain-class combinations, ensuring that each row or column contains both and so that the training data includes all classes and domains, while some domain-class combinations may be absent. The complementary binary mask (Mond 1 Maa represents OOD domain-class combinations.
CG Curation of DomainNet. We form F. x K class-environment matrix for DomainNet and evaluate the zero-shot accuracy for every domain-class combination. The resulting data distribution is visualized in Figure 4b We designate the combinations that fall within the lowest 20% of zero-shot accuracies the out-of-distribbution (OOD) set, while the top 80% constitute the in-distribution (ID) set. To ensure comprehensive representation, each row/colume contains at least one entry from the ID set. The ID data is then further segregated into training set and an ID validation set at 9:1 ratio. Meanwhile, the OOD data is divided between OOD validation and test sets. When evaluating a model trained on the main training dataset, we assess its performance across the ID validation, OOD validation, and OOD test subsets. For this benchmark, key metric is the average top-1 accuracy for both the ID and OOD sets.
3.2 Evaluations
Baselines We take the OpenAl's VIT-B CLIP (Radford et al. 2021 and the Meta's VIT-B/14 DINOv2
2023) as the pretrained model for each benchmark and implement three finetuning strategies as the baseline; full finetuning, linear probing then full finetuning (LP-FT) and reweighting. LP-FT [Kumar et al. 2022 a simple two-stage finetuning strategy that addresses the problem that full finetuning can distort pretrained features and often underperforms linear-probing on out-of-distribution data. Note: Our proposed approach bears similarities with LP-FT (Kumar al.) 2022). both methods start with linear probing stage and then proceed with finetuning. However, two critical differences are: our approach employs two heads with a multi-label cross-cntropy loss and imposes an orthogonality constraint during the linear probing stage, and ii) we keep the heads frozen during the second finetuning stage. Reweighling (Buda) et al. 2018) balances the number samples from each group in each batch during training. It robust group shifts in OOD generalization tasks. For the CG benchmarks, we implemented two versions of reweighting strategies: Reweight-Ei E, which does ro-sampling according to the domain labels, and Reweight- Y F, which balances according to the domain-class combinations.
Oquab et al.

---

Published in Transactions on Machine Learning Research (05/2024)
Table 1: Test accuracy(%%) and Fl-macro scores (%) different methods on CG-Bench. Bold values mark the highest accuracy F1 score. The OOD accuracy of FMoW is the worst region accuracy. The highest accuracy values and those within a range of 0.2 are in bold.
OfficeHome
DomsinNet
IWidCCam
FMoW
Model
Methods
II) Ace
OOD Arc
II) Arr
OOD Acc
ID Aec
OOD Acc
ID F1
0OD F1
Are:
OOD Ace
Zero-Shot
89.2
50.3
61.7
6.6
13.7
6,9
11.7
9.2
20.4
18.8
Probing
90.9
72.7
14.4
27.8
Fine-Tuning
94.3
51.0
82.0
7.5
74.5
16.5
13.8
22.2
65.8
38.7
Fine-Toning (WISE)
93.7
52.5
76.4
8.7
67.0
13.7
31.6
17.0
19.5
10.6
LP-FT
$1.5
17.0
40.2
LP-FT (WISE)
93.0
42.8
79.1
5.3
74.4
18.2
44.5
28.7
56.6
36.3
CLIP
Reweight-E
91.0
51.9
81.2
7.4
75.3
17.2
45.2
24.3
62.4
41.8
Rewoight-E (WiSK)
75.9
13.7
18.0
Reweight-YxE
93.7
52.2
81.0
7.6
72.2
17.4
11.6
30.0
58.0
11.1
Reweight-Yai (WISE)
93.4
53.4
75.5
8.5
55.9
15.1
22.7
42.1
37.5
81.6
18.3
31.0
CFA (WISE)
93.1
56.9
76.5
9.2
74.6
19.7
45.6
32.5
53.5
36.6
Fine-Tuming
91.8
38.6
82.4
5.3
76.4
14.4
17.6
18.3
66.1
38.4
Linear Probing
93.3
40.0
75.4
4.8
77.0
19.6
50.7
27.9
45.5
25.5
LP-FT
93.1
38.2
82.5
5.1
77.6
23.1
30.8
67.1
37.1
LP-FT (WISE)
94.0
39.7
81.6
6.2
77.9
22.3
53.2
31.0
61.0
33.7
DINO<2
Reweight-E
91.2
38.9
81.8
5.2
76.9
13.1
17.9
62.3
38.4
Reweight-Yxtt
91.3
81.5
5.3
72.2
17.4
41.6
30.0
57.5
37.6
CFA
92.8
39.2
82.6
5.6
78.1
22.8
52.5
30.8
67.2
38.5
CPA (WiSE)
93.1
40.4
79.6
6.4
78.2
23.8
52.6
33.4
59.8
34.3
Postprocessing with WiSE-FT (Wortsman et al.) 2022) After finetuning using the three baseline methods and our proposed CFA, we also apply WiSE-FT Wortsman et al. (022)) with a 0.5 to postprocess the model, which just takes an average of initial and finetuned model parameters in the parameter space. It has been shown that Wisc-FT can improve the model performance (especially OOD performance) in some cases WWortsman et al.] 2022). Note: For the CLIP experiments, interpolate the finetuned model with the zero-shot CLIP encoder and classification head. For the DINOv2 experiments, since there is no zero- shot classification head available, we interpolate the finetuned models with the linear probing/stage- CFA results. Consequently, we do not perform WiSE-FT on full finetuning and reweighting basclincs since they do have linear-probed heads.
Implementation of CFA Empirically, we make small modifications to Stage to divide it into two steps: i) Train W2 on the domain labels with reweighting until it converges with W1 fixed to its zero-shot weight; ii) then train W1 on the class label with reweighting, and encourage the orthogonality between Wi and the fixed W2 with lo regularization term, W213- We also use reweighting for the class labels when performing the linear probing for LP-FT for fair comparison. Following Sec. 2.1 we normalize the row vectors of W, and W2 to the unit norm and normalize the outputs of & also to the unit norm. In constrain W1 and W2 to subspace determined
addition, in the linear probing stage of CFA and LP-FT, by the zero-shot linear classifier of CLIP, as we find it can improve the final OOD performance. Besides, we find that in Stage 2 it is empirically sufficient to train the encoder with very small À value, which is a loss coefficient in (3). and we deploy X = 0 in Stage 2 to reduce compute cost. In both Stage and Stage 2 we use the AdamW (Loshchilov Hutter 2017 optimizer with a cosine annealing scheduler [Loshchilov
Hutter 2016). More details and hyperparameters can be found in Appendix
Empirical Conclusions The results of our empirical experiments are presented in Table 1 Obscrving the results, we conclude that: a) Compared with full finetuning, LP-FT and reweighting, our CFA can improve the performance of pretrained models on OOD data for compositional generalization; b) WiSE-FT can further improve the OOD performance of all methods in most cases (when WiSE-FT fails, it will fail on both ID and OOD); c) While CFA enjoys superior OOD performance, its in-distribution (ID) performance is maintained around the same level of full finetuning, which is a desired property. Also, notice that,
8

---

Published in Transactions on Machine Learning Research (05/2024)
Donwin
Class
Class
Damahn
Damain
Dosain Ninseest
Claso
Class Zcnol (Ukeeen)
62
0.3
0.4
8.0
04
-16
(a) Pretrained Features
(b) CFA-Fimetuned Features
Figure 5: Visualization of the features for CLIP ViT-B/16 before and after the finetunning of CFA. Left: Feature extracted using pre-trained CLIP ViT-B/16 image encoder. Right: Feature extracted using CFA- finetuned CLIP ViT-B/16 image encoder.
although our CFA increases the performance of models OOD data in CG tasks, there is still a gap between its ID and OOD performance, indicating that CG is quite a challenging task and needs future algorithm advancement to be further addressed.
3.3 Feature Visualization Real-World Data
To empirically show that our two-stage algorithm promotes the encoder's ability to learn the compositional feature, we conduct a feature visualization study for the CLIP ViT-B/16 image encoder on the DomainNet [Peng et al. 2019) dataset. Specifically, we take the encoder both before and after finetuning with CFA, visualizing features for domains and 3 classes. This resulted in 6 unique domain-class combinations (5 out of which present in the training set). The visualization in Figure 5b clearly shows that the features finetuned with CFA conform to compositional feature structure. In contrast, the features of the original CLIP model (Figure 5ab do not exhibit this structure. This visualization not only demonstrates the feature alignment ability of CFA but also provides further evidence of its effectiveness for large neural networks trained on real-world data.
3.4 Partial Availability of Domain Labels
In real-world applications, training samples may not be provided with domain labels, which may affect the applicability of CFA. We divide this problem into two scenarios: i) Domain labels are partially available; ii) Domain labels are completely unavailable.
Corresponding experiments on the Office-Home dataset are designed and conducted to demonstrate that the CFA remains effective even in the absence of domain labels. Under the first condition, we conduct experiments by only using 10%, 20%, and 50% of the domain labels for Stage-1 to learn the linear head for finetuning in Stage-2. We repeat each experiment 3 times using different subsets of the data and present the average result to mitigate the randomness caused by subsampling. For the second scenario, we propose that we can leverage the zeroshot ability of CLIP to predict the domain labels. To be specific, we first manually design a set of labels that can describe the domains of the data (for the experiment we just use the original domain names). Then, we perform the zeroshot classification on the training data and the domain labels using CLIP model. Finally, we take the predicted domain labels as ground truth domain labels and use them in 1t We adopt the hyrrpaaraeetess shown in the manuscript. and present the mean accuracy of CFA
9

---

Published in Transactions on Machine Learning Research (05/2024)
Table 2: CFA with partial availability of domain labels. Experiments are conducted on Offce-Home
Method
Domain label ratio
ID
ID (WiSE)
OOD
OOD (WiSE)
100%
(Original)
94
93.1
54.3
56.9
50%
94
93.3
53.8
56.9
CFA
20%
93.9
93.1
53.8
56.7
10%
94.0
93.2
53.4
55.9
0% (CLIP Prodict)
94.1
93.0
52.0
53.6
Finetune LP-FT
0%
94.3
93.7
51
52.5
0%
93.5
93.0
43.9
42.8
Table 3: Domain label prediction with partial training label availability. Experiments are conducted on Homee
Domain Label Ratio
0%
(Zero-Shot)
10%
20%
50%
100%
Avg. #Data per Domain 0 Accuracy
308
616
1541
3082
61.5
82.0
84.6
85.5
86.3
and WiSE-FT over 3 seeds on the Office-Home datasel. From results in Table
2
we can see our CFA works
well as domain labels are partially available.
The mild reliance of CFA's performance on domain label availability beneficial for the practical application of this method but may raise questions about why the reliance is so mild. We believe there are two main reasons: the number of domains is small (4-6 domains for each 4 dataset of CG-Bench), so the available data per domain is relatively abundant, and ii) domain labels are casier to predict than class labels since image style or background information is highly indicative of the domain label, and these visual features are easy to capture by neural networks.
We clarify this by conducting a simple ablation study: in the Office-Home dataset with 4 domains, we fit linear classifiers to CLIP features for different amounts of domain labels (1%, 2%,..., 100% of the training data), and show the test prediction accuracy for domains below. Fach experiment repeated three times with different sampling seeds, and mean accuracy is reported in Table
From the results, can observe that as the domain labeling ratio increases from 10% to 100%, the domain prediction accuracy modestly improves from 82.0% to 86.3%, indicating relatively small enhancement. This suggests that the domain label is indeed quite easy to predict, which could explain why our CFA can work well with a partial availability of domain labels. Furthermore, the zero-shol prediction accuracy of the CLIP ViT-B/16 16 model for domain labels 61.5%, significantly higher than a random guess (25%). This outcome supports the notion that CFA with CLIP-predicted domain labels can also improve over vanilla finctuning.
Class Label Availability As for class labels, we do require their full availability for the training set, since our work focuses on supervised finetuning of pretrained encoders. Mcanwhile, although CLIP is not explicitly supervised by domain and class labels, the texts in the image-text, pairs (used for CLIP pretraining) provide abundant class and environment information. Furthermore, CLIIP's unsatisfactory zero-shot performance our CG-bench indicates the need for supervised finetuning to improve its effectiveness. On the other hand, DINOv2 is self-supervised pretraining method, thus it needs to be supervised finetuned before applying to downstream classification tasks. Overall, the above statement wants to justify that the class label availability cannot be waived for the compositional generalization task considered in our paper.
Ablation Studies We conducted ablation studies on i) the choice of normalized head v.s. unconstrained head, ii) the choice of frozen head v.s. trainable head, and iii) the loss coefficient in Stage on CLIP model. Due to the page limit, we defer results and details to Appendix C Also, we discuss the training stability, overhead and performance gain of CFA in Appendix C
4 Related Works
OOD Generalization In OOD generalization, the model is trained on labeled data from a limited number of known domains, and the goal is to improve the performance of the models so that they can better generalize to previously unseen or new test domains (Blanchard et al. 2011 A common approach to tackle
10

---

Published in Transactions on Machine Learning Research (05/2024)
OOD generalization is domain-invariant learning, which aims to learn representation of data that has an invariant distribution over different training domains. Previous works taking this approach match domains in feature space either by aligning moments (Sun & Saenko 2016) or using an adversarial loss (Ganin et al. 2016). However, these methods are later pointed out to be generally inadequate (Zhao 2019 Another popular approach is to learn the optimal invariant. predictors. Taking this approach, an invariant risk minimizer (TRM) optimizes a highly non-convex bi-level objective and simplifies the optimization using a penalty regularized objective. However, [Rosefeld et al. (2021); Kamath al. (2021); Ahuja et al. (2021) theoretically show that these algorithms fail even in simple data models. Similarly, Wang et al. (0022)
proposed Invariant-feature Subspace Recovery (ISR), which recovers the subspace spanned by the invariant features, and then fits predictors in this subspace. Distributionally robust optimization (Sagawa et al, 2020) is also a choice to tackle OOD generalization, which optimizes models over a worst-case distribution that is perturbed around the original distribution. In addition to the methods designed for training a model from scratch, recent works (Kumar al. 2022 Wortsman et al., 2022 Goyal ct al.] (2022) also discuss increasing the OOD accuracy over the pretrained model. While these methods provide impressive empirical improvements on pretrained models, theoretical explanations are yet to be provided.
Composition Generalization In the computer vision literature, previous research has investigated attribute-object compositions (also referred to as compositional zero-shot learning) (Misra et al. 2017
Nagarajan & Grauman 2018 Purushwalkam et al. 2019 Naeem et al., 2021 Nayak et al. 2023 Hao et al, 2023), with the goal of predicting attributes of an object in addition to its class. For instance, in a binary image classification task involving cat versus tiger, a classifier might be required to predict whether the animal is old or young alongside the conventional class prediction. In contrast, compositional generalization (CG) has different focus. It concentrates purely prodicting the object class (c.g., cat versus tiger). Specifically, CG tasks aim to accurately identify young cat images as cat, even when the training data consisting only of young tiger and old cat instances and lacking young cat images. This scenario intro- duces an out-of-distribution (OOD) shift, and the challenge lies in developing models robust to such shifts, While compositional zero-shot learning (CZSL) can decouple attributes from objects, it is not universally adaptable for OOD generalization tasks. This approach relies on powerful vision-language models (VT.ss) such as CLIP, while CG does limit the type of image classificrs. However, certain real-world domains, such remote sensing or medical imaging, there is lack of paired image-lexi data to train strong VLMs. Therefore, adopting self-super models such as MAE (He et al. 2022) and DINO (Caron et al 2021 presents a more practical strategy for these domains (Cong et al.
2022 Wanyan et al. 2023). As shown in Table [ our CFA can work with both VLMs and self-supervised models. In contrast, CZSL can not be directly applied to self-supervised models. Besides, Sivaprasad et al. (2022) explores CG under a slightly simplified premise, where only one random domain is masked for each class. In addition, their method is built upon ResNet (He et al. 2016b) and does not scale well with modern transformer architoctures.
5 Conclusion
This paper delves into the challenge of Compositional Generalization (CG) in machine learning, focusing on generalization to unseen domain-class combinations. By developing CG-Bench, a suite of benchmarks from real-world image datascts, we highlighted the shortcomings of prevalent pretraining-finctune in tackling this challenge. Our proposed solution, the Compositional Feature Alignment (CFA), offers a promising approach to improve the CG performance of pretrained models, as evidenced in our experiments. Despite these advances, our study is not without limitations. Our experiments are currently limited to the base-sized ViT models, and our empirical studies draw from a restricted number of datasets of limited size. As we strive to overcome these limitations in future work, we look to include larger models and diversify our benchmark suite, exploring alternative data sources beyond imagcs. We invite the broader machine lcarning community to join us in the ongoing exploration of the important challenge of CG.
frameworks
11

---

Published in Transactions on Machine Learning Research (05/2024)
References
Kartik Ahuja, Jun Wang, Amit Dhurandhar, Karthikeyan Shanmugam, and Kush R. Varshney. Empirical or invariant risk minimization? sample complexity perspective. In International Conference on Learning Representations, 2021. URL
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprinl 02893, 2019.
Sara Beery, Elijah Cole, and Arvi Gjoka. The iwildcam 2020 competition dataset. arXi::2004.10340 2020.
arXiv preprint
Gilles Blanchard, Gyemin Lee, and Clayton Scott. Generalizing from several related classification tasks to a new unlabeled sample. Advances in neural information processing systems, 24:2178 2186, 2011.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Mateusz Buda, Atsuto Maki, and Maciej A Mazurowski. A systematic study of the class imbalance problem in convolutional neural networks. Neural networks, 106:249-2550 259,, 2018.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in supervied vision transformers. In Proceedings of the IEEE/CVP international conference on computer vision, pp. 9650-9660, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597 1607. PMLR, 2020.
Gordon Christie, Neil Fendley, James Wilson, and Ryan Mukherjee. Functional map of the world. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018
Yezhen Cong, Samar Khanna, Chenlin Meng, Patrick Liu, Erik Rozi, Yutong He, Marshall Burke, David Lobell, and Stefano Ermon. Satmae: Pre-training transformers for temporal and multi-spectral satellite imagery. Advances in Neural Information Processing Systems, 35:197-211, 2022.
J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In CVPR, 2009.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidi- rectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Vol- ume (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 11/8655/11/N191
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkorcit, and Neil Houlsby. An image is worth 16x16 16 words: Transformers for image recognition at scale. ICLR, 2021.
Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 118 (43):e2100091118, 2021.
Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In Interna- tional conference machine learning, pp. 1180 1189. PMLR, 2015.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. ammain-ddversarill training of neural networks. The journal of machine learning research, 17(1):2096-200) -2300, 2016.
12

---

Published in Transactions on Machine Learning Research (05/2024)
Sachin Goyal, Ananya Kumar, Sankalp Garg, Zico Kolter, and Aditi Raghunathan. Finetune like you pretrain: Improved finetuning of zero-shot vision models. arXio preprint wrXiv:2112.00633 2022.
Ishaan Guirajani and David Sopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021.
Shaozhe Hao, Kai Han, and Kwan- Yee K Wong. Learning attention disentangler for compositional zero- shot lcarning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 15315 15324, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 770 770-778, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016b.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729 97389,977---7 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 16000 16009, 2022.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, :58 589,, 2021.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? In International Conference on Artificial Intelligence and Statistics, pp. 4069-4077. PMLR, 2021.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Trena Gao, et al. Wilds: A benchmark of in- the-wild distribution shifts. In International Conference on Machine Learning, pp. 2021.
PMLR,
Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang, Fine-tuning can distort pretrained features and underperform distr---iittibttionn. In International Conference on Learning Representations, 2022.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint w:Xiv16608.03933 2016.
Ulya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ar preprint 0rXiv:1711.05101, 2017.
Ishan Misra, Abhinav Gupta, and Martial Hebert. From red wine to red tomato: Composition with context. In Proceedings of the IEEE Conference Compuler Vision and Pallern Recognilion, pp. 1792-1880 2017.
Dustin G. Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis, 20, 2020.
Muhammad Ferjad Naeem, Yongqin Xian, Federico Tombari, and Zeynep Akata. Tearning graph embeddings for compositional zero-shot learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pallern Recognition, pp. 953 962, 2021.
Tushar Nagarajan and Kristen Grauman. Attributes as operators: factorizing unseen attribute-object com- positions. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 169-185, 2018.
13

---

Published in Transactions on Machine Learning Research (05/2024)
Nihal V. Nayak, Peilin Yu, and Stephen Bach. Learning to compose soft prompts for compositional zero-shot learning. In The Eleventh International Conference on Learning Represenlations, 2023.
Maxime Oquab, Timothée Darcet, Theo Moutakanni, Huy V. Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Russell Howes, Po- Yao Huang, Hu Xu, Vasu Sharma, Shang-Wen Li, Wojeiech Galuba, Mike Rabbat, Mido Assran, Nicolas Ballas, Gabriel Syn- naeve, Ishan Misra, Herve Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. Dinov2: Learning robust visual features without supervision, 2023.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40):24652-244 24662,, 2020.
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IREE/CVF international conference on computer vision, pp. 1406 1415, 2019.
Senthil Purushwalkam, Maximilian Nickel, Abhinav Gupta, and Aurelio Ranzato. Task-driven modular networks for zero-shot compositional learning. In Proceedings of the IEEE/CVE International Conference on Compuler Vision, pp. 3593 3602, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pp. 8748 8763. PMLR, 2021.
Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. In International Conference on Learning Representations, 2021. URL httpps://penrreiiew.ee nww.wewwwwwe/frrr
id BbNIBVJJ4344
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2020.
Shibani Santurkar, Dimitris Tsipras, and Aleksander Madry. (BREEDS): Benchmarks for subpopulation
shift. In International Conference on Learning Representations, 2021.
Sarath Sivaprasad, Akshay Goindani, Mario Fritz, and Vineet Gandhi. Class-wise domain generalization:
A novel framework for evaluating distributional shift. In NewrIPS 2022 Workshop Distribution Shifts: Connecting Methods and Applications, 2022.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In Furopean conference computer vision, pp. 443 450. Springer, 2016.
Christos Thrampouidiis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisiting neural-collapse geometry. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Hemanth Venkateswara, Jose Euscbio, Shayok Chakraborty, and Sethuraman Panchanathan, Deep hashing network for unsupervised domain adaplation. In Proceedings of the IEEE Conference on Compuler Vision and Pattern Recognition, pp. 5018 0-227, 2017.
Haoxiang Wang, Haozhe Si, Bo Li, and Han Zhao. Provable domain generalization via invariant-feature subspace recovery. In International Conference on Machine Learning, pp. URL httpp::/prooeeedings mlr, m.t..ress//1112//aag2222 hhtml
PMI.R, 2022.
Peisong Wang, Weihan Chen, Weixiang Xu, Qinghao Hu, and Jian Cheng. Rethinking the value of prompt learning for vision-language models, 2023. URL
Xinye Wanyan, Sachith Seneviratne, Shuchang Shen, and Michael Kirley. Dino-me: Sclf-supervised con- trastive learning for remote sensing imagery with multi-sized local crops. arXiv preprint 0rXiv:2308.06670, 2023.
14

---

Published in Transactions on Machine Learning Research (05/2024)
Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca. Roclofs, Raphael Gontijo Topes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7959-7971, 2022.
Yibo Yang, Shixiang Chen, Xiangtai Li, Liang Xie, Zhouchen Lin, and Dacheng Tao. Inducing neural collapse in imbalanced learning: Do we really need a learnable classifier at the end of deep neural network? In Al- ice H. Oh, Alekh Agarwal, Damielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Yuzhe Yang, Haoran Zhang, Dina Katabi, and Marzych Ghassemi. Change is hard: A closer look subpopulation shift. In International Conference on Machine Learning, 2023.
Han Zhao. On learning invariant representations for domain adaptation. ICML, 2019.
Han Zhao, Shanghang Zhang, Guanhang Wu, José MF Moura, Joao P Costeira, and Geoffrey J Gordon. Adversarial multiple source domain adaptation. Advances in neural information processing systems, 31: 8559-88 8570, 2018.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geomet- ric analysis of neural collapse with unconstrained features. Advances in Newral Information Processing Systems, 34:29820 29834, 2021.
15

---

Published in Transactions on Machine Learning Research (05/2024)
A Benchmark Curation
We introduce the curation process of DomainNet. (Peng et al. 2019) in Section 3.1 In this section, we provide more details of the proposed benchmark on the remaining datasets.
OfficeHome (Venkateswara et al., 2017) consists of classes of objects depicted in four domains: (Art, Clipart, Product Image, Real World Image]. The ITD and OOD splits are created in the same process the DomainNet. The average top-1 accuracy is measured to evaluate the model performance on both sets.
iWildCam (Beery et al., 2020) is dataset designed for classifying 182 classes of wild animals in camera traps. In the original iWildCam, each camera trap is defined to be domain, resulting in 324 domains that are labeled as camera IDs. Such a definition of domains ignores the semantic similaritics among them. For example, two camera traps may both be set in forests, however, the original iWildCam assigns them different domains. Consequently, images taken in the forests are distributed in both ID and OOD sets regardless of the fact that they are from very similar environments. Observing such deficiencies in the original iWildCam, we are motivated to reassign the domain labels to the data. We propose 6 domains for the dataset: {colored forest, gray-scale forest, colored trail, gray-scale trail, colored savanna, gray-scale savanna]. To assign the new domains to the images, we take advantage of the zero-shot classification ability of the CLIP models. Specifically, we design 10 prompts for cach domain and perform the zero-shot prediction of the domains for each image using the assembled prompts. New domain labels are generated for both training and test sets, resulting a dataset that is consistent with our CG setting: the E' K matrix is sparse, and the test set contains domain-class pairs that have few or no samples in the training set. We take the domain-class pairs with less equal to 20 samples in the training set the OOD samples in the test set. In addition to measuring the average top-1 accuracy, we also measure the F1-macro score for both the ID and OOD scts.
FMoW (Christie et al,, 2018) dataset that aims to classify utility types from satellite images. This dataset has two sets of domains: years and regions. The ID and OOD of the dataset are split according to the years, i.c., all the satellite images in the test set taken before 2016 are ID samples, and images after 2016 are OOD samples. In our CG-Bench, we only perform evaluations on region domains. To be specific, the dataset has K 62 classes and E 5 domains: {Asia, Europe, America, Africa, Oceania). The E x K matrix is naturally sparse because certain utilitics were only introduced post-2016 in some underdeveloped regions, leading to the creation of OOD entries. Therefore, FMOW is readily suited for our CG setting without the need for additional curation. We measure the average top-1 accuracy for the ID set and the worst region accuracy for the OOD set as suggested by previous works.
B Proof of Theorem I
We first restate the setting, and introduce some notations and definitions. Then we present two lemmas useful for the proof of our main theorem. Finally, we restate our main theorem (Theorem 1 and provide a proof.
Notation and Setting We demote Z = ((xn)
] E R Y = [v,,..., yN], and E = [ei...., eN]
as the stack of features, class labels, and environment labels, respectively. In the context of the unconstrained feature model, the optimization objective of Stage 2 is transformed to:
1 2 KN (cr(81 WiZ. + EN
min
(ce(33 W&Z, E) s.t. ZEU(d)
(6)
Furthermore, use Wi to refer to the j-th row vector of matrix W, and 2 represents the i-th column vector of Z, which is also the learned leature of the i-th training sample.
Below, we define simplex-encoding label (SEL) matrices, which are introduced in Thrampoulidis et al. 2022 and also consider SVD of the two heads, W1 and W2.
16

---

Published in Transactions on Machine Learning Research (05/2024)
Definition 2 (Simplex-Encoding Label Matrices). S1 RKXN and S2 RxN are simplez-encodingg label (SEI.) motrices for classes and domains, respectively, such that.
1/K
c= yi
VC [K], E E [ Si
(7)
-1/K
y
AE [ S2(c
1 1-1/E
c Ej
(8)
-1/E c c+ Ci
In other words, the i-th column of S1 is a K --dimensional vector, (ey. k 1), where ej represents the i-th slandard basis vector. Similarly, the i-th column of S2 is E-dimensiooall vector, (ee, 11.
Definition 3 (SVD of Heads). For W1 and W2, we consider their compact SVD as
W1 UN and W2 U2A2V2
(9)
Specifically, A1. A2 are positive diagonal matrices and U1. V1, U2, V2 have orthonormal columns.
Now, we focus the first term of 6 and study its optimum with theoretical tools from Thrampoulidis
et al. 2022p.
Lemma 1 (Optimum of Class Loss). Assuming the feature dimension is at least K, and training data exists for each class, and W1 is normalized such that W1 E E(dd) Addilionalll, we assume. 31 is sufficiently large. Then, for any constont norm magnitude O, the global minimum of the following objective,
min Z KN lcE(3: W2Z,Y) Y)
s.t. 1
(10)
satisfies
z n nKAT
(11)
where 7/1 = I-1/K*
Proof. From Theorem 1 of Thrampoulidis et al. 2022 we know that the optimum of 10 shall satisfy
w S
(12)
where 7 > 0 is a scaling factor and S is defined in Definition
With 12 and the norm constraint || a, it is obvious that Z should span the same feature subspace as Wi. since the cross-entropy loss is monotone such that a larger magnitude of W z* leads to smaller loss. Therefore, we can express * as
Z = 7i -
(13)
and the scaling factor can be easily determined by considering ||W(() 12 and | a.
We study the optimum of the second term of G in the same fashion, and obtain similar results. Lemma (Optimum of Domain Loss). Assuming the feature dimension d is at least E, and training data exists for each domain, and W2 is normalized such that W2 E U(d) Additionally, we assume B2 is sufficiently large. Then, for any constant norm magnitude > O, the global minimum of the following objective,
min z EN
( W2Z,E) E)
st. | b.
(14)
satisfies
z = -2vA,,
(15)
where 7 T-1/E'
17

---

Published in Transactions on Machine Learning Research (05/2024)
Proof. The proof is the same as that of Lemma
With Lemma T and Lemma we can prove that the learned feature vector of each training sample can be decomposed linear combination of two vectors depending on the class label and domain label, respectively, while the two vectors live in orthogonal feature subspaces. This indicates that the learned features conform to compositional feature structure satisfying Definition
Notice that Theorem 2 (the restatement of Theorem 1 is slightly different from Theorem This is because we found a small issue in the original proof for Theorem T We fixed the issue before the supplementary material submission deadline, which led to a slightly different expression of zf. However, the same main conclusion still applies to the updated theorem (Theorem 2 the learned features comply with the compo- sitional feature structure in Definition [I which justifies the algorithm design of our algorithm design. We will revise Theorem 1 in main text correspondingly in the revision.
Theorem (Feature Aligmment (Theorem Restated)). Assuming the feature dimension is no smaller than E, and training data erists for each class and domain (though not necessarily for each domain-class combination), and W1 and W2 are normalized and span orthogonal subspoces such that W1 U(d) K W2 u(d) and W.W.T 0. Additionally, we assumne 31,82 sufficiently large. The global minimam of 6 resulls in the following: for any i [N], denole zι as the i-th column veclor of Z, we have
2 = ay: W.Tbe:
(16)
where azy E IRK is a veclor depending on the class label yi, and be, R.E. is veclor depending on the domain label 6z
Proof. Lemina and Lemma derive the optimum the two loss terms of 6 respectively. For the linear combination of the two loss terms with coefficient X, straightforward see that the optimum Z should also be linear combination of (11) and (1). since these two optima stay in orthogonal subspaces spanned by W1 and W2, respectively (Stage ensures that W1 and W₂ live in orthogonal subspaces). In other words, the first loss term 1 (Cr W, Z, Y) does not affect the converged direction of Projw, 7*,i.e. Z* projected into the subspace spanned by W2, and vice versa. However, the first term affects the magnitude of Projw, z" since we have the norm constraint Z E Udd) N in 6 Specifically, the larger À in 6 leads to larger magnitude of |Projw, 11/ Hence, we can express Z as
z" = a7 VAA 122
(17)
where a, b are scaling factors (depending on A) that ensure Z* =(d))
By plugging in Definition 2 sample, as
express the i-th column vector of Z*, the learned feature of the i-th
= Vi - 1 b2 UT, USee 1 -
(18)
Clearly, the first term of 18 depends on the class label yi only and the sccond term relies only on the domain label eş. Besides, the two terms live in the feature space spanned by Vi and V22, respectively, which are determined by W, and W2. Hence, can be re-expressed (18 as
z = W.T., ay: W.Tber
(19)
where ay UA 1) and b by2U223UT (ee, 1 1).
C Details of Experiment
Here, we supplement. Section 3 with more details about our experimental studies.
18

---

Published in Transactions on Machine Learning Research (05/2024)
Computational Expenses The experiments described in this paper are executed NVIDIA RTX A6000 GPUs with 48GB memory, utilizing a total of 12 GPUs. The individual training runs for Office- Home, DomainNet, iWildCam, and FMoW require 0.2, 0.8, 0,5, and 0.2 GPU hours respectively. Each experiment is repeated ten times with varying random sceds, which alter the order of batch shuffling in the training loader.
Hyperparameter Choices We present the hyperparameter settings for our CFA models in Table [ The parameters for each stage are chosen based on model performances on the OOD validation set. Note that À is the domain loss coefficient in (3] and Aortho is the coefficient for the orthogonality regularization loss [W] W2ll f that we use to ensure orthogonality of heads in Stage The hyper-parameters in Stage are also used for the two baseline algorithms, Full finetuning (FT) and LP-FT.
Ablation Studies We conduct ablation studies on i) the choice of normalized head v.s. unconstrained head, ii) the choice of frozen head v.s. trainable head, and iii) the loss coefficient À in 3 The results are discussed as follows:
Head Normalization. In our implementation, we normalize the weight of our classification heads that has bias terms, following recent works (Goyal et al.) 2022 Wang et al. 2023) which show that head normalization is useful for CLIP finetuning. In this ablation study, we repeat the vanilla full finetuning experiment using unconstrained classification heads (i.e., with bias and without normalization). The Lest. accuracies(%) are present in Table 5 From the results, we can see that the models finetuned with constrained heads have slightly better OOD performances compared to the unconstrained ones, justifying the use of the head normalization technique.
Trainable Heads u.s. Frozen Heads. In Stage 2 we finetune the models with frozen heads, W1 and W2, which are linearly probed in Stage [ In this ablation study, we perform the full finetuning during Stage i.c., training all the layers including the classification heads. The test accuracics are shown in Table The results indicate that using trainable heads or using frozen heads does not vary the uned model performance much. Therefore, this ablation study justifics that our algorithmic design of frozen classification heads is reasonable.
Domain Loss Coefficient. In our empirical implementation, we consider == 0 in (3) and only finetune with respect to W1 in Stage 2 to reduce the training cost. The motivation for doing so the observation that the coefficient. À Stage 2 does not. affect the final performance. To justify this, we conduct the following ablation study: we run Stage with the original two-term loss, (3), while adjusting the coefficient X for the domain prediction loss. The test accuracics(") (%) are shown in Figure 6 From the figure, one observe that the value of Ả can be set to zero in Stage 2 without negatively affecting the final performance. Therefore, it is sufficient to consider X 0 in Stage 2 and only fine-tune the model with the first loss
term in 3
C.1 Discussion on Training Stability and Overhead
Here we conduct an additional study to examine the training stability of our two-stage method, CFA. All of the following experiments are conducted for CLIP ViT-B/16 on the Office- dataset and are an average over 3 sceds.
Table 4: Hyperparameters for our algorithm.
Model
CLIP
DINOv2
Stage
[Linear Probing)
Stage
Vine-Tuuiigg
Stage
[Linear Probing)
Stage
(Fin-TTunigg
Dataset
A
Auctho
Epochs
Epochs
Learning Rate
A
Aarcho
Epochs
Epochs
Learning Rate
OfficeHome
1
100
200
3
10-5
1
100
200
3
x 10-5
DomainNet
1
10000
200
3
10-5
1000
10
200
10
5 10
IWildCam
10
10
200
5
10-5
1
100
200
10-5
FMoW
10
100
200
3
10-5
100
1000
200
4
5 x 10-5
19

---

Published in Transactions on Machine Learning Research (05/2024)
Table 5: The ablation study on the constraints on classification heads. Blue cells represent the normalized classification head with no bias; Orange cells indicate the heads are unconstrained.
Offe
DomainNet
iWildCam
FMoW
Methods
ID
0OD
ID
0OD
ID
0OD
ID
OOD
FT
94.3
51.0
82.0
74.5
16.5
65.8
38.7
FT-Wiie
93.7
52.5
76.4
8.7
67.0
13.7
49.5
40.6
FT
94.3
50.8
82.0
74.9
16.2
65.8
38.5
FT-Wise
93.7
52.3
76.5
66.8
13.5
49.5
40.1
Table 6: The ablation study on the frozen classification heads v.s. trainable heads in Stage represent the fronzen heads; Orange cells indicate the heads are trainable.
Blue cells
OflecHome
DomainNet
iWidCaam
FMoW
Methods
Zeroshot
89.2
50.3
61.7
6.6
13.7
6.9
20.4
18.8
FT
94.3
51.0
82.0
7.5
74.5
16.5
65.8
38.7
FT (Wise)
93.7
52.5
76.4
8.7
67.0
13.7
29.5
40.6
PT
94.3
51.4
82.1
7.3
75.3
16.4
65.8
38.5
FT (WISE)
93.7
52.9
76.6
8.4
67.7
13.6
49.7
39.7
IP-PT
74.0
65.9
40.2
(WISE)
LP-FT (WISE)
93.1
79.4
74.4
58.5 8
36
CFA (LP-FT)
93.7
53.2
81.6
7.3
71.0
18.3
65.3
41.6
CFA (WISE)
93.1
56.1
76.5
9.2
7/.6
19.7
58.5
36.00
CFA (I.P-TT)
93.7
52.8
81.7
7.0
73.9
18.4
65.5
41.0
CFA (WISE)
93.0
56.2
76.6
74.5
19.6
.3.7
36.0
Ablation on Loss Coefficient (FMoW)
67.0
43.0
OOD Worst Region Accuracy
66.5
42.5
66.0
42.0
Accuracy
65.5
41.5
À 1>
65.0
41.0
==
64.5
40.5
64.0
40.0
63.5
39.5
63.0
39.0
10-4
-0 Loss Coefficient (A)
10-3
10-1
10-4
10-2 Loss Coefficient (A)
10-3
10
Figure 6: The ablation study domain loss coefficients. The experiments conducted the FMoW dataset example. The red dashed lines are the performance of the model without supervision the domain labels.
We first. show that our method is stable the number of iterations in Stage-1. From Table 7 that as the linear heads are trained longer in Stage-1, the ID (WiSE) performance slightly increases and the OOD accuracy reaches the peak at 4000 to 6000 iterations. Except for the trade-off between ID and OOD performance, our method is stable with no drastic drop in accuracy. The 6000 iterations of Stage- was the hyperparameter used for experiments in Table
we can see
We then demonstrate how training epochs Stago-2 will affect the final performance of our model. Table shows that for longer training epochs, the ID accuracy increases and the OOD accuracy decreases. Hence, in our paper, used 3 epochs for the best OOD performance.
20

---

Published in Transactions on Machine Learning Research (05/2024)
Table 7: Study on the training stability of CFA's 's Stage-1.
LP Iterations
ID
ID (WISE)
OOD
OOD (WiSE)
2000
94.0
92.8
54.3
56.9
4000
94.0
93.1
54.3
57.3
6000 (Chosen)
94.0
93.1
54.3
56.9
8000
94.0
93.6
53.3
56.5
10000
94.0
93.7
52.1
55.0
Table 8: Study on the training stability of CFA's Stage-2.
Epochs
TD Acc
ID Acc (WiSE)
OOD Acc
OOD Acc (WiSE)
3 (Chosen)
94.0
93.1
54.33
56.9
5
94.1
98.5
52.3
56.3
10
94.1
93.7
52.2
55.7
Parameter Overhcad Compared with the vanilla full fine-tuning, our CFA introduces an additional linear classifier for domain labels in Stage-1. In terms of model parameters, we provide the parameter count of CLIP and DINOv2 DomainNet, in Table g One see that the additional parameters (e-classifier) introduced by CFA is negligible compared with the total number of parameters.
Table 9: Parameter count for CLIP and DINOv2 on DomainNet. y-classifier is the linear classifier for class labels, and e-classifier is the linear classifier for domain labels. Vision encoders, language encoders and y-classifier are modules shared between the vanilla fine-tuning and CFA, and the last column is additional parameters introduced by CFA.
Parameters
Vision Encoder Language Encoder
y-classifier 177 10"
e-classifier (CFA only)
CLIP VTT-B
86 x 10"
10"
3.1 10"
DINOv2 ViT-B
10
265 10a
10
Compute Overhead The training cost of Stage-1 (linear probing) is much smaller compared to Stage-2 (backbone finetuning). Stage-1 requires only a single forward pass over training samples for feature gathering (features are stored in CPU memory), and then the linear probing can be achieved fastly on CPU. The training cost for Stage-2 of CFA slightly less than standard full finetuning, as the last layer remains frozen during Stage-2 of CFA. Overall, we conclude that the additional training costs of CFA are mild (mostly coming from the forward pass over training samples), which is less than that of one finc-tuning epoch. In other words, compared with vanilla fine-tuning for N epochs, our CFA's training compute cost. is less than fine-tuning for N + 1 epochs. The modest compute overhead of CFA making it practical method for addressing compositional generalization problems in real-world applications.
Memory and Speed Overhead For CFA, the memory requirement of Stage-1 is less than Stage-2 since Stago-1 only involves forward passes of the image encoder (no backward pass). For Stago-2 of CFA, the memory usage and training speed are the same as that of the varella fineluning, since Stage-2 of CFA does not introduce additional parameters or loss for training.
C.2 Discussion of the Performance Gain
From Table [ one may think the performance gain of CFA over reweighting, a strong baseline method, is not significant. However, our main contribution lies in designing a principled algorithm for compositional generalization, a novel challenge OOD generalization. As our theoretical analysis and feature visualization in Fig. 5 demonstrate (on Color MNIST and DomainNet, respectively), our method aligns fcatures with a compositional structure suitable for compositional generalization. In OOD generalization research, it's typ- ical for specialized algorithms to only modestly outperform baselines like ERM, as evidenced by extensive benchnarking in studies like DomainBed (Guirajani & Lopez-Paz] 2021 and WILDS (Koh et al. 2021). Thus, we consider CFA's performance gain to be meaningful and convincing. Additionally, the mild imple- mentation complexity and training costs (discussed in Sec. C.1 make CFA a viable method for enhancing compositional generalization.
Morcover, we want to emphasize that woal-world datasets used in our paper have lots of labeling issues (for both class and domain labels), which may prevent CFA from obtaining greater performance gain. Notably,
21

---

Published in Transactions on Machine Learning Research (05/2024)
CFA employs both class and domain labels for training, leading increased susceptibility to label noise, particularly from domain labels. Below, we illustrate two types of labeling noise encountered:
Mislabeling: We observe that the DomainNet dataset contains many mislabeled images. Many images in the "Clipart" domain are actually real photos of a single object on a white background, which should be categorized into the "Product" domain. On the other hand, numerous class labels also mislabeled: images in the "bush (plant)" class contain pictures of President George Bush of the USA; the "cooler" class includes electric fan images, which are incorrectly categorized; the "square (shape)" class contains lable, which should be placed into the "table" class instead. Note: We only manually examined a very tiny subset of the DomainNet dataset (Gulrajani & Lopez- Paz. 2021). which comprises 0.6 million images, and have already found many mislabeled images. Therelore, overall, we believe the mislabeling ratio is not negligible.
Ambiguous Labels: We observe that certain domains contain large number of images that are visually similar to those in another domain. For example, in both ffficeHHome and DomainNet, the images in the "Product" domain are real photos of objects, making them almost indistinguishable from their counterparts in the "Real" domain. The distinguishing feature of the "Product" domain is that its images all have a white background; however, some images in the "Real" domain also share this characteristic. Additionally, in DomainNet, the "Infograph" domain contains many images stylistically similar to those in "Clipart" "Real"; some images in the "Painting" domain are sketches, despite the presence of a separate "Sketch" domain. This ambiguity issue extends to class labels as well. In DomainNet, the classes "cup," "coffee cup," " and "mug" lack clear stylistic distinctions.
In addition to these issues, other aspects of these datasets may affect learning. For instance, the iWildCam dataset includes class labeled "emply"; signifying the absence of animals, which comprises significant portion of the dataset.
C.3 Effect of the Orthogonality Regularization
To better understand the impact of the orthogonality regularization on the optimization process, we con- ducted an ablation study Stage-1 of our CFA method using the DomainNet datasct. In this study, we optimized the objective in Equation 1 with different orthogonal regularization coefficients (Northo) for 20 epochs. The results of this ablation study are visualized in Fig. 7
As shown in Fig. 7 the orthogonality loss (measured by (w/W)) decreases as the linear probing epochs progress, indicating that the orthogonality regularization effectively enforces the orthogonality constraint between the two classifier heads W1 and W2. The rate of decrease in orthogonality loss is proportional to the value of the regularization coefficient Aortho, with larger values resulting in faster convergence towards orthogonal classifier heads.
This ablation study demonstrates that the orthogonality regularization term in our CFA method is crucial for obtaining orthogonal classifier heads W1 and W,, which are subsequently used in Stage-2 of our approach. By iveec enforcing the orthogonality constraint, our method ensures that the learned features are well- suited for compositional generalization tasks.
22

---

Published in Transactions on Machine Learning Research (05/2024)
Effectiveness of Orthogonal Regularization (DomainNet)
10a
100
Aortho
Orthogonal Loss
0.0
10-1
1.0
10
10.0
10
10
10-5
Linear-Probing Epochs
Figure 7: Effectiveness of the orthogonality regularization during Stage-1 linear probing on DomainNet. The orthogonality loss, WF"W2/P is plotted against the number of epochs for different values of Aortho. Larger Aortho leads to faster convergence towards orthogonal W1 and W2.
23


Reviews:
[{"review_id": "9CP2zxaNih", "reviewer": "Reviewer_L4Jk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strength\n1. Paper is well written and motivations are clearly stated. Good problem setup with helpful visualizations that clearly convey the message of the paper\n2. Method is relatively simple but effective as results from Table 1 demonstrate\n3. Experiments using Clip to generate domain labels show compelling performance and an ability to leverage the approach even in settings without domain labels.\n\n\n## Weaknesses\n1. The benchmark is a bit overstated as a contribution since 2/3 of the sub-tasks are basically wholesale lifting of existing datasets with no real modification.\n2. The theoretical guarantees given are with respect to training data but say nothing about  whether the aligned featurizer also exhibits the compositional property for data not in the training set.\n\n## Questions : \n1. How do you decide what numbers to bold in Table 1 ?  On FMoW, the performance of Reweight-E with clip (41.8) is better than your method (41.6) but your method's result is instead bolded. \n2. Any idea hypotheses about why WISE tends to hurt performance on the FMoW dataset ?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a method to achieve compositional generalization across domains and item classes via a two stage learning process. \nFirst, orthogonal domain and classes discriminators are learned given a fixed set of features and then in the second stage, a new set of features are learned given fixed class and domain discriminators.\nThe paper shows some improvement of using this method over reasonable baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The clip-CFA row in Table 1 has some incorrectly formatted #s in the first 2 columns\n2. Since the paper is currently under the 12 page limit, the Ablations could have been included in the main body of the paper instead of the appendix. \n3. **[Requested Experiment]** It would be good to have an ablation with a single stage version of the whole fine-tuning process to confirm that the 2 stage process is indeed needed though it is more complex and has extra overhead. Specifically,  having a single stage of training where the body parameters and heads are all jointly trained ---> with and without the orthogonality constraints on the heads.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "lOIRqFwvZr", "reviewer": "Reviewer_kZdn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n1. The paper addresses a significant and relatively unexplored aspect of machine learning: CG, highlighting its practical implications.\n\n2. The introduction of CFA is innovative and shows promise, as evidenced by empirical evaluations.\n\n3. The creation of CG-Bench marks a valuable contribution, providing a resourceful benchmark for future research in this area.\n\nCons:\n1. While the theory presented in the paper aims to learn features with robustness and interpretability, there are doubts about the generalizability of the results. Figure 5, although impressive, may represent a selectively favorable scenario.\n\n2. The methods to maintain orthogonality between W_1 and W_2 during training, and the integration of constraints in Equation (2) into the overall optimization objective, are not clearly explained. More detail on the optimization process and the impact of the orthogonal constraint on loss optimization would be beneficial. Visualizations of loss curves with orthogonal contrast ablations would enhance clarity.\n\n3. There are typographical errors in Table 1, such as \"943.07\" or \"54.33.2\".\n\n4. The overall improvements presented in the paper, while positive, are not markedly significant.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a novel approach to tackle the a new challenge in machine learning named as compositional generalization (CG), which aims at enhancing generalization to unseen domain-class combinations. The study proposes a method called Compositional Feature Alignment (CFA) to boost the CG performance of pretrained models via two-stage training. It also contributes to the field by developing CG-Bench, a benchmark suite derived from real-world image datasets, to assess the effectiveness of the proposed approach.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The paper introduces a promising new approach to CG in machine learning, demonstrating potential in empirical tests. However, there are areas where improvements are needed, particularly in terms of clarity and detail in the methodology.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HsmYXLD5eC", "reviewer": "Reviewer_kFXZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n1-the presented method seems relatively simple and therefore reproducible.\n\nWeaknesses: \n1-The main weakness seems to be that the method is specifically tailored to the dataset and may have limited impact on real-world problems. Specifically there seems to be two key assumptions: (i) - domain labels are available, and (ii) - the domain labels and class labels are mutually independent. The second assumption is implicit in the algorithmic constraint $W_1 W_2^T = 0$. Yes, section 3.4 has a discussion on having limited domain data labels, nevertheless the fact that the data generative process cleanly consist of just a few domains is already a limiting factor. For many real world applications (ii) is not a valid assumption, e.g. imagine different species being recorded with different instruments, in this case the domains are not independent with the classes so it would seem a key design decision is violated. \n2-As the authors mention in the paper, the evaluations were conducted with \"based-size\" (i.e. small) ViT models. Therefore, the relevance of this algorithm at large scales is dubious.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposed a, two-stage, fine-tuning approach for achieving compositional generalization between class labels & domains. \nThe approach involves : 1-multi-head training on classes & domains, 2-finetuning the features through the heads. It also makes use for weight averaging (WiSE-FT) to improve the performance further. Results are presented on a combined dataset consisting of Office-home, DomainNet, and iWildCam and WILDS datasets.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "-The legend in Figure 2 is broken (some shapes should be triangles instead of squares), axes are not labeled -- do the axes represent various features?\n-Definition 1 seems arbitrary and overly constrained. It seems tailored to the datasets under study, this definition does not seem suitable for realistic data.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers found the paper well-written with motivations clearly stated (Reviewer L4Jk), the proposed method is simple (Reviewer L4Jk, Reviewer kFXZ) and effective (Reviewer L4Jk) and “shows promise, as evidenced by empirical evaluations” (Reviewer kZdn). The paper addresses a “significant and relatively unexplored aspect of machine learning” (Reviewer kZdn). While Reviewer kFXZ is concerned about limiting assumptions made that aren't met in (all) real-world applications, the authors have made claims within the scope of that setting and have clearly outlined the assumptions made -- e.g. "Empirically, we find that if the learned features (i.e., the outputs of the last hidden layer) conform to a compositional structure where the subspace of domain related features is orthogonal to that of class-related features, the corresponding model can generalize across unknown domain-class pairs."


During the rebuttal, the authors have made clarifications, expanded on the discussion of limitations and conducted an additional ablation experiment to verify the necessity of the two-stage design of their method.


I encourage the authors to add the additional discussion of limitation to the revised manuscript as well.
Decision: Accept (poster)