Embracing Unknown Step by Step: Towards Reliable Sparse Training in Real World
Bowen Lei, Dongkuan Xu, Ruqi Zhang, Bani Mallick
TL;DR
Sparse training offers resource efficiency but hurts real-world reliability, particularly in OOD detection. The authors introduce MOON, an unknown-aware sparse-training method that adds a $K+1$ output dimension, modifies loss with an unknown-aware term, applies an auto-tuning schedule, and uses epoch-wise weight averaging to guide exploration toward reliable solutions without requiring OOD data. Theoretical results relate MOON to reduced unreliability in a Gaussian Mixture Feature Space, and extensive experiments across benchmarks show up to $8.4\%$ AUROC improvement while preserving accuracy and calibration. The approach is plug-in, broadly compatible with existing sparse-training methods, and offers practical gains for deploying sparse DNNs in resource-constrained environments. $MOON$ thus advances the readiness of sparse models for real-world reliability demands and open avenues for broader robustness research.
Abstract
Sparse training has emerged as a promising method for resource-efficient deep neural networks (DNNs) in real-world applications. However, the reliability of sparse models remains a crucial concern, particularly in detecting unknown out-of-distribution (OOD) data. This study addresses the knowledge gap by investigating the reliability of sparse training from an OOD perspective and reveals that sparse training exacerbates OOD unreliability. The lack of unknown information and the sparse constraints hinder the effective exploration of weight space and accurate differentiation between known and unknown knowledge. To tackle these challenges, we propose a new unknown-aware sparse training method, which incorporates a loss modification, auto-tuning strategy, and a voting scheme to guide weight space exploration and mitigate confusion between known and unknown information without incurring significant additional costs or requiring access to additional OOD data. Theoretical insights demonstrate how our method reduces model confidence when faced with OOD samples. Empirical experiments across multiple datasets, model architectures, and sparsity levels validate the effectiveness of our method, with improvements of up to \textbf{8.4\%} in AUROC while maintaining comparable or higher accuracy and calibration. This research enhances the understanding and readiness of sparse DNNs for deployment in resource-limited applications. Our code is available on: \url{https://github.com/StevenBoys/MOON}.