PaperlayerView exact revision on arXiv

Source-linked AI summary

Model Patching: Closing the Subgroup Performance Gap with Data Augmentation

Karan Goel, Albert Gu, Yixuan Li, Christopher Ré

arXiv:2008.06775v1cs.LGcs.AIcs.CVstat.ML

TL;DR

Classifiers can perform unevenly across subgroups, and hand-crafted augmentation may be infeasible when subgroup differences are hard to express. Model patching learns within-class subgroup transformations and trains on the resulting augmentations; CAMEL outperforms robust-optimization baselines on CelebA and Waterbirds and improves robust accuracy on ISIC.

  • Problem

    Classifiers optimized for average performance can be inaccurate on important subgroups, and hand-crafted augmentation may be impossible when subgroup differences are difficult to express.

    Source: §1 · p. 1

  • Method

    Model patching learns transformations between subgroups within a class and uses them as augmentations to train a classifier for subgroup robustness.

    Source: §1 · pp. 1–2

  • Results

    On CelebA and Waterbirds, CAMEL consistently outperforms state-of-the-art robust-optimization approaches, reducing subgroup performance gap by up to 10%; on ISIC, robust accuracy improves 11.7% over a group robustness baseline.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that model patching may be a useful tool for domain experts addressing unequal subgroup performance.

    Source: §5 · p. 10

    The theoretical analysis assumes correspondence between examples in different subgroups, formalized as coupled sets.

    Source: §3 · p. 6

Abstract

from arXiv · show

Classifiers in machine learning are often brittle when deployed. Particularly concerning are models with inconsistent performance on specific subgroups of a class, e.g., exhibiting disparities in skin cancer classification in the presence or absence of a spurious bandage. To mitigate these performance differences, we introduce model patching, a two-stage framework for improving robustness that encourages the model to be invariant to subgroup differences, and focus on class information shared by subgroups. Model patching first models subgroup features within a class and learns semantic transformations between them, and then trains a classifier with data augmentations that deliberately manipulate subgroup features. We instantiate model patching with CAMEL, which (1) uses a CycleGAN to learn the intra-class, inter-subgroup augmentations, and (2) balances subgroup performance using a theoretically-motivated subgroup consistency regularizer, accompanied by a new robust objective. We demonstrate CAMEL's effectiveness on 3 benchmark datasets, with reductions in robust error of up to 33% relative to the best baseline. Lastly, CAMEL successfully patches a model that fails due to spurious features on a real-world skin cancer dataset.

1 Introduction

Model patching learns transformations between subgroups within a class and uses them as augmentations to train classifiers for subgroup robustness. CAMEL improves subgroup performance across benchmark and real-world settings, while addressing cases where hand-crafted subgroup edits may be impossible.

1 Introduction

  • Model patching learns transformations that change subgroup identity while preserving class labels, then uses them as augmentations to encourage robustness to subgroup-feature variation.The framework has a transformation-learning stage followed by classifier training with controlled augmentations.Source: §1 · p. 1
  • Unlike hand-crafted augmentation heuristics, model patching learns subgroup differences and directly targets subgroup robustness with prescribed transformations.Hand-crafted augmentations may be impossible when subgroup differences are difficult to express.Source: §1 · pp. 1–2
  • The training stage combines a subgroup robust objective with a consistency regularizer that constrains predictions on original and augmented examples to be similar.The analysis relates the consistency loss to a mutual-information bound on classifier subgroup invariance.Source: §1 · p. 2
  • On CelebA and Waterbirds, CAMEL consistently outperforms state-of-the-art robust-optimization approaches, reducing subgroup performance gap by up to 10%; on ISIC, robust accuracy improves 11.7% over a group robustness baseline.The ISIC evaluation concerns benign lesions split by colored-bandage presence.Source: §1 · p. 2
  • On MNIST, CAMEL cuts robust error rate to a third of other approaches; its subgroup consistency regularizer improves robust accuracy by up to 2.5% over prior consistency losses.CAMEL also learns representations that are more invariant, as measured by mutual information estimates.Source: §1 · p. 2
  • Replacing CycleGAN with heuristic augmentations worsens the subgroup performance gap by 3.35%, while combining CAMEL with heuristic augmentations adds 1.5% robust accuracy.The consistency regularizer also improves robust accuracy by up to 2.5% over prior consistency losses.Source: §1 · p. 2

2 CAMEL: CycleGAN Augmented Model Patching

CAMEL patches a classifier by learning class-conditional transformations between subgroups, then training on original and augmented data with a subgroup-robust objective and consistency regularization.

