On Harmonizing Implicit Subpopulations
Authors: Feng Hong
Jiangchao Yao
Yueming Lyu
Zhihan Zhou
Ivor Tsang
Ya Zhang
Yanfeng Wang
Venue: ICLR 2024 poster (2024)
Keywords: imbalanced learning
subpopulation imbalance
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose a novel method for subpopulation imbalanced learning, which succeeds in discovering and balancing the latent structures in training data.

Abstract: Machine learning algorithms learned from data with skewed distributions usually suffer from poor generalization, especially when minority classes matter as much as, or even more than majority ones. This is more challenging on class-balanced data that has some hidden imbalanced subpopulations, since prevalent techniques mainly conduct class-level calibration and cannot perform subpopulation-level adjustments without subpopulation annotations. Regarding implicit subpopulation imbalance, we reveal that the key to alleviating the detrimental effect lies in effective subpopulation discovery with proper rebalancing. We then propose a novel subpopulation-imbalanced learning method called Scatter and HarmonizE (SHE). Our method is built upon the guiding principle of optimal data partition, which involves assigning data to subpopulations in a manner that maximizes the predictive information from inputs to labels. With theoretical guarantees and empirical evidences, SHE succeeds in identifying the hidden subpopulations and encourages subpopulation-balanced predictions. Extensive experiments on various benchmark datasets show the effectiveness of SHE.

