Generalizing to any diverse distribution: uniformity, gentle finetuning & rebalancing
out-of-distribution generalization;ood generalization;diversity;rebalancing;finetuning
Yes
No ethics review needed.
Please refer to the weaknesses.
- The authors propose the distributionally diverse risk and derive risk bounds for the DD risk.
- The authors demonstrate the optimality of the uniform distribution under their uncertainty set characterized by the entropy.
- Various experimental results are given to support their methods.
This paper proposes the distributionally diverse risk to address the out-of-distribution generalization problem. The DD risk takes the similar formulation as the DRO risk with the uncertainty set captured by entropy. Based on the DD risk, the authors demonstrate the optimality of the uniform distribution. Experimental results validate the effectiveness of the proposed reweighting/rebalancing methods.
Given the theoretical and experimental results, my major concerns include:
- Relationship with Distributionally Robust Optimization: The proposed DD risk has similar formulation as the DRO risk, and the only difference is that the authors use entropy to capture the difference between distributions. In DRO literature, there are works using $f$-divergence, Wasserstein distance, MMD distance, etc. For different distance metrics used in DRO, there are also works to derive the risk bounds (for example, see [1,2]). What is the advantage of this specific kind of distance metric? Why can it lead to better generalization performance compared with others? I would suggest the authors providing a more comprehensive illustration on this. 
- As the authors demonstrate, the optimality depends on the "entropy" selected for the uncertainty set. And the result is more like: the uniform distribution is optimal if we consider this specific kind of DRO risk (or the DD risk). Can the authors demonstrate why this specific formulation is better? 
- As for the experiments on real-world datasets, can the authors demonstrate in detail how the density estimator is learned? (I read the appendix but found it is not so clear.) Are the embeddings coming from a pre-trained ResNet-50? How is the effect of the density estimator on other datasets? I think the density estimation itself may be as hard as the original prediction task. How can the authors guarantee that the learned estimator is good? And when it is bad, will it even hurt the model performance?
- What if there are noisy data? In the DRO literature, there is a phenomenon named "over-pessimism", i.e. the worst-case risk may be too unrealistic and thus hurt the model learning. I do think it will hold in the formulation proposed in this work. For example, when doing reweighting or rebalancing, what if some noisy data or outliers are given a much higher sample weight?



[1] Variance-based regularization with convex objectives. John Duchi, Hongseok Namkoong.
[2] Learning models with uniform performance via distributionally robust optimization. John Duchi, Hongseok Namkoong.
Yes
No ethics review needed.
Please see weaknesses.
1. The paper introduces a fresh perspective on OOD generalization by focusing on the DD risk, which is a significant departure from traditional domain adaptation and generalization approaches.

2. The theoretical analysis provides mathematical grounding for the optimality of training on uniform distribution and the impact of non-uniformity on generalization.

3. The paper offers practical solutions such as gentle finetuning and rebalancing to mitigate the effects of non-uniformity when uniform samples are unavailable, which are valuable for real-world applications.

4. With the development of basic models and large-scale models, it becomes increasingly important to maintain robustness when the distribution of training and test data is very different. The emphasis on uniform training in this paper provides new ideas for large model training. At the same time, fine-tuning and rebalancing strategies also have application potential in the actual training of basic models, providing an important reference for future basic model design and training optimization.
The paper aims to address the challenge of out-of-distribution (o.o.d.) generalization in machine learning by proposing a framework centered around the concept of distributionally diverse (DD) risk, which aims to minimize the worst-case error across all sufficiently diverse test distributions within a known domain. The paper presents theoretical findings and practical remedies for training models that generalize well under any diverse test distribution. Key contributions include proving the optimality of training on a uniform distribution, analyzing the impact of non-uniformity through finetuning and rebalancing, and providing empirical evidence supporting the theory across various tasks involving distribution shifts.
1. Given the problem formulation in this paper, the solution seems rather straightforward. The stated objective is to control "the worst-case error across all distributions with sufficiently high entropy." The approach borrows from TTA by adjusting the training distribution to control Distributional Discrepancy (DD). However, the paper does not adequately justify why DD is a truly reasonable metric and what practical physical interpretations it has in real-world applications.
2. The practical effectiveness of the proposed method is questionable. For example, in Tables 2 and 3, rebalancing does not show a significant improvement over ERM. The final performance gains seem to stem from various tricks (e.g., UMAP, label conditioning, etc.), which are inconsistently applied across different experiments. This suggests to me that the results might be due to data tuning rather than a genuinely effective method applicable to diverse scenarios.
3. The quantitative analysis of different strategies is not in-depth enough. For example, there is no detailed data to support the applicability and effect of the rebalancing strategy under different types of distribution shifts. In addition, there is a lack of systematic analysis of the sensitivity of fine-tuning parameters (such as fine-tuning step size, number of iterations, etc.) and their effects.
4. The paper's presentation could benefit from several improvements. For instance, Table 2 references "UMAP" and Figure 3 mentions "WDL2" without providing explanations in the corresponding captions or within the main text. This lack of clarification might hinder reader comprehension.
Yes
No ethics review needed.
1. In theorem 3.2, what is the choice of $\alpha$ that minimizes the upper bound?
This paper follows a clear and logical organization:
1. Worst-case risk should be considered for prediction under covariate shift.
2. The distribution diverse risk is a good surrogate for the worst-case risk.
3. The risk on uniform distribution is a good surrogate for the distribution diverse risk.

