Generalizing to any diverse distribution: uniformity, gentle finetuning & rebalancing

Authors: Andreas Loukas, Karolis Martinkus, Edward Wagstaff, Kyunghyun Cho
Primary area: learning theory
Keywords: out-of-distribution generalization, ood generalization, diversity, rebalancing, finetuning

Abstract:
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.

Decision: reject
Reviewer scores: 5, 5, 3, 5

Reviews and discussion:
<thread>
creation date:
1730493342151

thread_type:
official_review

<thread_object>
creator: reviewer_b5w2

<content>
summary:
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.

strengths:
The proposed theoretical framework is novel. The paper is well written.

weaknesses:
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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730541169210

thread_type:
official_review

<thread_object>
creator: reviewer_3tma

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
1. In theorem 3.2, what is the choice of $\alpha$ that minimizes the upper bound?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730644070658

thread_type:
official_review

<thread_object>
creator: reviewer_6w8l

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730669078050

thread_type:
official_review

<thread_object>
creator: reviewer_dzlp

<content>
summary:
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.

strengths:
- 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.

weaknesses:
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.

questions:
Please refer to the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732403572357

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General reply cont.

comment:
### 2. Explaining how DD risk and DRO differ

**TLDR**: The DD risk and DRO approaches make starkly different assumptions about the relation between the training and test distribution and lead to different theoretical arguments and practical solutions. 

As a reminder, the “distributionally robust stochastic optimization (DRO) framework (..) learns a model providing good performance against perturbations to the data-generating distribution" [3], whereas we here advocate for learning a model that exhibits good performance against any high entropy distribution in the same domain.   

**Similarities.** Both approaches consider the worst case behaviour over a set of distributions. Further, at the limit, both DRO and DD converge to the worst case risk over the domain.  

**Differences.** The assumptions posed by the two frameworks about the relation of the training and test distribution are starkly different: 

- DRO assumes that the test distribution is close to the training distribution. The small train-test discrepancy assumption, though simple & intuitive is not always valid. To account for actual distribution shifts encountered in the wild, often very wide uncertainty sets are needed around the training distribution. Further, as also mentioned in the justification of DD risk, train-test discrepancy has been empirically found to correlate poorly with generalization [1] and concurrent theoretical work also advocates against it [2]
- We don't consider the worst-case bounded distribution shift (as is done in DRO) but the worst case risk under any distribution of sufficient entropy. As such, the training distribution holds no special role in our framework, whereas in practice DRO defines the “center of mass” of the set of admissible test distributions.

The above are consequential conceptual and practical differences:

- Conceptually, whereas in DRO one needs to think about how similar is the test distribution to the training distribution, we here consider arbitrary distributions on the test domain and pose a constraint on entropy. This is a change in paradigm.
- Practically, the two approaches lead to very different solutions to the problem and, as we show empirically, our approach can yield quantifiable benefits.

We have significantly expanded the related work discussion in Appendix A (lines 820-828) to highlight these differences and to include the references kindly provided by the reviewers. We have also included Figure 5 which illustrates the conceptual differences between DD and DRO by providing examples of admissible test distributions for DD that are inadmissible under DRO and vice-versa.     

[1]  An empirical investigation of
domain generalization with empirical risk minimizers. Ramakrishna Vedantam, David Lopez-Paz, and David J Schwab. NeurIPS 2021.

[2] Beyond Discrepancy: A Closer Look at the Theory of Distribution Shift. Robi Bhattacharjee, Nick Rittler, Kamalika Chaudhuri. ArXiv *29 May 2024.*

[3] Learning Models with Uniform Performance via Distributionally Robust Optimization. John Duchi, Hongseok Namkoong. 2018
</content>
</thread_object>

</thread>
<thread>
creation date:
1732403593523

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General reply

comment:
We thank all reviewers for their work and comments. 

The reviews have acknowledged the novelty of our perspective within the imporant problem of o.o.d. generalization, found the paper well-written, commented positively on the practical applicability of our theory, and have not identified any errors in our theorems. 

There were two especially important questions that we received: 

1. to provide further justification of our "distributionally diverse (DD) risk” formalism  
2. to clarify the relation of our approach with distributionally robust optimization (DRO)