2.1 Stage 1: Learning Inter-Subgroup Transformations

  • CycleGANs learn class-conditional transformations that translate examples between subgroups while preserving their class, generating augmented examples for training.Pairwise models can be used when there are more than two subgroups; multi-domain models are another option.Source: §2.1 · p. 3

2.2 Stage 2: Subgroup Robustness with Data Augmentation

  • Stage 2 trains the classifier on original and augmented data, combining subgroup robustness with consistency regularization to target worst-case subgroup performance and invariance to subgroup features.The learned transformations serve as black-box augmentations.Source: §2.2 · p. 4

2.2.1 A Subgroup Robustness Objective

  • SGDRO optimizes class-conditional worst-case subgroup robustness across classes, preserving the class-subgroup hierarchy that flat group robustness loses.The subgroup performance gap measures the difference between a class’s best- and worst-performing subgroups.Source: §2.2.1 · p. 4

2.2.2 Subgroup Invariance using a Consistency Regularizer

  • Self-consistency aligns predictions across subgroup-augmented examples, while translation consistency aligns original-example predictions with the average translated predictions.Self-consistency is the more important component; translation consistency also helps ignore potential CycleGAN artifacts.Source: §2.2.2 · p. 5

2 CAMEL: CycleGAN Augmented Model Patching

  • The method combines the SGDRO objective with consistency loss, weighted by the consistency-strength hyperparameter λ.The section presents CAMEL as a two-stage framework and outlines a theoretical analysis of its invariance guarantees.Source: §2 · p. 3, §2.2.2 · p. 5

3 An Information Theoretic Analysis of Subgroup Invariance

The analysis models correspondences among same-class examples from different subgroups as coupled sets. It relates prediction invariance across these sets to mutual information and shows how CAMEL’s losses bound that quantity.

3 An Information Theoretic Analysis of Subgroup Invariance

  • A coupled set contains one example per subgroup within a class, with examples sharing class features and differing in subgroup features.The paper notes that such corresponding examples may be hidden from the training distribution.Source: §3 · p. 6
  • The target invariance requires identical predictions within coupled sets, a stronger condition than CDAT’s class-conditional invariance because I(Ŷ; Z | [X]) ≥ I(Ŷ; Z | Y).CDAT minimizes a lower bound on I(φ(X); Z | Y).Source: §3 · pp. 5–6
  • Theorem 1 relates I(Ŷ; Z | [X]) to the Jensen-Shannon divergence of predictions on coupled sets.For two subgroups per class, the quantity can be upper-bounded by the CycleGAN and consistency losses.Source: §3 · p. 6
  • At the global optimum of trained CAMEL, predictions are conditionally independent of subgroup identity given the coupled set, and Stage 1 and Stage 2 losses together upper-bound the mutual-information objective.The consistency loss minimizes a proxy for the coupled-set divergence using augmented examples.Source: §3 · p. 6

4 Experiments

Experiments evaluate CAMEL’s subgroup robustness on benchmarks and skin cancer classification, then isolate the contributions of learned augmentations and the subgroup consistency regularizer. CAMEL improves benchmark accuracy and subgroup balance, while the ISIC results show improved robust accuracy and malignant-class accuracy.

4.1 Subgroup Robustness and Invariance on Benchmark Datasets

  • On all datasets, CAMEL improves aggregate and robust accuracy by up to 5.3%; on Waterbirds, it reduces the landbird subgroup gap by 10.22% versus GDRO.The comparison is summarized in Table 2.Source: Table 2 · p. 7, §4.1 · p. 7

Methods.

  • The experiments compare CAMEL with ERM, GDRO, and dataset-specific robustness baselines, using ResNet-50 classifiers fine-tuned from pretrained ImageNet weights.The evaluation spans four datasets and includes ablations of learned transformations and consistency losses.Source: §4 · p. 6, Methods · p. 7

4.1 Subgroup Robustness and Invariance on Benchmark Datasets

  • CAMEL is the only method reported to make the model invariant to dataset subgroups, assessed using estimated prediction–subgroup mutual information on MNIST-Correlation.Table 3 reports the estimated mutual information, for which lower is better.Source: §4.1 · p. 7, Table 3 · p. 7

4.2.1 Effect of Learned Augmentations

  • As λ rises from 20 to 200, CAMEL’s robust accuracy rises by 7% and its subgroup gap is 9.37% lower; subgroup pairing is 14.84% worse, while strong heuristic consistency harms performance.The heuristic augmentations can miss subgroup changes or distort class information.Source: §4.2.1 · p. 8

