DAST: Context-Aware Compression in LLMs via Dynamic Allocation of Soft Tokens
- 長い文脈をそのまま LLM に読ませると計算量が増えるため、本文を soft token に圧縮して扱う方法が用いられる。
- DAST は、各チャンクに同じ数の soft token を割り当てるのではなく、perplexity と attention から情報量を見積もり、重要なチャンクに多く割り当てる。
- LongBench と Long-Term Memory MSC で、Beacon などの既存手法を上回り、とくに強い圧縮をかけた場合の性能低下が小さかった。
Abstract(日本語訳)
大規模言語モデル(LLM)は、長い文脈入力を扱う際に計算上の非効率と冗長な処理に直面しており、そのため圧縮技術が注目されている。既存の意味ベクトルに基づく圧縮手法は有望な性能を示しているが、文脈チャンク間にある本来的な情報密度の違いを考慮せず、soft token を各チャンクに一様に割り当てている。この一様な配分は、情報上重要な領域への割り当てを避けがたく減らしてしまう。これに対して本論文は、LLM が内部にもつ文脈的関連性の理解を利用して圧縮を導く、単純で有効な手法である Dynamic Allocation of Soft Tokens(DAST)を提案する。DAST は、perplexity に基づく局所的情報と、attention によって捉えられる大域的情報を組み合わせ、情報量の多いチャンクへ soft token を動的に割り当てることで、有効な文脈依存の圧縮を可能にする。複数のベンチマークにおける実験結果は、DAST が state-of-the-art 手法を上回ることを示している。
論文の面白いところ
長文入力の圧縮では、どの部分を残すかが性能を左右する。従来の soft token 型の圧縮では、長い文書をチャンクに分けたあと、各チャンクにほぼ同じ圧縮容量を与えることが多かった。この仮定は扱いやすいが、実際の文書では答えに関わる段落と背景説明だけの段落が同じ密度をもつとは限らない。DAST の着眼点は、圧縮容量そのものを固定するのではなく、その配分を入力ごとに変える点にある。しかも、外部の重要度推定器を別に使うのではなく、LLM 自身の perplexity と attention を使って判断する。これは、検索や要約のように明示的に文を削る方法とは異なり、soft token による連続表現の圧縮と相性がよい。論文中の NaturalQuestions の分析では、答えを含む位置に DAST がより多くの soft token を割り当てる様子が示されている。長文 RAG や長い対話履歴の保持では、入力の全体を等しく扱うより、情報の濃淡に合わせて予算を配るほうが自然である。この論文は、その直観を比較的単純な機構で実装している。
問題設定
LLM は長い文脈を扱えるようになったが、入力が長くなるほど計算量とメモリ使用量が増える。文脈圧縮は、この問題を緩和するために、元の長い入力を短い表現へ置き換える技術である。初期の方法には、文を削る pruning や要約があるが、入力の文面を直接変えるため、必要な情報を落とす危険がある。近年は、元のテキストを少数の soft token に圧縮する意味ベクトル型の方法が使われる。soft token は通常の語彙トークンではなく、モデル内部で扱う連続的な表現であり、短い長さで文脈情報を保持するために使われる。問題は、文書を複数チャンクに分ける場合、各チャンクの情報密度が同じではないことである。既存の Multi-chunks Compression は、推論時に各チャンクへ固定数の soft token を割り当てるため、重要なチャンクが圧縮不足になりうる。DAST は、総 soft token 数を一定に保ったまま、チャンクごとの割り当て数を動的に変える問題としてこれを定式化する。
提案手法
DAST は Dynamic Allocation of Soft Tokens の略で、長い入力をチャンクに分けたうえで、それぞれに割り当てる soft token 数を入力に応じて決める。各チャンクの重要度は、局所的な情報と大域的な情報の二つから推定される。局所的な情報には perplexity を用いる。論文では、低い perplexity は現在の文脈情報の関連性が高いことを示すものとして扱われ、チャンク単位で計算される。大域的な情報には attention を用いる。圧縮後の表現に対する attention が大きいチャンクは、系列全体の理解に寄与する部分と見なされる。DAST はこの二つを係数 alpha で重み付けして組み合わせ、softmax で正規化したスコアからチャンクごとの soft token 数を求める。実験では alpha を 0.5 としており、感度分析では値を変えても性能は比較的安定していた。必要に応じて、訓練時の圧縮制約と整合するように割り当てを再調整する reallocation アルゴリズムも使う。
結果
評価は LongBench の Single-Document、Multi-Document、Few-Shot タスクと、Long-Term Memory MSC で行われた。Llama-2-7B では、文書・例圧縮の平均スコアが Beacon の 41.3 に対し、DAST は 46.4 だった。Qwen-2-7B でも、Beacon の平均 49.7 に対し、DAST は 51.6 だった。差が大きいのは Multi-Document で、Llama-2-7B では 27.5 から 37.4、Qwen-2-7B では 40.3 から 45.6 に伸びている。Long-Term Memory MSC では、約 4 倍圧縮で DAST が 55.9、約 24 倍圧縮でも 51.9 を保った。Beacon は同じ条件で 39.0 から 32.3 に下がっており、DAST は強い圧縮下での劣化が小さい。ablation では、ランダム割り当てが 34.52、一様割り当てが 34.90、動的割り当てが 38.14 であった。perplexity または attention の片方を外すと性能が下がるため、局所情報と大域情報の併用が効いていると読める。一方で latency は Beacon より大きく、Llama-2-7B で Beacon 1.2 に対し DAST 1.6、Qwen-2-7B で Beacon 3.3 に対し DAST 3.8 である。外部モデルを使う LongLLMLingua よりは低いが、単純な直接圧縮よりは追加の計算が必要になる。
具体例
たとえば、長い企業年次報告書と「今年の営業利益が減った主な理由は何か」という質問を LLM に与える場面を考える。報告書には沿革、役員紹介、地域別の販売概況、財務諸表、リスク要因などが続き、質問に直接関係する記述は数段落だけである。通常の一様なチャンク圧縮では、沿革の段落にも財務の段落にも同じ数の soft token が割り当てられる。DAST では、各チャンクの perplexity と、圧縮表現に対する attention から、営業利益の減少理由に関わる財務説明や事業別コメントのチャンクを相対的に重く見る。その結果、広告費の増加、原材料費の上昇、特定地域の売上減少といった答えに必要な箇所へ、より多くの soft token が配られる。モデルは圧縮後の短い表現を読み、営業利益減少の理由を要約して答える。間違えやすいのは、数値表の近くにある注記や前年との比較文が圧縮で薄くなる場合である。DAST はこの危険を完全に消すものではないが、重要そうなチャンクに圧縮容量を寄せることで、一様配分より答えの根拠を残しやすくする。