概要
- 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して最終的な値を得る．