We provide concrete answers to these two questions below. All other questions are answered at individual replies. We have also updated the manuscript with more detailed discussion of these main points as well as with clarifications and new empirical results requested by the reviewers.  

We hope that the reviewers read our answers with open mind, giving our work a fair and equitable treatment.

### 1. Justifying the DD risk formalism

The DD risk entails measuring the behavior of a classifier across any test distribution of sufficiently high-entropy within a domain. In the following, we justify the definition from a mathematical, intuitive, and empirical perspective.  

1. **Mathematical perspective.** Our definition of DD risk follows from two *desiderata* on how to model the generalization of a predictive model under covariate shift:
    1. **We can assume to know very little about the relation between the training and test distributions.** Unfortunately, the distance between training and test distribution (also referred to as train-test discrepancy) is often large in practice, discouraging stronger assumptions. This is supported by the literature where, train-test discrepancy has been empirically found to correlate poorly with generalization [1]. Concurrent/independent theoretical work also advocates against it [2].
    2. **Since we wish to avoid judging the behavior of a classifier based on pathological / overly specific examples, we care most about distributions that do not assign overly high likelihood to any small set of examples.** For context, desideratum 2 directly addresses the reason why worst-case analysis is non-meaningful in learning: if no such constraint is placed, an adversary may always construct a pathological test distribution that concentrates all its mass on a small set of inputs where the predictive model is wrong (in general, such a set can always exist when learning from examples). We directly avoid this situation by asserting that no such small high likelihood set can exist.   
    
    It can be easily confirmed that asserting that the test distribution entropy is high satisfies both these desiderata. Lower bounding test entropy ensures that the spread of a distribution is large, without imposing any further assumptions about the relationship between train and test distributions.     
    
2. **Intuitive motivation.** A generally performant model is one that generally performs well on the many test instances that the world throws at it—even if it might fail in specific instances. In other words, though we cannot hope that our model will generalize to any arbitrary sharp test distribution, we may expect that it will perform well on test distributions that are more spread out and thus more typical of the data domain.   
    
    Our assumption of high test entropy exactly corresponds to the minimal and intuitive assumption that the classifier will be evaluated on a diverse set of possible inputs, rather than on few pathological examples. 
    
3. **Empirical justification.** Finally, to emphasize the practical applicability of our assumption in real-world applications, we remark that, as mentioned in the introduction, test entropy was identified [1] as the most predictive measure for o.o.d. generalization in a recent thorough empirical study, significantly outperforming many other intuitive measures across multiple datasets.  

These points justify our assumption as grounded in previous work, minimal, intuitive, and practically applicable. 

Aiming to make these points more easily accessible to future readers, we have included a relevant new section in Appendix B of the revised document.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732715356680

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Further evidence supporting the significance of the claimed empirical improvements

comment:
Some of the reviewers posed doubts about the competitiveness of our approach with respect to the state-of-the-art (sota). Though our purpose is not to claim sota results, it is still relevant to answer this question. To this end, in the following we present our results in context of the public leaderboard from the WILDS benchmark. 


| **Algorithm**  | **PovertyMap** (Worst-U/R r) | **iWildCam** (Macro F1) |
| --- | --- | --- |
| Fish | - | 22.0 (1.8) |
| IRMX (PAIR opt) | 0.47 (0.09) | 27.9 (0.9) |
| IRMX | 0.45 (0.05) | 26.7 (1.1) |
| CORAL | 0.44 (0.07) | 32.7 (0.2) |
| CGD | 0.43 (0.04) | - |
| Group DRO | 0.39 (0.06) | 23.8 (2.0) |
| C-Mixup | **0.53 (0.07)** | - |
| ABSGD | - | 33.0 (0.6) |
| ERM w/ data aug | 0.49 (0.06) | 32.2 (1.2) |
| ERM (more checkpoints) | - | 32.0 (1.5) |
| ERM (grid search) | 0.45 (0.06) | 30.8 (1.3) |
| ERM (rand search) | 0.5 (0.07) | 30.6 (1.1) |
| IRM | 0.43 (0.07) | 15.1 (4.9) |
| ARM-BN | - | 23.3 (2.8) |
| Test-time BN adaptation | - | 13.8 (0.6) |
| MixUp | - | 13.8 (0.8) |
| JTT | - | 11.0 (2.5) |
| rebalancing (ours) | 0.47 (0.10) | **35.5 (0.8)**  |

