技術概要：KLAS - 類似性を用いたニューラルネットワークの継ぎ接ぎによる精度と効率のトレードオフ改善

1. 問題定義

大規模な事前学習済みビジョン・トランスフォーマー（ViT）や畳み込みニューラルネットワーク（CNN）の展開には、根本的な課題が存在する。環境は強力なクラウドサーバーからリソース制約の厳しいエッジデバイスまで多岐にわたり、特定の計算予算内で柔軟なモデル選択を必要とする。ニューラルアーキテクチャ探索（NAS）は精度と効率のバランスを図ろうとするが、従来の 1 対 1 の NAS は計算コストが過大であり、1 対多またはゼロショット NAS 手法は単一のモデル設計空間に限定されることが多く、既存の事前学習済みモデルの豊富なプールを活用できていない。

モデル・ステッチングは、異なる事前学習済み「アンカー」モデルのブロックを組み合わせ、補間されたネットワークを作成することで有望な代替手段を提供する。しかし、SN-Net などの既存の多対多ステッチング手法は、ステッチ構成を選択するために（隣接するアンカーの「最も近いステッチ」や固定されたブロックペアリングなどの）単純なヒューリスティックに依存している。これらのヒューリスティックは、しばしば最適ではない精度と効率のトレードオフをもたらすだけでなく、多様なモデルファミリー間での汎用性を欠いている。さらに、ステッチ構成の探索空間は広大（深度 $n$ の $k$ 個のモデルの場合 $O(k^2n^2)$）であり、網羅的な探索は非現実的である。核心的な問題は、高いトレーニングコストを伴わずに性能を最大化する互換性のあるステッチ構成を特定するための、原理的かつ自動化された手法の欠如である。

2. 手法：KLAS フレームワーク

著者は、事前学習済みモデルの中間表現間の類似性を活用してステッチ選択を自動化するフレームワークであるKLAS（KL 発散に基づくアンカー・ステッチング）を提案する。

中核概念：類似性指標としての KL 発散

KLAS は、成功するステッチングには表現の類似性（活性化の幾何学的な整合性）と機能的な類似性（タスクに関連する情報および決定境界の維持）の両方が必要であると仮定する。著者らは、カルバック・ライブラー（KL）発散が両方を捉えるのに唯一適しているとの見解を示す。

• 表現の整合性: 活性化分布間の統計的距離を定量化し、軽量な変換がソース活性化をターゲット活性化にマッピングできるかどうかを示す。
• 機能的な整合性: ソースブロックの出力がターゲットモデルの決定境界をどの程度維持するかを評価し、ステッチされたネットワークが最小限のファインチューニングで下流タスクの性能を維持することを保証する。

ステッチングの成功を予測できないことが多い Centered Kernel Alignment (CKA)、平均二乗誤差 (MSE)、または交差エントロピー (CE) などの指標とは異なり、KL 発散は線形プローブによって生成される確率分布の互換性を直接測定する。

KLAS アルゴリズム

このフレームワークは主に 2 つの段階で動作する。

1. アンカー選択:

   • KLAS は、候補となる事前学習済みモデルの最後のブロックの softmax 確率分布間の KL 発散を計算する。
   • 最後のブロックの KL 発散が低いモデルは、類似した決定境界とクラス信頼度分布を共有しているため、互換性のあるアンカーと見なされる。これは先行研究で使用されていたヒューリスティックな「最も近いステッチ」に取って代わる。

2. ブロック選択:

   • アンカーが選択されると、KLAS は以下で定義されるステッチスコア $\Gamma(i, j)$ を用いてブロックのペアを評価する。
     $$\Gamma(i, j) = \frac{\Theta(P^f_i, P^g_j)}{\Theta(P^g_j, P^g_{j+1})}$$
     ここで:
     • $\Theta(P^f_i, P^g_j)$ はクロスアンカー活性化距離（ソースブロック $i$ とターゲットブロック $j$ 間の KL 発散）である。
     • $\Theta(P^g_j, P^g_{j+1})$ はアンカー内ブロック容量（ターゲットアンカー内の連続するブロック間の KL 発散）である。
   • このスコアは、仮想的なステッチ層にかかる変換の負担と、ターゲットブロックが分布の違いを吸収する能力とのバランスを取る。
   • ProbeNet: 完全なステッチ層をトレーニングすることなくこれらのスコアを効率的に計算するために、KLAS は ProbeNet という統一されたアーキテクチャを採用する。これはアンカーモデルのすべてのブロックに線形分類プローブを接続し、単一の順伝播・逆伝播パスで共同でトレーニングする。これにより、個々のプローブをトレーニングする場合と比較して、無視できる程度のワンタイムコスト（例：Swin-B の場合 0.25 GPU 日）で済む。

