Adaptive Stochastic Weight Averaging

TL;DR

Adaptive Stochastic Weight Averaging (ASWA) improves model generalization by selectively updating parameter averages based on validation performance, combining SWA with early stopping to address overfitting and underfitting issues.

Contribution

This work introduces ASWA, a novel method that adaptively updates model parameters only when validation performance improves, enhancing generalization over traditional SWA.

Findings

ASWA outperforms baseline models on 11 benchmark datasets.

ASWA achieves statistically significant improvements in generalization.

The method effectively balances overfitting and underfitting issues.

Abstract

Ensemble models often improve generalization performances in challenging tasks. Yet, traditional techniques based on prediction averaging incur three well-known disadvantages: the computational overhead of training multiple models, increased latency, and memory requirements at test time. To address these issues, the Stochastic Weight Averaging (SWA) technique maintains a running average of model parameters from a specific epoch onward. Despite its potential benefits, maintaining a running average of parameters can hinder generalization, as an underlying running model begins to overfit. Conversely, an inadequately chosen starting point can render SWA more susceptible to underfitting compared to an underlying running model. In this work, we propose Adaptive Stochastic Weight Averaging (ASWA) technique that updates a running average of model parameters, only when generalization performance is improved on the validation dataset. Hence, ASWA can be seen as a combination of SWA with the early stopping technique, where the former accepts all updates on a parameter ensemble model and the latter rejects any update on an underlying running model. We conducted extensive experiments ranging from image classification to multi-hop reasoning over knowledge graphs. Our experiments over 11 benchmark datasets with 7 baseline models suggest that ASWA leads to a statistically better generalization across models and datasets