Beyond Perplexity: A Geometric and Spectral Study of Low-Rank Pre-Training

技術的サマリー：パープレキシティを超えて：低ランク前学習の幾何学的およびスペクトル的研究

問題提起
大規模言語モデル（LLM）の前学習は、フルランクの重み、勾配、オプティマイザ状態の保存に伴うメモリコストによって制約されています。これらのコストを削減する解決策として低ランク前学習が登場し、その手法はオプティマイザベース（例：GaLore、Fira）、アーキテクチャベース（例：CoLA、SLTrain）、アダプターベース（例：ReLoRA）の 3 つのカテゴリーへと進化しました。しかし、これらの手法の評価は依然として限定的です。現在の慣行は、単一のシード実行からの検証パープレキシティにほぼ完全に依存しています。著者らは、パープレキシティは解の質の代理指標として不適切であると主張します。パープレキシティが同一のモデルでも、損失ランドスケープの根本的に異なる領域に収束し、固有の内部表現を持つ可能性があるからです。さらに、既存の比較では、ある手法で前学習されたモデルをファインチューニングする能力と、その手法で前学習されたモデルが下流タスクでどの程度汎化するかを区別できていないことが多いです。

手法
著者らは、3 つのモデル規模（60M、130M、350M）において、5 つの低ランク前学習手法（GaLore、Fira、CoLA、SLTrain、ReLoRA）がフルランク学習に対して発見する解を特徴づける診断フレームワークを導入しました。本研究は、4 つの次元にわたる 16 の指標を用いてこれらの手法を評価します。

1. 1 次元損失ランドスケープ（ランダム方向）： 著者らはモデルパラメータを 100 のランダムな方向に沿って摂動させ、検証損失の平均と分散を測定します。これにより、チェックポイントからの損失増加率を表す「シャープネス（Sharpness）」と、損失面の異方性を表す「方向分散（Direction Variance）」というスカラー指標を導出します。
2. 1 次元損失ランドスケープ（Top-K PCA 方向）： 重み行列の主要な特異値方向に沿って摂動を適用し、最も重要な軸に沿った幾何学構造を分析します。
3. 1 次元補間： 著者らはチェックポイント間の直線経路に沿った損失を測定します。ここでは 2 つの特定の指標を導入します。
   • 連続チェックポイントバリア高さ（CCBH）： 同一手法の連続するチェックポイント間のバリアを測定します。
   • 手法間バリア高さ（IMBH）： 同一学習ステップにおいて異なる学習手法間の幾何学的分離を定量化する新規指標です。これは、異なる手法が同じ盆地に収束するのか、幾何学的に異なる解に収束するのかを明らかにします。
4. スペクトルおよびランク指標： 本研究は、重みと更新のスペクトル構造を追跡し、「実効ランク（Effective Rank）」、「安定ランク（Stable Rank）」、「スペクトルギャップ（Spectral Gap）」、「閾値ランク（Threshold Rank）」を含みます。
5. 活性化の類似性： 著者らは、活性化 L2 距離、コサイン類似度、線形 CKA（Centered Kernel Alignment）を用いて、低ランクモデルの隠れ状態をフルランクのベースラインと比較します。

最後に、著者らはこれらの幾何学的およびスペクトル特徴が、パープレキシティ単独よりも下流のゼロショット性能をよりよく予測できるかどうかを、LOSO（Leave-One-Size-Out）およびLOMO（Leave-One-Method-Out）交差検証を用いた線形回帰でテストします。

主要な貢献

1. 診断フレームワーク： 本論文は、単なるシードやエポック間だけでなく、学習アルゴリズム間の幾何学的分離を定量化する新規指標「IMBH」を含む包括的な評価スイートを提案します。
2. クロスファミリーの経験的知見： 本研究は、低ランク手法がフルランク学習にも互いにも等価ではないことを確立します。主要な知見には以下が含まれます。
   • 幾何学的分離： 手法は幾何学的に異なる盆地に収束し、IMBH は学習が進むにつれて単調に増加します。
   • スペクトル特性： GaLore はフルランク更新のスペクトル構造と密接に一致しますが、CoLA はそうではありません。
   • 表現の発散： 低ランクの活性化は、後続の層においてフルランクから次第に発散します。CoLA は直交する活性化を示し、Fira は最終層でのみ整合性を回復します。
3. 改善された下流予測： 検証パープレキシティ単独では下流性能の不完全な予測指標です。幾何学的およびスペクトル特徴でパープレキシティを補完することで予測精度が大幅に向上します（LOSO ピアソン相関が 0.873 から 0.913 に増加）。
4. 統合されたコードベース： 著者らは、前学習、指標計算、下流評価を統合したオープンソースのパイプラインを公開し、さらなる研究を促進します。

結果

• 損失ランドスケープの幾何学： フルランク学習はランダム方向ではよりシャープな盆地に落ち着きますが、Top-1 PCA 方向ではより平坦な盆地に落ち着きます。対照的に、Fira や SLTrain などの低ランク手法はランダム方向ではより平坦な盆地に収束し、CoLA や ReLoRA は最もシャープな盆地に到達します。GaLore は混合プロファイルを示し、ランダム方向では中程度に平坦ですが、Top-1 PCA 方向ではシャープです。
• 手法間バリア： すべての手法は幾何学的に分離された盆地を占有します。60M 規模では、フルランクと任意の低ランク手法間のバリアは、低ランク手法同士のバリアよりも高くなります。Fira はフルランクに幾何学的に最も近い盆地に収束し、SLTrain は最も孤立したままです。
• スペクトル分析： フルランクはすべてのランク指標で優位です。GaLore は実効ランクと安定ランクにおいてフルランクを密接に追跡しますが、スペクトルギャップは小さくなります。Fira と GaLore はフルランクの重み更新のスペクトル分布を最も忠実に再現します。
• 活性化の整合性： ReLoRA はすべての規模でフルランク表現を最も密接に追跡します。CoLA は最大の乖離を示し、後続の層では活性化がベースラインに対して直交します。Fira は層平均のドリフトは大きいですが、驚くべきことに最終層では乖離が低いです。
• 下流性能： 検証パープレキシティは下流のランキングを一貫して予測しません。例えば、Fira はすべての規模で最低のパープレキシティを達成しますが、130M と 350M ではパープレキシティの上位 2 位に入らないにもかかわらず、複数の下流カテゴリーで CoLA が主導しています。

重要性
本論文は、低ランク前学習手法がフルランク学習の交換可能な近似ではなく、到達する解を根本的に変化させるものであると結論付けています。本研究は、検証パープレキシティがモデルの質の不完全な図示であることを示しています。幾何学的およびスペクトル指標を取り入れることで、研究者は下流性能をよりよく予測し、異なる低ランク戦略によって導入される固有の帰納的バイアスを理解できます。この研究は、Fira のような一部の手法はフルランクに幾何学的に近い盆地に収束する一方で、CoLA のような他の手法は異なる谷を占有することを強調しており、低ランク手法の選択がモデルが占める損失ランドスケープの特定の領域を決定づけることを示唆しています。