概要
- ICLR23
- 状態空間モデル(state-space model; SSM)は様々なモダリティにおいて有用性が検証されてきたが,未だ言語系においては確認できていない.
- また,SSMは $\mathcal{O}(L)$であるにも拘らず, $\mathcal{O}(L^2)$であるTransformerよりも遅い
- 実験によって,SSMが①前方にあるトークンの記憶と②トークン間の比較が苦手なことを発見し,この二つの難点を乗り越える新たなSSMとしてH3 (Hungry Hungry Hippos)を提案する.
- Transformerに替わるモデルとも言われている(要出典)
SSMにおける二つの問題点
- 問題点①②を検証するために,二つのタスクInduction HeadとAssociative Recallを実施
- Induction Head : 特殊なトークン
|-
で囲まれた部分文字列の先頭の文字を出力させるタスク- 前方のトークンを如何に覚えているかを測ることができる
- Associative Recall : key-valueでセットになってるアルファベットと数字の組に対して,与えられたkeyに対応するvalueを出力させるタスク
- この場合
a 2 c 4 b 3 d 1
に対して入力がa
なので2
が答え (間違ってたら教えてくれ) - トークン間の関係を覚えているかどうかを測ることができる
- この場合
- Induction Head : 特殊なトークン
- 結果は以下の通り
- Attentionは100%成功しているが,従来手法はほとんどできていない
- 提案手法であるH3はほぼ100%成功
- Attentionは $QK^\top$によりトークン間の関係を記憶可能であり(②), $\mathrm{softmax}(QK^\top)V$によりトークン自体を直接記憶可能(①)
SSMの改善
-
SSM
- $x_i, u_i,y_i$をそれぞれ状態信号,入力信号, 出力信号とすると,
$$x_{i} = Ax_{i-1} + Bu_{i}$$
$$y_{i} = Cx_i+Du_i$$
- $x_i, u_i,y_i$をそれぞれ状態信号,入力信号, 出力信号とすると,
-
①前方トークンの記憶
- shift演算(e.g., $(a,b,c) → (0,a,b)$)を使うことで記憶
- 例えば,常に $A$がshift演算として機能するなら, $B=e_1$の時,連鎖的に $m$ステップ前までの $u_i$が $x_i$に格納される.→ $x_i = \lbrack u_i ,… , u_{i-m+1} \rbrack$
-
②トークン間の比較
- Attentionと同様, $QK^\top V$のように乗算することで記憶
- $K^\top V$部分はHiPPOの対角行列versionによって初期化された対角行列によるSSMが通される
- 対角行列の初期化はこちらを参照.
- HiPPOについては以下を参照
-
最終的には以下のように設計
- 計算量の観点からEfficient Transformer系列に倣って,以下のように設計
$$Q \odot \mathrm{SSM_{diag}}(\mathrm{SSM_{shift}(K) \odot V})$$ - すなわち, $K^\top V$を先に計算しておく
- 計算量の観点からEfficient Transformer系列に倣って,以下のように設計
The shift SSM can detect when a particular event occurs, and the diagonal SSM can remember a token afterwards for the rest of the sequence
- H3の流れ
- 入力 $u$に対して $Q = uW_Q, K = uW_K, V = uW_v$を得る.
- $K$を $\mathrm{SSM_{shift}}$に通して $\bar{K}$を得る.
- $Q,K,V$をmulti-head化 (すなわちdim方向で分割)
- 各headごとに $KV := \mathrm{SSM_{diag}}(\bar{K}V^\top)$を計算.
- ${Q_i \in \mathbb{R}^d | i = 1,…,N}$ごとに $Q_i(KV)_i$を計算してconcat→ $Q \odot \mathrm{SSM_{diag}}(\mathrm{SSM_{shift}(K) \odot V})$を得る.
- headをconcatして最終的な値を得る.