Published as a conference paper at ICLR 2024
ON HARMONIZING IMPLICIT SUBPOPULATIONS
Feng Hong'
Jiangehao Yao 1,2,80
Yueming Lyu x.
Zhihan Zhou
Ivor W. Tsang" 3.4,5
Ya Zhang 1.2
Yanfeng Wang 1.2.5
"Cooperative Medianet Innovation Center, Shanghai Jiao Tong University
"Shanghai Artificial Intelligence Laboratory *CARR, Agency for Science, Technology and Research 'IHPC, Agency for Science, Technology and Research "Nanyang Technological University
[feng.hong, Sunarker, zhihanzhou, ya zhang, wangyanfeng) @sjtu.edu.cn (Lyu Yueming, ivor tsang) @cfaraasstar.edu.gg
ABSTRACT
Machine learning algorithms learned from data with skewed distributions usually suffer from poor generalization, especially when minority classes matter as much as, or even more than majority ones. This is more challenging on class-balanced data that has some hidden imbalanced subpopulations, since prevalent techniques mainly conduct class-level calibration and cannot perform subpopulation- level adjustments without subpopulation annotations. Regarding implicit subpopulation imbalance, reveal that the key alleviating the detrimental effect lies in effective subpopulation discovery with proper rebalancing. We then propose a novel caopopulation-imaalaeee
learning method called Scatter and HarmonizE (SHE). Our method is built upon the guiding principle of optimal data partition, which involves assigning data to subpopulations in a manner that maximizes the predictive information from inputs to labels. With theoretical guarantees and empirical evidences, SHE succeeds in identifying the hidden subpopulations and encourages subpopulation-ballaeee predictions. Extensive experiments on various benchmark datasets show the effectiveness of SHE. The code is available.
INTRODUCTION
The imbalance nature inherent in real-world data challenges algorithmic robustness especially when minority classes matter as much as, or even more than majority ones (Reed, 2001; Zhang et al., 2023b). It becomes more exacerbated in scenarios where the observed categories are apparently balanced but the implicit subpopulations' remain imbalanced (Zhang et al., 2020). Specifically, such imbalance stays not in the class level but in the implicit subpopulation level, giving rise to the subpopulation imbalance problem. is ubiquitous in some sensitive applications, e.g. medical diagnosis with ethnic minorities or auto-driving decisions in weathers, yielding severe fairness concerns and generalization impairments (Yang et al., 2023).
Typical studies in imbalanced learning (Buda et al., 2018; He & Garcia, 2009; Wang et al., 2021; Menon et al., 2021; Cui et al., 2021) focus on the class-imbalance setting like Fig. I(a), employing the explicit class distribution to calibrate the training of majority and minority classes, which cannot handle implicit subpopulation imbalance like Fig. ((b). Other efforts for spurious correlations, which arise from discrepancies in class distribution across specific attributes compared the overall class distribution, aim to make predictions by causally relevant features, while excluding these spuriously correlated attributes (Nam et al., 2020: Zhang et al., 2022: Seo et al., 2022: Taghanaki et al., 2022). Our goal for implicit subpopulation imbalance, shares the similar rebalancing spirit with these works for class imbalance and spurious correlations, but differs in the underlying problems and mechanisms. We present a comprehensive comparison of these three concepts of imbalanced learning in Tab. 1.
The key challenges to cope implicit subpopulation imbalance problems are twofold. First, the mixed distribution of multiple subpopulations makes predictions more difficult (compared to single
In this paper, the term "subpopulations" pertains to some implicit attributes that differentiate the "classes" concept and contribute intra-class variations.

---

Published as a conference paper at ICLR 2024
BKX
STEE
Viaiane
ASahpepulatiies
Spiis
(a) Class imbalance
(b) Subpopulation imbalance (c) Feature variance
(d) Per-split accuracy
Figure I: (a) The number of samples for each category in CheXpert (Irvin et al., 2019). The class
index is sorted by sample numbers in descending order. The imbalance phenomenon of classes evident. (b) The imbalanced age subpopulation distribution in CheXpert (Irvin et al., 2019) with the prediction target of diseases. (c) Within-class feature variance at different subpopulation numbers. All experiments are conducted on CIFAR- 100 with an imbalance ratio IR- 100, and the within-class variance is calculated as in Papyan et al. (2020). As comparison, the within-class variance of our method for the learned subpopulations is much lower than ERM under the mixed distribution. (d) Many/Mediun/Few accuracies of ERM and SHE in COCO. The performance of minority subpopulations is poor, and our method relatively alleviates this phenomenon.
distribution). This is because vanilla classification models tend to map all training samples of the same class to identical features (Papyan et al., 2020; Han et al., 2022). However, when there are significant discrepancies within class (I.e., sampling from different subpopulations), foreing them to identical features encounters more obstacles (as illustrated in Fig. 1(c) and would impair generalization (Ma et al., 2023a). Second, different subpopulations might have different prediction mechanisms (i.e., rely on different features) and machine learning algorithms tend to ignore minority subpopulations. resulting in degraded performance on these data (as in Fig. I(d).. Besides, the implicit nature of subpopulations makes it harder to conduct rebalancing among subpopulations. These difficulties restrict existing methods from achieving practical effectiveness directly.
To address the above challenges, we propose novel method to handle implicit subpopulation imbalance, namely, Scatter and HarmonizE (SHE). Intuitively, we seek to decompose complex mixed training data into multiple simpler subpopulations, where the prediction mechanisms within each subpopulation are consistent (Scatter), and then conduct subpopulation balancing (Harmonize). Specifically, we first introduce the concept of optimal data partition, which divides training data into subpopulations that can bring the maximum additional prediction ability (Def. 3.1). Then, an empirical risk that theoretically consistent with the pursuit of optimal data partition (Eq. (1) and Thm. 3.3), proposed. To account for the imbalance nature of subpopulations, we obtain aubpopulation-balancee predictions w.r.t. the learned data partition by simply applying the LogSumlXxt operation to outputs (Thm. 3.4). Finally, a practical realization that can be optimized end-to-end without increasing model capacity is provided (Sec. 3.4). We summarize the contributions as follows:
We study the practical yet under-explored subpopulation imbalance learning problem that cannot be efficiently solved by existing methods, and identify the unique challenges, whose key lies in exploring the implicit subpopulations to facilitate prediction and subpopulation balancing.
We proposed a novel SHE method that uncovers hidden subpopulations by optimizing the predic- tion ability and achieves subpopulation-balancee operation. Theoretical analysis shows promise of SHE under implicit subpopulation imbalance. We conduct extensive experiments to comprehensively understand the characteristics of our proposed SHE, and verify its superiority in improving subpopulation imbalance robustness.
predictions by simply applying a LogSumExp
RELATED WORK
In this section, we briefly review the related works developed for the typical class imbalance and spurious correlations, which we summarize comparison with our work in Tab. 1.
Class Imbalance. Re-sampling (Buda et al., 2018: Wallace et al., 201 1) and Re- Re-weighting (Menon et al., 2013; He & Garcia, 2009) are the most widely used methods to train on class-imbalanced datasets. Explorations inspired by transfer learning (Chu et al., 2020; Wang et al., 2021) seek to transfer knowledge from head classes to tail classes to obtain more balanced performance. Menon et al. (2021); Ren et al. (2020) propose logit adjustment (LA) techniques that modify the output logits by the conditiona offset terms. The vector-scaling (VS) loss (Kini et al., 2021) instead of considering the simple additive operation, uses multiplicative factors to adjust the output logits. Ma et al. (2023b) proposes use the semantic scale measured by the feature volume rather than
2

---

Published as a conference paper at ICLR 2024
Table I: A comparison of different types of imbalance problems, including class-level shifts, subpopulation-level shifts, assumptions underlying the problem and possible negative impacts. For class imbalance, the training class distribution is skewed, i.e. py (v) > pr where y = arg max Py (v))
arg minyey Py For spurious correlation, is assumed that sub- populations and classes are causally independent but there exists E 8 that is spuriously correlated with class y E y in training. For subpopulation imbalance, the subpopulation distribution of training data is imbalanced, L.., (s(s) > ps (s'),, where s = max,es Ps (s) arg minges (s(s). For simplicity, we use p(-) without subscripts in the following sections adapt various variables.
Imbalance type
Subpopulation shift
Class shift
Assumption
Detrimental Impact prediction
Class Imbalance
pv(u) x pv(x')
Predict minority classes majority classes
Spurious Correlation
Pr (r(s(ale) x py (v)
Su Y
Predict relying on irrelevant features
Subpopulation Imhhllnnee
(*( > ps(s')
Ignore features for minority subpopulations
the sample size of classes to guide the class rebalancing. Cui et al. (2021); Zhu et al. (2022) further improve the prediction performance under class imbalanced data by combining the contrastive learning techniques. Some work (Zhou et al., 2022; 2023; Hong et al., 2023; Zheng et al., 2024) has explored overcoming class imbalance in the context of unsupervised or weakly supervised learning.
Spurious Correlations. The distributionally robust optimization (DRO) framework (Ben- Tal et al., 2013; Gao al., 2017: Duchi al., 2021) has been proposed to improve the worst case generalization. However, the DRO objective results in excessive attention to worst cases, even if they are implausible. Group DRO (GDRO) (Sagawa et al., 2019) optimizes soft version of worst-case performance over set of subgroups, which despite effectiveness requires prior subgroup labels available. Some efforts (Nam et al., 2020; Zhang et al., 2022; Seo et al., 2022) have been made to reduce the reliance on the group-level supervision, but primarily focus on mitigating spurious correlation instead of the imbalance among causal factors, namely, removing the false associations between labels and irrelevant features in training samples. The typical scheme is first detecting minority group and then designing an algorithm to promote the detected minority group. Following this framework, a series of works (Nam al., 2020; Liu et al., 2021; Zhang al., 2022) explore the minority discovery, which assumes that ERM models are prone to rely on spuriously correlated attributes for prediction, and therefore the failure samples are the minority ones. Some other works (Sohoni et al., 2020; Seo et al., 2022; Liu et al., 2023) treat the model predictions or feature clustering results directly as spuriously correlated features, which in combination with ground truth can yield more fine- -nn-grrined subgroup labels. MaskTune (Taghanaki et al., 2022) forces the trained model for more feature exploration by masking, to indirectly mitigate spurious correlations.
3 METHOD
3.1 PROBLEM FORMULATION
Let X' be the input space and y = (1.,,....CC C} be the class space. We denote the underlying space of subpopulations as S = {1, 2.., K}. The overall data distribution
be formulated mixture of distributions of latent subpopulations, i.e., p( x y) = E.es p(s) p(x, yls). The training set can be
denoted as D {(x;, V/. s; >;));= ()) where any input æ; is associated with classifica- tion label y and an unobserved subpopulation label s,. Here we focus on the implicit subpopulation imbalance problem, i.e., p()) is skewed. We assume that subpopulations are heterogeneous with inconsistant predictive mechanisms. That is, data distribution p( xc, y/s) differs across subpopulations, and p(y| az, s) may vary among certain subpopulations. For fair evaluation among all subpop- ulations, a subpopulationnbaaaaee
test distribution (x, y) = Exes Poal (s)p((, yls), where (s) = k: Vs E S, is used for evaluation following imbalanced learning literatures (Menon et al., 2021; Cao et al., 2019). In nutshell, the goal is to learn deep model f : X p()) D that minimizes the following aubpopplation-ballnee
error rate (SBER):
min SBER(J) = E( -Poar 1 (y # arg max l (x)). x'E)
In our experiments, we use a subpopulation test set as an unbiased estimator for SBER.
3.2 MOTIVATION
In Fig. 2, we visualize a toy motivating example whose prediction goal is to distinguish between circles (semi-transparent) and triangles tnon-transparent) For training data, they are sampled from
3

---

Published as a conference paper at ICLR 2024
Training
Class I
XL
Class
Subpopulation
Subpopulation 2
Test
X1
X1
SHE predictions learned subpopulations (train)
Ground truth
ERM predictions
Visualization of toy motivating example, which is a 2D ceapopulation-imalllnnn learning problem. The left column illustrates the data distribution of the training set and that of the test set under 2 classes consisting of 2 subpopulations. The middle column exhibits the model prediction of ERM. The right column shows the predictions and the learned subpopulations of SHE on the training set and predictions on the test set. The training set is highly subpopulationiimallee the imbalance ratio IR- 20 and the test set is balanced (referring to Appx. F.1 for more details).
Figure 2:
with
both Subpopulation I (blue) and Subpopulation 2 (red), and the training samples of Subpopulation 2 are much less than those of Subpopulation I, i.e. under subpopulation imbalance. About the test set, it is balanced sampled from both subpopulations, i.e... under subpopulation balance". According to the visualization in Fig. 2y a more important feature in the class prediction for Subpopulation 1, while in terms of Subpopulation 2, 22 can be a more effective feature in the class prediction. Unfortunately, due to the subpopulation imbalance, ERM's s predictions rely heavily on æ and perform poorly in Subpopulation 2. However, if we can accurately identify the latent subpopulations in the training data, such a classification problem a mixed distribution can transformed into two simple linear classification problems, and the key features in Subpopulation 2 will not be ignored. Therefore, the key to alleviating subpopulation imbalance is to discover the potential subpopulations in the training data that promote prediction and subpopulation rebalancing. In the right column of Fig. 2, we present the predictions and the learned subpopulations of SHE on the training set and the corresponding predictions on the test set. As can be seen, SHE successfully discriminates between two subpopulations on the training data, with the aid of which more accurate predictions are obtained.
3.3 SCATTER AND HARMONIZE
Optimal Data Partition. For data with implicit heterogeneous structures, we resort to a proper data partition so that each partition has consistent predictive mechanism during training. Such way promotes the prediction ability and helps protect vulnerable subpopulations. the following, we first introduce the optimal data partition in Def. 3.1 that learns assign samples to subpopulations.
Definition 3.1 ((Optimal) Data Partition). Let X and Y be random variables that take values in X' x y following fixed joint distribution px,y. A data partition defined as mapping v of the training data and its labels to the subpopulation space, i.e., / X' x ) - S. So u(X, Y) is random K. We then define the optimal data partition based on
variable taking values from S and |S] information theory as
v² arg [(X; Y; (xx, Y)) arg (XX
where I(X;Y; denotes the interaction information (MeGill, 1954) of X, X,Y,((X,Y) I(X;)) Y) denotes the mutual information of X and Y. and I(X:Y|(XX)) (|x,X, Y)) denotes the conditional mutual information between X and Y given
In information theory, the mutual information I(X;)) can characterize the prediction ability from input X to class label Y (Cover& & Thomas, 2006). The interaction information [(X; Y; ( Y)) means the gain of correlation between X and Y given data partition A larger I(X;; ( Y) indicates greater improvement in the prediction ability of a data partition v from input to label
In practice, it is common to have citpopulationimmaalanece set for training. And for the test set, need build cappoppllaiionaalaneee counterpart evaluate the algorithmic rohustness wet. latent suhpopulations.
4

---

Published as a conference paper at ICLR 2024
Y. Due to the hierarchical nature of semantics (Deng al., 2009), the data partition usually comes with multiple possibilities. Def. 3.1 helps us pursue the optimal data partition v to maximize the prediction ability of the training data. Intuitively, the optimal data partition decomposes the prediction in complex mixed distribution into several classification problems in multiple simple distributions partitioned by v" In the following, we remark advantageous property of the optimal data partition.
Proposition 3.2. The optimal data partition at least does not inhibits the prediction ability, i.e., T(X; Y:u (X,Y) 2 O.
Prop. 3.2 shows that the optimal data partition can help to improve the prediction ability, and at least has no negative impact even in the worst case. Please refer to Appx. C.I for the proof.
Objective. After introducing the above concept and analysis, we explore incorporating the idea of optimal data partition to improve the prediction ability and achieve subpopulation -bllnnced model. For this reason, we propose the following empirical risk with respect to the training set D, whose relation with the optimal data partition will be proved and discussed in the subsequent theorem.
(I) where 1()) denotes the indicator function, with value when is true and 0 otherwise, (a: is the prediction of for subpopulation s, i.e., Se X' p()), and (y>([u(( Y)) is the empirical entropy of labels conditioning on the data partition v with respect to the training set D. We use the following Thm. 3.3 discuss the consistency between Eq. (1) and the optimal data partition.
R(S, u1 = - E Σ [v( 14) 8) log f (a.)) HD Y)),
2=1 #ES
Theorem 3.3. Let f = arg miny R(S, v; D) be the optimal solution for the empirical risk R(D) in Eq. for any 'D and v. Assume that the hypothesis space H satisfies Vx E X, Vy V. V H. log fx(x) > -m, where m > 0. Define mapping family G = {g X x ) - Rlg( y)
log f (x), f HY and RN (G) O( (tx)) denotes the Rademacher complexity of G with the sample
size N (Bartlett & Mendelson, 2002) (detailed in Appx. B.3). Then for any ō E (0, 1), have: ( (
(-R(F", vV + BB} = with probability at least 1 8, where B -I(;YY) is constant, and K is the number of subpopulations.
m V 2log 8 + 4K Rx (G),
VN
Thm. 3.3 presents an important implication that minimizing the empirical risk R in Eq. (1) asym- potitically aligns with the direction of maximizing ((x;Y:((X))) in Def. 3.1 in a sense of statis- tical consistency. We kindly refer the readers to Appx. C.2 for the complete proof. To further verify this, we trace the Normalized Mutual Information (NMI) score (Strehl & Ghosh, 2002) between the learned sub- populations and the true subpopulation annotations dur- ing training in each epoch and visualize it in Fig. 3. It can be seen that our method gradually learns the sub- populations that correlates well to the true annotations. learned by our method in COCO in Fig. 5 in Appendix. It can be observed that our method uncovers meaningful subpopulations, i.e., Subpopulation 1: cut up apples bananas; Subpopulation 2: the whole apples or bananas. Fig. 3 and Fig. 5 demonstrate the promise of SHE to discover the latent subpopulation structure inherent in the training samples.
1.0
NMI
0.2
0.0g
Epoch
Figure 3: NMI scores between the learned subpopulations and the true annotations on the toy dataset in Fig. 2 during training. We also visualize the two subpopulations
Subpopulation-bable balanced prediction. With the inferred subpopulations, we discuss how to achieve subpopulation- balanced predictions. Let (x) be the output logits of for any subpopulation S and fx( = softmax(2, (x)) We show that the overall prediction f(x) = ((ntmax(((()))) with xxx) = log Eoes e²((m) is ciopooplationn-allneee according the following Thm. 3.4.
Theorem 3.4. Supposing that for any subpopulation E S, 2, can perfectly fit the data distribution of a given subpopulation 8, i.e., p(æ, yl oc ="((), then = llog Eses e² can perfectly fit the subpopulation eiblopullaion-bbalaeeedddd overall distribution, i.e., (bn v) xc ="(x)
Thm. 3.4 implies that alongside pursuing the optimal data partition, the LogSumExp operation on the logits of the learned subpopulations can be directly aggregated into balanced prediction. We
5

---

Published as a conference paper at ICLR 2024
kindly refer readers to Appx. C.3 for more details. By contrast, the ordinary learning methods will fit the distribution p(z, y) = Ees p(s) p(2, y/s), which is non-robust to subpopulation imbalance.
Discussion. We would like to briefly discuss the core differences between SHE and some related techniques. Classic clustering methods (Cheng, 1995; Asano et al., 2020; Caron et al., 2020) divide the input space X' into several disjoint clusters, with the goal that the clusters match as closely to the target classes. Our method, on the other hand, divides the data in a subpopulation level instead of the class level, with the goal that the partition maximally intervenes with predictions from input classes. Some works for spurious correlations (Sohoni et al., 2020; Seo et al., 2022; Liu et al., 2023) use the predictions of ERM or their feature clustering results as subpopulations, based on an underlying assumption that data from the same subpopulation will have the same ERM predictions or features and conversely not. Such an assumption might not be valid, especially when there are not many spurious associations captured during training. In this case, the clustering learned by these methods remains at the class level, as the ERM model uses the given classes as supervision. comparison, SHE has theoretically and empirically been oriented to learn meaningful subpopulation structures.
3.4 REALIZATION
Optimization for the data partition v. We use a subpopulation-wett matrix V E (VIV e K
R ,.. Σx= Uis 1, Vi = 11,,,,,.... 1,2,...,,N) to represent data partition in Eq. (1) with respect to the training set D. Each us in V denotes the probability of the i-th data point being sampled from the subpopulation s, i.e., Diss p(v(2,, y1) s). To accelerate the optimization of V. we further propose
a diversity regularization term Div( = Σ.,, s5p,,, #sg fl (ar - Saz (z))l2- which prevents the collapse together of different subpopulations. Increasing the diversity among the outputs can also foree the model to learn richer features to help prediction (Brown et al., 2005; Krogh & Vedelsby, 1994; Tang et al., 2006). Thus, the final loss function of our method can be formulated as follows:
1 N Σ Ui. log f.d (x.) H) -
N
-=-N N
C C=
Σ Div(z;)
2=1 SES
==1
where 3 is hyperparameter that controls the weight of the diversity regularization term.
Multi-head strategy. A classical classification model f parameterized by 0 consists of deep feature extractor 1o and a linear classifier 9 with the parameter matrix W. The final prediction is denoted as f(x) = (lttmax((xx))... where z is the output logits of æ, i.e., x(x) = ((v(x))
w WT vxz). Since we need to obtain separate prediction results for each subpopulation in Eq. (2). we apply a multi-head strategy following Tang et al. (2020); Vaswani et al. (2017). Specifically, we equally divide the channels of the feature and the classifier weight into K groups. i.e. (xx)
[1/1 (ar),
..... (x)]. W = [[11, W.. Wx] and the outputs logits for any subpopulation E S is denoted as (x) W, (w)) Thus the final subpopulation-ballnee
prediction is obtained
by f(x) = softmax(2(x:)), where =(x) = log Lees e (=((a))) according to Thm. 3.4. Note that, our multi-head strategy does not introduce any additional parameters the network compared with the network counterpart without considering the subpopulation imbalance. That is to say, we just split the output features of the penultimate layer and the classifier weights of the last layer into different groups, and use them to generate the corresponding predictions for multiple subpopulations.
4 EXPERIMENT
4.1 EXPERIMENTAL SETUP
Datasets. We evaluate our SHE on COCO (Lin et al., 2014), CIFAR- (Krizhevsky et al., 2009). and ticredImageNet (Ren et al., 2018). For COCO, we follow the ALT-protocol (Tang et al., 2022) to conduct subpopulation-imen -nattionnimbllaaeed training set and balanced test set. For CIFAR-100, we take the 20 superclasses as classification targets and generate subpopulation imbalances by sampling in the subclasses of each superclass. Following Cui et al. (2019), we use the exponential sampling with
max.es L(o..n 1(s.-s)
imbalance ratio IR E (20, 50, 100), where IR- = min,es Eio. 1(a,=a) For tieredImageNet,
we take the 34 superclasses as classification targets and generate subpopulation imbalances by imbalanced sampling in 10 subclasses of each superclass with the imbalance ratio IR. = 100.
Baselines. We consider extensive baselines: 1) empirical risk minimization (ERM); 2) imbalanced learning methods: PaCO (Cui et al., 2021), BCL. (Zhu et al., 2022), IFI. (Tang et al., 2022), DB (Ma

---

Published as a conference paper at ICLR 2024
Table 2: Performance (Mean i Std) of methods on COCO, CIFAR-100 with the imbalance ratio IR E (100, 50, 20) (marked as CIFAR-IRIR), and magred Bold indicates the best results.
Method
COCO
CIFAR-IRI00
CIFAR-IRSO
DIFAR-IR20
tieredImageNet
ERM
62.52 0.32%
52.49 + 0.27%
55.20 + 0.41%
58.92 + 0.62%
48.23 0.27%
PaCO
62,59 + 0.24%
52.89 + 0.39%
55.47 : 0.29%
59.15 + 0,44%
48.72 + 0,45%
BCL
62.83 + 0.42%
53.02 + 0.26%
55,50 +:20 0.33%
59.07 + 0.23%
48,56 + 0.61 %
IFL
62.57 + 0.15%
52.45 + 0.33%
55.16 f 0.42%
59.07 = 0.51%
48.64 0.18%
DB
62.72 0.48%
52.96 + 0.21%
55.52 4 0.27%
59.19 + 0.37%
48.52 0.13%
TDE
62.64 + 0.27%
52.67 I 0.12%
55.34 + 0.17%
59.10 + 0.22%
48.36 0.54%
ETE-DR
62.45 2 0.37%
52.43 = 0.18%
55.27 = 0.13%
58.87 + 0.17%
48.51 + 0.66%
L.f
62.06 0.83%
52.13 + 0.52%
54.78 + 0.64%
58.54 + 0.52%
47.87 + 0.23%
Focal
61.67 0.53%
$1.77 + 0.63%
54.64 + 0.62%
58.33 + 0.73%
47.68 0.62%
EIIL
62.61 0.33%
52.82 + 0.17%
55.55 + 0.32%
59.02 + 0.35%
48.56 0.33%
ARL
62.48 1 0.22%
52.67 + 0.36%
55.32 + 0.17%
59.03 + 0.24%
48.55 0.38%
GRASP
62.73 + 0.25%
52.92 + 0.41%
55.62 + 0.30%
59.12 I 0.27%
48.37 + 0.24%
JTT
62.32 + 0.75%
52.37 + 0.48%
55.02 + 0.32%
58.61 + 0.64%
48.04 + 0.39%
MaskTune
60.23 + 0.73%
51.63 + 0.31%
54.35 t 0.49%
58.03 + 0.36%
47.56 + 0.54%
SHE
64.56 0.24%
54.52 1 0.35%
56.87 + 0.17%
60.72 + 0.41% %
50.14 + 0.18%
et al., 2023b), TDE (Tang et al., 2020), and ETF-DR (Yang et al., 2022); 3) methods for spurious correlations that do not require subpopulation annotation on the training and validation set: LfF (Nam et al., 2020). Focal (Lin et al., 2017), EIIL (Creager et al., 2021), ARL (Lahoti et al., 2020), GRASP (Zeng et al., 2022), JTT (Liu et al., 2021), and MaskTune (Taghanaki et al., 2022). Note that, some imbalance learning methods like LA (Menon et al., 2021), LDAM (Cao et al., 2019), and CB (Cui et al., 2019) will degrade to the ERM performance when the class level is balanced.
Implementation details. We use 18-layer ResNet as the backbone. The standard data augmentations are applied as in Cubuk et al. (2020). The mini-batch size is set to 256 and all the methods are trained using SGD with momentum of 0.9 and weight decay of 0.005 as the optimizer. The pre-defined K is set to 4 if not specifically stated and the hyper-parameter 33 in Eq. (2) is set to 1.0. The initial learning rate is set 0.1. We train the model for 200 epochs with the cosine learning-rate scheduling.
4.2 PERFORMANCE EVALUATION ON SUBPOPULATION IMBALANCE
Overall performance. Tab. 2, we summarize the top- test accuracies on three datasets, COCO, CIFAR-100 with imbalance ratio IR {100, 50, 20} and tieredImageNet. As can be seen, SHE achieves consistent improvement over all baselines on these benchmark settings. Specifically, we achieve the gains 1.72% on COCO, 1.50%, 1.35%, 1.53% on CIFAR- 100 with three imbalance ratios, and 1.42% on ticredImageNet compared to the best baseline. In comparision, imbalanced basclines usually show marginal improvement or perform comparably with ERM, whose gains mainly come from contrastive representation learning (e.g., PaCO), invariant representation learning (e.g., IFI.), and robust classifier design (e.g. ETF-DR), etc. The baselines regarding spurious correlations, on the other hand, usually assume that the model tends to fit spurious correlations, leading to performance degradation when there are no obvious spurious correlations captured by the model during training.
Many/Mediun/Fert analysis. In Tab. 3, we show the fine-grained per-split accuracies of different methods on COCO. Note that, the Many/Mediun/Few three splits correspond to the training sample number of the subpopulation that ranks in the top, middle and bottom partitions. As expected, baselines generally have higher accuracy in dominant subpopulations but perform poorly in tails. On the Few-split, a gap of 4.42% is achieved between SHE and the best baseline, and we achieve the best results on Many-split and Medium-split. This shows a merit of SHE that enhances the performance of minority subpopulations without sacrificing the performance of head subpopulations.
4.3 PERFORMANCE EVALCATION ON RICHER IMBALANCE CONTEXTS
Training under subpopulation imbalance coupled with class imbalance. It is practical to see how SHE performs when both class and subpopulation imbalances coexist in the data. To verify this, follow (Tang et al., 2022) to construct a class and subpopulation imbalanced training set. For CIFAR and tieredlmageNet, we construct the training set by imbalanced sampling with an imbalance ratio
7

---

Published as a conference paper at ICLR 2024
Table 3: Per-split accuracies on COCO. Many, Medium, and Few are the three splits of the test set based on the training imbalancedness. Overall means the full test set. MT: the short for MaskTune. The complete experimental results (Mean t Std) of all baselines can be found in Appx. F.3.
Method
ERM
PaCO
BCL
IFL
DB
TDE
EII
ARL
GRASP
JIT
MT
SHE
Many
67.21%
67.45%
66.89%
67.71% %
67.35% 66.32% 66.87% 67.32%
67.13%
66.93%
64.48%
67.71% %
Medium
52.22%
53.33%
53.21%
52.17%
52.11%
53.23% 52.79%
53.34%
53.26%
51.24%
50.11%
53.50%
Few
37.10%
36.23%
37.67%
36.82%
37.47%
37.02% 37.06% 37.18% 37.29%
36.48%
33.27%
42.09%
Overall
62.52% 62.59% 62.93% 62.57% 62.72% 62.64% 62.61% 62.48% 62.73% 62.32% 60.23% 64.56%
Table 4: Performance under more imbalance settings. Bold indicates superior results. (Left) Performance on COCO, CIFAR-100 (IR = 100), and magred where both class imbalance and subpopulation imbalance co-exist (Mean t Std). (Right) Performance on datasets for spurious correlations. The worst group accuracy (Worst Acc) and the average accuracy (Mean Acc) is reported.
The second column means whether using the group annotation on the training or validation set.
Setting: subpopulation and imbalance (Mean Std)
Setting: sporious correlation (Worst Ace Mcan Acc)
Method
COCO
CIFAR-IR
ticredImageNet
Method
(Train/ Val)
Group Info
CelebA
Waterbirds
ERM
63.57 + 0.34%
59.24
0.46%
53.65 1 0.46%
LA
66.47 + 0.27%
59.73 : 0.27%
54.12 t 0.35%
GDRO
Yes / Yes
88.3% / 99.8%
91.4% / 99.5%%
LDAM
66.32 2 0.33%
59.66 0.26%
54.01 1 0.51%
LIF
No Yes
77.2%% /85. 1%%
82. 1% / 99.33%
CB
66.17 : 0.21%
59.45 0.36%
53.78 + 0.21%
SD
No Yes
83.2%
/91.6%
87.3% / /90.3%
PaCO
66.78 1: 0.41%
59.87 : 0.51%
54.15 : 0.39%
JTT
No Yes
888.0% 81.1%
86.7% / 993.3%
BCL
59.78 0.37%
54.23 : 0.27%
CIM
No/ Yes
81.3%
889.2%
77.2% / 995.6%
IFL
65.34 + 0.52%
59.44 + 0.41%
53.88 + 0.43%
ERM 1.FF
No
47.2% /95 6%
74.9% /8, 1%
DB
66.43 0.15%
59.81 0.29%
54.14 + 0.30%
No/ No
24.4% / 1%
67.55% / 87.55%
TDE
66.12 : 0.44%
59.63 0,34%
53.91 + 0.28%
No
40.6%
88.00%
71.8% 992.3%
ETF.DR
65.92 1 0.26%
59.71 + 0.18%
54.07 1 0.31%
MaskTune
No/ No
78.0%
/ 99.3%%
80.7% / 92.11%
SHE
68.11 # 0.27%
61.67 0.31%
55.73
0.22%
No/ No
77.9% /91.7%
81.9%
IR = 100 on both classes and subpopulations. The classes and subpopulations are both balanced on the test set. According to the results in Tab. 4 (left), we can see that the imbalance learning baselines consistently improve test accuracy compared to ERM when class imbalance also exists. When we combine SHE with a classical imbalanced learning baseline LA (Menon et al., 2021), our SHE achieves 1. 1.19% improvement on COCO, 1.80% on CIFAR and 1.50% on magred compared to the best baseline, showing the potential of SHE on more complex imbalance lcarning problems.
Training under spurious correlations. We directly apply SHE into GDRO (Sagawa et al., 2019) (using the learned subpopulations instead of the prior subgroup annotations) to verify the effectiveness on spurious correlation datasets, CelebA (Liu et al., 2015) and Waterbirds (Sagawa et al., 2019). In Tab. 4 (right), we compare SHER/GDRO with a series of baselines, and our method achieves the promising performance in mitigating spurious correlations when there is no group information avail- able. Methods that require group annotations (e.g.. SD (Pezeshki et al., 2021) and CIM (Taghanaki al., 202 are also exhibited for reference. Interestingly, more visualization results in Appx. F2 show that the performance comes from dividing the training data into two meaningful subpopulations: data w/ and w/o spurious correlations, which is actually different from the prior group annotations.
4.4 ABLATION STUDY AND ANALYSIS
Ablation on varying the latent subpopulation number K. To study the effect of the latent subpopulation number K in SHE, we conduct ablation on COCO as shown in Fig. 4(a). When K = 1. Eq. (2) degenerates to the cross-entropy loss, and so is performance. When K 1, SHE shows a significant improvement over ERM and is robust to K. At K 4, our SHE achieves the best results on average. Similar phenomenon on CIFAR and tieredImageNer can be found in Appx. F.5.
Effect of (a) the diversity term and (b) the entropy term. To study the effect of the diversity term Div(#) Eq. (2), we conduct experiments on 3 on COCO. As shown Fig. 4(b), even without the diversity term (3- 0), SHE still significantly outperforms the ERM baseline. The addition of the diversity term continually enhances the performance to the best on average at 3 = 1.0, and SHE is generally robust to the choice of B. We also conduct comparison with SHE without the entropy term HD in Eq. (2) (termed as SHEW entropy) in Tab. 5, which confirms that the entropy term consistently and effectively enhances the performance.
Effect of pursuing (a) the optimal data partition and (b) the cnoopulation-allanee predic- tion. In Tab. 5, we present the performance of ERM, ERM with the multi-head strategy (namely
8

---

Published as a conference paper at ICLR 2024
ABL
6.5
IIIL
Acunzyy
NM0 mooee
SHE ERM
-B-
SHE
665
0,2
8.3
6:00
0000
CHAIBAR2 CHAABAREE WWA.BCCAAKAAA CIAAMIHUX
ABL
GRASP 5M0
Dotiset
Metlodk
(a) Ablation on K
(b) Ablation on B (c) Optimization approach for V. (d) NMI score
Figure 4: (a) Performance of SHE and ERM on COCO with varying subpopulation number
K. (b) Performance of SHE and ERM on COCO with varying 3. (c) Performance of ERM, mHEml based V, and SHE COCO, CIFAR-IR20, CIFAR-IR50, and CIFAR-IR 100.0 (d) NMI scores between the learned subpopulations and the true annotations on Waterbird.
Table 5: Performance ERM, SHE, and some of their variants on COCO and CIFAR- 100.
Method
ERM
ERMaultd-hhadd
SHEgmL
SHE entropy
SHE
COCO
62.52 + 0.32%
62.47 = 0.28%
62.82 + 0.27%
64.15 + 0.27%
64.56
0.24%
CIFAR-IR IR 100
52,49 + 0.27%
52.53
0.17%
52.63 t 0.22%
53.96 + 0.37%
54.52
0.35 %
JIFAR-IR50
55.20 I 0.41%
55.16 0.47%
55,36 t 0.37%
56,31 : 0.23%
56,87
0.17 %
TEAR-IR20
58.92 + 0.62%
58.88 0.36%
59.21 t 0.48%
60.03 + 0.38%
60.72 0.41 %
ERMmutttheadd. and SHE by removing the multi-head network but following the way of EIIL to utilize the learned subpopulations (namely SHEguL): SHE achieves significant improvement over ERM and ERMmuiti-haath while ERMmulti-heud achieves only comparable results to ERM, showing the necessarity of pursuing the optimal data partition. The component of SHE pursue subpopulation- ballnned predictions is better (SHE vs. SHEEuLL), which confirms its effectiveness.
Analysis on the optimization approach for subpopulation-weight matrix V. We construct a variant of SHE uses a model-based approach to learn the data partition from image features, namely SHEmodel based v. As can be seen in Fig. 4(c), SHE based shows a clear performance degradation compared to SHE. A possible reason is that v in Def. 3.1 is a function of both the input ar and the label y,. but mHEee based can only learn the data partition from æc.
Quality of the recovered subpopulations. To investigate the capability of SHE in discovering subpopulations in training data, we conduct a comparative analysis between SHE and baselines based on subgroup inference (EIIL, ARL, GRASP). Specifically, Fig. 4(d) presents the NMI scores on Waterbird between the recovered subpopulations and the ground truth annotations. Our SHE exhibits remarkable capability to accurately discover the latent structures within the training data.
Fine-tuning from pre-trained models. Foundation models have achieved impressive performance in numerous areas in recent years (Radford et al., 2021; Rogers et al., 2020; Brown al., 2020). Fine-tuning from these pre-trained models using downstream training data is gradually becoming a prevalent paradigm. In Tab. we exhibit the results of different methods fine-tuned on the COCO dataset with three multimodal pre-training models, i.e., CLIP (VIT- B/32) (Radford et al., 2021). ALIGN (EfficientNet- L2 BERT-Large) (Jia et al., 2021), and AltCII (ViT-I.) (Chen et al., 2022). The LoRA (Hu et al., 2022) technique used for fine-tuning to speed up training and prevent overfitting, Despite the notable improvements obtained through fine-tuning compared to training from scratch, SHE consistently surpasses all baselines with different large-scale pre-trained models.
Table 6: LoRA fine-tuning of different meth- ods under three popular pre-trained models COCO. The complete results (Mean t Std) can be found in Appx. F.4
Method
CLII
ALIGN
AIICLIP
Zero-shot ERM
76.59%
78.45%
82.55%
84.46%
83.23%
84.93%
BCL
84.43%
83.42%
85.01%
IFL 84.49%
83.36%
84.89%
LIF 84.27%
83.05%
84.17%
JTT 84.37%
83.07%
84.55%
MaskTune
83.37%
82.66%
83 83.92%
CONCLUSION
SHE 85.34%
84.19%
85.76%
In this paper, we focus on a hidden subpopulation imbalance scenario and identify its several critical challenges. To alleviate the subpopulation imbalance problem, we first introduce the concept of optimal data partition, which splits the data into the subpopulations that are most helpful for prediction. Then, novel method, SHE, is proposed to uncover and balance hidden subpopulations in training data during training. partition and makes balanced predictions. Empirical evidence likewise demonstrates that our method uncovers meaningful latent structures in the data. Extensive experiments under diverse settings and different configures consistently demonstrate the effectiveness of SHE over range of baselines.
theoretically demonstrated that our method converges to optimal data

---

Published as a conference paper at ICLR 2024
ETHICS STATEMENT
By discovering the latent subpopulations the training data and encouraging subpopulation-balanee predictions, the paper aims to improve the generalization and performance parity of machine learning models across different subpopulations of data. This can have important implications for various social applications, such as medical diagnosis, auto-driving, and criminal justice, where subpopulation imbalance may exist and lead to biased or inaccurate outcomes, Our proposed method does not require the annotation of subpopulation cost of data annotation and on the other hand avoids the serious fairness consequences of annotation omission. Negative impacts may also occur when the proposed subpopulation discovery technology falls into the wrong hands, for example, it can be used to identify minorities for malicious purposes. Therefore, is the responsibility ensure that such technologies are used for the right purposes.
even the predefined semantics of subpopulation, which reduces the
REPRODUCIBILITY STATEMENT
All the experiments are conducted NVIDIA GeForce RTX 3090s with Python 3.7.10 and Pytorch 1.13.I. We provide experimental setups and implementation details in Sec, 4.1 and Appx. F.I. The theoretical proofs are given in Appx. C.
LCKNOWLEDGEMENT
This work is supported by the National Key R&D Program of China (No. 2022ZD0160702), STCSM (No. 2251 106101, No. 22511 1057000 No. 20D0I100100), 111 plan (No. BP0719010) and National Natural Science Foundation of China (No. 62306178).
REFERENCES
Alessandro Achille and Stefano Soatto. Information dropout: Learning optimal representations through noisy computation. IEEE Trans. Pattern Anal. Mach. Intell., 40(12):2877-2905 2018.
Alexander A. Alemi, lan Fischer, Joshua V. Dillon, and Kevin Murphy, Deep variational information bottleneck. In ICLR. OpenReview.net, 2017.
Yuki Markus Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labelling via simultancous clustering and representation learning. In ICLR. OpenReview.net, 2020.
Philip Bachman, R. Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In NeurlPS, pp. 15509-15519, 2019.
Yogesh Balaji, Swami Samkaranarayanan, and Rama Chellappa. Metareg: Towards domain general- ization using meta-regularizztion. In NeurlPS, pp. 1006-1016, 2018.
Pierre Baldi and Peter J. Sadowski. Understanding dropout. In Christopher J. C. Burges, Léon Bottou, Zoubin Ghahramani, and Kilian Q. Weinberger (eds.), NeurlPS, pp. 2814-2822, 2013.
Peter L. Bartlett and Shahar Mendelson. Rademacher and gaussian complexities Risk bounds and structural results. J. Mach. Learn. Res., 3:463-482, 2002.
Aharon Ben-Tal, Dick den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Manag. Sci., 59(2): 341-357, 2013.
Gavin Brown, Jeremy L. Wyatt, Rachel Harris, and Xin Yao. Diversity creation methods: survey and categorisation. Inf. Fusion, 6(1):5-20, 2005.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter. Christopher Hesse, Mark Chen. Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. In NewrIPS, 2020.
10

---

Published as a conference paper at ICLR 2024
Mateusz Buda, Atsuto Maki, and Maciej A. Mazurowski. A systematic study of the class imbalance problem in convolutional neural networks. Neural Networks, 106:249-230 259, 2018.
Flávio P. Calmon, Dennis Wei, Bhanukiran Vinzamuri, Karthikeyan Natesan Ramamurthy, and
Kush R. Varshney. Optimized pre-processing for discrimination prevention. In NeurIPS, pp.
1992-4001, 2017.
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Aréchiga, and Tengyu Ma. Learning imbalanced datasets with label-distribution-aaaae margin loss. In NeurIPS, pp. 1565-1576, 2019.
Fabio Maria Carlucci, Antonio D' D'Innocente, Silvia Bucci, Barbara Caputo, and Tatiana Tommasi. Domain generalization by solving jigsaw puzzles. In CVPR, pp. Foundation IEEE, 2019.
Computer Vision
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In NeurIPS, 2020.
Woong-Gi Chang, Tackgeun You, Seonguk Seo, Suha Kwak, and Bohyung Han. Domain-specific batch normalization for unsupervised domain adaptation. In CVPR, pp. 7354-7362. Computer Vision Foundation IEEE, 2019.
Prithvijit Chattopadhyay, Yogesh Balaji, and Judy Hoffman. Learning to balance specificity and invariance for and out of domain generalization. In ECCV, volume 12354 of Lecture Notes in Computer Science, pp. 301-318. Springer, 2020.
Mengxi Chen, Linyu Xing, Yu Wang, and Ya Zhang. Enhanced multimodal representation learning with cross-modal kd. In CVPR, pp. 11766-11775, 2023.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. In ICML, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 2020.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In ICCV, pp. 6620-9629) IEEE, 2021.
Zhongzhi Chen, Guang Liu, Bo-Wen Zhang, Fulong Ye, Qinghong Yang, and Ledell Wu. Altelip: Altering the language encoder in CLIP for extended language capabilities. CoRR, 2022.
Yizong Cheng. Mcan shift, mode seeking, and clustering. IEEE Trans. Pattern Anal. Mach. Intell., 1788):79--79, 1995.
Peng Chu, Xiao Bian, Shaopeng Liu, and Haibin Ling. Feature space augmentation for long-tailed data. In ECCV, volume 12374 of Lecture Notes in Computer Science, pp. 694-710. Springer, 2020.
Thomas M. Cover and Joy A. Thomas. Elements of information theory (2. ed.). Wiley, 2006. ISBN 978-0-471-24195-9.
Elliot Creager, learning. Marina Meila and Tong Zhang (eds.), ICML, volume 139 of Proceedings of Machine Learning Research, pp. 2189 2200. PMLR, 2021.
Jacobsen, and Richard S. Zemel. Environment inference for invariant
Ekin Dogus Cubuk, Barret Zoph, Jonathon Shlens, and Quoc Le. Randaugment: Practical automated data augmentation with reduced search space. In NeurlPS, 2020.
Jiequan Cui, Zhisheng Zhong, Shu Liu, Bei Yu, and Jiaya Jia. Parametric contrastive learning. In
ICCV, pp. 695-704, IEEE, 2021.
Yin Cui, Menglin Jia, Tsung-Yi Yi Lin, Yang Song, and Serge J. Belongie. Class-balanced loss based on effective number of samples. In CVPR, pp. 2019.
Computer Vision Foundation / IEEE,
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, pp. 248-255 IEEE Computer Society, 2009.
11

---

Published as a conference paper at ICLR 2024
John C. Duchi, Peter W. Glynn, and Hongseok Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Math. Oper. Res., 46(3):946 9999999699, 2021.
Rui Gao, Xi Chen, and Anton Kleywegt. Wasserstein distributional robustmess and regularization
in statistical learning. CoRR, abs/1712.06050, 2017.
Ziv Goldfeld and Yury Polyanskiy. The information bottleneck problem and its applications in machine learning. IEEE Sel. Areas Inf. Theory, 1(1):19-38, 2020.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your latent A new approach to self-supervised learning. In NeurIPS, 2020.
X. Y. Han, Vardan Papyan, and David L. Donoho. Neural collapse under MSE loss: Proximity to and dynamics the central path. In ICLR. OpenReview.net, 2022.
Moritz Hardt, Eric Price, and Nati Srebro. Equality of opportumity in supervised learning. In NeurIPS, pp. 3315-3323, 2016.
Tatsunori B. Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Faimesss without demographics in repeated loss minimization. In ICML, volume 80 of Proceedings of Machine Learning Research, pp. 1934-1943. PMLR, 2018.
Haibo He and Edwardo A. Garcia. Learning from imbalanced data. IEEE Trans. Knowl. Data Eng.. 21(9):123-11284 2009.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, pp. 9726-9735. Computer Vision Foundation IEEE, 2020.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In ICLR. OpenReview.net, 2017.
R. Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Philip Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In ICLR. OpenReview.net, 2019.
Feng Hong, Jiangchao Yao, Zhihan Zhou, Ya Zhang, and Yanfeng Wang. Long-tailed partial label learning via dynamic rebalancing. In ICLR. OpenReview.net, 2023.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu Yuanzhi Li, Shean Wang, Lu Wang. and Weizhu Chen. Lora: Low-rank adaptation of large language models. In ICLR. OpenReview.net, 2022.
Jeremy Irvin, Pranav Rajpurkar, Michael Ko, Yifan Yu, Silviana Ciurea-Heus, Christopher Chute, Henrik Marklund, Behzad Haghgoo, Robyn L. Ball, Katie S. Shpanskaya, Jayne Seekins, David A. Mong, Safwan S. Halabi, Jesse K. Sandberg, Ricky Jones, David B. Larson, Curtis P. Langlotz, Bhavik N. Patel, Matthew P. Lungren, and Andrew Y. Ng. Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison. In AAAI, pp. 590-597. AAAI Press, 2019.
Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc V. Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy text supervision. In Marina Meila and Tong Zhang (eds.). ICML, volume 139 of Proceedings of Machine Learning Research, pp. 4904 4916. PMLR, 2021.
Ray Jiang, Aldo Pacchiano, Tom Stepleton, Heinrich Jiang, and Silvia Chiappa. Wasserstein fair classification. In UAI, volume 115 of Proceedings of Machine Learning Research, pp. 862-872.
AUAI Press, 2019.
Kenji Kawaguchi, Zhun Deng, Xu Ji, and Jiaoyang Huang. How does information bottleneck help deep learning? In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), ICML, volume 202 of Proceedings of Machine Learning Research, pp. 16049-16096. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2024
Michael J. Kearns, Seth Neel, Aaron Roth, and Zhiwei Steven Wu. Preventing faimess gerrymander- ing: Auditing and learning for subgroup fairness. In Jennifer G. Dy and Andreas Krause (eds.), ICML, volume 80 of Proceedings of Machine Learning Research, pp. 2569-2577. PMLR, 2018.
Ganesh Ramachandra Kini, Orestis Paraskevas, Samet Oymak, and Christos Thrampoulidis. Label- imbalanced and group- sooupssnnsitvve classification under overparameterization. In NeurlPS, pp. 18970- 18983, 2021.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Anders Krogh and Jesper Vedelsby. Neural network ensembles, cross validation, and active learning. In NIPS. pp. 2-38. MIT Press, 1994.
Preethi Lahoti, Alex Beutel, Jilin Chen, Kang Lee, Flavien Prost, Nithum Thain, Xuezhi Wang, and Ed H. Chi. Faimess without demographics through adversarially reweighted learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan -isun-Tien Lin (eds.), NeurlPS, 2020.
Liang Li, Junpu Zhang, Siwei Wang, Xinwang Liu, Kenli Li, and Kegin Li. Multi-view bipartite graph clustering with coupled noisy feature filter. TKDE, (28122):28422112
2023.
Liang Li, Yuangang Pan, Jie Liu, Yue Liu, Xinwang Liu, Kenli Li, Ivor W. Tsang, and Keqin Li.
Bgae: Auto-encoding multi-view bipartite graph clustering. 7KDE, pp. 1-14, 2024.
Tian Li, Maziar Sanjabi, Ahmad Beirami, and Virginia Smith. Fair resource allocation in federated learning. In ICLR. OpenReview.net, 2020.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In ICML, volume 139 of Proceedings of Machine Learning Research, pp. 6357-633 6368. PMLR, 2021.
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In ECCV, volume 11219 of Lecture Notes in Computer Science, pp. 647-663. Springer, 2018.
Yiying Li, Yongxin Yang. Wei Zhou, and Timothy M. Hospedales. Feature-critic networks for heterogeneous domain generalization. In ICML,, volume 97 of Proceedings of Machine Learning Research, pp. 3915-39322 PMLR, 2019.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft COCO: common objects in context. In ECCV, volume 8693 of Lecture Notes in Computer Science, pp. 740-755. Springer, 2014.
Tsung- Yi Lin, Priya Goyal, Ross B. Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In ICCV, pp. 2999-3007. IEEE Computer Society, 2017.
Ralph Linsker. Self-organization in perceptual network. Computer, 21(3):105-117, 1988.
Evan Zheran Liu, Behzad Haghgoo, Annie S. Chen, Aditi Raghunathan. Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In ICML,, volume 139 of Proceedings of Machine Learning Research, pp. 6781-6792. PMLR, 2021.
Sheng Liu, Xu Zhang, Nitesh Sekhar, Yue Wu, Prateek Singhal, and Carlos Fernandez-Grandee Avoiding spurious correlations via logit correction. In ICLR. OpenReview.net, 2023.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In ICCV, pp. 1730-3735 IEEE Computer Society, 2015.
Jiawei Ma. Chong You, Sashank J. Reddi, Sadeep Jayasumana, Himanshu Jain, Felix Yu, Shih-Fu Chang, and Sanjiv Kumar. Do we need neural collapse? learning diverse features for fine-grained and loig-tail classification, 2023a.
Yanbiao Ma, Licheng Jiao, Fang Liu, Yuxin Li, Shuyuan Yang, and Xu Liu. Delving into semantic scale imbalance. In ICLR. OpenReview.net, 2023b.
13

---

Published as a conference paper at ICLR 2024
David Madras, Elliot Creager, Toniann Pitassi, and Richard S. Zemel. Learning adversarially fair and transferable representations. In ICML, volume 80 Proceedings of Machine Learning Research, pp. 3381-3390. PMLR, 2018.
Colin McDiarmid et al. On the method of bounded differences. Surveys in combinatorics, 141(1): 148-188, 1989.
William J. McGill. Multivariate information transmission. Trans. IRE Prof. Group Inf. Theory, 4: 93-111, 1954.
Aditya Krishna Menon, Harikrishna Narasimhan, Shivani Agarwal, and Sanjay Chawla. On the statistical consistency of algorithmss for binary classification under class imbalance. In JCML, volume 28 of JMLR Workshop and Conference Proceedings, pp. 603-61 611. J.L..org. 2013.
Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. In ICLR. OpenReview.net, 2021.
Jun Hyun Nam, Hyuntak Cha, Sungsoo Ahn, Jacho Lee, and Jinwoo Shin. Learning from failure: De-biasing classifier from biased classifier. In NeurlPS, 2020.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. PNAS, 117(40):2465-22466. 2020.
Mohammad Pezeshki, Sékou-Oumar Kaba, Yoshua Bengio, Aaron C. Courville, Doina Precup, and Guillaume Lajoie. Gradient starvation: A learning proclivity in neural networks. In NeurIPS, pp. 1256-1272,20 2021.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In ICML, volume 139 of Proceedings of Machine Learning Research, pp. 8748-8763. PMLR, 2021.
William J Reed. The pareto, zipf and other power laws. Economics Letters, 74(1):15-19) 2001. ISSN 0165-1765.
Jiawei Ren, Cunjun Yu, Shunan Sheng, Xiao Ma, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Balanced meta-softmax for long-tailed visual recognition. In NeurlPS, 2020.
Mengye Ren, Eleni Triantafillou, Sachin Ravi, Jake Snell, Kevin Swersky, Joshua B. Tenenbaum, Hugo Larochelle, and Richard S. Zemel. Meta-leaming for cumi-supervieed few-shot classification. In ICLR. OpenReview.net, 2018.
Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in bertology: What know about how BERT works. Trans. Assoc. Comput. Linguistics, 8:842-866, 2020.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. CORR, Abs//191.08877 108831, 2019.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In ICLR. OpenReview.net, 2020.
Sconguk Seo, Joon- Young Lee, and Bohyung Han. Unsupervised learning of debiased representations with pssudo-attributes. CVPR, pp. 16721-16730. IEEE, 2022.
Baifeng Shi, Dinghuai Zhang, Qi Dai, Zhanxing Zhu, Yadong Mu, and Jingdong Wang. Informative dropout for robust representation learning: A shape-bias perspective. In ICML, volume 119 of Proceedings of Machine Learning Research, pp. 8828-882 8839. PMLR, 2020.
Noam Slonim and Naftali Tishby. Document clustering using word clusters via the information bottleneck method. In Emmanuel J. Yannakoudakis, Nicholas J. Belkin, Peter Ingwersen, and Mun-Kew Leong (eds.), SIGIR, pp. 208-215. ACM, 2000.
Nimit Sharad Sohoni, Jared Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Vine-grained robustness in coarse-grained classification problems. In NeurlPS, 2020.
14

---

Published as a conference paper at ICLR 2024
Jiaming Song, Pratyusha Kalluri, Aditya Grover, Shengjia Zhao, and Stefano Ermon. Learning controllable fair representations. In AISTATS, volume 89 of Proceedings of Machine Learning Research, pp. 2022114-211 PMLR, 2019.
Alexander Strehl and Joydeep Ghosh. Cluster ensembles A knowledge reuse framework for combining multiple partitions. J. Mach. Learn. Res., 3:583-617, 2002.
Sacid Asgari Taghanaki, Kristy Choi, Amir Hosein Khasahmadi, and Anirudh Goyal. Robust representation learning via perceptual similarity metrics. In ICML, volume 139 of Proceedings of Machine Learning Research, pp. 10043-10053. PMLR, 2021.
Sacid Asgari Taghanaki, Aliasghar Khani, Fereshte Khani, Ali Gholami, Linh Tran, Ali Mahdavi- Amiri, and Ghassan Hamarnh. Masktune: Mitigating spurious correlations by forcing explore. In NeurlPS, 2022.
E. Ke Tang, Ponnuthurai N. Suganthan, and Xin Yao. An analysis of diversity measures. Mach. Learn., 65(1):247-271, 2006.
Kaihua Tang, Jianqiang Huang, and Hanwang Zhang. Long-tailed classification by keeping the good and removing the bad momentum causal effect. In NeurlPS, 2020.
Kaihua Tang, Mingyuan Tao, Jiaxin Qi, Zhenguang Liu, and Hanwang Zhang. Invariant feature learning for generalized long-tailed classification. In ECCV, volume 13684 of Lecture Notes in Computer Science, pp. 709-726. Springer, 2022.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. In ECCV, volume 12356 of Lecture Notes in Computer Science, pp. 776-794. Springer, 2020.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In ITW, pp. 1-5. IÉEE, 2015.
Naftali Tishby, Fernando C. N. Percira, and William Bialek. The information bottleneck method. CoRR, physies/0004057, 2000.
Michael Tschannen, Josip Djolonga, Paul K. Rubenstein, Sylvain Gelly, and Mario Lucic. On mutual information maximization for representation learning. In ICLR. OpenReview.net, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurlPS, pp. 5998 0008-600,,400 2017.
Byron C. Wallace, Kevin Small, Carla E. Brodley, and Thomas A. Trikalinos. Class imbalance, redux. In ICDM, pp. 754-763. IEEE Computer Society, 2011.
Jianfeng Wang, Thomas Lukasiewicz, Xiaolin Hu, Jianfei Cai, and Zhenghua Xu, RSG: simple but effective module for learning imbalanced datasets. In CVPR, pp. 3784-3793, Computer Vision Foundation/ IEEE, 2021.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron C. Courville, Ruslan Salakhutdinov, Richard S. Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In Francis R. Bach and David M. Blei (eds.), ICML, volume 37 of JMLR Workshop and Conference Proceedings, pp.
IMLR.org. 2015.
Qinwei Xu, Ruipeng Zhang, Ya Zhang, Yanfeng Wang, and Qi Tian. A Courier-based framework for domain generalization. In CVPR, pp. 14383-14392. Computer Vision Foundation IEEE, 2021.
Qinwei Xu, Ruipeng Zhang, Ya Zhang, Yi-Yan Wu, and Yanfeng Wang. Federated adversarial domain hallucination for wrivacy-preserving domain generalization. IEEE Transactions Multimedia, pp. 1-13, 2023.
Yilun Xu, Shengjia Zhao, Jiaming Song, Russell Stewart, and Stefano Ermon. A theory of usable information under computational constraints. In ICLR. OpenReview.net, 2020.
Yibo Yang, Shixiang Chen, Xiangtai Li, Liang Xie, Zhouchen Lin, and Dacheng Tao. Inducing neural collapse in imbalanced learning: Do we really need a learnable classifier at the end deep neural network? In NewrIPS, 2022.
15

---

Published as a conference paper at ICLR 2024
Yuzhe Yang, Haoran Zhang, Dina Katabi, and Marzyeh Ghassemi. Change is hard: A closer look subpopulation shift. In ICML, 2023.
Rui Ye, Zhenyang Ni, Chenxin Xu, Jianyu Wang, Siheng Chen, and Yonina C Eldar. Fedfm: Anchor-based feature matching for data heterogeneity in federated learning. TSP, 2023.
Muhammad Bilal Zafar, Isabel Valera, Manuel Gomez -Rodiizz. and Krishna Gummadi. Fairness constraints: A flexible approach for fair classification. J. Mach. Learn. Res., 20:75:1-75:42, 2019.
Yuchen Zeng, Kristjan H. Greenewald, Kangwook Lee, Justin Solomon, and Mikhail Yurochkin. Outlier-obbust group inference via gradient space clustering. CoRR, abs/2210.06759, 2022.
Haoran Zhang, Amy X. Lu, Mohamed Abdalla, Matthew B. A. McDermott, and Marzyeh Ghassemi. Hurtful words: quantifying biases in clinical contextual word embeddings. In CHII, pp. 110-120.
ACM, 2020.
Lei Zhang, Yingjun Du, Jiayi Shen, and Xiantong Zhen. Learning to learn with variational inference for cross-domain image classification. IEEE Trans. Multim,, 25:3319-3328, 2023a.
Michael Zhang, Nimit Sharad Sohoni, Hongyang R. Zhang, Chelsea Finn, and Christopher Ré. Conrect-n-contrasst contrastive approach for improving robustness spurious correlations. In ICML, volume 162 of Proceedings of Machine Learning Research, pp. 26484-26511 PMLR, 2022.
Yifan Zhang, Bingyi Kang, Bryan Hooi, Shuicheng Yan, and Jiashi Feng. Deep long-tailed learning: A survey. IEEE Trans. Pattern Anal. Mach. Intell., 2023b.
Han Zhao, Amanda Coston, Tameem Adel, and Geoffrey J. Gordon. Conditional learning of fair representations. ICLR. OpenReview.net, 2020a.
Shanshan Zhao, Mingming Gong, Tongliang Liu, Huan Fu, and Dacheng Tao. Domain generalization via entropy regularization. In NeurlPS, 2020b.
Zihua Zhao, Mengxi Chen, Tianjie Dai, Jiangehao Yao, Bo Han, Ya Zhang, and Yanfeng Wang. Mitigating noisy correspondence with geometrical structure consistency learning, In CVPR, 2024.
Hongwei Zheng, Linyuan Zhou, Han Li, Jinming Su, Xiaoming Wei, and Xu Xiaoming. Bem:
Balanced and entropy-based mix for long-tailed semi-supervised learning. CVPR, 2024.
Kaiyang Zhou, Yongxin Yang. Yu Qiao, and Tao Xiang. Domain generalization with mixstyle. In ICLR. OpenReview.net, 2021.
Zhihan Zhou, Jiangchao Yao, Yan-Feng Wang, Bo Han, and Ya Zhang. Contrastive learning with boosted memorization. In ICML, volume 162 of Proceedings of Machine Learning Research, pp. 27367-27377. PMLR, 2022.
Zhihan Zhou, Jiangchao Yao, Feng Hong, Ya Zhang, Bo Han, and Yanfeng Wang. Combating representation learning disparity with geometric harmonization. In NeurlPS, 2023.
Jianggang Zhu, Zheng Wang, Jingjing Chen, Yi-Ping Phoebe Chen, and Yu-Gang Jiang. Balanced contrastive learning for long-tailed visual recognition. CVPR, pp. 1898-6907, IEBE, 2022.
16

---

Published as a conference paper at ICLR 2024
APPENDIX
CONTENTS
Broader Related Work
19
A.1
Information Theory-Guided Objective Design
19
A.2 Domain Generalization
A.3 Algorithmic Fairness
20
A.4
Comparision with Creager et al. (2021) and Lahoti et al. (2020)
20
B Supplementary Equations
20
B.1 Mutual Information and Entropy
20
B.2 Normalized Mutual Information (NMI) score
21
B.3 Rademacher complexity in Thm. 3.3
B.4 McDiarmid's Inequality
22
Theoretical Proofs
22
C.I Proof of Prop. 3.2
22
C.2 Proof of Thm. 3.3
22
C,3 Proof of Thm. 3.4
25
C,4 An Extension of Thm. 3.4
26
Additional Discussions of Eq. (2)
26
The efficiency and scalability of optimizing V
26
F
Detailed Supplement for Experiments
27
F1
Supplemental Description of the Experimental Setup
27
F2
More Visualization of Learned Subpopulations
................................
28
F3
Complete Results of Tab. 3
29
F.4
Complete Resuls of Tab. 6
29
E5
More Results on Varying the Latent Subpopulation Number K
29
F.6
Comparison with More Clustering and Rebalancing Components
31
F7
Linear Probing Performance
32
F88 Worst Case Performance.
32
F9 More Exploration on Alternative Ways of Doing Inference
32
E.10
More Exploration on the Optimization Approach for V
33
F11
More Analysis on Richer Imbalance Contexts
34
E12 Computational Cost
17

---

Published as a conference paper at ICLR 2024
F.13
Balanced-Can Case Performance
34
E.14
In-distribution Performance...
35
F.15
Reverse-Distribution Performance
35
F.16
T-SNE Feature Visualization of Subpopulations
35
G Limitations and Future Explorations
36
18

---

Published as a conference paper at ICLR 2024
A BROADER RELATED WORK
A.1 INFORMATION THEORY-GUIDER OBJECTIVE DESIGN
Information Bottleneck. Information Bottleneck Theory (Tishby et al., 2000; Slonim & Tishby, 2000) in deep learning is a concept that has been extensively researched and developed over the years. It focuses on optimizing neural networks by maximizing the relevant information about the target while minimizing redundant data. The information bottleneck has gained widespread attention in recent years within the field of deep learning (Tishby & Zaslavsky, 2015; Goldfeld & Polyanskiy, 2020; Kawaguchi et al., 2023). For instance, gradient-based methods were employed in optimizing Deep Neural Network (DNN) to tackle the Information Bottleneck Lagrangian (Alemi al., 2017). This approach, known as the deep variational IB (VIB), enables the system to learn stochastic representation rules, showcasing enhanced generalization capabilities and robustness to adversarial examples. A similar objective was explored in Achille & Soatto (2018), where the emphasis was on promoting minimality, sufficiency, and disentanglement of representations. This disentanglement property was also harnessed for generative modeling purposes. leading to the development of the 3-variational autoencoder (Higgins et al., 2017).
Mutual Information Maximization: Mutual Information Maximization (InfoMax) princi- ple (Linsker, 1988) is a common training objective designed to enhance the information sharing between model outputs and target variables. This approach is particularly popular in self-supervised learning and representation learning and have demonstrated promising empirical results (Hjelm et al., 2019; Tschannen et al., 2020; Chen et al., 2020; 2021; 2023; Zhao et al., 2024). In general, their objective is to maximize the mutual information between representations from different views of the same image. For instance, in DecpInfoMax (Hjelm et al., 2019), 91 extracts overall features from the entire image, and captures local features from patches, where 91 and 92 are activations different layers of the same convolutional network. Extending this idea, Bachman et al. (2019) generate the two views by using different augmentations of the same image. Contrastive Multiview Coding (Tian al., 2020) extends the objective to incorporate multiple views, with each view corresponding to different image modality.
Comparision with our work. The information bottleneck and mutual information maximization techniques involve the mutual information between input, representation, and label variables, aiming to optimize the network for learning effective classifiers or generalizable representations. In contrast to these methods, our method, distinctively, models mutual information maximization (Def. 3.1) with the direct purpose of learning an effective data partition, which further serves the subpopulation harmonization.
A.2 DOMAIN GENERALIZATION
The objective of domain generalization is to extract knowledge that is invariant across diverse source domains and generalize it to novel, unseen target domains. A multitude of methods have emerged for domain generalization, broadly categorized into five groups: domain alignment, meta learning, domain hallucination, architecture-basee methods, and regularization-basee methods (Xu al., 2023). Domain alignment methods (Li et al., 2018; Zhao et al., 2020b; Grill et al., 2020; Ye et al., 2023) target on minimizing the discrepancies between source domains to learn domain-invariant features. Mleta-larniing-bsse based methods (Balaji et al., 2018; Li et al., 2019; Zhang et al., 2023a) enhance the generalizability of models domain shifts by partitioning training domains into distinct meta-train and meta-test domains. Through mntiioopti these methods simulate unseen domain shifts, thereby improving the models' adaptability to such shifts. Domain hallucination (Zhou et al., 2021; Xu et al., 2021) aims to augment training samples by transforming the original samples into specific unseen domains while preserving their underlying semantics. Architecture-baseee methods (Chattopadhyay et al., 2020; Chang et al., 2019) typically involve designing domain-speciff modules for different domains. During final testing, these methods aggregate results inferred from all source domains to achieve comprehensive outcome. Regularization-based domain generalization methods (Shi et al., 2020; Carlucci et al., 2019) involve learning general and universal features across domains through various regularization techniques.
Comparision with our work. The main distinctions between domain generalization and the sub- population imbalance problem discussed in our paper are as follows: (1) In domain generalization,
19

---

Published as a conference paper at ICLR 2024
domain labels are accessible during training, whereas subpopulation annotations not visible. (2) The goal of domain generalization is to exhibit strong generalization performance on unseen domains, while the problem of subpopulation imbalance aims for a comprehensive performance across all encountered subpopulations. Furthermore, the concepts of domain and subpopulation differ that domains are more akin image styles unrelated to semanties, while subpopulations represent kind of semantic abstraction distinct from the class dimension. Therefore, domain generalization methods typically aim to learn domain-invariant features. In contrast, our method separates the learning of subpopulations with different prediction mechanisms and then balancing them.
A.3 ALGORITHMIC FAIRNESS
Fairness is a critical and extensively studied aspect in algorithmic decision-makinn When dealing with biased data, algorithms tend to make decisions based on attributes that is sensitive or should be protected(e.g. g., race and gender), raising concerns about fairness (Kearns et al., 2018). Various approaches in algorithmic fairness aim to mitigate this issue by introducing fair constraints during the training procedure, such as demographic parity or equalized odds (Hardt al., 2016; Jiang et al., 2019; Calmon et al., 2017). Additionally, alternative fairness criteria include accuracy parity (Zhao et al., 2020a; Sagawa et al., 2020) (ensuring uniform accuracy across subgroups), small prediction variance (Li et al., 2020; 2021) (maintaining minimal prediction variations among subgroups) and small prediction loss for all subgroups (Zafar et al., 2019; Hashimoto et al., 2018). Some work introduces independence constraints to the objective to ensure that decisions do not rely sensitive attributes (Madras et al., 2018; Song et al., 2019).
Comparision with our work. In the algorithmic faimess problem, the protected attribute annotations are sometimes visible and sometimes not. When attribute annotations are invisible, the algorithmic fairness problem bears some similarities our problem. The key difference is that algorithmic faimesss often aims to protect specific sensitive attributes, preventing the model from using them in decision-making In our case, different subpopulations have different decision mechanisms, and we want to preserve features of minority subpopulations. In essence, while algorithmic faimeess tends to learn fewer features (excluding protected attributes), we aim to learn more features (safeguarding features of the disadvantaged subpopulations) Additionally, the evaluation metrics differ between the algorithmic fairness and our problem.
A.4 COMPARISION WITH CREAGER ET AL. (2021) AND LAHOTI ET AL.. (2020)
EIIL(Creeger et al., 2021) and ARL(Lahoti et al., 2020) infer subgroup membership based on the violation degree under the invariant learning principle or the stability of the loss space, and then perform group-invariant learning or reweighting. In comparison, SHE aims to optimize the predictive ability, specifically the interaction information, by partitioning data into subpopulations and concurrently rebalancing predictions among these subpopulations. Besides, these works assume that the subpopulation distribution causally independent the predicted target, which may not always hold in our scenario.
B SUPPLEMENTARY EQUATIONS
B. MUTUAL INFORMATION AND ENTROPY
Mutual Information. In probability theory and information theory, the mutual information (MI) of two random variables is measure of the mutual dependence between the two variables. The mutual information of two jointly random variables X (continuous) and Y (discrete) is defined as
I(xY)) Σ { p(z, v) log( p(x.v) )dx. yEy
(3)
p(x)p(y)
Relation to entropy. Mutual information can be equivalently expressed as
T(X;Y) H(Y) H(IIX) H(X) H(XY)),
20

---

Published as a conference paper at ICLR 2024
where H(Y) and (X) are (marginal) entropies, and H())X) and HX conditional entropies.
Entropy. The entropy H for discrete continuous variable X and continuous variable Y can be defined as follows, respectively.
H(X) - p(ar) log(p(x))ddrr
(5)
plv) H(Y) =- [((v) log(p(u))
(6)
yEy
Empirical Entropy. Here we provide the equation for the empirical entropy in Eq. (1):
HD(()= = Σ
TD(V) log TD(y),
yEy
Hb(Y|)(X,))) = Σ
(s)) Σ
log log
SES
ycy
where Tv) Σ 1(y v). Tp(s) 1 E 1(v(,, wi) s). mD(uls)
IN (o)=)))(((())))))-)))))) are empirical frequencies in the training set. En (((())))==))
B.2 NORMALIZED MUTUAL INFORMATION (NMI) SCORE
The Normalized Mutual Information (NMI) (Strehl & Ghosh, 2002; Li al., 2024; 2023) between the learned data partition v (for simplicity here we will use to denote v(X, v)X, and the ground truth subpopulation defined as:
21(S:v) NMI(S, H(S) + H())
(8)
which is a normalization of the Mutual Information (MI) score scale the results between 0 (no correlation) and (perfect correlation). We use the empirical score of NMI in our experiments.
B.3 RADEMACHER COMPLEXITY IN TIM. 3.3
The Rademacher complexity (Bartlett & Mendelson, 2002) is a concept used in statistical learning theory and machine learning measure the complexity of a class of functions. It provides way to quantify how well a function class can fit random noise, which in turn helps in understanding the capacity of the class overfit training data.
Consider a sample space X x y and a class of real-valued functions G' defined on X x y - R. For sample set D' = {(x1, v1 (x2, (22,U2), U2),... (æNN yN)} drawn i.i.d. from X x V, introduce Rademacher random variables σ1, 02 ... NN, which are independent and take values +1 or with equal probability 1/2.
The empirical Rademacher complexity of the function class with respect to the sample is defined as the expected value of the supremum (maximum) of the average sum of the product of g( v.) and o, over all functions in G. Mathematically, it's expressed as:
(G) E, sup Σ (x, 91) 1=1
9EG
The Rademacher complexity Ry (G) of the class G is the expectation of the empirical Rademacher complexity over all sample set of size N drawn from the space X x ),
Rx(G) Ep [Q)]
(10)
21

