Exploring Optimal Substructure for Out-of-distribution Generalization via Feature-targeted Model Pruning

TL;DR

This paper introduces SFP, a feature-targeted model pruning method that enhances out-of-distribution generalization by removing reliance on spurious features, supported by theoretical analysis and superior experimental results.

Contribution

The paper proposes a novel feature-targeted pruning framework, SFP, which improves OOD generalization by focusing on invariant features and providing theoretical guarantees.

Findings

SFP outperforms state-of-the-art OOD methods with up to 4.72% accuracy improvement.

Theoretical analysis links model sparsity to OOD structure robustness.

SFP effectively reduces dependence on spurious features in biased datasets.

Abstract

Recent studies show that even highly biased dense networks contain an unbiased substructure that can achieve better out-of-distribution (OOD) generalization than the original model. Existing works usually search the invariant subnetwork using modular risk minimization (MRM) with out-domain data. Such a paradigm may bring about two potential weaknesses: 1) Unfairness, due to the insufficient observation of out-domain data during training; and 2) Sub-optimal OOD generalization, due to the feature-untargeted model pruning on the whole data distribution. In this paper, we propose a novel Spurious Feature-targeted model Pruning framework, dubbed SFP, to automatically explore invariant substructures without referring to the above weaknesses. Specifically, SFP identifies in-distribution (ID) features during training using our theoretically verified task loss, upon which, SFP can perform ID targeted-model pruning that removes branches with strong dependencies on ID features. Notably, by attenuating the projections of spurious features into model space, SFP can push the model learning toward invariant features and pull that out of environmental features, devising optimal OOD generalization. Moreover, we also conduct detailed theoretical analysis to provide the rationality guarantee and a proof framework for OOD structures via model sparsity, and for the first time, reveal how a highly biased data distribution affects the model's OOD generalization. Extensive experiments on various OOD datasets show that SFP can significantly outperform both structure-based and non-structure OOD generalization SOTAs, with accuracy improvement up to 4.72% and 23.35%, respectively.