4.2.2 Analyzing the Subgroup Consistency Regularizer

  • On the small Waterbirds (landbird, water) subgroup, consistency regularization shrinks the train–validation loss gap versus GDRO; triplet-JS and KL replacements reduce robust accuracy by −2.5%.The subgroup consistency regularizer also improves over prior consistency losses when used with heuristic augmentations.Source: Figure 4 · p. 9, §4.2.2 · p. 9

4.2.3 Additional GAN Ablations

  • CAMEL improves the robust performance of each of three additional GAN base models, whether used alone for augmentation or integrated into model patching.The comparison is described in Appendix D.8.Source: §4.2.3 · p. 9

4.3 Real-World Application in Skin Cancer Classification

  • On ISIC, GradCAM illustrates that CAMEL removes the model’s reliance on the spurious bandage feature, shifting attention to the skin lesion instead [70].The experiment augments only the benign class, split into subgroups by colored-bandage presence; the malignant class has no subgroups.Source: §4.3 · p. 9

5 Conclusion

The conclusion presents model patching as a framework for addressing unequal classifier performance across subgroups by encouraging invariance to subgroup features. The authors say theoretical analysis and empirical validation suggest it may be useful to domain experts.

5 Conclusion

  • Model patching is presented as a framework for improving classifier subgroup robustness by encouraging subgroup-feature invariance.It addresses the question of how domain experts can fix classifiers with unequal subgroup performance.Source: §5 · p. 10
  • The authors support the framework with theoretical analysis and empirical validation.Source: §5 · p. 10
  • The authors conclude that model patching may be a useful tool for domain experts.Source: §5 · p. 10

Broader Impact

The broader-impact discussion connects subgroup performance failures to potential real-world bias and frames model patching as a step toward helping domain experts address them. The related-work material situates the approach among data-augmentation methods while distinguishing its focus on subgroup robustness.

Broader Impact

  • Unequal subgroup performance can contribute to discrimination or bias in applications; the authors cite underrepresented darker skin tones as a possible concern in skin-cancer classification.They note that skin-cancer datasets overwhelmingly contain images of light-skinned individuals.Source: Broader Impact · p. 11

B.1 Overview of Data Augmentation

  • Data augmentation is used in vision, NLP, and audio; its theoretical motivation is often framed as tangent constraints expressing desired model invariances.The paper’s evaluation is on image classification, though the authors say the ideas could extend to speech and text.Source: §B.1 · p. 17

B.2 Augmentation Primitives and Pipelines

  • Unlike policies built from generic image transformations, this work learns augmentation primitives aimed at subgroup robustness and reports that heuristic augmentation can complement CAMEL.The paper also distinguishes its robustness focus from prior learned-augmentation work aimed at generating examples for underrepresented classes or domains.Source: §B.2 · p. 18

B.3 Data Augmentation and Model Robustness

The section contrasts heuristic robustness methods with learned semantic subgroup transformations, then develops coupling-conditioned invariance and its information-theoretic motivation. It relates CycleGAN and Jensen–Shannon divergence to subgroup invariance objectives.

B.3 Data Augmentation and Model Robustness

  • Semantic data augmentation remains relatively underexplored for robustness, and CAMEL learns subgroup transformations rather than specifying heuristic augmentations [33].The paper distinguishes its goal—worst-case subgroup robustness—from prior augmentation work focused on out-of-distribution examples, and identifies subgroup-feature invariance as its novelty.Source: §B.3 · p. 19

B.4 Learning Robust Predictors

  • GDRO targets worst-case accuracy over predefined groups, but need not prevent subgroup-specific features; robust optimization can also face small-group generalization limits and trade off aggregate accuracy.The cited comparison notes that stronger feature assumptions may be needed for invariance, while those approaches suit specialized extreme out-of-domain settings.Source: §B.4 · p. 19

C.1 Background: CycleGAN

  • CycleGAN learns mappings between two domains from unpaired samples; cycle-consistency encourages near-inverse mappings, while identity loss encourages preservation of the original domain.In CAMEL, the domains correspond to pairs of subgroups within a class.Source: §C.1 · p. 19

C.2 Background: Properties of the Jensen-Shannon Divergence

  • For a uniform mixture of distributions, mutual information between the sample and its mixture indicator equals their Jensen–Shannon divergence; the optimal discriminator loss is JSD minus log 2.The JSD is also the square of a metric, giving its square root a triangle-inequality property.Source: §C.2 · p. 20

