ICLR2022で発表された、新しい時系列モデリング手法としてS4(Structured State Space Sequence model)というものがある。S4は長距離ベンチマークで従来手法を圧倒的性能で破って話題となった。

S4の論文はいくつかの研究の集大成となっており、核となる技術としてHiPPO (Recurrent Memory with Optimal Polynomial Projections)というものがある。そこで論文と著者実装を通じて、HiPPOとはどんなものなのかを理解することを目指す。

実装編はこちら。 dosuex.com

概要

まずはHiPPOの概要を記載する。

時系列データからの学習における中心的な問題は、累積する履歴に応じて、圧縮された「表現」を逐次的に更新することである。

本研究では、連続信号や離散時系列を多項式基底への射影により逐次的に圧縮するためのフレームワーク（HiPPO）が開発された。

HiPPOは、過去の各時間ステップの重要度を与える測度(重み)が与えられると、逐次関数近似問題に対する最適解を求めることができる。さらにHiPPOを拡張して、すべての履歴を記憶するために、時間スケールの変化に対して不変性を持つ新しい記憶更新機構（HiPPO-LegS）を作ることができる。 HiPPO-LegSは、タイムスケールに対するロバスト性、高速なアップデート、有界な勾配といった利点を持っており、本研究の集大成に当たる。

ベンチマークである、Permuted MNISTにおいて、HiPPO-LegSは98.3%の最新精度を達成した。また、時間スケールや欠損データに対するロバスト性をテストする軌道分類タスクにおいて、HiPPO-LegSはRNNとニューラルODEのベースラインを25～40%の精度で上回った。

背景

時系列データからの学習は現代の機械学習の基本的な課題であり、言語モデル、音声認識、動画処理、強化学習のようなタスクに通じるものである。

長期間にわたる複雑な時間依存関係をモデル化する中心的な課題は「記憶」をどのような表現として保持するかである。記憶を内部状態として扱う、深層学習時系列モデリングの手法は次のようなものがある。

RNN：長時間にわたる時系列情報では勾配消失問題が発生することが知られている。
LSTM：サンプリングレートなどの変化によって起きる、時間スケールの変化に弱いことが知られている。

本研究で記憶の扱いに関して新しく得られた洞察は、記憶の保持・更新を逐次関数近似の問題として定式化したことである。

これにより、時系列関数 $f(t)$ はいくつかの基底関数に対する最適な係数を保存することによって低次元情報で近似される。この近似は、過去の各時刻の重要度を表す測度とともに評価される。

HiPPO(high-order polynomial projection operators)とは、任意の関数を与えられた測度のもとで直交関数系の係数空間に射影するものであり、これが時系列情報を圧縮して記憶として保持する操作に相当する。さらに時系列情報が更新されるに応じて、係数も時間発展させることができ、これが記憶を更新する操作に相当する。

記憶のモデリング

ここから、本研究での記憶の扱い方を数式を使って少しづつ具体化していく。

問題設定

時系列関数 $f(t) \in \mathbb{R}, t \ge 0$ において、過去の履歴 $f(x)_{x\le t}$ を元にして未来を予測したい。しかし過去の履歴を完全に記憶することは難しいため、何かしら履歴を圧縮した記憶として保持する必要がある。

関数空間での距離

ふたつの関数 $f$ と $g$ の距離を計算するために、以下のように関数の内積を定義する。

$\langle f, g \rangle_{\mu} = \int_{0}^{\infty} f(x) g(x) d\mu(x)$

ここで $\mu(x)$ は測度と呼ばれる量である。測度は点 $x$ における「重み」とイメージすると良い。これにより関数の $L_2$ ノルムも

$\| f \|_{L_2 (\mu)} = \langle f, f \rangle_{\mu}^{1/2}$

定義することができる。

多項式展開

時系列関数を何かしらの基底関数の多項式を用いて、関数

$g^{(t)} = \sum^{N-1}_{n=0} c_n(t) g_n$

で近似することを考える。

これにより複雑な時系列関数の振る舞いを、有限個の係数列 ${c_n}$ の情報へと圧縮することができる。なお本論文では基底関数として直交多項式を用いている。直交多項式は、直交性（次数の異なる基底関数同士の内積がゼロとなる）性質を持ち関数近似の領域でよく使われる手法である。

関数の多項式展開の身近な例として、フーリエ多項式などが挙げられる。

逐次近似

各時刻 $t$ において $f(x)_{x\le t}$ を逐次的に近似関数 $g^{(t)}$ を求めたい。これは $( - \infty, t )$ において定義される測度 $\mu^{(t)}$ を用いて