To ensure fairness, we have included only leaderboard entries that no not deviate from the official benchmark guidelines. 

We observe that: 
- **on iWildCam our approach achieves a statistically significant improvement over the leading method ABSGD** (p-value of 0.015 based on a Welch t-test);
- **on PovertyMap, our approach is not significantly worse that the the best perfoming method C-Mixup** (p-value of 0.448 that these samples come from the same distribution according to a Welch t-test).

As for ColorMNIST, though there is no public leaderboard, in our paper we show that our method significantly outperforms 14 other methods, which constitutes a reasonably throrough treatment.

*To summarize, we have presented statistical evidence that our approach is competitive or better than previous methods on these datasets.* We hope that this new data, put our empirical results in better perspective of the literature addressing the doubts expressed. We also re-iterate that, despite this good performance, our main point was to support our theory empirically by showing improvement over ERM, which the above results do show.

As for the criticism that our method's performance depends on how the density estimator is employed (e.g., use of dimensionality reduction). We acknowledge that density estimators are imperfect and that tuning them can affect the final result. However, we point out that most modern machine learning methods have hyperparameters and that, thankfully, there is a generally applicable method to select them by using a validation set (as we do here).
</content>
</thread_object>

</thread>
<thread>
creation date:
1733249336639

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional results on Waterbirds and MultiNLI

comment:
As reviewer DzLP suggested testing our method on two aditional experiments (Waterbirds and MultiNLI) in the table below we provide the preliminary results for them, when no aditional information (group labels) are used. This also addresses the final concern raised by reviewer B5w2. Due to time constraints we did not perform the hyperparameter tuning for the embedding projection and normalization we did for the rest of the experiments. Instead we used a concervative choice of using UMAP to down-project the embedding space to 64 dimensions, together with normalizing the embeddings by the maximum vector lenght. We exactly followed the experimental setup from [1] with their proposed hyperparameter sweep for ERM.

| **Method** | **Waterbirds (Worst Acc)** | **MultiNLI (Worst Acc)** |
| --- | --- | --- |
| ERM | 85.5 $\pm$ 1.0 | 67.6 $\pm$ 1.2 |
| JTT | 85.6 $\pm$ 0.2 | 67.5 $\pm$ 1.9 |
| RWY | 86.1 $\pm$ 0.7 | 68.0 $\pm$ 1.9 |
| SUBY | 82.4 $\pm$ 1.7 | 64.9 $\pm$ 1.4 |
| Rebalancing (UMAP 64) | 85.4 $\pm$ 0.2 | 69.1 $\pm$ 2.9 |

Consistently with the rest of our experiments, the improvement was significant when density fit was good: the negative log likelihood of the MultiNLI training set was approximately two times smaller than of Waterbirds for the trained density model. Accordingly, in MultiNLI we saw a notable improvement over ERM and other baselines that do not use aditional information repored in [1]. 

[1] Idrissi, Badr Youbi, et al. "Simple data balancing achieves competitive worst-group-accuracy."
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work proposes a new method to enhance the out-of-distribution (OOD) performance of a prediction model in the covariate shift setting, by trying to mimic a uniform covariate distribution in training through the proposed distributionally diverse (DD) risk minimization. The authors pointed out the effect of considering a uniform distribution for worst-case risk minimization, and connected the DD risk to the uniform case. Two practical methods are developed to implement the DD risk minimization.

I agree with reviewers' appreciation on the importance of the topic, the insight contribution, development of the practical methods, and the additional results and differentiation with DRO in the rebuttal phase. Nevertheless, there are also concerns raised. Theoretically, the reviewers are still concerned on the rationale of the DD risk from the proposed risk under uniformity perspective, and the practice of preferring the uniform distribution as a reference. The authors argued that uniformity on manifold can be achieved by density estimation, but I'm not sure if the estimated density is already expressed under the base measure of the manifold instead of that of the ambient space. Practically, the method has quite a few degrees of freedom in the implementation detail, which could have a significant effect on the performance. There are concerns on the solidity of the presented results. Reviewers and I expect the authors to make a systematic and convincing reasoning or investigation for a guideline for each implementation choice (e.g., density estimation).

Given this situation, I tend to recommend a reject for the current version, and encourage the authors to further enhance the paper accordingly.