Domain-Aware Pruning: Sparsity and Domain Generalization via Regularized Probabilistic Masking

2026-08-09

AI Summary

This work addresses the challenge of simultaneously improving model efficiency and out-of-distribution generalization by proposing Domain-Aware Pruning (DAP), a novel framework that uniquely treats neural network pruning as an intrinsic mechanism for domain generalization. DAP employs differentiable pruning via continuous probability masks and introduces a regularization term to suppress domain-sensitive weights, thereby automatically uncovering domain-invariant sparse subnetworks. Notably, it integrates seamlessly into existing domain generalization pipelines without requiring fine-tuning and is agnostic to the underlying algorithm. Extensive experiments across five benchmarks demonstrate that DAP achieves high sparsity while preserving or even surpassing the out-of-distribution performance of the original dense models, along with enhanced adversarial robustness and interpretability.

Abstract

Domain generalization (DG) and neural network pruning are conventionally treated as distinct objectives, targeting out-of-distribution (OOD) robustness and model efficiency, respectively. In this work, we bridge this gap by introducing Domain-Aware Pruning (DAP), a framework that leverages network sparsity as a mechanism to implicitly enhance generalization to unseen domains. Diverging from standard binary mask optimization, DAP learns a continuous parameter retention probability $p \in [0, 1]$, framing network compression as a continuous probabilistic masking problem. By introducing a regularization objective that actively penalizes the retention of domain-sensitive weights during the mask training, DAP identifies a domain-invariant subnetwork. Empirical results across five DG benchmark datasets demonstrate that DAP achieves significant sparsity while consistently matching or exceeding the OOD performance of its dense counterparts. Crucially, DAP is an algorithm-agnostic framework that integrates seamlessly with existing DG pipelines without necessitating post-hoc fine-tuning. Beyond efficiency and generalization, we show that DAP natively provides increased robustness to adversarial perturbations and yields highly interpretable models, where the retained weights reliably encapsulate the most domain-invariant and task-critical representations.

Problem

• domain generalization
• neural network pruning
• out-of-distribution robustness
• model sparsity
• domain-invariant representation

Innovation

• Domain Generalization
• Neural Network Pruning
• Probabilistic Masking
• Domain-Invariant Subnetwork
• Regularized Sparsity