技術的概要：損失の整合性時：効率的な事前学習のための勾配ベースの複合損失重み付け

問題定式化

現代の深層学習システムは、大規模な（しばしばラベルなしの）データによる事前学習に続き、下流タスクへの適応を行うという、2 段階のパラダイムを頻繁に採用しています。多くの事前学習パイプラインにおいて、目的関数は単一の損失ではなく、複数の損失項（例えば、次のトークン予測、対照損失、または回帰目的）の重み付き組み合わせです。これらの項の相対的な重みは、学習された表現の質と下流タスクのパフォーマンスに大きく影響する重要なハイパーパラメータです。

現在、これらの重みの選択は、手動調整、ランダムサーチ、またはベイズ最適化に依存しています。これらの方法は、多数の独立した学習実行を必要とするため、大規模な事前学習には計算コストが過大です。さらに、均一重み付けのような単純なデフォルト値は、異なる損失項間の多様な貢献度と競合を考慮できないことがよくあります。本論文は、繰り返しフルトレーニング実行のコストを伴わずに下流タスクのパフォーマンスを最大化するため、複合事前学習目的における線形損失重みを調整する効率的な手法の必要性に対処します。

手法：勾配整合事前学習（GraP）

著者は、損失重みをオンラインで学習する勾配ベースのバイレベル最適化手法である**勾配整合事前学習（GraP）**を提案します。その中核的なアイデアは、複合事前学習勾配を下流目的の勾配と整合させることです。

1. バイレベル最適化の定式化

この問題は、外側ループで損失重み $w$ を最適化して下流損失 $L_{down}$ を最小化し、内側ループでモデルパラメータ $\theta$ を最適化して複合事前学習損失 $L_{pre}(w)$ を最小化する、バイレベル最適化タスクとして定式化されます。
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{subject to} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. 勾配ベースの更新則

切り捨てられた逆伝播アプローチ（単一ステップ SGD）を用いて、著者は重みのための更新則を導出します。重みに関する下流損失の勾配は、下流勾配と個々の事前学習勾配との内積に比例します。
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
ここで、$g_{down}$ は下流勾配、$g_k$ は $k$ 番目の事前学習損失の勾配です。目標は、複合勾配 $\sum w_k g_k$ と下流勾配との整合性を最大化することです。

3. 埋め込み空間近似

素朴な実装では、$g_k$ を取得するためにモデル全体を $K$ 回フルバックワードパス計算する必要があります。計算コストを削減するため、GraP は全パラメータ空間ではなく、共有バックボーン埋め込み $z$ に関する勾配を計算します。
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
損失は同じ埋め込みの上に配置された個別のヘッドを通じて適用されるため、これらの勾配を計算することは、フルバックワードパスよりもはるかに安価です。最適化目的は、この共有表現空間における整合性を最大化することになります。
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. 正規化戦略

本論文は、重み自体（例えば、$\sum w_k = 1$）やその $L_2$ ノルムを正規化することに反対しています。これらは複合勾配の最適方向を歪める可能性があるためです。代わりに、GraP は複合勾配自体の$L_2$ ノルム（$\|\sum w_k g_k\| = 1$）に制約を課します。これにより、最適化は更新ベクトルの大きさではなく方向に焦点を当て、局所目的の下で最適な整合方向を回復します。

5. アルゴリズムの効率性

結果として得られるアルゴリズムは、反復ごとにバックボーンに対して1 回のフォワードパスと 1 回のバックワードパスのみを必要とします。ヘッドに対して計算された勾配は、損失重みの更新とバックボーンパラメータの更新の両方に再利用されます。これにより、ハイパーパラメータ調整のオーバーヘッドは、単一の学習実行に対して約 30% 増加する程度に抑えられ、ベイズ最適化の桁違いのコストと比較されます。

主要な貢献

