LoSA: ブロック単位の拡散言語モデルに対する局所性対応スパース注意

arXiv cs.CL / 2026/4/15

💬 オピニオンIdeas & Deep AnalysisModels & Research

共有:

要点

ブロック単位の拡散言語モデルは自己回帰的コストを削減できる一方で、スパース注意の非効率さにより長い文脈においては注意（attention）がメモリバウンドになってしまう。
本論文は、DLMに対する素朴なスパース注意においてKVインフレーション問題を特定する。ここでは、クエリ固有のプレフィックス選択が、KVキャッシュページの過剰なユニオン（結合）を読み込む原因となる。
LOSA（Locality-aware Sparse Attention）は、連続する復元（denoising）ステップの間で大部分のトークンはほとんど変化しないという観察を活用し、安定したトークンに対してはキャッシュ済みのプレフィックス注意を再利用し、活性のあるトークンに対してのみスパース注意を用いる。
複数のブロック単位DLMに対する実験により、LOSAは高い精度をほぼ密（dense）な場合に近い水準で維持しつつ効率を改善することが示される。具体的には、攻撃的なスパース性において平均精度が最大+9ポイント向上し、RTX A6000上で最大4.14×の注意速度向上が得られる。
報告されている利益は、復元ステップ間における局所性（locality）と時間的安定性を、品質を大きく損なうことなく、KVの読み込みと注意計算を削減するために活用できることを示唆している。

Abstract

ブロック単位の拡散言語モデル（DLMs）は、任意の順序で複数のトークンを生成し、自回帰的デコーディングのパイプラインに代わる有望な手段となります。しかし、長いコンテキストのシナリオでは、メモリ帯域に制約された注意（attention）によって依然としてボトルネックになっています。素朴な疎注意は、KVインフレーション問題のためDLMsではうまく機能しません。これは、異なるクエリが異なるプレフィックス位置を選択し、その結果、アクセスされるKVページの集合が大きくなるためです。これに対処するために、連続する復元（denoising）ステップの間では、アクティブなトークンのうち有意な隠れ状態の変化を示すのはごく一部であり、大多数の安定したトークンはほぼ一定のままであることを観察します。この洞察に基づき、LOS A（Locality-aware Sparse Attention）を提案します。これは、安定したトークンに対してキャッシュされたプレフィックス注意の結果を再利用し、疎注意をアクティブなトークンにのみ適用します。これにより、ロードする必要のあるKVインデックス数が大幅に削減され、その結果、高いスピードアップと高い精度の両方が得られます。複数のブロック単位DLMおよび各種ベンチマークにおいて、LOSAは密な場合に近い精度を維持しつつ、効率を大きく改善します。さらに、強い疎化レベルで平均精度を最大+9ポイント向上させながら、注意密度を1.54倍低く保ちます。また、RTX A6000 GPU上で注意の速度を最大4.14倍向上させており、提案手法の有効性を示しています。