PaperlayerView exact revision on arXiv

Source-linked AI summary

No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems

Nimit S. Sohoni, Jared A. Dunnmon, Geoffrey Angus, Albert Gu, Christopher Ré

arXiv:2011.12945v2cs.LGcs.CV

TL;DR

Hidden stratification occurs when coarse labels conceal subclasses with uneven performance, while subclass labels are often unavailable in practice. GEORGE estimates subclasses by clustering learned features and uses those estimates in GDRO, achieving theoretical sample-complexity guarantees and empirical reductions in hidden stratification across four datasets.

  • Problem

    Coarse labels can hide meaningful subclasses, producing uneven performance on important groups, while subclass labels are frequently unavailable.

  • Method

    GEORGE clusters features from an ERM model to estimate subclass labels, then uses the estimated groups in GDRO.

  • Results

    Across four datasets, GEORGE provides evidence of reduced hidden stratification and, under suitable assumptions, achieves the same asymptotic sample-complexity rates as using true subclass labels.

  • Takeaways & Limitations

    GEORGE offers a way to measure and mitigate worst-case subclass performance without access to subclass labels.

  • Takeaways & Limitations

    Theoretical guarantees require sufficiently strong conditions on the data distribution and recovered clusters, and small total-variation error can be difficult to ensure in practice.

Abstract

from arXiv · show

In real-world classification tasks, each class often comprises multiple finer-grained "subclasses." As the subclass labels are frequently unavailable, models trained using only the coarser-grained class labels often exhibit highly variable performance across different subclasses. This phenomenon, known as hidden stratification, has important consequences for models deployed in safety-critical applications such as medicine. We propose GEORGE, a method to both measure and mitigate hidden stratification even when subclass labels are unknown. We first observe that unlabeled subclasses are often separable in the feature space of deep neural networks, and exploit this fact to estimate subclass labels for the training data via clustering techniques. We then use these approximate subclass labels as a form of noisy supervision in a distributionally robust optimization objective. We theoretically characterize the performance of GEORGE in terms of the worst-case generalization error across any subclass. We empirically validate GEORGE on a mix of real-world and benchmark image classification datasets, and show that our approach boosts worst-case subclass accuracy by up to 22 percentage points compared to standard training techniques, without requiring any prior information about the subclasses.

1 Introduction

Hidden stratification arises when coarse labels conceal meaningful subclasses, causing uneven and sometimes clinically dangerous performance. GEORGE estimates these subclasses from feature representations and uses them to improve worst-case performance without subclass labels.

  • Coarse class labels can conceal semantically distinct subclasses, leading models optimized for average performance to underperform on important groups.Examples include rare abnormal skin-lesion subclasses and drain-free pneumothorax images.
  • Feature representations of deep neural networks often contain information about unlabeled subclasses, motivating methods that estimate and address hidden stratification.
  • ERM can achieve arbitrarily poor performance on underrepresented subclasses in the proposed generative setting.
  • GEORGE estimates subclasses by clustering superclass features, then trains with GDRO to optimize worst-case performance over the estimated groups.
  • Under sufficiently strong distributional and cluster-recovery conditions, GEORGE attains the same asymptotic sample-complexity rates as using true subclass labels.
  • Across four image-classification tasks, GEORGE’s clusters often align with true subclasses and support measurement of worst-case subclass performance.

2 Background

The paper frames hidden-subclass robustness as a group-robust optimization problem with unobserved group labels. Its setup contrasts average empirical-risk minimization with worst-case subclass objectives.

  • Prior GDRO methods optimize worst-case performance over known groups, but group labels are often unavailable in practice.
  • The formal setup observes datapoints and superclass labels while treating each datapoint’s subclass label as latent and unobserved.
  • Standard training typically maximizes overall population accuracy or minimizes empirical risk rather than directly optimizing worst-case subclass performance.
  • Robust accuracy is defined as the worst-case expected accuracy over all subclasses.
  • If subclass labels were known, the worst-case per-subclass training risk could be optimized with GDRO; the paper instead seeks an approximation without those labels.

3 Modeling Hidden Stratification

