BOSCH：LLMにおける短いコンテキストの注意ヘッド選択のためのブラックボックス・バイナリ最適化

Abstract

大規模言語モデル（LLM）のポストトレーニングにおけるハイブリッド化では、二次的な自己注意をスライディングウィンドウ注意（SWA）に置き換えることが多く、これによりKVキャッシュの使用量を削減し、レイテンシを改善します。既存のハイブリッド化方式は、典型的には層レベル（例：インタリービング）で定義されるか、あるいは局所から大域へと至る静的ランキングによってヘッドレベルで定義されます。しかし、層レベルの方式は、同一層内のヘッドを介して局所依存と大域依存がルーティングされるという事実を無視しており、一方で静的なヘッドレベルのランキングは絡み合い（entanglement）に悩まされます。つまり、ハイブリッド化後にあるヘッドの局所／大域の振る舞いが変化し得るのです。そこで本研究では、BOSCH（Black-box Binary Optimization for Short-context Head Selection）を提案します。これは訓練不要の手法であり、この問題を大規模近傍探索（Large Neighborhood Search）として定式化し、3つの下位問題に分解します：（i）小さな予算でのブラックボックス・プローブによる層の重要度の検出、（ii）これらの感度に基づく層ごとの適応的なSWA比率の割り当て、（iii）比率バケット内でのヘッドレベル最適化のグループ化。パラメータ数1.7Bから30Bまでの4つのLLMに対し、4つのSWA比率にわたって大規模に実験を行った結果、BOSCHは層レベルのヒューリスティックおよび強力な静的ヘッドレベル手法6種のいずれに対しても一貫して優れた性能を示し、特にSWA比率が高い場合ほど改善が大きいことが分かりました。継続的な事前学習（continual pretraining）のもとでは、BOSCHは元の長コンテキスト性能をより速く、かつより高い水準まで回復します。選択されたヘッドの分析からは、異なるSWA比率間でBOSCHが選ぶヘッドが大きく入れ替わる（ターンオーバーが大きい）ことが明らかになり、固定した局所性ランキングに頼るのではなく、対象の比率ごとにヘッドレベルの選択を行うことの重要性が示されています。