SeWA: Selective Weight Average via Probabilistic Masking

TL;DR

SeWA introduces an adaptive, probabilistic checkpoint selection method for weight averaging that improves model generalization and convergence with minimal hyperparameter tuning, validated across multiple domains.

Contribution

The paper presents SeWA, a novel probabilistic framework for checkpoint selection in weight averaging, reducing manual tuning and improving performance.

Findings

SeWA achieves better generalization with fewer checkpoints.

Theoretical bounds show sharper stability guarantees than SGD.

Experimental results confirm effectiveness across domains.

Abstract

Weight averaging has become a standard technique for enhancing model performance. However, methods such as Stochastic Weight Averaging (SWA) and Latest Weight Averaging (LAWA) often require manually designed procedures to sample from the training trajectory, and the results depend heavily on hyperparameter tuning. To minimize human effort, this paper proposes a simple yet efficient algorithm called Selective Weight Averaging (SeWA), which adaptively selects checkpoints during the final stages of training for averaging. Based on SeWA, we show that only a few points are needed to achieve better generalization and faster convergence. Theoretically, solving the discrete subset selection problem is inherently challenging. To address this, we transform it into a continuous probabilistic optimization framework and employ the Gumbel-Softmax estimator to learn the non-differentiable mask for each checkpoint. Further, we theoretically derive the SeWA's stability-based generalization bounds, which are sharper than that of SGD under both convex and non-convex assumptions. Finally, solid extended experiments in various domains, including behavior cloning, image classification, and text classification, further validate the effectiveness of our approach.