一般化線形モデルに対する信念伝搬法 (belief propagation, BP) を、大自由度極限で近似することで、一般化近似メッセージ伝搬法 (generalized approximate message passing, GAMP) を導出する計算ノートです。

計算ノート（PDF）

GAMPについて

一般化線形モデルにBPを直接適用すると、測定ノード $\mu$ と変数ノード $i$ の各辺についてキャビティメッセージを保持する必要があります。また、因子から変数へのメッセージ更新には、多数の変数についての積分が現れます。変数数を $N$、観測数を $M$ とすると、辺に依存する量だけでも $O(MN)$ 個存在します。

GAMPでは、測定行列の各成分が $A_{\mu i}=O(N^{-1/2})$ である高次元極限を利用します。まず、多数の弱い寄与の和として現れるキャビティ場をガウス分布で近似します。さらに、辺 $(\mu,i)$ ごとに異なるキャビティ量を、因子ごと・変数ごとの量のまわりで展開します。その結果、辺ごとのメッセージを保持せずに更新を閉じることができます。

問題設定

測定行列を $\bm{A}=(A_{\mu i})\in\mathbb{R}^{M\times N}$、観測データを $\bm{y}\in\mathbb{R}^M$ とし、次の目的関数を考えます。

$$ H_{\bm y,\bm A}(\bm x) = \sum_\mu \ell(y_\mu,[\bm A\bm x]_\mu) + \sum_i J(x_i). $$

$\ell(y,z)$ はoutput側の損失関数、$J(x)$ は各変数に作用するinput側の正則化関数です。この目的関数に対応するGibbs分布を構成し、その周辺分布をBPによって近似するところから出発します。

通常のAMPでは二乗誤差型の観測モデルを考えることが多いのに対し、GAMPではoutput側を一般の損失関数 $\ell(y,z)$ として扱えます。このoutput側の一般化によって、変数側の正則化だけでなく、観測過程に応じた非線形な応答も更新へ組み込めます。

GAMPのアルゴリズム

標準ガウス測度を $D\xi=e^{-\xi^2/2}d\xi/\sqrt{2\pi}$ とします。有限温度におけるoutput自由エントロピーとinput自由エントロピーを、それぞれ次のように定義します。

$$ \begin{aligned} \phi^{\mathrm{out}}(y,\omega,V) &\coloneqq \log\int D\xi\, \exp\!\left\{-\beta\ell\!\left(y,\omega+\sqrt{V}\xi\right)\right\}, \\ \phi^{\mathrm{in}}(B,\Lambda) &\coloneqq \log\int dx\, \exp\!\left\{-\frac{\Lambda}{2}x^2+Bx-\beta J(x)\right\}. \end{aligned} $$

このとき、GAMPの更新は次のアルゴリズムとしてまとめられます。

アルゴリズム (GAMP):

  1. 初期値 $\widehat{\bm{x}}^{[0]}$ と $\bm{v}^{[0]}$ を選び、$\bm{g}^{[-1]}=\bm{0}$ とする。

  2. 推定値の変化が十分に小さくなるまで、$t=0,1,2,\ldots$ に対して以下を繰り返す:

    1. output側の平均と分散を更新する。

      $$ \begin{aligned} V_\mu^{[t]} &=\sum_i A_{\mu i}^2v_i^{[t]}, \\ \omega_\mu^{[t]} &=\sum_i A_{\mu i}\widehat{x}_i^{[t]} -g_\mu^{[t-1]}V_\mu^{[t]}. \end{aligned} $$

    2. output側の応答を更新する。

      $$ \begin{aligned} g_\mu^{[t]} &=\partial_\omega\phi^{\mathrm{out}} \!\left(y_\mu,\omega_\mu^{[t]},V_\mu^{[t]}\right), \\ \Gamma_\mu^{[t]} &=-\partial_\omega^2\phi^{\mathrm{out}} \!\left(y_\mu,\omega_\mu^{[t]},V_\mu^{[t]}\right). \end{aligned} $$

    3. input側の有効場と精度を更新する。

      $$ \begin{aligned} \Lambda_i^{[t]} &=\sum_\mu A_{\mu i}^2\Gamma_\mu^{[t]}, \\ B_i^{[t]} &=\sum_\mu A_{\mu i}g_\mu^{[t]} +\widehat{x}_i^{[t]}\Lambda_i^{[t]}. \end{aligned} $$

    4. 変数ごとの平均と分散を更新する。

      $$ \begin{aligned} \widehat{x}_i^{[t+1]} &=\partial_B\phi^{\mathrm{in}} \!\left(B_i^{[t]},\Lambda_i^{[t]}\right), \\ v_i^{[t+1]} &=\partial_B^2\phi^{\mathrm{in}} \!\left(B_i^{[t]},\Lambda_i^{[t]}\right). \end{aligned} $$

  3. 収束した $\widehat{\bm{x}}^{[t+1]}$ を推定値として出力する。

