Compositional Risk Minimization
Authors: Divyat Mahajan
Mohammad Pezeshki
Charles Arnal
Ioannis Mitliagkas
Kartik Ahuja
Pascal Vincent
Venue: ICML 2025 poster (2025)
Keywords: Compositional Generalization
Out of Distribution Generalization
Provable Extrapolation
Energy Based Models
Primary area: deep_learning->theory
TL;DR: Provable method for extrapolating classifiers to novel combinations of attributes (a.k.a. compositional generalization)

Abstract: Compositional generalization is a crucial step towards developing data-efficient intelligent machines that generalize in human-like ways. 
In this work, we tackle a challenging form of distribution shift, termed compositional shift, where some attribute combinations are completely absent at training but present in the test distribution. This shift tests the model's ability to generalize compositionally to novel attribute combinations in discriminative tasks. We model the data with flexible additive energy distributions, where each energy term represents an attribute, and derive a simple alternative to empirical risk minimization termed compositional risk minimization (CRM). We first train an additive energy classifier to predict the multiple attributes and then adjust this classifier to tackle compositional shifts. We provide an extensive theoretical analysis of CRM, where we show that our proposal extrapolates to special affine hulls of seen attribute combinations. Empirical evaluations on benchmark datasets confirms the improved robustness of CRM compared to other methods from the literature designed to tackle various forms of subpopulation shifts.

