PaperlayerView exact revision on arXiv

Source-linked AI summary

Can Subnetwork Structure be the Key to Out-of-Distribution Generalization?

Dinghuai Zhang, Kartik Ahuja, Yilun Xu, Yisen Wang, Aaron Courville

arXiv:2106.02890v1cs.LGstat.ML

TL;DR

The paper examines whether neural-network structure can reduce reliance on spurious correlations and improve OOD generalization. It probes functional subnetworks, proposes MRM to learn their structure, and reports that suitable subnetworks and MRM can improve OOD performance.

  • Problem

    Models can exploit spurious correlations that do not generate the label, contributing to failures under distribution shifts; a structure-based approach beyond linear cases remains of interest.

    Source: §1 · p. 1

  • Method

    The paper probes functional subnetworks in deep models and proposes MRM to learn subnetwork structure from data for OOD tasks.

    Source: §1 · pp. 1–2, §3.1 · p. 3

  • Results

    The authors report that biased trained models contain less shortcut-susceptible subnetworks, and that a subnetwork can achieve better OOD with commensurate in-distribution accuracy in comparable iterations.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Subnetwork selection can improve OOD generalization, and MRM can be used with other methods such as IRM.

    Source: §1 · p. 2

    The oracle subnetwork uses out-domain data, so its comparison with baselines lacking that access is acknowledged as unfair and serves to illustrate structure's importance.

    Source: §4 · p. 6

Abstract

from arXiv · show

Can models with particular structure avoid being biased towards spurious correlation in out-of-distribution (OOD) generalization? Peters et al. (2016) provides a positive answer for linear cases. In this paper, we use a functional modular probing method to analyze deep model structures under OOD setting. We demonstrate that even in biased models (which focus on spurious correlation) there still exist unbiased functional subnetworks. Furthermore, we articulate and demonstrate the functional lottery ticket hypothesis: full network contains a subnetwork that can achieve better OOD performance. We then propose Modular Risk Minimization to solve the subnetwork selection problem. Our algorithm learns the subnetwork structure from a given dataset, and can be combined with any other OOD regularization methods. Experiments on various OOD generalization tasks corroborate the effectiveness of our method.

1. Introduction

The paper asks whether model structure can reduce reliance on spurious correlations under distribution shifts. It finds useful subnetworks inside biased models and proposes Modular Risk Minimization (MRM) to learn structures for better OOD generalization.

1. Introduction

  • The oracle subnetwork and MRM significantly surpass the full network on OOD performance in FULLCOLOREDMNIST.Source: Figure 1 · p. 1, §1 · p. 2
  • Even large trained networks that exploit spurious correlations can contain subnetworks that are less susceptible to those shortcuts.Source: §1 · p. 2
  • The functional lottery ticket hypothesis proposes that a full network may contain a subnetwork with better OOD and commensurate in-distribution accuracy when trained in isolation.The proposed comparison is in a comparable number of iterations.Source: §1 · p. 2
  • MRM learns subnetwork structure to improve OOD generalization and can be combined with other methods such as IRM.The authors report that this combination boosts performance as well.Source: §1 · p. 2

2. Invariant Prediction

The section frames OOD generalization as learning from seen environments for unseen distributions, where spurious correlations can make invariant prediction difficult. A motivating example argues that sparse classifiers can match spurious classifiers in training performance yet perform better OOD.

2.1. Out-of-distribution (OOD) generalization problem

  • OOD generalization seeks good performance on unseen environmental distributions despite training data coming only from seen environments.The paper uses “domain” and “environment” interchangeably.Source: §2.1 · p. 2
  • When spurious features correlate strongly with labels in training environments, models may favor them over invariant features, even when that correlation vanishes at test time.The example sets the spurious-label correlation to chance in unseen environments; high spurious-feature dimensionality can also promote reliance on those features.Source: §2.1 · p. 2