The first point is an important objective and recognized challenge for covariate shift generalization. The paper has also well addressed the third point by (1) showing reweighting data towards a uniform distribution yields the lowest distribution diverse risk within the framework of reweighted empirical risk minimization (Theorem 3.1), and (2) studying the generalization gap between distribution diverse risk and uniform distribution risk at both population and sample levels (Theorems 3.2 and 4.3).

The paper also does a good work in explicitly acknowledging possible limitations, which improves overall clarity. For example, the author notes that Theorem 3.1 is less applicable to the setting of supervised domain adaptation, which is just fine because learning under uncertainty of target distribution is actually more challenging. But the statement helps accurately contextualize the results in the literature.
This paper addresses risk minimization conditioned on covariates, i.e., worst-case risk as named in the paper, by introducing the learning objective of distributionally diverse risk that considers the worst-case risk over high-entropy distributions. The paper proceeds to consider reweighted risk on uniform distribution as a surrogate for the distributionally diverse risk, which leads to the classic importance weighting technique. Experiments are conducted on datasets with prevalent covariate shift.
I am not convinced by the paper that the distribution diverse risk is a good surrogate for the worst-case risk. The introduction has a short explanation that high-entropy test distributions as considered by the distribution diverse risk reflect natural distribution shifts, and high-entropy training and test distribution empirically yield good generalization in literature, which will benefit from more detailed elaboration. In the current submission, I find the argument is not supported, or even contradicted by the theory.

The distribution diverse risk (DD) converges to the worst-case risk until $\gamma \to \infty$ (eq 4), when DD risk reduces to the worst-case over **all** possible distributions X. This asymptotic behavior is satisfied by a large family of "uncertainty sets" of test distributions, e.g., when the radius of the uncertainty set of Distributionally Robust Optimization (DRO) tends to infinity, or when we redefine DD as distributions with a differential entropy near any real value by a gap $\gamma$. This asymptotic argument does not address why  DD is superior to the alternative designs of uncertainty sets, in terms of a surrogate for the worst-case risk, which is pivotal for the position of this paper in the broad literature of OOD generalization under covariate shift, e.g., DRO [1].

It is also revealed in Theorem 3.2 that the DD risk scales with $\Omega(\gamma)$ as $\gamma \to \infty$, indicating that the DD risk cannot be estimated by uniform distribution risk at all when we actually consider the objective of worst-case risk. The author also acknowledges that the theorem only applies to $\gamma = O(1)$, contradicting the setup for the worst-case risk. Therefore, I interpret from the paper that the DD risk is another heuristic design of the uncertainty set of target distribution, as similar in DRO. However, a stronger theoretical or empirical reasoning for this design is expected in order to consolidate this paper.

Minor points:

1. Implications for LLM and foundation models (L61-69) in the introduction can be over-claimed since they are not supported by theory and the LLM-free experiments. Moreover, importance weighting in general is ineffective for over-parameterized models [2]. The intuition is that over-parameterized models interpolate the samples with an almost zero loss on each sample, corresponding to $\epsilon = 0$ in Theorem 3.1. In this case, the sample risk does not uniformly bound the population risk because there are infinitely many interpolators on training data among which most are not generalizable. In this case, considering the maximum risk of predictors does not make sense. As is also stated by the author, Theorem 3.1 does not consider the inductive bias of optimizers, but generalization for over-parameterized models are only possible when implicit regularization is imposed by the learner. 

2. 0-1 loss and the existence of a ground truth model $f^\star(x)$ is assumed throughout the theory. Combined together, they imply an impractically strong assumption that the classification task is noiseless, i.e. $f^\star(x)$ is binary rather than a probability, which is not even satisfied by the simple classic model of mixture of gaussians, which also appears as the first experiment.