各更新量の役割は次のとおりです。

  • $\widehat{x}_i,v_i$: 変数 $x_i$ の平均と分散。
  • $\omega_\mu,V_\mu$: output側へ渡される平均と分散。
  • $g_\mu,\Gamma_\mu$: 損失関数が定めるoutput側の応答。
  • $B_i,\Lambda_i$: input側の有効場と精度。

これらを上から順に更新し、再び $\widehat{x}_i$ と $v_i$ を求める循環を繰り返します。

導出過程

GAMPの導出では、BPの因子更新に含まれる多変数積分をガウス分布での近似によって簡略化し、その後も残る辺依存のキャビティ量を展開によって消去します。大まかな流れは以下のとおりです。

$$ \text{BP} \quad \underset{\text{ガウス分布での近似}}{\longrightarrow} \quad \text{rBP} \quad \underset{\text{キャビティ量の消去}}{\longrightarrow} \quad \text{GAMP} $$

ここでrBPは、キャビティ場をガウス分布で近似したBP、いわゆる relaxed BP を表します。BPからrBPへの移行によって多変数積分が1次元積分へ置き換わり、rBPからGAMPへの移行によって辺ごとの量が消去されます。

BPからrBPへ: ガウス分布での近似

BPの因子メッセージには、次のキャビティ場が現れます。

$$ u_{\mu\to i} = \sum_{j\ne i}A_{\mu j}x_j. $$

$A_{\mu j}=O(N^{-1/2})$ のもとでは、この量は多数の弱い寄与の和なのでガウス分布で近似できます。キャビティ分布に関する平均と分散を $\widehat{x}_{j\to\mu}$、$v_{j\to\mu}$ と書けば、キャビティ場は次のように表せます。

$$ u_{\mu\to i} \simeq \omega_{\mu\to i}+\sqrt{V_{\mu\to i}}\,\xi, \qquad \xi\sim\mathcal{N}(0,1). $$

ここで、キャビティ場の平均と分散を次のように定義しています。

$$ \omega_{\mu\to i} =\sum_{j\ne i}A_{\mu j}\widehat{x}_{j\to\mu}, \qquad V_{\mu\to i} =\sum_{j\ne i}A_{\mu j}^2v_{j\to\mu}. $$

これにより、BPの因子更新に含まれていた変数 $j\ne i$ に関する多変数積分が、単一の標準ガウス変数 $\xi$ に関する1次元積分へ置き換わります。このガウス分布での近似後の更新がrBPです。

rBPからGAMPへ: キャビティ量の消去

rBPでは積分自体は簡単になりましたが、依然として次のような辺 $(\mu,i)$ に依存する量を保持する必要があります。

$$ \omega_{\mu\to i},\quad V_{\mu\to i},\quad g_{\mu\to i},\quad \Gamma_{\mu\to i},\quad B_{i\to\mu},\quad \Lambda_{i\to\mu}. $$

そこでGAMPでは、これらを辺に依存しない $\omega_\mu,V_\mu,g_\mu,\Gamma_\mu,B_i,\Lambda_i$ のまわりで展開します。

例えば、有効場とキャビティ有効場の差、およびoutput場とキャビティoutput場の差は、次のように表されます。

$$ B_{i\to\mu}=B_i-A_{\mu i}g_{\mu\to i}, \qquad \omega_{\mu\to i}=\omega_\mu-A_{\mu i}\widehat{x}_{i\to\mu}. $$

どちらの差も $O(N^{-1/2})$ です。この差を保って応答とモーメントを展開すると、次の式が得られます。

$$ g_{\mu\to i} =g_\mu+A_{\mu i}\widehat{x}_i\Gamma_\mu+O(N^{-1}), \qquad \widehat{x}_{i\to\mu} =\widehat{x}_i-A_{\mu i}g_\mu v_i+O(N^{-1}). $$

