HackSick vol.7 LT資料【LLMアーキテクチャ入門・事前学習時の躓き所解説】 スパースなAttention・状態空間モデル

LLMアーキテクチャの概要や、状態空間モデル（SSMs）の要点、学習時の躓きどころなどを解説した資料です。

Transcript

1. 【LLMアーキテクチャ入門・事前学習時の躓き所解説】 スパースなAttention・状態空間モデル An Introduction to Large Language Model Architectures: Understanding Challenges in Pretraining with Sparse Attention and State Space Models
2. Vol.7
3. 自己紹介 / About us り っ か ぼ た ん 六花 牡丹 Rikka Botan 独立研究者（機械学習 / 代数学 / 数理論理学） Independent researcher (machine learning / algebra / mathematical logic) ◆趣味 お菓子作り・紅茶・クラシック鑑賞・お洋服 ◆最近の活動 Silver Award: Liquid AI Hackathon Series | Tokyo 記事執筆（Mamba, LFM2 (LTCs) 関連） SSE Modelシリーズの公開 X(Twitter) Portfolio
4. 目次 1 LLMの演算ってSelf Attention以外に何があるの？ 2 Self Attention以外の演算を使うメリットは何だろう？ 3 最近の研究の流れと有名な手法のそれぞれの違いを理解してみよう！ 4 実際に構築するときのポイント！細かなポイントを解説 5 Self Attentionと組み合わせて複合モデルを構築してみよう！
5. （前提知識復習）Self Attentionって何？ ▪Self Attentionのイメージ 現代の機械学習技術の基盤になっている演算。 T1 T2 T3 T4 T5 T6 T1 0.09 0.41 0.05 0.03 0.07 0.35 T2 0.11 0.04 0.14 0.35 0.11 0.25 T3 0.31 0.01 0.01 0.10 0.28 0.29 T4 0.20 0.33 0.02 0.01 0.43 0.01 T5 0.21 0.10 0.14 0.36 0.03 0.16 T6 0.51 0.13 0.06 0.12 0.05 0.13 入力からQuery, Key, Valueという ３つのベクトルを生成。 QueryとKeyの間で類似度のような計算 を実行する。 トークン間の距離に依存せず、 トークン同士のそれぞれの関係性（類似度） を並列で学習可能である。 （１度の行列演算ですべてのトークンの 計算が完結し、計算待ちが発生しない） これにより長い入力でもすべての関係性を 参照でき、高い性能を示す。 また、並列処理が可能なため 高い学習速度を有する。 𝐴𝑡𝑡𝑒𝑛𝑡𝑖𝑜𝑛(𝑄, 𝐾, 𝑉) = 𝑆𝑜𝑓𝑡𝑚𝑎𝑥 𝑄𝐾 𝑇 𝑑 𝑉
6. （前提知識復習）Self Attentionって何？ ▪Attention is all you need それまで畳み込みや再帰を使用したモデルが 最高記録として報告されていた中で、 Attentionのみを用いたモデルが翻訳タスクで最高性能を示し、 「Attentionさえあればいい！」 という言及が研究者に衝撃を与えた。 ➢ 以降、AttentionのみからなるTransformerが 言語モデルにおけるデファクトスタンダードになる。 [1] Ashish Vaswani et al., Attention Is All You Need, arXiv:1706.03762
7. 効率化のための手法 ▪ 系列演算を系統に整理してみよう 【１】Efficient Full Attention メモリの読み書きを減らしたり、 演算自体を削減 また低ランクに圧縮して 演算を抑えるなども行われる。 【２】Sparse Sequence Modeling トークン同士の関係性の内、 重要なもののみを計算 単純に閾値で制限するものや 軽い演算で選別してから 演算を行う手法などがある。 【３】Linear Sequence Modeling 再帰や畳み込みを利用して Self Attentionとは違う演算を実施 Self Attention の O(N²) の計算量 ・KVキャッシュの課題を解消 [2] Weigao Sun et al., Speed Always Wins: A Survey on Efficient Architectures for Large Language Models, arXiv:2508.09834
8. LLMの演算ってSelf Attention以外に何があるの？ ▪重要な４つの系統 【１】Linear Attention Softmax Attention の簡素化。 Softmax を外して QKᵀ の結合則 を使い、行列状態 の逐次更新に変形。 【２】線形RNN（再帰） 非線形な再帰をあえて線形化した手法。 非線形 RNN では不可能だった 時間方向の並列学習（並列スキャン） が可能になる。 【３】SSMs（状態空間モデル） 状態フィードバック制御を 離散化・時変化した手法。 構造化された状態行列 を用いて 入力を状態に圧縮する。 【４】TTT RNN（Test-Time Training） 状態の更新を「テスト時のオンライン学習」 とみなす手法。 単純な加算より表現力の高い更新則が特徴。 [2] Weigao Sun et al., Speed Always Wins: A Survey on Efficient Architectures for Large Language Models, arXiv:2508.09834
9. Self Attention以外の演算を使うメリットは何だろう？ ▪本当にSelf Attentionさえあればいい・・・？ TransformerベースのLLMは様々な評価で高い性能を示している。 しかし、近年課題があるとも言及されている。 ▪長い入力時の非効率性と性能面の弱点 【１】メモリの2次的な増大 入力の長さがNの時、Attentionスコア行列がN×Nになるため系列長に対して2乗でメモリが増える。 ⇒メモリの使用量が大きく、デバイスで扱いきれなくなる。 【２】推論速度の低下 KVキャッシュ（前の時刻までの計算結果）が線形に肥大化し、 1トークンあたりの生成時間が伸びていく（過去のすべてのトークンとの関係を計算する副作用）。 【３】Attention Sink 先頭の少数トークンに注意が集中する現象。先頭に引っ張られて性能が発揮されにくくなる。 【４】注意の希釈 Softmax関数は合計が一定。⇒入力が長い場合に重要な少数のトークンに注意を集中させにくくなる。 ➢ これらを解消するためにSelf Attention以外の演算が効果的。
10. 最近の研究の流れ ▪各手法を理解する上でのポイント Self Attention以外の手法は、 過去の状態を中間状態に圧縮し、 その圧縮した状態を逐次更新していく。(メモリの２次的な増大の回避) この更新の仕方が各手法の主な違い。 （右表のMemorizing with Gateの部分） 例えば・・・ Mamba2では、入力に重みを付けて足すことで状態を更新する。 この時に、これまでの入力が長いほど更新量が小さくなりやすい ようにすることで自然な記憶の減衰を可能にしている。 （過去の記憶を適度に忘却しながら更新する） Gated Delta Netでは、Delta則（一般化したHouseholder変換）を用いて、 特定の古い記憶を消してから新しい記憶を状態に加える。 結果として必要な情報を効率的に記憶することが可能になった （加算のみだと忘却が不十分で、記憶が飽和する） 直交的更新で状態が暴れず安定＝正確な想起も可能になっている。 ➢ つまり、過去の記憶をどのように忘却するのか、 新しい記憶をどのように圧縮して更新するのかがポイント！ [2] Weigao Sun et al., Speed Always Wins: A Survey on Efficient Architectures for Large Language Models, arXiv:2508.09834
11. 最近の研究の流れ ▪各手法を理解する上でのポイント より複雑な記憶の制御は高い性能を達成しやすい一方で、速度とのトレードオフがある。 いかに高速に複雑な記憶の制御を行うかが研究の主な課題となっている。 近年では、Attentionとの組み合わせで単体のみからなるモデルより高い精度を発揮することが示されており、 複合モデルの研究が盛んに行われている。（代表例：Qwen3, Nemotron-H, LFM2, Granite-H） [3] Songlin Yang et al., Gated Delta Networks: Improving Mamba2 with Delta Rule, arXiv:2412.06464
12. 実際に構築するときのポイント ▪ HiPPO 初期化 / A 行列のパラメータ化 / ステップ Δ SSMsは初期化が性能を大きく左右する。 ・A 行列：HiPPO由来の値で初期化し、対角＋負の実部で安定化（発散防止）。 ・離散化ステップ Δ：softplus 等で正に制約し、fp32 で計算して数値安定を確保。 ・状態のスケールや初期値の取り違えは、学習の停滞・発散に直結する（すべての記憶が直ちに忘却されるなど）。 これらの初期化を一つでも誤ると、 初期値が発散（exp）し、収束しなくなる。 Lossが一定値から一切下がらなくなる。 もし学習が進まない場合は、 一度初期化の妥当性をチェックした方が良い。 [4] https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba2.py
13. 実際に構築するときのポイント ▪Inference Cacheのクラス外部化とCUDA Graphの最適化 SSMsは過去の状態から次の時刻の状態を逐次効率的に計算するために、Inference Cacheというメモリを用意する。 この時、Inference CacheをSSMs Class内部に定義してしまうと、 torch.compile使用による最適化が一部無効になり、速度低下につながる。 SSMsのClass内部にInference Cacheを定義するのではなく、 Class外部にInference Cacheを定義する。 Class内部に定義するとCUDA Graphで処理を 最適化できなくなり、速度が低下する。 （NVIDIA GPUを使用する際に特に注意） [5] https://github.com/tommyip/mamba2-minimal/blob/main/mamba2.py
14. 複合モデルを構築してみよう！ ▪実はTransformer層の非線形性は弱い？ 既存のモデルにおける層ごとの性質を知らべると、ほぼ線形であることがわかっている。 ⇒Attentionは少数で良い可能性が高い。 [6] Anton Razzhigaev et al., Your Transformer is Secretly Linear, arXiv:2405.12250
15. 複合モデルを構築してみよう！ ▪AttentionとSSMsの組み合わせ方 AttentionとSSMsの複合モデルは下記のように構築すると精度・速度が高くなると言及されている。 (1) Attention層は少数(約7〜30%)でよい。 (2) Attention層は先頭には置かず深さ方向に均等分散。 (3) 深い層ほどAttentionの比率を高める。 ⇒精度を維持または向上させつつ、長い入力で数倍程度の高速化が可能。 [7] NVIDIA, NVIDIA Nemotron 3: Efficient and Open Intelligence, arXiv:2512.20856
16. 結論 言語モデルにおいてアーキテクチャによる性能への影響は限定的だと度々述べられる。 しかしそれはアーキテクチャの改善が無意味という意味ではなく、 性能を維持したまま高速化が可能ということを意味する。 モデル使用時に重要視されるのは、性能と速度であり、 アーキテクチャの改善はそうした意味で重要な要素となる。