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

arXiv: https://arxiv.org/abs/2608.08624 · HTML: https://mlsanity.com/paper/2608.08624/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 31 of its 40 references that are in the map.
- First posted on arXiv 9 Aug 2026.

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Moment Matching for Multi-Source Domain Adaptation](https://mlsanity.com/paper/1812.01754/) (2018)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [Deep Hashing Network for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1706.07522/) (2017)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Pruning Filters for Efficient ConvNets](https://mlsanity.com/paper/1608.08710/) (2016)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Deeper, Broader and Artier Domain Generalization](https://mlsanity.com/paper/1710.03077/) (2017)
- [Channel Pruning for Accelerating Very Deep Neural Networks](https://mlsanity.com/paper/1707.06168/) (2017)
- [Learning Efficient Convolutional Networks through Network Slimming](https://mlsanity.com/paper/1708.06519/) (2017)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [Importance Estimation for Neural Network Pruning](https://mlsanity.com/paper/1906.10771/) (2019)
- [SWAD: Domain Generalization by Seeking Flat Minima](https://mlsanity.com/paper/2102.08604/) (2021)
- [Can Subnetwork Structure be the Key to Out-of-Distribution Generalization?](https://mlsanity.com/paper/2106.02890/) (2021)
- [The Generalization-Stability Tradeoff In Neural Network Pruning](https://mlsanity.com/paper/1906.03728/) (2019)
- [In Search of Forgotten Domain Generalization](https://mlsanity.com/paper/2410.08258/) (2024)
- [Rethinking Domain Generalization: Discriminability and Generalizability](https://mlsanity.com/paper/2309.16483/) (2023)
- [Accelerating Deep Unsupervised Domain Adaptation with Transfer Channel Pruning](https://mlsanity.com/paper/1904.02654/) (2019)
- [Model Compression using Progressive Channel Pruning](https://mlsanity.com/paper/2507.04792/) (2025)

## Abstract (arXiv)

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 ∈ [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.