線形逆問題に対する信念伝搬法から近似メッセージ伝搬法を導出する

線形逆問題に対する信念伝搬法を大自由度で近似することで近似メッセージ伝搬法 (AMP) を導出する計算ノートです。

近似メッセージ伝搬法について

問題設定

$N$ 次元の信号 $\bm{x}^0$ を、$M$ 回の無雑音線形観測から復元する問題を考えます。観測データは $y_\mu=\sum_{i=1}^N A_{\mu i}x_i^0$ によって生成されるとします。測定行列 $\bm{A}\in\mathbb{R}^{M\times N}$ と観測データ $\bm{y}\in\mathbb{R}^M$ が与えられたとき、ここでは次の制約付き最小化問題を扱います。

\begin{aligned} \min_{\bm{x}\in\mathbb{R}^N} \quad \sum_{i=1}^N J(x_i) \quad \text{s.t.} \quad \bm{y}=\bm{A}\bm{x}. \end{aligned}

$J$ は各成分に作用する正則化関数です。例えば $J(x)=|x|$ とすると、$\ell_1$ ノルム最小化が得られます。

測定率は典型的に $M/N \approx \alpha$ であるとします。また、測定行列の各成分は独立に $A_{\mu i}\sim\mathcal{N}(0,1/N)$ に従うと仮定します。この規格化では、各行の二乗ノルムは $1$ に、各列の二乗ノルムは $\alpha$ に集中します。

\begin{aligned} \sum_{i=1}^N A_{\mu i}^2 &\underset{N \to \infty}{\to} 1, & \sum_{\mu=1}^M A_{\mu i}^2 &\underset{N \to \infty}{\to} \alpha. \end{aligned}

AMPアルゴリズム

一変数の非線形関数 $\eta$ を次式で定義します。

\eta(h;\lambda) \coloneqq \operatorname{prox}_{\lambda J}(h) = \argmin_{x\in\mathbb{R}} \left\{ \frac{1}{2\lambda} (x-h)^2 +J(x) \right\}.

上の最小化問題の最小点が一意である場合、$\eta$ は単一値の関数です。また、ベクトル $\bm{h}=(h_1,\ldots,h_N)^{\mathsf T}\in\mathbb{R}^N$ に対しては、$\eta$ が各成分に作用するものとし、同じ記号を用いて次のように定義します。

\eta(\bm{h};\lambda) \coloneqq \left( \eta(h_1;\lambda),\ldots,\eta(h_N;\lambda) \right)^{\mathsf T} \in\mathbb{R}^N.

アルゴリズム (AMP):

初期値 $\widehat{\bm{x}}^{[0]}$ と $\lambda^{[0]}>0$ を選び、$\bm{z}^{[0]}=\bm{y}-\bm{A}\widehat{\bm{x}}^{[0]}$ とする。
以下を推定信号 $\widehat{\bm{x}}$ が収束するまで $t=0,1,2,\ldots$ に対して更新する。
$\begin{aligned} \bm{h}^{[t]} &= \widehat{\bm{x}}^{[t]} + \frac{1}{\alpha}\bm{A}^{\mathsf T}\bm{z}^{[t]}, \\ \widehat{\bm{x}}^{[t+1]} &= \eta\!\left(\bm{h}^{[t]}; \lambda^{[t]}\right), \\ \bm{z}^{[t+1]} &= \bm{y} - \bm{A}\widehat{\bm{x}}^{[t+1]} + \frac{\bm{z}^{[t]}}{\alpha}\frac{1}{N}\sum_{i=1}^N \eta'\!\left(h_i^{[t]}; \lambda^{[t]}\right), \\ \lambda^{[t+1]} &= \frac{\lambda^{[t]}}{\alpha}\frac{1}{N}\sum_{i=1}^N \eta'\!\left(h_i^{[t]}; \lambda^{[t]}\right). \end{aligned}$
$\widehat{\bm{x}}^{[t+1]}$ を推定信号として出力して終了する。

残差更新に含まれる次の第3項はOnsager反作用項などと呼ばれます。

\frac{\bm{z}^{[t]}}{\alpha}\frac{1}{N}\sum_{i=1}^N \eta'\!\left(h_i^{[t]};\frac{\chi^{[t]}}{\alpha}\right).

解釈

ISTAとの関係

得られた更新式は、反復縮小しきい値法 (iterative shrinkage thresholding algorithm, ISTA) に補正を加えたものとして解釈できます。学習率 $\eta^\text{lr}>0$ とデノイザーの実効的な正則化係数 $\lambda>0$ を用いると、ISTA型の更新は次式で表されます。

