Learning domain-invariant features through channel-level sparsification for Out-Of Distribution Generalization

TL;DR

This paper introduces Hierarchical Causal Dropout (HCD), a novel method that enforces feature sparsity at the channel level to improve out-of-distribution generalization by isolating causal features from spurious correlations.

Contribution

The paper proposes HCD, which uses channel-level causal masks and a mutual information objective to better separate causal from non-causal features for OOD robustness.

Findings

HCD outperforms existing methods on OOD benchmarks.

The approach effectively isolates causal features from spurious ones.

Experimental results demonstrate improved generalization across diverse datasets.

Abstract

Out-of-Distribution (OOD) generalization has become a primary metric for evaluating image analysis systems. Since deep learning models tend to capture domain-specific context, they often develop shortcut dependencies on these non-causal features, leading to inconsistent performance across different data sources. Current techniques, such as invariance learning, attempt to mitigate this. However, they struggle to isolate highly mixed features within deep latent spaces. This limitation prevents them from fully resolving the shortcut learning problem.In this paper, we propose Hierarchical Causal Dropout (HCD), a method that uses channel-level causal masks to enforce feature sparsity. This approach allows the model to separate causal features from spurious ones, effectively performing a causal intervention at the representation level. The training is guided by a Matrix-based Mutual Information (MMI) objective to minimize the mutual information between latent features and domain labels, while simultaneously maximizing the information shared with class labels.To ensure stability, we incorporate a StyleMix-driven VICReg module, which prevents the masks from accidentally filtering out essential causal data. Experimental results on OOD benchmarks show that HCD performs better than existing top-tier methods.