The model explains hidden stratification through latent attributes that generate features and coarse labels, separating inherent hardness from imbalance. Imbalance can make ERM’s worst-case subclass accuracy collapse while GDRO remains robust.

  • 3.1 A Model of the Data Generating and Labeling Process: Latent attributes generate feature vectors and superclass labels, while only some attributes influence the observed label; the remaining attributes define hidden subclasses.
  • 3.1 A Model of the Data Generating and Labeling Process: The model assumes meaningful subclass structure represented by Gaussian components in a canonical feature space.
  • 3.2 What Causes Hidden Stratification, and When Can It Be Fixed?: Hidden stratification has two causes: inherent hardness and dataset imbalance.
  • 3.2 What Causes Hidden Stratification, and When Can It Be Fixed?: Inherent hardness lower-bounds the worst-case subclass error because some subclasses are intrinsically more similar to other superclasses.
  • 3.2 What Causes Hidden Stratification, and When Can It Be Fixed?: Unlike inherent hardness, imbalance-driven performance gaps can be addressed with known subclass labels by minimizing the robust GDRO objective.
  • 3.2 What Causes Hidden Stratification, and When Can It Be Fixed?: As the rare-subclass proportion α approaches 0, ERM worst-case subclass accuracy is O(α), whereas GDRO achieves 1 − O(α).

4 George: A Framework for Mitigating Hidden Stratification

GEORGE uses feature-space clustering to estimate hidden subclasses and then applies GDRO to those estimated groups. The procedure includes representation and clustering choices intended to retain rare, poorly performing subclasses.

  • GEORGE is a two-step algorithm that first estimates subclasses and then uses them to train a robust classifier.
  • The estimated clusters are used as groups in GDRO to minimize maximum per-cluster average loss and improve worst-case estimated-subclass performance.
  • An ERM model’s superclass features are dimensionally reduced and clustered to produce proxy subclass labels for training and validation data.
  • UMAP is applied before clustering because the authors find that it improves results.
  • Loss components orthogonal to the decision boundary are used as an alternative representation because subclasses with larger loss differences are more separable.
  • Standard k-means and Gaussian-mixture clustering often miss small clusters, so GEORGE uses over-clustering with automatically selected hyperparameters.
  • GEORGE selects model checkpoints using validation worst-case per-cluster performance rather than overall validation performance.

5 Analysis of George

The analysis shows that GEORGE can approximate robust training over latent subclasses when its estimated subclass labels are sufficiently accurate. Under Gaussian feature-space assumptions and related regularity conditions, it achieves the same asymptotic sample-complexity rates as robust training with true subclass labels.

  • Analysis of GEORGE: GEORGE uses estimated subclass labels to approximate the true GDRO objective for minimizing maximum per-subclass loss.The procedure uses probabilistic group assignments as noisy supervision for robust optimization.
  • Theoretical guarantee: Accurate subclass estimates allow GEORGE to achieve the same optimal asymptotic sample-complexity rates as GDRO trained with true subclass labels.This result applies when each subclass is represented by a different Gaussian in feature space and soft group predictions are used.
  • Loss estimation: Subclass loss estimates based on probabilistic weighting are unbiased and converge to the true sample-average subclass losses at rate O(1/√n).The weights use the probability that each example belongs to a given subclass.
  • Theoretical guarantee: Under Lipschitzness, bounded parameters, and unique Gaussian subclass conditionals, the robust-risk estimate is defined by the maximum estimated subclass loss.The theorem analyzes ˆRrobust := maxc ˆRc under these distributional assumptions.
  • Scope and dependence: GEORGE’s guarantee depends on recovering a good feature space, and model architecture can dramatically affect subclass recovery.The analysis links robust generalization to feature-space quality and reports this architectural effect empirically in Appendix C.

6 Experiments