\begin{aligned} \bm{r}^{[t]} &= \bm{y} - \bm{A}\bm{x}^{[t]}, \\ \bm{h}^{[t]} &= \bm{x}^{[t]} + \eta^\text{lr} \bm{A}^{\mathsf T}\bm{r}^{[t]}, \\ \bm{x}^{[t+1]} &= \eta\!\left(\bm{h}^{[t]};\lambda\right). \end{aligned}

一方、AMPの更新は次式です。

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

両者を比べると、本記事のAMPは次の3つの変更を加えた形と解釈することができます。

ISTAの学習率を $\eta^\text{lr}=1/\alpha$ に固定
通常の残差 $\bm{r}^{[t]}$ の代わりにOnsager反作用項を含む $\bm{z}^{[t]}$ を使用
デノイザーの正則化係数を固定せず $\lambda^{[t]}$ として反復ごとに更新

正則化係数のスケジューリング

正則化係数 $\lambda^{[t]}=\chi^{[t]}/\alpha$ の更新規則は次式で与えられます。

\lambda^{[t+1]} = \frac{\lambda^{[t]}}{\alpha}\frac{1}{N}\sum_{i=1}^N \eta'\!\left( h_i^{[t]}; \lambda^{[t]} \right).

これはBPとの対応から導かれたスケジューリングですが、正則化係数の選び方はこの更新則に限られません。例えば、正則化係数を固定する方法、あらかじめ定めた規則に従って減少させる方法、残差や反復の進行に応じて調整する方法などが考えられます。ただし、別のスケジューリングを用いる場合、その反復法は本節で導出したBP近似とは異なるアルゴリズムになるため、収束性や推定性能は別途検討する必要があります。

正則化線形回帰への拡張

ここまでは線形制約 $\bm{y}=\bm{A}\bm{x}$ を課した問題を扱いましたが、同じ考え方は次の正則化線形回帰にも適用できます。

\min_{\bm{x}\in\mathbb{R}^N} \left\{ \frac{1}{2 \gamma}\| \bm{y}-\bm{A}\bm{x} \|_2^2 + \sum_{i=1}^N J(x_i) \right\}, \qquad \gamma>0.

この場合は、因子ノードに用いたデルタ関数を、$\gamma$ を分散スケールとする次のGaussian因子へ置き換えればよいです。

\delta\!\left( y_\mu-\sum_{i=1}^N A_{\mu i}x_i \right) \quad\longrightarrow\quad \sqrt{\frac{\beta}{2\pi\gamma}}\exp\!\left(-\frac{\beta}{2\gamma}\left( y_\mu-\sum_{i=1}^N A_{\mu i}x_i \right)^2\right).

このGaussian因子を用いると、因子メッセージの分散には $\gamma$ が加わるため、デノイザーの実効的な正則化係数は $(\chi^{[t]}+\gamma)/\alpha$ となります。したがって、正則化線形回帰に対するAMP更新式は次のように書けます。

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

$\gamma=0$ とすれば、デノイザーの引数と $\chi^{[t+1]}$ の更新は線形制約付き問題と同じ形になります。無雑音で線形制約が実行可能な場合、$\gamma\to0$ ではGaussian因子がデルタ関数へ集中するため、線形制約付きAMPが回復します。

固定点と最適解

$J(x_i)$ が凸関数の場合、AMPの固定点 $(\bm{x}^*,\bm{z}^*,\chi^*)$ は、元の最適化問題の停留点条件を満たします。このことを見てみましょう。

先に有限の $\gamma>0$ に対する固定点条件から調べます。以下、$J(\bm{x})=\sum_{i=1}^N J(x_i)$ とおきます。まず、固定点におけるデノイザーの更新式と $\bm{h}^*=\bm{x}^*+\alpha^{-1}\bm{A}^{\mathsf T}\bm{z}^*$ から、デノイザーの停留点条件が得られます。

\bm{x}^*=\eta\!\left(\bm{h}^*;\frac{\chi^*+\gamma}{\alpha}\right) \quad\longrightarrow\quad \bm{0} \in \partial J(\bm{x}^*) - \frac{1}{\chi^*+\gamma}\bm{A}^{\mathsf T} \bm{z}^*.

次に、$\chi$ の固定点条件から、Onsager反作用項の係数が得られます。

\chi^*=\frac{\chi^*+\gamma}{\alpha} \frac{1}{N} \sum_{i=1}^N\eta'\!\left(h_i^*;\frac{\chi^*+\gamma}{\alpha}\right) \quad\longrightarrow\quad \frac{1}{\alpha N}\sum_{i=1}^N\eta'\!\left(h_i^*;\frac{\chi^*+\gamma}{\alpha}\right)=\frac{\chi^*}{\chi^*+\gamma}.