Compositional Risk Minimization
Divyat Mahajan 12
Mohammad Pezeshki' Charles Arnal'
Ioannis Mitliagkas 2
Kartik Ahuja
Pascal Vincent 12x
Abstract
have never seen before, such "a blue elephant riding a bi- cycle on the Moon." While most works on compositionality have focused on its generative aspect, i.e., imagination, as seen in diffusion models (Yang et al., 2023a), composition- ality is equally important in discriminative tasks. these tasks, the goal make predictions in novel cireumstances that are best described as combinations of circumstances seen before. In this work, we dive into this less-explored realm of compositionality in discriminative tasks.
Compositional generalization is crucial step to- wards developing data-efficient intelligent ma- chines that generalize human-like ways. In this work, we tackle a challenging form of dis- tribution shift, termed compositional shift, where some attribute combinations are completely ab- sent at training but present in the test distribution. This shift tests the model's ability to generalize compositionally to novel attribute combinations in discriminative tasks. We model the data with flexible additive energy distributions, where each energy term represents an attribute, and derive a simple alternative to empirical risk minimization termed compositional risk minimization (CRM). We first train an additive energy classifier to pre- dict the multiple attributes and then adjust this classifier to tackle compositional shifts. We pro- vide an extensive theoretical analysis of CRM, where we show that our proposal extrapolates to special affine hulls of seen attribute combinations. Empirical evaluations benchmark datasets con- firms the improved robustness of CRM compared to other methods from the literature designed to tackle various forms of subpopulation shifts.
We work with multi- autti-at data, where each input (e.g.. an image) associated with multiple categorical attributes, and the task is to predict an attribute or multiple attributes. During training, we observe inputs from only a subset of all possible combinations of individual attributes, and during test we will novel combinations of attributes never seen at training. Following Liu et al. (2023), we refer to this distribution shift as compositional shift. Towards the goal of tackling these compositional shifts, we develop an adap- tation of naive discriminative Empirical Risk Minimization (ERM) tailored for aulti-at data under compositional shifts. We term our approach Compositional Risk Min- imization (CRM). The foundations of CRM are built on additive energy distributions that are studied in generative compositionality (Liu et al., 2022a), where each energy term represents one attribute. In CRM, we first train an additive energy classifier to predict all the attributes jointly, and then we adjust this classifier for compositional shifts.
1. Introduction
Our main contributions are follows:
The ability to make sense of the rich complexity of the sen- sory world by decomposing it into sets of elementary factors and recomposing these factors in new ways is a hallmark of human intelligence. This capability is typically grouped un- der the umbrella term compositionality (Fodor & Pylyshyn, 1988; Montague, 1970). Compositionality underlies both semantic understanding and the imaginative prowess of hu- mans, enabling robust generalization and extrapolation. For instance, human language allows us to imagine situations we
Theory of discriminative compositional shifts: For the family of additive energy distributions, prove that additive energy classifiers generalize compositionally to novel combinations of attributes represented by a special mathematical object, which we call discrete affine hull. Our characterization sharp, i.c., we show that is not possible generalize beyond discrete affine hull. We show that the volume of discrete affine hull grows very fast in the number of training attribute combinations thus generalizing to many attribute combinations. The proof techniques developed in this work are very different from exist- ing works on distribution shifts and hence may be of independent interest.
extrapolation
"Work done Meta Joint last author 'Meta FAIR MMila, Uni- versité Montréal. Correspondence to: Divyat Mahajan <di- vyat.mahajan @ mila.quebec>.
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Compositional Risk Minimization
Train log prior log p(y.a)
Bias fleamed) Biy.e)
Components for
Representaann (ec 0)
F((s)
only
Not
Linear Layer W
Train adjusted logits
Backbone (.)
Test adjusted logits
Components Cor y
log (
-y(y. Bias (estrapolatio))
Test log prior
Figure 1. The additive energy classifier trained in CRM computes the logits for each group (v, a) by adding the energy components of each attrihute via boradcasting. For the train logits, we add the log of the prior probahilities and learned hias B(y, a) for the groups present in train data. At test time, the log prior term is replaced with the log of the test prior (if available, otherwise assumed to be uniform), and the biases for novel test groups, 3" (y, a), are extrapolated using 11. Finally, we obtain p(y. alæ) by applying softmax function the adjusted logits. This adaptation from train to test is possible because of the additive energy distribution p(aly, a), which allows the model to factorize the distribution into distinct components associated with each attribute.
A practical method: CRM is a simple algorithm for training classifiers, which first trains an additive energy classifier and then adjusts the trained classifier for tack- ling compositional shifts. We empirically validate the superiority of CRM to other methods previously pro- posed for addressing subpopulation shifts. Our code repository be accessed via the link in the footnote
where we observe data from all the groups but some groups present much more data than the others.
Formally, let p(x; =) = p(=)p(x(z)) denote the train distri- bution, and ((x,,) = (==)((=(=)) the test distribution. We denote the support of each attribute component 2; under training distribution training distribution as gtrain, The corresponding supports for the test distribution are denoted as Z and Ztest, We define the Cartesian product of marginal support under train- ing as 2 == Ztrain x Z x Ztrain
Ztrain and the support of under
2. Problem Setting
2.1. Generalizing under Compositional Shifts
In this work, we study compositional shifts from a training distribution p to a test distribution q, characterized by:
In compositional generalization, we aim to build a classifier that performs well new contexts that are best described as novel combination of seen contexts. Consider an input i (e.g., image), this input belongs to group that is character- ized by an attribute vector
1. p(x|2) g(<l2), V2 2x.
(Z1..... Zm) (e.g., class la- bel, background label), where z; corresponds to the value of tth attribute. There are m attributes and each attribute 2; can take possible values. So 2 E with = =1 (1.....d)"" dd)
2. tes z ztrain but test C 2
The first point states that the conditional density of inputs conditioned on attributes remains invariant from train to test, which can be understood the data generation mecha- nism from attributes to the inputs remains invariant. What changes between train and test is thus due to only shift- ing prior probabilities of attributes from p(z) to g(2). The second point specifies how these differ their support: at test we observe novel combinations of individual attributes but not a completely new individual attribute. The task of compositional generalization then to build classifiers that are robust to such compositional distribution shifts. Also, we remark that the above notion should remind the reader of the notion of Cartesian Product Extrapolation (CPE) from Lachapelle et al. (2024). Specifically, if a model succeeds on test distributions q(z) with support equal to the full Carte- sian product (2test 2 2x), then it is said to achieve CPE.
We use the Waterbirds dataset as the running example (Sagawa et al., 2019). Each image x has two labeled at- tributes summarized in the attribute vector (y, a), where y tells the class of the bird Waterbird (WB) or Landbird (LB), and a tells the type of the background Water (W) or Land (L). Our training distribution consists of data from three groups - (WB, W). (LB.L) (LB, W). Our test dis- tribution also consists of points from the remaining group (WB, L) as well. We seek build class predictors that per- form well on such test distributions that contain new groups. This problem setting differs from the commonly studied problem in Sagawa et al. (2019); Kirichenko et al. (2022),
'Github ffcebooknessearchccoopsssii -ris miniimmmiiiiiiiiiiiiiiiiiioo minimization
2

---

Compositional Risk Minimization
2.2. Additive Energy Distribution
This is the right juncture contrast AEDs with distribu- tional assumptions in recent provable approaches compo- sitional generalization (Dong & Ma, 2022; Wiedemer et al., 2023; 2024; Brady et al., 2023; Lachapelle et al., 2024), These works assume labeling functions or decoders that are deterministic and additive over individual features, proving generalization over the Cartesian product of feature sup- ports (further discussion in Appendix A). While insightful, this assumption is restrictive, as each attribute combination corresponds to single observation with limited generative interactions. In contrast, AEDs capture stochastic decoders, offering a more flexible way to model inputs as a conjunc- tion of characteristics (see Appendix B).
We assume that p(z]2) is of the form of an additive energy distribution (AED):
m
p(z]2)= (()) exp (-
Σ
(x, 21)
(I)
where Z(2) = exp E
m ==
Ea(e, 2:) de is the parti-
tion function that ensures that the probability density p(al)) integrates to one. Also, the support of p(z[z) is assumed to be IR, Vz zx
We thus have one energy term Es associated to each attribute 2/. Note that we do not make assumptions on E, except Z(z) < oc, leaving the resulting p(æ]z) very flexible. This form is a natural choice to model inputs that must satisfy a conjunction of characteristics (such as being natural image of a landbird AND having a water background), correspond- ing our attributes.
3. Provable Compositional Generalization
Our goal is to learn a model yielding ((2[2)) that will match the test distribution g(=[ar) and thus allow us to predict the attributes at test time in a Bayes optimal manner. If we successfully learn the distribution g(> 2)), then we can straightforwarddy predict the individual attributes q(2, (z), e.g., the bird class in Waterbirds dataset, by marginalizing over the rest, e.g., the background in Waterbirds dataset. Observe that (==(ar) differs from the training p(= which can be estimated through standard ERM with cross-entropy loss. Since some attributes 2 observed at test time are never observed at train time, the distribution learned via ERM assigns a zero probability to these attributes and thus cannot match the test distribution ((]2x))
Recall (zl, (. =m) is a vector of m categorical at- tributes that can each take d possible values. We will denote as (=) the representation of this attribute vector as a con- catenation of m one-hot vectors, i.e.
(z) = [onehot(z]), .., (
Thus (=) will be a sparse vector of length mnd containing m ones. We also define a vector valued map E(x) (I, 1)...., , E1 d)....... ..... Fm (Ir, 1),...,
. Em (x. d))
In what follows, we first introduce a novel mathematical object termed Discrete Affine Hull over the set of attributes. We then describe generative approach for classification that requires us to learn p(z|2) including the partition func- tion, which is not practical. Next, we describe a purely discriminative approach that circumvents the issue of learn- ing p<>|==) and achieves the same extrapolation guarantees as the generative approach. We present the generative ap- proach as it allows to understand the results more easily. Building generative models based on our theory is out of scope of this work but is an exciting future work.
where (x, 2x) is the energy term for 7th attribute taking the value 2,. This allows us to reexpress equation 1 using simple dot product, denoted (:,)):
p(alz)= = 2(3) exp (o(())
where 7(z)= f exp (- (a(2), F(x)) )dx da is the partition function.
There are two lines of work that inspire the choice of addi- tive energy distributions. Firstly, these distributions have been used to enhance compositionality in generative tasks (Du et al., 2020; 2021; Liu et al., 2021) but they have not been used in discriminative compositionality. Secondly, for readers from the causal machine learning community, it may be useful to think of additive energy distributions from the perspective of the independent mechanisms principle (Janz- ing & Schölkopf, 2010; Parascandolo et al., 2018). The principle states that the data distribution is composed of independent data generation modules, where the notion of independence refers to algorithmic independence and not statistical independence. In these distributions, we think of energy function of an attribute as an independent function.
3.1. Discrete Affine Hull
We define the discrete affine hull of a set of attribute vectors A (2(1) ....,,(K)] where 2(0) Z, defined (AA) as:
{z E 2 (
Σ aua
04 = 1}
In other words, the discrete affine hull of A consists of all attribute vectors whose one-hot encoding lies in the (regular) affine hull of the one-hot encodings of the attribute vectors of A. This construct helps characterize which new attribute combinations we can extrapolate to.
3

---

(4) Compositional Risk Minimization
As illustration, consider the Waterbirds dataset, where observe three of four possible groups. In one-hot coding, WB is [1, 0). LB is [0, IJ. Water is [1, OJ, and Land We show that the missing attribute vector WB on
For any 2' E DAff( (trri)), by definition there exists
such that a(2)) = Σ DE 2mn (=)(z). Thus (a(2'), F(x)) Ezegtal 02 (o()), E(x)), by linearity of the dot product.
is [0, 1]. I., represented as [1 O 0 1]. can be expressed as an affine combination of the observed vectors, meaning the discrete affine hull of three one-hot concatenated vectors contains all four possible combinations.
Substituting the expression for (a(2), from equa- tion 4, this becomes
(o(z'), = E a: ((((E))) + C()))
2g2
(5)
=o(('), + Σ o.C(),) =E
P
D 1
(+1) 0o0
(-1)
(+1)-
(3)
From equation 5, we can conclude that (o(z', B(x)) es- timates (ot(z", Exx) perfectly up to a constant error that does not depend on æ. This difference of constant is ab- sorbed by the partition function and hence the conditional densities match: p(x=|=)) p(x|2²).
In Section D.5, we generalize this finding formal mathe- matical characterization of discrete affine hulls, providing visualization method. In Section D.7, we show how discrete affine hulls generalize the extrapolation of additive func- tions studied in Dong & Ma (2022); Lachapelle et al. (2024) over discrete domains. We also show how our results lead to sharp characterization of extrapolation of these functions. Throughout, "affine hull"" refers to the discrete affine hull.
Classifier based on conditional density p(a[=). If, on data from training distribution p. we were able to train a good conditional density estimate p(x]2), V2 E gtrain, then Theorem implies that p(x|2') will also be a good esti- mate of p(x]=²) for new unseen attributes 2' E DAff gtaai))). then straightforward to obtain a classifier that generalizes to compositionallly test distribution q. Indeed, we have
Provided Z C DAff (Zttran,.
3.2. Extrapolation of Conditional Density
We learn set of conditional probability densities p(x+l=) 2(x) exp (o()), Yz E Ztrain by maximizing
((<)22)9(<3))
<(=1<x) = Exrezion ((x)=")q<=")
the likelihood over the training distribution, where É de- notes the estimated energy components and 2 denotes the estimated partition function. Under perfect maximum like- lihood maximization p(x|z) p(z[z) for all the train- ing groups E Ztrain We can define p(x[z) for all E 2 beyond Ztrain in natural way as follows. For each 2 E 2x, we have estimated the energy for ev- cry individual component z; denoted Éx(x, z.). We set 2(x) foa exp (- (a()), E(x)) dr and the density for 1
p( 2')((''²)
Ex p(z))="01(=")
E zie2 p(1|2 p(a(2"))(>")
((("")
(6)
where we used the property of compositional shifts ((+22) = p(x|z). If we know test group prior q(=') (or e.g. assume it to be uniform), we can directly use the expres- sion in RHS to correctly compute the test group probabilities ( [z), even for groups never seen at training.
each 2 E z=, [[((=)))= - exp (- (o(x),
2
3.3. Extrapolation of Discriminative Model
Theorem 1. If the true and learned distribution (p(-12) and p(-12)) are AED, then 2)) = p(-12), E gtrain
In Section 3.2, we saw how we could, in principle, obtain classifier that generalizes under compositional shift, by first training energy based conditional probability density models p(x(2). However learning such a model requires dealing with the problematic partition function throughout training. Indeed making a gradient step to maximize its log likelihood with respect to parameters 0 involves estimating the gradient of its log partition function Vo log 2(2; 0) =
p(-12/) = v( E DAff (2train
The result above argues that so long as the group z' is in the discrete affine hull of gtrain, the estimated density ex- trapolates to it. We provide a proof sketch ahead, with the complete proof in Appendix D.
Proof sketch: Under perfect maximum likelihood maxi- mization p(x|z) = p(z[z), V2 Ztrain Replacing these densities by their expressions and taking their log we obtain
Va log f exp - (a(=,, B(a; 0)) dr which is typically intractable. This difficulty in training energy-based models is a well known open problem. While crude stochastic approximations of this gradient might be obtained via e.g. Contrastive Divergence (Hinton, 2002) or variants of more expensive MCMC sampling, no unbiased computationally efficient solution known in the general case.
(o(2), B()) = (o(x), E(τ)) + C(z).Vz traiin
where C(2) = log (Z(=)/((=))
4

---

Compositional Risk Minimization
But is it really necessary to precisely model the conditional density of high dimensional x, when our goal is simply to predict few classes and attributes 2, given x? We will now develop an alternative approach, Compositional Risk Mini- mization (CRM), that achieves a similar extrapolation result as Theorem 1, while being based on simple discriminative classifier training. It sidesteps the need and difficulties of explicitly modeling p(z/2z) and doesn't 't require dealing with the partition function throughout training.
where, B is the extrapolated bias defined (²)=)=
exp ( (o())),
log Ex Σ Eggavin exp (-(0(E), log B(2)
where F., B are the solutions from optimization equation 9.
Note that B(2) was learned for all E 2train but never for = E geest, hence the necessity of extrapolation B*. Each of these steps is easy to operationalize. We explain the process and provide pseudocode in Section 4.
Observe that if we apply Bayes rule to the AED p(z[2) in equation 2. we get
Theorem 2. Consider the setting where follows AED
p(z|2)p(2) p(2(2)= = Emeggiin p(x(z")pp())
v/2 E 2x, the test distribution Q saaiisfie compositional shift characterization and Ztest - (2 If p(=(z) p(=| V22 E Ztrain Vx RR" and ((2) ((2), Vz E gtest then the output of CRM (equation 10) matches the test dis- tribution, i.e., ((=|<c) = ((=|xr) V2 E grest Vx E RR..
exp (-(((2)) B(x)) + logp(z) - log 2(z)) Ex'ginn exp (o('),, E(x)) + logp(z)) - log 2(1))
A complete proof is provided in Appendix D.2. Observe that p(-[z)- p(-[z) is condition that even model trained via ERM can satisfy (with sufficient capacity and data) but it cannot match the true g(-l2r). In contrast, CRM optimally adjusts the additive-energe classifier for the compositional shifts. CRM requires the knowledge of test prior q(z) but the choice of uniform distribution over all possible groups is reasonable one to make in the absence of further knowledge. Notice how learned bias B(2) can only be fitted for x E gtrain remaining undefined for z' & gtrain But we can compute the extrapolated bias B (z"), Vz E Zuust, based remarkably on only data from the train distribution.
We thus define our additive energy classifier as follows. To guarantee that we can model this p(z|z), we use a model with the same form. For each gtrain
exp (- (o(t)) B(x)) logf(z) - B(x))
p(2|z))
Σ Erggturin exp( (o(2),,, B(x)) + log p(z) B(x²) (7)
where p(2) is the empirical estimate of the prior over z, i.., p(2), E R" + Rmn is function to be learned, bias B is a lookup table containing a learnable offset for each combination of attribute. Given a data point (x, x), loss {(2. ((=))) = - log p(2(zz) measures the prediction performance of p(-lz)). The risk, defined as the expected loss, corresponds to the negated conditional log-likelihoode
Illustrating CRM's adaptation to test distribution. To better convey how CRM can adapt to the Bayes optimal classifier of the test distribution, we provide example. Consider a two-dimensional setting, where the distribution of x E R2 conditioned on the attributes 21 E {-1, 1} and 22 E 1) is a Gaussian with mean (21, 22) and identity covariance. Suppose the training groups are drawn with equal probability and can take one of the following three possible values (+1, +1), observe data from the group (-1, -1) during training, but at test time we draw samples from all the four groups with equal probability. First, we can show that the above distribu- tion can be expressed as an additive energy distribution, as Ear, (#1, 22)) = || (221, 0) 1|2 ] || (0, (0,222)]]).. For the task of classifying groups, the Bayes optimal classifier has a closed form solution, where each decision region an intersection of two half-spaces. Figure 2 shows that CRM learns the Bayes optimal classifier on the training distribu- tion, enables shifting from train prior to test prior to yield the Bayes optimal classifier for the test distribution, and cor- rectly generalize to the unseen (-1 - 1) group. For further details, including illustration of the failure of ERM-trained (10) binary classifier on this problem, see Appendix E. Beyond
R(p) = ((,,,)) [
E(e,))= log (o(((()))
(8)
In the first step of CRM, we minimize the risk R.
-1). We do not
E, B E arg min R(p)
(9)
B
If the minimization is over arbitrary functions, then p(-1z) p(-Iz), Va E IR". In the second step of CRM, we compute our final predictor ((2(x) as follows. Let ((2) be an estimate of the marginal distribution over the attributes q=) with support Żtest For each 2 E ttest
exp ( (- (d(z) A(x)) + log (() - - B²))))
((z|>)
Σ Exegum exp (a(x1), B(x)) (og((z)) - B- (x²)
5

---

Compositional Risk Minimization
Train Prior
Uniform Prior
tains all the possible drd combinations with high probabil- ity (greater than 1 - )) hence CRM achieves CPE. We want to emphasize this surprising finding: with almost linear growth in m and d, CRM generalizes to exponentially many dmd groups
Theorem 3. Consider the setting where p(-[z) follows AED
Vz E 2x, Ztrain comprises of s attribute vectors & drawn uniformly at random from 2*, and the test distribution q satisfies compositional shift characterization with Z
2 If 2 2c(md d'log(d), where d is sufficiently large, p(z|z) p(z|z), Vz E Ztrain Vx E "R",, q(()) q(z), Va e z then the output of CRM (equation 10) matches the test distribution, i.e., ((2[z) = q(2[xr) Vz E 2*, 2x,Vx E R", with probability greater than 6'
Figure 2. Extrapolating to an unseen test group. The distribu- tion to model corresponds mixture of 4 Gaussians. But the (-1,-1) group (pink dashed) has zero prior probability in the training distribution, i.e. absent from the training set. As result, discriminative training (left) learns only three decision re- gions and would misclassify test points from (-1, CRM adjusts the prior at test time (right) to uniform distribution over all attribute combinations, cnabling it to recover four decision re- gions and correctly generalize to the unseen (-1,-1) group. The additive energy form composes the attributes' information to yield this unseen group's location. Decision regions were obtained from finite-data simulations, leading to minor imperfections.
In Appendix D.3, we first present the proof for the case with m = 2 attributes and provide visual illustrations to assist the reader. It is then followed by the more involved proof for the general case of m attributes in Appendix D.4.
3.4. Further Insights CRM
Does test distribution belong to the affine hull of train distributions? A key implication of the AED assump- tion is that the energy for novel group 2' E DAff (trai)) at test time can be expressed as an affine combination of the energies of the training groups, i.e., (o(='), F(x)) = Eggzvai 0g (a()), E(x)) (check Lemma I in Appendix D for details). However, this does not imply that the condi- tional density q<<|=') is an affine combination of the train- ing conditional densities {p(x|>) 2 Ztrain). Instead, we have the following relationship.
this toy example, we highlight that the additive energy form supports modeling nearly arbitrarily complex distributions.
Analyzing growth of Discrete Affine Hull. In the dis- cussion so far, we have relied on crucial assumption that the attribute combinations in the test distribution are in the affine hull. Is this also a necessary condition? Can we generalize to attributes outside the affine hull? We consider the task of learning p(-12) from Theorem and the task of learning (-|1z) from Theorem 2. In Appendix D.6, we show that the restriction to affine hulls is indeed necessary.
log (q(x)2²))
Σ zu
ag lglogp(z/
Under the assumption of compositional shifts 2 is only restricted to be a subset of the Cartesian product set zx but our results so far have required us to restrict the sup- port further by confining it to the affine hull, i.e., gtest C DAff(Ztraing C xx. This leads us to a natural question. If the training groups that form gtrain are drawn random, then how many groups do need such that the affine hull captures xx, i.e., DAff (Ztrain) = zx, at which point CRM can achieve Cartesian Product Extrapolation (CPE). Another way to think about this is to say, how fast does the affine hull grow and capture the Cartesian product set 2 xX? Next, we answer this question.
- -log|
exp
az log p((e)) dr 2
Σ
(12)
Please check Appendix D.2 for the derivation. In contrast, prior works impose stronger assumption that the test dis- tribution should lie in the convex/affine hull of the train distributions (Krueger et al., 2021; Qiao & Peng, 2023; Yao et al., 2023).
Why Additive Energy Classifier? In Appendix D.2, for novels group z' E DAff(z) we derive the following relationship between the test classifier for the novel group q(=1|a |z:) and the training classifiers {r(z|0r | E gtrain).
Consider the the general setting with m attributes, where each attribute takes d possible values, leading to dmm possible attribute combinations. Suppose we sample attribute vec- tors 2 that comprise the support gtrain uniformly at random (with replacement) from these ddd possibilites. In the next theorem, we show that if the number of sampled attribute vectors exceeds 20(md + d d(oog(d) then (2tran) con-
q(>'(1²) [<) Softmax (log ((²g(())+++)))))))))))))) +
Σ z
ag log p(212)
log Exmp(x) exp
Σ
a, log log p(z)x)))
(13)
6

---

Compositional Risk Minimization
This equation is central in deriving the classifier d(= [==] (equation 10) in the second step of CRM, where substi- tute p(z(ær) with the learned additive energy classifier ((=|æ) (equation 9). While alternative methods could be used to estimate p(=|zz), but they would have to separately infer the affine combination weights az.. To address this, we adopt the additive energy classifier (equation 7), which simplifies the computation of (()) Crucially, avoids the need to explicitly estimate the affine weights (ag), rather the re-
Algorithm 1 Compositional Risk Minimization (CRM)
Input: Training set Dirain = [(x,,)) Output: Classifier parameters 0, W. B Training:
Estimate train prior p(2) based on group counts Dirain
Compute the energy terms as E(x) = ((;;()
Train additive energy classifier (e.q. 7) by emprirical risk minimization: 0, W, B E argming ", , I(())
CRM Step 1:
quired adjustment is absorbed into single updated bias term B* (equation 11).
Estimate extrapolated bias B (e.q. 11) via an average over training examples.
CRM Step 2:
Inference on test point
Why Compositional Risk? To better understand the com- positional risk formulation, note that the classical ERM objective (8) can be restated as follows,
Set ((2) as uniform prior over all groups.
Compute test group probabilities q(z l2) via e.q. 10, using E and B learned during training.
R(D) Σ ((2))( ZE2X
5. Experiments
where R(p]z) = Ex-ppar =) ((=.)( Note that in the above summation p(z) is zero on all groups that are not in the support of the training distribution. However, to tackle compositional shifts, we want to learn predictors that instead minimize the following compositional risk.
5.1. Setup
We evaluate CRM on widely recognized benchmarks for subpopulation shifts (Yang et al., 2023b), that have attributes x = (y. a). where y denotes the class label and a denotes the spurious attribute (y and a are correlated). However, the standard split between train and test data mandated in these benchmarks does not actually evaluate compositional gen- eralization capabilities, because both train and test datasets contain all the groups (2train Z = _*). Therefore, we repurpose these benchmarks for compositional shifts by discarding samples from one of the groups (2) in the train (and validation) dataset; but we don't 't change the test dataset, i.e., & Ztrain but = E Ztest, Let us denote the data splits from the standard benchmarks (Dtrain: Dval, Drest) Then we generate multiple variants of compositional shifts { (DDrs train D Drest) 2 E _x*, where Drra and D are generated by discarding samples from Drrain and Dval that belong to the group z.
Reeap (p)= = Bzng(=) oz= ((-x)) (())))
Σ (()l(() *
(14)
as we have p(z|z)= = (<||)) Vz E zx. In the above ob- jective, g(2) can be zero on groups that have zero probability under p(=). Hence, the minimizer of expected risk R(p) can be different from the minimizer of composi- tional risk (T). In Theorem 3, we show that our approach (CRM) outputs the Bayes optimal predictor and hence it provably minimizes Reomp
Following this procedure, we adapted Waterbirds (Wah et al., 2011), CelebA (Liu et al., 2015), MetaShift (Liang & Zou, 2022), MultiLL (Williams et al., 2017), and CivilCom- ments (Borkan et al., 2019) for experiments. We also exper- iment with the NICO++ dataset (Zhang et al., 2023), where we already have gtrain Ç gtest = 2* as some groups were not present in the train dataset. However, these groups are still present in the validation dataset (Zval = Z*). Hence, the only transformation we apply to NICO+ ples from the validation dataset so that 2train Zual, Note that our benchmarks cover diverse scenarios, with binary (Waterbirds, CelebA, MetaShift) and non-binary attributes (MulüNLI, CivilComments, NICO++), resulting in total groups varying from 4 to 360 (NICO++).
4. Algorithm for CRM
In a nutshell, CRM consists of: a) training additive energy classifier ((=(ar (e.q. 7) by maximum likelihood (e.q.9) for trainset group prediction; b) compute extrapolated biases 8" (e.q.11); c) infer group probabilities compositionally shifted test distribution using ((2[<z) (e.q 10). Algorithm provides the associated pseudo-code, where we have a ba- sic architecture using a deep network backbone p(z; 0) fol- lowed by a linear mapping (matrix 2 For the case where we have attributes = (g. a) (illustrated in Figure 1), we provide a detailed algorithm and PyTorch implementation in Appendix C.
drop sam-
Instead of linear, we could use separate non-linear heads to obtain the energy components for each attribute.
For baselines, we train classifiers via ERM, Group Dis-
7

---

Compositional Risk Minimization
tributionally Robust Optimization (GroupDRO) (Sagawa
Average Acc
WGA
al., 2019), Logit Correction (LCC) (Liu et al., 2022b), pervised logit adjustment (sLA) (Tsirigotis al., 2024), Invariant Risk Minimization (IRM) (Arjovsky et al., 2019), Risk Extrapolation (VREx) (Krueger et al., 2021), and Mixup (Zhang et al., 2017). all cases we employ a pre- trained architecture as the representation network Φ, fol- lowed by a linear layer W to get class predictions, and fine-tune them jointly (see Appendix F.3 for details).
Dataset
Method
WGA
(No Gronpe Droppod)
ERM
77.9
43.0
62.3
G-DRO
77.9
42.3
87.3
Waterbirds
LC
88.3
75.5
88.7
sLA
89.3
77.3
89.7
CRM
87.1
78.7
86.0
ERM
85.8
39.0
52.0
G-DRO
89.2
67.8
91.0
CelebA
LC
91.1
57.4
90.0
For evaluation metrics, we computed the average accuracy, group- bouupbbllanced accuracy, and worst-group accuracy (WGA) on the test set. Due to imbalances in group distribution, a method can obtain good average accuracy despite having bad vorst-group accuracy. Therefore, WGA more indica- tive of robustness to spurious correlations (Appendix F.2).
SLA
90.9
57.4
86.7
CRM
91.1
81.8
89.0
ERM
85.7
60.5
63.0
G-DRO
86.0
63.8
80.7
MetaShift
LC
88.5
68.2
80.0
sLA
88,4
63.0
80.0
CRM
87.6
73.4
74.7
5.2. Results
ERM
68.4
7.5
68.0
G-DRO
70.4
34.3
57.0
Table shows the results of our experiment, where due to space constraints, we only compare with the best performing baselines and don't oon't report the group balanced accuracy and standard error over random seeds. Complete results with comparisons with IRM, VREx, and Mixup are provided in Appendix G.I (Table 5). For each dataset and metric, we re- port the average performance over its various compositional shift scenarios ((D D Drest) ) | = E 2*) (detailed results for all scenarios are Appendix G.2). In all cases, CRM either outperforms is competitive with the baselines in terms of worst group accuracy (WGA).
MultiNLI
LC
75.9
54.3
74.3
SLA
76.4
55.0
71.7
CRM
74.3
58.7
74.7
ERM
80.4
55.9
61.0
G-DRO
80.1
61.6
64.7
Civil
LCC
80.7
65.7
67.3
Comments
sLA
80.6
65.6
66.3
CRM
83.7
67.9
70.0
ERM
85.0
35.3
35.3
G-DRO
84.0
36.7
33.7
NICO+
LC
85.0
35.3
35.3
sLA
85.0
33.0
35.3
CRM
84.7
40.3
39.0
Further, for Waterbirds and MultiNLI, while the logit adjust- ment baselines appear competitive with CRM on average, we look more closely at the worst case compositional shift scenario, we find these baselines fare much worse than CRM. For Waterbirds, LC obtains 69.0% WGA while CRM obtains 73.0% WGA for the worst case scenario of dropping the group (0, 1) (Table 6). Similarly, for MultiNLI, sLA ob- tains 19.7% WGA while CRM obtains 31.0% WGA for the worst case scenario of dropping the group (0, 0) (Table 9).
Table 1. Robustness under compositional shift. We compare the proposed CRM method to haseline ERM classifier training with no group information, and to robust methods that leverage group labels: G-DRO, LC, and sLA. We report test Average Accuracy and Worst Group Accuracy (WGA) (mean over random seeds), averaged a group is dropped from training and validation sets. Last column is WGA under the dataset's 's standard subpopulation shift benehmark, i.c. with no group dropped. All methods have harder time to generalize when groups are absent from training, but CRM appears consistently more robust. Full results with balanced group accuracy, standard error across random seeds, and baselines IRM, VREx, and Mixup in Appendix (Table 5).
We also report the worst group accuracy (other metrics in Table 14, Appendix G.5) for the original benchmark (Dtriin: Dval, Prrain), which was not transformed for compo- sitional shifts, denoted WGA (No Groups Dropped). This can be interpreted as the "oracle" performance for that benchmark, and we can compare methods based on the performance drop in WGA due to discarding groups in com- positional shifts. ERM and GroupDRO appear the most sensitive to compositional shifts, and the logit adjustment baselines also show sharp drop for the CelebA benchmark; while CRM is more robust to compositional shifts.
additional attributes (total m = 5 attributes). We provide setup details in Appendix G.3, and Table 12 presents the associated results. We find that CRM still continues to be the superior method (as per WGA), in fact, the difference in WGA between CRM and runner up bascline increases to 28.9%. as compared to 14% for the case of CelebA 2- attribute (Table 2). Hence, the benefits with CRM become more pronounced for multiple spurious attributes settings.
Multiple spurious attributes case. We benchmark CRM for class label prediction with multiple spurious attributes. For this, we augment the CelebA benchmark to have three
Importance of extrapolating the bias. We conduct an ablation study for CRM where we test variant that uses
8

---

Compositional Risk Minimization
Method
Waterbirds
CelebA
MetaShift
MulitNII
CivilComments
NICO++
CRM (B)
55.7 (1.0)
58.9 (0.4)
58.7 (0.6)
30.4 (2.6)
52.4 (0.7)
31.0 (1.0)
CRM
78.7 (1.6)
8188 (1.2)
0.74
58.7 (1.4)
67.9 (0.5)
40.3 (4.3)
Table Importance bias extrapolation. We report Worst Group Accuracy, averaged as a group dropped from training and validation (standard error based on 3 random seeds). CRM (B) is ablated version of CRM where we use the trained bias B instead of the extrapolated bias B" mandated by our theory. The extrapolation step appears crucial for robust compositional generalization. Merely adjusting logits based on shifting group prior probabilities does not suffice.
the learned bias B (e.q. 9) instead of the extrapolated bias 8" (e.q. 11). Results are presented in Table 2. They show significant drop in worst-group accuracy if we use the learned bias instead of the extrapolated one. Hence, our theoretically grounded bias extrapolation step is crucial to generalization under compositional shifts. In Appendix G.4 (Table 13) we conduct further ablation studies, showing the impact of different choices of the test log prior.
and sLA (Tsirigotis et al., 2024) that use logit adjustment for group robustness. There are many other interesting ap- proaches that were proposed, see the survey (Zhou et al., 2022) for details. The theoretical guarantees developed for these approaches (Arjovsky et al., 2019; Rosenfeld et al., 2020; Ahuja et al., 2020) require a large diversity in terms of the environments seen the training time. In our setting, we incorporate inductive biases based on additive energy distributions that help us arrive at provable generalization with limited diversity in the environments.
Further, we analyze CRM's performance with varying num- ber of train groups, details in Appendix G.6.
7. Conclusion
6. Related works
We provide a novel approach (CRM) based on flexible ad- ditive energy models for compositionality in discriminative tasks. CRM can provably extrapolate to novel attribute combinations within the discrete affine hull of the training support, where the affine hull grows quickly with the train- ing groups to cover the Cartesian product extension of the training support. Our empirical results demonstrate that the additive energy assumption is sufficiently flexible to yield good classifiers for high gigh-dimessional images, and CRM is able extrapolate novel combinations in DAff gtrain), without having model high-dimensional p(x (z) nor hav- ing to estimate their partition function. CRM is a simple and efficient algorithm that empirically proved consistently more robust to compositional shifts than approaches based on other logit-shifting schemes and GroupDRO.
Due to space constraints, we briefly deseribe the relevant prior works, with a detailed discussion in Appendix A.
Compositional Generalization Compositionality has long been seen as an essential capability (Fodor & Pylyshyn, 1988; Hinton, 1990; Plate et al., 1991; Montague, 1970) on the path to building human-level intelligence. The history of compositionality being too long to cover in detail here, we refer the reader to these surveys (Lin et al., 2023; Sinha et al., 2024). Most prior works have focused on generative aspect of compositionality, where the model needs to recombine individual distinct factors/concepts and generate the final output in the form of text (Gordon et al., 2019; Lake & Baroni, 2023) or image (Liu et al., 2022a; Wang et al., 2024). For image generation in particular, a fruitful line of work rooted in additive energy based models (Du et al., 2020; 2021; Liu al.. 2021: Nie et al., 2021), which translates naturally to additive diffusion models (Liu et al., 2022a; Su et al., 2024). Our present work also leverages an additive energy form, but our focus is on learning classifiers robust under compositional shifts, rather than generative models.
Limitations. A key limitation of our work is the reliance labeled attributes, which may not always be available in practice. Prior works (Pezeshki et al., 2023; Tsirigotis et al., 2024) propose to first infer spurious attributes and then use domain generalization methods, in the context of subpopulation shifts. A promising direction for future research is to extend such methods to the compositional shift setting and integrate them with our approach.
Domain Generalization Generalization under subpopu- lation shifts, where certain groups or combinations of at- tributes are underrepresented in the training data, is well- known challenge in machine learning. GroupDRO (Sagawa et al., 2019) is a prominent method that minimizes the worst-case group loss to improve robustness across groups. IRM (Arjovsky et al., 2019) encourages the model to learn invariant representations that perform well across multi- ple environments. Closely related to our proposed method are the logit adjustment methods, LC (Liu al., 2022b)
Further, based on AED, we provided compositional general- ization guarantees that are applicable in both discriminative and generative tasks. However, we only offer tractable approach using CRM for the discriminative task. Extending these ideas to the generative counterpart, e.g. by starting from equation 12, remains promising direction for future work and could further leverage our theoretical framework.
9

---

Compositional Risk Minimization
Acknowledgements
Dong, K. and Ma, T. First steps toward understanding the extrapolation of nonlinear models to unseen domains. arXiv preprint arXiv:
We thank Yash Sharma for his contribution to early explo- ration of compositional generalization from generative perspective, and Vinayak Tantia for having many years ago helped shed light on the challenges posed by composition- ality in discriminative training. This research was entirely funded by Meta, in the context of Meta's AI Mentorship program with Mila. Ioannis Mitliagkas in his role as Di- vyat Mahajan's ss academic advisor, acknowledges support by an NSERC Discovery grant PRGPIN-2019-41 2019-06512), and Canada CIFAR AI chair. Pascal Vincent CIFAR Associate Fellow in the Learning in Machines & Brains pro- gram. Divyat Mahajan acknowledges support via FRQNT doctoral scholarship (https:: http///ooooog...o..oo////.. http:/o.og//0.111.111 354785) for his graduate studies.
2022.
Du, Y., Li, S., and Mordatch, I. Compositional visual gen- eration with energy based models. Advances in Neural Information Processing Systems, 33:6637-6664 2020.
Du, Y., Li, S., Sharma, Y., Tenenbaum, J., and Mordatch, I. Unsupervised learning of compositional energy concepts. Advances in Neural Information Processing Systems, 34, 2021.
Fodor, J. A. and Pylyshyn, Z. W. Connectionism and cogni- tive architecture: A critical analysis. Cognition, 28(1-2): 3-71, 1988.
Gordon, J., Lopez-Paz, D., Baroni, M. and Bouchacourt, D. Permutation equivariant models for compositional generalization language. In International Conference on Learning Representations, 2019.
Impact Statement
This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
References
Hinton, G. E. Mapping part-whole hierarchies into connec- tionist networks. Artificial Intelligence, 46(1-2):47-75, 1990.
Ahuja, K., Wang, J., Dhurandhar, A., Shanmugam, K., and Varshney, K. R. Empirical or invariant risk minimiza- tion? a sample complexity perspective. arXiv preprint arXiv:2010.1641 2020.
Hinton, G. E. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771- 1800, 2002.
Arjovsky, M., Bottou, L., Gulrajani, I. and Lopez- Paz, D. Invariant risk minimization. arXiv preprint arXiv: 1907.02893, 2019.
Hupkes, D., Dankers, V., Mul, M., and Bruni, E. Compo- sitionality decomposed: How do neural networks gen- eralise? Journal of Artificial Intelligence Research, 67: 757-795, 2020.
Bayat, R., Pezeshki, M., Dohmatob, E., Lopez-Paz, D., and Vincent, The pitfalls of memorization: When memo- rization hurts generalization. Proceeding international conference on representaiton learning (ICLR), 2025.
Idrissi, B. Y., Arjovsky, M., Pezeshki. M., and Lopez-Paz, D. Simple data balancing achieves competitive worst- group-accuracy. In Conference Causal Learning and Reasoning, pp. 336-351. PMLR, 2022.
Borkan, D., Dixon, L., Sorensen, J., Thain, N., and Vasser- man, L. Nuanced metrics for measuring unintended bias with real data for text classification. In Companion pro- ceedings of the 2019 world wide web conference, pp. 491-500, 2019.
Janzing, D. and Schölkopf, B. Causal inference using the algorithmic markov condition. IEEE Transactions Information Theory, 56(10):5168 -1994 2010.
Kang, B., Xie, S., Rohrbach, M., Yan, Z., Gordo, A., Feng, J., and Kalantidis, Y. Decoupling representation and classifier for long-tailed recognition. arXiv preprint arXiv: 10X0009910009/177 2019.
Brady, J., Zimmermann, R. S., Sharma, Y., Schölkopf. B., Von Küggelgen, J., and Brendel, W. Provably learning object-centric representations. In International Confer- ence on Machine Learning, pp. 3038-3062. PMLR, 2023.
Kim, N. and Linzen, T. Cogs: A compositional gener- alization challenge based on semantic interpretation. In Proceedings of the 2020 conference on empirical methods in natural language processing (emnlp), pp. 9087-9105, 2020.
Devlin, J., Chang, M.-W., Lec, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. arXiv preprint arXiv:1810.04805, 2018.
10

---

Compositional Risk Minimization
Kirichenko, P., Izmailov, P., and Wilson, A. G. Last layer re-training is sufficient for robustness to spurious correla- tions. arXiv preprint arXiv:2204.02937 2022.
Menon, A. K., Jayasumana, S., Rawat, A. S., Jain, H., Veit,
A., and Kumar, S. Long-tail learning via logit adjustment. arXiv preprint arXiv:2007.07314 2020.
Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Bi- nas, J., Zhang, D., Le Priol, R., and Courville, A. Out-of- distribution generalization via risk extrapolation (rex). In International conference on machine learning, pp. 5815- 5826. PMLR, 2021.
Montague, R. Pragmatics and intensional logic. Synthese, 22(1):68-92) :6--94, 1970.
Nie, W., Vahdat, A., and Anandkumar, A. Controllable and compositional generation with latent-space energy-based models. Advances in Neural Information Processing Systems, 34, 2021.
Lachapelle, S., Mahajan, D., Mitliagkas, I., and Lacoste- Julien, S. Additive decoders for latent variables identifi- cation and cartesian- product extrapolation. Advances in Neural Information Processing Systems, 36, 2024.
Nikolaus, M., Abdou, M., Lamm, M., Aralikatte, R., and Elliott, Compositional generalization in image cap- tioning, arXiv preprint arXiv: 1909. 04402, 2019.
Parascandolo, G., Kilbertus, N., Rojas-Carulll Carulla, M., and
Lake, B. and Baroni, M. Generalization without systematic- ity: On the compositional skills of peatuence-t--equeeee recurrent networks. In International conference on ma- chine learning, pp. 2873-2882. PMLR, 2018.
Schölkopf, B. Learning independent causal mechanisms. In International Conference Machine Learning, pp. 4036-4044 PMLR, 2018.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E.,
Lake, B. M. and Baroni, M. Human-like systematic general- ization through meta-learning neural network. Nature, 623(7985):115 115-))1)111-11111 121, 2023.
De Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A. Automatic differentiation in pytorch. 2017.
Pezeshki, M., Bouchacourt, D., Ibrahim, M., Ballas, N., Vincent, P., and Lopez-Paz, D. Discovering environments with xrm. arXiv preprint arXiv:2309. 167448 2023.
Liang, W. and Zou, J. Metashift: A dataset of datasets for evaluating contextual distribution shifts and training conflicts. arXiv preprint arXiv:2202.06523: 2022.
Plate, T. et al. Holographic reduced representations: Con- volution algebra for compositional distributed representa- tions. In IJCAI, pp. 30-35, 1991.
Lin, B., Bouneffouf, D., and Rish, I. A survey on com- positional generalization in upplications. arXiv preprint arXiv:2302.01067, 2023.
Qiao, F. and Peng. X. Hopology-awart robust optimiza- tion for out-of-distributio generalization. arXiv preprint arXiv:2307. 13943, 2023.
Liu, N., Li, S., Du, Y., Tenenbaum, J., and Torralba, A. Learning to compose visual relations. Advances in Neural Information Processing Systems, 34, 2021.
Ren, J., Yu, C., Ma, X., Zhao H., Yi, S., al. Bal- anced meta-softmax for long-tailed visual recognition. Advances in neural information processing systems, 33: 4175-4186, 2020.
Liu, N., Li, S., Du, Y., Torralba, A., and Tenenbaum, J. B. Compositional visual generation with composable dif- fusion models. European Conference on Computer Vision, pp. 423-439. Springer, 2022a.
Rosenfeld, E., Ravikumar, P., and Risteski, A. The risks of invariant risk minimization. arXiv preprint arXiv:2010.05761, 2020.
Liu, S., Zhang, X., Sekhar, N., Wu, Y., Singhal, P., and Frrnandez-Granda, C. Avoiding spurious correlations via logit correction. arXiv preprint arXiv:22/2.01433, 2022b.
Sagawa, S., Koh, W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case gener- alization. arXiv preprint 1rXXiv1911.0080 2019.
Liu, Y., Alahi, A., Russell, C., Horn, M., Zietlow, D., Schölkopf, B., and Locatello. F. Causal triplet: An open challenge for intervention-centrr causal representation learning. In Conference on Causal Learning and Reason- ing, pp. 553-573. PMLR, 2023.
Schug, S., Kobayashi, S., Akram, Y., Wolczyk, M., Proca, A., Von Oswald, J., Pascanu, R., Sacramento. J., and Steger, A. Discovering modular solutions that generalize compositionally. arXiv preprint veXiv:2312.1500) 2023.
Liu, Z., Luo, P., Wang, X., and Tang, X. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pp. 3730- 3738, 2015.
Sinha, S., Premsri, T., and Kordjamshidi, P. A survey on compositional learning of ai models: Theoretical and experimetnal practices. arXiv preprint arXiv:2406.08797 2024.

---

Compositional Risk Minimization
J., Liu, N., Wang, Y., Tenenbaum, J. B., and Du, Y. Com- positional image decomposition with diffusion models. arXiv preprint arXiv:2406. 19298, 2024.
Tsirigotis, C., Monteiro, J., Rodriguez, P. Vazquez, D., and Courville, A. C. Group robust classification without any group information. Advances in Neural Information Processing Systems, 36, 2024.
Wah, C., Branson, S., Welinder, P., Perona, P., and Belongie, S. The caltech-uesd birds-200-2011 dataset. 2011.
Wang, Z., Gui, L., Negrea, J., and Veitch, V. Concept alge- bra for ascore-basede text-controlled generative models. Advances in Neural Information Processing Systems, 36, 2024.
Wiedemer, T., Brady, J., Panfilov, A., Juhos, A., Bethge, M., and Brendel, W. Provable compositional gener- alization for object-centric learning. arXiv preprint arXiv:2310.05322, 2023.
Wiedemer, T., Mayilva,anan, P., Bethge, M., and Brendel, W. Compositional generalization from first principles. Advances in Neural Information Processing Systems, 36, 2024.
Williams, A., Nangia, N., and Bowman, S. R. A broad- coverage challenge corpus for sentence understanding through inference. arXiv preprint 1eXiv:1704.05444 2017.
Yang, L.., Zhang, Z., Song, Y., Hong, S., Xu, R., Zhao, Y., Zhang, W., Cui, B., and Yang, M.-H. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1-39, 2023a.
Yang, Y., Zhang, H., Katabi, D., and Ghassemi, M. Change is hard: A closer look at subpopulation shift. In Interna- tional Conference on Machine Learning, 2023b.
Yao, H., Yang, X., Pan, X., Liu, S., Koh, P. W., and Finn, C. Improving domain generalization with domain relations. arXiv preprint arXiv:2302.02609 2023.
Zhang, H., Cisse, M., Dauphin, Y. N., and Lopez-Paz, D. mixup: Beyond empirical risk minimization. 2017.
Zhang, X., Hc, Y., Xu, R., Yu, H., Shen, Z., and Cui, Nico++: Towards better benchmarking for domain gen- eralization. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 16036- 16047, 2023.
Zhou, K., Liu, Z., Qiao, Y., Xiang, T., and Loy, C. C. Do- main generalization: A survey. IEEE Transactions Pattern Analysis and Machine Intelligence, 43(4):4396- 4415, 2022.
12

---

Compositional Risk Minimization
Appendix
List of Contents
The content in the Appendix has been organized as follows.
Appendix A: Further Discussion on Related Works
Appendix B. Additive Energy Distributions versus Additive Decoder Models: A Closer Look
Appendix C: CRM Implementation Details
Appendix C.I Algorithm for Two Attribute Case
Appendix C.2 PyTorch Implementation for Two Attribute Case
Appendix D: Proofs
Appendix D.I: Proof for Theorem 1: Extrapolation of Conditional Density
Appendix D.2: Proof for Theorem 2: Extrapolation of CRM
Appendix D.3: Growth of Discrete Affine Hull for 2 Attribute Case
Appendix D.4: Proof for Theorem 3: Growth of Discrete Affine Hull for m Attribute Case
Appendix D.5: Discrete Affine Hull: Closer Look
Appendix D.6: No Extrapolation beyond Discrete Affine Hull: Proof for Theorem 10
Appendix D.7: Extrapolation of Discrete Additive Functions via Discrete Affine Hulls
Appendix E: Additional Details on CRM's Adaptation to Test Distribution
Appendix F; Experiments Setup
Appendix F.l: Dataset Details
Appendix F:2: Metrie Details
Appendix E3: Method Details
Appendix G: Additional Results
Appendix G.1: Results Averaged over all the Compositional Shift Scenarios
Appendix G.2: Detailed Results for all the Compositional Shift Scenarios
Appendix G.3: CelebA Multiple Suprious Attributes
Appendix G.4: Results for Ablations with CRM
Appendix G.5: Results for the Original Benchmark
- Appendix G.6: CRM's Analysis with Varying Group Size
13

---

Compositional Risk Minimization
A. Further Discussion on Related Works
Compositional Generalization. Compositionality has long been seen as an important capability on the path to building (Fodor& & Pylyshyn. 1988; Hinton, 1990; Plate et al., 1991; Montague, 1970) human-level intelligence. The history of compositionality very long to cover in detail here, refer to these surveys (Lin et al., 2023; Sinha et al., 2024) for more detail. Compositionality is associated with many different aspects, namely systematicity, productivity, substitutivity, localism, and overgeneralization (Hupkes et al., 2020). In this work, we are primarily concerned with systematicity, which evaluates a model's s capability to understand known parts or rules and combine them in new contexts. Over the years, several popular benchmarks have been proposed to evaluate this systematicity aspect of compositionality. (Lake & Baroni, 2018) proposed the SCAN dataset, (Kim & Linzen, 2020) proposed the COGS dataset. These works led to development of several insightful approaches to tackle the challenge of compositionality (Lake & Baroni, 2023; Gordon et al., 2019). Most of these works on systematicity have largely focused on generative tasks, (Liu et al., 2022a; Lake & Baroni, 2023; Gordon al., 2019; Wang et al., 2024), i.., where the model needs to recombine individual distinct factors/concepts and generate the final output in the form of image or text. There has been lesser work on discriminative tasks (Nikolaus et al., 2019), i.e., where the model is given an input composed of a novel combination of factors and it has to predict the underlying novel combination. In this work, our focus is to build an approach that can provably solve these discriminative tasks.
On the theoretical side, recently, there has been a growing interest to build provable approaches for compositional gen- eralization (Wiedemer et al., 2023; 2024; Brady et al., 2023; Dong & Ma, 2022; Lachapelle et al., 2024). These works study models where the labeling function or the decoder is additive over individual features, and prove generalization guarantees over the Cartesian product of the support of individual features. The ability of a model to generalize to Cartesian products of the individual features is an important form of compositionality, which checks the model's capability to correctly predict in novel circumstances described as combination of contexts seen before. (Dong & Ma, 2022) developed results for additdynamicallyive models, i.e., labeling function is additive over individual features. While in (Wiedemer et al., 2023), the authors considered a more general model class in comparison to Dong & Ma (2022). The labeling function/decoder in (Wiedemer et al., 2023) takes the form f ,In) = C( (z)). 4/ (x..). However, they require strong assumption, where the learner needs to know the funetion C that is used to generate the data. (Lachapelle al., 2024; Brady et al., 2023) extended the results from (Dong & Ma, 2022) to the unsupervised setting. (Lachapelle et al., 2024; Brady et al., 2023) are inspired by the success of object-centric models and show additive decoders enable generative models (autoencoders) to achieve Cartesian product extrapolation. While these works take promising and insightful first steps for provable compositional guarantees, the assumption of additive deterministic decoders (labeling functions) may come as quite restrictive. In particular a given attribute combination can then only correspond to a unique observation, produced by a very limited interaction between generative factors, not to a rich distribution of observations. By contrast an additive energy model can associate an almost arbitrary distribution over observations to a given set of attributes. Hence, we take inspiration independent mechanisms principle (Janzing& & Schölkopf, 2010; Parascandolo al., 2018) for our setting based on additive energy models. In the spirit of this principle, we think of each factor impacting the final distribution through an independent function, where independence is in the algorithmic sense and not the statistical sense. Based on this more realistic assumption of additive energy, our goal is to develop an approach that provably enables zero-shot compositional generalization in discriminative tasks, where the model needs to robustly predict never seen before factor combinations that the input composed of. These additive energy distributions have also been used in generative compositionality (Liu et al., 2022a) but not in discriminative compositionality.
Finally, in another line of work (Schug et al., 2023), the authors consider compositionality in the task space and develop an approach that achieves provable compositional guarantees over this task space and empirically outperforms meta-learnin approaches such as MAML and ANIL. Specifically, they operate in cuudent-teacher framework, where each task has a latent code that specifies the weights for different modules that are active for that task.
Domain Generalization. Generalization under subpopulation shifts, where certain groups or combinations of attributes are underrepresented in the training data, is a well-known challenge in machine learning. Group Distributionally Robust Optimization (GroupDRO) (Sagawa et al., 2019) is prominent method that minimizes the worst-case group loss to improve robustness across groups. Invariant Risk Minimization (IRM) (Arjovsky et al., 2019) encourages the model to learn invariant representations that perform well across multiple environments, Perhaps the simplest methods are SUBG and RWG (Idrissi et al., 2022), which focus on constructing a balanced subset or reweighting examples to minimize or eliminate spurious correlations. There are many other interesting approaches that were proposed, see the survey for details (Zhou et al., 2022). The theoretical guarantees developed for these approaches (Rosenfeld et al., 2020; Arjovsky et al., 2019; Ahuja et al., 2020)
14

---

Compositional Risk Minimization
require large diversity in terms of the environments seen the training time. In our setting, we incorporate inductive biases based additive energy distributions that help arrive provable generalization with limited diversity in the environments.
Closely related to our proposed method are the logit adjustment methods (Kang et al., 2019: Menon et al., 2020; Ren et al., 2020; Bayat et al., 2025) used in robust classification. (Kang et al., 2019) introduced Label-Distribution-Aie (LDAM) loss for long-tail learning, proposing a method that adjusts the logits of a classifier based on the class frequencies in the training set to counteract bias towards majority classes. Similarly, (Menon al., 2020) and (Ren et al., 2020) (Balanced Softmax), modify the standard softmax cross-entropy loss account for class imbalance by shifting the logits according to the prior distribution to prevent memorization towards robust generalization. Closest to our work are the Logit Correction (LC) (Liu et al., 2022b) and Supervised Logit Adjustment (sLA) (Tsirigotis et al., 2024) methods that use logit adjustment for group robustness. LC adjusts logits based on the joint distribution of environment and class label, reducing reliance on spurious features in imbalanced training sets. When environment anmotations are unknown, a second network infers them. Supervised Logit Adjustment (sI.A) adjusts logits according to the conditional distribution classes given the environment. In the absence of environment annotations, Unsupervised Logit Adjustment (uLA) uses self-supervised learning (SSL) pre-train model for general feature representations, then derives biased network from this pre- pre-trained model to infer the missing environment annotations.
Margin
the classes. (Bayat et al., 2025) shifts the logits according to the held-out prediction of examples
B. Additive Energy Distributions versus Additive Decoder Models: A Closer Look
We will explore examples demonstrating the flexibility of additive energy distributions. Consider a set of images, each containing a distinct object that varies in shape, size, and color. We can associate an energy term with each attribute: one detecting shape, another detecting color, and third detecting size. Together, these energy terms define the distribution of images conditioned on an object's shape, color, and size.
Now, contrast this with an additive decoder- based model. At any given pixel, it is unlikely that shape, color, and size attributes interact additively-rather, their interactions are complex and non-trivial. This example illustrates why additive energy distributions naturally model the conjunction of object characteristics, whereas additive decoders struggle with this task.
Another example from different data modality is the CivilComments benchmark (Borkan et al., 2019), where the attributes toxic language (class label) and demographic identity (spurious attribute) interact non-trivially in text space. However, under additive energy distributions, we can model their interactions via energy component that checks whether the language is toxic, and another energy component checks the demographic identity.
We now provide a simple example that can be modeled both by an additive decoder and additive energy distribution. Consider a two dimensional x, where each dimension of x controlled by a different attribute.
(11, x2) = (9(21), (22))
(15)
Observe that we can express the above in terms of an additive decoder x = (g(=)),0) + (0, (z2)). We can also express the above in terms of an additive energy distribution as follows using Dirac delta functions ō as follows.
p(z|2) = 6(x - (9(21), h(22))) = 8(x1 - g(z1))6(z2 - h(22)) = exp (zz( - ((x)))) + Log(6(22-g(6) - 9(22)))).
The above example is meant illustrate the point that if there are different subsets of the pixels, where each subset is controlled or impacted by different set of attributes, then they can be modeled via additive energy distribution or an additive decoder equivalently.
15

---

Compositional Risk Minimization
C. CRM Implementation Details
C.1. Algorithm for Two Attribute Case
Algorithm 2 Compositional Risk Minimization (CRM) for 2 Attribute Case
Input: training set Dirain with examples (x, y, a), where y is the class to predict and a an attribute spuriously correlated
with y
Output: classifier parameters 0, W, B
Let I,, B Xdd be the log prior and the bias terms.
Define logits: (L) == ((W Φ( id, + (W ((zz0)) +1id +lddytdd )
Define log probabilities: log ply a z; 0, W. L. B) = (FL,B(x) logsumexpl (F
Training:
Estimate log prior traii from Dtrain,
Ltrain <- -0 if (v,, a) absent from Dirain
Optimize 0, W, and B to maximize the log-likelihood over Dirain; 0, W, B <- arg max,,w.B E (j..a)E logp(y, alæ; 0, W, I,train B)
Extrapolate bias: B" + log ( ExeDimin exp(Fu,o(z) logsumexpl ( Frrtain B (z))))
Inference on test point z:
Compute group probabilities, using B", and Lumff log dyd=
_ aiming for shift uniform prior:
q(y, a|z) + exp(log p(y, al2; 0, W, Lmif, _B))
Marginalize over get class probabilities: ((ylx)) Ea Elvlvv a/x)
C.2. Porch Implementation for Two Attribute Case
# Pseudo-code for Compositional Risk Minimizat ion (CRM)
2 import: torch
3 import: torchvision
4 from torch import nn
from torch.mn import functional
6 import load dataloaders
7 import
metric
class CRM (MM.MMccllc)
ннв Architcture of CRM Layer
Args:
dhi Input feature dimension
d_y: Total categories for class labels
daa Total cateogies for spurious attributes a
16
def
init (self, é_phi: int,
d_y: int,
d_a: int,
21
): super (CRM, sel5). init_ self.d phi= a phi syl d_y solf.d_a= a_a
25
self. falf.iiner_iie -iyeer_layer= Inn d_ytd_a)
balf.biag- Parameter (torch.zeres. (self.dy, self.d_a) )) requirress_grra
()
16

---

Compositional Risk Minimization
def forward x)
return self.linear_layen (x)
def CRM CRV_logits (x, phi, layer, log_prier):
u.n Computes the logit as per the additive energy classifier Inputs:
35
Batch of Input images, Expected Shape (batch size, 3, 224, 224) phi: Representation network
ermllayer: CRM layer with linear layer and bias
log_prior: Expected Shape (d_y, d_a)
Returns:
Logits for all groups of shape (batch size, d.y, aa)
"".
44
y erm layer.d
da ca
#Obtain features via the representation network features phi (x)
#E_x is a vector of size d_ytd_a
E_x = erm_ lim_laaer.iiner_Baae
(features)
#Energy components for y, vector of size d_y
E E
53
#Energy components for a, vector of size a H_2 k_x[d_y:d_ytd__a
54
# Energy for all groups
logits -EEI. unsl..zzsueeze(() (1) E_2. unsqueeze )0)) # Integrate log prior and bias
57
28
logits += return logits
crm lim_layyrbbis
def estimate_prion (dataset):
62
"""Return counts of different groups in the "atasot"** d_y= dataset.num_y
d_a- dataset.
counts torch.zeros (d_y,, d_a))
for x, y, countsly,d ay +=
dataset:
67
priors = counts / (datasett)
return priors
71 def CRM_extrapolate_bbie (trainset, phi, crm_layer, log_prior): """Compute extrapolated bias (B (lstar)) as per equation (11)
72
Inputs:
trainset: Train Dataloader
phi: Representation Network
76
crmllayer: CRM layer with linear layer and bias log_prior: Expected Shape (d_y, da)
Returns:
Updated CRM layer with extrapolated bias of shape (d_y, d_))
81
y= trainset.num
a trainset.num
85
CComuute logits on all samples in trainset logits=[]
for logits.append append CRM CRM_logits (x, phi, crm_layer, log_prior) logits= torch. (logits, dim=0)
in trainset:
#Compute extrapolated bias (B" (\star)))
energy_tr= logits prgor + layer.bias
Log probs= jog_probe=
torch. sum( torch.ext (logits) vilw(-1), din=l
17

---

Compositional Risk Minimization
extrapolated_bias- corch.logi torch. m torch. corch.exp(ee (xe(eggggyy_ dim=0)
log_probs
erm erm waw.a..eet.apolaeee...... bias= extrapolated bias return samer
99
def
CRs_test (testset, phi, crn_layer); """Module for evaluating CRM
100
Inputs:
102
testset: Test Dataloader
103
phi: Representation Network
104
crm_layer: CRM layer with linear layer and bias
105
Returns:
107
List of values corresponing to evaluation metric on each test batch
108
109
dyy testset.
111
d_ testset.
112
113
#Set test prior to be uniform log_prior=
114
(1d__ggggdddd____________ ttrrch.
115
116
Cinal_ressl
117
for x, a in testset:
118
#Forward Pass
119
logit CRM_logits (x, phi, crm_layer, log_prior)
120
#Marginalize over attribute to get class label predictions
121
y pred= dogit.sunte
122
#Compute metric of choices (empute_eetrrc(( fy_pre,, y) final_ appp_eeddepppddeeeecc (res)
123
= compute metric
125
126
return final_res
127
128
def
CRM_train (trainset, phi, crm_layer); """Module for training CRM
129
130
Inputs:
131
trainset: Train Dataloader
132
phi: Representation Network
133
erm_layer: CRM layer with linear layer and bias
134
135
Returns:
136
Learned models phi and ormlayer
137
138
139
#Initalize Optimizer
140
opt= oarch.optim..o
141
[phi parameters () +orm parmatttrr.. parameters ())), 1r-0.001,
142
144
145
priors - est asstmmat matemamriiorree (trainset)
146
#Since some group counts can be 0, clamp them to tiny value to avoud inf with log
147
log prior priors. clamp (min=1e-10) ()g(()
148
149
#CRM Step 1 training loop
for epoch in range (1000): for x, y, in trainset: #Forward Pass opt.zero_grad()
152
153
154
logit CRM_logits (x, phi, cam_layer, log_prior)
#Comuute group label
156
getrainset.
18

---

Compositional Risk Minimization
157
Compute cross entropy loss
158
loss F. --coss_entrop (logit. view g. gonong (), reduction- mean') imitimize parameters
139
loss backward ste.
162
163
return phi, layer prior
164
def example_main())
106
167
#lload training and test dataloaders trainset, testset = 1oad_dataloaders, () #Actributes (y,a)
168
160
y trainset.num.
171
a trainset.num.
172
173
#Pretrained representation network
phi= torchvision. models. resnet. resnet50 (pretrained-True) #Feature Dimension
175
176
d_phi phi.feat_din
177
#ORM Layer
178
crn_layer= CRM(d_phl d_y, d_a)
180
#Train
181
phi, crm Jayer, log prior= CRLttrai phi, erm crm_layer)
182
#hhis is where the magic is happening: extrapolating to biases of unseen combinations
184
ern layer CRM extrapolate bias (trainset, phi, erm layer, prior)
185
186
#Test
CRr_test (testset, phi, crm_layer)
19

---

Compositional Risk Minimization
D. Proofs
Remark on proofs We want to emphasize that the proofs developed here are quite different from related works on compositionality (Dong & Ma, 2022; Wiedemer al., 2023). The foundation of proofs is built on a new mathematical object, discrete afline hull. The proof of Theorem 2 cleverly exploits properties of softmax and discrete affine hulls to show how can learn the correct distribution without involving the intractable partition function in learning. The proof of Theorem 3, uses fundamental ideas from randomized algorithms arrive at the probabilistic extrapolation guarantees.
We start with a basic lemma.
Lemma 1. If z' DAff(Z i.e., o(2) ExEzuan 0,0(2), where (1,az) = 1, then (o(2'),E(2)) ExEguai O2 (a(2),
Proof. (o(z'), = EEggtm 0,0(2), = Esgguin 02 ((x))
D.I. Proof for Theorem 1: Extrapolation of Conditional Density
Theorem 1. If the true and learned distribution (p(-1z) and p(-lz)) are AED, then p(-1z) = p(-12), Va E Ztrain p(-12' = p(-l2/) Vz E (
Proof. We start by expanding the expressions for true and estimated log densities below
log [<<[=) = (xx)) + log (z)))
(16)
log [(x|z)] = (o()), ((r) + log (()).
We equate these densities for the training attributes graain For fixed E garain, we obtain that for all x E I
(o(z), = (a(2), E(x)) + C(=),
(17)
where C(=) = log (Z)=)/((=)). Since z' E DAff( (ttrain,,,,) we can write L' Exezun 022, (1, = From Lemma 1, know that (a(='), E(x)) Exegion 0₂ (a(2), (x:).
We use this decomposition and equation 17 to arrive at the key identity below. For all x E R
((zz))E(z)) = Σ as (o(2)
ztain
Σ
[(x) + +(z))
Σ
a (((()), + (Σ Σ 0()(
(18)
ztai
g
= (o(z'), + Σ a C())
2tinn
From this we can infer that
p(x(2))
2(2²) exp
- (o(x'), f(x)))
(19)
- Σ 0()
exp
((2))
ZE train
20

---

Compositional Risk Minimization
We now the fact that density integrates to for continuous random variables (or alternatively the probability sums for discrete random variables). Thus
I f(xl2')dx=11
/ 1 exp ( (o(())),, E (=((z))dr== dr
2(22)
(- E a,((z)) / exp
(- B(x)) )dr do== =1
exp
(12))
Σ a. a,C(z) 2(2')
exp
2(x)
=E Z:wnn
1 exp
Σ a=C(z)) Z(z²)
(20)
2(21)
2
We substitute equation 20 into equation 19 to obtain
(x(2² ((2) exp (=(=',,, p(())2)), Var e R"
(21)
D.2. Proof for Theorem 2: Extrapolation of CRM
Theorem 2. Consider the setting where 2)) follows AED V2 E Z* the test distribution q satisfies compositional shift characterization and 2e C (Ztrain. If p(z])) the output of CRM (equation 10) matches the test distribution, i.e. ((=|>) = <=( V2 E g Vo E Rd.
p(z[2), Yz E Ztrain Vx E IR" and @(2) = q(2). V2 tes then
Proof. Since q follows compositional shifts,
log ((+lz) [gg((z/>) = ((x)), - log 2(2)
(22)
We can write it as - (x),),, E(x)) = logp(z/z) + log Z(z).
Consider z' DAff We can express as a (z')((==)))))))))) = Eeezum (2)(2), where (1,az)=1 =1
We use equation 22 and show that the partition function z' can be expressed as affine combination of partition of the individual points and correction term. We obtain the following condition. Yz gtest, where recall 2 DAff (ttrain,
log (Z(²') = log exp ( (o(2)) F(())) dx
eexp
( (- Σ a (o(2), E(x)) ) dx
log
ZE Z
(23)
= log
exp
Σ (x(og(((122) + log Z(z)) dæ
2tain
ag log Z(z) -
log
exp
ag logg(z(>)) dx ZE 2trin
Σ
2tain
Denote the latter term in the above expression as
R( (((as)eezm)) = log
/ exp
as logg(z(z)) dæ aain
Σ
(24)
21

---

Compositional Risk Minimization
We now simplify log (1(x)|=)) using the property of partition function from equation 23 below. Vz gtes
log (g(x|2/) =
log Z(z))
Σ az [log((x|>) + log 2())) log 2(z')
Σ Oz p(rl)) +
Σ
0₂ log (()) - Σ
o=log - R(f0z):E200)
(25)
2ua
Oi ag_ogp(z|2) - R(fas)
ZE Ztenn
We now simplify the first term in the above expression, i.c. Exeguin 02 log p(x]2)) in terms of p(z|z).
p(z|z)p(z))
Σ
az (p(x|=))
Σ Oz log
p(z)
t
#E 2t
(26)
Σ
(lo (ogg(z//z) - logp(z)) + logp(r)
Olz
xE 2toin
Similarly, R(fa=) R(fa=)azzmin can be phrased in terms of p(z|ar) follows.
R(fa:)eezvnn = log
/ exp
Σ az log p(xlz)) de
Σ
a₂ logg(p)) log (mwp(t)
exp
Σ
a log p p(2)))))
(27)
ZE
ZE
Σ
a log logp(2)+ + 5((ax)22ezmm)
ZE vrain
where S ((faz)ezzn)) =log log (Bexp(a) exp (Esezum O₂ log rog)(=)=))) and Bewp(a) is the expectation w.r. distribution p(z). We equation 26, equation 27 to simplify equation 25 as follows.Vz/ E grest
log (()g(())
Σ
as - S(fa))(((ua))))) + logp(x)
2tai
log 9(>1212 ((21(z)9(z)
Σ
as logp((2)))) ((1az)zezmn) logp(x)
((21)
(28)
- 5((22126210n) + log (p(x)
log (g(='[z)) =
Σ az
((x))
ZE train
We use translation invariance of softmax to obtain
q(21110 (x) Softmax log (o')(=')+ +
Σ logp(=/<) - S((a))
ZE Ztnin
(29)
q(z1 (x)) = Softmax logg(z')+
Σ
as logg(z/x)) log (Eanp(z)
exp
Σ
as log log p(2l)))))]] p))))((())))
2tiin
2E taain
22

---

Compositional Risk Minimization
To avoid cumbersome notation, took the liberty to show only one input to softmax, other inputs bear the same parametrization, they are computed other z's. From the above equation it clear that if the learner knows the marginal distribution over the groups at test time, i.e., g(=) and estimates ((2/x) for all 2 in the training distribution's 's support, i.e., Ztrain, then the learnr can successfully extrapolate to q(2 (x).
Let us now use the additive energy classifier of the form we defined in equation 7 and whose energy E and bias we optimized (equation 9) to match p(z|x) (z), so that:
-
[()) (())) + B(z)
exp
p(2(z))
Σ Zz exp
(o)) B(x)) + ((2) - B()))
Consequently
Σ Og llz(())))
2tamn
Σ ax (o(),, + logp(z) -ogg(t)-(((s))
B(=)))
log
Σ
exp
(a)) ((x)) +ogp(() B()))
Ztrain
(30)
where we used the property that a1) = Let us use this to simplify the last term of equation 29:
Σ a₂ log p(ela)))]
log
Exna(x)
exp
#E 2tan
exp
Exezuai as
(o(=), F(x)) + log p(=) - B(x))
log
Brap(at)
Ezezuin exp
- (a()), F(x))) + log B(3))
Exezrsi as (- (o(),, B(x)))
exp
Σ
az (1ogn())-B(e)))
log
Exmp(a)
cxp
(0(3), B(x)) + log -B(E))
BiEzz zuan exp
=E Ztain
(31)
exp
Exgzrni Os (o(x), B(x))
1)
log
Bxup(())
(- (o(3), B(x)) + logp(=) B(z)) B(z))
+
Σ
as (1og((2)- - B(z))
Σ EE Zunin exp
#E ntiin
exp
(-((()), B(x)) + logp(z) B(x)) 1)
(o(o2)))
Σ a (10g(())- B(2))
log
Bxmp(a)
Σ Esezuin exp
ZE Zown
= +x))+++++++
Σ as (logp(2)- -B(x)
where we used Lemma 1, and B* is as defined in equation 11.
Let us also define c(x) log Exee exp ( (o(),, B(x)) logr(3) - B(z))) so that we can reexpress equation 30
as:
(o(2,, Σ as (= E(x)) + logp(z) - B(=)) c(x)
Σ as log p(zlz) =
(32)
ZE ttrain
2E Ztrain
23

---

Compositional Risk Minimization
Subtracting equation 31 from equation 32 we get:
Σ 02 logg(=/<)) log Eemple) exp
Σ
az log log
3E2
Σ Ot2 (- B(x)) + logp(z) - B(x)) dx)
(2'22)
Σ (1ogp(()))Bit
zE
Σ as ( (x() - (2) -B)))
- (a(z), B(x)) c(x) - (²))
(33)
Substituting this inside equation 29 yields
<<21<x) =
Softmax loga(=")-(a(z") - (0((2) E(x)) - c(x)) B²()))
(34)
Softmax (o(z), B(x)) + logg(=") B* (z)))
where removed the c(x) term softmax is invariant addition of terms that do not depend on 2/. If ((=)) = V2' E gesst then the expression in RHS corresponds to ((=1(zr), as we had defined it in equation 10, before stating our theorem. Thus (==(<z) = ((²'lx). This completes the proof.
D.3. Growth of Discrete Affine Hull for Attribute Case
Before presenting the proof of Theorem 3, we specifically deal with the case m = attributes and present a proof guided with visual intuitions. We believe this proof is easier to understand and helps build intuition for a reader to better grasp the proof in the next section for the general case. We first establish some basic lemmas. In the first lemma below, we consider a setting with two attributes, where each attribute takes two possible values, i.c., 2 and d 2. In this setting there four possible one-hot vectors z 22 23 24. We first show that each z can be expressed as affine combination of the remaining three.
Lemma 2. 1fm 2, d = then there are four possible concatenated one-hot vectors denoted 2 ,,,2 23, 2A, Each z can be expressed affine combination of the remaining.
Proof. Below we explicitly show how each zt can expressed terms of other zd's.
(+1)
(-1)-
+
(+1)
(35)
o
1
(-1)
+ (+1)
+ (+1)
(36)
100
(+1)
+ (+1)
+ (-1)
(37)
24

---

Compositional Risk Minimization
22
21
1
2
1
Figure 3. Setting of two attributes and two possible values per attribute. Illustration of extrapolation from three groups the remaining fourth group. Three dark colored groups indicate the observed groups and the light colored shaded group indicates the group that is the affine combination of the three observed groups.
O (-1)
iin + (+1)-
+ (+1)
(38)
0
We illustrate the setting of Lemma 2 in Figure 3. We now understand implication of Lemma 2. Let consider the setting where m = 2 and d > 2. Consider subset of four groups {(t, j), (t', (), (i. (')).. Under one-hot concatenations these groups are denoted as =² 0..... 1,, 0,0, 0, 1j. OJ. =² 10, .. 1y, 0,0, 0, 1, 0].
first attribute
second attribute
23 =3 10.... 1,....0.. 0,0 0, 21 =- +23 -21. Similarly, we can express every other 2f in terms of rest of 2 in the the set ((i.)), (i', 5), (i,J"), (')}.
....... 0j, and 2 = =0,.... 1e. 0,0, 1 0]. Observe that using Lemma 2, we get
In the setting when m = 2 and d 2 2, the total number of possible values takes is d², Each group recall is associated with attribute vector 2 = [2] where 21 E {1..... .d} and 2n2 E {1....d) d). The set of all possible values of 2 be visualized as x d grid in this notation. We call this groups for training, which shows that with just 2d 1 possible to affine span all the possible d2 groups in the grid G. We leverage the insights from this approach and show that with randomized approach of selecting groups, we continue groups.
d grid as G. We will first describe specific approach of selecting observed
to affine span d' groups with
Denote the set of observed groups as N. Suppose that their affine hull contains all the points in a subgrid S S G of size m x n. Let the subgrid S = {x1,''' ... .m x fyi, ... }. Without loss of generality, we can permute the points and make the subgrid contiguous as follows S {1, (1.... m) x [1....,n]] Next, we add a new point g = (gas 9y) E but g & S. We argue that if IN E (1.... m). then the affine hull of N U {g} contains a larger subgrid of size m x (n + 1). Similarly, we want argue that if By E (I..... ......... n),, then the affine hull of N U {g) contains a larger subgrid of size (m + 1) x n. Define Cz as the Cartesian product of {9x) with {1,... m). i.e., C. ((92, 1), 2). 2).... n).. Define C,, as the Cartesian product of (1, } with {gy), ). i.e. C. {11, gy), (2, 9z), -.. (m,,g))
Theorem Suppose the affine hull of the observed set N contains a subgrid S of size m n. If the new point g (ga my)
shares the x-coordinate with a point in S, and g S, then the the affine hull of N U (g) contains S SUC Cy-
Proof. We write the of observed groups N N (2 The affine hull of N contains S = (1, ... m) x {1, nn.. We observe new group g E S, which shares its coordinate with one of the points in S. Without loss of generality let this point be g (1, n. 1) (if this were not the case, then we can always permute the columns and rows to achieve such a configuration). Consider the triplet- - (2), 1,,,2 2, 2²) = ((1, n), (2, n), (1, n + 1)). Observe that 21, 2² 23 form a 2 x 2 subgrid, where 24 = (2, n 1). We use Lemma 2 to infer that the fourth point 2 = (2, n + 1) on this 2 x subgrid can be obtained as an affine combination of this triplet, i.c., 2² = az" + 322 + 723 Since 2!, 2122 are in the affine hull of N, they can be written as an affine combination of seen points in N follows 21 ΣkEN at 0x x² EKEN by 8c, As result,
25

---

Compositional Risk Minimization
(a) Step 1
(b) Step 2
(c) Step
Figure 4. Illustration of steps of the deterministic sampling procodure for a 4 x grid. (a) shows the base set, (b) add a group in blue and the affine hull extends to include all the blue cells, (c) add group in yellow and the affine hull extends include all yellow cells.
we obtain
24 az + +322 + aa 2 +
(
ty2
(39)
(xd + 3bb) )00 722
KEN
a 3. Since a + 7 = 1, is an affine combination of points in .N J {g). Thus we have shown the claim for the point (2, n + 1). We can repeat this claim for point (3, + 1) and so on until we reach (m, n + 1) beyond which there would be no points in S that are expressed as affine combination of N. We make this argument formal through induction. We have already shown the base case above. Suppose all the points (nn + 1) in j < i < are in the affine hull of N J {g)- Consider the point 2² (i + 1, + 1). Construct the triplet 22 = ((i, m), n 1), (i)1,n)). Again from Lemma 2, it follows that 24 = az 1 +y23, We substitute 21 and 2 with their corresponding affine combinations. 2 = a EkENU()) 0+2" + ExENULO) 2 +y ΣKENU(s) c22",, Since ΣKENUKS) eaş + Bbk + 7C follows that 24 an affine combination of =1, =1, and 23, This completes the proof.
Observe that Ex (oas + (ba)
We now describe simple deterministic procedure that helps us understand how many groups need to see before we are guaranteed that the affine hull of seen points span the whole grid G = [l, x
(1,...,d)
We start with base set of three points B (1L, (2, From Lemma the affine hull contains (2,
For each E (2.... d 1} add the points (1, 1), 1, to the set B. From Theorem 4, it follows that affine hull of BU BU i 1)} (i + 1)] contains (i (i+11 x i- subgrid (1, , i 1} {1,.... "....,,+111 13 (here we apply Theorem 4 in two steps once for the addition of (1, + 1) and then for the addition of (i + 1,
At the end of the above procedure B contains 2d 1 points and its affine hull contains the grid G. We illustrate this procedure in Figure 4.
now discuss a randomized procedure that also allows us to span the entire grid with (d groups. The idea of the procedure is to start with base set of groups and construct their affine hull S. Then we wait to sample group g that outside this affine hull. If this sampled group shares the coordinate with affine hull of B denoted as S, then we expand the subgrid by one along with y coordinate. Similarly, also wait for point that shares y coordinate and then expand the subgrid by one along the coordinate.
We use S₂ to denote the distinct set of z-coordinates that appear in S and same goes for Sy, We write = (gasgy)) The procedure goes as follows.
Set 0, B and Flag T.
26

---

Compositional Risk Minimization
Sample a group y from G uniform at random. Update B B = Su
While + G, sample group g from G uniform at random.
If Flag =- 2. 9N Sax & S, then update B = B u (g), S = SU (S,, x {gy)) and Flag = y- If Flag =- y. By Sy g $ S, then update B BU (g). = Su (lgx) Syy) and Flag = 2.
In the above procedure, in every step in the while loop a group 9 is sampled. Whenever the Flag flips from x to y, then following Theorem 4, the updated set S belongs to the affine hull of B. We can say the same when Flag flips from to r. In the next theorem, we will show that the while loop terminates after Sed log( d) steps with a high probability and the affine hull of B contains the entire grid G. We follow this strategy. We count the time it takes for Flag to flip from x to y (from y to I)) as it grows the size of S from a x l subgrid
x (I + 1) (k: x (I: 1) subgrid to (k + 1) x (k + 1)) subgrid. Theorem 5. Suppose we sample the groups based on the randomized procedure described above. If the number of sampled groups is greater than Scd log(d), then G S (B) with probability greater than equal to 1 1
Proof. We take the first group g that is sampled. Without loss of generality, say this group is (1,
Suppose the Flag is set to x. Define an event A:: newly sampled g (g shares z-coordinate with a point in S (size k x k). g E S. Under these conditions Flag flips from x to To compute the probability of this event let us count the number of scenarios in which this happens. If gN takes one of the k values in S. and gy takes one of the remaining (d - k). then the event happens. As result, the probability of this event is P(At) (b)(d-)
d²
Suppose the Flag is set to Define an event A newly sampled g = (gxs 9y) shares y-coordinate with a point in S (size x (k+ + 1)) and g & S. Under these conditions Flag flips from y to x. The probability of this event is.( ( R)
d²
Define T as the number of groups that need to be sampled before Al occurs. Define The as the number of groups that need to be sampled before A occurs. Observe that after Tf + The number of sampled groups the size of the current subgrid S, which is in the afline hull of B, grows to (k + 1) x (k + 1).
(Tf : T.)) Tsum is the total number of groups sampled before the affine span of the observed groups Define Tsum Ei=
B contains the grid G.
We compute
E ET sum
sum!
Σ [BT ( 1
k=1
d
d-1 /k((d (b))d-)))- =2
(d-1)/2
E
Σ
Σ
2//kdd - -k))
(40)
A=1
k=
A=1
(d
1)/2
(d-1)/2
2
Σ
d²
/k((d -k)) = 2d
Σ
[
+
=
- (d))((d))))2)
k d - k
k=1
k=1
Similarly, we obtain a similar bound for _x=1 B
d
d-1
(d-1)/2
Σ
EETT Σ 1=1
& /((( + ()(d k)) =2
Σ
de @/((k + 1)(d - kk))
k=1
k=1
(41)
(d- 1)/2
(d- (d-1)/2 1)/2
2
Σ
d /((k + () k)) s 2d
Σ
4d 1)/2)
k
k=1
Overall E ~ ldd(/2)) From Markov inequality, it immediately follows that P(T S 8od log(//2) 2 2 ! In
the above approximations, we use Σi=1 x logg + 7, where y is Euler's 's constant. We drop ? as its a constant, which can always be absorbed by adapting the constant c.
27

---

Compositional Risk Minimization
Theorem 6. Consider the setting where p(.2z) follows AED Yz Z*, Z comprises of attribute vectors drawn uniformly at random from 2X, and the test distribution satisfies compositional shift characterization. If 2 Scd log d, where d is sufficiently large, f(z]z) p(z|z), Vz E Z vz E R", ((2) q(2), Vz E Z then the output of CRM (equation 10) matches the test distribution, i.e., ((2[x) = (<|2x), 2/2 E giest Vx E IR" with probability greater than 6*
Proof. Suppose the support of training distribution p(=) contains groups. know that these s groups are drawn uniformly random. From Theorem 5, is clear that if grows as Od d), then with high probability the entire grid of dl combinations is contained in the affine span of these observed groups. This can be equivalently stated 3x DAff( (2train) with probability greater than equal to 1 L. If z C DAff(Ztrain then from the assumption of compositional shifts, it follows that test C DAff (trrin)
). We can now use Theorem 2 and arrive at our result. This completes the proof.
D.4. Proof for Theorem 3: Growth of Discrete Afline Hull for Attribute Case
Let {I,...,,d, ,d)" and let us consider the groups = [= Fm]] E G. G contains d points. We are given set of groups S C G. (Sf is the discrete affine hull of S, where recall (Sf = Aff(S) G is the set of points in the affine hull of S that lie on the grid G.
Theorem 7. If the total number of groups sampled is greater than 2ed(m + 1 + In(d)), then DAf G with probability at least
Proof. Let us write == 2ed(m + +1 +1 In(d)) is easy to see that Aff is an affine vector space of dimension m(d - 1). Consider an infinite sequence (zt) of i.i.d. uniformly sampled groups in G, the increasing sequence sets St {2 ..'......'), 2'), and the corresponding increasing sequence of affine spaces Aff (S"). The theorem is equivalent to stating that have Af* Aff(G) with probability at least -
For every 2 1, if the "newly sampled" point z' does not belong to Aff(S--1), then necessarily dim Aff(S') dim (ff( ) + 1. For (S) to be equal to Aff(G), there needs to be m(d 1) dim Aff such increases in dimensionality when going from 5 = 0 to S (and there cannot be more increases than that). Note that as long as (Aff(S*-1) + G, the probability that 2 does not belong to IS is at least 1/|G], hence with probability 1 we have (Sf = G for large enough.
For 1, ..... (d 1), define the random variable
T = Ki : dim Aff(S*) = =i]l
(as noted above, Tj is well-defined and finite with probability 1) and the random index
t; = min dim Aff(S*) = =).
The random variable T; counts the number points z' sampled before a point not yet Aff(S ( is sampled (thus leading to an increase in dimension). Define also
m(d-1)-1
Tsum := 1
Σ
T.
1=1
Note that Tsum min SS) Af) Hence have show that with high probability, the random variable
Tsum is smaller than s.
The probability of a newly sampled point not belonging to DAff(S*) (and thus leading to an increase in dimension) is equal to (IG) (AAT(())))))(GG) hence T, is geometric variable of mean |GI/(IG| - (S = d /dd - (DAff((()) )))
and
m(d-1)-1
m(d-1)-1
d
E [Tamm
= +
Σ
ETI = 1+
Σ
(42)
d - DA
1=1
The following lemma, whose proof given further below, is the key ingredient to bound this sum:
28

---

Compositional Risk Minimization
Lemma 3. Let A G be such that DAf( A. 1f Al following inequality holds:
2 f.e. if contains more than half the points of G, then the
2d d
|A|
dim Aff(G) dim (ff
d
Intuitively, the lemma states that if the set A contains most points in G' (whose cardinality is d''),, then the dimension of the affine space it spans is almost that of Aff(G).
Let i" = minfi (DAff(S"))/4a > 1/2]. Then Lemma 3 applies to DAff(S') for all 2 4.,, and we see in particular that
dim Aff(G) 1- = dim Aff(G) dim Afr(5 S 2d dd
(Sf )
d,
dmd
hence: 1" must be greater than dim Aff(G) d = m(d 1) d. Thus we can split the sum as follows:
mid-1
E T [um]
1
E[T]]
m(d-1)
dm
d
< 1
Σ
dm
S
dm
DDfff
11
2d
1++ Σ +
Σ
dim Aff(G) dim Aff(S4)
1=1
m(d-1)-1
2d
S $1+ 1) +
Σ
m(d - 1) dim Aff(
i= 1)-d
2d Σ j=1
5 2md +
2d(m + 1 In(d))
j
where we apply both Lemma 3 and the definition of i" get the third line.
From Markov's inequality, we then see that P(Tsum I clB 2 1 1-1 P(Tsum S = 2c((m+1+ln(d)) > 1-
Proof of Lemma 3. We prove the lemma by induction on m N. If m 1, then dim Aff |A| and the statement
reduces the almost trivial incquality
dim Aff(G) dim Aff = d |A| 2d |A
Let us now consider m E N and assume that the lemma has been proved for m - 1. Let := (d - -IAl)/d" be the ratio of missing points, and let us define the subsets
G = {2EG am = k}
for & [d]. Consider the set of indices
1: === E AnG* #0).
As
U KEY
the average number of points A within each set Gh (for le E J) must be |A]/[T], and particular there exists k/ E I such that LAnG*1 2 [A/[[]]. Up to reordering the elements (1...., .d, we can assume without loss of generality that k' 1. Consider now the sets G and A == AnG ! There a trivial isomorphism between the one hot encoding representations
29

---

Compositional Risk Minimization
of G and that of the set [d]"-1, and the set A does verify Aff(A') nGl = AA. Moreover, 7 by assumption. Hence we can apply our recursive hypothesis to the sets A and Gt to conclude that
14
14|
AA
17
1 2m
dim Aff( - dim (ff(A & 2d d |A_
ddh LA/d/I1
S 2d
(43)
dm-1
dm
As A' C.A. A. we can write {a1..... } AA,A' for some 1, and we see that
dim Aff(G') dim Af'(A) dim Aff(G' (a13) - dim Aff(A' U {a1))
2 dim Aff U (a1,a2)) dim Aff U (2), 2...
2 dim Af U fai......00 a1))
dim Aff(A" U {ai ;..))).).)))))))))))))))
dim Aff(G² U A dim (ff
where each successive inequality is true because the newly added point 0x+1 either belongs to both Aff U (ai,..., and Aff(A² U [aI,...,a,, to neither of them, or only to (ff U {ai...., az}). Hence we find that
dim Aff(G' U A) - dim (ff S dim Aff(G") dim (f")
(44)
Note also that UKE/ Gh Aff(G' U.A (in fact, the inclusion is an equality). Indeed, it is easy to see that for any E G have G Af=(G) with our vectorization of G, the G is composed of all vectors of the shape [ui,..., Um ei
where the u, Rd are one-hot encodings and ex [0, .......... I.... 0] is one-hot encoding of [d]. i.e. d- zeroes and
one in position i. Let write E G as [V1..... Um k. and consider := [V1,.... cm- edi E G Note that 2 2
[O..... eg - ef.. Then any z' = [ui..... Um 1. eg] E G can be written as [M1...., um ei] + 2- & E Afr(G" U {z]),
which shows that G* C Af U (=)). As A o 0 for any & E I (by definition), this means that G C Aff(GG UA)
for all & E I. Hence we get
(UG²) dim Aff(A) S 2d d | Kke
dim Aff
(45)
dm
by combining equation 43 and equation 44. Furthermore, the same argument as above shows that if 2 G for [d][L then G Aff(GG U {=*)), which means that
((Uce) U ((t"beta))) (eer
G == Aff Aff
This means that adding d |4| vectors to Aff (UREs G*) is enough for the resulting set to affinely generate G; this is equivalent to saying that
(UG² KKET
dim Aff (G) dim Aff
S sd-
(46)
By combining equation 45 and equation 46, find that
dim Aff (G) dim Aff (A) < 2d m
LAld/II|
+ d -
dm
m
|A|
A
2d
+ 2d
(1 d/ll)) + d- [Z).
dm
dm
Now we only need to show that 2d /// (1 d/l41) + |2| S 0 complete the recurrence and prove the lemma. Remember
[A| that we have assumed that 7 > 1/2. Note also that as A = (G* A) < [[|dm=1, we have |Aldd" < ////d. Then
All and = [l/d) to showing that for any a E [1/2, 1] and any the desired inequality is equivalent (by setting a == am c E [a, 1], we have
2a 1/c) + S 0,
30

---

Compositional Risk Minimization
22
2
3
5
Steps for completing the Cartesian product via small connected components of size three:
3
2
4
5
6
Cartesian Product for Cj
Cartesian Product for
Observed Combinations
Unobserved combirations to which we generalize
Figure 5. Illustration of the discrete affine hull. Each cell in the 6 grid represents attribute combination, where observed combinations are solid-colored. The elements in blue form one connected component, Ct, and the elements in yellow form another connected component, C2. Extrapolation is possible for unobserved combinations, represented by the crosshatched cells, as long the test distribution samples from the Cartesian products of the connected components. The steps for completing the Cartesian product visually shows the intuition behind the extrapolation process.
which is simple exercise (one sees that the expression is an increasing function of on [a, 1] by deriving with respect to and that it is equal to 0 when c = 1).
Theorem 8. Consider the selling where p(-1z) follows AED Vz E Z*, Zrrain comprises of altribule vectors drawn uniformly at random from 2x, and the test distribution q satisfis compositional shift characterization. If 2 2d(m + + In)d),, p(=|x)= = p(2[x), , x), Vz E Zrrain Vx R", ((2) = q(),, V2 E ges then the output of CRM (equation 10) matches the test distribution, i.e., ((2|z) ((22)), V2 gesst Vx R", with probability greater than
Proof. Firstly, we can use Theorem to conclude that DAff (ttran) = 2x with probability at least compositional shifts Zuest C (Atf(z We can now use Theorem 2 arrive at the result.
1/c. Owing to
D.5. Discrete Affine Hull: Closer Look
In the next result, we aim to give a characterization discrete affine hull that helps us give two-dimensional visualization of ( Ztrain). Before we even state the result, we illustrate discrete affine hull of 6 grid. Consider the 6 grid shown in Figure 5. The attribute combinations corresponding to the observed groups shown solid colored cells (blue and yellow). The light shaded elements (blue and yellow) denote the set of groups that belong the affine hull of the solid colored groups. We now build the characterization that helps explain this visualization.
We introduce graph on the attribute vectors observed. Each vertex corresponds to the attribute vector, i.e., [21, 22]. There an edge between two vertices if the Hamming distance between the attribute vectors is one. connected component is a subgraph which all vertices are connected, i.e., between every pair the subgraph there exists path. Let us start by making an observation about the connected components in this graph.
We consider partition of observed groups into K maximally connected components, (C1, Cx }. Define Cij as the set of values the jth component takes in the ith connected component. Observe that Cjj n Ci 0 for +f 1. Suppose this was not that case and City n Ci * 0. In such a case, there exists point in C. and another point in C₁2 that share the jth component. As result, the two points are connected by an edge and hence that would connect C, and Cj. This contradicts the fact that C, and Cj are maximally connected, i.c., we cannot add another vertex to the graph while maintaining that there a path between any two points in the component. In what follows, we will show that the afine hull of C Cj x Cz2.
31

---

Compositional Risk Minimization
which is the Cartesian product extension of Cj. Next, we give some definitions and make simple observation that allows us think of sets C₁ x C22 as subgrids, which are easier visualize.
Definition Contiguous connected component: For each coordinate j E {1, consider the smallest value and the largest value assumed by it in the connected component C and call it min; and maxj. We say that the connected component C' is contigous if each value in the set {ming, min, +1,-- maxx, -1, is assumed by some point in C for all j E (1,2).
Smallest subgrid containing a contigous connected component C:: The range of values assumed by 5th coordinate in C, where j E 21, are {miny.... max,) }. The subgrid {minj. - ... max1] } x {ming, -- -mxx) } is the smallest subgrid containing C. Observe that this subgrid is the smallest grid containing C' because if we drop any column or row, then some point taking that value in C' will not be the subgrid anymore.
The groups observed training time can be divided into K maximally connected components (C1. .C}.. We argue that without any loss of generality each of these components contiguous. Suppose some of the components in {Ci, ,Ck) boر
not contiguous. We relabel the first coordinate as (T(])) Σラ< |Cyal r, where dj is the r point in Cn. We can similarly relabel the second coordinate as well. Under the relabeled coordinates, each component is maximally connected and contiguous. Also, under this relabeling the Cartesian products Cý1 x C52 correspond to the smallest subgrid containing Cj. Let us go back to the setting of Figure 5. The sets of observed groups shown solid blue and solid yellow form two connected components C1 and C2 respectively. Their Cartesian product extensions are shown as well in the Figure 5. Since the connected components were contiguous the Cartesian product extensions correspond to smallest subgrids containing the respective connected component.
Theorem 9. Given the partition of training support as Ztrain (C1,,... ;Ck, we have:
The affine span of contiguous connected component C' is the smallest subgrid that contains that connected component C.
The affine span of the union over disjoint contiguous connected components is the union of the smallest subgrids that contain the respective connected components.
Proof. C denotes the connected component under consideration and the smallest subgrid containing it is S. Denote the affine span of C' as A. We first show that the subgrid S S A.
We start with target point = (f1,tz) f2) inside S. We want show that the one-hot concatention of this point t can be expressed as an affine combination of the points in C.
Firstly, t is already in C,, then the point is trivially in the affine span. If that is not the case, then let us proceed to more involved cases. Consider the shortest path joining a point of the form (t1, 82) E C, where $2 # t2, and a point of the form (81,12 C, where 81 41 If such points do not exist, then I cannot be in S, which is contradiction.
We assign weight of (+1) to the concatenation of one-hot encodings of the point (2), We then traverse the path until we encounter point where $2 changes, note that such a point has to occur because of existence of ($1, (s1,ty) on the path. We call this point = ($1, s2). The point before v on the path is w = (81, 82). We assign a weight of (-1) to w. We summarize the path seen so far below. We also write the weights assigned to the points
8 = (t1, 82) u (s 82)
(+1)
(47)
w (5 82)
(-1)
0 (3), 82)
After w. we have a weight of +1 assigned to f1, -1 assigned to sy (note that s cannot be t,, this follows from the fact that we are on shortest path between points of the form (11, 82) and (81, f2))). We call this state S,. After w, we wait for point
32

---

Compositional Risk Minimization
on the path where 81 changes or we reach the terminal state (81, (2). The latter can happen if 81 81. In the latter case, we assign weight (+1) the terminal state and thus the final weights are (+1) for 1 and l2 and zero for everything else. This leads to the desired affine combination. We call this state Ti, corresponding terminal state.
Now suppose we were in situation where we reach a point (si", The point before @ is = (51, 82). We assign a weight of (+1) to r. We summarize the path seen after encountering w below.
v (5),8) 82)
:
(48)
(5) k2)
(+1)
q = (sf (22)
After r, we have a weight of +1 assigned to & and weight of +1 assigned szz We call this state Sg. After r, we wait for point where š2 changes. It could be that 52 changes to t2. The state before it is say u (2),5) and last state ($1, tg). Assign a weight of -1 1 to and assign weight of +1 to e. Thus we achieve the target as affine combination of points on the path. We call this state T2, corresponding to the terminal state.
Now let us consider the other possibility that the terminal state has not been reached. We call such a point m = (8T
The point that occurs before this point is / = (5) 82). We assign a weight of (-1) to (. We summarize the path taken below.
9 (st,8)
(49)
(5 .2) (3 $5)
(-1)
After I, 1 is assigned a weight of +1 and 8 is assigned a weight of 1. We reach the state Si again. From this point on, the same steps repeat. We keep cycling between S and S₂ until we reach the terminal state from either S₁ or S₂ at which point we achieve the desired affine combination. The cycling of states only goes on for a finite number of steps as the entire path we are concerned with has finite length. We show the process in Figure 6. Thus S C A.
We now make observation about the set A. which is the affine hull of set C. Suppose the first coordinate takes values between (min),.... mmxx) ). The corresponding one-hot encodings of the first coordinate are written as (onehot( (onehot(min.)..... ................................ onehot affine combination of vectors in (onchot(min:), .... onchot ] can lead to onchot(c). We justify this claim as follows. Observe that no vector in (minh)) ... (oneoo(mmxx))) has a non-zero entry in the same coordinate where (c) is also non-zero. Hence, any affine combination of vectors always have zero weight in the entry where onehot(c) is non-zero. is now clear that the first component of affine hull of A is always between (min1,".. max1). }. Similarly, the second component of affine hull of A is always between {ming, max2). Therefore, A C S. As result, A = S. Another way to say this is that DAff(C)) = C₁ x Cy2-
Now consider a value which is not in {minj,... .... maxx]. }. We claim that no
(michot(minn).
onehot onchot(max,)) (max,)) }} will
We now move to the second part of the theorem. We have already shown that DAff(C,) = C₁₁ x Cy2. We now want to show that
K DAff UC,)- - ( (C1x
j=l
jm1
Observe that (Af C (AU and (Bf C DAff(AU B), which implies DAff(A) S DAF((AUB). Therefore, from the first part and this observation it follows that U (C/r x C12) S DAfF(U, Cz). We now show DAfF UK C3) C U (Cn x Cz2).
33

---

Compositional Risk Minimization
T
S
S₂
T2
Figure 6. Illustration of state transition in proof of Theorem 9.
Take the K maximally connected components {CI.-.. Cx} and let the set of respective smallest subgrids con- taining them be 1S Define point z' as the affine combination of points across these components as z' E No
1== aijzij. where Zij is the jth point in C₁, which contains N, points. We can also write 2' as N.
2' = E
ai Zij- Define of = E Ciyy Zij- Observe that of is in the affine combi- j== Onty) 1= Σكك azs
Σgك ai
nation of points in C2 and hence of is point in S.. Let 04 Σド=1 Oij. In this notation, we can see z' is affine combination of 2/'s denoted as Σ ā;2ī. In this representation, there is at most one point per S; in the affine combination. There are two cases to consider. In the first case, exactly one component ă; is non-zero and rest all components are zero. In the second case, at least two components ā;'s are non-zero. In this setting, we can only keep the non-zero â 's in the sum denoted as Σ; =. Suppose z/- (ep: 6g) (without loss of generality), where ep is one-hot vector that is one on the pd coordinate. Observe that no other point in the sum Σ; á will have non-zero contribution on the nth coordinate. As result, in the final vector the pth coordinate of the first attribute will take the value 0 < à; 1. This point is not valid point in the set of all possible one-hot concatenations 2 and hence it does not belong the affine hull DAff ( U 1=1 C₁). Thus we are left with the first case. Observe that in the first case, we will always generate point in one of the DAff(C;),
where j E {1,.... K}. Thus DAf C2) S , (Aff which implies DAfF U C2) S U C₁ X Cj2. This completes the proof.
D.6. No Extrapolation beyond Discrete Affine Hull: Proof for Theorem 10
In this section, we rely on the characterization of discrete affine hulls shown in the previous section in Theorem 9. Suppose we learn an additive energy model to estimate p(æ+|z) and estimate the density p(z[2) for all training groups using maximum likelihood. In this case, we know that p(z[2) = p(z|z) for all x E ( In the next thcorem, we show that such densities that satisfy p(x|2) = p(x|2) for all E DAff (Ztrain) may not match the true density outside the affine hull. In the next result, assume that V2 x p(-1z) is not uniform.
Theorem 10. Suppose we learn an additive energy model to estimate p(x:2) and estimate the density p p(x|z) for all training groups. There exist densities that maximize likelihood and exactly match the training distributions but do not extrapolate to distributions outside the affine hull of ztrain, i.e., ]z zx, where f(-1z) + p(-1z).
Proof. We first take graain and partition the groups into K maximally connected components denoted (Ci,..., ..., Cx.. From Theorem 9, we know that the affine hull of zrrain is the union of subgrids {Si,... ,Sk) where each subgrid Sy is the Cartesian product C₁ x Cj2-
Let us consider all points (ž,, 22) in some subgrid Sg- For each such (ã,, 32) E Sk, define Er ž1) = E (z. ă1) + (a)(), E2(x, 52) = E(x, 22) az (z). Note that regardless of choice of ag the density, p(x|2) 2(2) ((()) matches the true density 22)| for all groups in , S..
Select any group žfef (2) that is not in the union of subgrids. From the definition of z it follows that there are points of the form (21, in one of the subgrid Sj and points of the form (2/, 22) are in some subgrid S,. Let (x) Ei(,,s) ) (x(w.2) and (xx) (x (2.,2 Observe that Ê. (z, 21) Fa 22) = Ei 21) E2(2, 22)+ 2
34

---

Compositional Risk Minimization
az(x) - (x) 0. Thus this choice of (z(x)) - (x) ensures that ((ar) 21 #2) is uniform and hence cannot match the true p(r|21,22)) 22).
This completes the proof.
Based on the above proof, we now argue that there exist solutions to CRM that do not extrapolate outside the affine hull. Let us consider solutions to CRM denoted E, B, which satisfies the property that (a))), B(=)) = (a))), F(x)) B(=) = B(z)Vz E grrain, Following the proof above, we can choose É's in such way that the sum of energies a certain reference point outside the affine hull is zero and at all points inside the affine hull the sum of energies achieve a perfect match. For the group 2ref (21, #2) not in the affine hull of gtrain, we set Eg (ar, 21) E2 22) (a(2), F(())) O.
Suppose ((=|x)) - (=|ar), V2 E DAff( 2train) U(zret). We now compute the likelihood ratio at žref and point E Ztrain We obtain
((2ret xx) ((21z) ((2(z)
g(zret x)
((<ret(z) (2) log [ (z)
(50)
llgg
((2(<))
g(2(<))
(or(zre) B(x)) - (z))), E(x)) (o E(x)) (())), (0(=)) ((ret)
where (=) corresponds to collection of all terms that only depend on z. We already know that (o(x), B(x)) = E(x)) and (Zrett, B(x)) = 0. Substituting these into the above expression we obtain
(z(zre)) E(x)) 0(2) ((ret)
(51)
From the above condition, it follows that (x]) is uniform. This implies that p(z]2ret) is also uniform, which contradicts the condition that p(x|22e) is not uniform. Therefore, ((2[z) = ((=[x), Va E 2train Uf cannot be true.
D.7. Extrapolation of Discrete Additive Functions via Discrete Affine Hulls
Define a real-valued function f(21, 2m) Σー= (z)), where each 2j f1,..... .d., Vj {1.-.., Define f= (1), ff [d)] and f [f1, fm]. We can re-express the function f in terms of one-hot encoding notation. [(21,-.. =m) ff. Given data {(2), y²) }}ام where y' = -(2²). Denote S = {='} Given a new point ž, which is not in the training data, we seek to predict the label ÿ. We can exploit the additive structure of the function to predict y. Suppose ă (S). From this (3) = Σÿ= (3)())), where E, aj = We can express ğ in terms of the Σj=1
seen data as follows. Observe that ÿ = (a(2), f) (Ej= az)()"), = aj it follows that we can perfectly predict the labels for points in the discrete affine hull. Thus from Theorem 7, it follows that points in S are sampled uniformly at random, then O(md + dloog d) suffice to extrapolate to the entire grid of d"" points and achieves CPE. In the work of (Dong & Ma, 2022), the authors also studied such discrete functions. However, their analysis does not propose the crucial object discrete affine hull, which gives a sharp characterization of extrapolation, and also do not provide sharp bounds the number of samples needed for CPE. In (Dong & Ma, 2022), show that non-trivial extrapolation is achievable provided the bipartite graph induced the probability distribution over the seen data is connected.
(x()). = E== ==1
ayy². From this,
35

---

Compositional Risk Minimization
Attribute
Attribute
(a) ERM
(b) CRM (Uniform Prior)
Figure 7. Failure of ERM to generalize to unseen test group. We consider the task in Figure 2, where model distribution with four Gaussian components. The group (-1, 11) (pink dashed) has zero prior probability in the training distribution. Subfigure (a) shows the decision boundary obtained by ERM-trained hinary classifier for predicting the attrihute z1 (attribute on x-axis), attribute 22 (attribute on y-axis) respectively. ERM fails to Iearn the decision boundary that would generalize well on samples from the missing group -1)... Subfigure b) shows the classifiers obtained using CRM with uniform priors (and where we marginalize the predicted group probabilities to get g(>1 (z) and q( x)). By contrast to ERM, CRM extrapolates and yiclds decision boundaries that generalize well on samples from unseen group (-1, 11). Decision regions were obtained from finite-data simulations, leading to minor imperfections.
E. Additional Details on CRM's Adaptation to Test Distribution
E.1. Derivation of Bayes Optimal Classifier
Three group setting. (Figure 2 left, train prior) Suppose the data is drawn from three groups +11, (-1,+ (+1, which are sampled with equal probability.
The Bayes optimal classifier predicts (+1,+1) if
||- (1.1)1 > - -1,--(11-1111" | (1, 1)12 < | (1, (1)11 = -2(21 + x2) < -2 x2)
22 > 0
||- (1.)) >e ||-- (-1,11111)11 = | - (1, < - (-1, = - + x2) < -2(-2) + = x1 > 0
(52)
The Bayes optimal classifier predicts (+1, if
-1,- > e (1.1)1²
= || (1, 1) > | (1, -1)12 = -2(21 + x2) > -2 12) = 32 <0
1| -1 11,11)11 > (--(==
= | - (1, 1)1 < 1| (-1, 1)12 = 2((11 v2) < -2(-x) x2) = 16 > 822
(53)
From same calculation follows that the Bayes optimal classifier predicts (-1, +1) an < and x2 > T1.
Four group setting. (Figure 2 right, uniform prior) Suppose the data is drawn from four groups which are sampled with equal probability.
(-1,-1)), (1,,
The Bayes optimal classifier for predicting the groups be obtained using exactly same calculations above.
The Bayes optimal classifier predicts: (+1,+1) if x1 > 0 and > 0. (-1,1) if < 0 and x2 > 0. (-1, if < and x2 < 0, and (-1, 1) if æ < 0 and Z2 > 0.
E.2. Comparison with ERM
We also train ERM the same training datasets where we observe data sampled uniformly from the groups (+1, (-1, (+1, -1), but we don't observe data from the group (-1,-1) Figure 7a shows the decision boundary of ERM for predicting the attribute 2 (attribute on x-axis) and attribute 22 (attribute on y-axis). Note that ERM fails to generalize to the novel group at test time (bottom left quadrant), while CRM with uniform test prior (Figure 7b) can adapt to the test distribution and extrapolate to the missing group
36

---

Compositional Risk Minimization
F. Experiments Setup
E.. Dataset Details
Waterbirds (Wah et al., 2011). The task is to classify land birds (y = 0) from water birds (y = 1), where the spurious attributes are land background (a = 0) and water background (a 1). Hence, we have a total of 4 groups 2 (y. a) the
dataset.
CelebA (Liu et al., 2015). The task to classify blond hair (y 1) from non-blond hair (y 0), where the spurious attribute is gender, female (a 0) and male (a 1). Hence, we have a total of groups (y, a) in the dataset.
MetaShift (Liang & Zou, 2022). The task is classify cats (y 0) from dogs (y 1), where the spurious attribute is background, indoor (a 0) and outdoor (a 1). Hence, have total of groups z (y. a) in the dataset.
MultiNLI (Williams et al., 2017). The task is to classify the relationship between the premise and hypothesis in a text document into one of the classes: netural (y 0), contradiction = 1), and entailment (y 2). The spurious attribute are words like negation (binary attribute a), which are correlated with the contradiction class. Hence, we have total of 6
groups (y. a) in the dataset.
CivilComments (Borkan et al., 2019). The task is to classify whether a text document contains toxic language (y = 0) versus it doesn't contain toxic language (y 1), where the spurious attribute corresponds to 8 different demographic identities (Male, Female, LGBTQ, Christian, Muslim, Other Religions, Black, and White). Hence, we have total of 16 groups z = (y. a) in the dataset.
NICO+ (Zhang et al., 2023). This is a large scale (60 classes with spurious attributes) domain generalization benchmark, and we follow the procedure in Yang al. (2023b) where all the groups with less than 75 samples were dropped from training. This leaves us with 337 groups during training, however, the validation set still has samples from all the 360 groups. Hence, additionally discard these groups from the validation set as well design the compositional shift version.
Dataset
Total Classes
Total Groups
Train Size
Val Size
Test Size
Waterbirds
2
4
4795
1199
5794
CelebA
2
4
162770
19867
19962
MetaShift
2
4
2276
349
874
MultiNLI
3
6
206175
82462
123712
CivilComments
2
16
148304
24278
71854
NICO+
60
360
62657
8726
17483
Table 3. Statitics for the different benchmarks used in our experiments.
F.2. Metric Details
Given the test distributions z = (y) ~ q(z) and x ~ g(z[2), lets denote the corresponding class predictions as 1 = (x) as per the method M. Then average accuracy defined as follows:
Average Acc Bewgttt) [1y (1y
Mx)]]
Hence, this denotes the mean accuracy with groups drawn as per the test distribution g(2). However, if certain (majority) groups have a higher probability being sampled than others (minority groups) as per the distribution q(=[z), then the average accuracy metric more sensitive clasifiisisss in majority groups compared to the minority groups. Hence, method can achieve high average accuracy though its accuracy for the minority groups might be low.
Therefore, we use the worst group accuracy metric, defined as follows.
Worst Group Acc 2 [1yy M (x)]
Essentially we compute the accuracy for each group (y. a) ~ ((x) as E Ex=g(=]s) [ == M [x)]] and then report the worst performance over all the groups. This metrics has been widely used for evaluating methods for subpopulation shifts (Sagawa
37

---

Compositional Risk Minimization
al., 2019; Yang et al., 2023b).
Similarly, we define the group balanced accuracy (Tsirigotis et al., 2024) as follows, where we compute the average of all per-group accuracy [x- [ly == M (x)]].
Group Balanced Ace =
Σ
[x) [1ly = (x)]]
(2.4)
E.3. Method Details
For all the methods have pre-trained representation network backbone with linear classifier heads. We use ResNet- 50 (He et al., 2016) for the vision datasets (Waterbirds, CelebA, MetaShift, NICO++) and BERT (Devlin al., 2018) for the text datasets (MultiNLI, CivilComments). The parameters of both the representation network and linear classifier are updated with the same learning rate, and do not employ any special fine-tuning strategy for the representation network. For vision datasets we use the SGD optimizer (default values for momemtum 0.9), while for the text datasets we use the AdamW optimizer (Paszke et al., 2017) (default values for beta (0.9, 0.999) ).
Hyperparameter Selection. We rely on the group balanced accuracy on the validation set to determine the optimal hyperparameters. We specify the grids for each hyperparameter in Table 4, and train each method with randomly drawn hyprraarameters. The grid sizes for hyperparameter selection were designed following Pezeshki et al. (2023).
Dataset
Learning Rate
Weight Decay
Batch Size
Total Epochs
Waterbirds
10 -5)
10 (niiform 6,-3)
20iform(5.7)
5000
CelebA
10 Unitorrm -3) 10 (0iifoormm 6,-3) 2 Uniform(5.7)
10000
MetaShift
10 Uniformm -3)33) 10 Unifform( -6,33 2 Uniform(5.7)
5000
MulitNLI
10 (0uittorm -4) 10 Uniforrm( 6,-3)
2 Uniform 4,6)
10000
CivilComments
10 Uniforrmm 16) 10 Uaiiform (60fformm-6,3)))
Uniform 1,6)
10000
NICO+
10 Uniiform -5.-3) 10 Uniifoommm-66 6,-
2 2miform(557
10000
Table 4. Details about the grids for hyperparameter selection. The choices for grid sizes were taken from Pezeshki et al. (2023).
38

---

Compositional Risk Minimization
G. Additional Results
G.I. Results Averaged over all the Compositional Shift Scenarios
We provide complete results on benchmarking CRM for compositional shifts on the datasets Waterbirds (Wah et al., 201 1), CelebA (Liu al., 2015), MetaShift (Liang & Zou, 2022), MultiNLI (Williams al., 2017), CivilComments (Borkan et al., 2019), and NICO++ dataset (Zhang al., 2023). We compare CRM against 7 baselines; ERM, Group Distributionally Robust Optimization (GroupDRO) (Sagawa al., 2019), Logit Correction (LCC) (Liu al., 2022b), supervised logit adjustment (sLA) (Tsirigotis al., 2024), Invariant Risk Minimization (IRM) (Arjovsky et al., 2019), Risk Extrapolation (VREx) (Krueger et al., 2021), and Mixup (Zhang et al., 2017).
Table 5 provides full results for the same, with groups balanced accuracy, standard across 3 random seeds, and baselines IRM, vREx, and Mixup, that were excluded from Table in the main paper. We find that CRM is either competitive or outperforms all the baselines w.r.t. both group balanced accuracy and the worst group accuracy. Also, CRM outperforms the baselines IRM, VREx, and Mixup by wide margin w.r.t WGA. Hence, we had decided to drop these baselines from the Table 1 in the main paper to show succinet comparison with strong baselines.
Further, in scenarios where CRM is competitive with baselines on WGA (Waterbirds, MultiNLI, CivilComments) please note that here we aggregate performance over multiple compositional shift scenarios. When we analyze the worst case com- positional shift scenario, we find that CRM outperforms all the baselines. Please check the next subsection (Appendix G.2) for detailed results each compositional shift scenario.
G.2. Detailed Results for all the Compositional Shift Scenarios
We now present resutls for all the compositional shift scenarios associated with each benchmark; Waterbirds (Table 6), CelebA (Table 7), MetaShift (Table 8), MultiNLI (Table 9), and CivilComments (Table 10, Table 11). Here we do not aggregate over the multiple compositional shift scenarios of a benchmark, and provide a more detailed analysis with results for each scenario. For each method, further highlight the worst group accuracy amongst all the compositional shift scenarios. This helps us easily compare the performance of methods for the respective worst case compositional shift scenario, as opposed to the average over all scenarios in Table 1. An interesting finding is that CRM outperforms all the baselines in the respective worst case compositional shift scenarios by wide margin, especially w.r.t worst-group accuracy.
scenario for it, i.e, the scenario with the lowest worst
39

---

Compositional Risk Minimization
Worst Group Acc
Dataset
Method
Average Acc
Balanced Acc
Worst Group Acc
(No Groups Dropped)
ERM
77.9 (0.1)
75.3 (0.1)
43.0 (0.2)
62.3 (1.2)
G-DRO
77.9 (0.9)
78.8 (0.7)
42.3 (2.6)
87.3 (0.3)
LCC
(0.9)
86.9 (0.6)
75.5 (1.8)
88.7 (0.3)
sLA
89.3 (0.4)
87.5 (0.4)
77.3 (1.4)
89.7 (0.3)
Waterbirds
IRM
70.8)
(0.3)
28.7 (2.2)
72.3 (1.2)
VREx
81.0
80.0 (0.5)
45.6 (1.1)
84.3 (0.7)
Mixup
81.6 (0.1)
79.9 (0.1)
52.2 (0.4)
69.7 (0.9)
CRM
87.1 (0.7)
87.8 (0.1)
78.7 (1.0)
86.0 (0.6)
ERM
85.8 (0.3)
75.6 (0.1)
39.0 (0.3)
52.0 (1.0)
G-DRO
89.2 (0.5)
86.8 (0.1)
67.8 (0.8)
91.0 (0.6)
LC
91.1 (0.2)
83.5 (0.0)
57.4 (0.5)
90.0 (0.6)
SLA
90.9 (0.2)
83.6 (0.3)
$7.4
86.7 (1.9)
CelebA
IRM
80.4 (1.3)
76.7 (1.1)
40.1 (2.4)
67.7 (3.5)
VREx
86.2 (0.3)
82.8 (0.5)
49.2 (2.1)
89.0 (0.6)
Mixup
84.9 (0.2)
77.9 (0.2)
12.8 (0.9)
62.0 (1.0)
CRM
(0.2)
(0.0)
81.8 (0.5)
89.0 (0.6)
ERM
85.7 (0.4)
81.7 (0.3)
60.5 (0.5)
63.0 (0.0)
G-DRO
86.0
82.6 (0.2)
63.8 (1.1)
80.7 (1.3)
LC
88.5 (0.0)
85.0 (0.0)
68.2 (0.5)
80.0 (1.2)
SLA
(0.1
84.0 (0.0)
63.0 (0.5)
80.0 (1.2)
MctaShift
IRM
(3.7
80.3 (0.4)
55.8 (1.0)
69.3 (2.4)
VREx
84.9 (0.4)
81.7 (0.3)
59.9 (0.2)
75.3 (2.2)
Mixup
86.8 86.8 (0.0)
82.8 (0.1)
62.8 (0.7)
68.3 (2.7)
CRM
87.6 (0.3)
84.7 (0.2)
73.4 (0.4)
74.7 (1.5)
ERM
(2.1)
(1.1
7.5 (1.3)
68.0 (1.7)
G-DRO
70.4
(0.2)
34.3 (0.2)
57.0 (2.3)
LC
(0.1)
77.3 (0.2)
54.3
74.3 (1.2)
sLA
(0.3)
77.4 (0.2)
51.5
71.7 (0.3)
MultiNII
(0.1)
(0.4)
8.1
54.3 (2.4)
VREx
69.0
68.8 (0.2)
4.1 (0.3)
69.7 (0.3)
Mixup
70.2 (0.1)
69.7 (0.1)
14.6 (1.0)
63.7 (2.9)
CRM
(0.3)
76.1 (0.3)
58.7 (1.4)
74.7
ERM
80.4 (0.2)
78.4 (0.0)
55.9 (0.2)
61.0
G-DRO
80.1
78.9 (0.0)
61.6 (0.5)
64.7 (1.5)
LC
80.7 (0.1)
(0.0)
65.7 (0.5)
67.3
sLA
80.6
(9.0)
65.6
66.3 (0.9)
CivilComments
IRM
79.7 (0.2)
78.0 (0.0)
53.5 (0.5)
60.3 (1.5)
VREx
79.8 (0.1)
78.7 (0.1)
57.5 (0.4)
63.3 (1.5)
Mixup
(0.1)
78.2 (0.0)
55.4 (0.6)
61.3 (1.5)
CRM
83.7 (0.1)
78.4 (0.0)
67.9 (0.5)
70.0 (0.6)
ERM
(0.0)
(0.0)
35.3 (2.3)
35.3 (2.3)
G-DRO
84.0
83.7 (0.3)
36.7 (0.7
33.7 (1.2)
LC
85.0
85.0 (0.0)
35.3 (2.3)
35.3 (2.3)
sLA
85.0
85.0 (0.0)
33.0 (0.0)
35.3 (2.3)
NICO+
IRM
64.0
62.7 (0.3)
0.0
0.0 (0.0)
VREx
86.0 (0.0)
(0.0)
37.3 (4.3)
38.0 (5.0)
Mixup
85.0 (0.0)
84.7 (0.3)
33.0 (0.0)
33.0 (0.0)
CRM
84.7 (0.3)
84.7 (0.3)
40.3 (4.3)
39.0 (3.2)
Table Robustness under compositional shift. We compare the proposed Compositional Risk Minimization (CRM) method with 7 haselines. We report various metrics, averaged as group is dropped from training and validation sets. Last column is WGA under the dataset's standard subpopulation shift benchmark, i.c. with no group dropped. All methods have harder time generalize when groups are absent from training, but CRM appears consistently more robust (standard error based on random seeds).
40

---

Compositional Risk Minimization
Discarded Group (y, a)
Method
Average Acc
Balanced Acc
Worst Group Acc
ERM
74.0 (0.0)
82.3 (0.3)
67.0 (0.0)
G-DRO
77.3 (0.7)
83.0 (0.6)
59.7 (1.9)
LC
85.7 (0.3)
88.7 (0.3)
82.0 (0.6)
sLA
86.0 (0.0)
89.0 (0.0)
82.3 (0.3)
(0, 0)
IRM
72.7 (2.0)
81.3 (1.2)
58.7 (3:2)
VREx
80.0 (1.5)
85.7 (0.9)
69.3 (2.4)
Mixup
87.3 (0.3)
89.3 (0.3)
84.3 (0.3)
CRM
86.7 (0.9)
88.7 (0.3)
83.0 (1.5)
ERM
67.3 (0.3)
71.7 (0.3)
28.0 (1.2)
G-DRO
58.3 (3.2)
70.7 (2.0)
11.7 (1.6)
LC
82.7 (3.2)
86.0 (1.7)
72.0 (5.8)
SLA
86.3 (1.7)
88.0 (1.0)
78.7 (3.3)
1)
IRM
55.7 (2.2)
67.7 (0.9)
(3.7)
VREx
66.0 (1.0)
74.0 (0.6)
22.7 (1.9)
Mixup
65.7 (0.3)
73.0 (0.0)
19.7 (0.3)
CRM
86.0 (2.1)
86.7 (0.7)
73.0 (4.2)
ERM
84.0 (0.0)
78.0 (0.0)
38.3 (0.3)
G-DRO
90.0 (0.0)
86.0 (0.6)
67.0 (3.6)
LC
93.0(0.0)
89.0 (0.6)
79.0 (1.2)
sLA
93.0 (0.0)
89.0 (0.6)
79.3 (1.5)
(1,0)
IRM
87.7 (0.3)
81.3 (0.7)
48.0 (4.0)
VREx
89.7 (0.3)
82.7 (0.3)
50.3 (1.7)
Mixup
84.3 (0.3)
80.7 (0.3)
51.7 (2.0)
CRM
86.7 (0.3)
89.0 (0.0)
83.7 (0.3)
ERM
86.3 (0.3)
69.3 (0.3)
38.7 (0.7)
G-DRO
86.0 (0.6)
75.7 (2.2)
31.0 (9.2)
11,
LC
92.0 (0.0)
84.0 (0.6)
69.0 (1.5)
sLA
92.0 (0.0)
84.0 (0.6)
69.0 (1.5)
IRM
78.3 (0.3)
51.3 (0.9)
(0.3)
VREx
88.3 (0.7)
77.7 (1.8)
40.0 (5.5)
Mixup
89.0 (0.0)
76.7 (0.3)
53.0 (0.6)
CRM
89.0 (0.6)
86.7 (0.7)
75.0 (3.2)
Table 6. Results for the various compositional shift scenarios for the Waterbirds benchmark. For each metric, report the mean (standard error) over random seeds on test dataset. We highlight worst case compositional shift scenario for each method, i.e, scenario with the lowest worst group accuracy amongst all the compositional shift scenarios. CRM outperforms all the baselines in the respective worst case compositional shift scenarios.

---

Compositional Risk Minimization
Discarded Group (y, a)
Method
Average Acc
Balanced Acc
Worst Group Acc
ERM
68.7 (0.3)
74.0 (0.0)
37.7 (0.3)
G-DRO
85.0 (0.6)
88.0 (0.0)
75.0 (1.2)
(0,0)
LC
88.0 (0.0)
90.3 (0.3)
82.3 (0.3)
sLA
87.7 (0.3)
90.3 (0.3)
82.3 (0.7)
IRM
65.0 (2.5)
72.0 (0.6)
31.3 (4.7)
VREx
82.7 (0.3)
87.7 (0.3)
70.7 (0.9)
Mixup
64.7 (0.3)
72.0 (0.0)
29.3 (0.7)
CRM
91.7 (0.3)
89.3 (0.3)
81.0 (2.0)
ERM
91.3 (0.9)
91.0 (0.6)
86.7 (1.3)
G-DRO
85.0 (1.5)
88.7 (0.7)
72.7 (3.7)
(0,
LC
93.0 (0.6)
87.7 (0.9)
71.0
SLA
92.7 (0.3)
88.0 (0.0)
71.3 (0.9)
IRM
77.0 (2.5)
83.7 (1.3)
53.0 (1.5)
VREx
78.7 (0.9)
85.0 (0.6)
55.7 (1.9)
Mixup
92.0 (0.6)
91.3 (0.3)
87.3 (0.9)
CRM
88.3 (0.9)
91.0 (0.6)
85.0 (2.0)
ERM
87.0 (0.0)
59.3 (0.3)
4.0 (0.0)
G-DRO
91.7 (0.3)
86.3 (0.7)
71.7 (0.9)
(1,0)
LC
88.3 (0.3)
70.7 (0.7)
(2.1)
sLA
88.3 (0.3)
71.0
21.3 (1.9)
IRM
84.7 (1.5)
72.7 (3.9)
47.3 (10.3)
VREx
88.3 (0.7)
79.3 (1.8)
40.3 (7.9)
Mixup
88.0 (0.0)
67.7 (0.3)
17.3 (1.7)
CRM
93.0 (0.0)
85.7 (0.3)
73.3 (1.8)
ERM
96.0 (0.0)
78.0 (0.6)
27.7 (2.0)
G-DRO
95.0 (0.0)
84.3 (0.3)
51.7 (1.2)
(1,1)
LC
(0.0)
85.3 (0.3)
55.3 (1.9)
sLA
95.0 (0.0)
85.0 (0.6)
54.7 (2.3)
IRM
95.0 (0.0)
78.3 (0.7)
28.7 (3.0)
VREx
95.0 (0.0)
79.0 (0.0)
30.3 (0.9)
Mixup
95.0 (0.0)
80.7 (0.3)
37.0 (1.2)
CRM 91.3 (0.3)
91.0 (0.0)
88.0 (0.6)
Table 7. Results for the various compositional shift scenarios for the CelebA benchmark. For each metric, report the mean (standard error)
over random seeds on test dataset. We highlight worst case compositional shift scenario for each method, i.e, scenario with the lowest worst group accuracy amongst all the compositional shift scenarios. CRM outperforms all the baselines in the respective worst case compositional shift scenarios.
42

---

Compositional Risk Minimization
Discarded Group (y, a)
Method
Average Acc
Balanced Acc
Worst Group Acc
ERM
84.3 (0.3)
84.0 (0.6)
80.3 (0.9)
G-DRO
84.0 (0.6)
83.3 (0.7)
78.0 (0.6)
(0,0)
LC
89.0 (0.0)
85.7 (0.3)
74.3 (1.8)
sLA
90.0 (0.0)
85.0 (0.0)
67.3 (1.9)
IRM
$1.0
81.7 (0.3)
70.0 (1.5)
VREx
83.7 (0.9)
83.0 (0.6)
76.7 (1.9)
Mixup
85.0 (0.0)
84.0 (0.0)
77.0 (0.6)
CRM
87.3 (0.3)
84.3 (0.3)
73.3 (0.7)
ERM
85.0 (0.0)
79.0 (0.0)
49.0 (0.0)
G-DRO
86.0 (1.0)
81.7 (0.3)
55.3 (3:2)
(0,
LC
86.0(0.0)
84.0 (0.0)
63.7 (0.3)
sLA
86.0 (0.0)
84.0 (0.0)
64.0 (0.6)
IRM
83.7 (0.3)
78.7 (0.9)
44.7 (1.9)
VREx
84.0 (0.0)
81.0 (0.0)
48.3 (0.7)
Mixup
86.0 (0.0)
80.0 (0.0)
52.7 (0.3)
CRM
88.3 (0.3)
85.7 (0.3)
78.0 (1.0)
ERM
90.0 (0.0)
82.0 (0.0)
48.3 (0.3)
G-DRO
90.3 (0.3)
82.7 (0.9)
52.7 (2.3)
(1,0)
LC
90.0 (0.0)
84.3 (0.3)
62.0 (0.0)
sLA
88.7 (0.3)
81.0 (0.0)
46.7 (0.7)
IRM
90.0 (0.0)
81.7 (0.3)
49.3 (1.7)
VREx
90.0 (0.0)
82.0 (0.0)
48.3 (0.3)
Mixup
90.3 (0.3)
82.7 (0.3)
52.3 (1.7)
CRM
87.0 (1.2)
83.3 (0.7)
70.0 (1.0)
ERM
83.3 (1.2)
81.7 (0.9)
64.3 (1.2)
G-DRO
83.7 (0.9)
82.7 (0.9)
69.3 (2.0)
(1,1)
LC
89.0 (0.0)
86.0 (0.0)
72.7 (0.7)
SLA
89.0 (0.0)
86.0 (0.0)
74.0 (0.0)
IRM
80.0 (0.6)
79.0 (0.6)
59.0 (1.0)
VREx
82.0 (0.6)
(0.6)
66.3 (0.9)
Mixup
85.7 (0.3)
84.3 (0.3)
69.3 (0.9)
CRM
87.7 (0.3)
85.3 (0.3)
72.3 (1.7)
Table 8. Results for the various compositional shift scenarios for the MetaShift benchmark. For each metric, report the mean (standard error) over random seeds on test dataset. We highlight worst case compositional shift scenario for each method, i.e, scenario with the lowest worst group accuracy amongst all the compositional shift scenarios. CRM outperforms all the baselines in the respective worst case compositional shift scenarios.
43

---

Compositional Risk Minimization
Discarded Group (y, a)
Method
Average Acc
Balanced Acc
Worst Group Acc
ERM
62.7 (0.3)
66.7 (0.3)
0.7 (0.3)
G-DRO
63.3 (0.3)
68.0 (0.0)
1.7 (0.7)
(0
LC
68.0 (0.0)
72.0 (0.0)
20.0 (0.0)
sLA
67.7 (0.3)
72.0 (0.0)
19.7 (1.5)
IRM
61.3 (0.7)
64.3 (0.3)
4.3 (0.9)
VREx
63.3 (0.3)
68.7 (0.3)
7.0 (1.0)
Mixup
63.0 (0.0)
64.3 (0.3)
1.0 (0.0)
CRM
64.7 (0.9)
70.7 (0.9)
31.0 (5.6)
ERM
77.7 (0.3)
71.7 (0.3)
14.0 (1.0)
G-DRO
80.7 (0.7)
80.7 (0.7)
74.0 (1.0)
(0,
LC
81.0 (0.0)
81.0 (0.0)
75.3 (0.3)
sLA
81.3 (0.3)
80.7 (0.3)
69.0 (0.6)
IRM
74.0 (0.0)
68.7 (0.3)
10.0 (1.5)
VREx
76.0 (0.0)
69.7 (0.3)
(0.6)
Mixup
78.3 (0.3)
74.3 (0.3)
34.0 (2.0)
CRM
80.0 (0.6)
78.0 (1.2)
62.3 (8.2)
ERM
58.0 (0.0)
67.0 (0.0)
0.0 (0.0)
G-DRO
57.7 (0.3)
67.7 (0.3)
0.0 (0.0)
(1,0)
LC
70.7 (0.9)
74.3 (0.3)
47.3 (4.3)
sLA
73.3 (2.7)
76.3 (1.7)
58.3 (9.7)
IRM
49.0 (0.6)
55.3 (1.5)
0.7 (0.3)
VREx
55.3 (0.3)
65.7 (0.3)
1.0 (0.0)
Mixup
57.0 (0.0)
66.0 (0.0)
0.0 (0.0)
CRM
69.5 (0.5)
74.0 (0.0)
63.5 (0.5)
ERM
82.0 (0.2)
73.0 (0.2)
20.0 (1.2)
G-DRO
80.3 (0.3)
79.3 (0.3)
72.7 (0.9)
(1,1)
LC
81.7 (0.3)
81.3 (0.3)
74.3 (1.5)
sLA
82.0 (0.0)
81.0 (0.0)
75.3 (0.7)
IRM
76.3 (0.3)
67.7 (1.2)
23.0 (7.0)
VREx
79.7 (0.3)
69.0 (0.0)
3.0 (1.5)
Mixup
80.7 (0.3)
72.3 (0.3)
28.3 (3.8)
CRM
81.3 (0.3)
80.7 (0.3)
71.3 (1.8)
ERM
62.0 (0.0)
68.3 (0.3)
0.0 (0.0)
G-DRO
60.0 (0.0)
67.7 (0.3)
0.0 (0.0)
(2
LC
72.3 (0.3)
74.7 (0.3)
48.7 (0.7)
sLA
72.7 (0.7)
74.3 (0.3)
48.3 (0.9)
IRM
57.0 (0.6)
57.3 (0.3)
0.0 (0.0)
VREx
59.7 (0.3)
67.7 (0.3)
0.0 (0.0)
Mixup
61.0 (0.0)
66.3 (0.3)
0.0 (0.0)
CRM
68.7 (0.3)
72.7 (0.3)
50.0 (0.6)
ERM
81.3 (0.3)
74.3 (0.3)
17.3 (2.4)
G-DRO
80.7 (0.3)
79.0 (0.0)
57.3 (2.2)
12)
LC
82.0 (0.0)
80.7 (0.3)
60.0 (1.2)
sLA
81.7 (0.3)
80.3 (0.3)
59.3 (0.9)
IRM
76.3 (0.3)
69.0 (0.0)
10.7 (0.9)
VREx
80.0 (0.0)
72.3 (0.3)
9.7 (1.2)
Mixup
81.0 (0.0)
75.0 (0.6)
24.3 (3.0)
CRM
81.3 (0.3)
80.0 (0.6)
72.7 (0.9)
Table Results for the various compositional shift scenarios for the MultiNLI benchmark. For each metric, report the mean (standard error) over 3 random seeds the test datasct. We highlight the worst compositional shift scenario for each method, i.c, the scenario with the lowest worst group accuracy amongst all the compositional shift scenarios. CRM outperforms all the baselines in the respective worst case compositional shift scenarios.
44

---

Compositional Risk Minimization
Discondal Group (y; e)
Mahhed
Avenge Ace
Rolanced Ace
Wons Group Ace
ERM
(9.0
78.7 (0.3)
(1.3(1.5)
G-DRO
(1.3
79.0
64.7(3.0) (3.0)
(0,0)
(0.3)
(8.0)
(0.3)
IRM
17.7 (0.3)
(0.0)
00.3 (0.3)
VREX
(0.3)
(9.0)
66.0 (0.0)
Mixup
(0.3)
(8.3
58.7
84.0.00.0
(0.3)
(8.6)
(1.2(
(8.0
T5.7 (0.3)
64.3
(0. 1)
(9.3
79.0 (0.0)
64.3 (0.9)
<LA
79.7
79.0
(0.3
(7.3
(2.6)
VREx
77.3
(9.0
06.0
Mixap
(8.6
78.3 (0.3)
62.3 (2.4)
CRM
83.3
78.7 (0.3)
(1.6
ERM
(d.3)
17.7 (0.3)
(8.0(1.)
(0,2)
GRO
(9.0
(0.3)
(3.7
sLA
(9.3
79.0 (0.0)
(2.1
IRM
(..3
71.7 (0.3)
34.7
VREx
78.3
77.7
30.7
Mixap
T1.7 (0.3)
40.0 (0.6)
CRM
(3.3
78.7 (0.3)
(1.00
(0.3)
79.0 (0.0)
64.3 (2.0)
G-DRO
$1.3 (0.6)
(0.0)
67.3 (2.7)
(0.3)
LCC
79.0
60.0
TRM
(0.7.
(0.0)
(5.7
VREx
77.7 (0.3)
79.0 (0.0)
64.7 (0.9)
Mixup
78.7 (0.9)
78.7
62.7
(0.3)
(0.3)
(0.3)
ERM
(8.0
(0.3)
38.0
(0.9)
75.7 (0.3)
$2.0 (3.2)
(0. 4)
(9.0
79.0 (0.0)
60.7 (1.5)
sLA
(0.3)
(9.0
62.0
IKM
(1.3
78.3 (0.3)
Mixup
(7.7
77.7 (0.3)
(4.0(4.4
CRM
(0.7)
79.0 (9.0)
69.7 (1.9)
FRM
(0.0
(0.0)
(0.6
G-DRO
(0.0
79.0
67.3 (7.8(1.8)
(0.5)
<LA
(0.0
(0.3)
66.7 (0.3)
IRM
(8.3
(0.3)
59.7 (0.9)
VREx
(0.3(
(0.3)
59.0
(9.3
(3.4)
CRM
(4.0
78.7 (0.3)
(1.0
(0.3)
75.0 (0.0)
36.3 (1.2)
G-DRO
(8.3
78.3 (0.3)
46.3
(0,6)
LE
$0.7
79.0 (6.0)
58.7 (2.3)
(9.9
79.0
(3.1)
(0.0)
VREX
(8.7
78.0 (6.0)
37.3 (2.2)
Mixup
(0.3)
78.0 (0.0)
33.7 (U.3)
(3.3
(0.3)
(1.0)
ERM
79.0.(0.0)
(0.3)
(0.0
(0.
79.7 (0.3)
(9.0)
(2.3
sLA
(0.3)
(9.0
36.3
IRM
(7.06
(0.3)
33.0
(0.0)
Mixup
(7.7
77.7 (0.3)
40.7 (3.5)
CRM
33.3.00.3)
78.3 (0.3)
64.0
Table 10. Results for the various compositional shift scenarios for the CivilComments benchmark (Part 1). For each metric, report the (standard error) over random seeds the test dataset. We highlight the worst case compositional shift scenario for each method. i.e the scenario with the lowest worst group accuracy amongst all the compositional shift scenarios. CRM outperforms all the baselines in the respective worst compositional shift scenarios.
45

---

Compositional Risk Minimization
Discondal Group (y; e)
Mahhed
Avenge Ace
Bolanced Ace
Wons Group Ace
ERM
(0.3)
79.0 (0.0)
80.3(0.3)
G-DRO
42.3
79.0
(1.3)
(1,0)
(1.3
(9.0)
(1.2)
IRM
81.7 (0.3)
(0.3)
64.0(0.6)
VREX
(2.0
79.0
68.3 (1.2)
Mixup
(0.3
(8.0)
63.3 (1.2)
(4.0.
(0.0)
(1.7
60.3
(2.0
79.0 (0.0)
67-3 (0.9)
(1.1)
10.7(0.3)
79.0 (0.0)
69.3 (0.9)
<LA (0.7
(0.3)
79.0 (0.0)
(1.2(
VREx
82.3
79.0 (0.0)
67.3 (1-2)
Mixap
(0.7
78.0 (0.0)
58.3 (0.3)
CRM
84.0(0.0)
78.3 (9.3)
70.0 (0.6)
ERM
(0.3
78.7 (0.3)
61.3 (0.7)
DRO
80.7
(9.0)
63.7
(1,2)
sLA
(2.0
79.0 (0.0)
89.7 (1.8)
IRM
(1.3
78.0 (0.0)
36.3
VREx
(0.3
(0.0)
(2.6)
Mixup
T8.7 (0.3)
62.8
CRM
83.7 (0.3)
78.3 (0.3)
63.7 (3.2)
(0.9)
78.0 (0.0)
30.0 (1.6)
G-DRO
(0.6)
(0.0
(2.6
(1,3)
LCC
(2.0
79.0
70.0
$2.7
(1.6)
TRM
(2.0
(0.0)
58.7 (1.5)
VREx
(1.3
79.0 (9.0)
66.7 (0.3)
Mixup
(2.0
(0.0)
57.3
(0.0)
ERM
(0.3)
78.7 (0.3)
58.3 (1.6)
79.0 (8.0)
68.0 (0.6)
(1.
$2.0
79.3 (0.3)
7U.7 (0.3)
sLA (2.0
(0.3)
70.0 (0.6)
IKM
(1.0
54.3 (0.9)
(0.0)
Mixup
82.0
78.0 (0.0)
57.3 (1.2)
CRM
(3.3)
78.3 (0.3)
GU.D (1.6)
FRM
82.0
78.7 (0.3)
63.7 (0.3)
G-DRO
(1.7
(9.0)
64.7 (1.3)
(1,5)
81.3(0.3)
<LA
(2.0
(9.0)
(1.3(
IRM
(0.3
(8.0
62.0
VREx
(1.3(0.3)
(9.0
(2.0
(2.0
CRM
(0.3)
78.3 (0.3)
70.0 (1.00
79.0 (0.0)
65.3 (2.4)
G-DRO
81.0
79.3 (0.3)
66.0
(1,6)
LC
(0.7
79.0 (6.0)
60.7 (2.3)
(0.3)
63.3
VREX
(2.D
79.0 (0.0)
66.3 (1.9)
Mixup
(0.)
79.0 (0.0)
64.7 (1.7)
CRM
(4.0
(0.3)
ERM
82.0(112)
(0.3)
63.3 (1.6)
61.0
(1.7)
(0.7
(9.0)
(1.0
sLA
(0.3)
(9.0
(7.0
IRM
(1.3
(0.3)
61.7 (0.7)
Mixup
(0.3
(0.0)
06.3 (1.2)
CRM
84.3 (0.3)
77.0 (0.0)
63.0 (1.0)
Table 11. Results for the various compositional shift scenarios for the CivilComments benchmark (Part 2). For each metric, report the (standard error) over random seeds the test dataset. We highlight the worst case compositional shift scenario for each method. i.e the scenario with the lowest worst group accuracy amongst all the compositional shift scenarios. CRM outperforms all the baselines in the respective worst compositional shift scenarios.
46

---

Compositional Risk Minimization
G.3. CelebA Multiple Suprious Attributes
We augment the CelebA dataset (Liu al., 2015) with three more binary spurious attribute (a2, 03, a4), described follows:
02: Determines whether the person is wearing eyeglasses not
Ag: Determines whether the person is wearing hat not
04 Determines whether the person is wearing earring not
Hence, we have a total of 2 = 32 groups with three binary attributes (y, a1, Q2, 43, a1); with denoting blond hair and ai denoting the gender, same as in our prior experiments with CelebA. Since CRM models each attribute with a different energy component, we incorporate additional energy layer for new attributes as compared to our prior experiments with two attributes. However, all the baselines would treat the two spurious attributes (a) U22 ag 04) as a single "meta" spurious attribute a' that takes 16 possible values, and aim to predict y. Table 12 presents the results for the multi-attribute CelebA dataset, where we generate multiple benchmarks with compositional shift by dropping one of the 32 groups from the training & validation dataset (similar to the setup for prior experiments). We find that CRM outperforms all the baselines w.r.t worst group accuracy and balanced accuracy, hence, remains superior for the case of multiple attributes well.
Worst Group Acc
Method
Average Acc
Balanced Acc
Worst Group Acc
(No Groups Droppd)
ERM
94.8 (0.1)
83.5 (0.1)
0.9 (0.5)
0.0 (0.0)
G-DRO
93.2 (0.0)
89.1 (0.0)
11.2 (1.5)
0.0 (0.0)
LC
93.1 (0.1)
91.7 (0.3)
34.3 (5.9)
52.0 (26.1)
sLA
93.4 (0.1)
91.7 (0.0)
33.0 (1.3)
26.3 (26.3)
IRM
90.2 (0.6)
85.6 (0.2)
26.6 (0.4)
40.0 (2.0)
VREx
90.0 (0.3)
89.1 (0.1)
10.6 (1.3)
25.0 (25.0)
Mixup
94.9 (0.0)
87.0 (0.0)
2.7 (0.5)
14.7 (14.7)
CRM
91.0 (0.1)
90.1 (0.1)
63.2 (2.6)
73.7 (2.9)
Table 12. CelebA with Multiple Spurious Attributes. We compare CRM to baselines on CelebA datasct with attributes. Similar to the prior setup (Table 1), report the Average Accuracy, Group Balanced Accuracy, and Worst Group Accuracy (WGA), averaged as group is dropped from the training and validation sets. Last column is WGA under the standard subpopulation shift scenario where no groups were dropped. CRM is the best approach w.r. the worst group accuracy.
G.4. Results for Ablations with CRM
In the implementation of CRM in Algorithm C, we have the following two choices; 1) we use the extrapolated bias 8* (equation 11); 2) we set ((=) as the uniform distribution, i., q(= = (y. a)) = x, We now conduct ablation studies by varying these components as follows.
Bias B* + Emp Prior: We still use the extrapolated bias B" but instead of uniform q(2). we use test dataset to obtain the counts of each group, denoted as the empirical prior. Note that this approach assumes the knowledge test distribution of groups, hence we expect this improve the average accuracy but not the necessarily the worst group accuracy.
Bias B Unf Prior: We still use the uniform prior for ((z) but instead of the extrapolated bias B*, we use the learned bias B (equation 9). This ablation helps us understand whether extrapolated bias 3* crucial for CRM to generalize to compositional shifts.
Bias B+ Emp Prior: Here change both aspects of CRM use the learned bias B and empirical prior from the test. dataset for q(z).
Table 13 presents the results of the ablation study. We find that extrapolated bias crucial for CRM as the worst group accuracy with learned bias is much worse! Further, using empirical prior instead of the uniform prior leads to improvement in average accuracy at the cost of worst group accuracy.
47

---

Compositional Risk Minimization
Dataset
Ablation
Average Ace
Balanced Ace
Worst Group Acc
CRM
87.1 (0.7)
87.8 (0.1)
78.7 (1.6)
Bias B* Emp Prior
91.6 (0.2)
87.4 (0.3)
75.2 (1.3)
Waterbirds
Bias B + Unf Prior
81.2 (0.6)
82.7 (0.2)
55.7 (1.0)
Bias B + Emp Prior
84.3 (0.6)
81.6 (0.3)
51.3 (1.0)
CRM
91.1 (0.2)
89.2 (0.3)
81.8 (1.2)
Bias 8* Emp Prior
94.3 (0.1)
75.8 (0.4)
34.1 (1.0)
CelebA
Bias B Unf Prior
83.6 (0.1)
84.7 (0.2)
58.9 (0.4)
Bias B + Emp Prior
90.9 (0.1)
77.2 (0.3)
35.4 (0.7)
CRM
87.6 (0.2)
84.7 (0.1)
73.4 (0.7)
Bias 8 + Emp Prior
89.2 (0.2)
84.0 (0.4)
65.1 (1.4)
MetaShift
Bias B + Unf Prior
87.2 (0.3)
82.9 (0.4)
58.7 (0.6)
Bias B + Emp Prior
88.1 (0.1)
82.1 (0.1)
56.1 (0.4)
CRM
74.3 (0.3)
76.1 (0.3)
58.7 (1.4)
Bias + Emp Prior
74.7(0..)
72.3 (0.4)
41.4 (1.5)
MultiNLI
Bias B + Unf Prior
72.5 (0.6)
74.0 (0.4)
30.4 (2.6)
Bias B + Emp Prior
73.2 (0.5)
70.8 (0.1)
22.2 (0.9)
CRM
83.7 (0.1)
78.4 (0.0)
67.9 (0.5)
Bias 8* Emp Prior
87.0 (0.1)
74.0 (0.2)
48.1 (0.6)
CivilComments
Bias B + Unf Prior
76.9 (0.3)
77.8 (0.1)
52.4 (0.7)
Bias B + Emp Prior
83.5 (0.2)
77.9 (0.1)
62.3 (0.8)
CRM
84.7 (0.3)
84.7 (0.3)
40.3 (4.3)
Bias B* Emp Prior
85.0 (0.0)
85.0 (0.0)
41.0 (4.9)
NICO+
Bias B Unf Prior
85.0 (0.0)
85.0 (0.0)
31.0 (1.0)
Bias B + Emp Prior
85.0 (0.0)
85.0 (0.0)
27.7 (3.9)
Table 13. Ablation study with CRM. We consider the average performance over the different compositional shift scenarios for each benchmark, and report the mean (standard error) over random seeds on the test dataset. CRM corresponds to the usual implementation with extrapolated bias 8" and uniform prior for d(=). CRM obtains better worst group accuracy than all the ablations, highlighting the importance of both extrapolated bias and uniform prior! Extrapolated bias is critical for generalization to compositional shifts as the performance with learned bias is much worse
G.5. Results for the Original Benchmarks
We present results the original benchmarks (Derain; Dval, Dtrain) Table 14, which corresponds to the standard subpop- ulation shift case for these benchmarks. For Waterbirds, CelebA, MetaShift, and MultiNLI, subpopulation shift implies all the groups change from train test, inducing spurious correaltion between class labels and attributes For the NICO+ dataset, we have a total of 360 groups in the test dataset but only 337 of them are present in the train dataset. But still this is not a compositional shift as the validation dataset contains all the 360 groups. We find that CRM is still competitive to the baselines for the standard subpopulation shift scenario each benchmark!
(y, a) are present in both the train and test dataset (2train gess 2* ), however, the groups sizes
48

---

Compositional Risk Minimization
Dataset
Method
Average Acc
Balanced Acc
Worst Group Acc
ERM
87.3 (0.3)
84.0 (0.0)
62.3 (1.2)
G-DRO
91.7 (0.3)
91.0 (0.0)
87.3 (0.3)
LC
92.0 (0.0)
91.0 (0.0)
88.7 (0.3)
sLA
92.3 (0.3)
91.0 (0.0)
89.7 (0.3)
Waterbirds
IRM
87.3 (0.3)
36.0(0.0)
72.3 (1.2)
VREx
92.0 (0.0)
(0.3)
84.3 (0.7)
Mixup
85.0 (0.0)
36.0
69.7 (0.9)
CRM
91.3 (0.9)
91.0
86.0 (0.6)
ERM
95.7 (0.3)
84.0 (0.0)
52.0 (1.0)
G-DRO
92.0 (0.6)
93.0 (0.0)
91.0 (0.6)
LC
92.0 (0.6)
92.0 (0.0)
90.0 (0.6)
sLA
92.3 (0.3)
91.7 (0.3)
86.7 (1.9)
CelebA
IRM
87.0 (2.5)
85.3 (1.2)
67.7 (3.5)
VREx
92.0 (0.0)
92.0 (0.0)
89.0 (0.6)
Mixup
95.0 (0.0)
86.7 (0.3)
62.0 (1.0)
CRM
93.0 (0.0)
92.0 (0.0)
89.0 (0.6)
ERM
90.0 (0.0)
84.0 (0.0)
63.0 (0.0)
G-DRO
90.3 (0.3)
88.3 (0.3)
80.7 (1.3)
LC
89.7 (0.3)
87.7 (0.3)
80.0 (1.2)
sLA
90.0 (0.6)
87.7 (0.3)
80.0 (1.2)
MetaShift
IRM
90.7 (0.3)
85.3 (0.9)
69.3 (2.4)
VREx
90.0 (0.0)
86.7 (0.3)
75.3 (2.2)
Mixup
90.7 (0.3)
85.3 (0.7)
68.3 (2.7)
CRM
88.3 (0.7)
85.7 (0.3)
74.7 (1.5)
ERM
81.7 (0.3)
80.7 (0.3)
68.0 (1.7)
G-DRO
80.7 (0.3)
78.0 (0.0)
57.0 (2.3)
LC
82.0 (0.0)
(0.0)
74.3 (1.2)
sLA
82.0 (0.0)
82.0 (0.0)
71.7 (0.3)
MultiNLI
IRM
75.7 (0.3)
74.3 (0.3)
54.3 (2.4)
VREx
79.0 (0.0)
79.0 (0.0)
69.7 (0.3)
Mixup
81.3 (0.3)
80.0 (0.0)
63.7 (2.9)
CRM
81.7 (0.3)
81.7 (0.3)
74.7 (1.3)
ERM
80.3 (0.3)
79.0 (0.0)
61.0 (2.5)
G-DRO
79.7 (0.3)
79.0 (0.0)
64.7 (1.5)
LC
80.7 (0.3)
(0.3)
67.3 (0.3)
sLA
80.3 (0.3)
79.0
66.3 (0.9)
CivilComments
IRM
80.3 (0.7)
79.0
60.3 (1.5)
VREx
80.3 (0.7)
(9.0)
63.3 (1.5)
Mixup
80.0 (0.6)
(0.0)
61.3 (1.5)
CRM
83.3 (0.3)
78.0 (0.0)
70.0 (0.6)
ERM
85.3 (0.3)
85.0 (0.0)
35.3 (2.3)
G-DRO
83.7 (0.3)
83.3 (0.3)
33.7 (1.2)
NICO+
LC
85.0 (0.0)
85.0 (0.0)
35.3 (2.3)
sLA
85.0 (0.0)
85.0 (0.0)
35.3 (2.3)
IRM
63.7 (0.3)
62.7 (0.3)
(0.0)
VREx
86.0 (0.0)
86.0 (0.0)
38.0 (5.0)
Mixup
85.0 (0.0)
84.7 (0.3)
33.0 (0.0)
CRM
85.0 (0.0)
84.7 (0.3)
39.0 (3.2)
Table 14. Results for the standard subpopulation shift case for each benchmark. Here we not transform the datasets for compositional
shifts, hence all the groups present in both the train and the test dataset (except the NICO++ benchmark). CRM is still competitive with the baselines for this scenario where no groups were discarded additionally.
49

---

Compositional Risk Minimization
G.6. CRM's Analysis with Varying Group Size
Test
Test Acc.
Test
1.00
0.95
0.95
0.9
0.90
0.90
0.8
0.85
0.85
0.7
0.80
0.80
100 90 80 70 60 50 40 20
100 90 80 70 50 40 30 20
100 90 80 70 50 40 30 20 10 Percentage ge of Total Groups at training
Percentage of lotal Groups at training
Percentage of Total Groups at Training
(a) 2 d= 10
(h) m 2 d 20
(c) m 2 d 30
Test Acc.
Acc
0.95
0.95
0.90
0.90
0.85
0.85
0.80
0.80
100
60 40 10 Percentage Groups Training
80
100 90 80
60 40 10
Percentage Groups Training
(d) m = 2, d 40
(c) m = 2, d = 50
Figure 8. Varying Group Size Analysis (m 2 attributes). We analyze the rate of growth of total groups required achieve cartesian- product extrapolation. For each scenario, we evaluate CRM's generalization capabilities as we discard more groups from the training dataset. X-axis denotes the percentage of total groups available for training, and y-axis denotes the test average accuracy (mean standard error over 3 random seeds) obtained by CRM. We find that observing at least 20% of total train groups is sufficient for good generalization.
Setup. We conduct experiments to understand the rate of growth of total groups required in order to achieve Cartesian- Product extrapolation, as we vary the total number of attributes (m) and the total number of categories (d) for each attribute. Given attributes (21, 22, 2m), we sample data sample data from the following (additive) energy function.
E(x,2) = > || m(2))l
where z. ((=;) E IR for aol i E f1, .m) Note that the energy function can be rewritten as follows:
1
E(x,2) 2 - ((c=m))) m(z)) + C(z1, 22)
with m=z)= = 1 Em ((zz) and Σ-1 = 2mI,. Hence, the resulting distribution is essentially multi-variate gaussian distribution p(x|2)= = Z(z) exp E(x, 2)) = N E))
To generate data from a particular configuration (d, m, n) we first sample d m orthogonal vectors to get mean vectors for the different realizations of each attribute, i.e, {u(2 k) i E [1.m] & k E [1, d]. Then we sample from the resulting normal distribution x N (u(2), Σ) to create dataset with uniform support over all the dmh groups We fix the data dimension as n = 100 and have the following two setups.
m 2 Attribute Case. We fix m and vary in the following range, [10, 20, 40, 50). This results in groups with sizes [100, 400, 900, 1600, 2500].
m > 2 Attribute Case. We fix d and vary m in the following range, [7,8, 9). This results in groups with sizes [128, 256, 512].
For both these setups, we analyze how the performance of CRM degrades as we discard more groups from the training dataset. Note that the test dataset contains samples from all the groups and there are no group imbalances. Hence, average accuracy in itself is a good indicator of generalization performance.
50

---

Compositional Risk Minimization
Test ACC.
Test Acc.
Test Acc.
0.78
0.80
0.8
0.7
0.75
0.76
0.6
0.70
0.74
100
40
20
100 60
40
20 10
100
40
20
10 5
Percantage of Total Groups
Porcentage Tatal Groups at Training
Percentage Total Groups Training
(a) m 7, 2
(b) m = .d =
(c) m =
Figure 9. Varying Group Size Analysis (m > 2 attributes). We analyze the rate of growth of total groups required to achieve cartesian- product extrapolation.
each scenario, we evaluate CRM's generalization capabilities we discard more groups from the training dataset. X-axis (log scale) denotes the percentage of total groups available for training, and y-axis denotes the test average accuracy (mean standard error over random seeds) obtained by CRM. We find that observing at least 10% of total train groups sufficient for good generalization.
Results. Figure 8 and Figure presents the results for the m 2 and m 2 attribute case respectively. For m attribute case, find that CRM trained with 20% of the total groups (0.2d?) still shows good generalization for predicting 21 (~ 90% test accuracy), and the drop in test accuracy as compared to the oracle case of groups dropped is within 10%. Similarly, for m > 2 attribute case, we find that CRM trained with 10% of the total groups (0.1 x 2 still shows good generalization for predicting 21, and the drop in test accuracy as compared to the oracle case is within 10%.
51


Reviews:
[{"review_id": "4VdxQSRLdS", "reviewer": "Reviewer_uM4Q", "summary": "This paper addresses the compositional shifts, a hard type of sub-population shifts, and proposes compositional risk minimization. The method is well-motivated and some theoretical analyses are provided. Results on the sub-population shift benchmark are shown to support the proposed method.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive and insightful feedback! We now address the concerns raised by the reviewer ahead.\n\n**1. The analysis and proposed algorithm are designed to handle multiple attributes, with the theoretical advantages being most relevant for this multi-attribute context. However, the experiments are limited to only 2~3 attributes. I suggest that the authors include empirical results with multiple attributes to better align with the theoretical analysis.**\n\nThanks for raising this issue, we would like to provide some clarifications regarding this. In our theoretical results (Theorem 2 & 3), the key finding is that if we observe $O(md + dlogd)$ groups at training time, then CRM would generalize to test distributions over all the $d^m$  groups. Note that the result considers the setting with multiple groups, which can arise from multiple attributes $m$ or from multiple values per attribute $d$. Hence, the theoretical advantages are not restricted to the multi-attribute scenario.\n\nIn our experiments, we already consider several scenarios that go beyond few attributes and few groups. NICO++ dataset has 360 groups, CelebA (multiple spurious attribute case, Table 12, Appendix G.3) consists of 5 attributes and a total of 32 groups. In addition to this, we also provided experiments on synthetic data with varying $d$ ($m=2$ with $d$ up to 50 leading to 2500 groups) and varying $m$ ($d=2$ and $m$ up to $9$ leading a total of 512 groups) in Figure 9, Appendix G.6. In all these experiments CRM performs well thus aligning the behavior of the method with the theoretical claims.\n\nFinally, we also want to point the guarantees in Theorem 2 are applicable to settings with small number of groups as well, i.e., in the setting of datasets like Waterbirds CRM offers a non-trivial Bayes optimality guarantee when only three out of the four groups are observed at training time.\n\n**2. I understand that the paper’s chosen disjoint setting adds a level of complexity. However, in real-world scenarios, it is often feasible to obtain a small number of samples for different attribute combinations (particularly with only two attributes, as in these experiments). The proposed method should also be evaluated in traditional settings where all attribute combinations have some representation. This would confirm that the method performs well without requiring group-dropping.**\n\nIn the paper, we had already carried out a comparison in the traditional setting where all attribute combinations have some representation. These were presented in the rightmost column in Table 1, \"WGA (no groups dropped)\", as well as in Table 14 in Appendix G.5 which contains additional metrics. CRM remains competitive with the baselines in this scenario. This confirms that the method performs well also without requiring group-dropping. We also want to point that as the number of groups grow (due to increase in $d$ or $m$), it is natural to expect that we will be disjoint setting where no samples are available from any group. \n\nAlso, given the page limit, several results currently are in the supplementary material and we will move some results to the main body in the future revision. If you want us to move some specific result to the main body, please let us know.\n\nThanks again for your constructive comments! We are open to further discussion and would be happy to address any remaining concerns.", "claims_and_evidence": "- The compositional risk minimization method is reasonable and well-motivated.\n- The formulation of the compositional shift setting provides a foundation for further research.\n- Experimental results are good to support the method.", "methods_and_evaluation_criteria": "The analysis and proposed algorithm are designed to handle multiple attributes, with the theoretical advantages being most relevant for this multi-attribute context. However, the experiments are limited to only 2~3 attributes. I suggest that the authors include empirical results with multiple attributes to better align with the theoretical analysis.", "theoretical_claims": "Correct", "experimental_designs_or_analyses": "I understand that the paper’s chosen disjoint setting adds a level of complexity. However, in real-world scenarios, it is often feasible to obtain a small number of samples for different attribute combinations (particularly with only two attributes, as in these experiments). The proposed method should also be evaluated in traditional settings where all attribute combinations have some representation. This would confirm that the method performs well without requiring group-dropping.", "relation_to_broader_scientific_literature": "This paper proposes a more efficient method for domain generalization.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "I would suggest the authors moving additional results into the main body.", "questions_for_authors": "N/A", "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": "", "extra_scores": {}}, {"review_id": "AtZPWvSaeU", "reviewer": "Reviewer_ypJd", "summary": "This paper introduces a method for addressing compositional shifts in discriminative tasks. The authors propose a theoretical framework built on additive energy distributions, where each energy term represents an attribute. They introduce the discrete affine hull concept to characterize extrapolation capabilities. Their two-step algorithm first trains an additive energy classifier to predict attributes jointly, then adjusts this classifier for compositional shifts. Theoretical guarantees show that the proposed method can extrapolate to test distributions within the discrete affine hull of training distributions. Experiments on several benchmarks demonstrate the effectiveness of the proposed method.\n\nPros:\n+ The proposed method is well-motivated and reasonable, and builds on additive energy distributions that are studied in generative compositionality.\n+ The proposed algorithm is practical and easy to implement. And the authors provide detailed implementation in the appendix.\n+ The extensive empirical evaluation demonstrates consistent improvements across diverse benchmarks.\n\nCons:\n- The additive energy assumption may be too limited for many real-world situations where different factors interact in complex ways rather than simply adding together. This could reduce how useful the approach is in practice.\n- The additive energy distributions were previously studied in generative compositionality, while the authors extend this framework to discrimination tasks. This is an incremental contribution and the novelty appears limited in scope.\n- The paper assumes access to attribute labels during training, which might not always be available in practice.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive and insightful feedback! We are glad they appreciate the technical soundness of our work, on both the theoretical and empirical front. We now address the concerns raised by them.\n\n> Additive Energy Distribution (AED) Limitations\n\nWe emphasize that the benchmarks used to evaluate CRM are both realistic and widely adopted in the subpopulation shift literature. Since CRM consistently outperforms baselines, this suggests that AED assumption is not overly restrictive and can model realistic datasets effectively. \n\nWe now clarify how AED models complex interactions between attributes. Note that AED does not imply additive interactions in data space ($x= \\sum_{i} g(z_i)$) as in additive decoders (Lachapelle et al. 2024). Instead, it models the AND operation between attributes, as illustrated below via examples (check Appendix B for details).\n\ni) Consider images that contain a distinct object varying in shape, size, and color. At any pixel, it is unlikely that shape, color, and size attributes interact additively, rather their interactions are complex which can't be captured via additive decoders. However, under AED, interactions are modeled via an energy component per attribute: one detecting shape, AND another detecting color, AND a third detecting size. Together, these energy terms define the distribution of images conditioned on attributes.\n    \nii) An example from a different data modality is the the CivilComments benchmark, where the attributes toxic language (class label) and demographic identity (spurious attribute) interact non-trivially in text space. However, under AED, we can model their interactions via an energy component that checks whether the language is toxic, AND another energy component checks the demographic identity.\n    \n> Novelty of the work\n\nWe explain the key features that set this work apart.  \n\na) *Discrete Affine Hull, a novel mathematical object:* Existing AED works in generative compositionality lack theoretical guarantees for generalization beyond training data. To address this, we introduce a novel mathematical object, the discrete affine hull, which precisely characterizes extrapolation to new distributions for both discriminative and generative tasks. For instance, our theoretical guarantee states that it is possible to generalize from $O(m*d)$ groups to $d^m$ groups in both discriminative and generative tasks. \n\nb) *Discriminative training without estimating partition function:*  Compositionality in discriminative tasks is a major problem, and our work makes key advances. One way to learn a classifier is via generative classification (lines 171-190, right column), where we first train densities $\\hat{p}(x|z)$ on observed groups, estimate new densities via affine combinations, and then use Bayes rule to derive $\\hat{p}(z|x)$. While this guarantees generalization, it is impractical due to the intractable gradient estimation of log partition functions (line 193-207, right column). While CRM circumvents these issues and retains the same guarantees.\n\n> Prior work regarding convex hulls\n\nWe thank the reviewer for pointing us to these references, and we are happy to cite and contrast with them. However, note that the densities from a new group $q(x|z')$ in the affine hull cannot be expressed as a convex or even an affine combination of the densities. For details, see e.q. (22) in Appendix D.2, summarized below. \n\n$\\log\\big(q(x|z')\\big)= \\sum_{z\\in \\mathcal{Z}^{\\mathsf{train}}}\\alpha_z \\log p(x|z) - R\\big(\\{\\alpha_z\\}_{z\\in \\mathcal{Z}^{\\mathsf{train}}}\\big)$\n\nThus, in our setting, only energy terms are expressed as an affine combination, and our guarantees apply to distributions that are outside the convex hull of the training distributions. \n\n> Missing attribute labels scenario\n\nWe believe this is an exciting future work. Existing works such as XRM (Pezeshki et al. 2024) show how one can discover the environments and then use existing domain generalization methods that require environment labels. We believe it would be exciting to extend these works to infer spurious attributes directly in combination with our approach. \n\n> Computational complexity of CRM\n\nFor CRM training stage 1 (e.q. 7), the cost for each step is similar to that of training an ERM-based classifier for predicting the group $z$, which is proportional to  $| \\mathcal{Z}_{\\mathsf{train}} | \\times m \\times d$.\n\n In the training stage 2, we compute the extrapolated bias $B^{*}(z)$ (e.q. 11), and the cost is proportional to $| \\mathcal{Z}_{\\mathsf{test}} | \\times \\text{number of training samples} \\times m \\times d$. \n \nObserve that in the worst case scenario the number of classes for $z$ at test time is $d^m$, making inference cost $O(d^{m+1}m)$. Any method that predicts $z$ would have to compute a probability vector of size $d^m$ and thus spend at least $O(d^m)$ per inference.\n\nThanks again for your constructive comments, and please let us know if there are any remaining concerns.", "claims_and_evidence": "1. The authors provide rigorous theoretical analysis with detailed proofs showing that the proposed method can generalize to novel attribute combinations within the discrete affine hull.\n\n2. Empirical validation on several common benchmarks show the proposed method outperforms other baselines.", "methods_and_evaluation_criteria": "1. The proposed two-step algorithm is consistent with the theoretical framework. The proposed method is well-motivated and reasonable, and builds on additive energy distributions that are studied in generative compositionality.\n\n2. The authors use average accuracy, group-balanced accuracy, and worst-group accuracy to evaluate performance. The evaluation is comprehensive.", "theoretical_claims": "The proofs appear mathematically sound. The theoretical analysis provides a sharp characterization of extrapolation, demonstrating that generalization beyond the discrete affine hull is fundamentally impossible. This establishes clear boundaries on what can be achieved in this domain.", "experimental_designs_or_analyses": "The experiments systematically demonstrate the advantages of the proposed method, especially on worst-group accuracy. The ablation studies clearly show the importance of the extrapolated bias term, which aligns with the theoretical framework.", "relation_to_broader_scientific_literature": "1. This work is built on additive energy distributions for generative tasks. The authors extend it to discriminative tasks.\n\n2. The problem connects to out-of-distribution generalization and subpopulation shifts. The authors clearly articulate how compositional generalization relates to these established research areas.", "essential_references_not_discussed": "The idea that compositional generalization can only be achieved within the discrete affine hull is analogous to the assumption in the following papers, which posit that the test distribution should lie within the convex hull of the training distributions:\n[1] Qiao, F., & Peng, X. Topology-aware Robust Optimization for Out-of-Distribution Generalization. ICLR 2023.\n[2] Yao, H., Yang, X., Pan, X., Liu, S., Koh, P. W., & Finn, C. Improving Domain Generalization with Domain Relations. ICLR 2024.", "other_strengths_and_weaknesses": "Please see summary.", "other_comments_or_suggestions": "None.", "questions_for_authors": "1. The paper assumes attribute labels are available during training. Can the proposed method be adapted to settings where attribute labels are only partially available or must be inferred?\n\n2. How does the computational complexity of the proposed method scale with the number of attributes and classes?", "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": "", "extra_scores": {}}, {"review_id": "d4ivuyKcpS", "reviewer": "Reviewer_TU2X", "summary": "This paper proposes compositional risk minimization (CRM), an approach to compositional generalization that is based on additive energy distributions. The intuition is to train an energy-based classifer on the training set, then modify it to account for known bias between the observed training and test distributions. The authors show a number of theoretical results as well as empirical results on benchmarks for subpopulation shifts.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive and insightful feedback! We will fix the typo in the caption of Figure 1, thanks for pointing this. We now address the concerns raised by the reviewer ahead.\n\n**My main concern is whether the additive energy distribution assumption is realistic (beyond the particular subpopulation-shift setting considered in the experiments). It would be great to have additional discussion on this point. What practical settings is the additive energy distribution assumption applicable to? For example, does it apply to the blue elephant on the Moon example laid out in the introduction? It's reasonable if it does not hold, but discussing the boundaries of when it holds would be good to include.** \n\nWe believe that additive energy assumption is practical for settings where the image is aptly described by an AND operation among the attributes (and this applies to the blue elephant example). The summation of additive energies leads to a product of exponentials which act as a soft AND operation. Hence, each energy term contributes to checking one of the conditions in the AND operation.\n\nRegarding the specific image \"blue elephant on the moon\", let us start with a simplification of this image, \"elephant on the moon\". Then we have one energy term that detects elephant and the other energy term detects moon. However, let us now consider the original image \"blue elephant on the moon\". If we do a simple AND between detecting blue color, elephant, and the moon, then even the image \"elephant on blue moon\" will have the same energy (density) as the image \"blue elephant on the moon\", which is not desirable. \n\nBut we can model this scenario using additive energy distribution by having energy components for each object-specific attribute. Hence, the final energy function becomes $E(x|o, z)= \\sum_{o_i} \\sum_{z_{ij}}  E_{ij}(x, o_i, z_{ij})$, where $o_i$ refers to the location of the $i^{th}$ object and $z_{ij}$ refers to the $j^{th}$ attribute for the $i^{th}$ object. This essentially allows us to \"bind\" the attribute information to an object and still model the overall distribution with additive energies. In the example above, we have one energy term for the object elephant with attribute blue, which gets added with the energy term for the object moon (with some default attribue value). \n\nThanks again for this interesting question! A thorough investigation of this is a fruitful future direction. We are very open to further discussion and would be happy to address any remaining concerns.", "claims_and_evidence": "Yes, all claims are supported with adequate theoretical and empirical evidence.", "methods_and_evaluation_criteria": "The method is well constructed and the benchmarks are well chosen.", "theoretical_claims": "The proofs appear correct; no issues found.", "experimental_designs_or_analyses": "The experimental setup described in section 5.1 appears sound.", "relation_to_broader_scientific_literature": "This work is related to prior work on compositionality with energy-based models. However, these works typically consider a generative setting; by contrast, the authors consider a discriminative setting and produce a novel set of theoretical results.", "essential_references_not_discussed": "No missing references as far as I am aware.", "other_strengths_and_weaknesses": "The paper is quite well written and presented. Experiments are conducted over a number of datasets. Overall, the authors present an interesting, fresh approach to compositional generalization.\n\nMy main concern is whether the additive energy distribution assumption is realistic (beyond the particular subpopulation-shift setting considered in the experiments). It would be great to have additional discussion on this point.", "other_comments_or_suggestions": "Typos\n\n- \"boradcasting\" in Figure 1 caption", "questions_for_authors": "What practical settings is the additive energy distribution assumption applicable to? For example, does it apply to the blue elephant on the Moon example laid out in the introduction? It's reasonable if it does not hold, but discussing the boundaries of when it holds would be good to include.", "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": "", "extra_scores": {}}, {"review_id": "P7hPKf32KT", "reviewer": "Reviewer_C2eF", "summary": "This paper addresses compositional generalization by tackling compositional shift, where test data contains unseen combinations of attributes. The authors propose Compositional Risk Minimization (CRM), using additive energy distributions to model attributes and providing an alternative to empirical risk minimization. Their approach involves training an additive energy classifier and adjusting it for compositional shifts, with theoretical analysis showing extrapolation capabilities to affine hulls of seen attribute combinations. Experimental results on benchmark datasets demonstrate improved robustness compared to existing methods for handling subpopulation shifts.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive and insightful feedback! We now address the concerns raised by the reviewer ahead.\n\n**1. Limited empirical comparisons against existing baseline methods**\n\nWe have done extensive benchmarking of CRM with 6 widely used baselines in the literature of subpopulation shifts apart from ERM. Please check Table 5 in Appendix G.1, where we compare CRM with GroupDRO, LC, sLA, IRM, VREx, and Mixup. We find that CRM outperforms all the baselines across diverse benchmarks w.r.t. the worst group accuracy.\n\nNote that in the main body of the paper (Table 1), due to space constraints where we only compared with the best performing baselines (GroupDRO, LC, sLA). Also, note that our primary comparison baseline has been with GroupDRO, as it is the most effective method for addressing subpopulation shifts, hence it was the main focus of our work. We also included baselines like LA (Logit Adjustment) that share conceptual similarities with our approach. \n\n**2. The authors use 'Compositional Risk Minimization' as the title. Could you explain the core concept of Compositional Risk Minimization using simple mathematical formulations? Please refer to the classical mathematical formulation of Empirical Risk Minimization (ERM).**\n\nThe classical ERM objective can be stated as follows. \n\n$$ R_p(f) = \\sum_{z \\in \\mathcal{Z}} p(z)R(f|z)$$ \n\nwhere $R(f|z) = \\mathbb{E}_{p(x|z)}[\\ell(f(X),Y)|z]$ and $p(z)$ is the training prior probability and $\\mathcal{Z}$ is the set of all $d^m$ groups. If $\\ell$ is the cross-entropy loss, then output of ERM (with no capacity constraints) matches the true $p(z|x)$. Also, note that in the above summation $p(z)$ is zero on all groups that are not in the support of the training distribution. \n\nTo tackle compositional distribution shifts, we want to learn predictors that minimize\n\n$$ R_q(f) = \\sum_{z \\in \\mathcal{Z}} q(z)R(f|z)$$ \n\nand not $R_p(f)$, which ERM minimizes. In the above objective $q(z)$ can be non-zero on groups $z$ that have zero probability under $p(z)$. Our results show that our approach outputs a predictor that provably minimizes risk under any compositional shift $R_q(f)$ and hence the name *compositional risk minimization*.\n\nSpecifically, in Theorem 3, we showed that CRM outputs the Bayes optimal predictor and hence it provably minimizes $R_q(f)$ with a high probability, where $\\ell$ is cross-entropy loss or $0/1$ loss, as long as the number of training groups grow as $O(md + dlog d)$. \n\nTo clarify, our approach does not require us to explicitly compute the risk on the test distribution $R_q(f)$. For additive energy distributions, in the second step of CRM we adapt the $\\hat{p}(z|x)$ to $\\hat{q}(z|x)$ with the extrapolated bias $B^{\\star}$, which equals the true test predictor $q(z|x)$ (Theorem 2). Hence, CRM avoids the computation of $R_q(f)$.\n\nThanks again for this interesting question! We will add this discussion to the paper as well. We are very open to further discussion and would be happy to address any remaining concerns.", "claims_and_evidence": "The claims are supported by both theoretical and empirical evidence.", "methods_and_evaluation_criteria": "Yes, the evalucation criteria (Average Accuracy and Worst Group Accuracy) makes sense.", "theoretical_claims": "The mathematical derivations and theoretical arguments seem rigorous upon initial examination.", "experimental_designs_or_analyses": "Based on my review of the experimental section, the empirical results support the paper's claims.", "relation_to_broader_scientific_literature": "The authors provide a thorough and systematically structured overview of previous work in this research area.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n\n1. The research tackles an important problem in compositional generalization with novel insights.\n2. Clear implementation details and reproducibility through provided pseudocode.\n\nWeakness:\n\nLimited empirical comparisons against existing baseline methods", "other_comments_or_suggestions": "The authors use 'Compositional Risk Minimization' as the title. Could you explain the core concept of Compositional Risk Minimization using simple mathematical formulations? Please refer to the classical mathematical formulation of Empirical Risk Minimization (ERM).", "questions_for_authors": "No", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a novel attribute classification method that achieves Bayes-optimal classification for the test distribution, even when it includes attribute combinations not observed in the training data. Assuming that features follow an additive energy distribution, the authors derive a provable method that classifies, in a Bayes-optimal manner, any attribute composition lying within the discrete affine hull of the attribute compositions present in the training set. 

Reviewers agree that the paper is well written, with all the claims supported by rigorous theoretical analysis and extensive experiments on benchmark datasets.
Decision: Accept (poster)