Across four datasets, GEORGE measures and mitigates hidden stratification by clustering learned representations and optimizing over estimated subclasses. It improves worst-case subclass performance over standard training, while cluster-based evaluation helps estimate robust performance without subclass labels.

  • 6 Experiments: GEORGE’s theoretical guarantees require strong distributional and cluster-recovery conditions, while estimating the data distribution to small total variation error is difficult in practice.The paper notes that such estimation is achievable in special cases including the mixture-of-Gaussians setting.
  • 6.2 End-to-End Results: GEORGE improves robust accuracy by up to 22 points over ERM without requiring subclass labels.The method also substantially reduces the gap between ERM and subclass-GDRO robust error.
  • 6.2 End-to-End Results: GEORGE significantly improves worst-case subclass accuracy over ERM on Waterbirds, U-MNIST, and CelebA, while ISIC methods perform similarly.On CelebA, clustering pretrained BiT embeddings raises robust accuracy to nearly match subclass-GDRO.
  • 6.3 Clustering Results: GEORGE’s clusters align with poorly performing human-labeled subclasses, helping its GDRO stage target the subclasses that determine robust performance.Clusters also identify meaningful subclasses absent from the human-provided schema.
  • 6.3.2 Unlabeled Subclass Discovery: On ISIC, clustering separates benign no-patch images into histopathology and non-histopathology groups, revealing progressively harder clinically relevant examples.At least 77% of images in one no-patch cluster required histopathology, compared with under 7% in each other cluster.
  • 6.3.3 Estimating Robust Accuracy: Cluster-robust performance is usually closer to true robust performance than overall performance, enabling hidden-stratification detection without subclass labels.On ISIC, it estimates robust performance on the histopathology subclass better than the patch/no-patch subclass performance does.
  • 6.4 Effects of Validation Metric: Using unreweighted cluster-robust validation for checkpoint selection raises GEORGE’s true robust performance to 83.3%, over 22 points above ERM.The experiment retains unknown validation subclass labels and reflects differing training and validation distributions.

7 Conclusion

George measures and mitigates hidden stratification without subclass labels by clustering ERM representations and applying group DRO to the resulting groups. The paper provides theoretical guarantees and empirical evidence, while identifying representation and clustering selection as open directions.

  • 7 Conclusion: George clusters ERM features to approximate unlabeled subclasses, then uses those assignments as groups in group DRO.The two-stage procedure targets both measurement and mitigation of hidden stratification.
  • 7 Conclusion: Under suitable assumptions, George achieves the same asymptotic sample complexity rates as methods using true subclass labels.
  • 7 Conclusion: George is empirically validated on four datasets and shows evidence of reducing hidden stratification on real-world machine learning tasks.
  • 7 Conclusion: Future work includes learning better first-stage representations, improving unsupervised selection metrics, and characterizing when ERM representations separate subclasses.
  • 7 Conclusion: Hidden stratification concerns variable performance across unlabeled subclasses within the same superclass and has been observed in computer vision and medical machine learning.

B.2.3 ISIC

The ISIC dataset is constructed from publicly archived skin images with explicit benign or malignant metadata and fixed train, validation, and test splits. Subclass information is derived retrospectively from metadata, including patch and histopathology distinctions.

  • B.2.3 ISIC: The ISIC archive contained 23,906 images with corresponding metadata at the time of collection.
  • B.2.3 ISIC: Only images explicitly described as benign or malignant were used to formulate the task as binary classification.
  • B.2.3 ISIC: The dataset was split randomly into 80% training, 10% validation, and 10% test images.
  • B.2.3 ISIC: Benign patch images were identified as images from the SONIC repository, while histopathology examples were identified from diagnosis confirmation metadata.
  • B.2.3 ISIC: CelebA contains over 200,000 celebrity-face images annotated with 40 attributes, with the task defined as blond versus not blond classification.

B.3.1 Result Reporting

Results use repeated trials, dataset-specific metrics, and clustering choices selected through unsupervised scores. The procedure combines ERM representations, dimensionality reduction, clustering, and overclustering before robust evaluation.

  • B.3.1 Result Reporting: Each method was evaluated in ten trials per dataset, except CelebA, which used five trials because of its larger size.
  • B.3.1 Result Reporting: Reported X ± Y intervals and plot error bars are 95% confidence intervals computed from per-trial variability.
  • B.3.1 Result Reporting: George trains an ERM model, clusters its activations, and uses the assignments in a second-stage group DRO model.
  • B.3.1 Result Reporting: ISIC robustness is evaluated with worst per-benign-subclass AUROC rather than accuracy.The metric compares each benign subclass against the malignant superclass.
  • B.3.1 Result Reporting: The clustering procedure is computationally efficient and supports clusters of varied sizes, but its optimization remains an avenue for future work.
  • B.3.1 Result Reporting: The clustering pipeline uses UMAP or a loss component, searches clustering methods and k values, and applies overclustering with factor 5.
  • B.3.1 Result Reporting: Using the loss component can be an effective heuristic because rare subclasses are often misclassified at higher rates on imbalanced or spurious-feature tasks.