---

Published as a conference paper at ICLR 2024
B.4 MCDIARMID'S INEQUALITY
McDiarmid's Inequality (McDiarmid et al., 1989) is a concentration inequality which provides bounds the probability that function of independent random variables deviates significantly from its expected value.
Let 11 X Xn be independent random variables taking values in X. Consider function f: X satisfying the following bounded difference condition:
For each i e (1, n) and for any #1, Ina a' E Xn,
[ T--1 Cs I. v.) (T11 1x-1, thi Iit1 ;.,
xn)) s c
(11)
where is a constant.
Then, for any 0,
2 ?
Pr[ X2...., .,,) -
exp
(12)
-i=1
and similarly,
2
2
X2 -X) X,)]. -c]
exp
(13)
Σ
This inequality is particularly useful scenarios where one wishes to control the deviations of function of several independent variables from its expected value, especially in the context of machine learning and statistical learning theory.
C THEORETICAL PROOFS
C.1 PROOF PROP. 3.2
Proof of Prop. 3.2.
For data partition z' that satisfies v' X, Y, we have
(X::)>((X,))
T(X;Y)
(14)
For the optimal data partition v according to Def. 3.1, we have
I(X;Y; (X, z (xx;Y;
(15)
C.2 PROOF OF THM. 3.3
Lemma C.1. H(YIX) = infy Evyy - log (x)), where the infinum is achieved when ( p(yl
Proof for Lem. follows the same strategy as Proposition in Xu et (2020).
Proof of Lem. C.I.
inf Ev.yy - log (x)
inf E. E. log p(v12=)
(x))((()z)
(16)
=inf En (KTT HY)))))
22