3. Section 4.1 and 4.2 are less relevant to the sample-level analysis of the generalization behavior of DD risk. This is because the author characterizes the risk gap by the l1 distance between the empirical measure and the population measure. However, this distance $\delta(u,p)$ always equals 2 for any continuous measures u and discrete measures p, which makes eq. 8 ineffective when p is an empirical measure, the sum of n dirac measures. This happens to Theorem 4.1 as well. In fact, a standard PAC-Bayes bound will characterize the risk gap as the ratio between some divergence measure and a power of n. The absence of sample numbers in eq.8 and Theorem 4.1 cannot reflect the non-asymptotic behavior as $n \to \infty$.

4. ColoredMNIST is a dataset that contains large concept shift instead of covariate shift as claimed by the paper. This synthetic dataset flips the label of digits (y) with a probability that is correlated with the color (part of x, denoted by $x_2$). This correlation varies across splits, leading to concept shift ($y|x_2$ shift). If we denote the shape of the digits as $x_1$, the dataset features $y \not \perp x_2 | x_1$. This causes $y|x$ shift.

5. An exhaustive investigation of baselines is performed for ColoredMNIST. Reporting their results for iWildCam and PovertyMap can be more convincing. For example, the public benchmark of PovertyMap shows that the SOTA C-Mixup, some simple variant of Mixup in Table 3, is surprisingly effective (0.53 v.s. 0.47 as the best result in the paper).

[1] Duchi, J. C., Hashimoto, T., & Namkoong, H. (2019). Distributionally robust losses against mixture covariate shifts.

[2] Zhai, R., Dan, C., Kolter, J. Z., and Ravikumar, P. K. Understanding why generalized reweighting does not improve over ERM.
Yes
No ethics review needed.
Is uniform distribution best in real applications? I believe image data are located in the low dimensional manifold. Why does this uniform distribution help generalization?
The proposed theoretical framework is novel. The paper is well written.
This paper proposes a new learning framework for studying the out-of-distribution generalization. They construct the uncertainty set by using the entropy of the data distribution. They define the worst-case risk in this uncertainty set by Distributionally diverse (DD) risk. They build a series of theorems containing: A classifier optimized for the uniform distribution will yield the smallest DD risk; the generalization bound of DD risk under the zero-one loss. They further provide two ways to improve the generalization ability, fine-tuning with penalty and rebalance through weighting. They verify the effectiveness of the proposed method on IwildCam, PovertyMAp, and ColorMNIST datasets.
The experimental results are not solid enough to make a conclusion. Although rebalancing (UMAP-8, label cond., WDL2) achieves good performance in "-90%" and therefore performs well in average score, it remains unclear whether weight distance to initialization is practically effective, as no results for this metric are provided on other datasets in Tables 1 and 2. Furthermore, Tables 1 and 2 exclude many relevant baseline methods—such as ARM, which performs well on ColoredMNIST, as well as VREx, RSC, and others.

Additionally, there is no universal criterion for determining which method is best suited to practical application. The authors apply different modeling approaches (e.g., WDL2 for ColorMNIST and UMAP for PovertyMap) depending on the dataset, yet provide no general guidance on how to select methods for other contexts.
We put forth a new framework for out-of-distribution generalization that focuses on the worst case error across any diverse test distribution within a domain; we propose practical remedies to improve performance.
@inproceedings{
anonymous2024generalizing,
title={Generalizing to any diverse distribution: uniformity, gentle finetuning \& rebalancing},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=sRb9sddBy2},
note={under review}
}
As training datasets grow larger, we aspire to develop models that generalize well to any diverse test distribution, even if the latter deviates significantly from the training data. Various approaches like domain adaptation, domain generalization, and robust optimization attempt to address the out-of-distribution challenge by posing assumptions about the relation between training and test distribution. Differently, we adopt a more conservative perspective by accounting for the worst-case error across all sufficiently diverse test distributions within a known domain. Our first finding is that training on a uniform distribution over this domain is optimal. We also interrogate practical remedies when uniform samples are unavailable by considering methods for mitigating non-uniformity through finetuning and rebalancing. Our theory provides a mathematical grounding for previous observations on the role of entropy and rebalancing for o.o.d. generalization and foundation model training. We also provide new empirical evidence across tasks involving o.o.d. shifts which illustrate the broad applicability of our perspective.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
out-of-distribution generalization ood generalization diversity rebalancing finetuning
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/d7c813a9a403b6139c47b8f2e5399d1aa1890695.pdf
learning theory
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Generalizing to any diverse distribution: uniformity, gentle finetuning & rebalancing
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission