💬 NLP

Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting

本論文は、Sharpness-Aware Minimization、大きな学習率、または短縮されたアニーリング期間といった手法を通じて事前学習の最適化をより平坦な極小値へと偏らせることが、さまざまなモデルサイズおよび事後学習タスクにおいて、破滅的な忘却を大幅に軽減し、下流タスクのパフォーマンスを向上させることを示している。

原著者： Ishaan Watts, Catherine Li, Sachin Goyal, Jacob Mitchell Springer, Aditi Raghunathan

公開日 2026-05-06

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Ishaan Watts, Catherine Li, Sachin Goyal, Jacob Mitchell Springer, Aditi Raghunathan

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

マスターシェフ（AI モデル）を、数百万種類もの異なる料理（事前学習）を味わわせることで訓練していると想像してください。目標は、特定の飲食店のキッチン（事後学習/微調整）で働く前に、可能な限り最高のシェフに育てることです。

長らく、研究者たちは初期の訓練段階において、シェフが「最高」の味覚を持つことだけが重要だと考えていました。シェフが最初から「完璧な」万能職人としてスタートすれば、キッチンでどのような変化が起きても、後から自動的に優れた専門職人になれると想定していたのです。

問題点：「硬直した」シェフ
この論文は、その仮定が誤っていることを主張しています。実は、一般的な味覚においてあまりにも硬直した「完璧さ」で訓練されたシェフは、新しいことを求められたときに簡単に崩壊してしまうことが判明しました。

これを谷に置かれたボールに例えて考えてみましょう。

標準的なアプローチ（AdamW）： これはシェフを非常に深く、狭く、鋭い谷に座らせるように訓練します。彼らは正確にその場所に留まることには非常に優れています。しかし、わずかに押されると（新しいレシピを学ばせたり、記憶を圧縮して小さなポケットに収めさせたりするように）、すぐに谷から転がり落ちてしまいます。彼らは知っているすべてのことを「忘却」してしまいます。
この論文の解決策（シャープネスアウェア）： これはシェフを広く、平坦な谷に座らせるように訓練します。彼らは谷の絶対的な最深点にいるわけではありませんが、平坦な地面に囲まれています。彼らを少し押しても、転落することはありません。少し転がって安全な場所に留まります。彼らは「頑健」です。

研究者たちが行ったこと
チームは、「鋭い谷」のシェフではなく、「平坦な谷」のシェフを訓練するための 3 つの方法をテストしました。

「揺らす」方法（SAM）： シャープネスアウェア・ミニマゼーションと呼ばれる特殊な技術を使用しました。シェフが学習している間、テーブルを優しく揺らしていると想像してください。テーブルが揺れている間でもシェフが正しく料理の味を識別できれば、安定して学習していることになります。これにより、モデルは広く平坦な谷を見つけるように強制されます。
「高速レーン」方法（高い学習率）： シェフが学習の初めに、より速く、より攻撃的に学習するようにしました。これはレースを走るようなものです。十分に速く走れば、道にできた小さな深い穴に自然とハマることを避けることができます。
「早期停止」方法（短いアニーリング）： 通常、訓練はシェフの速度を徐々に落として終了します。研究者たちは、この減速を早期に停止させる（ペースをより長く維持する）ことが、シェフをその安定した平坦な領域に留めるのに役立つことを発見しました。

結果
彼らは異なるサイズのモデル（小規模から中規模大規模まで）でこれをテストし、以下の結果を得ました。

忘却の減少： これらの「平坦な谷」モデルが後に特定のタスク（数学やコーディングなど）を学ぶよう求められたり、実行速度向上のために記憶が圧縮（量子化）されたりした際、元の知識を大幅に忘れませんでした。
「訓練途中」のハック： 高価な特殊な「揺らし」方法を訓練プロセス全体で使用する必要はありません。彼らは、訓練の最後の 10%（「アニーリング」段階）だけでこれを使用すれば、コストのほんのわずかな割合でほぼすべての恩恵を得られることを発見しました。
実世界での証明： 彼らは 10 億パラメータの大型モデル（OLMo-2-1B）でこれをテストしました。この方法で訓練されたモデルは、最初は一般的なタスクにおいてわずかに「弱く」見えたものの、数学を学習した後のスキル保持において31% 向上し、効率化のために圧縮された後のスキル維持において40% 向上しました。