C.3 Subgroup Invariance using Coupled Distributions

  • Coupled sets pair equal-probability examples across a class’s subgroups, enabling a stronger invariance target: matching learned features across subgroup counterparts, conditioned on coupled-set identity.The paper assumes examples can be translated into other subgroups; the coupled-set construction separates shared class identity from subgroup identity.Source: §C.3 · p. 21

C.4 MI Bounds for Class-conditioned Invariance

  • Class-conditional domain adversarial training minimizes a lower bound on I(φ(X); Z | Y), using a conditional subgroup predictor whose cross-entropy corresponds to the bound up to a data-distribution constant.Attaching the domain-prediction head to the prediction layer yields a lower bound on I(Ŷ; Z | Y).Source: §C.4 · p. 22

C.5 MI Bounds for Coupling-conditioned Invariance

The analysis connects prediction consistency across coupled sets to subgroup information conditioned on coupled-set identity. It shows how generated coupled sets and CycleGAN losses bound the gap between the consistency objective on augmented examples and the true-set objective.

C.5 MI Bounds for Coupling-conditioned Invariance

  • The stronger coupling-conditioned objective includes a nonnegative term relating subgroup identity to coupled-set identity beyond class-conditioned feature invariance.The analysis interprets conditioning on a coupled set as revealing more information than conditioning on class alone.Source: §C.5 · pp. 22–23
  • In practice, generated coupled sets are noisy proxies for true sets and may contain artifacts; translation loss regularizes average augmented predictions toward the original example.The discrepancy between augmented-set and true-set JSD objectives can be bounded using pair-conditioned CycleGAN discriminator losses and JSD’s metric properties.Source: §C.5 · p. 23
  • For model predictions, the mutual information I(Ŷ; Z | [X]) equals the expected JSD of predictions across true coupled sets, and CAMEL’s global optimum induces Ŷ ⊥ Z | [X].With two subgroups per class, the quantity can be upper-bounded using CycleGAN and consistency losses.Source: §C.5 · pp. 23–24
  • The combined CycleGAN subgroup modeling and consistency regularizer minimize the desired identity-conditioned mutual information in the proof.The proof relates discriminator loss to divergence between original and generated subgroup distributions, then combines those bounds with consistency loss.Source: §C.5 · pp. 23–24

D Experimental Details

The appendix points to experimental-protocol and dataset information provided for reproducibility, including a table summarizing dataset split sizes.

D Experimental Details

D.1 Dataset Information

The appendix describes dataset construction and preprocessing for the experiments, including a controlled MNIST correlation setup and subgroup undersampling in CelebA.

D.1 Dataset Information

  • The appendix describes dataset preparation and says all datasets will be made available for download.Table 7 summarizes training, validation, and test example counts by dataset.Source: §D.1 · p. 25
  • The controlled MNIST-Correlation task predicts digit parity, with clean and zigzag subgroups drawn from MNIST and MNIST-Corrupted.The setup uses 40,000 examples with correlation ρ = 0.98, making class and subgroup highly correlated.Source: §D.1 · p. 25
  • CelebA-Undersampled reduces the training subgroup of non-blonde females from 71,629 examples to a random subset of 4,054.The subset size was chosen to equalize the ratio of subgroup sizes across the two classes.Source: §D.1 · p. 25
  • Waterbirds images are resized to 224 × 224 × 3 without center-cropping, retaining the full image and its background.The paper contrasts this with prior preprocessing that resized to 256 × 256 × 3 before center-cropping to 224 × 224 × 3.Source: §D.1 · p. 25
  • ISIC images are resized to 224 × 224 × 3 before use.Source: §D.1 · p. 26

D.2 CycleGAN Training Details

CycleGANs use Adam with a constant learning rate of 0.0002 and the default hyperparameters suggested by [97]. Training data, schedules, batch sizes, and augmentations vary by dataset.

D.2 CycleGAN Training Details

  • Across datasets, CycleGAN training uses Adam (β1 = 0.5), batch normalization, and a constant learning rate of 0.0002.The learning rate applies to both generators and both discriminators.Source: §D.2 · p. 26
  • MNIST-Correlation training uses 200 images each from MNIST and MNIST-Corrupted for 2500 epochs, with random rotation, padding, and cropping.Training uses batch size 25, cycle-loss coefficient 10.0, and identity-loss coefficient 1.0.Source: §D.2 · p. 26
  • CelebA-Undersampled uses separate class-specific CycleGANs trained for 4000 epochs on 1000 images per subgroup.Training uses batch size 16, cycle-loss coefficient 10.0, identity-loss coefficient 1.0, random flips, and crops up to 10%; generated samples were visually inspected at checkpoints because training was unstable.Source: §D.2 · p. 26
  • Waterbirds uses separate class-specific CycleGANs trained for 4000 epochs on 56 landbird and 184 waterbird images per subgroup.Both classes use batch size 4, cycle-loss coefficient 10.0, identity-loss coefficient 1.0, random flips, and crops up to 10%.Source: §D.2 · p. 26
  • ISIC benign subgroups each contribute 100 images to 4000-epoch CycleGAN training with batch size 4 and random flips and crops.The cycle-loss and identity-loss coefficients are both 10.0.Source: §D.2 · p. 26