$g^{(t)} := \arg \min \| f_{x\le t} - g^{(t)} \|_{L_2 (\mu^{(t)})}$

となるような近似関数 $g^{(t)}$ を求めることに相当する。

HiPPOフレームワーク

以上で立てた方針を、具体的な手順にまとめていく。

逐次関数近似の方針

時系列関数 $f(t)$ を近似する関数 $g^{(t)}$ は、任意の基底におけるその展開のN個の係数で表現することができる。

まず、適切な基底を選択する。近似理論の古典的な手法を活用すると、自然な基底は測度 $\mu^{(t)}$ に対する直交多項式の集合であり、これは直交基底を形成する。このとき直交性により各基底の係数は

$c_n^{(t)}:=\left\langle f_{\leq t}, g_n\right\rangle_{\mu^{(t)}}$

で得ることができる。

次に、係数 $c_n^{(t)}$ に対する時間微分を計算する。 $c_n^{(t)}$ は $f(t)$ のダイナミクスによって決められる常微分方程式(ODE)に従って時間発展する。

逐次関数近似アルゴリズム

まずHiPPOを以下のように定義する。

$( - \infty, t )$ で定義される時間変化する測度 $\mu^{(t)}$ 、多項式のN次元部分空間 $\mathcal{G}$ 、および連続関数 $f:\mathbb{R}_{\ge 0} \rightarrow \mathbb{R}$ が与えられたとき、HiPPOは以下の性質を持つ射影演算子 $proj_t$ と係数抽出演算子 $coef_t$ を時間 $t$ ごとに定義する。

(1) $proj_t$ は時間 $t$ までに制限された関数 $f$

$f(x)_{x\le t} := f(x)|_{x\le t}$

を取り、近似誤差

$\| f_{x\le t} - g^{(t)} \|_{L_2 (\mu^{(t)})}$

を最小化する多項式 $g(t)\in \mathcal{G}$ に写像する。(図1 (1)~(2))

(2) 係数抽出演算子 $coef_t : \mathcal{G} \rightarrow$ は多項式 $g(t)$ を、測度 $\mu^{(t)}$ に対して定義された多項式直交基底の係数ベクトル $c^{(t)} \in \mathbb{R}^N$ に写像する。(図1 (2)~(3))

演算子 $proj$ と $coef$ を連鎖させてできる演算子を $hippo$ と呼ぶことにする。すなわち $hippo(f)(t) = c^{(t)} = coef_t(proj_t(f))$ である。

筆者らの導出により係数関数 $c^{(t)} = coef_t(proj_t(f))$ はODE

$\frac{d}{d t} c(t)=A(t) c(t)+B(t) f(t), \qquad A(t)\in \mathbb{R}^{N \times N}, B(t)\in \mathbb{R}^{N \times 1} \tag{1}$

に従って時間発展する(図1 (3))。

このODEを離散化して解くことで $c(t)$ を逐次的に求めることができる(図1 (4))。なお、この微分方程式を線形時間不変(linear-time-invariant; LTI)方程式と呼ぶ。

様々な測度族とHiPPO

本研究の主たる理論的結果は、様々な測度族 $\mu(t)$ (図2)に対してHiPPOを導出したことである。

translated Legendre (LegT) 測度

LegT測度

$\mu^{(t)}(x)=\frac{1}{\theta} \mathbb{I}_{[t-\theta, t]}(x)$

は最近の履歴 $[t-\theta, t$ ] に対して一様な重みを付与する(図2：左)。要約される履歴の長さを表すハイパーパラメータ $\theta$ が存在する。

LegT測度のもとでLTI方程式(式(1))の $A(t)\in \mathbb{R}^{N \times N}, B(t)\in \mathbb{R}^{N \times 1}$ は次のように与えられる。(導出は参考文献1の付録Dを参照)

