JavaScriptを有効にしてください

【論文メモ】Hungry Hungry Hippos: Towards Language Modeling with State Space Models

 ·  ☕ 3 min read


概要

  • 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が答え (間違ってたら教えてくれ)
      • トークン間の関係を覚えているかどうかを測ることができる


  • 結果は以下の通り
    • 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$$
  • ①前方トークンの記憶

    • 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$
  • ②トークン間の比較

  • 最終的には以下のように設計

    • 計算量の観点からEfficient Transformer系列に倣って,以下のように設計
      $$Q \odot \mathrm{SSM_{diag}}(\mathrm{SSM_{shift}(K) \odot V})$$
    • すなわち, $K^\top V$を先に計算しておく

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


参考

共有

YuWd (Yuiga Wada)
著者
YuWd (Yuiga Wada)
機械学習・競プロ・iOS・Web