1. 問題定式化: 著者は、下流パフォーマンスを明示的にターゲットとしたバイレベル最適化問題として、事前学習損失重み選択を定式化しました。
2. 正規化の洞察: 彼らは正規化戦略を分析し、重み自体を制約するのではなく、更新ベクトル（複合勾配）のノルムを制約することが、最適な整合方向を保持する点で優れていることを実証しました。
3. 効率的なアルゴリズム: 彼らは、ハイパーパラメータの数に対するフォワード・バックワードステップの数に関して $O(1)$ の複雑さを持つアルゴリズム、GraP を提案しました。これは、複数のバックワードパスを回避するために共有表現空間における勾配整合を利用します。
4. 経験的効率性: この手法は、ベイズサーチなどの標準的なハイパーパラメータ最適化と同等かそれ以上のパフォーマンスを達成しながら、ほぼ50 倍効率的であり、従来の手法が計算的に実行不可能な設定での調整を可能にします。
5. 冗長性の特定: このアプローチは、冗長な損失項を特定し、重みを下げる能力を実証しています（例えば、ビジョンタスクにおいて特定の対照損失にほぼゼロの重みを割り当てること）。これは効果的に自動損失選択メカニズムとして機能します。

実験結果

この手法は、イベントシーケンスモデリングと自己教師ありコンピュータビジョンの 2 つのドメインで評価されました。

イベントシーケンスモデリング

• データセット: Churn, AgePred, AlfaBattle, MIMIC-III, および Taobao。
• 設定: 事前学習は、複数のフィールドにわたる次のトークン予測（NTP）と対照損失の複合目的で行われました。
• パフォーマンス: GraP は、GradNorm、DWA、MGDA、PCGrad などの慎重に調整されたベースラインおよび Optuna で調整された重みと同等かそれ以上のパフォーマンスを達成しました。RNN、Transformer、HT-Transformer の各バックボーンにおいて、Churn および MIMIC-III データセットで一貫して最良の結果を達成しました。
• 効率性: 数十の独立した実行を必要とするベイズ最適化（Optuna）と比較して、GraP は単一の学習実行と中程度のオーバーヘッド（スループット平均 44% の減速）で同様のパフォーマンスを達成しました。

自己教師ありコンピュータビジョン

• 設定: CIFAR-10、CIFAR-100、および ImageNet-100 において、All4One フレームワーク（Barlow Twins、BYOL、NNCLR 目的を組み合わせる）を使用して評価されました。
• パフォーマンス: 目的がしばしばよくバランスされているこの設定において、GraP はデフォルトの All4One パラメータおよび他のマルチタスク学習手法と同等のパフォーマンスを達成しました。
• 損失選択: 注目すべきは、GraP が Attn-NNCLR 損失にほぼゼロの重みを割り当てたことです。この損失を手動で除去しても同様のパフォーマンスが得られたことは、この手法が低影響または冗長な目的を特定する能力を有していることを検証しました。

除去実験

• ラベル漏洩: 重み調整に使用されるラベル付きデータの割合を変化させる実験は頑健性を示し、ラベルはバックボーンに対する直接的な教師信号ではなく、弱いガイダンス信号を提供することを示唆しました。
• 正規化: 提案された複合勾配正規化は、標準的な重み合計正規化よりも優れていました。
• 補助学習: 有限のラベル付きデータでの過学習の減少が原因と考えられ、GraP による事前学習 followed by 教師あり微調整は、結合補助学習よりも優れていました。

意義と主張

本論文は、勾配ベースの損失重み最適化が、古典的なマルチタスク学習および検索ベースのハイパーパラメータ最適化に対する実用的かつ効率的な代替手段であると主張しています。その主な意義は、従来の手法が計算的に過大となる複数の損失項を伴う大規模な事前学習シナリオにおいて、ハイパーパラメータ調整を実行可能にすることにあります。

著者は、GraP が複数の学習実行を必要としないため、計算予算が限られた設定に適していることを強調しています。この手法は各損失に個別のヘッドを仮定していますが、事前学習の更新方向を下流目標と整合させることで、競合する目的間のトレードオフを効果的に処理します。この研究は、複雑な事前学習パイプラインにおいて、損失重みの自動的かつオンラインな整合が、高価な手動調整やブラックボックス検索によって得られるものと同様、あるいはそれ以上の転移性能を持つ表現をもたらす可能性を示唆しています。