Optimal Transport Reweighting for Robust Learning under Spurious Correlations and Label Noise
Machine learning models often fail under subpopulation shift when spurious correlations mislead them into relying on shortcut features. This paper introduces POTER, a method that derives per-sample importance weights from optimal-transport (OT) dual potentials between the training distribution and a small validation-based reference distribution. By avoiding loss-based proxy selection entirely, POTER aims to remain robust under label noise while requiring only a single standard ERM training stage. The authors report state-of-the-art worst-group accuracy on several benchmarks and identify a previously overlooked failure mode: subgroup-concentrated label noise.
POTER offers a conceptually clean and empirically strong framework for robust learning under spurious correlations and label noise. Its use of OT dual potentials to sidestep loss-based selection is well motivated, and the class-conditioned decomposition keeps computational overhead modest. However, the theoretical support stops at reweighting derivations: the paper provides no finite-sample generalization bound for the final reweighted classifier, and the method remains reliant on a representative, group-annotated validation set. The empirical gains are largest under label noise, whereas in the clean setting the improvements over recent validation-only baselines are more incremental.
The strongest aspect of POTER is its principled derivation of sample importance from transport geometry rather than empirical loss. The class-conditioned OT decomposition in Proposition A.1 prevents cross-class contamination from shared spurious attributes and yields efficient, parallelizable computations. Empirically, the reweighting signal behaves as intended: bias-conflicting samples receive higher weights, and mislabeled examples are naturally downweighted because their feature embeddings misalign with the reference distribution.
The paper's theoretical contributions are limited to the reweighting derivation and do not extend to generalization guarantees for the learned classifier. Additionally, POTER depends on fixed pretrained feature extractors that are not trained on the downstream task; if these embeddings are poorly aligned with the target classes, the OT distances may be uninformative, yet the authors do not ablate this design choice. Another concern is that the class-conditioning step partitions training data by observed labels, which are themselves corrupted under label noise, but the effect of this misclassification is not analyzed. Finally, while subgroup-concentrated noise is introduced as an important failure mode, the evaluation is restricted to two image datasets, leaving open questions about broader applicability.
The empirical evidence supports the claim that POTER outperforms loss-based and retraining-based methods under label noise, with large margins over KNN-RAD and Group DRO at 30% symmetric noise. The comparison to validation-only baselines is generally fair, though POTER is not uniformly best: on CivilComments its worst-group accuracy ($71.0\%$) trails GIC ($72.5\%$), and its average accuracy on CelebA ($91.2\%$) is noticeably lower than ERM ($94.9\%$). The characterization of JTT and related retraining pipelines is accurate: these methods indeed rely on a preliminary model's mistakes to identify minority samples, which becomes unreliable when mislabeled samples also incur high loss. The introduced subgroup-concentrated noise setting is valuable, but additional datasets and noise types would strengthen the claim of broad robustness.
The authors release code at a public GitHub repository and report hyperparameters, optimizer settings, and data augmentation protocols in Appendix E. OT computation uses the standard Sinkhorn algorithm with specified regularization and convergence thresholds, and the training configurations are taken from prior protocols. A potential barrier to independent reproduction is the reliance on a limited validation set with group annotations; while this is standard for the problem setting, the exact reference distribution construction and the median-MAD normalization of dual potentials are method-specific details that must be followed precisely. No compute budget is reported beyond a single NVIDIA RTX 3090.
Machine learning models often suffer performance degradation under subpopulation shift, particularly when spurious correlations cause models to rely on shortcut features that fail to generalize across subgroups. A recent line of work mitigates this issue by using loss-based signals to identify informative samples, but these signals can become severely distorted under label noise: mislabeled samples may also incur large losses and contaminate subsequent reweighting or retraining. Despite its practical importance, this intersection remains largely underexplored. We propose POTER, a reweighting framework based on optimal transport that derives sample importance from the transport geometry between the training distribution and a reference distribution constructed from limited validation group annotations. By measuring alignment at the individual-sample level rather than relying on loss, POTER downweights mislabeled or strongly bias-aligned samples while assigning higher importance to samples better aligned with the reference distribution. In addition, POTER requires only a single ERM training stage, moving beyond the retraining paradigm common in recent work. Across standard benchmarks and noisy-label settings, POTER achieves state-of-the-art worst-group accuracy, including cases where label corruption is concentrated within minority subgroups.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.