Learning Latent Space Energy-Based Prior Model https://arxiv.org/abs/2006.08205 この論文の末尾に付いているPyTorch codeを理解できるようになるための解説です。

1. Learning Latent Space Energy-Based Prior Model [2] の解説
正田 備也 @ Rikkyo University
2022 年 7 月 15 日
1 Model
xi は画像やテキストなどのインスタンス、zi ∈ RM
は潜在確率変数で、 xi の低次元空間での表現とする。
以下のような確率モデルを考える。
pθ(xi, zi) = pα(zi)pβ(xi|zi) (1)
α は事前分布 pα(zi) のパラメータ、β は観測データ xi を生成する確率分布 pβ(xi|zi) のパラメータである。θ
は (α, β) のことで、α と β をまとめて θ と書いているだけである。
この論文では、事前分布として isotropic Gaussian などの単純な分布を使うのではなく、EBM (energy-based
model) を使って事前分布の “表現力” を上げようとしている。事前分布は以下のように定式化される。
pα(z) =
1
Z(α)
exp(fα(z))p0(z) (2)
p0(z) は isotropic Gaussian など単純な分布で、それを exp(fα(z)) によって “exponential tilting” している。
Z(α) は規格化定数なので Z(α) =
∫
exp(fα(z))p0(z)dz を満たす。
式 (2) の z の関数 fα(z) は MLP として実装され、α はこの MLP の重みおよびバイアスを表す。
観測データを生成する確率分布 pβ(x|z) は、入力を z とするニューラルネットワーク gβ(z) を使って実装さ
れる。モデリングしたい観測データの種類に応じて、転置畳み込みや RNN などを使う。β はこのニューラル
ネットワークのパラメータを表す。そして、この gβ(z) の出力を使って、インスタンス x は以下のようにモデ
リングされる。
x = gβ(z) + ϵ (3)
ϵ は、例えば画像の場合なら、画素数を D として ϵ ∼ N(0, σ2
ID) とモデリングする。すると、観測データ x
を x ∼ N(gβ(z), σ2
ID) とモデリングしていることになる。
周辺分布は
pθ(xi) =
∫
pα(zi)pβ(xi|zi)dzi (4)
事後分布は
pθ(xi|zi) =
pθ(xi, zi)
pθ(xi)
=
pα(zi)pβ(xi|zi)
pθ(xi)
(5)
1

2. 2 Maximum Likelihood Estimation Learning Gradient
データセットの対数周辺尤度は、以下のように表される。
L(θ) =
N
∑
i=1
log pθ(xi) (6)
この対数周辺尤度を最大化することによって、パラメータ θ = (α, β) を推定する。この最大化に必要な勾配は
∇θ log pθ(x) だが、これは、以下のように書き直すことができる。
∇θ log pθ(x) = Epθ(z|x)[∇θ log pθ(x, z)] (7)
このように書き直せることを示すために、以下の等式を使う。x ∼ pθ(x) を任意の確率分布とすると、
Epθ(x)[∇θ log pθ(x)] = 0 (8)
この等式が成り立つことは、次のように証明できる。
Epθ(x)[∇θ log pθ(x)] =
∫
(
∇θ log pθ(x)
)
pθ(x)dx （期待値の定義）
=
∫ (∇θpθ(x)
pθ(x)
)
pθ(x)dx （対数関数の微分）
=
∫
∇θpθ(x)dx （分母と分子でキャンセル）
= ∇θ
∫
pθ(x)dx （積分と微分の順序交換）
= ∇θ1 （密度関数の積分は 1 に等しい）
= 0 （定数の微分は 0） (9)
では、等式 (7) を証明する。
Epθ(z|x)[∇θ log pθ(x, z)] = Epθ(z|x)[∇θ log pθ(z|x)pθ(x)]
= Epθ(z|x)[∇θ log pθ(z|x)] + Epθ(z|x)[∇θ log pθ(x)]
= 0 + Epθ(z|x)[∇θ log pθ(x)] （式 (8) より）
= ∇θ log pθ(x) （pθ(x) は z に非依存）
(10)
等式 (7) は、以下のように証明することもできる。まず、Jensen の不等式より
log pθ(x) = log
∫
pθ(x, z)dz = log
∫
q(z)
pθ(x, z)
q(z)
dz ≥
∫
q(z) log
pθ(x, z)
q(z)
dz (11)
この不等式は、q(z) = pθ(z|x) のとき等号で成立する（EM アルゴリズムに関連する議論）
。つまり、
log pθ(x) =
∫
pθ(z|x) log
pθ(x, z)
pθ(z|x)
dz = Epθ(z|x)[log pθ(x, z)] − Epθ(z|x)[log pθ(z|x)] (12)
よって、
∇θ log pθ(x) = ∇θEpθ(z|x)[log pθ(x, z)] − ∇θEpθ(z|x)[log pθ(z|x)]
= Epθ(z|x)[∇θ log pθ(x, z)] − Epθ(z|x)[∇θ log pθ(z|x)] （積分と微分の順序交換）
= Epθ(z|x)[∇θ log pθ(x, z)] − 0 （式 (8) より）
= Epθ(z|x)[∇θ log pθ(x, z)] (13)
2

