Let the Code LLM Edit Itself When You Edit the Code

要約

この研究では、開発者が既存のコードをリアルタイムで編集し、コード・アシスタント、例えば大規模言語モデルに、次のトークンや次の行をその場で再予測するよう要求する、コード生成における典型的なシナリオを調査する。単純に考えると、LLMは正確な予測を行うためにKVキャッシュ全体を再エンコードする必要がある。しかし、このプロセスは、特にシーケンス長が長い場合、計算コストがかかる。編集された部分シーケンスをエンコードし、元のKVキャッシュに統合するだけでは、時間的な混乱が生じ、性能が著しく低下する。我々は、この効率と精度のトレードオフに対処するために、位置エンコーディ ング(PIE)を導入する。PIEは、ロータリ・ポジション・エンコーディングの上に構築され、まず、時間的混乱を引き起 こすKeyキャッシュ内のロータリ行列を削除し、次に正しいロータリ行列を再適用する。このプロセスにより、トークン間の位置関係が正しいことが保証され、行列の乗算は1ラウンドで済む。RepoBench-C-8kデータセットで、1.3B、6.7B、および33Bのパラメータを持つDeepSeek-Coderモデルを使用した広範な実験を通じて、PIEの有効性を検証する。評価には、コード挿入、コード削除、マルチプレースコード編集の3つの実コーディングタスクが含まれます。その結果、PIEはすべてのモデルサイズとタスクにおいて、標準的な完全再計算アプローチと比較して計算オーバーヘッドを85%以上削減し、同時にモデル性能をよく近似していることが実証されました。

要約(オリジナル)

In this work, we investigate a typical scenario in code generation where a developer edits existing code in real time and requests a code assistant, e.g., a large language model, to re-predict the next token or next line on the fly. Naively, the LLM needs to re-encode the entire KV cache to provide an accurate prediction. However, this process is computationally expensive, especially when the sequence length is long. Simply encoding the edited subsequence and integrating it to the original KV cache meets the temporal confusion problem, leading to significantly worse performance. We address this efficiency and accuracy trade-off by introducing \underline{\textbf{Positional \textbf{I}ntegrity \textbf{E}ncoding} (PIE). Building upon the rotary positional encoding, PIE first removes the rotary matrices in the Key cache that introduce temporal confusion and then reapplies the correct rotary matrices. This process ensures that positional relationships between tokens are correct and requires only a single round of matrix multiplication. We validate the effectiveness of PIE through extensive experiments on the RepoBench-C-8k dataset, utilizing DeepSeek-Coder models with 1.3B, 6.7B, and 33B parameters. Our evaluation includes three real-world coding tasks: code insertion, code deletion, and multi-place code editing. Results demonstrate that PIE reduces computational overhead by over 85% compared to the standard full recomputation approach across all model sizes and tasks while well approximating the model performance.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL