RidgeFill、再学習なしで長文コンテキストのPrefillを高速化

導入

長いプロンプトを扱う大規模言語モデルでは、回答生成そのものより、最初のトークンが返るまでのPrefillが大きな待ち時間になることがあります。RidgeFillは、モデルの重みを再学習せずに、このPrefillの計算量を減らす近似手法です。単純に層やプロンプトのトークンを捨てるのではなく、後段層が必要とする入力とキャッシュを予測します。

仕組み

モデルはまず、指定した分割層まで通常どおりプロンプトを処理します。その分割層の状態を基に、後続する各層について専用のマップを使い、本来その層へ渡される入力を推定します。推定された入力から、後段層のキャッシュを生成する仕組みです。

マップは校正用テキストを1回前向き計算し、その結果から閉形式のリッジ回帰で求めます。勾配学習やモデル重みの更新はありません。後段層自身の射影処理も変更せず、入力だけを予測値に置き換えます。また、選択した行、プロンプト末尾の数千トークン、そしてすべてのデコードステップは完全なモデル計算に戻します。つまり、全トークンを一律に近似するのではなく、誤差に敏感な部分を残す設計です。

この発想は、投影機構を学習段階で組み込むYOCOやDeepSeek-V4.1-Flashに近いものですが、RidgeFillはその能力を、対応する訓練を受けていない既存モデルへ後付けしようとします。

評価で分かったこと

論文は、全トークンの平均損失だけでは近似誤差を見落とすと指摘します。そこで、キャッシュを埋めた領域でのneedle retrieval、遠距離コンテキストに依存するトークンの損失、予測したkeyやvalueを正確な値に差し替えるsplice testを使い、どこに誤差があるかを調べました。

その結果、影響する要素はモデルごとに異なります。Llama-3.1-8Bではkey、特にそのスケールが重要でした。一方、Qwen3.5-9Bではvalueと再帰的な状態の影響がより大きいとされています。

llama.cppの派生版をA100で動かしたQwen3.8-27Bでは、16K〜120Kトークンの入力に対してTTFTが1.25〜1.31倍高速化されました。120Kトークンでは約28秒の短縮と報告されています。RULERは約1.2〜1.4ポイント、LongBench v2は約1.4ポイント低下しました。95%区間はゼロを含むものの、数ポイントの低下を排除できません。vLLM上のQwen3.8-Flash-Nextでは、9K〜32Kトークンで1.50〜1.60倍の高速化が報告されています。

意義と注意点

RidgeFillは、公開済みモデルに対して再学習なしで近似Prefillを試せる点に価値があります。マップを一度作成し、対応する推論実装を使えば、長文入力の待ち時間を調整可能な形で削減できます。

ただし、校正テキストと大きく異なるプロンプトでも同じ精度が得られるかは、実運用で検証が必要です。速度はモデル、ハードウェア、推論フレームワーク、完全計算に残すトークンの数にも左右されます。したがって、現時点では無損失の置き換えではなく、速度と品質のバランスを選ぶ手法と見るのが適切です。

出典：Hugging Face Daily Papers