B.3.5 Dimensionality Reduction and Clustering: Further Details

Further experiments show that clustering quality and representation choice vary substantially across datasets. George can track robust performance, but difficult subclasses and weak representations limit gains in some settings.

  • B.3.5 Dimensionality Reduction and Clustering: Further Details: U-MNIST clustering consistently identifies a cluster enriched for the low-frequency “8” subclass and sometimes a low-performing crossed-“7” subset.
  • B.3.5 Dimensionality Reduction and Clustering: Further Details: Waterbirds performs best with the one-dimensional loss component, which yields sufficiently precise clusters of poorly performing subclasses despite limited minority-class frequency.
  • B.3.5 Dimensionality Reduction and Clustering: Further Details: On ISIC, patch and non-patch examples separate over 99% of the time, while histopathology and patch-color distinctions also appear in clusters.
  • B.3.5 Dimensionality Reduction and Clustering: Further Details: Despite high-fidelity ISIC clustering, neither the non-patch nor histopathology subset shows significant performance improvement.
  • B.3.5 Dimensionality Reduction and Clustering: Further Details: Subclass-GDRO also fails to improve ISIC non-patch performance, suggesting underperformance may involve factors beyond optimization, such as model capacity.
  • B.3.5 Dimensionality Reduction and Clustering: Further Details: CelebA clustering poorly identifies subclasses without BiT, whereas George-BiT improves robust performance and distinguishes men and women in the non-blond superclass at approximately 95% accuracy.
  • B.3.5 Dimensionality Reduction and Clustering: Further Details: BiT embeddings perform worse than standard George on other datasets, indicating that task-specific information matters for learning clusterable representations.
  • B.3.5 Dimensionality Reduction and Clustering: Further Details: Silhouette score selection would not have chosen BiT for CelebA, exposing a limitation of the current unsupervised representation-selection heuristic.

C.3 Runtime

GEORGE adds clustering and a second robust-training stage, increasing runtime, but shortened training can retain much of its robustness improvement. Its clustering is sensitive to scale and cluster-count choices, while noisy labels remain tolerable below a subclass-size threshold.

  • Runtime: 2–3× is the approximate runtime of standard GEORGE relative to simply training an ERM model.The default pipeline trains an ERM model, reduces and clusters activations, then trains a robust model.
  • Runtime: Less than 1.5× is the total runtime of George-BiT relative to GEORGE on the evaluated datasets.George-BiT omits ERM training and dimensionality reduction, retaining clustering and GDRO training.
  • Runtime: Over 70% of GEORGE’s worst-case performance improvement can be recovered at 1.3× ERM runtime on Waterbirds and CelebA by shortening second-stage training.U-MNIST also requires an adjusted learning-rate decay schedule and fixing k to 5.
  • Robustness to label noise: +3 points robust accuracy on MNIST and +4 points robust AUROC on ISIC are achieved over ERM below the stated label-noise threshold.This behavior uses a minor modification that discards small clusters; the threshold is that corrupted labels in each superclass remain fewer than the smallest subclass size.
  • Choosing k: 95.0%, 96.3%, 95.9%, 94.4%, and 90.8% are the U-MNIST robust accuracies for 2, 5, 10, 25, and 100 clusters per superclass.Robust performance initially improves with k, then decreases as large k fragments clusters.

C.7 Additional Classification Metrics

Additional metrics show that GEORGE generally improves per-subclass averaged accuracy while slightly reducing average precision, reflecting a trade-off between average-case and worst-case performance. The section also evaluates checkpoint selection using validation robust performance and reports convergence behavior.

  • Additional classification metrics: GEORGE typically increases per-subclass averaged accuracy while slightly decreasing average precision compared with ERM.The stated exception is U-MNIST, where per-subclass averaged accuracy decreases very slightly.
  • Checkpoint selection: Validation robust accuracy selects the “best” checkpoint for the reported test average-performance comparisons.ISIC uses AUROC instead of accuracy for checkpoint selection.
  • Convergence analysis: Theoretical and simulated convergence of ˜R_c − R_c is compared in Figure 10 using a synthetic mixture-of-Gaussians experiment.The experiment varies sample size and averages results over 20 trials with newly sampled per-subclass distributions.

