ARIMAモデルをJavaScriptでスクラッチ実装した

はじめに

機械学習モデルをブラウザ上で動かすことができるai-on-browserを作っています。ここに時系列モデルのARIMAを実装しました。

• arima.js

statsmodelsなどの実装と比べると、かなり簡易的なものになっており色々足りていないので安定した動作は保証できませんが、基本的な理解はできるかなと思います。

ARIMAモデルとは

ARIMA(p, d, q) は次の3つの要素からなるモデルです。

• AR(p): 自己回帰。過去 p 個の値の線形結合で現在の値を説明する
• I(d): 和分。d 回差分を取ることで非定常な時系列を定常にする
• MA(q): 移動平均。過去 q 個の誤差項（残差）の線形結合で現在の値を説明する

差分を取った系列 $y_t$ に対して、ARMA(p, q) は次のように書けます。

$$
y_t = \sum_{i=1}^{p} \phi_i y_{t-i} + \sum_{j=1}^{q} \theta_j u_{t-j} + u_t
$$

ここで $u_t$ は残差です。この $\phi$（AR係数）と $\theta$（MA係数）を推定します。

コンストラクタ

コンストラクタは p, d, q を受け取り、$\phi$ は 0 で、$\theta$ は 0.3 で、それぞれ初期化します。

constructor(p, d, q) {
	this._p = p
	this._d = d
	this._q = q
	this._rate = 0.1
	this._beta = 1.0e-5
	this._phi = Array(this._p).fill(0)
	this._the = Array(this._q).fill(0.3)
}

また、パラメータ更新ステップ幅 rate を 0.5、正規化項 beta を 1.0e-5 として用意します。

差分と積分

d 階差分を計算します。

_diff(data) {
	let y = data
	const lastValues = [y[y.length - 1]]
	for (let t = 0; t < this._d; t++) {
		const ny = []
		for (let i = 0; i < y.length - 1; i++) {
			ny[i] = y[i + 1] - y[i]
		}
		y = ny
		lastValues.push(y[y.length - 1])
	}
	return [y, lastValues.slice(0, -1)]
}

y[i+1] - y[i] を d 回繰り返すだけです。差分を取るたびに系列の最後の値を lastValues に保存しておき、これを後で差分の逆演算（積分）に使います。

積分側は次のようになります。

_integrate(data, lastValues) {
	let c = data
	for (let t = lastValues.length - 1; t >= 0; t--) {
		const integrated = []
		let acc = lastValues[t]
		for (let i = 0; i < data.length; i++) {
			acc += data[i]
			integrated.push(acc)
		}
		c = integrated
	}
	return c
}

_diff で計算した lastValues を初期値として累積和を取り、差分を取った回数分だけこれを繰り返すことで元のスケールに戻します。

パラメータ推定（fit）

fit メソッドでは、まず this._diff(data) で d 階差分を取った系列 y を求めます。以降の処理は、この y に対する ARMA(p, q) モデルとして扱うことができます。

ARMAのパラメータ推定は非線形最小二乗問題になります。ここではガウス・ニュートン法に正則化を加えた形で、1ステップ分のパラメータ更新を行います。

fit(data) {
	const [y] = this._diff(data)
	const n = y.length
	const pq_max = Math.max(this._p, this._q)

	for (let k = 0; k < 1; k++) {
		this._u = [y[0]]
		for (let i = 1; i < n; i++) {
			let v = y[i]
			for (let j = 0; j < Math.min(i, this._p); j++) {
				v -= this._phi[j] * y[i - j - 1]
			}
			for (let j = 0; j < Math.min(i, this._q); j++) {
				v += this._the[j] * this._u[i - j - 1]
			}
			this._u[i] = v
		}

		let J = Matrix.zeros(n, this._p + this._q)
		for (let i = 0; i < n; i++) {
			for (let j = 0; j < Math.min(i, this._p); j++) {
				J.set(i, j, -y[i - j - 1])
			}
			for (let j = 0; j < Math.min(i, this._q); j++) {
				J.set(i, j + this._p, this._u[i - j - 1])
			}
		}
		J = J.slice(pq_max)

		const f = new Matrix(n - pq_max, 1, this._u.slice(pq_max))
		const H = J.tDot(J)
		H.add(Matrix.eye(H.rows, H.cols, this._beta))
		const d = H.solve(J.tDot(f)).value

		let e = d.reduce((s, v) => s + Math.abs(v), 0)
		e /= this._phi.reduce((s, v) => s + Math.abs(v), 0) + this._the.reduce((s, v) => s + Math.abs(v), 0)
		if (Number.isNaN(e) || e < 1.0e-12) break

		for (let i = 0; i < this._p; i++) {
			this._phi[i] -= this._rate * d[i]
		}
		for (let i = 0; i < this._q; i++) {
			this._the[i] -= this._rate * d[i + this._p]
		}
	}
}

処理の流れは次のようになっています。

1. 残差 _u の計算：現在の _phi, _the を使い、AR項とMA項からその時点での残差 $u_i$ を逐次的に計算する
2. ヤコビアン J の構築：残差 $u_i$ を $\phi_j$, $\theta_j$ で偏微分した値を並べた行列。AR部分の偏微分は $-y_{i-j-1}$、MA部分の偏微分は $u_{i-j-1}$
3. 正規方程式を解いて更新量 d を求める：$J^\top J$ に微小な正則化項（_beta × 単位行列）を加えてから $J^\top f$ に対して解く。$J^\top J$ が特異に近い場合でも解けるようにするための正規化項
4. 収束判定：更新量とパラメータの絶対値の比 e が十分小さい、またはNaNになったら打ち切る
5. パラメータの更新：_rate を掛けて _phi, _the を更新する

ここまでが更新一回分です。複数回 fit を呼び出すことで徐々に収束していきます。

予測（predict）

パラメータが求まれば、予測はAR項・MA項の定義通りに将来の値を1ステップずつ計算していくだけです。

predict(data, k) {
	const [y, lastValues] = this._diff(data)
	const preds = []
	const lasts = y.slice(y.length - Math.max(this._p, this._q))
	lasts.reverse()

	for (let t = 0; t < k; t++) {
		let pred = 0
		for (let i = 0; i < this._p; i++) {
			pred += this._phi[i] * lasts[i]
		}
		pred += this._u[this._u.length - 1]
		for (let i = 0; i < this._q; i++) {
			pred -= this._u[this._u.length - i - 2] * this._the[i]
		}
		preds.push(pred)
		lasts.unshift(pred)
		lasts.pop()
	}

	return this._integrate(preds, lastValues)
}

• 差分後の系列の末尾 Math.max(p, q) 個を初期状態として保持する
• k ステップ分、AR項＋直近の残差＋MA項から予測値を計算し、lasts を1つずつスライドさせながら繰り返す
• 予測が終わったら、fit 時に取っておいた lastValues を使って _integrate で元のスケールに戻す

不足点

• 定数項（drift）は、扱いがよくわかっていないため未実装

動かし方

サイト上で、時系列データに対してARIMAモデルを選択すると、パラメータ（p, d, q）を変えながら挙動を確認できます。

• https://ai-on-browser.github.io/?data=air&task=TP&preprocess=&model=arima

p = 12, d = 1, q = 0 で綺麗に予測できることが確認できると思います。

おわりに

ai-on-browserのARIMA実装について紹介しました。ai-on-browserには他にも色々なモデルが実装されているので、よければ触ってみてください。