$A_{n k}=\frac{1}{\theta}\left\{\begin{array}{ll} (-1)^{n-k}(2 n+1) & \text { if } n \geq k \\ 2 n+1 & \text { if } n \leq k \end{array}, \quad B_n=\frac{1}{\theta}(2 n+1)(-1)^n\right.$

translated Laguerre (LagT) 測度

LagT測度

$\mu^{(t)}(x) = e^{-(t-x)} \mathbb{I}_{(-\infty, t]}(x)= \begin{cases}e^{x-t} & \text { if } x \leq t \\ 0 & \text { if } x>t\end{cases}$

は指数関数的に減衰する重みを使用し、最近の履歴をより重要視する(図2：中央)。

LagT測度のもとでLTI方程式(式(1))の $A(t)\in \mathbb{R}^{N \times N}, B(t)\in \mathbb{R}^{N \times 1}$ は次のように与えられる。(導出は参考文献1の付録Dを参照)

$A_{n k}=\left\{\begin{array}{ll} 1 \quad \text{ if } n \geq k \\ 0 \quad \text{ if } n \lt k \end{array}, \quad B_n=1\right.$

HiPPO-LegS: Scaled Measures for Timescale Robustness

これまで見てきたように、測度関数に何を選択するかが記憶機構の特性を決定していた。したがって、より適切な測度を選択することで、より優れた理論的特性を持つ記憶機構を作ることができそうだ。

信号処理ではスライド窓関数が一般的であるが、記憶に対するより直感的なアプローチは、忘却を避けるために時間と共に窓関数をスケーリングすることである(図1：右図)。これを scaled Legendre measure (LegS) 測度と呼ぶ。

scaled Legendre measure (LegS)

LegS測度

$\mu^{(t)}=\frac{1}{t} \mathbb{I}_{[0, t]}$

は全履歴に対して一様な重みを付与する。 LegSのもとでのLTI方程式

$\frac{d}{d t} c(t)=-\frac{1}{t} A c(t)+\frac{1}{t} B f(t)$

の係数行列は次のように与えられ

$A_{n k}=\left\{\begin{array}{ll} (2 n+1)^{1 / 2}(2 k+1)^{1 / 2} \quad \text { if } n>k \\ n+1 \quad \text { if } n=k, \\ 0 \quad \text { if } n \lt k \end{array} \quad B_n=(2 n+1)^{\frac{1}{2}}\right.$

さらにLTI方程式は以下のように離散化される。

$c_{k+1} = \left( 1 - \frac{A}{k} \right)c_k + \frac{1}{k}B f_k .$

微分方程式を離散化すると一般に時間ステップ $\delta t$ が式に含まれることになるが、上の式では $\delta t$ が含まれていないことに注目して欲しい。これはすなわちLegSが時間スケールに対して不変であることを反映している。この特徴によって、LegSの時間スケールに対するロバスト性が保証される。

詳細は割愛するが、その他にもLegSは

計算の効率性が高い
勾配消失が起こらない
誤差が入力の滑らかさに応じて減少する

などの特徴を持つ。

図3にHiPPOによる逐次関数近似の様子を示す。現時刻に近い領域は微細な構造まで近似されているが、遠い過去の領域は大まかな傾向だけつかんだトレンドライン近似となっていることがわかる。

性能評価

評価実験によって、HiPPOの系列モデリングの評価およびHiPPO-LegSの時間スケールロバスト性が確認されている。

Sequential Image Classification on Permuted MNIST

Permuted MNISTとはMNISTのピクセルの並び順を別の順序で置換し、この置換を全てのデータに対して施したデータセットである(図4)。予測モデルに対してはピクセルを並べたベクトルが入力される。置換操作により、クラス数(1~9の9個)は保たれたまま局所構造が破壊されるので、ラベルの予測のためにはピクセル間の長期的な相関関係を捉える必要がある。

表1がPermuted MNISTの予測結果である。HiPPOがLSTMやTransformerといった、現在の主流となっている手法を超えていることがわかる。

HiPPO-LegSの時間スケールロバスト性

Character Trajectories dataset では固定のサンプリングレートで記録したのペン先軌道の測定値から、文字を分類することを目的としている(図5)。

時間スケールのロバスト性を確認するために、学習データとテストデータのサンプリングレートを変えて性能を評価する。表2の100Hz→200Hzがテストデータでサンプリングレートを引き上げた場合であり、200Hz→100Hzが引き下げた場合である。いずれも、従来の手法がサンプリングレートの変化にほとんど対応できていないのに対して、HiPPO-LegSが強いロバスト性を持っていることがわかる。

また欠損値(Missing values)を含む場合に対するテスト結果が表2の下部に記載されている。こちらも従来手法たちを大幅に超える結果を出していることがわかる。

まとめ

HiPPO(high-order polynomial projection operators)とは、任意の関数を与えられた測度のもとで直交関数系の係数空間に射影するものであり、この係数を時間発展させることで、逐次的な関数近似を実行する。

どすえのブログ

ソフトウェア開発ブログ

時系列モデリング手法 HiPPO を読み解く(1)

概要

背景