D.1 Analysis of Example 3.1

The constructed example separates superclass labels from a rare latent attribute, exposing how regularized ERM can neglect rare subclasses. In contrast, GDRO optimizes the worst subclass loss and retains high rare-subclass accuracy as rarity increases.

  • Example setup: The superclass label Y equals Z2, while the latent attribute Z1 creates two subclasses within each superclass.The observed data are the underlying features V, with each latent configuration producing a distinct Gaussian distribution.
  • Example setup: As rare-subclass proportion α approaches 0, ERM worst-case subclass accuracy is O(α), whereas GDRO achieves 1 − O(α).The example uses linear models with regularized logistic loss and Gaussian subclass distributions.
  • ERM behavior: Only O(α) of the Z1 = −1 subclass is classified correctly by ERM when α is sufficiently small.The result follows from the regularized ERM direction concentrating on the dominant subclass.
  • GDRO behavior: 1 − O(α) is the accuracy of the GDRO decision boundary as α decreases to zero.GDRO chooses the boundary separating superclass means with maximum distance to any subclass mean.
  • Risk estimation: ˜R_c is an unbiased estimate of the per-subclass sample-average loss R_c, with their difference converging at O(1/√n).The estimator reweights superclass examples using subclass-to-superclass posterior likelihood ratios.

D.2.2 Proof of Theorem 1

Theorem 1 bounds GEORGE’s estimated robust risk by combining distribution-estimation error, estimated per-subclass risk error, and uniform convergence. The proof specializes these steps to unique Gaussian subclass distributions under regularity assumptions.

  • Theorem statement: Theorem 1 assumes Lipschitz loss and predictor, bounded predictor parameters, and unique Gaussian P(x|z = c) for each subclass.The theorem concerns a high-probability bound for the estimated robust risk.
  • Estimated risk: ˆR_c uses weights ˆw(x,c) computed from estimated subclass and superclass conditional distributions.The estimated risk mirrors ˜R_c but replaces true reweighting quantities with estimates from ˆP.
  • Proof strategy: The proof first relates total-variation distribution error to estimated-risk error, then bounds Gaussian-mixture estimation error and applies uniform convergence.This sequence yields the final high-probability robust-risk bound for the returned model.
  • Risk-estimation error: 3M ˆπ_min^-1 ϵ + O(1/√n) bounds |ˆR_c − ˜R_c| with high probability when estimated superclass and subclass distributions are within TV distance ϵ.M bounds the loss, and ˆπ_min is the minimum estimated subclass proportion.

Total variation in estimated per-subclass distributions: Gaussian case

The Gaussian-case analysis connects mixture estimation to per-subclass distributional accuracy, then transfers these bounds to robust training and GEORGE’s worst-case subclass generalization risk.

  • Mixture-component recovery: When mixture components are distinct and have nonzero weights, sufficiently small mixture-level error implies component-level error O(ϵ) after a permutation.Lemma 3 establishes correspondence between true and estimated Gaussian components under these assumptions.
  • Risk transfer: Applying component recovery and robust-loss bounds yields per-subclass risk estimation error ˜O(1/√n) for a fixed prediction function.The proof transfers total variation bounds from estimated mixture components to subclass risks and then uses triangle inequality.
  • Robust optimization: Uniform convergence extends the robust-risk approximation to every classifier in the hypothesis class, including the minimizer of the estimated robust objective.The argument combines uniform convergence, triangle inequality, and a union bound.
  • Guarantee: Under the stated assumptions, GEORGE’s excess robust generalization risk is ˜O(1/√n), matching the generic worst-case lower-bound rate up to logarithmic factors.The comparison applies even though subclass labels are estimated rather than directly observed.
  • Technical requirement: The proof requires independent samples for distribution estimation and robust-loss optimization, although the practical implementation reuses all examples.Randomly splitting the data preserves the asymptotic dependence on n.
Loading 2011.12945v2…