---

Published as a conference paper at ICLR 2024
where the infimum is achieved when (x) -(yl=))
Corollary C.2.
1. H(Y) = y Bu.y log (0), where 0 denotes the empty input and the infimum is achieved when (0) = p(y).
2. H(YIX;u(X, Y)) = inl ExI Eses - -1(v(ar,)) = 8) log 12 (x)), where the infimum is achieved when (a:) = [yyl2::1
3. H(Y Y)) inf, Ez Eoes (1(( s) log 5% where the infimum is achieved when (0) = p(uls).
4. HD(Y| Y)) = inf, N achieved when (0) Σi. (= ΣN (=(),
i-1 Eves - (1((2,)))) s) log (0), where the infimum is
The proof of Cor. C.2 is similar to proof of Lem. C.I.
Proof Thm. 3.3. Define a function T for any training set D:
T(D)=
,1;D) B)1
(17)
According to Lem. C.I and Cor. C.2, we have:
N
T(D) =
M
Σ
1(((,;)))) = 8) 1t HD(Y|u(X,Y)) BI
N
SES
N N
-1((; Y))
Σ Σ (V((,,,))) vi) = =) log f:th - (a)))
i-1 #ES
= 1-innf inf E m,y
Σ
-1(v(x,)) v) = s) log f (x) + inf En. Σ -1((((0))) = $) log (og
RES
8ES
+ inf
Σ
1(((,,)))) = x) log (mx) inf
Σ
(1(((04 = s) log 1 (0)
1=1 SES
i=1 sES
N
$ supp E Σ
-1 -1(v(x,)) = log (x)
Σ
Σ
((v(2, v:) =s) - -log (x:;
N
#ES
i-1 #ES
Ee..y
Σ
((v(x, = s) - log (0) +
1 Σ Σ
[(( v.) = =) -log 1 (0)
N
sES
1=1 EE
(18)
We define Q(D) as
Q(D) := sup
E Σ
-1(v(x,)) = 8) log (x)
Σ [(((,,)))) =) llog (xxx)
N
AES
i=1 BES
E Ex,y Σ 1(v(x,))= = log (0) + N
Σ
1(v(0t))1) =s) -log f (0)
SES
1=1 SES
(19)
23

---

Published as a conference paper at ICLR 2024
Let D and D' be two identical data sets except that the j-th data point is different.
Q(D)-Q(D)
1 N
Sup([Ex,
Σ -1(((xx,) 8) f"
Σ
Σ
(v(x;;)))) 8) -log S (x:))
N
SES
Em.y
Σ
1(v(x,)) y) = 8) log (0)
Σ 1(v(x;1))) = =) - log (0)l
N
#ES
SES
- -E Σ
-1(((, v) = 8) log (x)
Σ
[u(az 1/) = =8) -log S (a²)
N
SES
N
Emyy Σ
I(v(x,)) 8) log (0)
Σ
Σ
1(v(at1 = =s) - llo f.f (0)1)
#ES
SES
sup f
1
log 2
333 N
1
log Uu( mp (0) N
log 2o
1 M u(e, v/) (a/)) N log Syaee 2/) (0))
N
2m
N
(20)
According McDiarmid's inequality (McDiarmid al., 1989) (also introduced in Appx. B.4),
V8 (0, 1) we have:
Q(D) EDQ(D)+ + VN
V22 8
(21)
with probability at least - 6. For simplicity of writing, we define
A = Exzy Σ -1(v(z, = 5) log (x)) Ex,,y Σ
[(v(x,))) s) (0)
SES
SES
Ap
1
Σ
1V((,, vi) =) log 1 (x+:)
1
Σ
1(v(z,, v:) =)) -loo f (0)
N
N
1= SES
1= 1
(22)
24

