HeteroFold、異なるモデル間のKV Cache共有を実現

異なるモデルを役割ごとに組み合わせるマルチエージェントLLMでは、計画、検索、実行、検証などを別々のモデルに担当させられます。一方で、エージェント間の通信をテキストで行うと、受信側は送信側がすでに処理した長いコンテキストをもう一度Prefillしなければなりません。モデルを組み合わせる柔軟性が、重複計算というコストによって損なわれる可能性があります。

KV Cacheの共有は、この問題への有力なアプローチです。モデルは入力を処理する過程でKeyとValueの中間状態を保存し、後続の生成で再利用できます。しかし、Cacheはモデルに依存した内部表現です。モデルファミリーが異なれば、トークナイザー、層数、注意機構、KVの次元、表現空間が異なります。そのため、送信側のテンソルを受信側へそのまま渡しても機能しません。

論文が提案するHeteroFoldは、こうした差異をまたぐPrefill-free転送を目指します。まず異なるトークナイザーによるトークン境界を対応付け、次に異なるアーキテクチャの層を整合させます。そのうえで送信側のK/V状態を受信モデルの表現空間へ写像し、受信側の注意パターンや出力を保つよう校正します。送信モデルと受信モデル自体は凍結されたままで、両者を接続する写像を学習する構成です。

評価では、Llama、Qwen、Ministralの間で6方向の転送を調べ、長文脈、短文脈、マルチエージェントの設定を比較しています。作者の報告によれば、4つの長文脈ベンチマークすべてで最良のCache転送性能を示し、短文脈の多くの設定でも比較手法を上回りました。32Kコンテキストでは、Llama 3.1 8BからMinistral 3 14Bへの転送が、受信側の通常Prefillより約10.7倍高速で、Dense LatentおよびKV Ridgeより1.18〜1.47倍高速だったとされています。マルチエージェント評価では、テキスト通信に匹敵する結果も報告されています。

この研究の意義は、KV Cacheを単一モデル内部だけの一時的な状態ではなく、モデル間で交換可能な計算資産として捉え直す点にあります。長い文書、タスク履歴、環境情報を繰り返し共有するエージェントでは、再計算を減らし、異なるモデルの組み合わせを使いやすくできる可能性があります。

ただし、実運用では変換と校正そのもののコスト、Cacheの転送量、未評価のタスク分布での安定性を確認する必要があります。提供された素材だけでは、すべてのモデル規模や本番環境で同じ効果が得られるとは判断できません。HeteroFoldは、モデルの境界を越えたCache共有が可能であることを示す一つの設計案として位置づけるのが適切です。

出典：Hugging Face Daily Papers