[DL Hacks] Deterministic Variational Inference for RobustBayesian Neural Networks (ICLR2019)

DEEP LEARNING JP Deterministic Variational Inference for Robust Bayesian Neural Networks (ICLR2019) [DL Papers] Makoto Kawano, Matsuo Lab. http://deeplearning.jp/

• http://deeplearning.jp/

書誌情報 • 著者   Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard E. Turner, Jose Miguel Hernandez-Lobato, Alexander L. Gaunt MSR Cambridge でのインターンの成果 (プリンストン大学の学生) • ICLR2019 oral presentation: 7, 7, 7 • 選定理由：ICLR & ベイズ勉強のため • TL;DR：ベイズニューラルネットワークの学習を安定化・高速化 した • 2

アウトライン 論文について ベイズニューラルネットワーク 変分推論 提案手法 1：決定論的変分近似 近似における仮説 提案手法 2：経験ベイズによる近似分布のパラメータ推定 実装について 再現実装レポジトリ 実装の方針 レイヤの実装 3

ベイズニューラルネットワーク • 一般的なニューラルネットワークの場合   重み (W /b) は，最尤推定 ⇒ 値が一意に決まっている 同じ入力なら出力も毎回同じ • ベイズニューラルネットワークの場合  重みに事前分布 p(W ) を設定  予測時に不確実性を捉える  モデルに対して正則化する } [MacKay, 1992, Graves, 2011] ⇒ 意思決定や小さなデータセットでの過剰適合を減らす 5

ベイズニューラルネットワークの欠点 • 事後分布が解析的に求まらない  モンテカルロ法や変分推論がよく使われる  計算資源による制約で分散に苦しめられる (サンプリング数) ⇒ アプローチ 1：変分推論の決定論的近似による解決 • 事前分布の選択に性能が大きく左右される  ニューラルネットワークが表現する「関数」の事前知識があっても， 重みの事前分布に反映するのは難しい ⇒ アプローチ 2：経験ベイズによる解決 6

