【疑問】Mamba って Transformer の代わりになるん？状態空間モデルと「選択」のしくみ【AI】

Mamba っていうモデルが、Transformer より長い文章を速く扱えるって聞いた
中身が状態空間モデルらしいんやけど、それって何

＞＞1
いちばん簡単な形は、隠れた状態 $h$ を1歩ずつ更新して、そこから出力を作る式

前の状態に $\bar{A}$ をかけて、今の入力を $\bar{B}$ 倍して足す。RNN とよく似た形

＞＞2
数字で回してみて

＞＞3
全部1次元で $\bar{A}=0.9$、$\bar{B}=1$、$C=1$ とする。入力が $1,0,0,0$ なら
$h$ は $1\to 0.9\to 0.81\to 0.729$ と、少しずつ薄まりながら残る
入力が $0,2,0,0$ なら $0,2,1.8,1.62$

＞＞4
それやと RNN と何が違うん

＞＞5
RNN は更新の中に非線形の関数が入るけど、この形は線形で、$\bar{A},\bar{B},C$ が時刻によらず同じ
そうすると出力は、入力に決まった重み $C\bar{B},C\bar{A}\bar{B},C{\bar{A}}^2\bar{B},\dots$ をかけて足したもの（畳み込み）になる
＞＞4 なら重みは $1,0.9,0.81,\dots$。学習のときは全部の時刻をまとめて計算できる

＞＞6
それが S4 とかいうやつ？

＞＞7
そう。この形に、長い依存を覚えやすい $A$ の作り方を組み合わせたのが S4（Gu et al., 2022）で、長い系列のベンチマークで強かった
でも文章のような、とびとびの記号の並びでは Transformer に届かなかった

＞＞8
なんで文章は苦手やったん

＞＞9
Mamba の論文（Gu & Dao, 2023）は、パラメータが入力によらず同じなのが弱点やと言ってる
「この単語は覚える、このつなぎの言葉は無視する」という、中身を見た取捨選択ができない

＞＞10
Mamba はそこをどう変えたん

＞＞11
$B$、$C$ と、時間の刻み $\Delta$ を、その時刻の入力から計算するようにした。これを選択（selection）と呼んでる
$\Delta$ が大きいと状態をリセットして今の入力を重く見て、小さいと今の入力を無視して状態を保つ。RNN のゲートに近い働き

＞＞12
その代わり、パラメータが時刻ごとに変わるから、＞＞6 の畳み込みの計算は使えなくなる
そこを、GPU のメモリの使い方を工夫した並列のスキャンで速く計算した、というのが論文のもう1つの柱

Transformer より速いって、どれくらい？

＞＞14
Attention は系列の長さ $n$ に対して、全部の組み合わせ $n^2$ を計算する。$n=1000$ なら100万組
状態空間モデルは長さに比例する $n$ で済む。生成のときも、過去の全部を持っておく（KVキャッシュ）必要がなく、1歩あたりの計算は一定
論文では、同じくらいの大きさの Transformer と比べて、生成のスループットが約5倍と報告してる

＞＞15
じゃあもう Transformer いらんやん

＞＞16
そうとも言い切れない。論文の言語モデルの比較は30億パラメータ程度までで、もっと大きい規模でも同じか確かめる必要がある、と論文自身が書いてる
状態が固定の大きさやから、長い文脈の中から特定の1語を正確に取り出すような作業は、全部を見られる Attention のほうが向いてる場面もある

＞＞17
それなら両方まぜたらええんちゃう

＞＞18
実際、Mamba の層と Attention の層を組み合わせたモデルも出てきてる
どっちが勝つかというより、長さと精度の釣り合いで使い分ける方向

統計の状態空間モデル（カルマンフィルタ）と同じものなん？

＞＞20
「見えない状態が1歩ずつ変わり、そこから観測が出る」という発想は同じ
ただ、統計の状態空間モデルは確率のモデルとして状態を推定するのが目的で、Mamba はニューラルネットの層として入力を出力に変換する部品。名前は同じでも使い方が違う