これを残差の固定点条件へ代入すると、$\bm{z}^*$ と通常の残差の関係が得られます。

\bm{z}^*=\bm{y}-\bm{A}\bm{x}^*+\frac{\chi^*}{\chi^*+\gamma}\bm{z}^* \quad\longrightarrow\quad \frac{\bm{z}^*}{\chi^*+\gamma}=\frac{\bm{y}-\bm{A}\bm{x}^*}{\gamma}.

この関係式をデノイザーの停留点条件に代入することで次式が得られます。これは正則化線形回帰の目的関数の劣微分がゼロになることと等価です。

\boxed{ \begin{aligned} \bm{0} \in \partial J(\bm{x}^*)-\frac{1}{\gamma}\bm{A}^{\mathsf T}(\bm{y}-\bm{A}\bm{x}^*). \end{aligned} }

$\gamma=0$ のバージョンでも同様に収束先を求めることができます。この条件ではOnsager反作用項の係数が1へ収束し、残差の固定点条件から $\bm{y}-\bm{A}\bm{x}^*=\bm{0}$ が得られ、AMPの収束先は次式を満たすこととなります。これは元の制約付き最適化問題のKKT条件に一致します。

\boxed{ \begin{aligned} \bm{0}\in\partial J(\bm{x}^*)-\frac{1}{\chi^*}\bm{A}^{\mathsf T}\bm{z}^*, \qquad \bm{y}-\bm{A}\bm{x}^*=\bm{0}. \end{aligned} }

$J$ が凸であれば、有限 $\gamma$ の停留点条件と $\gamma\to0$ のKKT条件はいずれも最適性の十分条件であるため、AMPの固定点 $\bm{x}^*$ は対応する最適化問題の最適解に一致します。注意すべき点として、非凸な $J$ では、固定点は一般に停留点の候補であって、大域的最適解とは限りません。

補足: 行列のスケーリングについて

本ノートでは $A_{\mu i}\sim\mathcal{N}(0,1/N)$ と仮定しています。この規格化では、各列の二乗ノルムが次の値へ集中します。

\sum_{\mu=1}^M A_{\mu i}^2 \longrightarrow \frac{M}{N} = \alpha.

そのため、有効場の更新には $\bm{A}^{\mathsf T}\bm{z}$ とともに $1/\alpha$ が現れます。

\bm{h}^{[t]} = \bm{x}^{[t]} + \frac{1}{\alpha} \bm{A}^{\mathsf T}\bm{z}^{[t]}.

文献 ^[1] などでは $A_{\mu i}\sim\mathcal{N}(0,1/M)$ という規格化が使われます。この規格化では列の二乗ノルムが1へ集中するため、有効場やOnsager反作用項に現れる係数の配置が変わります。

「伝搬」「伝播」問題

Message Passing の日本語表記には「メッセージ伝搬法」と「メッセージ伝播法」があり、現在の研究文献でも統一されていません。

「伝搬」は圧縮センシングや信号処理の文献で用例が見られ、たとえば文献^[2]では「近似メッセージ伝搬法」と表記されています。一方、SITA2023のプログラム^[3]では、セクション名として明確に「伝播法」の表記が採用されています。また、Belief Propagation にも「確率伝搬法」という用例があります^[4]。SITA2025のプログラム^[5]に至っては「直交近似的メッセージ伝播法」と「近似メッセージ伝搬法」が並んでいます。

このように、アルゴリズムの種類によって訳語が明確に分かれているわけでもないようです。

D. L. Donoho, A. Maleki, and A. Montanari, Message Passing Algorithms for Compressed Sensing: I. Motivation and Construction, in 2010 IEEE Information Theory Workshop on Information Theory (ITW 2010, Cairo) (2010), pp. 1–5. ↩︎
早川諒, 林和則, 近似メッセージ伝搬法を用いた離散値ベクトルの再構成, 電子情報通信学会2017年総合大会. www.ieice.org/publications/conferences/summary.php ↩︎
電子情報通信学会, 「第46回情報理論とその応用シンポジウム (SITA 2023) プログラム」, セッション 2.3「メッセージ伝播法」, 2023. www.ieice.orf/ess/sita/SITA2023. ↩︎
林和則, 「確率伝搬法とその応用」, 数理解析研究所講究録 1616, 16–40 (2008). ↩︎
電子情報通信学会, 「第48回情報理論とその応用シンポジウム (SITA 2025) プログラム」, セッション 5.4「信号処理基礎」, 2025. www.ieice.org/ess/sita/SITA2025. ↩︎