State Space Modelとは
SSM / 状態空間モデル
系列データを連続的な状態変数の遷移として捉え、系列長に対して線形の計算量で処理できるモデルアーキテクチャ
ひとことで言うと
文章を長くしても計算の負担があまり増えない、Transformerの代わりになりうる新しいモデルの作り方。
概要
State Space Model(SSM)は、制御工学における状態空間表現の考え方を深層学習に応用し、入力系列を隠れ状態の連続的な遷移として捉えて出力を計算する系列モデリングのアーキテクチャ。 Self-Attentionが系列長の2乗に比例する計算量を必要とするのに対し、SSMは系列長に対して線形の計算量で処理できるため、非常に長い系列を扱う際の計算効率で優位性を持つ。 2023年に発表されたMambaは、入力の内容に応じて状態遷移のパラメータを動的に変化させる選択的機構(Selective SSM)を導入し、それまでのSSMが苦手としていた文脈に応じた選択的な情報の保持・破棄を可能にした。 Transformerを置き換える、あるいは組み合わせて使う系列モデリング手法として、LLMを中心に研究が活発化している。
背景
Self-Attentionは系列内の全トークンペアの関連度を計算するため、系列が長くなるほど計算量とメモリ使用量が急増するという課題がある。 SSMは、隠れ状態を介して過去の情報を要約しながら逐次的に処理するRNN的な発想を、より効率的かつ並列学習しやすい数理的枠組みで再構成し、長い系列を線形の計算量で扱えるようにすることを目指している。
歴史
2021年: GuらがS4(Structured State Space Sequence Model)を提案し、SSMが長い系列のモデリングで有効であることを示す。 2023年: GuとDaoがMambaを発表し、入力依存の選択的機構によりSSMの表現力を高め、Transformerに匹敵する性能を一部のタスクで実証。 2024年以降: SSMとTransformerを組み合わせたハイブリッドアーキテクチャの研究が進展。
アーキテクチャ
各タイムステップの隠れ状態は、前の状態と現在の入力から状態遷移行列・入力行列を用いた線形な更新式によって計算される。 出力は隠れ状態から出力行列を通じて計算される。 Mambaなどの選択的SSMでは、これらの行列を固定パラメータではなく入力トークンの内容に応じて動的に変化させることで、重要な情報を選択的に記憶し不要な情報を忘れる能力を獲得している。
コード例
NumPyでの単純な状態空間モデルの再帰計算
import numpy as np
A, B, C = 0.9, 1.0, 1.0 # スカラー版の状態遷移・入力・出力行列
h = 0.0
inputs = [1.0, 0.5, -0.2, 0.3]
outputs = []
for u in inputs:
h = A * h + B * u # 状態更新
outputs.append(C * h) # 出力計算
print(outputs)利点
欠点
- Self-Attentionのような、系列内の任意の2点間を直接比較する能力は構造上制限される
- 比較的新しい技術であり、大規模モデルでの実運用実績やエコシステム(ツール・ライブラリ)がTransformerほど成熟していない
- タスクによってはTransformerベースのモデルと同等以上の性能を安定して得るための工夫がまだ研究途上にある
比較
- Transformer — TransformerのSelf-Attentionは系列長の2乗の計算量を要するのに対し、SSMは線形の計算量で長い系列を処理できる
- RNN — SSMはRNNと同様に隠れ状態を介して系列を逐次処理するが、並列学習しやすい数理的定式化を持つ点が異なる
- コンテキストウィンドウ — SSMの線形計算量という特性は、非常に長いコンテキストウィンドウを扱う用途と相性がよい
関連用語
よくある質問
MambaとTransformerのどちらが優れているのか？
タスクや系列長によって異なる。非常に長い系列の処理効率ではMambaに優位性があるとされる一方、Transformerは実績とエコシステムの成熟度で勝る。両者を組み合わせたハイブリッドモデルの研究も進んでいる。
選択的SSMとは何か？
状態遷移に使うパラメータを、固定値ではなく入力トークンの内容に応じて動的に変化させる仕組み。文脈に応じて重要な情報を保持し、不要な情報を忘れることができる。