これらを場の定義に戻すと、$O(MN)$ 個の辺依存量が消え、因子ごとの $O(M)$ 個の量と変数ごとの $O(N)$ 個の量だけで更新が閉じます。

ポイント

Onsager反作用項

キャビティ量の展開を通して $O(1)$ の項を残していくと、次のOnsager反作用項が現れます。

$$ -g_\mu V_\mu, \qquad \widehat{x}_i\Lambda_i. $$

これらは、ある辺を除いたキャビティ量と、すべての辺を含む量との差を補正する効果があります。単純に辺添字を落とすだけでは、自分自身の寄与が更新を介して戻ってくる効果を数え込んでしまいます。Onsager反作用項はその自己反作用を相殺する役割を持ちます。

零温度GAMP

有限温度GAMPの $v_i$ と $V_\mu$ は、$\beta\to\infty$ では消失します。一方で、これらを逆温度でスケーリングした次の量は有限な値として残ります。

$$ \chi_i=\beta v_i, \qquad \tau_\mu=\beta V_\mu. $$

このスケーリングのもとで、有限温度の自由エントロピーは1変数の最小化問題へ帰着します。output側の最小点は次のように表されます。

$$ r_\mu^* = \argmin_r \left\{ \ell(y_\mu,r)+\frac{(r-\omega_\mu)^2}{2\tau_\mu} \right\} = \operatorname{prox}_{\tau_\mu\ell(y_\mu,\cdot)}(\omega_\mu). $$

input側の最小点は次のように表されます。

$$ \widehat{x}_i = \argmin_x \left\{ \frac{\widetilde{\Lambda}_i}{2}x^2 -\widetilde{B}_i x +J(x) \right\}. $$

すると、零温度GAMPではoutput側とinput側の局所更新を近接演算子によって表せることになります。有限温度での平均・分散の更新が、零温度では最小点とその感受率の更新に移り変わると見ることができます。

AMPとの関係

GAMPはAMPと別系統のアルゴリズムではなく、output側を一般化した枠組みとして見ることができます。ここでは、次の二乗誤差損失を選びます。

$$ \ell(y,u) = \frac{1}{2\lambda}(y-u)^2. $$

この場合、ガウス積分を実行することでoutput側の応答は次のように簡単になります。

$$ \Gamma_\mu = \frac{\beta}{\lambda+\beta V_\mu}, \qquad g_\mu = \Gamma_\mu(y_\mu-\omega_\mu). $$

さらに、残差 $z_\mu$ と有効場 $h_i$ を次のように定義します。

$$ z_\mu=\frac{g_\mu}{\Gamma_\mu}=y_\mu-\omega_\mu, \qquad h_i=\frac{B_i}{\Lambda_i}. $$

これらと $A_{\mu i}^2\simeq 1/N$ を用いて分散と精度をスカラー化し、さらに各成分に作用するデノイザーを

$$ \begin{aligned} \eta_\beta\!\left(h;\frac{\lambda+\chi}{\alpha}\right) \coloneqq \frac{ \displaystyle \int d x\,x \exp\!\left\{ -\beta\left[ \frac{\alpha}{2(\lambda+\chi)}(x-h)^2+J(x) \right] \right\} }{ \displaystyle \int d x\, \exp\!\left\{ -\beta\left[ \frac{\alpha}{2(\lambda+\chi)}(x-h)^2+J(x) \right] \right\} }. \end{aligned} $$

と定義することによって、通常のAMP更新式を得ることができます。

$$ \begin{aligned} \bm{z}^{[t]} &= \bm{y} - \bm{A}\bm{x}^{[t]} + \frac{\bm{z}^{[t-1]}}{\alpha}\frac{1}{N}\sum_{i=1}^N \eta_\beta'\!\left(h_i^{[t-1]};\frac{\chi^{[t-1]}+\lambda}{\alpha}\right), \\ \bm{h}^{[t]} &= \bm{x}^{[t]} + \frac{1}{\alpha}\bm{A}^{\mathsf T}\bm{z}^{[t]}, \\ \bm{x}^{[t+1]} &= \eta_\beta\!\left(\bm{h}^{[t]};\frac{\chi^{[t]}+\lambda}{\alpha}\right), \\ \chi^{[t+1]} &= \frac{\chi^{[t]}+\lambda}{\alpha}\frac{1}{N}\sum_{i=1}^N \eta_\beta'\!\left(h_i^{[t]};\frac{\chi^{[t]}+\lambda}{\alpha}\right). \end{aligned} $$