D.3 Architectures and Training Information

The appendix specifies classifier training, model selection, hyperparameter sweeps, the mutual-information measurement, and baseline implementations. Classifier training avoids augmentation to isolate the method's effect.

D.3 Architectures and Training Information

  • Classifier training uses SGD with momentum 0.9 and no data augmentation; MNIST-Correlation uses a randomly initialized CNN, while other datasets fine-tune ImageNet-pretrained ResNet-50s.The MNIST-Correlation architecture has convolutional, pooling, dropout, and dense layers; other datasets share ImageNet normalization.Source: §D.3 · pp. 26–27

D.4 Hyperparameters

  • Models are selected by validation robust accuracy, then the selected hyperparameters are run three times and averaged for Table 2.ISIC additionally uses AUROC for model selection.Source: §D.4 · p. 27, §D.4.3 · p. 27

D.4.1 CelebA-Undersampled

  • CelebA-Undersampled sweeps run for 50 epochs, with CAMEL varying its consistency penalty while fixing learning rate, weight decay, and adjustment coefficient.CAMEL penalties are {5.0, 10.0, 20.0, 50.0}; ERM and GDRO sweep learning rates, and GDRO also sweeps adjustment coefficients.Source: §D.4.1 · p. 27

D.4.2 Waterbirds

  • Waterbirds sweeps run for 500 epochs, with CAMEL varying consistency penalties and learning rates across two configurations.ERM and GDRO sweep learning rates and weight decays; GDRO also varies adjustment coefficients.Source: §D.4.2 · p. 27

D.4.3 MNIST-Correlation

  • MNIST-Correlation sweeps run for 100 epochs, except CDAT, which runs for 400 because it takes longer to converge.ERM, GDRO, IRM, CDAT, and CAMEL sweep their respective learning rates, weight decays, penalties, or domain-loss coefficients.Source: §D.4.3 · p. 27

D.4.4 ISIC

  • ISIC sweeps run for 75 epochs, varying weight decay for ERM and learning rate, weight decay, consistency penalty, and annealing rate for CAMEL.GDRO sweeps learning rates and weight decays.Source: §D.4.4 · p. 28

D.5 Mutual Information Measurement

  • The mutual-information experiment adds a domain head that predicts subgroup identity from the final feature layer without sending gradients to the main model.Average cross-entropy loss, added to log 2, is reported as the measurement.Source: §D.5 · p. 28

D.6.1 Methods

  • Baselines include ERM cross-entropy, GDRO worst-subgroup-loss optimization, adversarial domain-information removal, and IRM across subgroup environments.GDRO uses subgroup information and an adjustment coefficient; CDAT uses gradient reversal, while IRM treats subgroups as separate environments.Source: §D.6 · p. 28, §D.6.1 · p. 28

D.6.2 Ablations

The appendix describes subgroup-pairing and heuristic-augmentation ablations, an artificial ISIC evaluation, and comparisons with alternative GAN augmentation methods.

D.6.2 Ablations

  • Subgroup Pairing enforces consistency between paired examples from different subgroups, while Heuristic Augmentations randomly apply one to three image operations.The heuristic pipeline includes geometric, color, and image-mixing operations.Source: §D.6.2 · p. 28

D.7 ISIC Spurious Correlations

  • CAMEL exceeds the other methods on the ISIC validation set whether model selection uses robust accuracy or AUROC.Table 8 provides the detailed evaluation.Source: §D.7 · p. 29, Table 8 · p. 29
  • An artificial ISIC dataset tests bandage–benign-class correlation by placing segmented benign-class bandages onto malignant images.The evaluation is used to illustrate how CAMEL affects dependence on this spurious feature.Source: §D.7 · p. 29

D.8 Alternative GAN Augmentation Baselines

  • Alternative GANs perform poorly as standalone augmentation but improve when integrated into model patching.The comparison considers Augmented CycleGAN, DAGAN, and StarGAN-v2; model patching replaces its Stage 1 CycleGAN with an alternative model.Source: §D.8 · p. 29, Table 9 · p. 29
Loading 2008.06775v1…