変分推論導入の動機 • ベイズニューラルネットワークの問題設定 {   • 重み W (b も含む) でパラメータ化されたモデル M と • データセット D = (x, y) が与えられたときの事後分布 p(W |D) を求めたい (推論したい) 最後，予測分布 p(y |D, x ) = ′ ′  p(y ′ |x′ , W )p(W |D)dw を求めたい • p(W |D) がどんな形をしているかわからない  確率的順伝播法 [Hernández-Lobato and Adams, 2015]  MCMC：p(W |D) = p(D|W )p(W ) の右辺からサンプリングして推論  変分推論：知っている分布族 q(D; θ) で近似する 本研究はこれ 8

変分推論 • 知りたい事後分布 p(W |D) と近似分布 q(W ; θ) を似せる  KL ダイバージェンスを最小化 ⇒ELBO の最大化  q(W ; θ) dw q(W ; θ) log p(W |D)    p(D|W )p(W ) − log p(D) dw = − q(W ; θ) log q(W ; θ) DKL [q(W ; θ)∥p(W |D)] = ≥ EW ∼q [log p(D|W )] + DKL [q(W ; θ)∥p(W )]  VI でも解析的に求まらない - S 個のサンプリングによる MC 近似で求める  事前分布と近似分布をうまく選べば解析的に求まる - 正規分布族とか 9

MCVI(いわゆる reparameterization trick) • 推定量の分散が大きすぎると，学習が安定しない [Miller et al., 2017] S 1 EW ∼q [log p(y|W, x)] ≈ log p(y|W (s) , x), w(s) ∼ q(W ; θ) S s=1 ⇒ 事後分布を決定論的に推定することで学習が安定する 10

提案手法 1：決定論的変分近似 • EW ∼q [log p(y|W, x)] をニューラルネットワークで求める 1. 活性化信号を分布のまま，パラメータ化された層で順伝播させ， 2. パラメータ化されていない対数尤度で評価する 12

活性化信号 (モーメント) 伝播 • 層の区切り方を変える 一般的な区切り方 hl = al−1 W l + bl al = f (hl ) 本研究の区切り方 hl = f (al−1 ) a l = h l W l + bl • 中心極限定理  N 個の確率変数の和 ⇒ N → ∞ のとき，正規分布になる • a も確率変数 h の線型結合 ⇒ 要素 ai も正規分布に従うはず  非線形関数 f で h が複雑な分布になっても関係ない 13

近似精度について • 活性化信号の正規分布パラメータを近似で導出する ⟨ai ⟩ = ⟨hj ⟩⟨Wji ⟩ + ⟨bi ⟩ Cov(ai , ak ) = ⟨hj hl ⟩Cov(Wji , Wlk ) + ⟨Wji ⟩Cov(Wji , Wlk )⟨Wlk ⟩ + Cov(bi , bk )  それぞれの導出を頑張ると，al を al−1 で表せる ⟨hj hl ⟩ = Sjl′ {A(µ′j , µ′l , ρ′jl ) + exp[−Q(µ′j , µ′l , ρ′jl )]} A(µ1 , µ2 , ρ) Heaviside ReLU Φ(µ1 )Φ(µ2 ) SR(µ1 )SR(µ2 ) +ρΦ(µ1 )Φ(µ2 ) Q(µ1 , µ2 , ρ)    gh ρ  ρ 2ρ 2 2 − log 2π + 2g ρ̄ µ1 + µ2 − 1+ρ̄ µ1 µ2 + O(µ4 ) h   gr   ρ arcsin ρ−ρ 2 2 − log 2π + 2g (1+ρ̄) (µ1 + µ2 ) − µ1 µ2 + O(µ4 ) g ρ r r 14

近似における仮説 • 最終的な提案  正規分で近似された al の分布を al−1 の分布で表現 - 平均・分散共分散行列の両方を導出 ⇒ サンプリングが不要 • 3 つの仮説に基づいて近似している 1. 隠れ層のユニット間に (一般的な CLT が持つ)i.i.d はいらない 2. Q 関数は，二次項で打ち切っても十分に表現力がある 3. 層間には弱い相関しかない - 近似分布による独立変数を用いて表現しても問題ない 16

実験による仮説検証 • 近似分布からの MCMC と比較  ほぼ同じ ⇒ 仮説は成立している 17

パラメータ化されていない対数尤度 • 最後に出力された分布における対数尤度について  W についてではなく，aL についての対数尤度にする EW ∼q [log p(y|x, W )] = EW ∼q [log p(y|aL )] • EW ∼q [log p(y|aL )] は，解くタスクによって異なる  今回は回帰問題に取り組む  y は平均 m(x; W ) と不等分散 log σ 2 (x) = l(x; W ) に従う - ニューラルネットワークの出力は 2 次元：aL = (m, l) ⇒ p(y|aL ) = N (y; m, el ) 18

回帰問題時の対数尤度 • 具体的な対数尤度   2 Σ 1 + (⟨m⟩ − Σ − y) mm mℓ EaL ∼q̃(aL ) [log p(y|aL )] = − log 2π + ⟨ℓ⟩ + 2 e⟨ℓ⟩−Σℓℓ /2  等分散にしたい場合：⟨ℓ⟩ = Σℓℓ = Σmℓ = 0 • 事後予測分布 p(y) ≈  p(y|aL )˜(q)(aL )daL ≈ N (y|⟨m⟩, Σmm + e⟨ℓ⟩−Σℓℓ /2 ) 19

重みの初期値の重要性 • ELBO 第 2 項の KL ダイバージェンスに注目  d 次元の事前分布 p(W ) = N (µp , Σp ) と近似分布 q(W ) = N (µq , Σq )  |Σp | 1 ⊤ −1 log − d + Tr(Σ−1 Σ ) + (µ − µ ) Σp (µp − µq ) DKL [q∥p] = q p q p 2 |Σq |    よくある事前設定は，分布を左右対称にするため，µp = 0 と 活性化関数に合わせた分散（ [Glorot and Bengio, 2010] や [He et al., 2015] の初期値） ⇒ 分散 Σp の選択はニューラルネットワークの性能に大きく影響 21

提案手法 2：経験ベイズによる近似分布の パラメータ推定 • 事前分布を階層ベイズにする s ∼ p(s); W ∼ p(w|s)  手動で重みの事前分布のパラメータを決めるのは大変  正規分布の分散に対して共役事前分布である逆ガンマ分布を設定する sλ ∼ Inv − Gamma(α, β); Wiλ ∼ N (0, sλ ) 22

提案手法 2：経験ベイズによる近似分布の パラメータ推定 • 経験ベイズアプローチ (第二種 MAP 推定) の適用  事前分布の超パラメータ sλ をデータから点推定 ELBO = EW ∼q [log p(u|hL (W ))] − {DKL [q(W ; θ)∥p(w|s∗ )p(s∗ )]} ⇒ sλ∗ = arg min{DKL [q(W ; θ)∥p(w|s∗ )] − log p(s∗ ) λ λ λ ⊤ sλ Tr[Σ + µ (µ q q q ) ] + 2β λ s∗ = Ωλ + 2α + 2 ⇒ sλ を変分パラメータ µq と Σq で表すことが可能 23

再現実装レポジトリ • Github レポジトリ 25

実装の方針 • Pytorch で実装したい  一応著者らの実装コードはある  できる限り Pytorch の API に準拠させたい [Tran et al., 2018] - ただし，Tensorflow で実装 - nn.Sequential や model.parameters() をそのまま使いたい - 損失関数もモデルの外側で操作可能にしたい - 確率分布周りは torch.distributions.* で扱う 27

レイヤの実装 • nn.Module を継承する • 重み：単純に nn.Parameters にしてはダメ  事前分布と近似分布，変分パラメータを保持させる必要がある - 事前分布：基本的には distributions.Normal(0, 1) - 事後分布：事前分布のようにパラメータを渡さず，分布の形のみ - 変分パラメータ： nn.Parameters() にする必要がある • 順伝播 forward() ：ちょっと癖あり  活性化関数に合わせて，実装する必要がある  基本的な機能を VariationalLayer に集約，順伝播は継承させて実装 29

最終的なレイヤの実装：パラメータ編 • 基本的には nn.Linear と同じ書き方  make_weight_matrix と make_bias_vector がキモとなる 30

最終的なレイヤの実装：パラメータ編 31

最終的なレイヤの実装：順伝播編 • 活性化関数によって順伝播が異なるため，活性化関数の数だけレ イヤが増える  VariationalLinearCertainActivations  VariationalLinearReLU • 表に基づいて，実装をしていくのみ 32

まとめ • ベイズニューラルネットワークの学習を安定化させた  不確実性を含んだ活性化信号（分布）を伝播する決定論的手法 - 変分推論のサンプリングによる非効率性の解決を目指す - ReLU/Heaviside 活性化関数をもつニューラルネットワーク - 分布：CLT に基づき，多変量正規分布 • 回帰問題データセットで性能を検証  提案手法が精度上回りつつ，演算処理時間を短縮  不等分散を扱えるため，不確実性を捉えられる 33

参考文献 I Glorot, X. and Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249–256. Graves, A. (2011). Practical variational inference for neural networks. In Advances in neural information processing systems, pages 2348–2356. He, K., Zhang, X., Ren, S., and Sun, J. (2015). Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026–1034. Hernández-Lobato, J. M. and Adams, R. (2015). Probabilistic backpropagation for scalable learning of bayesian neural networks. In International Conference on Machine Learning, pages 1861–1869. 34

参考文献 II MacKay, D. J. (1992). A practical bayesian framework for backpropagation networks. Neural computation, 4(3):448–472. Miller, A., Foti, N., D’Amour, A., and Adams, R. P. (2017). Reducing reparameterization gradient variance. In Advances in Neural Information Processing Systems, pages 3708–3718. Tran, D., Mike, D., van der Wilk, M., and Hafner, D. (2018). Bayesian layers: A module for neural network uncertainty. arXiv preprint arXiv:1812.03973. 35