2.2. A Motivating Example

  • The motivating setup defines a sparse classifier as a maximum-margin classifier with at most d nonzero coefficients, a tractable constraint for linear models but challenging to learn in neural networks.The authors use this setup to motivate whether real neural networks contain sparse models that rely only on invariant features.Source: §2.2 · p. 3
  • When c or D is high, sparse and regular classifiers have similar training accuracy but the sparse classifier performs better OOD; margin comparisons imply logistic-loss gradient descent favors the spurious classifier (Soudry et al., 2018).The authors clarify that Proposition 1 is not intended to establish a general tradeoff between OOD performance and margin.Source: §2.2 · p. 3

3. A Functional Modularity Based Analysis

The paper probes pretrained networks for functional subnetworks and finds invariant digit modules even in models trained with spurious color correlations. In the analyzed first layer, color features are more pervasive than digit features.

3.1. Preliminaries

  • Modularity probing learns Bernoulli masks over network weights, using Gumbel-sigmoid training and logit regularization to identify task-specific subnetworks.Each layer's mask probabilities are sigmoid-transformed logits, thresholded at 0.5 to obtain a binary module.Source: §3.1 · pp. 3–4
  • On FULLCOLOREDMNIST, the two in-domains have color–digit bias coefficients of 1.0 and 0.9, while out-domain images have random colors.Digit shape is treated as invariant and color as spurious; the out-domain is used for module learning here and as a test distribution later.Source: §3.1 · p. 4
  • REx (76.17%) and DRO (78.56%) exceed ERM (58.04%) in OOD accuracy, while IRM reaches 59.55%.Figure 2 compares OOD accuracy across four algorithms and contrasts ERM with its digit module.Source: §3.1 · p. 4, Figure 2 · p. 4

3.2. Modular subnetwork introspection

  • Digit modules extracted from biased ERM models show good OOD generalization, indicating that trained networks can contain invariant subnetworks.The analysis searches for digit modules using one out-domain split and evaluates them on another.Source: §3.2 · p. 4
  • The first-layer visualization shows color features spread across more neurons than digit features, consistent with the paper's small-invariant-feature assumption.Figure 3 displays Bernoulli probabilities for the digit and color modules, from 0 to 1.Source: §3.2 · p. 4, Figure 3 · p. 5

4. Structure Matters: Towards Functional Lottery Tickets Hypothesis

The paper proposes that an appropriately initialized subnetwork can outperform its full network on OOD performance, then introduces MRM to learn such structures from training data. Its oracle subnetworks outperform ERM on OOD data while maintaining comparable in-domain performance, and ablations highlight the importance of initialization.

4. Structure Matters: Towards Functional Lottery Tickets Hypothesis

  • The functional lottery ticket hypothesis states that a suitably initialized subnetwork, trained in isolation on in-domain data, can surpass its full network's OOD performance.This is stronger than the original lottery ticket hypothesis, which requires similar full-network performance in the IID setting.Source: §4 · p. 5
  • On FULLCOLOREDMNIST, the oracle subnetwork substantially exceeds full-network ERM on OOD accuracy while maintaining indistinguishable in-domain performance.Its accuracy gap over random subnetworks remains considerable across sparsity levels; random subnetworks do not provide non-trivial benefit over ERM.Source: §4 · pp. 5–6, Figure 4 · p. 5
  • The oracle comparison uses out-domain data to select its subnetwork, unlike baselines without that privileged access, so it serves to illustrate the importance of structure.The authors identify this fairness issue and present MRM as a practical structure-searching approach that addresses the limitation.Source: §4 · p. 6, §4.1 · p. 6

4.1. Modular risk minimization

  • MRM pretrains a full model, probes it for a sparse predictive subnetwork using training data, then retrains only that subnetwork from scratch.The structure objective combines cross-entropy and sparsity regularization; MRM can also incorporate OOD losses such as IRM, REx, and DRO.Source: §4.1 · p. 6

4.2. Ablation for winning tickets learning

  • Reinitializing the winning-ticket subnetworks makes their OOD performance similar to full-network ERM, supporting the importance of reusing the original initialization.Resetting the color–digit bias relationship causes only a small accuracy drop for MRM, suggesting it finds a subnetwork with a robust inductive bias for the invariant function.Source: §4.2 · pp. 6–7, Figure 5 · p. 7

5. Experiments

