-- Views
June 29, 26
スライド概要
M2の千さんが、論文「LOG-LINEAR ATTENTION」の内容について発表を行いました。本論文では、Softmax Attentionの計算量(O(T²))と、線形アテンション/SSM(Mamba-2など)が抱える「固定サイズの隠れ状態」という表現力の限界を両立して解消するためのLog-Linear Attentionという手法を提案しています。本手法は、固定サイズの隠れ状態を、Fenwick木に基づき指数的にサイズが増大するバケット群(隠れ状態の集合)に置き換えるもので、直近のトークンは高解像度に、遠いトークンほど粗く要約して保持します。これにより、計算量O(T log T)・メモリ量O(log T)というSoftmax Attentionと線形アテンションの中間的なコストで、連想記憶(Associative Recall)など固定状態サイズのモデルが苦手としていたタスクへの対応力を高めています。また、既存のMamba-2やGated DeltaNetなど任意の線形アテンション系アーキテクチャに対し、アテンションマスクを階層化するだけで適用できる汎用的なフレームワークである点も特徴です。実験では、MQAR(連想記憶ベンチマーク)やNeedle-in-a-Haystackにおいて線形アテンションのベースラインを上回る性能を示し、長文脈処理に有効であることが確認されました。
立教大学大学院人工知能科学研究科における瀧雅人准教授が主催する研究室で2020年度からスタートしているまだ若い組織です。 最先端の深層学習について、高度化・説明性向上などをテーマに深く幅広く研究しています。 また医療や神経科学・物理学におけるデータ分析や、産業への社会実装にも携わっています。 研究室内のPaper Reading活動の記録として、研究室学生の発表資料を公開しています。 ご興味をお持ちの方は、HPをご確認ください。
ジャーナルクラブ論文紹介 Log-Linear Attention LLMにおける記憶と忘却のバランスについて 千 チゲン 2026/06/17
ICLR 2026 著者: 1
Background 2
Softmax Attentionは、⾧いシーケンス(文脈)へのスケールアップが難しい 計算量O(L^2)、メモリO(L) スケールアップ… https://people.csail.mit.edu/yoonkim/data/efficient_architectures_talk.pdf 3
Linear Attentionの強み:計算量O(L) 、メモリO(1) https://www.csail.mit.edu/news/transforming-transformers 4
固定サイズstateのイメージ 5
提案手法:Log-Linear Attention Linear Attention の限界は、全履歴を 1個の固定サイズ state に圧縮すること: • 訓練時の計算量: O(T) • 推論時の計算量・メモリ量: O(1) 「提案手法」 Log-Linear Attentionでは、stateサイズを 増やす代わりに、計算量とメモリ使用量のトレ ードオフを行う: • 訓練時の計算量: O(T log T) • 推論時の計算量・メモリ量: O(log T) 6
Log-Linear Attention のポイント: マスクMの構造化 7
Mの階層化は、効率的な計算のポイント • M を無階層にすると softmax と同じコストに戻る • Mの階層化 → 効率的な計算 8
階層化の1例: State Space Models 展開 行列変換 • SSM の recurrence を展開すると、入力列から出力列への変換は 𝑌 = 𝑀𝑋という行列積として書ける。 • この 𝑀の各成分は、入力 𝑥 が出力 𝑦 に与える影響 𝐶 𝐴 ⋯ 𝐴 𝐵 である。 • 下三角部分の任意の submatrix の rank が高々 𝑁である行列を 𝑁-semiseparable と定義する。 • 低ランク𝑀 ≈𝑈𝑉 • したがって、𝑌=𝑀𝑥を計算せず、 𝑌=𝑈(𝑉𝑥)で済む 9
Fenwick Treeによるトークン分割 10
Softmax v.s Linear v.s Log-Linear Attention Softmax Attention: 全トークンの状態(KVキャッシュ)を個別に保持 (O(t) メモリ) 11
Softmax v.s Linear v.s Log-Linear Attention Softmax Attention: 全トークンの状態(KVキャッシュ)を個別に保持 (O(t) メモリ) Linear Attention: 全トークンを1つの固定サイズ箱(state)に圧縮 (O(1) メモリ) 12
Softmax v.s Linear v.s Log-Linear Attention Softmax Attention: 全トークンの状態(KVキャッシュ)を個別に保持 (O(t) メモリ) Linear Attention: 全トークンを1つの固定サイズ箱(state)に圧縮 (O(1) メモリ) Log-Linear Attention 「提案手法」 : 直近トークンはそのまま保持。過去のトークンはFenwick木分割により粗い複数の箱 (state)にまとめる (O(log t) メモリ) 13
Fenwick Tree 分割 近い過去 各時刻、箱の数: 箱のサイズ: バケットサイズが大きい場合: • 多くのトークンが1つの固定サイズ状態 に圧縮されている • 低解像度 遠い過去 バケットサイズが小さい場合: • 少数のトークンが1つの固定サイズ状態 に圧縮されている • 高解像度 14
Fenwick Tree 分割 定式化: t = 7の時: 高解像度 低解像度 15
Fenwick Tree 分割 定式化: 全ての が等しい(=1)とき、Log-Linear Attention は Linear Attention になる。 したがって、 を異なる値に設定することが、 さまざまな時間の幅で起きるパターンを捉え るために、必要である。 16
推論:Fenwick Tree Recurrence 17
Fenwick Tree Recurrence 定式化 prefixの境界点𝑏𝒕 : バケット分割 状態𝑺 バケット𝐵 ℓ : 時刻 𝑡から見た過去トークンの集合 ℓ 18
prefixの境界点𝑏𝒕 : バケット分割 レベル0) 𝑏 𝑏 𝑏 𝑏 =7−2 =6−2 =4−2 =7 今のトークン自身、サイズ1、計算なし =7−2 =6 レベル1) コードの計算: 𝑡 = 7 は二進数で 0111 。最下位ビット(1)は 2 = 1。 次のインデックス: 𝑡 = 7 − 1 = 6 に移動。 =6−2 =4 レベル2) コードの 計算: 𝑡 = 6 は二進数で 0110 。最下位ビット (1) は 2 =2。 次のインデックス: 𝑡 = 6 − 2 = 4 に移動。 =4−2 =0 レベル3) コードの計算: 𝑡 = 4 は二進数で 0100 。最下位ビット (1) は 2 = 4。 次のインデックス: 𝑡 = 4 − 4 = 0 となり、ループが終了。 19
バケット𝐵 ℓ : 時刻 𝑡から見た過去トークンの集合 𝐵 𝐵 = 7 レベル0) 直近1個のトークンを細かく保存 𝐵 = 6 レベル1) 直近1個のトークンを細かく保存 𝐵 = 45 レベル2) 少し過去のトークンを 2 個まとめる = 0123 レベル3) さらに遠い過去のトークンを 4 個まとめる 20
Decoding時、状態𝑺 ℓ の更新ルール lssb(t): 「最下位の 1 ビットの位置」 • ℓ = 0の時: Level 0 には、今のトークン 𝑣 𝑘 を入れる。 • 0< ℓ < 𝑙𝑠𝑠𝑏(𝑡)の時: 低いレベルの bucket を空にする • lssb(t)+1の時: 低い level の箱(バケット)をまとめて一つ上の level に送る。 • それ以外の高い level は前の状態をそのまま保持する。 21
Hierarchical MatricesとしてのAttention Mask 22
Efficient Attentionの定義(Linear Attention, Mamba-2等々) Attentionに似ている行列 下三角行列 23
Pの展開 = 特徴: • 主対角成分は、同じ時刻 𝑡 の 𝑞 と 𝑘 の内積を表す • Fenwick tree による prefix 分割が、行列 𝑀 の係数パターンとして現れている • 行が深くなるにつれて、時刻 t の増加を示す 24
Pの展開 = 特徴: • 主対角成分は、同じ時刻 𝑡 の 𝑞 と 𝑘 の内積を表す • Fenwick tree による prefix 分割が、行列 𝑀 の係数パターンとして現れている • 行が深くなるにつれて、時刻 t の増加を示す 25
Pの展開 = 特徴: • 主対角成分は、同じ時刻 𝑡 の 𝑞 と 𝑘 の内積を表す • Fenwick tree による prefix 分割が、行列 𝑀 の係数パターンとして現れている • 行が深くなるにつれて、時刻 t の増加を示す 26
HODLR (Hierarchically Off-Diagonal Low-Rank) 行列: 低ランクで表す構造 = = 明示的に展開 4 × 4 = 16個の内積項 低ランク因子表示 (メモリ効率のよい表現) 𝟒個の 𝜆𝑞と 𝟒個の 𝑘 27
HODLRの階層構造を抽象的に見ると Level 0 Level 1 Level 2 Level 1 Level 3 Level 1 Level 2 Level 1 特徴: • Level 0 は圧縮せずに計算する。 • Level が高くなるほど、より遠い過去を大きなブロック単位でまとめる。 • Levelが高くなるにつれて、圧縮率が高くなる。 28
訓練:Chunkwise 計算 29
Recurrent v.s Chunkwise Parallel https://people.csail.mit.edu/yoonkim/data/efficient_architectures_talk.pdf 30
Recurrent v.s Chunkwise Parallel https://people.csail.mit.edu/yoonkim/data/efficient_architectures_talk.pdf 31
Recurrent v.s Chunkwise Parallel https://people.csail.mit.edu/yoonkim/data/efficient_architectures_talk.pdf 32
Chunkwise 計算による訓練の高速化 𝐷は block-diagonal matrix であり、各 chunk の内部だけを扱う行列。 𝑀(ℓ) は、 chunk をまたぐ依存関係を表す。ただし、すべての token 間を密に計算するのではなく、階 層レベル ℓごとに、特定の 箱(バケット) ℬ ℓ に属する位置だけを扱う。 33
Chunkwise 計算による訓練の高速化(アルゴリズム 元の系列⾧を 𝑇、chunk size を 𝐶とすると、chunk の数は 𝑇/𝐶 Chunk 内の計算(Level 0): 各 chunk 内の attention を直接計算 Chunk 間の計算(Level 1以上) : Fenwick tree の各 level について、過去 chunk の情報を状 態 𝑆として伝播させ、現在 chunk の出力 𝑌に加える 34
Log-Linear Mamba-2 & Log-Linear Gated DeltaNet 35
Mamba-2 と Gated DeltaNetへ適応 36
計算量・メモリの比較図 37
実験結果 38
シーケンス⾧ごとの学習スループット(左、高いほど良い)と順伝播・逆伝 播カーネル実行時間(右、低いほど良い) 実験条件: H100 GPU、batch size 2、48 head、head dimension 64、state dimension 128、chunk size 64 39
Log-Linear Attention導入によるMamba-2とGated DeltaNetの改善(1) 40
Log-Linear Attention導入によるMamba-2とGated DeltaNetの改善(2) トークン位置が進むにつれて損失が一貫して減少している場合、モデルがコンテキスト全体を 効果的に活用できていること(long-range context utilization)を示している。 一方、損失がある時点で頭打ちになる場合は、シーケンス内の遠く離れた過去の情報を十分に 利用できていないことを意味する。 41
しかし、Transformerとの差が依然として大きい NIAH は、⾧い文脈中に埋め込まれた特定情報を正確に取り出せるかを測るための⾧文検索タスクである。 42
結論・感想 43
結論・感想 • 「固定サイズという制約を緩和」 Linear Attention / State Space Model 系において、固定サイズの State の制約を緩和し、 状態数を系列⾧に対して対数的に拡張する点に取り組んだ論文である。 • 「Transformerとの差が大きい」 Log-Linear Attention は従来の Linear Attention の文脈保持能力を改善し、Mamba-2 や Gated DeltaNet に適用した場合にも性能向上が見られる。一方で、特に「⾧文における 情報検索」タスクでは、依然として Transformer との差が残っている。 • 「近傍バイアス」 また、Fenwick Tree に基づくトークン分割では、直近のトークンを高解像度で保持し、 遠い過去のトークンをより粗い粒度で圧縮して扱う。そのため、直近情報を優先するバ イアス、すなわち recency bias に近い性質が生じる可能性がある。(これは、SSM のボ トルネックを recency や over-smoothing の観点から分析した Wang et al. (ICLR 2025) の議論とも関連する。) • 今後は、このような階層的な圧縮構造によって、どのようなタスクで性能が低下 するのかを調べる価値がある。 Wang, Peihao, et al. "Understanding and mitigating bottlenecks of state space models through the lens of recency and over-smoothing." International Conference on Learning Representations. Vol. 2025. 2025. 44
本資料(スライド)の限界 本資料の限界として、低ランク構造が効率的な計算につながる理由を、数式レベルでは十 分に示せなかった点が挙げられる。特に、semiseparable matrix の off-diagonal block が低 ランクに分解できることを、Mamba-2 の SSD アルゴリズムと関連づけて説明できれば、 計算効率化の意義をより明確に示せたと考えられる。 また、Linear Attention から Linear Attention with Gates、Linear Attention with the Delta Rule、Long Convolution Model へ至る一連の流れを十分に説明できなかった。そのため、 recurrent form を parallel form に変換する意味や、SSM を matrix transformation として捉 える意義が伝わりにくかった可能性がある。 45