今回の論文
今回取り上げるのは、Pierre Foret、Ariel Kleiner、Hossein Mobahi、Behnam Neyshaburによる「Sharpness-Aware Minimization for Efficiently Improving Generalization」です。arXiv初版は2020年10月3日公開、会議発表はICLR 2021です。公開元はarXivおよびICLR、研究分野はニューラルネットワークの最適化と汎化です。論文の書誌情報と原論文v3の本文を参照しています。
選んだ理由は、モデル構造を大きくせずに、学習の進め方から未知データへの性能を改善する考え方を学べるためです。推論を軽く保ちたい専用の画像分類サービスでも、学習時に追加の計算を使う価値を検討する材料になります。
どんな技術か
SAM（Sharpness-Aware Minimization）は、現在の重みで損失が低いことに加え、重みを少し動かしても損失が急増しにくいことを目指す学習手法です。この記事のSAMは、画像の領域を切り出すSegment Anythingとは別の技術です。
損失を地形の高さに見立てると、通常の学習は低い場所を探します。SAMは、底だけが低い細い谷より、周囲も低い広い谷を探すイメージです。同程度の学習損失でも未知データへの性能が違うことに着目し、解の周辺まで学習の判断材料にします。
具体的には、まず損失が増える方向へ重みを一時的に動かし、移動先で計算した勾配を使って元の重みを更新します。入力画像を加工する方法でも、推論時に複数モデルを平均する方法でもありません。学習後は、得られた1組の重みで通常どおり推論できます。
課題
ニューラルネットワークは、学習データへ十分に適合できる自由度を持っています。しかし、学習損失を下げきったモデルが、未見の画像でもよく分類できるとは限りません。同じ程度に学習データを説明する重みが複数あり、それぞれの汎化性能が異なるためです。
データ拡張や正則化も過学習を抑えますが、通常の損失最小化だけでは、重みの周辺にある損失の急な増加を直接は評価しません。SAMは、こうした既存の工夫と組み合わせながら、学習で到達する解の性質へ働きかけます。
たとえば商品画像の分類器で、学習画像では高精度なのに、新しく撮影した画像では誤りが増える場面が考えられます。まずデータ分割や撮影条件の偏りを確認すべきですが、適切に分割しても過学習が残るなら、最適化方法を変える余地があります。
ただし、「重みの変化に強い」ことと「照明やカメラの変化に必ず強い」ことは同じではありません。SAMが直接扱うのは重み空間の近傍であり、業務上の分布変化への効果は別途評価する必要があります。
用語解説
- 汎化と過学習
- 汎化は、学習に使っていないデータでもよい予測ができることです。学習データに合わせすぎて未知データで性能が落ちる状態を過学習と呼びます。SAMが改善したいのは、学習損失そのものより最終的な未知データへの性能です。
- 損失地形と鋭さ
- モデルの重みを動かしたときの損失の変化を地形として捉えたものです。SAMでは、一定範囲で重みを動かした際に損失がどれだけ増えるかが鋭さの尺度になります。単に損失が低い解と、周囲も低い解を区別するための概念です。
- 勾配と逆伝播
- 勾配は、各重みを少し変えたときに損失が増える方向と変化率を表すベクトルです。逆伝播でこれを計算します。SAMは最初の勾配を一時的な移動に、2回目の勾配を実際の学習更新に使います。
- 摂動とL2ノルム
- 摂動は重みへ一時的に加える小さな変更です。L2ノルムは、各成分の二乗和の平方根で測るベクトル全体の大きさです。SAMの代表的な設定では、摂動全体の大きさを半径rhoで制限します。
- ミニバッチ
- 1回の更新に使う学習例のまとまりです。SAMの実装では、このまとまりから損失を増やす方向を推定します。全データで1つの摂動を作る場合と、小さなまとまりごとに作る場合では、学習する目的が変わります。
技術の仕組み
現在地の損失から、近傍で最大の損失へ目的を変えます
重みをw、学習データ上の平均損失をL(w)、摂動をepsilonとします。L2ノルムを使う場合、SAMの目的は次のように書けます。
通常の目的: L(w) を小さくする
SAMの目的: max L(w + epsilon) + lambda × ||w||_2^2 を小さくする
ただし ||epsilon||_2 <= rho内側のmaxは、「半径rhoの範囲で、最も損失が大きくなる重みの動かし方」を探します。外側の最小化は、その悪い場合の損失を下げます。lambdaの項は重みの大きさを抑えるL2正則化です。
近傍にはepsilon=0も含まれるので、周囲の最大損失だけを下げて現在地の損失を無視する設計ではありません。周囲まで含めて低い損失を求めることで、鋭い谷へ入り込むのを抑えようとします。原論文2節・式1
最悪方向を、1回の勾配から近似します
ニューラルネットワークについて、近傍の本当の最大損失を毎回厳密に求めるのは困難です。そこでSAMは、現在地で損失を一次近似します。L2制約なら、その近似に対して最も損失を増やす摂動は、勾配方向を長さrhoへそろえたものです。
g = grad L_B(w)
epsilon_hat = rho × g / ||g||_2L_Bは、今回のミニバッチBの損失です。たとえば勾配が(3, 4)、rho=0.05なら、勾配の長さは5なので摂動は(0.03, 0.04)になります。これは式を理解するための計算例です。各成分へ一律に0.05を加えるのではなく、ベクトル全体の長さを0.05にします。
ここで求めるのは、あくまで一次近似に対する最悪方向です。非線形な損失の真の最大点を見つけた保証はありません。また勾配がゼロに近い場合の除算には、実装上の数値的な保護が必要です。
移動先の勾配を、元の重みに適用します
次に同じミニバッチを使い、w+epsilon_hatで再び勾配を計算します。基本のSGD形式を、正則化を省略して書くと次の流れです。
1. 元の重みwで勾配gを計算する
2. gを正規化し、一時的にw + epsilon_hatへ移動する
3. 移動先で勾配g_samを計算する
4. 重みをwへ戻す
5. w_new = w - learning_rate × g_sam重要なのは、摂動を学習の更新として残さないことです。勾配を測る場所は移動先ですが、更新の起点は元のwです。モーメンタムなどのオプティマイザ状態を進めるのも、基本的には最後の更新です。
論文は計算を軽くするため、摂動の作り方自体をさらに微分する二階の項を省きます。移動先で通常の逆伝播をするため、巨大な二階微分行列を構築する必要はありません。それでも1更新に順伝播・逆伝播を2回ずつ要します。原論文2節・式2〜3、Algorithm 1
分散学習では、摂動を作るデータの単位も効きます
原論文の複数アクセラレータでの実装は、各装置の部分バッチから独立に摂動とSAM勾配を作り、最終的な勾配を平均します。最初の勾配を全装置で平均して共通の摂動を作る方法とは異なります。
論文は、摂動を推定する例数mに依存した鋭さを「m-sharpness」と呼びます。小さなResNetでの分析では、mが小さいほどよい汎化につながる傾向が見られました。したがって、GPU数や勾配蓄積を変える際に、実効バッチ数だけをそろえて再現できるとは限りません。原論文4.1節
実験と結果
CIFARでは、通常学習を長くする条件とも比較しています
論文はCIFAR-10・CIFAR-100の画像分類で、モデルとデータ拡張の条件ごとにSAMとSGDを比較しています。表1から、WideResNet-28-10の例を抜き出します。数値はテスト誤り率で、低いほど良好です。
| SAM側の学習条件 | データセット | SGD | SAM |
|---|---|---|---|
| 200エポック・基本拡張 | CIFAR-10 | 3.5 ± 0.1% | 2.7 ± 0.1% |
| 200エポック・基本拡張 | CIFAR-100 | 18.8 ± 0.2% | 16.5 ± 0.2% |
| 1,800エポック・AutoAugment | CIFAR-10 | 2.2 ± 0.1%未満 | 1.6 ± 0.1% |
| 1,800エポック・AutoAugment | CIFAR-100 | 16.1 ± 0.2% | 12.8 ± 0.2% |
値は5回の独立試行の平均と95%信頼区間です。基本拡張は左右反転、パディング、ランダム切り出しを含みます。AutoAugmentの行にも基本拡張を使用しています。
SAMは2回の逆伝播を使うため、このCIFAR比較では、SGD側にも標準の2倍のエポック数まで学習する機会を与え、標準と倍のうち良いスコアを報告しています。つまり、表の「200エポック」はSAM側の設定で、SGD側を必ず200エポックに固定した比較ではありません。また、実測の所要時間を完全にそろえた比較でもありません。
rhoは学習データの10%を使う検証で候補から選んでいます。CIFAR-100の基本拡張では2.3ポイント改善し、強いデータ拡張を使っても改善が残りました。既存の過学習対策があっても、最適化方法に追加の改善余地があったと読めます。原論文3.1節・表1
ImageNetと転移学習でも改善しています
ImageNetでは、ResNet-152を400エポック学習した同じ行で、Top-1誤り率が通常学習の20.9%からSAMの18.4%へ下がっています。通常学習は200エポックで20.3%、400で20.9%となる一方、SAMは18.5%から18.4%です。この設定では、長く学習しても精度を保つ効果が見られます。ただし同じエポック数でもSAMの勾配計算は多く、同一計算予算の比較ではありません。原論文表2
ImageNetで事前学習したEfficientNet-b7をFlowersへ適応する実験では、Top-1誤り率が1.16%から0.63%へ改善しました。これはゼロからの学習に限らず、既存モデルのファインチューニングにも効果があった例です。追加データで事前学習したEfficientNet-L2の結果とは、事前学習条件を分けて考える必要があります。原論文3.2節・表3
ラベルノイズへの効果にも条件があります
CIFAR-10の学習ラベルの一部をランダムに変更し、テストラベルは正しいままにする評価も行っています。ResNet-32を200エポック学習し、40%ノイズの条件でテスト正解率はSGDが68.8%、SAMが93.4%でした。一方、同じ表の既報手法MentorMixは94.2%で、SAM単独がすべての方法を上回るわけではありません。
また、ノイズ評価の設定選択には正しいラベルの検証データを使っています。実業務の系統的な誤ラベルや、あるクラスだけの取り違えまで検証した結果ではありません。「ラベルを確認しなくても大丈夫」という結論にはできません。原論文3.3節・表4、付録C.1
何に使える？
以下は原論文の仕組みと画像分類実験から考えられる応用案です。各業務での有効性が検証済みという意味ではありません。
推論負荷を増やしにくい商品・外観分類
商品カテゴリや製造品の外観を分類する専用モデルでは、推論のたびに大きなモデルを動かす余裕がない場合があります。SAMなら、同じ構造のまま学習方法を変えるため、追加の学習計算と引き換えに未知画像への性能改善を試せます。
実験時には、同じ製品や同じ撮影日に由来する似た画像を学習と評価へ混ぜないようにします。そのうえで、撮影条件ごとの誤り率や見逃しを比較すると、単なるデータ漏洩による高精度と区別できます。
少ない業務画像への転移学習
事前学習モデルを、自社の商品や部品などの分類へ適応する場面も候補です。Flowersなどで転移学習の改善が示されたため、学習画像への過適合が残るときに比較する技術的な根拠があります。
ただし、正例のない欠陥を発見する仕組みではありません。また、クラス不均衡が中心課題なら、Focal Lossによる簡単な例の影響の抑制のように、例ごとの損失を設計する方法も比較対象です。SAMが重みの近傍を扱うのに対し、Focal Lossは各例から学習へ伝わる影響を変えます。併用効果は、このSAM論文では示されていません。
ラベル品質にばらつきがある画像データの学習
過去の運用で蓄積した画像には、担当者間の判断差や入力ミスが含まれます。ラベルを監査しつつ、通常学習との比較でSAMを試す余地があります。ノイズを含むデータで過学習しにくい傾向が得られるかを、正しいラベルを確認した独立評価セットで測ります。
SAMは誤ラベルの位置や正しいラベルを返す技術ではありません。学習方法の変更とデータ修正を、別々の改善策として評価する設計が必要です。
開発や事業へのヒント
モデル規模より先に、同じ構造の学習を比較します
小規模なプロダクトでも、モデル構造とデータ分割を固定し、通常のオプティマイザとSAMを比較できます。rho=0.05は原論文で良好な初期候補とされていますが、採用を前提にせず、学習率や正則化も含めて検証します。
比べたいのは、最終精度だけではありません。同じエポック数の比較と同じGPU時間の比較を分けると、SAM自体の更新効果と、追加計算を払う価値を判断しやすくなります。精度差が小さい場合には、複数シードでのばらつきも確認します。
学習の追加費用を、推論回数と一緒に考えます
定期的に学習し、多数のリクエストへ同じモデルを配布するサービスでは、学習時に計算を増やし、推論構造を保つ選択に意味があります。一方、顧客ごとに頻繁に再学習するサービスでは、追加の学習時間が提供速度や費用を圧迫する可能性があります。
メモリ不足で逆伝播自体が難しい場合は、SAMとは別の制約です。MeZOの逆伝播を使わない重み更新は、その制約に対する異なる選択肢を説明しています。両者とも重みを一時的に動かしますが、SAMは勾配で悪い方向を探し、MeZOはランダム方向の損失差で更新方向を推定するため、目的も必要なメモリも異なります。
2回の計算を、単純な2回の学習更新にしません
実装では、摂動の追加・復元、勾配の消去、実際の更新を明確に分けます。2回とも通常のoptimizer.stepを実行すると、Algorithm 1とは違う処理になります。
さらにBatch Normalizationの移動統計やランダムなデータ拡張など、重み以外の状態も確認します。2回の順伝播で状態がどう変わるか、同じ学習例をどう再利用するかを決めないと、重みの更新式だけを合わせても比較条件が変わります。これは仕組みから導かれる実装上の確認事項です。
限界
最も直接的な負担は計算量です。標準的なSAMは1更新に2回の順伝播・逆伝播を必要とします。所要時間が厳密に2倍になるわけではありませんが、同じ更新数なら通常学習より重くなります。重みの復元用データなども必要で、推論相当のメモリへ学習を縮める技術ではありません。
rho、学習率、正則化、摂動を作る部分バッチの大きさによって結果は変わります。特に分散学習では、最初の勾配をいつ同期するかが目的に関わります。GPU数や勾配蓄積の変更を、性能に影響しない実装差と見なすことはできません。
また、実際のSAMは近傍全体の最大損失を厳密には解いていません。一次近似と二階項の省略を使う実用アルゴリズムです。論文の汎化上界にも仮定があり、「平らな解なら任意のデータで必ず高精度」という保証にはなりません。
原論文の主要な実証は画像分類です。日本語の文章分類、LLMの命令調整、RAGの回答品質に同じ改善幅があるとは言えません。外部APIにプロンプトを送るだけのアプリでは、その内部の学習更新をSAMへ変更することもできません。
最後に、ランダムなラベル変更への頑健性を、分布外の入力や画像への敵対的な変更に対する安全性へ読み替えてはいけません。運用で問題になる条件を別に定義し、正しいラベルの評価データで確認する必要があります。
よくある質問
Q. SAMは入力画像にノイズを加える技術ですか？
A. この論文では、損失が増える方向へモデルの重みを一時的に動かします。入力を変更する敵対的学習とは、摂動を加える対象が異なります。画像への攻撃に対する頑健性が、この学習だけで保証されるわけではありません。
Q. SGDやAdamをSAMへ完全に置き換えるのですか？
A. SAMは更新に使う勾配の求め方を変え、最後の更新にはSGDなどの基本オプティマイザを使います。最初の勾配は摂動の作成用、2回目は更新用です。単に学習率を変える操作とも異なります。
Q. 推論も2回実行する必要がありますか？
A. ありません。2回の計算が必要なのは学習更新です。学習後は通常の1組の重みで推論でき、SAMの摂動探索は行いません。同じモデル構造なら、この手順による推論計算の追加はありません。
Q. rhoは大きいほどよいですか？
A. そうとは限りません。大きくすると広い近傍を考えますが、現在地の勾配による一次近似が適切でなくなる可能性もあります。原論文の0.05は候補の一つで、評価データ上の誤り率と学習の安定性を見て選びます。
Q. ラベルの誤りが多くてもSAMで解決できますか？
A. ランダムな誤ラベルを加えたCIFAR-10では改善しましたが、誤ラベルを特定・修正する機能はありません。設定選択には正しい検証ラベルも使っています。特定のクラスだけで誤るような業務データでは、監査と独立評価を併せて行ってください。
今日の学び
この論文は、学習損失が低くても未知データへの性能が十分でない問題を扱いました。SAMは、重みを少し動かした近傍での悪い損失を、2回の勾配計算で近似して下げようとします。
得られるヒントは、モデルを大きくする前に、どのような性質の重みへ学習を導くかを見直すことです。推論構造を保ちながら改善を試せる一方、追加の学習計算、データ条件、更新手順まで含めて価値を評価する必要があります。