Experiments evaluate MRM and OOD variants on three biased datasets, each with two training in-domains and one out-domain. MRM methods improve OOD performance across the tasks, including the more difficult SCENEOBJECT setting.

5. Experiments

  • The evaluation compares methods using the same model architectures and training settings across three datasets, each with two training in-domains and one out-domain.The datasets are FULLCOLOREDMNIST, COLOREDOBJECT, and SCENEOBJECT.Source: §5 · p. 7
  • On FULLCOLOREDMNIST, ModDRO reaches 85.53% accuracy, versus 78.25% for DRO and 99.0% for the unbiased solution.The method promotes OOD accuracy by around 10% on this task.Source: Table 1 · p. 7, §5 · p. 7
  • On COLOREDOBJECT, all four proposed methods exceed 50% accuracy and move their respective baselines toward the optimal unbiased solution.The task pairs ten object classes with ten color backgrounds and uses bias coefficients (0.8,0.6,0.0).Source: Table 2 · p. 7, §5 · p. 7
  • On SCENEOBJECT, MRM and its variants improve OOD performance where previous OOD algorithms bring very limited benefit.The task is highly biased, and the unbiased solution reaches 56.46% accuracy.Source: §5 · p. 8, Table 3 · p. 8

6. Related Work

Related work covers invariant prediction and OOD generalization, debiasing without explicit multiple environments, modularity, and pruning. The paper situates its approach among methods that impose or identify structure in model weights.

Out-of-distribution generalization

  • OOD and invariant-prediction work seeks stable predictions across unseen distributions, with approaches including ICP, IRM, REx, and group DRO.The cited literature also discusses limits and failure modes of invariant prediction methods, including nonlinear settings.Source: §6 · p. 8

Debiasing

  • Debiasing work addresses spurious biases in realistic datasets without an explicit multiple-environments setting.Examples include texture bias in ImageNet-trained models and methods using explicit debiasing procedures.Source: §6 · p. 8

Modularity

  • Modularity research includes explicit module-level approaches and finer-grained weight-level approaches, with this work belonging to the latter.The cited weight-level work includes pruning and module-selection methods.Source: §6 · p. 8

Pruning

  • Pruning research spans early model-compression methods and the lottery ticket hypothesis, which highlights the importance of initialization.A contrasting view holds that inherited weights are not important in practical settings.Source: §6 · p. 8

7. Discussion

The discussion attributes MRM’s OOD performance to a label-noise-free setting and a parameterization that can zero out spurious weights. It also identifies data settings not covered by the study.

Data settings

  • The study does not consider settings where group attribution is available, which the authors leave for future work.The authors state that omitting label noise in FULLCOLOREDMNIST enables structure learning to mine the true invariant feature.Source: §7 · p. 9

Success of MRM

  • MRM is effective in the label-noise-free setting because invariant features are perfectly predictive, and its parameterization explicitly zeros out spurious weight components.The added structure parameters do not increase the network’s expressive power; random sparse models do not achieve the same effect.Source: §7 · p. 9

A.1. Proof for Proposition 1

The proof contrasts a regular classifier with an optimal sparse classifier, then the discussion distinguishes this work from related optimization and compression settings. The supplied discussion also connects MRM’s weight selection to other methods and causal discovery.

A.1. Proof for Proposition 1

  • For d = 2, the optimal sparse max-margin classifier uses invariant weight 1 and spurious weight 0, achieving perfect accuracy in both train and test environments.The invariant feature equals the label in both environments.Source: §A.1 · p. 13

B. More Discussion

  • IRM games and methods with extra parameters are excluded to keep parameter complexity matched across compared methods.IRM games use one network per environment, so their parameter count scales with the number of environments.Source: B. More Discussion · p. 14
  • MRM learns a subnetwork architecture and updates only selected weights, unlike Dropout’s random gradient masking and training-only effect.MRM keeps its selected subnetwork for inference.Source: B. More Discussion · p. 14
  • The authors distinguish their OOD setting from compression studies of underrepresented IID groups, where parameters may be useful rather than spurious.Their subnetwork search aims to maximize in-domain performance so useful parameters are retained.Source: B. More Discussion · p. 14
  • The authors connect good subnetworks to faster convergence and better performance, while stating that causal discovery is not the goal of this work.They frame the connection as capturing causal relationships hidden in data more effectively.Source: B. More Discussion · p. 14

