Blahut-Arimoto のアルゴリズム
概要
情報理論において歪みありの符号化を非可逆符号化と呼びます。非可逆符号化において、歪み度合い $D$ を変えたときの最小レート（1 symbol 当たりのビット数）*1 [bit/symbol] にはトレードオフの関係にあることが知られており、この関係はレート歪み関数 $R(D)$ によって記述されます。 Blahut-Arimoto のアルゴリズムは、レート歪み関数 $R(D)$ を逐次的に求める数値計算アルゴリズムです。
レート歪み関数
入力を符号化するエンコーダがあり、符号をノイズが入った通信路に入れて、通信路から受け取った符号を元の入力に復元するデコーダがある状況を考えます。 ただし、入力は確率分布 $p(x)$ に従う確率変数 $X$ 、復元された入力は確率変数 $\hat{X}$ で表します。 このエンコーダ・通信路・デコーダを通して復元された入力 $\hat{X}$ は条件付き確率 $q(\hat{x}|x)$ で表現されるとします*2。
このような状況の下でレート歪み関数を (1) 式で定義します。
\begin{align}
R(D) := \min_{q(\hat{x}|x) : \langle d(x, \hat{x}) \rangle \leq D } I(X; \hat{X}) \tag{1}
\end{align}
ここで $d(x, \hat{x})$ は $x$ と $\hat{x}$ の歪み（距離）関数であり、$\langle d(x, \hat{x}) \rangle$ は同時分布 $p(x, \hat{x})$ に関する期待値で、$\langle d(x, \hat{x}) \rangle = \sum_{x, \hat{x}} p(x, \hat{x})d(x, \hat{x}) = \sum p(x)q(\hat{x}|x)d(x, \hat{x})$ です。 また、$I(X; \hat{X})$ は相互情報量で (2) 式で定義されます。
\begin{aligned}
I(X; \hat{X}) = \sum_{x} \sum_{\hat{x}} p(x, \hat{x}) \log{\frac{q(\hat{x}|x)}{r(\hat{x})} } \tag{2}
\end{aligned}
ただし、$r(\hat{x})$ は $\hat{X}$ の従う確率分布です。ここでは (1) 式の右辺の最小化を各 $x$ に対する制約条件 $\sum_{\hat{x}} p(\hat{x} | x)=1$ がある変分問題として解きます。 具体的には 汎関数
\mathcal{F} \lbrack q(\hat{x}|x) \rbrack = I(X; \hat{X}) + \beta \langle d(x, \hat{x}) \rangle + \sum_{x} \lambda(x) \left( \sum_{\hat{x}} q(\hat{x} | x) - 1 \right) \tag{3}
を関数 $p(\hat{x}|x)$ について最小化します。$\beta, \lambda(x)$ はラグランジュ乗数です。 停留条件 $\frac{\partial F[q]}{\partial q} = 0$ を解くと、
\begin{aligned}
\frac{\partial \mathcal{F} \lbrack q \rbrack}{\partial q} &= \frac{\partial}{\partial q} \left\{ I(X; \hat{X}) + \beta \langle d(x, \hat{x}) \rangle + \sum_{x} \lambda(x) ( \sum_{\hat{x}} p(\hat{x} | x) - 1 ) \right\} \\
&= \frac{\partial}{\partial q} \left\{ \sum_{x, \hat{x}} p(x) q(\hat{x}|x) \log \frac{q(\hat{x}|x)}{r(\hat{x})} + \beta q(\hat{x}|{x})p(x)d(x, \hat{x}) + \lambda(x)(q(\hat{x}|x) - 1) \right\}\\
&= p(x)\log \frac{q(\hat{x}|x)}{r(\hat{x})} + p(x) \underbrace{-\sum_{\hat{x}}p(\hat{x})q(\hat{x}|x')\frac{1}{q(\hat{x})}p(x)}_{A} + \beta p(x) d(x, \hat{x}) + \lambda(x) = 0
\end{aligned}
となります。$A$ の部分が最初出てこなくて頭を抱えていたのですが、log を展開して（というか分母の方だけ見て）$q(\hat{x}|x)$ で偏微分すると
\begin{aligned}
- \frac{\partial}{\partial q} \sum_{x'}\sum_{\hat{x}} p(x) q(\hat{x}|x) \log{r(\hat{x})} &= -\frac{\partial}{\partial q} \sum_{x'}\sum_{\hat{x}} p(x') q(\hat{x}|x') \log{\sum_{x} p(x)q(\hat{x}|x) }\\
&= -p(x) \log r(\hat{x}) \underbrace{- \sum_{x'} p(x')q(\hat{x}|x') \frac{1}{r(\hat{x})} p(x)}_{A}
\end{aligned}
となり $A$ の項が出てきます。$r(\hat{x}) = \sum_{x'} q(\hat{x}|x') p(x')$ を使って整理すると、
\begin{aligned}
q(\hat{x}|x) &= Z \times r(\hat{x}) \exp{(-\beta d(x, \hat{x}))} \tag{4}
\end{aligned}
となります。ちゃんと $Z$ を求めると $Z(x, \lambda)=\exp{(-\lambda(x)/p(x))}$ ですが、$q(\hat{x}|x)$ は確率なので $Z$ はただの正規化係数（もしくは分配関数）です。そのため、(4) 式の正規化係数を除いた右辺について全ての状態の和を取った $\sum_{x} r(\hat{x}) \exp{(-\beta d(x, \hat{x}))} $ を $Z(x, \beta)$ として考えて良いはずです。 よって $\mathcal{F}[q]$ を最小にする $\hat{q}$ は、
\begin{aligned}
\hat{q}(\hat{x}|x) = \frac{ r(\hat{x}) \exp{(-\beta d(x, \hat{x}))} }{ \sum_{x} r(\hat{x}) \exp{(-\beta d(x, \hat{x}))} } \tag{5}
\end{aligned}
であることが分かりました。
Blahut-Arimoto のアルゴリズム
Blahut-Arimoto のアルゴリズムは 逐次的に $R(D)$ を求める数値計算アルゴリズムです。
手順
情報源の確率分布 $p(x)$、歪み関数 $d(x, \hat{x})$、$\hat{x}$ のとる値の集合と $\beta > 0$ が与えられているものとします。
- $r(\hat{x}), q(\hat{x}|x)$ を何らかの分布で初期化します (殆どは一様分布が使われるカモ？)。
- (6) 式の $q(\hat{x}|x)$ と $r(\hat{x})$ を交互に更新します。\begin{aligned} q(\hat{x}|x) &= \frac{r(\hat{x})\exp{(-\beta d(x, \hat{x}))}}{\sum{\hat{x}}r(\hat{x})\exp{(-\beta d(x, \hat{x}))}} \ r(\hat{x}) &= \sum{x}{p(x)q(\hat{x}|x)} \tag{6}\end{aligned}
- 十分 $q(\hat{x}|x)$ と $r(\hat{x})$ が収束するまで step 2 を回し、得られた確率分布から相互情報量 $I(X; X)$ を計算します。
やっていることは、(5) 式の $q$ （と $r$）に関する self-consistent 方程式を解いてるだけですね。 注意する点としては、歪み度合い $D$ を陽に与えることはせずに、逆温度 $\beta$ を与えることで Blahut-Arimoto アルゴリズムから求まる同時分布 $p(x, \hat{x})$ から $D=\langle d(x, \hat{x}) \rangle$ を計算しています。
証明
Cover T. の本の 10.8 節に一部証明が載っています。 https://www.inc.cuhk.edu.hk/InformationTheory/files/Abridged/Ch_9.pdf を見るのも良さそうです。
雑に何やっているのかを書くと、歪み制約を満たす同時分布 $p'(x, \hat{x})$ の集合 $A$ と分布の積 $q'(x, \hat{x}) = p(x)r(\hat{x})$ の集合 $B$ から $p'$ と $q'$ に関する最小化問題として$R(D)$ の右辺を書き直すと
\begin{aligned}
R(D) &= \min_{q} I(X; X)\\
&= \min_{q} p(x, \hat{x}) \log{\frac{q}{r}} \\
&= \min_{q} p(x, \hat{x}) \log{\left( \frac{q}{r} \cdot \frac{p}{p} \right) } \\
&= \min_{q' \in B} \min_{p' \in A} \sum_{x, \hat{x}} p' \log{\frac{p'}{q'}} \\
&= \min_{q' \in B} \min_{p' \in A} \sum_{x, \hat{x}} D_{\mathrm{KL}}(p' \| q') \\
\end{aligned}
となり、[Csiszár and Tusnády, 1984] で提案された Altenating Minimization を適用できて、最適解への収束が示せるようです（３行目から４行目の変形はてきとうにやっているので注意してください。$p'$ も $q'$ も $q(\hat{x}|x)$ に依存するので結局2つの最小化問題になる、と言われればそうかもしれませんが）。
例：ガウス型通信路で歪み関数が二乗誤差関数の場合
細かい説明は省いて結論を書くと、入力の確率分布が $X \sim \mathcal{N}(0, σ ^2)$ であるときのレート歪み関数は
\begin{aligned}
R(D) = \begin{cases}
\frac{1}{2}\log \frac{\sigma^2}{D}, & 0 \leq D \leq \sigma^2,\\
0, & D > \sigma^2
\end{cases}
\end{aligned}
で与えられます。それでは、このレート歪み関数が Blahut-Arimoto のアルゴリズムで得られるのかを確認しましょう。
実装
数式通り素直に書いたので多分実行速度は遅いです。本当はステップ2の終了条件をちゃんと設定したほうが良いと思いますが、今はステップ2を100イテレーション繰り返しています。簡単な例なのでこれだけで十分ですが、もう少し複雑な関数を扱う場合は注意する必要があるかもしれません。
のレート歪み曲線 (Cover T.の本のFigure 10.6 に対応)。赤線は解析的に求めたレート歪み曲線の関数をプロットしたもので、各点が $\beta$ を変えながら Blahut-Arimoto のアルゴリズムを実行した結果をプロットしたもの。
上の例は解析的に解けるような状況設定を考えましたが、一般にレート歪み関数 $R(D)$ を解析的に求めるのはかなり大変です。$p(x)$ に正規分布を使う場合や歪み関数に二乗誤差を使う場合などは $R(D)$ を解析的に求められますが、そうでない場合のほうが明らかに多いため、Blahut-Arimoto のアルゴリズムを使うと数値的にレート歪み関数 $R(D)$ の概形が分かって嬉しいです。
参考
- Cover T., "Elements of Information Theory", Chapter 10
- Blahut–Arimoto algorithm - Wikipedia
- https://arxiv.org/abs/physics/0004057
- http://curtis.ml.cmu.edu/w/courses/index.php/Alternating_Minimization
その他
- はてなブログで数式を書くのが本当に厳しいので、自前のブログサービスに切り替えるなどの違う方法を取る予定です。
*1:符号が値を取る空間の濃度に対数を取ったものをレートと呼ぶこともあり、これが同じなのか自分はよく分かっていませんでしたが、対数の底が 2 だと等価です。
*2:条件付き確率で表して良いのか、という疑問が浮かぶのは自然なように自分は思いますが、通信路のノイズが従う確率分布によって入力の従う確率分布が変化すると思えば自然な気がします。通信路ではないですが、変分オートエンコーダなどもエンコーダ、デコーダを条件付き分布で表現しており、考えは同じような気がします。参考: http://www.akita-pu.ac.jp/system/elect/ins/kusakari/japanese/teaching/InfoTheo/2009/note/7.pdf
