画像から実体の推移を予測して学習する論文を読んでみる

プラトンのイデア論では「本当にこの世に実在するのはイデアであって、我々が肉体的に感覚している対象や世界とはあくまでイデアの《似像》にすぎない」[1]としている。例えば3D迷路の場合、迷路内の自己位置が実体で、壁に囲まれた通路の視野が似像(画像)とするとプラトンのイデア論そのものである。３D迷路内を効率的に探査するため、画像から実体を推定し、その遷移予測から画像を復元する論文(以下TD_VAE)を読んでみる。

[1806.03107] Temporal Difference Variational Auto-Encoder

f:id:mabonki0725:20190202231633p:plain

(1)モデル

(a)潜在空間モデル

　3D迷路の自己位置が時系列で移動するとして、実体が $[z_1,\dots,z_T]$ と変動するに応じて観察 $[x_1,\dots,x_T]$ も変化する潜在空間モデルを考える。

　ここで実体の推移確率 $p(z_t|z_{t-1})$ を全観察 $[x_1,\dots,x_t]$ から推定するEncoder $q(z|x)$ を導入する。

　　　 $p(z_t|z_{t-1}) \approx q(z|x) = q(z_t|z_{t-1},\phi_t(x))$ (0)

ここで $\phi_t(x)$ は全観察 $[x_1,\dots,x_t]$ 上の関数である。

　また実体と観察の同時分布は尤度 $p(x_t|z_t)$ と推移確率 $p(z_t|z_{t-1})$ を使かうと次となる。

$p(x,z) = \Pi_t p_\theta(z_t|z_{t-1}) p(x_t | z_t)$

　この対数表現は

$\log p(x,z) = \sum_t \log p(z_t|z_{t-1}) + \log p(x_t|z_t)$

上式の実体 $z$ での期待値は

　　　 $\log p(x) = \int \log p(x,z) dz = \int \log p(x|z) p(z) dz = \mathbb{E}_{z \sim p(z|x)} [\log p(x|z) ]$

$z \sim p(z|x)$ の代わりに $z \sim q(z|x)$ を使うため(0)式のEncoderを使うと次の下位限界が得られる。

$\log p(x) \geq \mathbb{E}_{z \sim q(z|x)} [\sum_t \log p(x_t|z_t) + \{\log p(z_t|z_{t-1} ) - \log q(z_t|z_{t-1},\phi_t(x)) \} ]$ (1)

ここで中括弧内の(0)式の近似が一致する場合、両辺が一致することがわかる。

(b)ELBO(Evidence Lower Band Optimizer)モデル

(1)式で過去の観察 $x_{\lt t}$ に依存し、2回の $z_t,z_{t-1}$ の両方が $x_{\lt t}$ に依存すると(3)式に変形できる。

　 $\log p(x) = \sum_t \log p(x_t | x_{\lt t})$

$\log p(z_t | z_{t-1}) = \log p(z_t,z_{t-1} | x_{\lt t})$

　　　 $\log q(z_t | z_{t-1},\phi_t(x)) = \log q(z_t,z_{t-1} | x_{\lt t})$

を使って

$\log p(x_t | x_{\lt t} ) \geq \mathbb{E}_{z \sim q(z_t,z_{t-1}|x_{\lt t})} [\log p(x_t|z_t,x_{\lt t}) + \log p(z_t|z_{t-1} ) - \log q(z_t|z_{t-1},x_{\lt t})]$ (3)　　

　さらに以下の事実を使うと(4)式で示せる。

$p(x_t | z_{t-1},z_t,x_{\lt t}) = p(x_t | z_t)$ $x_{t}$ は $z_{t}$ のみ依存

$p(z_{t-1},z_t | x_{\lt t}) = p(z_{t-1} | x_{\lt t})p(z_t | z_{t-1})$ 実体遷移では $x_{\lt t}$ は無関係

　　　 $q(z_{t-1},z_t | x_{\lt t}) = q(z_t | x_{\lt t}) q(z_{t-1} | z_t,x_{\lt t})$ Bayes公式

　　　 $\log p(x_t | x_{\lt t} ) \ge \mathbb{E}_{z \sim q(z_t,z_{t-1}|x_{\lt t})} [\log p(x_t|z_t) + \log(z_t|x_{\lt t} ) + \log p(z_t | z_{t-1}) - \log q(z_t|x_{\lt t}) - \log q(z_{t-1} | z_t,x_{\lt t})]$ (4)

　ここで記憶の概念 $b_t$ をRNNで導入する。即ち $t$ 時点までの観測 $x_{\lt t}$ を使う代わりに過去の記憶を反映できるRNNモデル $b_t=f_B(b_t,x_t)$ とすると便利である。RNNを使って $p(z_t | x_t)$ を $p_B(z_t | b_t)$ とすると(4)式を変形して次のELBOの損失関数が定義できる。

　　　 $-\mathcal{L} = \mathbb{E}_{z_t,z_{t-1} \sim \psi(z,b)} [\log p(x_t|z_t) + \log p_B(z_t|b_t ) + \log p(z_t | z_{t-1}) - \log p_B(z_t|b_t) - \log q(z_{t-1} | z_t,b_{t-1},b_t)]$ 　　　(5)

　　　但し

　　　　 $\mathcal{L} = \log p(x_t | x_{\lt t})$ 負の対数尤度