最大の教訓
この論文は結論として、AI をスタートラインで「最高」になるように訓練するだけではいけないと述べています。後で環境を変えたときに壊れないよう、柔軟で安定したように訓練すべきです。

床が傾いたときに壊れてしまう、完璧で硬直したポーズを保持するよう体操選手を訓練することと、どんな新しい動きにも転倒せずに適応できる、強力でバランスの取れた重心を持つよう訓練することの違いです。「完璧さ（シャープネス）」ではなく、「バランス（平坦さ）」に焦点を当てることで、古い知識を忘れずに新しいことを学ぶのに優れた AI モデルが得られるのです。

技術的サマリー：シャープネスアウェアな事前学習は破滅的忘却を軽減する

問題提起

現在の事前学習最適化戦略は、通常、ベースモデルの品質を最大化するように調整されており、その指標は事前学習損失やベンチマーク性能である。これは、より強力なベースモデルが、微調整（ポストトレーニング）や量子化のようなその後の改変を経て、より強力なモデルを生み出すという仮定に基づいている。しかし、最近の研究は、この仮定が失敗しうることを示している：ある一定の点を超えて事前学習を延長すると、事前学習損失は改善されるものの、ポストトレーニング後の下流タスクのパフォーマンスは低下する可能性がある。

特定された核心的な問題は、ベースモデルの評価がパラメータ更新下におけるモデル能力の安定性を無視している点にある。これらの更新に対して極めて敏感なモデルは、「破滅的忘却」に陥り、初期の強さに関わらず事前学習された能力を失う。この感度は損失地形の幾何学と関連しており、具体的には、「シャープな」極小値に収束するモデルは、下流タスク中にパラメータが攪乱された際に、パフォーマンスの低下を受けやすくなる。

手法

著者らは、モデルをフラットな極小値へとバイアスし、局所パラメータ攪乱に対する感度を低下させるように設計された 3 つの事前学習最適化アプローチを検証する：

シャープネスアウェア最小化（SAM）： 損失の曲率にペナルティを課す明示的な最適化手法。SAM は、半径 $\rho$ の近傍内での最大損失を最小化するパラメータ $\theta$ を探索し、実質的に小さな攪乱下でも低損失のままとなる極小値を見つける。
ピーク学習率の増大： 「安定の縁（Edge of Stability）」現象に基づく暗黙的なアプローチ。ここで、より大きな学習率はヘッシアン行列の最大固有値を暗黙的に抑制し、それによって損失のシャープネスを低下させると仮定されている。
学習率減衰の短縮： 学習率スケジューリングを変更し、トレーニング終盤の減衰フェーズの期間を短縮する。これは、後期段階のダイナミクスが解の最終的な幾何学に大きく影響するという発見に動機付けられている。

本研究は、2000 万から 1 億 5000 万パラメータにわたるモデルサイズでこれらの介入を評価し、10 億パラメータモデル（OLMo-2-1B）においてスケールを拡大して検証する。評価は 2 つの主要なトレードオフに焦点を当てる：

学習 - 忘却トレードオフ： 新しい能力の獲得（微調整損失）と、事前学習された能力の保持（微調整後の事前学習損失）の間のバランス。
圧縮 - 忘却トレードオフ： 量子化（例：4 ビット）を通じてモデルが圧縮された際の能力の低下。

著者らはまた、これらの手法が実際に微調整方向に沿った方向性シャープネスを低下させるか、および損失の 2 次テイラー近似が微調整攪乱に対して有効かを検証するためにヘッシアンを分析する。

主要な貢献と結果

1. 改善された学習 - 忘却トレードオフ