3. 3 Maximum Likelihood Estimation Learning Gradient for α
θ は、α と β とをまとめて書いただけの記号だった。
まず α に関する勾配を考える。
∇α log pθ(x) = Epθ(z|x)[∇α log pθ(x, z)]
= Epθ(z|x)[∇α log pα(z)] + Epθ(z|x)[∇α log pβ(x|z)]
= Epθ(z|x)[∇α log pα(z)] （pβ(x|z) は α に非依存） (14)
ここで、∇α log pα(z) は、次のように書き換えられる。
∇α log pα(z) = ∇αfα(z) − Epα(z)[∇αfα(z)] (15)
というのも、まず
∇α log pα(z) = ∇α log
(
1
Z(α)
exp(fα(z))p0(z)
)
= ∇αfα(z) + ∇α log p0(z) − ∇α log Z(α)
= ∇αfα(z) − ∇α log Z(α) （p0(z) は α に非依存） (16)
と書き換えられ、さらに ∇α log Z(α) について
0 = Epα(z)[∇α log pα(z)] （式 (8) より）
= Epα(z)[∇αfα(z) − ∇α log Z(α)] （式 (16) より）
= Epα(z)[∇αfα(z)] − ∇α log Z(α) （Z(α) は z に非依存） (17)
より
∇α log Z(α) = Epα(z)[∇αfα(z)] (18)
という等式が成り立つことが言えるので、式 (16) と組み合わせて、式 (15) と書き換えられることが分かる。
したがって、式 (14) は、次のように書き換えられる。
∇α log pθ(x) = Epθ(z|x)[∇α log pα(z)]
= Epθ(z|x)[∇αfα(z) − Epα(z)[∇αfα(z)]]
= Epθ(z|x)[∇αfα(z)] − Epθ(z|x)
[
Epα(z)[∇αfα(z)]
]
= Epθ(z|x)[∇αfα(z)] − Epα(z)[∇αfα(z)] (19)
式 (19) の期待値は、pθ(z|x) および pα(z) からサンプルを得ることで、近似計算する。このサンプリングには、
ランジュバン・モンテカルロ [1] を、次のように利用する。s をステップ幅とし、例えば K = 60 などとして
z0 ∼ π0(z)
zk+1 = zk + s∇zk
log pα(zk) +
√
2sϵk+1 , ϵk+1 ∼ N(0, IM ) for k = 0, . . . , K − 1 (20)
π0(z) は initial distribution で、標準正規分布などを使う。ここで
∇z log pα(z) = ∇z log
(
1
Z(α)
exp(fα(z))p0(z)
)
= ∇zfα(z) + ∇z log p0(z) − ∇z log Z(α)
= ∇zfα(z) + ∇z log p0(z) (21)
例えば p0(z) =
∏M
m=1
1
√
2πσ0
exp(−
z2
m
2σ2
0
) とする場合は ∇z log pα(z) = ∇zfα(z) − z/σ2
0 となる。σ0 ≫ 0 と
考えるなら、∇z log pα(z) ≈ ∇zfα(z) としてよいだろう。
3

4. 4 Maximum Likelihood Estimation Learning Gradient for β
次に β に関する勾配を考える。
∇β log pθ(x) = Epθ(z|x)[∇β log pθ(x, z)]
= Epθ(z|x)[∇β log pα(z)] + Epθ(z|x)[∇β log pβ(x|z)]
= Epθ(z|x)[∇β log pβ(x|z)] （pα(z) は β に非依存） (22)
式 (22) の ∇β log pβ(x|z) は自動微分で計算する。例えば、式 (3) のように観測データをモデリングする場合
log pβ(xi|zi) = −
1
2σ2
∥xi − gβ(z)∥2
2 (23)
となる。また、式 (22) の期待値は、pθ(z|x) からサンプルを得ることで、近似計算する。このサンプリングに
は、ランジュバン・モンテカルロ [1] を次のように利用する。s をステップ幅とし、例えば K = 40 などとして
z0 ∼ π0(z)
zk+1 = zk − s∇zk
log pθ(zk|x) +
√
2sϵk+1 , ϵk+1 ∼ N(0, IM ) for k = 0, . . . , K − 1 (24)
π0(z) は initial distribution で、標準正規分布などを使う。ここで
∇z log pθ(z|x) = ∇z log
(
pβ(x|z)pα(z)
pθ(x)
)
= ∇z log pβ(x|z) + ∇z log pα(z)
= ∇z log pβ(x|z) + ∇zfα(z) + ∇z log p0(z) (25)
式 (20) の事前分布からのサンプリングとは ∇z log pβ(x|z) の部分が異なるだけである。
5 PyTorch code の誤り
論文のコードで、ランジュバン・サンプリングを行う関数に 1.0 / z.data という部分が見られるが、これ
はおそらく、∇z log p0(z) を ∇z log z と勘違いしたためだろう。
参考文献
[1] Langevin monte carlo 法には棄却が必要か. https://ktrmnm.github.io/blog/2016/08/25/201608-lmc/.
Accessed: 2022-07-14.
[2] Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based
prior model. In Proceedings of the 34th International Conference on Neural Information Processing Systems,
NIPS’20, Red Hook, NY, USA, 2020. Curran Associates Inc.
4