SAMPa: Sharpness-aware Minimization Parallelized

要約

Sharpness-aware Minimization (SAM) はニューラル ネットワークの一般化を改善することが示されています。
ただし、SAM の更新ごとに 2 つの勾配を \emph{順次} 計算する必要があるため、SGD などの基本オプティマイザーと比較して反復あたりのコストが事実上 2 倍になります。
我々は、SAMPa と呼ばれる SAM の簡単な修正を提案します。これにより、2 つの勾配計算を完全に並列化できるようになります。
SAMPa は、デバイス間の通信コストが無視できるという前提で、SAM の 2 倍の高速化を実現します。
経験的な結果は、SAMPa が計算時間の点で最も効率的な SAM バリアントの 1 つにランクされることを示しています。
さらに、私たちの方法は、視覚タスクと言語タスクの両方で一貫して SAM よりも優れています。
特に、SAMPa は理論的に \emph{固定} 摂動サイズであっても収束保証を維持します。これは新しいリアプノフ関数によって確立されます。
実際、私たちはこの収束保証を厳しい要件として扱うことで SAMPa に到達しました。このアプローチは一般に SAM ベースの手法の開発に有望であると考えられています。
コードは \url{https://github.com/LIONS-EPFL/SAMPa} で入手できます。

要約(オリジナル)

Sharpness-aware minimization (SAM) has been shown to improve the generalization of neural networks. However, each SAM update requires \emph{sequentially} computing two gradients, effectively doubling the per-iteration cost compared to base optimizers like SGD. We propose a simple modification of SAM, termed SAMPa, which allows us to fully parallelize the two gradient computations. SAMPa achieves a twofold speedup of SAM under the assumption that communication costs between devices are negligible. Empirical results show that SAMPa ranks among the most efficient variants of SAM in terms of computational time. Additionally, our method consistently outperforms SAM across both vision and language tasks. Notably, SAMPa theoretically maintains convergence guarantees even for \emph{fixed} perturbation sizes, which is established through a novel Lyapunov function. We in fact arrive at SAMPa by treating this convergence guarantee as a hard requirement — an approach we believe is promising for developing SAM-based methods in general. Our code is available at \url{https://github.com/LIONS-EPFL/SAMPa}.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google