　　　　 $z_t,z_{t-1} \sim \psi(z,b) = z_t \sim p_B(z_t | b_t) , z_{t-1} \sim q(z_{t-1} | z_t,b_t,b_{t-1})$

(c)TD_VAE Jumpyモデル　

　ある目的を達成する場合は道標(マイルストーン)を設定して進む場合が多い。目的を達成する強化学習でも例外でなく逐次的処理を効率化できる。されに道標を設けることにより、道標間の上位モデルと道標内の下位モデルで階層モデルを導入できる。TD_VAEでは２階層のRNN構造階層モデルを構築している(論文図8参照)。そこでTD_VAEではステップ間 $[t_1 \sim t_2]$ でのELBOモデルに変換している。

　　　 $\mathcal{L}_{t_1,t_2} = \mathbb{E}_{z_{t_1},z_{t_2}\sim \psi(z,b)} [\log p(x_{t_2} | z_{t_2}) + \log p_B(z_{t1} | b_{t_1}) + \log p(z_{t_2}|z_{t_1}) - \log p_B(z_{t_2}|b_{t_2}) - \log q(z{t_1} | z_{t_2},b_{t_1},b_{t_2}) ]$ 　　　(6)

(2) TD_VAEのアルゴリズム

アルゴリズムでは下記の(a)と(b)は同じ実体 $z_{t_1}$ を示しており、Encoder $q$ と実体の推定確率 $p$ とが同じになる様(下図ではSmoothing)に $\mathcal{KL}$ 距離最小化を損失関数に挿入している。

　　 $\mathcal{L} = \mathcal{KL}(q_S^{t_1|t_2}||p_B^{t_1}) + \log p_B^{t_2}(z_{t_2}) - \log p_T^{t_2}(z_{t_2}) - \log p_D^{t_2}(x_{t_2})$

　(6)式TD-VAEの損失関数とアルゴリズムの損失関数の対応を以下に示す。　　

(a) $\log q(z_{t_1} | z_{t_2},b_{t_1},b_{t_2}) \to \mathcal{KL}(q_S^{t_1|t_2}|| \cdot)$

(b) $\log p_B(z_{t1} | b_{t_1}) \to \mathcal {KL} (\cdot || p_B^{t_1})$

(d) $\log p_B(z_{t_2}|b_{t_2}) \to \log p_T^{t_2}(z_{t_2})$

(e) $\log (x_{t_2} | z_{t_2}) \to \log p_D^{t_2} (x_{t_2})$

　論文では上記のアルゴリズムの手順を下図の①から⑧に示している。

f:id:mabonki0725:20190203142435p:plain — TD-VAEアルゴリズム

　　以下の ①から⑧を繰返して損失関数 $\mathcal{L}$ を改善してEncoder $q(\phi)$ とDecoder $p(\theta)$ のパラメータを改善する。

　　十分改善した時、Encoderで画像を再生し強化学習より次の道標 $t_3$ に進む。

　　①RNNと観察 $x_{t_1}$ より $b_{t_1}$ を生成

　　②次の道標を選択

③RNNの $b_{t_2}$ より(d)式で次の実体を予測

　　④(d)式により実体 $z_{t_2}$ を予測

　　⑤(b)式のEncoder $q(\phi)$ と(a)式のDecoder $p(\theta)$ との $\mathcal{KL}$ を計算

　　⑥次の予測実体 $z_{t_2}$ を使ってEncoder $q(\phi)$ を計算

　　⑦予測された実体 $z_{t_2}$ をEncodeして観察 $x_{t_2}$ を生成

　　⑧損失関数を最小化するため $p(\phi)$ と $q(\theta)$ のパラメータを最適化

(3) 実験

　DeepMind-Labの3D迷路では２種類の実験を行っている。

(a) [1 ～40]をランダムの時点についてRNNで記憶 $b$ をさせて、適当に３つの記憶から実体 $z$ の推移５回予測してDecodeした画像を指名している。下図では異なった場面で連続した画像が得られていることが分る。

f:id:mabonki0725:20190203173036p:plain

　(b) TD-VAEで迷路探索を行った結果では、下図の４種類の画面遷移では全て前に進み通路に向かおうとしている事がわかる。これは実体を予測してDecodeした画像により強化学習を行っている事を示している。　

f:id:mabonki0725:20190203173639p:plain

(4) 感想

　(a)残念な事にTD-VAEでは(6)式の損失関数のみ定義されているだけで、次の分布の、詳細な記述が無い。KingmaのVAE[2]では $z$ は混合正規分布でしか収束が保障されないはずである（最近この拡張論文を知った[3]）。報告されているアルゴリズムで正しい収束ができているか不安である。ELBOの収束の問題は[4]に詳しい。

　 $p(z)$ 　実態の分布 (普通は正規分布)

$q(z|x)$ 変分Encoder (普通は混合正規分布　数千次元)

　 $p_B(b,x)$ 記憶呼出

(b)TD-VAEではJumpy（飛ばし）な時系列を導入して、効率的な実体の予測をしているが実験ではランダムに時系列をサンプリングしているに過ぎない。論文の4.3節では価値関数 $Q_{t_1}$ や $V_{t_1}$ に対して $t_2$ を探索すべきとある。また一方で記憶 $p_B(z_{t_1}|b_{t_1})$ と $p_B(z_{t_2}|b_{t_2})$ との訓練で得られるとあるが、具体的な方法が述べられていない。