C.1. Dataset details

The section describes three image datasets built to study relationships between invariant object or digit features and spurious backgrounds. FULLCOLOREDMNIST uses ten colors without the majority/minority color split used in related work.

C.1. Dataset details

  • FULLCOLOREDMNIST contains 60,000 images of size 3 × 32 × 32, with ten colors and a one-to-one digit–color bias relationship.Its bias coefficient is the ratio of biased data following that relationship.Source: §C.1 · p. 14
  • COLOREDOBJECT contains 10,000 images of ten object classes on ten color backgrounds, each image sized 3 × 64 × 64.Unlike Ahmed et al. (2021), this dataset does not include a minority group with other backgrounds.Source: §C.1 · p. 15
  • SCENEOBJECT contains 10,000 images of the same ten object classes on ten scenery backgrounds, each image sized 3 × 64 × 64.Its backgrounds include settings such as beaches, deserts, and forests.Source: §C.1 · p. 15

C.2. Experimental details

The experiments use a finite hyperparameter search for baseline regularizers and subnetwork learning, alongside dataset-specific training settings. Oracle-validation results are provided for all three datasets.

C.2. Experimental details

  • Oracle-validation results are reported for FULLCOLOREDMNIST, COLOREDOBJECT, and SCENEOBJECT.The paper says these results are consistent with the main-text conclusions.Source: §C.2 · p. 15, Table 6 · p. 16
  • IRM and REx regularization coefficients, their introduction steps, and loss scaling are searched over finite candidate sets.DRO group-proportion step size is also searched over a finite set.Source: §C.2 · p. 15
  • Subnetwork learning searches logit learning rates from 1e-2 to 1 and sparsity coefficients from 1e-8 to 1e-3.FULLCOLOREDMNIST uses five random-seed trials and a three-convolution-layer ConvNet; COLOREDOBJECT and SCENEOBJECT use Wide ResNet 28-2.Source: §C.2 · p. 16

C.3. More empirical results

Additional experiments compare learned modules, pruning, and MRM ablations. They report faster color-module convergence than digit-module convergence, limits of an intersection module alone, and the importance of MRM’s final stage.

C.3. More empirical results

  • The color module converges faster than the digit module under ERM, and digit-module learning with three OOD algorithms shows a similar pattern.The authors interpret this as neural networks favoring color or texture over digit or shape information.Source: §C.3 · pp. 16–17, Figure 7 · p. 16
  • The intersection of digit and color modules is not predictive for digit identification by itself and must be combined with other parts.Its complement behaves trivially in the reported logical-operation experiments.Source: §C.3 · p. 17, Figure 8 · p. 17
  • MRM surpasses magnitude pruning at every sparsity level in the reported ERM experiments on FULLCOLOREDMNIST.The comparison is described as a completeness check; magnitude pruning is not designed for OOD generalization.Source: §C.3 · p. 17, Figure 9 · p. 17

N2 ACCURACY

  • The ablation reports the importance of MRM’s third stage and presents this as supporting the functional lottery ticket hypothesis.The accompanying table is an ablation of the third stage.Source: N2 ACCURACY · p. 18, Table 8 · p. 18

C.4. Towards a different pruning method

The authors adapt MRM into a pruning procedure for IID tasks, while explicitly distinguishing it from a novel pruning method. They report stronger results than other baselines, especially at 98% pruning.

C.4. Towards a different pruning method

  • For pruning, MRM’s second stage jointly trains the subnetwork structure and model parameters until the specified sparsity is reached.The adaptation retains the three-stage training paradigm.Source: §C.4 · p. 18
  • The authors report that their method outperforms other baselines, especially at a 98% pruning ratio on CIFAR10 with ResNet32.Table 9 presents the pruning results across pruning ratios.Source: §C.4 · p. 18, Table 9 · p. 18
  • The authors do not claim novelty as a pruning method and note a similarity to Louizos et al. (2018).They say this pruning adaptation is for IID generalization.Source: §C.4 · p. 18
Loading 2106.02890v1…