NO TEARS アルゴリズムとは
NO TEARS はデータから有向非巡回グラフ（Directed Acyclic Graph; DAG）を推定するためのアルゴリズムです。因果推論ライブラリ CausalNex の中で使われています。
$d$ 個の変数間の関係を知るために DAG を推定しようとすると、$d$ に対して計算量が急増化することは容易に想像できます。実際、DAG 学習問題は素直に取り組むと NP 困難となります。これを解消するために、NO TEARS アルゴリズムでは「非巡回」という条件を滑らかな関数で表現し、DAG の学習を連続最適化問題に落とし込みます。つまり、重み $W \in M_d\left(\mathbb{R}\right)$ のグラフ $G\left(W\right)$ があり得る DAG の集合 $\mathbb{D}$ に含まれているかという条件のもとでの最適化問題
$$
\min_{W\in M_d\left(\mathbb{R}\right)} F\left(W\right) \\
\mathrm{subject\ to\ } G\left(W\right) \in \mathbb{D}
$$
の代わりに、滑らかな関数 $h : M_d\left(\mathbb{R}\right) \rightarrow \mathbb{R}$ で表現した同値な問題
$$
\min_{W\in M_d\left(\mathbb{R}\right)} F\left(W\right) \\
\mathrm{subject\ to\ } h\left(W\right) = 0
$$
を解きます。ここで、$F : M_d\left(\mathbb{R}\right) \rightarrow \mathbb{R}$ は重み $W$ に対する損失関数です。NO TEARS 論文では観測データの行列 $\mathbf{X}$ に対する $\ell_1$ 正則化二乗損失
$$F\left(W\right) = \frac{1}{2n}\left\|\mathbf{X} - \mathbf{X}W\right\|^2_{\mathrm{F}} + \lambda \left\|W\right\|_1$$
としています。フロベニウスノルム $\left\|A\right\|_{\mathrm{F}} = \sqrt{\sum_{i,j}\left|A_{ij}\right|^2} $ は行列 $A$ の成分を1列に並べたベクトルのユークリッドノルムであり、作用素ノルム $\left\|A\right\|_p$ は $ \max_{x\neq 0}\left(\left\|Ax\right\|_p / \left\|x\right\|_p\right)$ で定義されます（上式では $p=1$）。$h$ の台が巡回グラフの重みの集合と一致すなわち
$$\mathrm{supp}\left(h\right)=\left\{ W \in M_d\left(\mathbb{R}\right)\mid G\left(W\right) : \mathrm{cyclic} \right\}$$
が成り立ち、$h$ が何らかのしかたで「DAG っぽさ」を表し、$h$ とその勾配が簡単に計算できればこの言い換えは役に立ちそうです。そのような都合の良い $h$ はあるのかを見ていきます。
非巡回性の表現
上記のような連続最適化問題への言い換えを行うために、非巡回性を表現する $h$ の構成のしかたを考えてみます。
対角化可能で任意の固有値の絶対値が1未満の2値隣接行列
まず簡単な場合として、重みが $0$ か $1$ かしかない2値の場合を考えてみます。重みのないただの隣接行列 $B \in M_d\left(\left\{0,1\right\}\right)$ の任意の固有値が絶対値 $1$ 未満であり、$B$ は対角化可能であるとします。（後者の条件は論文に書いてませんでしたが、途中の計算が通らなかったので追加しました。線形代数苦手なので本当に必要なのかわかりません。DAG だったら自動的に対角化可能になるとかある？ わかる人いたら教えてください……。）この条件のもとで、以下の等式が成り立ちます。
$$\sum_{k \geq 0} B^k = \left(I-B\right)^{-1}$$
このことを証明してみましょう。$B$ の固有値が $\left(-1, 1\right)$ に含まれることから、$I - B$ の固有値が $\left(0, 2\right)$ に含まれるとわかります。なぜなら、$\left(\lambda I -B\right) x = 0$ を満たす固有値 $\lambda$ と固有ベクトル $x$ が存在するとき、$\lambda' = 1-\lambda$ とすれば $\left( \lambda' I - \left(I-B\right)\right) x = 0$ と書き換えられ、これは $I-B$ の固有値が
$$\lambda' = 1 - \lambda \in \left(1 - 1, 1 - \left(-1\right)\right) = \left(0,2\right)$$
であることを意味するからです。よって、$I - B$ の固有値が $0$ でないことから逆行列 $\left(I - B\right)^{-1}$ の存在が保証されます。任意の $n \geq 2$ について
$$\left(I - B\right)\sum_{k = 1}^{n-1} B^k = I - B^n $$
と展開できるため、両辺に逆行列を掛けて
$$\sum_{k = 1}^{n-1} B^k = \left(I-B^n\right)\left(I-B\right)^{-1}$$
が得られます。これは高校数学で学ぶ実数の等比数列の和の公式と同じ形です。$B$ が $D = P^{-1}BP$ と対角化できるとすると、これは
$$\sum_{k = 1}^{n-1} B^k = \left(PP^{-1} - \left(PDP^{-1}\right)^n\right)\left(I - B\right)^{-1} = P\left(I - D^n\right)P^{-1}\left(I-B\right)^{-1}$$
と変形できます。$B$ の固有値は絶対値 $1$ 未満なので、$B$ の固有値が並んだ行列 $D$ の冪乗 $D^n$ は $n \rightarrow \infty$ で零行列に収束し、右辺は $\left(I - B \right)^{-1}$ のみが残ります。これで所望の等式が得られました。
さて、これを使って $B$ が非巡回であることについて考えてみましょう。$B^k$ の $\left(i,j\right)$ 成分は変数 $i$ から変数 $j$ へ $k$ ステップで到達する道の数に対応します。実際、
$$x_1 \rightarrow x_2 \rightarrow x_3 \rightarrow x_1 \rightarrow \cdots$$
という閉路を持つ簡単な巡回グラフを考えてみると、隣接行列は
B =
\begin{pmatrix}
0 & 1 & 0 \\
0 & 0 & 1 \\
1 & 0 & 0 \\
\end{pmatrix}
となりますが、この冪乗を計算してみると
B^2 =
\begin{pmatrix}
0 & 0 & 1 \\
1 & 0 & 0 \\
0 & 1 & 0 \\
\end{pmatrix}
, B^3 =
\begin{pmatrix}
1 & 0 & 0 \\
0 & 1 & 0 \\
0 & 0 & 1 \\
\end{pmatrix}
と2ステップ先・3ステップ先で到達できる変数のところに $1$ が入っています。複数の行き方ができるようなグラフを考えると $1$ より大きい自然数が入るので、適当なグラフを考えて計算してみてください。話を戻しましょう。いま興味があるのは非巡回性、すなわち閉路がないことでした。閉路があるとは、何ステップ後かに元いたノードへ戻ってしまうことを指します。であれば、任意のステップ数 $k \geq 1$ に対して $B^k$ の対角成分 $\left(B^k\right)_{ii} \ \left(i = 1, 2, \cdots, d\right)$ がすべて $0$ でありさえすれば非巡回性は保証できます。対角成分だけに興味があるので、トレース（対角成分の和）を見ればよさそうです。先ほど得た等式を利用するために右辺 $\left(I-B\right)^{-1}$ のトレースを取ると、
$$\mathrm{tr}\left(\left(I-B\right)^{-1}\right) = \mathrm{tr} \left(\sum_{k\geq 0}B^k\right) = \mathrm{tr}\left(I\right) + \sum_{k\geq 1}\mathrm{tr}\left(B^k\right) =d + 0 = d $$
となります。つまり、固有値が絶対値 $1$ 未満の対角化可能な隣接行列に対しては
$$\mathrm{tr}\left(\left(I-B\right)^{-1}\right) = d$$
が非巡回性の必要十分条件というわけです。
一般の2値隣接行列
とはいえ、実際には一般の隣接行列についても考えたくなります。その場合についても、同じような議論で代わりの条件が考えられます。等比級数の公式を使わずとも、任意のステップ数 $k \geq 1$ について対角成分が $0$ であることさえ使えば
$$\mathrm{tr}\left(\exp B \right) = \mathrm{tr}\left(\sum_{k \geq 0} \frac{1}{k!} B^k \right) = \sum_{k\geq 0}\sum_{i=1}^d \frac{1}{k!}\left(B^k\right)_{ii} = \sum_{i=1}^d I_{ii} + \sum_{k\geq 1}\sum_{i=1}^d \frac{1}{k!}\left(B^k\right)_{ii} = d + 0 = d$$
と指数関数を用いることで非巡回性が特徴づけできます。
一般の重み行列
さらに一般化して、任意の重み $W \in M_d\left(\mathbb{R}\right)$ を持つグラフについても考えてみましょう。重みが非負の場合については、2値の場合と同じように非巡回性と任意の $k \geq 1$ に対し $W^k$ の対角成分がすべて $0$ になることが同値になります。負の重みがある場合は、実際には道があっても重みの差し引きで $0$ になってしまうことがあるため、同じようにうまくはいきません。こういうときは、2乗してしまえばいいのです。興味があるのは対角成分が $0$ であるかどうかだけなので、各成分を2乗してやってもその判定に影響はありません。よって、要素積 $W\odot W$ を考えて
$$\mathrm{tr}\left(\exp\left(W\odot W\right)\right) = d$$
となっているかどうかを見ることで非巡回性が判定できます。
以上より、この記事の最初で述べたような最適化問題の書き換え
$$
\min_{W\in M_d\left(\mathbb{R}\right)} F\left(W\right) \\
\mathrm{subject\ to\ } h\left(W\right) = 0
$$
は $h\left(W\right) = \mathrm{tr}\left(\exp\left(W\odot W\right)\right) - d$ としてやればよいことがわかります。$h$ は滑らかで、勾配は
$$\nabla h\left(W\right) = \left(\exp\left(W\odot W\right)\right)^\top \odot 2W$$
とコンピュータで計算しやすい形になります。また、閉路が増えると重み行列の各要素の絶対値が大きくなるため、$h$ は大きくなり、$h$ が小さいほど DAG に近づくという点において「DAG っぽさ」を表現できているといえます。
NO TEARS アルゴリズムの全体
NO TEARS アルゴリズムでは拡張ラグランジュ乗数法を用いてこの連続最適化問題を解いています。アルゴリズム全体は以下のようになっています。
[NO TEARS]
- 重み行列とラグランジュ乗数の初期値 $\left(W_0, \alpha_0\right)$、拡張ラグランジュ乗数法の向上率 $c\in\left(0,1\right)$、ペナルティパラメータの初期値 $\rho_0 > 0$、$h$ の許容誤差 $\varepsilon >0$、重みの下限 $\omega > 0$ を与える。
- c の条件を満たすまで以下を繰り返す。
a. 拡張ラグランジュ乗数法により $L^\rho\left(W,\alpha_t\right)$ を最小化する $W^\ast$ を見つけ、$W_{t+1}\leftarrow W^\ast$ と更新する。
b. ラグランジュ乗数を更新する：$\alpha_{t+1}\leftarrow\alpha_t + \rho h\left(W_{t+1}\right)$
c. $h\left(W_{t+1}\right) < \varepsilon$ であれば $\tilde{W}_{\mathrm{ECP}} = W_{t+1}$ として終了する。 - $\tilde{W}_{\mathrm{ECP}}$ の要素で絶対値が $\omega$ 以下であるものを $0$ に置き換えた重み行列 $\hat{W} = \tilde{W}_{\mathrm{ECP}}\odot 1 \left(\left|\tilde{W}_{\mathrm{ECP}}\right|>\omega\right)$ を返す。
ただし、$1\left(\varphi\right) \in M_d\left(\left\{0,1\right\}\right)$ は条件 $\varphi$ を満たす成分を $1$ として $\varphi$ を満たさない成分を $0$ とする行列を表します。また、NO TEARS ライブラリではペナルティパラメータの最大値 $\rho_{\max}$ を設定し、$\rho$ がこれを超えたら更新をストップできるようになっています。
拡張ラグランジュ乗数法
まずはラグランジュ（の未定）乗数法について説明しておきましょう。等式で表現される制約 $h\left(x\right) = 0$ と最小化したい関数 $F$ に対し、ラグランジュ関数 $L\left(x,\alpha\right) = F\left(x\right) + \alpha h\left(x\right)$ の微分が $\nabla L\left(x^\ast,\alpha^\ast\right) = 0$ となるような点 $x^\ast$ を解の候補として探索するのがラグランジュ乗数法です。
拡張ラグランジュ乗数法では、このラグランジュ関数にペナルティ項を加えた
$$L^\rho\left(x,\alpha\right) = F\left(x\right) + \alpha h\left(x\right) + \frac{\rho}{2}\left|h\left(x\right)\right|^2$$
を最小化するような点を解の候補と考えます。その際、各ステップで $\alpha_{t+1} \leftarrow \alpha_t + \rho h\left(x\right)$ とパラメータを更新していきます。NO TEARS ライブラリでは L-BFGS-B アルゴリズムを用いて最小化問題を解いているみたいです。その結果として向上率 $c$ に対し $h\left(W_{t+1}\right) < c h\left(W_t\right)$ が満たされていなければ、$\rho$ を10倍して再度 L-BFGS-B で最小化問題を解き、$h\left(W_{t+1}\right) < c h\left(W_t\right)$ となれば手順 b へ進みます。
以上が DAG 構造学習アルゴリズム NO TEARS の説明です。NO TEARS 以降にも NO BEARS や NO FEARS や NO CURL といった改良版が出ているみたいなので、読んで行間埋めができたらまとめるかもしれません。