80 回以上の事前学習実行と 3,500 回以上の微調整実験を通じて、著者らはシャープネスアウェアな介入が一貫して学習 - 忘却トレードオフを改善することを発見した：

SAM の性能： SAM で事前学習されたモデルは、標準的な AdamW に比べて優れたパレートフロンティアを達成する。例えば、StarCoder データセットにおいて、SAM はマッチングされた微調整損失において、忘却を**80%**削減する。
スケーリング効果： 事前学習トークン予算が増加するにつれて、AdamW に対する SAM の優位性は増大する。AdamW のチェックポイントはトレーニングが進むにつれてより敏感になり（破滅的過学習）、SAM はこれを軽減し、高トークン領域で最も顕著な利益を示す。
モデルサイズ： トークン対パラメータ比を制御する場合、モデル規模（2000 万から 1 億 5000 万パラメータ）が増大しても利益は持続し、さらに増大する。

2. 微調整を超えた堅牢性

シャープネスアウェアな事前学習によって付与される堅牢性は、他の下流攪乱にも及ぶ：

量子化： SAM でトレーニングされたモデルは、4 ビット量子化下でパフォーマンスの低下が著しく少ない。高トークン予算領域において、SAM はベースラインと比較して、量子化に起因する損失増加を2〜3 倍少なくする。
ガウス攪乱： SAM でトレーニングされたモデルは、重みに追加された等方性ガウスノイズに対して感度が低く、より堅牢な局所幾何学を示している。

3. 暗黙的手法と実用的なスケーラビリティ

学習率ダイナミクス： ピーク学習率の増大と減衰期間の短縮の両方が、ベースモデルの事前学習損失がわずかに悪化する場合でも、学習 - 忘却トレードオフを改善する。
トレーニング中期の介入： SAM の計算コスト（ステップあたりの計算量が約 2 倍になる）に対処するため、著者らは SAM を学習率減衰フェーズのみに適用することを提案する。
- OLMo-2-1B における実験では、トレーニングの最終 10%（トレーニング中期）のみに SAM を適用することで、標準的なベースラインと比較して、MetaMath ポストトレーニング後に31%、4 ビット量子化後に**40%**の忘却を削減した。このアプローチは、追加の計算をほぼ増やすことなく、利益の大部分を捉える。

4. ヘッシアン分析

本論文は、ヘッシアンの分析によってその理論的直感を検証する：

二次近似： 損失の 2 次テイラー展開は、特に SAM でトレーニングされたモデルにおいて、忘却に対する妥当な上限として機能する。
方向性シャープネス： SAM とより大きなピーク学習率の両方が、微調整方向に沿った方向性シャープネス（ $\Delta_{FT}^\top H \Delta_{FT} / \|\Delta_{FT}\|^2$ ）を明示的に低下させる。忘却の減少は、微調整のステップサイズの低下ではなく、損失地形の平坦化によって駆動されている。

意義と主張

本論文は、ポストトレーニング更新に対するパラメータの感度が、事前学習ダイナミクスと下流パフォーマンスの間の重要かつ見過ごされているリンクであると主張する。著者らは、ベースモデルの品質のみを最適化する標準的な慣行では不十分であると論じる。

主な主張には以下が含まれる：

事前学習目標の再定義： 事前学習レシピは、ベースモデルの品質と、下流パラメータシフトに対する低感度の両方に基づいて選択されるべきである。
パイプラインへの最適化： 言語モデルの事前学習は、ベースチェックポイントの孤立ではなく、微調整、アライメント、圧縮を含む完全なモデル開発パイプラインに対して最適化されるべきである。
実用的影響： 実務家は、事前学習全体を通じてではなく、トレーニング後期（減衰時）にシャープネスアウェアな手法（SAM など）を選択的に展開することで、顕著な堅牢性の利益を捉えることができ、このアプローチはスケーラブルである。

著者らは、適応性と低感度を第一級の評価基準として扱うことで、ベースモデル自体が従来の事前学習損失指標ではわずかに弱く見える場合でも、事前学習モデルを将来の改変に対するより強力な出発点にできることを結論づけている。