3. 構成選択:

   • 候補となるステッチ構成は、FLOPs バケット内の最小スコアに対する調整可能な閾値 $\tau$ に基づいてフィルタリングされる。これにより、精度と効率のスペクトル全体にわたって展開可能なモデルの密なパレートフロントが確保される。

3. 主要な貢献

1. ヒューリスティックの限界の特定: 本論文は、SN-Net などの単純なヒューリスティックが普遍的に最適ではなく、ヒューリスティックが見過ごす優れた「遠距離ステッチ」構成（例：Swin-Ti を直接 Swin-B に接続するなど）を見逃していることを実証する。
2. 原理的な類似性指標: ステッチ層のインスタンス化やトレーニングを必要とせずに、表現と機能の両方の整合性を捉えるステッチ選択の指標として KL 発散を導入する。
3. KLAS フレームワーク: 多様なモデルファミリー（ViT、CNN、LLM）にわたってアンカーとブロックの選択を自動化する、汎用性が高く低コストなフレームワーク。
4. 実証的検証: 複数のデータセット（ImageNet-1K、CIFAR-100、ADE20K、TruthfulQA）およびアーキテクチャにおいて、ベースラインに対して一貫した改善を示す包括的な実験。

4. 実験結果

著者は、KLAS を SN-Net および他のベースラインと比較し、同一のファインチューニング条件（ビジョンでは 50 エポック、LLM では 10 エポック）で評価した。

• 精度と効率のトレードオフ:

  • ImageNet-1K: KLAS は、SN-Net と比較して、同じ計算コストで最大1.21% 高い Top-1 精度を達成するか、同等の精度でFLOPs を 1.33 倍削減する。
  • AUC スコア: Swin ファミリーにおいて、KLAS は Area Under the Curve (AUC) で0.8950を達成し、SN-Net (0.8345) や CKA (0.8124) などの他の類似性指標を大幅に上回った。
  • クロスアーキテクチャ: KLAS は ResNet-18 と Swin-B を正常にステッチし、同程度の FLOPs でカスケード推論ベースラインを Top-1 精度で 2〜3% 上回った。

• 密な予測（セマンティックセグメンテーション）:

  • ADE20K における Mask2Former の使用において、KLAS によって選択されたステッチは、同程度の FLOPs バジェットで SN-Net ベースラインに対して最大**+0.9% の mIoU**向上を達成した。

• 大規模言語モデル（LLM）:

  • TruthfulQA ベンチマークにおいて Llama 1B および 3B アンカーを使用した場合、KLAS によるステッチモデルは、パラメータ数が一部の構成で少ないにもかかわらず、ヒューリスティックベースの LLM ステッチング手法である ESTA を最大0.017 の ROUGE-1および0.033 の ROUGE-2で上回った。

• 指標の有効性:

  • KL 発散は、実用的な Swin Ti-B ステッチ構成の**88.9%**を回復させたのに対し、CKA は 5.5%、SN-Net のヒューリスティックは 61.1% にとどまった。

5. 意義と主張

本論文は、KLAS が、アドホックなヒューリスティックに代わって原理的かつ類似性駆動のアプローチを採用することで、モデル・ステッチングにおける重要なギャップを埋めると主張している。その意義は以下の点にある。

• 汎用性: 単一のファミリーに限定された手法の限界を克服し、多様なアーキテクチャファミリー（トランスフォーマー、CNN、ハイブリッド）およびモダリティ（ビジョン、言語）で機能する。
• コスト効率: プローブトレーニングとステッチ層のファインチューニングのみという無視できる追加計算オーバーヘッドで、展開可能なモデルの連続的なパレートフロントを生成し、従来の NAS の過大なコストを回避する。
• 体系的な最適化: KL 発散によって分布の違いを明示的に捉えることで、ヒューリスティック手法が見逃す直感的ではない高性能なステッチ構成（例：「遠距離ステッチ」）の発見を可能にする。

著者らは、KLAS がゼロからトレーニングを必要とすることなく既存の事前学習済みモデルのエコシステムを効果的に活用し、展開シナリオにおける精度と効率のトレードオフを最適化するための堅牢で自動化されたソリューションを提供すると結論付けている。