---

Published as a conference paper at ICLR 2024
So Q(D) sup, IA - Apl. and we have:
EdQ(D) Ep spsup / A Apl
=Epsup E À Apl
/
Ep sup Evr Ad - Aol
Ep, n' sup A Apl
N
ED sup
Σ
oil
log pii ((a).e)) (x/)) + log ful v(at (0) log full Sula,,
log f (m (0))
N
ED, sup
1
o log f (a) (x;)) + ED,, sup
1
o log fih v(a) 2/) (x))
N
N
1
pn (( (2)) (0) o, log
+ E D. sup
a log ( w.) (0)l + ED' sup
N
i=1
=DED, sup
a, log (. (xx)) 2 D sup
o log f (m (0)
N
N
1=1
4ED,r sup
σy log (.. v.) (#:))
N
1=1
S 4
Σ ED, sup
o, log f (x:)
N
SES
1=1
4KR)(G),
(23)
where og, 1, .. the Rademacher variable that is uniformly sampled {-1,+1)
Combin-
ing Eq, (17), Eq. (18), Eq, (21), Eq. (23), Vo E 1), we have:
((XX, ; B)|
- 2og 5 4KK (G(G))
(24)
with probability at least 6.
C.3 PROOF OF THM. 3.4
Proof of Thm. 3.4.
If 2, can perfectly fit the data distribution of given subpopulation i.e., p(x, yls) oc (((( Let
p(x, y|s) w (xx) w We can get
p(æ, vls)
p(y(z;s)
2oftmax(2."
(25)
E, p( v'1s)
25

---

Published as a conference paper at ICLR 2024
For the subpopulation- cetpopulltioo-bblaneeddd balanced distribution, (sat(s) 1 K. Vs E S. We have
Pbat (bnt(2, v)
Σ
plz, pls)p
#ES
Σ
p(a,
yls)
(26)
Σ
e (t)
K SES
σ =<(x)
And we have
p(yl2) = p( u)
softmax( (x))
(27)
Σg ((,x1/))
C.4 AN EXTENSION OF TIM. 3.4
We can extend Thm. 3.4 to make SHE handle an arbitrary specified target distribution over the latent subpopulation Prest ( y) Eees Ptest())p(2: yls), in the form of weighted variant of LogSumExp.
Proposition C.3. Supposing that for any subpopulation s E S, 2, can perfectly fit the data dis- tribution of a given subpopulation i.e., p(a, y| e²((a) then log Eoes Ptost(s) can perfectly fi the subpopulation-balunn balanced overall distribution, i.e., (x v) e (a))
The proof shares the same spirit as the proof of Thm. 3.4.
D ADDITIONAL DISCUSSIONS OF EQ. (2)
Here we would like to explain more about the intuitive understanding of Eq. (2). For ease of reading, here we restate Eq. (2) as follows:
N
C=
1
Via log lo - (Y -8 Div(x,).
SES
With respect to V, the first term of Eq. (2) increases the weight of the subpopulation that predicts accurately for each sample. The second term makes the classes each subpopulation balanced as possible, which prevents the subpopulation from collapsing to the prediction target. And the third term prevents each subpopulation from collapsing to exactly the same prediction and accelerate the optimization. The second and third terms somewhat prevent V from falling into the trivial solutions.
E THE EFFICIENCY AND SCALABILITY OF OPTIMIZING V
As stated in Sec. 3.4, the subpopulation-wee weight matrix V has size of N K. To improve the efficiency of optimizing V and to increase the scalability of our method on large datasets (where N grows larger), we design a batch-specific update approach for V. This means that during each iteration, only the elements of V corresponding to the samples in each mini-batch (i.e., BatchSize x K elements) are updated, while the remaining elements are kept fixed. This approach significantly reduces the number of updated elements in V to BatchSSiz K, which is far smaller than the number parameters in a modern neural network, allowing for scalability even when dealing with large datasets.
26

---

Published as a conference paper at ICLR 2024
F DETAILED SUPPLEMENT FOR EXPERIMENTS
F.1 SUPPLEMENTAL DESCRIPTION OF THE EXPERIMENTAL SETUP
F.1.1 TOY MOTIVATING EXAMPLE (FIG. 2)
In Fig. 2, we visualize a toy motivating example whose prediction goal is to distinguish between circles (semi-transparent) and triangles nonntrrasppaennt. For training data, they are sampled from both subpopulation I (blue) and subpopulation 2 (red), and the training samples of subpopulation are much less than those of subpopulation I, i.e., under subpopulation imbalance. About the test set, it is sampled equally from both subpopulations, i.e., under subpopulation balance. Specifically, each class each subpopulation sampled from the following normal distributions: N[[1.3], [0.2, [0.2, 0.2]) (subpopulation 1, class 1), [[[,, 31, [0.2, 0.2]) (subpopulation I, class 2), N[[115,, [0.2, 0.2]) (subpopulation 2, class 1), and [[[1.5, 5, 2], [0.2, 0.2]) (subpopulation 2, class 2), respectively. The sample size of the training set is 200, 200, 10, 10 in the corresponding order. The sample size of the test set for each normal distribution is 100. We use a two-layer MLP with 5 hidden neurons as the model for the toy study. The batch size is set to 512. The toy models are trained using SGD with momentum of 0.9. We train the models for 60 epochs with initial learning rate 0.2.
F.1.2 TRAINING UNDER SUBPOPULATION IMBALANCE COUPLED WITH CLASS IMBALANCE
For COCO, we conduct training set with both subpopulation imbalance and class imbalance and both balanced test following the GLT-protocol in Tang al. (2022). For CIFAR and miredImageNet, shuffle the subcategories in the dataset randomly and then sample them in an imbalanced manner, so that they imbalanced at both the category and subpopulation levels.
F.1.3 TRAINING UNDER SPURIOUS CORRELATIONS
CelebA (Liu al., 2015) and Waterbirds (Sagawa al., 2019) are two datasets that have been widely used to benchmark the robustness of machine learning algorithms to spurious correlations. In the CelebA dataset, there is high correlation between gender (male, female) and hair color [blond, dark), meaning that the feature gender might be used as a proxy to predict the hair color. In Waterbirds, there high correlation between y = (land bird, water bird} and background {land, water). The baselines for comparison include GDRO (Sagawa et al., 2019), LfF (Nam al., 2020), SD (Pezeshki al., 2021), JTT (Liu al., 2021), CIM (Taghanaki al., 2021), and MaskTune (Taghanaki et al., 2022).
We strictly follow the experimental setup in Taghanaki et al. (2022). For the Waterbirds dataset, we train ImageNet pre-trained ResNet50 with batch size of 128 for 100 epochs decaying the learning rate by 0.1 after every 30 epochs. For the CelebA dataset, we train an ImageNet pre-trained ResNet50 with batch size of 512 for 20 epochs with a learning rate of 10 For our SHE, the pre-defined K set to 2. When the group information is unknown for all methods, we utilize the overall accuracy on a validation set that shares the same distribution as the training set as the selection metric.
1.1.4 LINEAR PROBING
We train linear classifier on a frozen pre-trained backbone and measure the quality of the represen- tation through the test accuracy. To eliminate the effect of the skewed distribution in the fine-tuning phase, the classifier is trained on a subpopulation- ballaace dataset. Specifically, the performance of the classifier is reported on the basis of pre-trained representations for different amounts of data, including full-shot, 100-shot, and 50-shot. In the fine-tuning phase, we train the linear classifier for 500 epochs with SGD of momentum 0.7 and weight decay 0.0005. The batch size set 1000. The learning rate decays exponentially from 10 to 10 The loss function is set to the ordinary cross-entropy loss.
27

---

Published as a conference paper at ICLR 2024
Apple
Banana
Subpopulation
Subpopulation
Subpopulation
Subpapalation
Figure 5: Visualization learned subpopulations in COCO. To case the visualization and analysis, set the number of subpopulations K Then, after training, randomly selected 8 images from the two subpopulations we learned about the classes "apple" and "banana" COCO.
Sahpoppun 2 Snppopliii
Apple
Banana
Pizza
Hotdog
Person
Figure 6: More visualization of learned subpopulations in COCO. To case the visualization and analysis, we set the number of subpopulations as K = 2. Then, after training, we randomly selected 9 images from the two subpopulations learned about the classes "apple", "banana", "pizza", "hotdog", and "person" in COCO.
F.2 MORE VISCALIZATION OF LEARNED SUBPOPULATIONS
F.2.1
COCO
We present the visualizations of the subpopulations discovered by SHE using the COCO dataset, as depicted in Fig. 5. We also make reference to these visualizations in Sec. 3.3. Here we show more visualizations in Fig. 6. It can be that samples from different subpopulations of the same class have obvious semantic differences, further demonstrating that our SHE able to discover meaningful hidden subpopulation structures in the training data.
F.2.2 WATERBIRDS
We show some visualizations of the subpopulations learned by SHE on the Waterbirds dataset in Fig. In the Waterbirds dataset, group annotations are constructed based on classes (land bird/water bird) and backgrounds (land/water). The spurious correlation is exhibited by that most of the birds on the land background are land birds and most of the birds on the water background are water birds. If we distinguish subpopulation based on their background, then category imbalance in each subpopulation will still lead the model learn spurious correlations. Fortunately, however, SHE divide the training data into two meaningful subpopulations: 1) data with spurious correlations and 2) data without spurious correlations, which actually different from the prior group annotations, as shown in Fig. 7. Such a result may be due to the second term HD (Y (YlV) of Eq. (2), which increases the entropy of the classes in each subpopulation to make the classes in each subpopulation as balanced
28

---

Published as a conference paper at ICLR 2024
Land Bird
Water Bird
Land Bird
Water Bird
Subpopulation
Subpopulation 2
Figure 7: Visualization of learned subpopulations in Waterbirds. We randomly selected 16 images from the two subpopulations we learned about the classes "land bird", and "water bird" Waterbirds.
Table 7: Per-split accuracies on COCO (Mean + std).
Method
Many
Medium
Few
Overall
ERM
67.21 + 0.24%
52.22 + 0.17%
37.10 i 0.42%
62.52 0.32%
PaCO
67.45 + 0.31%
53.33 + 0.21%
36.23 0.28%
62.59 0.24%
BCL
66.89 + 0.31%
37.67 + 0.24%
62.83 0.42%
IFL
67.71 0.13%
52.17 0.24%
36.82 O. 18%
62.57 0.15%
DB
67.35 + 0.32%
52.11 + 0.25%
37.47 + 0.36%
62.72 + 0.48%
TDE
66.32 + 0.22%
53.23 + 0.28%
37.02 0.33%
62.64 0.27%
ETF-DR
67.93 + 0.17%
51.34 + 0.22%
37.59 + 0.12%
62.45 + 0.37%
LfF
66.74 + 0.34%
52.34 0.26%
36.01 + 0.38%
62.06 0.83%
Focal
66.23 + 0.42%
22.11.0220
35.79 + 0.36%
61.67 + 0.53%
EIIL
66.87 + 0.18%
52.79 + 0.32%
37.06 + 0.28%
62.61 0.33%
ARL
67.32 + 0.25%
53.34 0.17%
37.18 + 0.24%
62.48 + 0.22%
GRASP
67.13 + 0.11%
53.26 + 0.16%
37.29 + 0.32%
62.73 0.25%
JTT
66.93 + 0.26%
51.24 + 0.35%
36.48 0.27%
62.32 0.75%
MaskTune
64.48 + 0.31%
50.11 + 0.35%
33.27 E 0.19%
60.23 0.73%
SHE
67.71 + 0.32%
53.50 0.26%
42.09 + 0.28%
64.56 0.24%
as possible. Such two subpopulations plus further subpopulation rebalancing can effectively prevent the model from relying on spuriously correlated features for prediction.
F.3 COMPLETE RESULTS OF TAB. 3
We provide the complete experimental results (Mean Std) of all baselines of Tab. 3 in Tab. 7.
F.4 COMPLETE RESULS OF TAB. 6
We provide the complete experimental results (Mean + Std) of Tab. 6 in Tab. 8.
F.5 MORE RESULTS ON VARYING THE LATENT SUBPOPULATION NUMBER K
In Fig. 4(a), we show ablation study on K on COCO. Here we give the complete experimental results more dataset, i.e., CIFAR-IR 100. CIFAR-IR50. CIFAR-IR20, and tiredImageNet. in Fig. When K 1, Eq. (2) degencrates to the cross-entropy loss, and our SHE degenerates to the ERM performance. When K SHE has significant improvement ERM and shows some robustness to the value of in general.
29

---

Published as a conference paper at ICLR 2024
Table LoRA fine-tuning under pre-trained models on COCO (Mean i std).
Method
CLIP
ALIGN
AICCLIP
Zero-shot
76.59 + 0.00%
78.45 0.00%
82.55 t 0.00%
ERM
84.32 + 0.14%
83.38 + 0.12%
84.85 + 0.07%
PaCO
84.38 + 0.11%
83.54 0.15%
85.06 + 0. 0.16
BCL
84.48 + 0.06%
83.32 + 0.11%
85.06 + 0.05%
IFL
84.55 + 0.06%
83,40 + 0.05%
84.83 : 0.06%
DB
84.46 : 0.11%
83.14 0.13%
84.14 0.40%
TDE
84.37 0.10%
83.42 + 0.05%
8 0.19%
ETF-DR
84.26 + 0.15%
83.36 0.09%
84.78 + 0.05%
LfF
84.17 + 0.05%
83.12 + 0.07%
84.20 + 0.02%
Focal
83.87 + 0.10%
82.77 0.12%
83.84 + 0.14%
EIIL
84.23 + 0.12%
83.21 + 0.07%
84.34 + 0.05%
ARL
84.02 0.08%
82.97 0.17%
84.07 + 0.07%
GRASP
84.32 + 0.16%
83.14 + 0.04%
84.44 + 0.13%
JTT
84.45 + 0.09%
83.09 0.02%
84.40 0.15%
MaskTune
83.27 + 0.10%
82.54 + 0.13%
83.52 + 0.20%
SHE
85,39 0.06%
84.23 + 0.06 %
85.69 0.11%
0.55
0.575
0.570
Accuracy
0.54
Accuracy
0.565
SHE
SHE
0.560
ERM
ERM
0.53
0.555
0.550
0.52
0.545
2
2
K
(a) CIFAR-IRIO0
(b) CIFAR-IRSO
0.505
0.61
0.500
Accuracy
Accuracy
0.495
SHE
SHE
0.60
ERM
0.490
ERM
0.485
0.59
0.480
2
2
K
K
(c) CIFAR-IR20
(d) tiredImageNet
Figure 8: Performance of SHE and ERM with varying subpopulation number K on CIFAR-IR100. CIFAR-IRSO, CIFAR-IR20, and tiredImageNet.
30

---

Published as a conference paper at ICLR 2024
Table 9: A series of baselines by combining different strategy permutations
Method abbreviation
Clustering
Rebalancing
Training strategy Sequential
K-means
OT
Prediction
Resample
Reweight
Simultaneous
K-1
K-2
K-3
K-4
0-1
0-2
/
0-3
044
P-
P-2
P-3
P4
Table 10: Comparison with more clustering and rebalancing components. Bold indicates superior results. The meaning of some method abbreviations can be found in Tab. 9.
Method
COCO
CIFAR-IR
CIFAR-IR50
CIFAR-IR20
ERM
62.52 + 0.32%
52.49 + 0.27%
55.20 + 0.41 %
58.92 + 0.62%
K-1
61.59 + 0.41 %
51.43 + 0.34%
54.56 + 0.47%
57.97 + 0.35%
K-2
62.63 + 0.35%
52.57 + 0.40%
55.03 + 0.18%
59.03 + 0.28%
K-3
61.77 + 0.57%
51.71 0.26%
54.29 + 0.37%
58.21 i 0.32%
K-4
62.48 t 0.27%
52.52 + 0.34%
55.27 + 0.27%
58.81 # 0.27%
0-1
62.56 + 0.39%
5 0.44%
55.13 + 0.42%
58.85 + 0.20%
0-2
62.49 + 0.21%
+ 0.42%
55.22 + 0.21% %
58.85 + 0.49%
0-3
62.54 + 0.32%
52.47 + 0.35%
55.15 + 0.38%
58.90 + 0.51%
0-4
62.58 + 0.36%
52.43 + 0.30%
55.08 + 0.58%
58.99 + 0.44%
P-1
61.32 + 0.28%
51.27 + 0.37%
54.31 + 0.42%
57.62 + 0.31%
P-2
62.47 + 0.46%
52.25 + 0.43%
54.96 + 0.49%
58.71 + 0.26%
P-3
61.58 + 0.14%
51.53 + 0.29%
54.00 + 0.18%
58.13 + 0.55%
P-4
62.30 + 0.38%
52.47 + 0.33%
55.25 + 0.26%
58.57 1 0.33%
SHE-RS
64.03 + 0.27%
53.97 + 0.38%
56. + 0.50%
60.25 : 0.54%
SHE-RW
63.82 + 0.22%
53.90 + 0.41%
56.05 + 0.27%
60.19 - 0.27%
SHE
64.56 + 0.24%
54.52 I 0.35%
56.87 + 0.17
60.72 + 0.41 %
F.6 COMPARISON WITH MORE CLUSTERING AND REBALANCING COMPONENTS
To further demonstrate the superiority of SHE, we compare it with a series of clustering and rebalanc- ing strategies. The clustering strategies include K-means, optimal transport clustering, and direct using the ERM predictions. The rebalancing strategies include reweighting and resampling. And the training strategies include: the simultaneous strategy: at each epoch training. the clustering results and rebalancing weights are updated: 2) the sequential strategy: first train an ERM model and obtain clustering results, then conduct rebalancing to retrain a model based on the clustering results. We construct a series of baselines by combining these strategy permutations, which are presented prediction by Thm. 3.4, we
Tab. To confirm the effectiveness of the subpopulation-balauee also construct two variants of SHE: SHE-RS (dynamic resampling based on the learned subpopu- lations during training), and SHE-RS RS (dynamic reweighting based on the learned subpopulations during training). We show the performance of all these baselines and SHE on COCO, CIFAR-RI R100, CIFAR-R50, and CIFAR-R22 in Tab. 10. Our method still achieves the best results very clearly, demonstrating the superiority of our method for subpopulation discovery and rebalancing.
31

---

Published as a conference paper at ICLR 2024
0.70
0.50
SHE ERM
SHE
0.65
0.45
ERM
Accuracy
Accuracy
0.60
0.40
0.55
0.35
0.50
0.30
0.45
0.25
full-shut
100-shol
50-shol
full-shot
100-shol
50-shol
Superclasses
Subclasses
Figure 9: Linear probing performance of SHE and ERM on CIAFR- 10000 superclasses (left) and subclasses (right) under different shot settings.
Table 11: Worst-case performance on CIFAR-IR 10100/50/20.
Mathod
FRM
PaCO
BCI.
IFI.
DB
TDE
ETF-DR
1.F
Focal
JTT
MT
SHE
CIFAR-IRI00
22.34%
23.49%
23.83%
22.17%
22.87%
22.53%
22.96%
20.74%
19.45%
21.38%
19.78%
27.48% 31.24%
CIFAR-IRS0
26.47%
27.17%
27.58%
26.71%
26.93%
26.68%
27.43%
23.67%
23.16%
24.58%
23.17%
CIFAR-IR20
37.10%
38.14% 38.07%
27.45%
37.77%
37.29%
37.93%
34.73% 34.25% 35.28%
34.64%
41.19%
Table 12: Compared with other alternative ways of doing inference.
Method
COCO
CIFAR-IR 100
CIFAR-IR50
CIFAR-IR20
ERM
62.52 0.32%
52.49 0.27%
55.20
0.41%
58.92 t 0.62%
SHEgm. SHESimAvg
62.82 + 0.27%
52.63 t 0.22%
55.36 t 0.37%
59.21 t 0.48%
64.54 0.17%
54.44 0.41%
56.78
0.28%
60.66 t 0.32%
SHE
64,56 + 0.24%
54.52 + 0.35 %
56.87 + 0.17% %
60.72 t 0.41 %
F.7 LINEAR PROBING PERFORMANCE
To quantitatively evaluate the representation quantity of different methods, conduct linear probing experiments on CIFAR-IRA 100 following the literature self-supervised learning (Chen et al., 2020; He al., 2020). To eliminate the subpopulation imbalance effect, the linear classifier is trained on balanced dataset on top of the fixed feature extractor. In Fig. 9, we show the linear probing performance of both superclasses and subclasses on CIFAR-IR 100 under different shots. As can be seen, our SHE consistently exceeds the ERM baseline for all settings, especially for the linear probing performance of fine-grained classes with the improvement 1313, 1.81%, %,, and 2.45% on full-shot, 100-shot, and 50-shot. This indicates that our SHE actually captures better and more generalized representations.
F.8 WORST CASE PERFORMANCE.
To further validate the efficacy of our SHE in enhancing the learning capability of rare samples, present the worst case performance on CIFAR-IRION 100, CIFAR-IR50, and CIFAR-IR20 in Tab. 11. evident that our method has achieved significantly superior results compared other baselines. further substantiating its remarkable effectiveness in mitigating the issue of subpopulation imbalance.
F.9 MORE EXPLORATION ON ALTERNATIVE WAYS OF DOING INFERENCE
Regarding alternative ways of doing inference, we validate some simple variants like applying group- invariant learning the learned subpopulation (marked as 'SHE ') or just averaging the logits across f. (marked as "SHEsimAvg" "). As shown in Tab. 12, SHE significantly outperforms SHEpmI. while achieving comparable results with SHESimAAg- This is because, applying the Normalized
32

---

Published as a conference paper at ICLR 2024
Table 13: Compatision between w/ or w/o LA COCO where both class imbalance and subpopula- tion imbalance co-exist.
Method
Acc
ERM SHE
63.57 0.34% 65.13 t 0.22 %
LA
66.47 + 0.27%
SHE LA
68.11 0.27 %
Table 14: Compared with other alternative ways of optimizing V.
Method
COCO
CIFAR-IR 100
CIFAR-IRSA
CIFAR-IR20
ERM
62.52 + 0.32%
52.49 t 0.27%
55.20 t 0.41%
58.92 + 0.62%
SHE model based V
63.22 0.13%
53.28 0.36%
55.81 0.17%
59.72 t 0.38%
SHEgM SHE
64.52 + 0.31%
54.47 + 0.26%
56.88 + 0.22%
60.65 + 0.33%
64.56 0.24%
54.52 + 0.35 %
56.87 + 0.17%
60.72 + 0.41 %
Table 15: Performance on datasets for spurious correlations.
Group Info
Method
CelebA
Waterbirds
(Train / Val)
GDRO
Yes Yes
88.3% / 99.8%
91.4% 93.5%
SHE
Yes / Yes
88.4% 99.3%
91.6% 993.2%
ERM
No/ No
47.2% / 95.6%
74.9% 98.1 18
SHE
No/ No
77.7% / 92.0%
82.0% 91.3%
SHEW//DRO
No No
77.9% 99177%
81.9% 91.3%
Weighted Geometric Mean (NWGM) approximation (Baldi Sadowski, 2013; Xu et al., 2015).
LogSumExp can be approximated as simple summation, which is equivalent to simple averaging the logits. Here, we use LogSumExp because of its advantage of being more theoretically rigorous with its numerical stability.
F.10 MORE EXPLORATION ON THE OPTIMIZATION APPROACH FOR V
In terms of alternative approaches for optimizing V. we examined different variations SHE. Firstly, we utilized 2-layer MLP learn V from image features, referred to SHEmodel based V. Sec- ondly, we employed an EM-style approach to alternately learn V and f., referred to as SHEEM. As indicated in Table 14, SHEmodel based V exhibits noticeable performance degradation compared to SHE. This can be attributed to the fact that in Definition 3.1 is dependent on both the input and the label y, whereas SHEmodel based V can only learn the data partition from æ. On the other hand, SHEEM demonstrates comparable results with SHE, yer SHE is simpler and superior, thus confirming the effectiveness of the proposed optimization approach.
Since SHE based v. which learns V solely from X, does not satisfy our formulation, we similarly construct SHEmoodel based V from y to learn V from both X and Y. The experimental results are presented in Tab. 16. SHEmooel based V from X.Y outperforms the variant that solely learn V from X, as aligns better with our formulation. However, due to the batch-wise updates of the network learning V under both variants, is challenging to accurately compute the subpopulation allocation for samples not in the batch, yielding the inaccuracy when calculating the entropy term in Eq. (2) and thus the performance degradation compared to SHE.
33

---

Published as a conference paper at ICLR 2024
Table 16: Performance of two variants of model-based methods to learn V.
Method
CIFAR-IRI00
CIFAR-IR50
CIFAR-IR20
ERM
52.49%
55.20%
58.92%
SHE based V from X
53.28%
55.81%
59.72%
SHemmodl based V from X,Y
54.08%
56.23%
60. 60.14%
SHE
54.52%
56.87%
60.72%
F.11 MORE ANALYSIS ON RICHER IMBALANCE CONTEXTS
SHE is primarily designed to address subpopulation imbalance issues and optimize overall per- formance under subpopulation balanced distributions. However, does not specifically focus on class imbalance and worst-case performance. Therefore, we combined SHE with LA or GDRO in the presence of class imbalance and spurious correlation, as discussed Sec. 4.3. this context, we present additional results of using SHE alone when dealing with class imbalance and spurious correlation scttings in Tab. 13 and Tab. 15, respectively.
From Tab. 13, it is evident that our SHE achieves significant improvements whether applied on top of ERM or LA. The results in Tab. 15 demonstrate that SHE achieves comparable performance to SHEw/ GDRO. We report the results of SHEw/ GDRO in Sec. 4.3 to maintain consistency with the objective of the spurious correlation task.
F.12 COMPUTATIONAL COST
Considering the computational cost, SHE (according Eq. (2)) and Appx. E)incurs additional computational overhead compared to ERM due to: 1) Weighted summation of the cross-entropy loss, 2) Computation of empirical entropy and the regularization term, and 3) Updating a matrix of size BatchSize x K. The computational expenses for these three components are all far smaller than the training cost of a modern deep neural network. Tab. 17 presents comparison of training duration for different methods.
Table 17: Training time of 200 epochs on CIFARI00-IRION
Method
Training time (Minutes)
A ERM
ERM
89
PaCo
104
+16.8%
BCL
110
+23.6%
DB
102
+14.6%
TDE
95
+6.7%
IFL
142
+60%
JTT
101
+13.5%
LfF
100
+12.4%
MaskTune
105
+17.9%
SHE
98
+11.1%
F.13 BALANCED-CASE PERFORMANCE
We provide the result of SHE and ERM under the subpopulation balanced senario (both train and test) in Tab. 18. Our SHE get slightly better performance even in the balanced case.
Table 18: Balanced case performance.
Method
CIFAR
tieredlmageNett
ERM
74.32% 74.75%
68.26%
SHE
68.83%
34

---

Published as a conference paper at ICLR 2024
10
10
-10
-30
-10
10
-30
-10
(a) ERM
(b) SHE
Figure 10: T-SNE visualization of 5 subpopulations whin the same class CIFARIO0. Models trained on CIFAR-IR20.
F.14 IN-DISTRIBUTION PERFORMANCE
We present results in Tab. 19 under the scenario where the test shares the same distribution as the training set, indicating that SHE achieves comparable or slightly better results than ERM in this situation.
Table 19: Performance when the test set shares the same distribution as the training set.
Method
CIFAR-IRIO0
CIFAR-IR50
CIFAR-IR20
ERM
71.21% 71.07%
70.51% 70.72%
68.94%
SHE
69.25%
F.15 DEVERSE-DIISTRRUUTIO PERFORMANCE
Although Prop. C.3 tells the direction of handling arbitrary specified target distribution. a challenge in application lies in correctly aligning the learned subpopulations with their corresponding test distributions. To validate the effectiveness of Proposition B.3, construct experimental scenario where the test distribution the reverse imbalanced distribution the training set. Therefore, we can correspondingly reverse-sort the learned subpopulations by sample size align with the test distribution. Tab, 20 illustrates the results.
Table 20: CIFAR-IR
Performance.
Method
CIFAR-IR50
CIFAR-IR20
ERM SHE
38.37%
42.96%
50.70%
43.53%
46.74%
53.82%
SHE distribution
45.64%
48.02%
54.64%
F.16 T-SNE FEATURE VISUALIZATION OF SUBPOPULATIONS
In Fig. 10. illustrate the t-SNE features of five subpopulations within the same category CIFAR100. This demonstrates that SHE is effective in preventing minority subpopulations( (oraaiiiioooooonneeeee orange and blue) from being overshadowed by others.
35

---

Published as a conference paper at ICLR 2024
G LIMITATIONS AND FUTURE EXPLORATIONS
Similar to some clustering methods (e.g., K-means), our method relies on predefined number of subpopulations K. Choosing an inappropriate value of K may lead overfitting underfitting of the subpopulation structure and affect the performance and interpretability of the method. We study the scenario in this paper where all subpopulation annotations are invisible. When part of the subpopulation annotations are visible, how to further improve the performance and robustness of machine learning algorithms by leveraging partially labeled data to refine the subpopulation structure needs to be further explored.
36


Reviews:
[{"review_id": "mNrN0SHF6E", "reviewer": "Reviewer_acbx", "summary": "This paper aims to solve the subpopulation imbalance problem. The authors propose a new method named Scatter and HarmonizE (SHE), which discovers and balances the latent subpopulation in training data. Specifically, it builds on the principle of optimal data partition from information theory, which approximately uncovers the hiddle subpopulations and assigns data to subpopulations. Then, it achieves subpopulation-balanced predictions by simply applying a LogSumExp operation. Theoretical analyses are provided to support the validity of the method. Finally, experimental results illustrate the superiority of the proposed method SHE.", "questions": "1. In Table 1, what is the formal definition of imbalance ratio IR? I have carefully checked this paper and have not found it.\n\n2. From the perspective of computational cost, what about the proposed method SHE against other baselines?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "First of all, I have to admit that I am not an expert in the area of subpopulation imbalance, and may miss some related work.\n## Originality\n* To my knowledge, the data partition method for subpopulation recovery based on information theory is somewhat novel although these techniques have been widely used in other sub-fields of machine learning.\n## Quality\n* The proposed method is reasonable. Extensive experiments illustrate its superiority. Besides, theoretical results are also provided to support the validity of the method. \n## Clarity\n* Overall, this paper is well-written and the motivation is very clarified.\n## Significance\n* The proposed method can contribute to the community of subpopulation imbalance.", "weaknesses": "## Originality\n* There are also many works inspired by the information theory to guide the design of training objectives in machine learning. More discussions can be added.\n\n## Quality & Clarity\n* The proposed method SHE heavily depends on the number of subpopulations $K$, which is unknown in practice.\n* For the equation at the end of Section 3.1, the goal to minimize the error rate of a specific test dataset is not proper because it should be the expected error rate w.r.t. the distribution, and the one for a specific test dataset is only its unbiased estimator.\n\n## Significance\n* The proposed method may have little effect on other sub-fields of machine learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "958BB8H0D8", "reviewer": "Reviewer_TvvJ", "summary": "This work addresses the subpopulation imbalance problem where the training data consists of multiple subpopulations and their proportion is imbalanced. A novel approach, referred to as scatter and harmonize, to identifying the subpopulations and minimizing risk for each subpopulation is proposed. The authors provide a theoretical analysis of the proposed approach and demonstrate its utility through extensive experiments.", "questions": "- I might have missed the detail, but I am not sure if the (optimal) data partition approach is completely novel. Could the authors clarify it? If not, could the authors kindly provide what approaches have been proposed? (maybe for some other problems)\n- Even if the authors focused on the subpopulation imbalance problem, domain generalization and subpopulation shift problems seem highly relevant to this work. Could the authors kindly clarify the relationship between the existing methodologies to tackle the domain generalization and subpopulation shift and the subpopulation imbalance problem? Also, this work is somewhat relevant to algorithmic fairness as well. It might be helpful for future readers to relate the problems conveniently.\n- For Thm 3.3, could the authors elaborate on why minimizing $\\hat{\\mathcal{R}}$ results in maximizing $I(X;Y,\\nu(X,Y))$? It seems like it is true asymptotically but not sure with finite samples.\n- Is it possible to establish an inequality between the (empirical) risk of SHE and that of ERM under the subpopulation-balanced distribution?\n- This work seems to be similar to [Lahoti et al. (2020)](https://proceedings.neurips.cc/paper/2020/hash/07fc15c9d169ee48573edd749d25945d-Abstract.html), which ensures fairness with respect to maximal heterogeneity. Might be interesting to investigate the differences and similarities. \n- Could the authors kindly explain why no method other than the proposed one outperforms the ERM across all datasets?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Considering a practically important problem with clear motivation.\n- Well-written and easy to follow.\n- Supported by extensive experimental results.", "weaknesses": "- I find that this study is relevant to various problems/applications. It could be informative if the authors can clarify the relationship to the existing work, including domain generalization, algorithmic fairness, or such.\n- Some details could be improved for completeness: e.g., the subpopulation imbalance problem makes sense only if $p(\\boldsymbol{x},y|s)$ differs by subpopulation $s$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FpqSsZYuNA", "reviewer": "Reviewer_FihB", "summary": "The paper studied the subpopulation imbalance problem in machine learning. The setting can be described as follows: let the dataset being consist of groups $S$, and under uniform sampling, the probability $\\Pr(S=s)$ varies for different realizations of $s$. If the functions $f_s: \\mathcal{X}\\rightarrow \\mathcal{Y}$ are quite different across different groups, the machine learning algorithm that assumes a uniform $f$ would overlook the groups that contain a small number of samples. \n\nTo overcome the issue, the paper proposed a method to incorporate the subpopulation annotation in the loss function. In particular, the paper proposed the information-theoretic objective to maximize the term of $I(X;Y| v(X,Y))-I(X;Y)$, where $v(X, Y)$ is the random variable for the subpopulation annotation. To learn the optimal partition function that minimizes the term, the paper used the ‘empirical’ version of the entropy terms as the loss function, and proved that the term convergence with an additive error of $O(1/\\sqrt{N})$ to the actual information ‘’gain’’ $I(X;Y| v(X,Y))$. Therefore, if we minimize the objective, we are essentially maximizing $I(X;Y| v(X,Y))-I(X;Y)$ given that the joint distribution of $(X,Y)$ is fixed. \n\nThe paper then conducted experiments on several real-world datasets to compare the proposed method with the benchmark algorithms. Experimental results show that their proposed method could outperform the benchmark algorithms in a majority of the settings, albeit the improvements are marginal.  \n\nI have a mixed feelings about this paper. On one hand, it studied a well-motivated problem, and proposed an objective function with some solid theoretical guarantees. On the one hand, the exposition of the objective and the main theorem has some non-trivial problems, and the experiment performance only offers slight improvements over the baselines. As such, I would recommend a ‘’weak accept’’ due to the pros and cons.", "questions": "Most of the questions are in the ''weakness'' section. A less technical question is as follows. For the toy example in Figure 2, the original (overall) dataset is quite balanced, and an extremely skewed sampling process obtains the imbalanced training data. I understand this is important to test the performances for classification under imbalanced settings. However, can you give practical motivations to consider such a setting where we could have obtained a balanced dataset but, for some reason, have to use a very imbalanced subsampling process for the training dataset?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I think the paper studied a well-motivated problem in machine learning. Subpopulation imbalance can be viewed as an extension of the class-imbalance problem in machine learning, and the problem is harder since the ‘imbalance’ is not easily visible from the dataset. Using a training-based method is a natural idea, and essentially, the core of the paper is to train an annotation algorithm $v(X,Y)$. The paper also contains the novelty in the design of the loss function, which resorts to the tools in information theory.\n\nThe design of the objective function is justified by a theoretical analysis. Although the proofs are mostly standard applications of information theory and concentration inequalities, I appreciate the fact that they are properly written. Due to a hectic review timeline, I did not get time to verify all the calculations in Appendix B. I believe the proofs are correct with a high-level read-through.\n\nThe experiments are conducted on various datasets, and a bulk of benchmark algorithms are used in the comparison. I do appreciate the report on the error range, which justifies that the improvement is not due to statistical fluctuations.", "weaknesses": "A main criticism I have for this paper is that many assumptions are not explicitly stated, and the quantifiers are not properly stated in the settings and theorems. Furthermore, the paper provided no intuition on how the analysis is conducted. These problems gave me a hard time parsing the result. For instance, my first impression was that mutual information is *not* the correct notion that should be used to measure the gain for ‘subpopulation harmonization’. In particular, if $f$ is a *deterministic* function of $x$ or even a randomized function whose randomness is *independent of* $x$, then $I(X;Y)$ and $I(X;Y| v(X,Y))$ are essentially the same. (Btw $v(X,Y)$ is a random variable, which I believe you never mentioned.) I think the key in your model is that $y$ is a randomized function of $x$, and the randomness is *dependent* on the choice of $s$. However, such a fact is only clear after carefully reading the proof (!), and the whole thing reads quite confusing at first. \n\nAdditional (hidden) assumptions for the theorem in this paper include the fixed distribution of $(X,Y)$ and a fixed number of groups of subpopulations. Apart from being sloppy and not stated properly, the assumption of a fixed number of groups of subpopulations also implies we should have considerable knowledge of the dataset.\n\nA concern about the experiment is that the improvements compared to the baseline, especially w.r.t. the very basic ERM, are too marginal. I understand that getting the SOTA performance in experiment-based machine learning is an interesting problem, however slight the improvement is. But for this specific problem, the small margin of improvement (which itself is in a low accuracy range) might have low impacts on practice. \n\nMinor: \n- In Table 1, the meaning of $p(\\cdot)$ is overloaded – the ‘’class-imbalance’’ distribution is supported on the labels, while the ‘’subpopulation-imblance’’ distribution is supported on the groups of the subpopulations. I’d suggest using a different notation. \n\n- In theorem 3.3, the Rademacher complexity of $G$ is not properly defined. Instead of simply pointing at the literature, I think the notion should be defined in the appendix with the proper quantifiers.\n\n- Inequality (16) uses McDiarmid’s inequality, but this technical tool was never introduced. :(", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XaDa2ExHq3", "reviewer": "Reviewer_Dvju", "summary": "This paper proposes an approach (SHE) to learning a classifier that performs well for a set of latent subpopulations that are not uniformly represented in the training data. The approach is to learn subpopulation structure through an “optimal data partition” with maximal conditional entropy of the label given the learned subpopulations, effectively identifying subpopulations with balanced label distributions. Evaluation is conducted with respect to a balanced marginal distribution over the subpopulations. An extensive empirical evaluation is conducted using the COCO, CIFAR-100, and ImageNet datasets. The experiments involve comparisons to alternative approaches and ablation studies.", "questions": "* How does SHE compare to baselines in-domain, i.e., in cases when both the training and testing data are imbalanced? If I understand correctly, the claim of Proposition 3.2 is that SHE should not underperform ERM (in the limit) but it seems that this is not evaluated in the experiments because of the focus on balanced test distributions.\n* SHE is motivated to address a particular distribution shift problem where there is subpopulation shift over a set of subpopulations with balanced label distributions. How would SHE perform under other notions of subpopulation shift, where it is not assumed that the label distributions are balanced within subpopulations?\n* Could SHE be extended to handle an arbitrary specified target distribution over the latent subpopulation? Would this be as simple as weighting the subpopulation components of the LogSumExp operation?\n* Is there a reason why the model-based $V$ variant of SHE could not depend on both $X$ and $Y$? My interpretation of the method was that $V$ is not used at test time, so this might allow for a scalable model-based subpopulation mapping that matches the performance of the full $N$x$K$ matrix.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The presentation of the toy example in Figure 2 is clear and compelling for motivating the problem.\n* The theoretical analysis in section 3 is sound (to the best of my knowledge) and the method straightforward to implement.\n* The empirical evaluation is extensive, with comparisons to several baseline approaches in several settings, including on computer vision datasets that have been variably rebalanced, and in the context of fine-tuning large multimodal foundation models.", "weaknesses": "* The paper is challenging to read, and borders on unreadable at times, due to numerous grammatical errors and unusual choices of terminology, particularly in the abstract and the first two sections. As a reader, I was not able to understand the problem that this paper aims to solve until reviewing the mathematical presentation in section 3. While I like this paper overall, this is enough for me to argue that this paper should not be published without substantial revision.\n* It is not clear how model complexity and finite-sample considerations interact with the core arguments of the work. For example, the toy example in Figure 2 relies on the use of a linear model for the ERM and subpopulation-specific models. However, a more complex, non-linear ERM model could still, in-principle, fit the data well, even if it might not due to data insufficiency for the underrepresented subpopulations. It’s not obvious how this consideration surfaces in the theoretical presentation nor in the experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)