Title: LAGGING INFERENCE NETWORKS AND POSTERIOR COLLAPSE IN VARIATIONAL AUTOENCODERS

Abstract: The variational autoencoder (VAE) is a popular combination of deep latent variable model and accompanying variational learning technique. By using a neural inference network to approximate the model's posterior on latent variables, VAEs efficiently parameterize a lower bound on marginal data likelihood that can be optimized directly via gradient methods. In practice, however, VAE training often results in a degenerate local optimum known as "posterior collapse" where the model learns to ignore the latent variable and the approximate posterior mimics the prior. In this paper, we investigate posterior collapse from the perspective of training dynamics. We find that during the initial stages of training the inference network fails to approximate the model's true posterior, which is a moving target. As a result, the model is encouraged to ignore the latent encoding and posterior collapse occurs. Based on this observation, we propose an extremely simple modification to VAE training to reduce inference lag: depending on the model's current mutual information between latent variable and observation, we aggressively optimize the inference network before performing each model update. Despite introducing neither new model components nor significant complexity over basic VAE, our approach is able to avoid the problem of collapse that has plagued a large amount of previous work. Empirically, our approach outperforms strong autoregressive baselines on text and image benchmarks in terms of held-out likelihood, and is competitive with more complex techniques for avoiding collapse while being substantially faster. 1 

1. INTRODUCTION:

Variational autoencoders (VAEs) (Kingma & Welling, 2014) represent a popular combination of a deep latent variable model (shown in Figure 1(a) ) and an accompanying variational learning technique. The generative model in VAE defines a marginal distribution on observations, x ∈ X , as:

p θ (x) = p θ (x|z)p(z)dz.

(1)

The model's generator defines p θ (x|z) and is typically parameterized as a complex neural network.

Standard training involves optimizing an evidence lower bound (ELBO) on the intractable marginal data likelihood (Eq.1), where an auxiliary variational distribution q φ (z|x) is introduced to approximate the model posterior distribution p θ (z|x). VAEs make this learning procedure highly scalable to large datasets by sharing parameters in the inference network to amortize inferential cost. This amortized approach contrasts with traditional variational techniques that have separate local variational parameters for every data point (Blei et al., 2003; Hoffman et al., 2013) .

While successful on some datasets, prior work has found that VAE training often suffers from "posterior collapse", in which the model ignores the latent variable z (Bowman et al., 2016; Kingma et al., 2016; Chen et al., 2017) . This phenomenon is more common when the generator p θ (x|z) is parametrized with a strong autoregressive neural network, for example, an LSTM (Hochreiter & Schmidhuber, 1997 ) on text or a PixelCNN (van den Oord et al., 2016) on images. Posterior collapse is especially evident when modeling discrete data, which hinders the usage of VAEs in p(z)

< l a t e x i t s h a 1 _ b a s e 6 4 = " B M l c n x o 7 s Q R J A 5 G g 0 p I 6 o I N x q O E = " > A A A C E n i c Z V D L S g M x F M 3 U V 6 2 v q k s 3 g 0 W o U M p U B H V X d O O y g m M L 7 V A y a a Y N T S Z j c k e o Q z 9 D 3 O p 3 u B K 3 / o C f 4 R + Y m c 7 C t h f C P T n 3 3 H B y / I g z D Y 7 z Y x V W V t f W N 4 q b p a 3 t n d 2 9 8 v 7 B g 5 a x I t Q l k k v V 8 b G m n I X U B Q a c d i J F s f A 5 b f v j m 3 T e f q J K M x n e w y S i n s D D k A W M Y D C U F 1 V 7 A s P I D 5 L n 6 W m / X H H q T l b 2 M m j k o I L y a v X L v 7 2 B J L G g I R C O t e 4 2 n A i 8 B C t g h N N p q R d r G m E y x k P a N T D E g m o v y U x P 7 R P D D O x A K n N C s D P 2 / 0 a C h U 6 9 G W X a 9 N w s Z U B K r m t G B S O R t v S Z 7 K 4 n w q / 5 o p a K l A 7 0 g h E I L r 2 E h V E M N C Q z H 0 H M b Z B 2 G p A 9 Y I o S 4 B M D M F H M f M U m I 6 w w A R N j y W T U W E x k G b h n 9 a u 6 c 3 d e a V 7 n Y R X R E T p G V d R A F 6 i J b l E L u Y i g R / S K 3 t C 7 9 W J 9 W J / W 1 0 x a s P K d Q z R X 1 v c f 9 s O f C w = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " B M l c n x o 7 s Q R J A 5 G g 0 p I 6 o I N x q O E = " > A A A C E n i c Z V D L S g M x F M 3 U V 6 2 v q k s 3 g 0 W o U M p U B H V X d O O y g m M L 7 V A y a a Y N T S Z j c k e o Q z 9 D 3 O p 3 u B K 3 / o C f 4 R + Y m c 7 C t h f C P T n 3 3 H B y / I g z D Y 7 z Y x V W V t f W N 4 q b p a 3 t n d 2 9 8 v 7 B g 5 a x I t Q l k k v V 8 b G m n I X U B Q a c d i J F s f A 5 b f v j m 3 T e f q J K M x n e w y S i n s D D k A W M Y D C U F 1 V 7 A s P I D 5 L n 6 W m / X H H q T l b 2 M m j k o I L y a v X L v 7 2 B J L G g I R C O t e 4 2 n A i 8 B C t g h N N p q R d r G m E y x k P a N T D E g m o v y U x P 7 R P D D O x A K n N C s D P 2 / 0 a C h U 6 9 G W X a 9 N w s Z U B K r m t G B S O R t v S Z 7 K 4 n w q / 5 o p a K l A 7 0 g h E I L r 2 E h V E M N C Q z H 0 H M b Z B 2 G p A 9 Y I o S 4 B M D M F H M f M U m I 6 w w A R N j y W T U W E x k G b h n 9 a u 6 c 3 d e a V 7 n Y R X R E T p G V d R A F 6 i J b l E L u Y i g R / S K 3 t C 7 9 W J 9 W J / W 1 0 x a s P K d Q z R X 1 v c f 9 s O f C w = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " B M l c n x o 7 s Q R J A 5 G g 0 p I 6 o I N x q O E = " > A A A C E n i c Z V D L S g M x F M 3 U V 6 2 v q k s 3 g 0 W o U M p U B H V X d O O y g m M L 7 V A y a a Y N T S Z j c k e o Q z 9 D 3 O p 3 u B K 3 / o C f 4 R + Y m c 7 C t h f C P T n 3 3 H B y / I g z D Y 7 z Y x V W V t f W N 4 q b p a 3 t n d 2 9 8 v 7 B g 5 a x I t Q l k k v V 8 b G m n I X U B Q a c d i J F s f A 5 b f v j m 3 T e f q J K M x n e w y S i n s D D k A W M Y D C U F 1 V 7 A s P I D 5 L n 6 W m / X H H q T l b 2 M m j k o I L y a v X L v 7 2 B J L G g I R C O t e 4 2 n A i 8 B C t g h N N p q R d r G m E y x k P a N T D E g m o v y U x P 7 R P D D O x A K n N C s D P 2 / 0 a C h U 6 9 G W X a 9 N w s Z U B K r m t G B S O R t v S Z 7 K 4 n w q / 5 o p a K l A 7 0 g h E I L r 2 E h V E M N C Q z H 0 H M b Z B 2 G p A 9 Y I o S 4 B M D M F H M f M U m I 6 w w A R N j y W T U W E x k G b h n 9 a u 6 c 3 d e a V 7 n Y R X R E T p G V d R A F 6 i J b l E L u Y i g R / S K 3 t C 7 9 W J 9 W J / W 1 0 x a s P K d Q z R X 1 v c f 9 s O f C w = = < / l a t e x i t > z < l a t e x i t s h a 1 _ b a s e 6 4 = " p L n q A s m T F W J Q S V A 6 9 c 6 x b 0 3 w v 3 M = " > A A A C D 3 i c Z V D L T g I x F O 3 g C / G F u n Q z k Z i 4 I G Q w J u q O 6 M Y l J o 4 Q g Z B O 6 U B D H 5 P 2 j g l O + A n j V r / D l X H r J / g Z / o G d Y R Y C N 2 n u 6 b n n N q c n i D g z 4 H k / T m F l d W 1 9 o 7 h Z 2 t r e 2 d 0 r 7 x 8 8 G B V r Q n 2 i u N L t A B v K m a Q + M O C 0 H W m K R c B p K x j f p P P W E 9 W G K X k P k 4 j 2 B B 5 K F j K C w V K P X Y F h F I T J 8 7 R f r n g 1 L y t 3 G d R z U E F 5 N f v l 3 + 5 A k V h Q C Y R j Y z p 1 L 4 J e g j U w w u m 0 1 I 0 N j T A Z 4 y H t W C i x o K a X Z I 6 n 7 o l l B m 6 o t D 0 S 3 I z 9 v 5 F g Y V J r V p k 2 M z d L G V C K m 6 p V w U i k L X 0 m u 5 u J C K q B q K Y i b U K z Y A T C y 1 7 C Z B Q D l W T m I 4 y 5 C 8 p N 0 3 E H T F M C f G I B J p r Z r 7 h k h D U m Y D M s 2 Y z q i 4 k s A / + s d l X z 7 s 4 r j e s 8 r C I 6 Q s f o F N X R B W q g W 9 R E P i J I o l f 0 h t 6 d F + f D + X S + Z t K C k + 8 c o r l y v v 8 A Q R i e L A = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " p L n q A s m T F W J Q S V A 6 9 c 6 x b 0 3 w v 3 M = " > A A A C D 3 i c Z V D L T g I x F O 3 g C / G F u n Q z k Z i 4 I G Q w J u q O 6 M Y l J o 4 Q g Z B O 6 U B D H 5 P 2 j g l O + A n j V r / D l X H r J / g Z / o G d Y R Y C N 2 n u 6 b n n N q c n i D g z 4 H k / T m F l d W 1 9 o 7 h Z 2 t r e 2 d 0 r 7 x 8 8 G B V r Q n 2 i u N L t A B v K m a Q + M O C 0 H W m K R c B p K x j f p P P W E 9 W G K X k P k 4 j 2 B B 5 K F j K C w V K P X Y F h F I T J 8 7 R f r n g 1 L y t 3 G d R z U E F 5 N f v l 3 + 5 A k V h Q C Y R j Y z p 1 L 4 J e g j U w w u m 0 1 I 0 N j T A Z 4 y H t W C i x o K a X Z I 6 n 7 o l l B m 6 o t D 0 S 3 I z 9 v 5 F g Y

V J r V p k 2 M z d L G V C K m 6 p V w U i k L X 0 m u 5 u J C K q B q K Y i b U K z Y A T C y 1 7 C Z B Q D l W T m I 4 y 5 C 8 p N 0 3 E H T F M C f G I B J p r Z r 7 h k h D U m Y D M s 2 Y z q i 4 k s A / + s d l X z 7 s 4 r j e s 8 r C I 6 Q s f o F N X R B W q g W 9 R E P i J I o l f 0 h t 6 d F + f D + X S + Z t K C k + 8 c o r l y v v 8 A Q R i e L A = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " p L n q A s m T F W J Q S V A 6 9 c 6 x b 0 3 w v 3 M = " > A A A C D 3 i c Z V D L T g I x F O 3 g C / G F u n Q z k Z i 4 I G Q w J u q O 6 M Y l J o 4 Q g Z B O 6 U B D H 5 P 2 j g l O + A n j V r / D l X H r J / g Z / o G d Y R Y C N 2 n u 6 b n n N q c n i D g z 4 H k / T m F l d W 1 9 o 7 h Z 2 t r e 2 d 0 r 7 x 8 8 G B V r Q n 2 i u N L t A B v K m a Q + M O C 0 H W m K R c B p K x j f p P P W E 9 W G K X k P k 4 j 2 B B 5 K F j K C w V K P X Y F h F I T J 8 7 R f r n g 1 L y t 3 G d R z U E F 5 N f v l 3 + 5 A k V h Q C Y R j Y z p 1 L 4 J e g j U w w u m 0 1 I 0 N j T A Z 4 y H t W C i x o K a X Z I 6 n 7 o l l B m 6 o t D 0 S 3 I z 9 v 5 F g Y V J r V p k 2 M z d L G V C K m 6 p V w U i k L X 0 m u 5 u J C K q B q K Y i b U K z Y A T C y 1 7 C Z B Q D l W T m I 4 y 5 C 8 p N 0 3 E H T F M C f G I B J p r Z r 7 h k h D U m Y D M s 2 Y z q i 4 k s A / + s d l X z 7 s 4 r j e s 8 r C I 6 Q s f o F N X R B W q g W 9 R E P i J I o l f 0 h t 6 d F + f D + X S + Z t K C k + 8 c o r l y v v 8 A Q R i e L A = = < / l a t e x i t > x < l a t e x i t s h a 1 _ b a s e 6 4 = " g X l W 4 o f y 6 W T d D o b d C q 1 l E t 4 6 O Q 0 = " > A A A C D 3 i c Z V D L T g I x F O 3 g C / G F u n Q z k Z i 4 I G T G m K g 7 o h u X m D h C h A n p l A 4 0 9 D F p 7 x j J h J 8 w b v U 7 X B m 3 f o K f 4 R / Y A R Y C N 2 n u 6 b n n N q c n S j g z 4 H k / T m F l d W 1 9 o 7 h Z 2 t r e 2 d 0 r 7 x 8 8 G J V q Q g O i u N K t C B v K m a Q B M O C 0 l W i K R c R p M x r e 5 P P m E 9 W G K X k P o 4 S G A v c l i x n B Y K n H j s A w i O L s e d w t V 7 y a N y l 3 G f g z U E G z a n T L v 5 2 e I q m g E g j H x r R 9 L 4 E w w x o Y 4 X R c 6 q S G J p g M c Z + 2 L Z R Y U B N m E 8 d j 9 8 Q y P T d W 2 h 4 J 7 o T 9 v 5 F h Y X J r V p k 3 M z f L G V C K m 6 p V w U D k L X 9 m c j c j E V U j U c 1 F 2 s R m w Q j E l 2 H G Z J I C l W T q I 0 6 5 C 8 r N 0 3 F 7 T F M C f G Q B J p r Z r 7 h k g D U m Y D M s 2 Y z 8 x U S W Q X B W u 6 p 5 d + e V + v U s r C I 6 Q s f o F P n o A t X R L W q g A B E k 0 S t 6 Q + / O i / P h f D p f U 2 n B m e 0 c o r l y v v 8 A P d K e K g = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " g X l W 4 o f y 6 W T d D o b d C q 1 l E t 4 6 O Q 0 = " > A A A C D 3 i c Z V D L T g I x F O 3 g C / G F u n Q z k Z i 4 I G T G m K g 7 o h u X m D h C h A n p l A 4 0 9 D F p 7 x j J h J 8 w b v U 7 X B m 3 f o K f 4 R / Y A R Y C N 2 n u 6 b n n N q c n S j g z 4 H k / T m F l d W 1 9 o 7 h Z 2 t r e 2 d 0 r 7 x 8 8 G J V q Q g O i u N K t C B v K m a Q B M O C 0 l W i K R c R p M x r e 5 P P m E 9 W G K X k P o 4 S G A v c l i x n B Y K n H j s A w i O L s e d w t V 7 y a N y l 3 G f g z U E G z a n T L v 5 2 e I q m g E g j H x r R 9 L 4 E w w x o Y 4 X R c 6 q S G J p g M c Z + 2 L Z R Y U B N m E 8 d j 9 8 Q y P T d W 2 h 4 J 7 o T 9 v 5 F h Y X J r V p k 3 M z f L G V C K m 6 p V w U D k L X 9 m c j c j E V U j U c 1 F 2 s R m w Q j E l 2 H G Z J I C l W T q I 0 6 5 C 8 r N 0 3 F 7 T F M C f G Q B J p r Z r 7 h k g D U m Y D M s 2 Y z 8 x U S W Q X B W u 6 p 5 d + e V + v U s r C I 6 Q s f o F P n o A t X R L W q g A B E k 0 S t 6 Q + / O i / P h f D p f U 2 n B m e 0 c o r l y v v 8 A P d K e K g = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " g X l W 4 o f y 6 W T d D o b d C q 1 l E t 4 6 O Q 0 = " > A A A C D 3 i c Z V D L T g I x F O 3 g C / G F u n Q z k Z i 4 I G T G m K g 7 o h u X m D h C h A n p l A 4 0 9 D F p 7 x j J h J 8 w b v U 7 X B m 3 f o K f 4 R / Y A R Y C N 2 n u 6 b n n N q c n S j g z 4 H k / T m F l d W 1 9 o 7 h Z 2 t r e 2 d 0 r 7 x 8 8 G J V q Q g O i u N K t C B v K m a Q B M O C 0 l W i K R c R p M x r e 5 P P m E 9 W G K X k P o 4 S G A v c l i x n B Y K n H j s A w i O L s e d w t V 7 y a N y l 3 G f g z U E G z a n T L v 5 2 e I q m g E g j H x r R 9 L 4 E w w x o Y 4 X R c 6 q S G J p g M c Z + 2 L Z R Y U B N m E 8 d j 9 8 Q y P T d W 2 h 4 J 7 o T 9 v 5 F h Y X J r V p k 3 M z f L G V C K m 6 p V w U D k L X 9 m c j c j E V U j U c 1 F 2 s R m w Q j E l 2 H G Z J I C l W T q I 0 6 5 C 8 r N 0 3 F 7 T F M C f G Q B J p r Z r 7 h k g D U m Y D M s 2 Y z 8 x U S W Q X B W u 6 p 5 d + e V + v U s r C I 6 Q s f o F P n o A t X R L W q g A B E k 0 S t 6 Q + / O i / P h f D p f U 2 n B m e 0 c o r l y v v 8 A P d K e K g = = < / l a t e x i t > p ✓ (x|z)

< l a t e x i t s h a 1 _ b a s e 6 4 = " W b l 6 / d 1 + N S e 4 D 7 

I 5 N Z + L i a V j j T 4 = " > A A A C L 3 i c Z V B N S w M x E M 3 6 b f 1 a 9 e g l W I Q K p W x F U G + i F 4 8 V r C 1 0 S 8 m m W R t M N k s y K 9 Z 1 / 4 I / R r z q 7 9 C L e P X o P z C 7 r W B 1 I M y b N 2 / C z A t i w Q 1 4 3 p s z N T 0 z O z e / s F h a W l 5 Z X X P X N y 6 N S j R l T a q E 0 u 2 A G C Z 4 x J r A Q b B 2 r B m R g W C t 4 P o 0 7 7 d u m D Z c R R c w j F l X k q u I h 5 w S s F T P r c S 9 1 A 9 k 6 s O A A c m y i i 8 J D I I w v c 3 w P f 4 p 7 r L d n l v 2 a l 4 R + D + o j 0 E Z j a P R c 7 / 8 v q K J Z B F Q Q Y z p 1 L 0 Y u i n R w K l g W c l P D I s J v S Z X r G N h R C Q z 3 b S 4 K M M 7 l u n j U G n 7 I s A F + 3 s i J d L k u 1 l l n s x E L 2 d A K W G q V g U D m a f 8 m 6 I 2 Q x l U A 1 n N R d q E 5 s 8 i E B 5 2 U x 7 F C b C I j v Y I E 4 F B 4 d w 9 3 O e a U R B D C w j V 3 J 6 C 6 Y B o Q s F 6 X L I e 1 f 8 6 8 h 8 0 9 2 p H N e 9 8 v 3 x 8 M j Z r A W 2 h b V R B d X S A j t E Z a q A m o u g B P a F n 9 O I 8 O q / O u / M x k k 4 5 4 5 l N N B H O 5 z d M W 6 s w < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " W b l 6 / d 1 + N S e 4 D 7 I 5 N Z + L i a V j j T 4 = " > A A A C L 3 i c Z V B N S w M x E M 3 6 b f 1 a 9 e g l W I Q K p W x F U G + i F 4 8 V r C 1 0 S 8 m m W R t M N k s y K 9 Z 1 / 4 I / R r z q 7 9 C L e P X o P z C 7 r W B 1 I M y b N 2 / C z A t i w Q 1 4 3 p s z N T 0 z O z e / s F h a W l 5 Z X X P X N y 6 N S j R l T a q E 0 u 2 A G C Z 4 x J r A Q b B 2 r B m R g W C t 4 P o 0 7 7 d u m D Z c R R c w j F l X k q u I h 5 w S s F T P r c S 9 1 A 9 k 6 s O A A c m y i i 8 J D I I w v c 3 w P f 4 p 7 r L d n l v 2 a l 4 R + D + o j 0 E Z j a P R c 7 / 8 v q K J Z B F Q Q Y z p 1 L 0 Y u i n R w K l g W c l P D I s J v S Z X r G N h R C Q z 3 b S 4 K M M 7 l u n j U G n 7 I s A F + 3 s i J d L k u 1 l l n s x E L 2 d A K W G q V g U D m a f 8 m 6 I 2 Q x l U A 1 n N R d q E 5 s 8 i E B 5 2 U x 7 F C b C I j v Y I E 4 F B 4 d w 9 3 O e a U R B D C w j V 3 J 6 C 6 Y B o Q s F

Z + L i a V j j T 4 = " > A A A C L 3 i c Z V B N S w M x E M 3 6 b f 1 a 9 e g l W I Q K p W x F U G + i F 4 8 V r C 1 0 S 8 m m W R t M N k s y K 9 Z 1 / 4 I / R r z q 7 9 C L e P X o P z C 7 r W B 1 I M y b N 2 / C z A t i w Q 1 4 3 p s z N T 0 z O z e / s F h a W l 5 Z X X P X N y 6 N S j R l T a q E 0 u 2 A G C Z 4 x J r A Q b B 2 r B m R g W C t 4 P o 0 7 7 d u m D Z c R R c w j F l X k q u I h 5 w S s F T P r c S 9 1 A 9 k 6 s O A A c m y i i 8 J D I I w v c 3 w P f 4 p 7 r L d n l v 2 a l 4 R + D + o j 0 E Z j a P R c 7 / 8 v q K J Z B F Q Q Y z p 1 L 0 Y u i n R w K l g W c l P D I s J v S Z X r G N h R C Q z 3 b S 4 K M M 7 l u n j U G n 7 I s A F + 3 s i J d L k u 1 l l n s x E L 2 d A K W G q V g U D m a f 8 m 6 I 2 Q x l U A 1 n N R d q E 5 s 8 i E B 5 2 U x 7 F C b C I j v Y I E 4 F B 4 d w 9 3 O e a U R B D C w j V 3 J 6 C 6 Y B o Q s F 6 X L I e 1 f 8 6 8 h 8 0 9 2 p H N e 9 8 v 3 x 8 M j Z r A W 2 h b V R B d X S A j t E Z a q A m o u g B P a F n 9 O I 8 O q / O u / M x k k 4 5 4 5 l N N B H O 5 z d M W 6 s w < / l a t e x i t > q (z|x)

< l a t e x i t s h a 1 _ b a s e 6 4 = " w 6 B v j H S y P

I f p o j R Q k d L H 6 p + a f X c = " > A A A C L X i c Z V B N S w M x E M 3 6 W e t X 1 a O X Y F E q l L I V Q b 0 V v X i s Y F X o l p J N s 2 1 o s l m T W b G u + w / 8 M e J V f 4 c g i F f P / g O z t Y K 2 A 2 H e v H k T Z p 4 f C W 7 A d d + c q e m Z 2 b n 5 3 E J + c W l 5 Z b W w t n 5 h V K w p a 1 A l l L 7 y i W G C h 6 w B H A S 7 i j Q j 0 h f s 0 u + f Z P 3 L G 6 Y N V + E 5 D C L W k q Q b 8 o B T A p Z q F 3 a u 2 4 n n y 8 S L e j x N S 5 4 k 0 P O D 5 C 7 F 9 / i 3 u E 1 3 2 4 W i W 3 G H g S d B d Q S K a B T 1 d u H L 6 y g a S x Y C F c S Y Z t W N o J U Q D Z w K l u a 9 2 L C I 0 D 7 p s q a F I Z H M t J L h P S n e t k w H B 0 r b F w I e s n 8 n E i J N t p t V Z s n 8 6 2 U M K C V M 2 a q g J 7 O U f T O s z U D 6 Z V + W M 5 E 2 g R l b B I L D V s L D K A Y W 0 p 8 9 g l h g U D j z D n e 4 Z h T E w A J C N b e n Y N o j m l C w D u e t R 9 V x R y Z B Y 6 9 y V H H P 9 o u 1 4 5 F Z O b S J t l A J V d E B q q F T V E c N R N E D e k L P 6 M V 5 d F 6 d d + f j R z r

l j G Y 2 0 L 9 w P r 8 B h 2 + q S A = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " w 6 B v j H S y P

I f p o j R Q k d L H 6 p + a f X c = " > A A A C L X i c Z V B N S w M x E M 3 6 W e t X 1 a O X Y F E q l L I V Q b 0 V v X i s Y F X o l p J N s 2 1 o s l m T W b G u + w / 8 M e J V f 4 c g i F f P / g O z t Y K 2 A 2 H e v H k T Z p 4 f C W 7 A d d + c q e m Z 2 b n 5 3 E J + c W l 5 Z b W w t n 5 h V K w p a 1 A l l L 7 y i W G C h 6 w B H A S 7 i j Q j 0 h f s 0 u + f Z P 3 L G 6 Y N V + E 5 D C L W k q Q b 8 o B T A p Z q F 3 a u 2 4 n n y 8 S L e j x N S 5 4 k 0 P O D 5 C 7 F 9 / i 3 u E 1 3 2 4 W i W 3 G H g S d B d Q S K a B T 1 d u H L 6 y g a S x Y C F c S Y Z t W N o J U Q D Z w K l u a 9 2 L C I 0 D 7 p s q a F I Z H M t J L h P S n e t k w H B 0 r b F w I e s n 8 n E i J N t p t V Z s n 8 6 2 U M K C V M 2 a q g J 7 O U f T O s z U D 6 Z V + W M 5 E 2 g R l b B I L D V s L D K A Y W 0 p 8 9 g l h g U D j z D n e 4 Z h T E w A J C N b e n Y N o j m l C w D u e t R 9 V x R y Z B Y 6 9 y V H H P 9 o u 1 4 5 F Z O b S J t l A J V d E B q q F T V E c N R N E D e k L P 6 M V 5 d F 6 d d + f j R z r

l j G Y 2 0 L 9 w P r 8 B h 2 + q S A = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " w 6 B v j H S y P 

I f p o j R Q k d L H 6 p + a f X c = " > A A A C L X i c Z V B N S w M x E M 3 6 W e t X 1 a O X Y F E q l L I V Q b 0 V v X i s Y F X o l p J N s 2 1 o s l m T W b G u + w / 8 M e J V f 4 c g i F f P / g O z t Y K 2 A 2 H e v H k T Z p 4 f C W 7 A d d + c q e m Z 2 b n 5 3 E J + c W l 5 Z b W w t n 5 h V K w p a 1 A l l L 7 y i W G C h 6 w B H A S 7 i j Q j 0 h f s 0 u + f Z P 3 L G 6 Y N V + E 5 D C L W k q Q b 8 o B T A p Z q F 3 a u 2 4 n n y 8 S L e j x N S 5 4 k 0 P O D 5 C 7 F 9 / i 3 u E 1 3 2 4 W i W 3 G H g S d B d Q S K a B T 1 d u H L 6 y g a S x Y C F c S Y Z t W N o J U Q D Z w K l u a 9 2 L C I 0 D 7 p s q a F I Z H M t J L h P S n e t k w H B 0 r b F w I e s n 8 n E i J N t p t V Z s n 8 6 2 U M K C V M 2 a q g J 7 O U f T O s z U D 6 Z V + W M 5 E 2 g R l b B I L D V s L D K A Y W 0 p 8 9 g l h g U D j z D n e 4 Z h T E w A J C N b e n Y N o j m l C w D u e t R 9 V x R y Z B Y 6 9 y V H H P 9 o u 1 4 5 F Z O b S J t l A J V d E B q q F T V E c N R N E D e k L P 6 M V 5 d F 6 d d + f j R z r l j G Y 2 0 L 9 w P r 8 B h 2 + q S A = = < / l a t

X X X 0 0 µ x,✓ < l a t e x i t s h a 1 _ b a s e 6 4 = " X W K g l M N T 3 f l z + H v E 5 P z z k x a y / d Q = " > A A A C J H i c Z V B N S w M x E M 3 W r 1 q / q h 5 F C B b B Q y l b E d R b 0 Y v H C t Y W u q V k 0 2 w b m m y W Z F Y s y 5 7 8 M e J V f 4 c n 8 e D F / + A / M L v t w d a B M C 9 v 3 g w z z 4 8 E N + C 6 X 0 5 h a X l l d a 2 4 X t r Y 3 N r e K e / u 3 R s V a 8 p a V A m l O z 4 x T P C Q t Y C D Y J 1 I M y J 9 w d r + + D q r t x + Y N l y F d z C J W E + S Y c g D T g l Y q l 8 + 9 G T c T z x J Y O Q H y W N a x Z 4 v E w 9 G D E i a 9 s s V t + b m g f + D + g x U 0 C y a / f K P N 1 A 0 l i w E K o g x 3 b o b Q S 8 h G j g V L C 1 5 s W E R o W M y Z F 0 L Q y K Z 6 S X 5 G S k + t s w A B 0 r b F w L O 2 b 8 d C Z E m W 9 Q q s 2 T m a h k D S g l T t S o Y y S x l Y / K / m U i / 6 s t q J t I m M A u L Q H D R S 3 g Y x c B C O t 0 j i A U G h T P L 8 I B r R k F M L C B U c 3 s K p i O i C Q V r b M l 6 V F 9 0 5 D 9 o n d Y u a + 7 t W a V x N T O r i A 7 Q E T p B d X S O G u g G N V E L U f S E X t A r e n O e n X f n w / m c S g v O r G c f z Y X z / Q s 8 k a a P < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " X W K g l M N T 3 f l z + H v E 5 P z z k x a y / d Q = " > A A A C J H i c Z V B N S w M x E M 3 W r 1 q / q h 5 F C B b B Q y l b E d R b 0 Y v H C t Y W u q V k 0 2 w b m m y W Z F Y s y 5 7 8 M e J V f 4 c n 8 e D F / + A / M L v t w d a B M C 9 v 3 g w z z 4 8 E N + C 6 X 0 5 h a X l l d a 2 4 X t r Y 3 N r e K e / u 3 R s V a 8 p a V A m l O z 4 x T P C Q t Y C D Y J 1 I M y J 9 w d r + + D q r t x + Y N l y F d z C J W E + S Y c g D T g l Y q l 8 + 9 G T c T z x J Y O Q H y W N a x Z 4 v E w 9 G D E i a 9 s s V t + b m g f + D + g x U 0 C y a / f K P N 1 A 0 l i w E K o g x 3 b o b Q S 8 h G j g V L C 1 5 s W E R o W M y Z F 0 L Q y K Z 6 S X 5 G S k + t s w A B 0 r b F w L O 2 b 8 d C Z E m W 9 Q q s 2 T m a h k D S g l T t S o Y y S x l Y / K / m U i / 6 s t q J t I m M A u L Q H D R S 3 g Y x c B C O t 0 j i A U G h T P L 8 I B r R k F M L C B U c 3 s K p i O i C Q V r b M l 6 V F 9 0 5 D 9 o n d Y u a + 7 t W a V x N T O r i A 7 Q E T p B d X S O G u g G N V E L U f S E X t A r e n O e n X f n w / m c S g v O r G c f z Y X z / Q s 8 k a a P < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " X W K g l M N T 3 f l z + H v E 5 P z z k x a y / d Q = " > A A A C J H i c Z V B N S w M x E M 3 W r 1 q / q h 5 F C B b B Q y l b E d R b 0 Y v H C t Y W u q V k 0 2 w b m m y W Z F Y s y 5 7 8 M e J V f 4 c n 8 e D F / + A / M L v t w d a B M C 9 v 3 g w z z 4 8 E N + C 6 X 0 5 h a X l l d a 2 4 X t r Y 3 N r e K e / u 3 R s V a 8 p a V A m l O z 4 x T P C Q t Y C D Y J 1 I M y J 9 w d r + + D q r t x + Y N l y F d z C J W E + S Y c g D T g l Y q l 8 + 9 G T c T z x J Y O Q H y W N a x Z 4 v E w 9 G D E i a 9 s s V t + b m g f + D + g x U 0 C y a / f K P N 1 A 0 l i w E K o g x 3 b o b Q S 8 h G j g V L C 1 5 s W E R o W M y Z F 0 L Q y K Z 6 S X 5 G S k + t s w A B 0 r b F w L O 2 b 8 d C Z E m W 9 Q q s 2 T m a h k D S g l T t S o Y y S x l Y / K / m U i / 6 s t q J t I m M A u L Q H D R S 3 g Y x c B C O t 0 j i A U G h T P L 8 I B r R k F M L C B U c 3 s K p i O i C Q V r b M l 6 V F 9 0 5 D 9 o n d Y u a + 7 t W a V x N T O r i A 7 Q E T p B d X S O G u g G N V E L U f S E X t A r e n O e n X f n w / m c S g v O r G c f z Y X z / Q s 8 k a a P < / l a t e x i t > µ x, < l a t e x i t s h a 1 _ b a s e 6 4 = " J r R 8 b p J T 9 W I D X N 5 J E d h m m w 5 B e V U = " > A A A C I n i c Z V B N S w M x E M 3 6 W e t X 1 W M v w S J 4 K G U r g n o r e v F Y w b W F b i n Z N N u G J p s l m R X L s g d / j H j V 3 + F J P A n + C f + B 2 b Y H 2 w 6 E e X n z Z p h 5 Q S y 4 A d f 9 d l Z W 1 9 Y 3 N g t b x e 2 d 3 b 3 9 0 s H h g 1 G J p s y j S i j d D o h h g k f M A w 6 C t W P N i A w E a w W j m 7 z e e m T a c B X d w z h m X U k G E Q 8 5 J W C p X q n s y 6 S X + p L A M A j T p 6 y K / U C m f j z k W d Y r V d y a O w m 8 D O o z U E G z a P Z K v 3 5 f 0 U S y C K g g x n T q b g z d l G j g V L C s 6 C e G x Y S O y I B 1 L I y I Z K a b T o 7 I 8 I l l + j h U 2 r 4 I 8 I T 9 3 5 E S a f I 1 r T J P Z q 6 W M 6 C U M F W r g q H M U z 5 m 8 j d j G V Q D W c 1 F 2 o R m Y R E I L 7 s p j + I E W E S n e 4 S J w K B w b h j u c 8 0 o i L E F h G p u T 8 F 0 S D S h Y G 0 t W o / q i 4 4 s A + + s d l V z 7 8 4 r j e u Z W Q V U R s f o F N X R B W q g W 9 R E H q L o G b 2 i N / T u v D g f z q f z N Z W u O L O e I z Q X z s 8 f g A 2 l p g = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " J r R 8 b p J T 9 W I D X N 5 J E d h m m w 5 B e V U = " > A A A C I n i c Z V B N S w M x E M 3 6 W e t X 1 W M v w S J 4 K G U r g n o r e v F Y w b W F b i n Z N N u G J p s l m R X L s g d / j H j V 3 + F J P A n + C f + B 2 b Y H 2 w 6 E e X n z Z p h 5 Q S y 4 A d f 9 d l Z W 1 9 Y 3 N g t b x e 2 d 3 b 3 9 0 s H h g 1 G J p s y j S i j d D o h h g k f M A w 6 C t W P N i A w E a w W j m 7 z e e m T a c B X d w z h m X U k G E Q 8 5 J W C p X q n s y 6 S X + p L A M A j T p 6 y K / U C m f j z k W d Y r V d y a O w m 8 D O o z U E G z a P Z K v 3 5 f 0 U S y C K g g x n T q b g z d l G j g V L C s 6 C e G x Y S O y I B 1 L I y I Z K a b T o 7 I 8 I l l + j h U 2 r 4 I 8 I T 9 3 5 E S a f I 1 r T J P Z q 6 W M 6 C U M F W r g q H M U z 5 m 8 j d j G V Q D W c 1 F 2 o R m Y R E I L 7 s p j + I E W E S n e 4 S J w K B w b h j u c 8 0 o i L E F h G p u T 8 F 0 S D S h Y G 0 t W o / q i 4 4 s A + + s d l V z 7 8 4 r j e u Z W Q V U R s f o F N X R B W q g W 9 R E H q L o G b 2 i N / T u v D g f z q f z N Z W u O L O e I z Q X z s 8 f g A 2 l p g = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " J r R 8 b p J T 9 W I D X N 5 J E d h m m w 5 B e V U = " > A A A C I n i c Z V B N S w M x E M 3 6 W e t X 1 W M v w S J 4 K G U r g n o r e v F Y w b W F b i n Z N N u G J p s l m R X L s g d / j H j V 3 + F J P A n + C f + B 2 b Y H 2 w 6 E e X n z Z p h 5 Q S y 4 A d f 9 d l Z W 1 9 Y 3 N g t b x e 2 d 3 b 3 9 0 s H h g 1 G J p s y j S i j d D o h h g k f M A w 6 C t W P N i A w E a w W j m 7 z e e m T a c B X d w z h m X U k G E Q 8 5 J W C p X q n s y 6 S X + p L A M A j T p 6 y K / U C m f j z k W d Y r V d y a O w m 8 D O o z U E G z a P Z K v 3 5 f 0 U S y C K g g x n T q b g z d l G j g V L C s 6 C e G x Y S O y I B 1 L I y I Z K a b T o 7 I 8 I l l + j h U 2 r 4 I 8 I T 9 3 5 E S a f I 1 r T J P Z q 6 W M 6 C U M F W r g q H M U z 5 m 8 j d j G V Q D W c 1 F 2 o R m Y R E I L 7 s p j + I E W E S n e 4 S J w K B w b h j u c 8 0 o i L E F h G p u T 8 F 0 S D S h Y G 0 t W o / q i 4 4 s A + + s d l V z 7 8 4 r j e u Z W Q V U R s f o F N X R B W q g W 9 R E H q L o G b 2 i N / T u v D g f z q f z N Z W u O L O e I z Q X z s 8 f g A 2 l p g = = < / l a t e x i t >

mean of true model posterior mean of approximate posterior

µ x, = µ x,✓ < l a t e x i t s h a 1 _ b a s e 6 4 = " 8 v n u d 1 + p s g y 2 E O 9 E Z d P k D c M D h n Q = " > A A A C Q n i c b V B N S 8 N A E N 3 4 b f 2 K e v S y W A Q P p a Q i q I h Q 9 O K x g l W h K W W z 3 Z j F 3 W z Y n Y g l 5 M f 4 Y 8 S r H v 0 N n s S r B z d p D l o d W O b t m z f D z A s S w Q 1 4 3 p s z N T 0 z O z e / s F h b W l 5 Z X X P X N 6 6 M S j V l X a q E 0 j c B M U z w m H W B g 2 A 3 i W Z E B o J d B 3 d n R f 3 6 n m n D V X w J o 4 T 1 J b m N e c g p A U s N 3 G N f p o P M l w S i I M w e 8 g b 2 A 5 n 5 S c T z H J / g / 6 s Q M S B 5 P n D r X t M r A / 8 F r Q r U U R W d g f v h D x V N J Y u B C m J M r + U l 0 M + I B k 4 F y 2 t + a l h C 6 B 2 5 Z T 0 L Y y K Z 6 W f l k T n e s c w Q h 0 r b F w M u 2 Z 8 d G Z G m W N Q q i 2 R + 1 Q o G l B K m Y V U Q y S I V Y 8 q / G c m g E c h G I d I m N B O L Q H j Y z 3 i c p M B i O t 4 j T A U G h Q t D 8 Z B r R k G M L C B U c 3 s K p h H R h I K 1 v W Y 9 a k 0 6 8 h d 0 9 5 p H T e 9 i v 9 4 + r c x a Q F t o G + 2 i F j p A b X S O O q i L K H p E z + g F v T p P z r v z 4 X y O p V N O 1 b O J f o X z 9 Q 0 l J b M X < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " 8 v n u d 1 + p s g y 2 E O 9 E Z d P k D c M D h n Q = " > A A A C Q n i c b V B N S 8 N A E N 3 4 b f 2 K e v S y W A Q P p a Q i q I h Q 9 O K x g l W h K W W z 3 Z j F 3 W z Y n Y g l 5 M f 4 Y 8 S r H v 0 N n s S r B z d p D l o d W O b t m z f D z A s S w Q 1 4 3 p s z N T 0 z O z e / s F h b W l 5 Z X X P X N 6 6 M S j V l X a q E 0 j c B M U z w m H W B g 2 A 3 i W Z E B o J d B 3 d n R f 3 6 n m n D V X w J o 4 T 1 J b m N e c g p A U s N 3 G N f p o P M l w S i I M w e 8 g b 2 A 5 n 5 S c T z H J / g / 6 s Q M S B 5 P n D r X t M r A / 8 F r Q r U U R W d g f v h D x V N J Y u B C m J M r + U l 0 M + I B k 4 F y 2 t + a l h C 6 B 2 5 Z T 0 L Y y K Z 6 W f l k T n e s c w Q h 0 r b F w M u 2 Z 8 d G Z G m W N Q q i 2 R + 1 Q o G l B K m Y V U Q y S I V Y 8 q / G c m g E c h G I d I m N B O L Q H j Y z 3 i c p M B i O t 4 j T A U G h Q t D 8 Z B r R k G M L C B U c 3 s K p h H R h I K 1 v W Y 9 a k 0 6 8 h d 0 9 5 p H T e 9 i v 9 4 + r c x a Q F t o G + 2 i F j p A b X S O O q i L K H p E z + g F v T p P z r v z 4 X y O p V N O 1 b O J f o X z 9 Q 0 l J b M X < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " 8 v n u d 1 + p s g y 2 E O 9 E Z d P k D c M D h n Q = " > A A A C Q n i c b V B N S 8 N A E N 3 4 b f 2 K e v S y W A Q P p a Q i q I h Q 9 O K x g l W h K W W z 3 Z j F 3 W z Y n Y g l 5 M f 4 Y 8 S r H v 0 N n s S r B z d p D l o d W O b t m z f D z A s S w Q 1 4 3 p s z N T 0 z O z e / s F h b W l 5 Z X X P X N 6 6 M S j V l X a q E 0 j c B M U z w m H W B g 2 A 3 i W Z E B o J d B 3 d n R f 3 6 n m n D V X w J o 4 T 1 J b m N e c g p A U s N 3 G N f p o P M l w S i I M w e 8 g b 2 A 5 n 5 S c T z H J / g / 6 s Q M S B 5 P n D r X t M r A / 8 F r Q r U U R W d g f v h D x V N J Y u B C m J M r + U l 0 M + I B k 4 F y 2 t + a l h C 6 B 2 5 Z T 0 L Y y K Z 6 W f l k T n e s c w Q h 0 r b F w M u 2 Z 8 d G Z G m W N Q q i 2 R + 1 Q o G l B K m Y V U Q y S I V Y 8 q / G c m g E c h G I d I m N B O L Q H j Y z 3 i c p M B i O t 4 j T A U G h Q t D 8 Z B r R k G M L C B U c 3 s K p h H R h I K 1 v W Y 9 a k 0 6 8 h d 0 9 5 p H T e 9 i v 9 4 + r c x a Q F t o G + 2 i F j p A b X S O O q i L K H p E z + g F v T p P z r v z 4 X y O p V N O 1 b O J f o X z 9 Q 0 l J b M X < / l a t e x i t > (b) Posterior mean space

Figure 1 : Left: Depiction of generative model p(z)p θ (x|z) and inference network q φ (z|x) in VAEs. Right: A toy posterior mean space (µ x,θ , µ x,φ ) with scalar z. The horizontal axis represents the mean of the model posterior p θ (z|x), and the vertical axis represents the mean of the approximate posterior q φ (z|x). The dashed diagonal line represents when the approximate posterior matches the true model posterior in terms of mean.

important applications like natural language processing. Existing work analyzes this problem from a static optimization perspective, noting that the collapsed solution is often a reasonably good local optimum in terms of ELBO (Chen et al., 2017; Zhao et al., 2017; Alemi et al., 2018) . Thus, many proposed solutions to posterior collapse focus on weakening the generator by replacing it with a non-recurrent alternative (Yang et al., 2017; Semeniuta et al., 2017) or modifying the training objective (Zhao et al., 2017; Tolstikhin et al., 2018) . In this paper, we analyze the problem from the perspective of training dynamics and propose a novel training procedure for VAEs that addresses posterior collapse. In contrast with other solutions, our proposed procedure optimizes the standard ELBO objective and does not require modification to the VAE model or its parameterization.

Recently, Kim et al. (2018) proposed a new approach to training VAEs by composing the standard inference network with additional mean-field updates. The resulting semi-amortized approach empirically avoided collapse and obtained better ELBO. However, because of the costly instancespecific local inference steps, the new method is more than 10x slower than basic VAE training in practice. It is also unclear why the basic VAE method fails to find better local optima that make use of latents. We consider two questions in this paper: (1) Why does basic VAE training often fall into undesirable collapsed local optima? (2) Is there a simpler way to change the training trajectory to find a non-trivial local optimum?

To this end, we first study the posterior collapse problem from the perspective of training dynamics. We find, empirically, that the posterior approximation often lags far behind the true model posterior in the initial stages of training (Section 3). We then demonstrate how such lagging behavior can drive the generative model towards a collapsed local optimum, and propose a novel training procedure for VAEs that aggressively optimizes the inference network with more updates to mitigate lag (Section 4). Without introducing new modeling components over basic VAEs or additional complexity, our approach is surprisingly simple yet effective in circumventing posterior collapse. As a density estimator, it outperforms neural autoregressive baselines on both text (Yahoo and Yelp) and image (OMNIGLOT) benchmarks, leading to comparable performance with more complicated previous state-of-the-art methods at a fraction of the training cost (Section 6).

2. BACKGROUND:

2.1 VARIATIONAL AUTOENCODERS VAEs learn deep generative models defined by a prior p(z) and a conditional distribution p θ (x|z) as shown in Figure 1(a) . In most cases the marginal data likelihood is intractable, so VAEs instead optimize a tractable variational lower bound (ELBO) of log p θ (x),

EQUATION

where q φ (z|x) is a variational distribution parameterized by an inference network with parameters φ, and p θ (x|z) denotes the generator network with parameters θ. q φ (z|x) is optimized to approximate the model posterior p θ (z|x). This lower bound is composed of a reconstruction loss term that encourages the inference network to encode information necessary to generate the data and a KL regularizer to push q φ (z|x) towards the prior p(z). Below, we consider p(z) := N (0, I) unless otherwise specified. A key advantage of using inference networks (also called amortized inference) to train deep generative models over traditional locally stochastic variational inference (Hoffman et al., 2013) is that they share parameters over all data samples, amortizing computational cost and allowing for efficient training.

The term VAE is often used both to denote the class of generative models and the amortized inference procedure used in training. In this paper, it is important to distinguish the two and throughout we will refer to the generative model as the VAE model, and the training procedure as VAE training.

2.2. POSTERIOR COLLAPSE:

Despite VAE's appeal as a tool to learn unsupervised representations through the use of latent variables, as mentioned in the introduction, VAE models are often found to ignore latent variables when using flexible generators like LSTMs (Bowman et al., 2016) . This problem of "posterior collapse" occurs when the training procedure falls into the trivial local optimum of the ELBO objective in which both the variational posterior and true model posterior collapse to the prior. This is undesirable because an important goal of VAEs is to learn meaningful latent features for inputs. Mathematically, posterior collapse represents a local optimum of VAEs where q φ (z|x) = p θ (z|x) = p(z) for all x.

To facilitate our analysis about the causes leading up to collapse, we further define two partial collapse states: model collapse, when p θ (z|x) = p(z), and inference collapse, when q φ (z|x) = p(z) for all x. Note that in this paper we use these two terms to denote the posterior states in the middle of training instead of local optima at the end. These two partial collapse states may not necessarily happen at the same time, which we will discuss later.

2.3. VISUALIZATION OF POSTERIOR DISTRIBUTION:

Posterior collapse is closely related to the true model posterior p θ (z|x) and the approximate posterior q φ (z|x) as it is defined. Thus, in order to observe how posterior collapse happens, we track the state of p θ (z|x) and q φ (z|x) over the course of training, and analyze the training trajectory in terms of the posterior mean space U = {µ : µ = (µ T x,θ , µ T x,φ )}, where µ x,θ and µ x,φ are the means of p θ (z|x) and q φ (z|x), respectively. 2 We can then roughly consider µ x,θ = 0 as model collapse and µ x,φ = 0 as inference collapse as we defined before. Each x will be projected to a point in this space under the current model and inference network parameters. If z is a scalar we can efficiently compute µ x,θ and visualize the posterior mean space as shown in Figure 1(b ). The diagonal line µ x,θ = µ x,φ represents parameter settings where q φ (z|x) is equal to p θ (z|x) in terms of mean, indicating a well-trained inference network. The collapsed local optimum is located at the origin,3 while the data points at a more desirable local optima may be distributed along the diagonal. In this paper we will utilize this posterior mean space multiple times to analyze the posterior dynamics.

3. A LAGGING INFERENCE NETWORK PREVENTS USING LATENT CODES:

In this section we analyze posterior collapse from a perspective of training dynamics. We will answer the question of why the basic VAE training with strong decoders tends to hit a collapsed local optimum and provide intuition for the simple solution we propose in Section 4. [Can you make the font slightly bigger? -GN] [Can you also add labels indicating the four training stages ("Stage 1", "Stage 2", ..., or even better "X Epochs", "X Epochs", ...). -GN] [Add labels to the top of "Vanilla VAE" and "Our Approach" (or whatever alternative name you come up with.) -GN] p ✓ (z|x) is log p ✓ (x). However, the dependence (if any) may be lost if q (z|x) is still close to p(z) because p ✓ (z|x) would be pressured to agree with q (z|x) [This is important, but I think not very clear. Could you be more explicit? -GN].

To examine the dynamic relationship between approximate posterior q (z|x), model posterior p ✓ (z|x), and prior p(z) during basic VAE training, we will next visualize the posterior mean-value space by training a basic VAE with a one-dimensional ["scalar"? -GN] latent variable on a relatively simple synthetic dataset.

3.2. OBSERVATIONS ON SYNTHETIC DATA:

As a synthetic dataset we use discrete sequence data, as posterior collapse has been found the most severe in text modeling tasks. To mimic the diversity of the real data distribution, we sample discrete data from a deep generative model, with a [[[four-mode mixture of Gaussian prior -four-modes of a Gaussian mixture as a prior?]]] and an LSTM generator. More details about this synthetic dataset and experiment details can be found in Appendix B.1.

We train a basic VAE with a one-dimensional latent variable, LSTM encoder, and LSTM decoder2 on our synthetic dataset. We sample 500 data points and show them on the posterior mean space plots at four different training stages from beginning to convergence in Figure 2 . The mean value of the approximate posterior distribution µ x, is from the output of the inference network, and µ x,✓ is approximated by:

EQUATION

where C is a grid set with small stride and sufficiently large coverage. We assume the density value outside this grid set is zero. p ✓ (z|x) is also approximated on this grid.

As illustrated in Figure 2 , all points are located at origin [[[after-upon] ]] initialization, which means z and x are both almost independent and zero mean at the beginning of training. In the second stage, the points starts to spread along the µ x,✓ axis. This phenomenon implies that for some data points p ✓ (z|x) moves far away from the prior p(z), and confirms that log p ✓ (x) is able to help escape model collapse. However, all of these points are still close to the µ x, axis, which suggests that q (z|x) fails to catch up to p ✓ (z|x) and these points are still in an inference collapsed state. 

iter = 0 iter = 200 iter = 2000 convergence Basic Aggressive µ x,✓ < l a t e x i t s h a 1 _ b a s e 6 4 = " X W K g l M N T 3 f l z + H v E 5 P z z k x a y / d Q = " > A A A C J H i c Z V B N S w M x E M 3 W r 1 q / q h 5 F C B b B Q y l b E d R b 0 Y v H C t Y W u q V k 0 2 w b m m y W Z F Y s y 5 7 8 M e J V f 4 c n 8 e D F / + A / M L v t w d a B M C 9 v 3 g w z z 4 8 E N + C 6 X 0 5 h a X l l d a 2 4 X t r Y 3 N r e K e / u 3 R s V a 8 p a V A m l O z 4 x T P C Q t Y C D Y J 1 I M y J 9 w d r + + D q r t x + Y N l y F d z C J W E + S Y c g D T g l Y q l 8 + 9 G T c T z x J Y O Q H y W N a x Z 4 v E w 9 G D E i a 9 s s V t + b m g f + D + g x U 0 C y a / f K P N 1 A 0 l i w E K o g x 3 b o b Q S 8 h G j g V L C 1 5 s W E R o W M y Z F 0 L Q y K Z 6 S X 5 G S k + t s w A B 0 r b F w L O 2 b 8 d C Z E m W 9 Q q s 2 T m a h k D S g l T t S o Y y S x l Y / K / m U i / 6 s t q J t I m M A u L Q H D R S 3 g Y x c B C O t 0 j i A U G h T P L 8 I B r R k F M L C B U c 3 s K p i O i C Q V r b M l 6 V F 9 0 5 D 9 o n d Y u a + 7 t W a V x N T O r i A 7 Q E T p B d X S O G u g G N V E L U f S E X t A r e n O e n X f n w / m c S g v O r G c f z Y X z / Q s 8 k a a P < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " X W K g l M N T 3 f l z + H v E 5 P z z k x a y / d Q = " > A A A C J H i c Z V B N S w M x E M 3 W r 1 q / q h 5 F C B b B Q y l b E d R b 0 Y v H C t Y W u q V k 0 2 w b m m y W Z F Y s y 5 7 8 M e J V f 4 c n 8 e D F / + A / M L v t w d a B M C 9 v 3 g w z z 4 8 E N + C 6 X 0 5 h a X l l d a 2 4 X t r Y 3 N r e K e / u 3 R s V a 8 p a V A m l O z 4 x T P C Q t Y C D Y J 1 I M y J 9 w d r + + D q r t x + Y N l y F d z C J W E + S Y c g D T g l Y q l 8 + 9 G T c T z x J Y O Q H y W N a x Z 4 v E w 9 G D E i a 9 s s V t + b m g f + D + g x U 0 C y a / f K P N 1 A 0 l i w E K o g x 3 b o b Q S 8 h G j g V L C 1 5 s W E R o W M y Z F 0 L Q y K Z 6 S X 5 G S k + t s w A B 0 r b F w L O 2 b 8 d C Z E m W 9 Q q s 2 T m a h k D S g l T t S o Y y S x l Y / K / m U i / 6 s t q J t I m M A u L Q H D R S 3 g Y x c B C O t 0 j i A U G h T P L 8 I B r R k F M L C B U c 3 s K p i O i C Q V r b M l 6 V F 9 0 5 D 9 o n d Y u a + 7 t W a V x N T O r i A 7 Q E T p B d X S O G u g G N V E L U f S E X t A r e n O e n X f n w / m c S g v O r G c f z Y X z / Q s 8 k a a P < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " X W K g l M N T 3 f l z + H v E 5 P z z k x a y / d Q = " > A A A C J H i c Z V B N S w M x E M 3 W r 1 q / q h 5 F C B b B Q y l b E d R b 0 Y v H C t Y W u q V k 0 2 w b m m y W Z F Y s y 5 7 8 M e J V f 4 c n 8 e D F / + A / M L v t w d a B M C 9 v 3 g w z z 4 8 E N + C 6 X 0 5 h a X l l d a 2 4 X t r Y 3 N r e K e / u 3 R s V a 8 p a V A m l O z 4 x T P C Q t Y C D Y J 1 I M y J 9 w d r + + D q r t x + Y N l y F d z C J W E + S Y c g D T g l Y q l 8 + 9 G T c T z x J Y O Q H y W N a x Z 4 v E w 9 G D E i a 9 s s V t + b m g f + D + g x U 0 C y a / f K P N 1 A 0 l i w E K o g x 3 b o b Q S 8 h G j g V L C 1 5 s W E R o W M y Z F 0 L Q y K Z 6 S X 5 G S k + t s w A B 0 r b F w L O 2 b 8 d C Z E m W 9 Q q s 2 T m a h k D S g l T t S o Y y S x l Y / K / m U i / 6 s t q J t I m M A u L Q H D R S 3 g Y x c B C O t 0 j i A U G h T P L 8 I B r R k F M L C B U c 3 s K p i O i C Q V r b M l 6 V F 9 0 5 D 9 o n d Y u a + 7 t W a V x N T O r i A 7 Q E T p B d X S O G u g G N V E L U f S E X t A r e n O e n X f n w / m c S g v O r G c f z Y X z / Q s 8 k a a P < / l a t e x i t > µ x, < l a t e x i t s h a 1 _ b a s e 6 4 = " J r R 8 b p J T 9 W I D X N 5 J E d h m m w 5 B e V U = " > A A A C I n i c Z V B N S w M x E M 3 6 W e t X 1 W M v w S J 4 K G U r g n o r e v F Y w b W F b i n Z N N u G J p s l m R X L s g d / j H j V 3 + F J P A n + C f + B 2 b Y H 2 w 6 E e X n z Z p h 5 Q S y 4 A d f 9 d l Z W 1 9 Y 3 N g t b x e 2 d 3 b 3 9 0 s H h g 1 G J p s y j S i j d D o h h g k f M A w 6 C t W P N i A w E a w W j m 7 z e e m T a c B X d w z h m X U k G E Q 8 5 J W C p X q n s y 6 S X + p L A M A j T p 6 y K / U C m f j z k W d Y r V d y a O w m 8 D O o z U E G z a P Z K v 3 5 f 0 U S y C K g g x n T q b g z d l G j g V L C s 6 C e G x Y S O y I B 1 L I y I Z K a b T o 7 I 8 I l l + j h U 2 r 4 I 8 I T 9 3 5 E S a f I 1 r T J P Z q 6 W M 6 C U M F W r g q H M U z 5 m 8 j d j G V Q D W c 1 F 2 o R m Y R E I L 7 s p j + I E W E S n e 4 S J w K B w b h j u c 8 0 o i L E F h G p u T 8 F 0 S D S h Y G 0 t W o / q i 4 4 s A + + s d l V z 7 8 4 r j e u Z W Q V U R s f o F N X R B W q g W 9 R E H q L o G b 2 i N / T u v D g f z q f z N Z W u O L O e I z Q X z s 8 f g A 2 l p g = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " J r R 8 b p J T 9 W I D X N 5 J E d h m m w 5 B e V U = " > A A A C I n i c Z V B N S w M x E M 3 6 W e t X 1 W M v w S J 4 K G U r g n o r e v F Y w b W F b i n Z N N u G J p s l m R X L s g d / j H j V 3 + F J P A n + C f + B 2 b Y H 2 w 6 E e X n z Z p h 5 Q S y 4 A d f 9 d l Z W 1 9 Y 3 N g t b x e 2 d 3 b 3 9 0 s H h g 1 G J p s y j S i j d D o h h g k f M A w 6 C t W P N i A w E a w W j m 7 z e e m T a c B X d w z h m X U k G E Q 8 5 J W C p X q n s y 6 S X + p L A M A j T p 6 y K / U C m f j z k W d Y r V d y a O w m 8 D O o z U E G z a P Z K v 3 5 f 0 U S y C K g g x n T q b g z d l G j g V L C s 6 C e G x Y S O y I B 1 L I y I Z K a b T o 7 I 8 I l l + j h U 2 r 4 I 8 I T 9 3 5 E S a f I 1 r T J P Z q 6 W M 6 C U M F W r g q H M U z 5 m 8 j d j G V Q D W c 1 F 2 o R m Y R E I L 7 s p j + I E W E S n e 4 S J w K B w b h j u c 8 0 o i L E F h G p u T 8 F 0 S D S h Y G 0 t W o / q i 4 4 s A + + s d l V z 7 8 4 r j e u Z W Q V U R s f o F N X R B W q g W 9 R E H q L o G b 2 i N / T u v D g f z q f z N Z W u O L O e I z Q X z s 8 f g A 2 l p g = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " J r R 8 b p J T 9 W I D X N 5 J E d h m m w 5 B e V U = " > A A A C I n i c Z V B N S w M x E M 3 6 W e t X 1 W M v w S J 4 K G U r g n o r e v F Y w b W F b i n Z N N u G J p s l m R X L s g d / j H j V 3 + F J P A n + C f + B 2 b Y H 2 w 6 E e X n z Z p h 5 Q S y 4 A d f 9 d l Z W 1 9 Y 3 N g t b x e 2 d 3 b 3 9 0 s H h g 1 G J p s y j S i j d D o h h g k f M A w 6 C t W P N i A w E a w W j m 7 z e e m T a c B X d w z h m X U k G E Q 8 5 J W C p X q n s y 6 S X + p L A M A j T p 6 y K / U C m f j z k W d Y r V d y a O w m 8 D O o z U E G z a P Z K v 3 5 f 0 U S y C K g g x n T q b g z d l G j g V L C s 6 C e G x Y S O y I B 1 L I y I Z K a b T o 7 I 8 I l l + j h U 2 r 4 I 8 I T 9 3 5 E S a f I 1 r T J P Z q 6 W M 6 C U M F W r g q H M U z 5 m 8 j d j G V Q D W c 1 F 2 o R m Y R E I L 7 s p j + I E W E S n e 4 S J w K B w b h j u c 8 0 o i L E F h G p u T 8 F 0 S D S h Y G 0 t W o / q i 4 4 s A + + s d l V z 7 8 4 r j e u Z W Q V U R s f o F N X R B W q g W 9 R E H q L o G b 2 i N / T u v D g f z q f z N Z W u O L O e I z Q X z s 8 f g A 2 l p g = = < / l a t

3.1. INTUITIONS FROM ELBO:

Since posterior collapse is directly relevant to the approximate posterior q φ (z|x) and true model posterior p θ (z|x), we aim to analyze their training dynamics to study how posterior collapse happens.

To this end, it is useful to analyze an alternate form of ELBO:

L(x; θ, φ) = log p θ (x)

marginal log data likelihood -D KL (q φ (z|x) p θ (z|x))

agreement between approximate and model posteriors ,

With this view, the only goal of approximate posterior q φ (z|x) is to match model posterior p θ (z|x), while the optimization of p θ (z|x) is influenced by two forces, one of which is the ideal objective marginal data likelihood, and the other is D KL (q φ (z|x) p θ (z|x)), which drives p θ (z|x) towards q φ (z|x). Ideally if the approximate posterior is perfect, the second force will vanish, with ∇ θ D KL (q φ (z|x)|p θ (z|x)) = 0 when q φ (z|x) = p θ (z|x). At the start of training, z and x are nearly independent under both q φ (z|x) and p θ (z|x) as we show in Section 3.2, i.e. all x suffer from model collapse in the beginning. Then the only component in the training objective that possibly causes dependence between z and x under p θ (z|x) is log p θ (x). However, this pressure may be overwhelmed by the KL term when p θ (z|x) and q φ (z|x) start to diverge but z and x remain independent under q φ (z|x). We hypothesize that, in practice, training drives p θ (z|x) and q φ (z|x) to the prior in order to bring them into alignment, while locking into model parameters that capture the distribution of x while ignoring z. Critically, posterior collapse is a local optimum; once a set of parameters that achieves these goals are reached, gradient optimization fails to make further progress, even if better overall models that make use of z to describe x exist.

Next we visualize the posterior mean space by training a basic VAE with a scalar latent variable on a relatively simple synthetic dataset to examine our hypothesis.

As a synthetic dataset we use discrete sequence data since posterior collapse has been found the most severe in text modeling tasks. Details on this synthetic dataset and experiment are in Appendix B.1.

We train a basic VAE with a scalar latent variable, LSTM encoder, and LSTM decoder on our synthetic dataset. We sample 500 data points from the validation set and show them on the posterior mean space plots at four different training stages from initialization to convergence in Figure 2 . The mean of the approximate posterior distribution µ x,φ is from the output of the inference network, and µ x,θ can be approximated by discretization of the true model posterior p θ (z|x) (see Appendix A). We note that the approximate posterior q φ (z|x) takes an aggressive step to catch up to the model posterior p θ (z|x).

As illustrated in Figure 2 , all points are located at the origin upon initialization4 , which means z and x are almost independent in terms of both q φ (z|x) and p θ (z|x) at the beginning of training. In the second stage of basic VAE training, the points start to spread along the µ x,θ axis. This phenomenon implies that for some data points p θ (z|x) moves far away from the prior p(z), and confirms that log p θ (x) is able to help move away from model collapse. However, all of these points are still distributed along a horizontal line, which suggests that q φ (z|x) fails to catch up to p θ (z|x) and these points are still in a state of inference collapse. As expected, the dependence between z and x under p θ (z|x) is gradually lost and finally the model converges to the collapsed local optimum.

4. METHOD:



4.1. AGGRESSIVE TRAINING OF THE INFERENCE NETWORK:

The problem reflected in Figure 2 implies that the inference network is lagging far behind p θ (z|x), and might suggest more "aggressive" inference network updates are needed. Instead of blaming the poor approximation on the limitation of the inference network's amortization, we hypothesize that the optimization of the inference and generation networks are imbalanced, and propose to separate the optimization of the two. Specifically, we change the training procedure to:

EQUATION

where optimizing the inference network q φ (z|x) is an inner loop in the entire training process as shown in Algorithm 1. This training procedure shares the same spirit with traditional stochastic variational inference (SVI) (Hoffman et al., 2013) that performs iterative inference for each data point separately and suffers from very lengthy iterative estimation. Compared with recent work that try to combine amortized variational inference and SVI (Hjelm et al., 2016; Krishnan et al., 2018; Kim et al., 2018; Marino et al., 2018) where the inference network is learned to be a component to help perform instance-specific variational inference, our approach keeps variational inference fully amortized, allowing for reverting back to efficient basic VAE training as discussed in Section 4.2. Also, this aggressive inference network optimization algorithm is as simple as basic VAE training without introducing additional SVI steps, yet attains comparable performance to more sophisticated approaches as we will show in Section 6.

4.2. STOPPING CRITERION:

Always training with Eq.4 would be inefficient and neglects the benefit of the amortized inference network. Following our previous analysis, the term D KL (q φ (z|x) p θ (z|x)) tends to pressure q φ (z|x) or p θ (z|x) to p(z) only if at least one of them is close to p(z), and thus we posit that if we can confirm that we haven't reached this degenerate condition, we can continue with standard VAE training. Since q φ (z|x) is the one lagging behind, we use the mutual information I q between z and x under q φ (z|x) to control our stopping criterion. In practice, we compute the mutual information on the validation set every epoch, and stop the aggressive updates when I q stops climbing. In all our experiments in this paper we found that the aggressive algorithm usually reverts back to basic VAE training within 5 epochs. Mutual information, I q can be computed by (Hoffman & Johnson, 2016) :

EQUATION

where p d (x) is the empirical distribution. The aggregated posterior,

q φ (z) = E x∼p d (x) [q φ (z|x)],

can be approximated with a Monte Carlo estimate. D KL (q φ (z) p(z)) is also approximated by Monte Carlo, where samples from q φ (z) can be easily obtained by ancestral sampling (i.e. sample

x from dataset and sample z ∼ q φ (z|x)). This estimator for I q is the same as in (Dieng et al., 2018) , which is biased because the estimation for D KL (q φ (z) p(z)) is biased. More specifically, it is a Monte Carlo estimate of an upper bound of mutual information. The complete algorithm is shown in Algorithm 1.

4.3. OBSERVATIONS ON SYNTHETIC DATASET:

By training the VAE model with our approach on synthetic data, we visualize the 500 data samples in the posterior mean space in Figure 2 . From this, it is evident that the points move towards µ x,θ = µ x,φ and are roughly distributed along the diagonal in the end. This is in striking contrast to the basic VAE and confirms our hypothesis that the inference and generator optimization can be rebalanced by simply performing more updates of the inference network. In Figure 3 we show the training trajectory of one single data instance for the first several optimization iterations and observe how the aggressive updates help escape inference collapse.

5. RELATION TO RELATED WORK:

Posterior collapse in VAEs is first detailed in (Bowman et al., 2016) where they combine a LSTM decoder with VAE for text modeling. They interpret this problem from a regularization perspective, and propose the "KL cost annealing" method to address this issue, whereby the weight of KL term between approximate posterior and prior increases from a small value to one in a "warm-up" period. This method has been shown to be unable to deal with collapse on complex text datasets with very large LSTM decoders (Yang et al., 2017; Kim et al., 2018) . Many works follow this line to lessen the effect of KL term such as β-VAE (Higgins et al., 2017) that treats the KL weight as a hyperparameter or "free bits" method that constrains the minimum value of the KL term. Our approach differs from these methods in that we do not change ELBO objective during training and are in principle still performing maximum likelihood estimation. While these methods explicitly encourage the use of the latent variable, they may implicitly sacrifice density estimation performance at the same time, as we will discuss in Section 6.

Another thread of research focuses on a different problem called the "amortization gap" (Cremer et al., 2018) , which refers to the difference of ELBO caused by parameter sharing of the inference network. Some approaches try to combine instance-specific variational inference with amortized variational inference to narrow this gap (Hjelm et al., 2016; Krishnan et al., 2018; Kim et al., 2018; Marino et al., 2018 ). The most related example is SA-VAE (Kim et al., 2018) , which mixes instance-specific variational inference and empirically avoids posterior collapse. Our approach is much simpler without sacrificing performance, yet achieves an average of 5x training speedup.

Other attempts to address posterior collapse include proposing new regularizers (Zhao et al., 2017; Tolstikhin et al., 2018; Phuong et al., 2018) , deploying less powerful decoders (Yang et al., 2017; Semeniuta et al., 2017) , using lossy input (Chen et al., 2017) , utilizing different latent variable connections (Dieng et al., 2017; 2018; Park et al., 2018) , or changing the prior (Tomczak & Welling, 2018; Xu & Durrett, 2018) . (Yang et al., 2017) ≤332.1 10.0 --≤359.1 7.6 --SA-VAE + anneal (Kim et al., 2018) ≤327.5 7.19 ------

Modified VAE Objective VAE + anneal 328.6 (0.0) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) 357.9 (0.1) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) β-VAE (β = 0.2) 332.2 (0.6) 19.1 (1.5) 3.3 (0.1) 20.4 (6.8) 360.7 (0.7) 11.7 (2.4) 3.0 (0.5) 10.0 (5.9) β-VAE (β = 0.4) 328.7 (0.1) 6.3 (1.7) 2.8 (0.6) 8.0 (5.2) 358.2 (0.3) 4.2 (0.4) 2.0 (0.3) 4.2 (3.8) β-VAE (β = 0.6) 328.5 (0.1) 0.3 (0.2) 0.2 (0.1) 1.0 (0.7) 357.9 (0.1) 0.2 (0.2) 0.1 (0.1) 3.8 (2.9) β-VAE (β = 0.8) 328.8 (0.1) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) 358.1 (0.2) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) SA-VAE + anneal 327.2 (0.2) 5.2 (1.4) 2.7 (0.5) 9.8 (1.3) 355.9 (0.1) 2.8 (0.5) 1.7 (0.3) 8.4 (0.9) Ours + anneal 326.7 (0.1) 5.7 (0.7) 2.9 (0.2) 15.0 (3.5) 355.9 (0.1) 3.8 (0.2) 2.4 (0.1) 11.3 (1.0)

Standard VAE Objective LSTM-LM * 328.0 (0.3) - - - 358.1 (0.6) - - - VAE

329.0 (0.1) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) 358.3 (0.2) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) SA-VAE 329.2 (0.2) 0.1 (0.0) 0.1 (0.0) 0.8 (0.4) 357.8 (0.2) 0.3 (0.1) 0.3 (0.0) 1.0 (0.0) Ours 328.2 (0.2) 5.6 (0.2) 3.0 (0.0) 8.0 (0.0) 356.9 (0.2) 3.4 (0.3) 2.4 (0.1) 7.4 (1.3) 

6. EXPERIMENTS:

Our experiments below are designed to (1) examine whether the proposed method indeed prevents posterior collapse, (2) test its efficacy with respect to maximizing predictive log-likelihood compared to other existing approaches, and (3) test its training efficiency.

6.1. SETUP:

For all experiments we use a Gaussian prior N (0, I) and the inference network parametrizes a diagonal Gaussian. We evaluate with approximate negative log likelihood (NLL) as estimated by 500 importance weighted samples5 (Burda et al., 2016) since it produces a tighter lower bound to marginal data log likelihood than ELBO (ELBO values are included in Appendix C), and should be more accurate. We also report D KL (q φ (z|x) p(z)) (KL), mutual information I q (MI), and number of active units (AU) (Burda et al., 2016) in latent representation. The activity of a latent dimension z is measured as A z = Cov x (E z∼q(z|x) [z] ). The dimension z is defined as active if A z > 0.01. (c) Yelp Figure 4 : NLL versus AU (active units) for all models on three datasets. For each model we display 5 points which represent 5 runs with different random seeds. "Autoregressive" denotes LSTM-LM for text data and PixelCNN for image data. We plot "autoregressive" baselines as their AU is 0. To better visualize the system difference on OMNIGLOT dataset, for OMNIGLOT figure we ignore some β-VAE baselines that are not competitive.

As baselines, we compare with strong neural autoregressive models (LSTM-LM for text and Pixel-CNN (van den Oord et al., 2016) for images), basic VAE, the "KL cost annealing" method (Bowman et al., 2016; Sønderby et al., 2016) , β-VAE (Higgins et al., 2017) , and SA-VAE (Kim et al., 2018) which holds the previous state-of-the-art performance on text modeling benchmarks. For β-VAE we vary β between 0.2, 0.4, 0.6, and 0.8. SA-VAE is ran with 10 refinement steps. We also examine the effect of KL cost annealing on both SA-VAE and our approach. To facilitate our analysis later, we report the results in two categories: "Standard VAE objectives", and "Modified VAE objectives". 6We evaluate our method on density estimation for text on the Yahoo and Yelp corpora (Yang et al., 2017) and images on OMNIGLOT (Lake et al., 2015) . Following the same configuration as in Kim et al. (2018) , we use a single layer LSTM as encoder and decoder for text. For images, we use a ResNet (He et al., 2016) encoder and a 13-layer Gated PixelCNN (van den Oord et al., 2016) decoder. We use 32-dimensional z and optimize ELBO objective with SGD for text and Adam (Kingma & Ba, 2015) for images. We concatenate z to the input for the decoders. For text, z also predicts the initial hidden state of the LSTM decoder. We dynamically binarize images during training and test on fixed binarized test data. We run all models with 5 different random restarts, and report mean and standard deviation. Full details of the setup are in Appendix B.2 and B.3.

6.2. RESULTS:

In Table 1 and Table 2 we show the results on all three datasets, we also plot NLL vs AU for every trained model from separate runs in Figure 4 to visualize the uncertainties. Our method achieves comparable or better performance than previous state-of-the-art systems on all three datasets. Note that to examine the posterior collapse issue for images we use a larger PixelCNN decoder than previous work, thus our approach is not directly comparable to them and included at the top of Table 2 as reference points. We observe that SA-VAE suffers from posterior collapse on both text datasets without annealing. However, we demonstrate that our algorithm does not experience posterior collapse even without annealing.

6.3. TRAINING TIME:

In Table 3 we report the total training time of our approach, SA-VAE and basic VAE training across the three datasets. We find that the training time for our algorithm is only 2-3 times slower than a regular VAE whilst being 3-7 times faster than SA-VAE. We analyze the difference between our approach and the methods that weaken the KL regularizer term in ELBO, and explain the unwanted behavior produced by breaking maximum likelihood estimation.

As illustrative examples, we compare with the KL cost annealing method and β-VAE.

Decreasing the weight of the KL regularizer term in ELBO is equivalent to adding an additional regularizer to push q φ (z|x) far from p(z). We set β = 0.2 in order to better observe this phenomenon.

We investigate the training procedure on the Yelp dataset based on: (1) the mutual information between z and x, I q , (2) the KL regularizer, E x∼p d (x) [D KL (q φ (z|x) p(z))], and (3) the distance between the aggregated posterior and the prior, D KL (q φ (z) p(z)). Note that the KL regularizer is equal to the sum of the other two as stated in Eq.5. We plot these values over the course of training in Figure 5 . In the initial training stage we observe that the KL regularizer increases with all three approaches, however, the mutual information, I q , in the annealing remains small, thus a large KL regularizer term does not imply that the latent variable is being used. Finally the annealing method suffers from posterior collapse. For β-VAE, the mutual information increases, but D KL (q φ (z) p(z)) also reaches a very large value. Intuitively, D KL (q φ (z) p(z)) should be kept small for learning the generative model well since in the objective the generator p θ (x|z) is learned with latent variables sampled from the variational distribution. If the setting of z that best explains the data has a lower likelihood under the model prior, then the overall model would fit the data poorly. The same intuition has been discussed in Zhao et al. (2017) and Tolstikhin et al. (2018) . This also explains why β-VAE generalizes poorly when it has large mutual information. In contrast, our approach is able to obtain high mutual information, and at the same time maintain a small D KL (q φ (z) p(z)) as a result of optimizing standard ELBO where the KL regularizer upper-bounds D KL (q φ (z) p(z)).

6.5. ANALYSIS OF INNER LOOP UPDATE:

We perform analysis to examine the tradeoff between performance and speed within the inner loop update in our approach, through fixing a budget of updates to the inference network instead of updat-ing until convergence. 7 In our implementation, we break the inner loop when the ELBO objective stays the same or decreases across 10 iterations. Note that we do not perform separate learning rate decay in the inner loop so this convergence condition is not strict, but empirically we found it to be sufficient. Across all datasets, in practice this yields roughly 30 -100 updates per inner loop update. Now we explore using a fixed budget of inner loop updates and observe its influence on performance and speed. We report the results on Yelp dataset from single runs in Table 4 . 8 We see that sufficient number of inner iterations is necessary to address posterior collapse and achieve good performance, but the performance starts to saturate near convergence, thus we believe that optimizing to a near-convergence point is important.

7. CONCLUSION:

In this paper we study the "posterior collapse" problem that variational autoencoders experience when the model is parameterized by a strong autoregressive neural network. In our synthetic experiment we identify that the problem lies with the lagging inference network in the initial stages of training. To remedy this, we propose a simple yet effective training algorithm that aggressively optimizes the inference network with more updates before reverting back to basic VAE training.

Experiments on text and image modeling demonstrate the effectiveness of our approach.

A APPROXIMATION OF THE MEAN OF THE TRUE MODEL POSTERIOR:

We approximate the mean of true model posterior p θ (z|x) by discretization of the density distribution (Riemann integral):

EQUATION

where C is a partition of an interval with small stride and sufficiently large coverage. We assume the density value outside this interval is zero. The model posterior, p θ (z|x), needs to be first approximated on this partition of interval. In practice, for the synthetic data we choose the interval [-20.0, 20.0] and stride equal to 0.01. This interval should have sufficient coverage since we found all samples from true model posterior p θ (z|x) lies within [-5.0, 5.0] by performing MH sampling.

B EXPERIMENTAL DETAILS:

In general, for annealing we increase the KL weight linearly from 0.1 to 1.0 in the first 10 epochs, as in Kim et al. (2018) . We also perform analysis for different annealing strategies in Appendix E

B.1 SYNTHETIC EXPERIMENT FOR SECTION 3 AND 4:

To generate synthetic data points, we first sample a two-dimensional latent variable z from a mixture of Gaussian distributions that have four mixture components. We choose dimension two because we want the synthetic data distribution to be relatively simple but also complex enough for a onedimensional latent variable model to fit. We choose mixture of Gaussian as the prior to make sure that the synthetic data is diverse. The mean of these Gaussians are (-2.0, -2.0), (-2.0, 2.0), (2.0, -2.0), (2.0, 2.0), respectively. All of them have a unit variance. Then we follow the synthetic data generation procedure in Kim et al. (2018) , where we sample data points from an one-layer LSTM conditioned on latent variables. The LSTM has 100 hidden units and 100-dimensional input embeddings. An affine transformation of z is used as the initial hidden state of LSTM decoder, z is also concatenated with output of LSTM at each time stamp to be directly mapped to vocabulary space. LSTM parameters are initialized with U(-1, 1), and the part of MLP that maps z to vocabulary space is initialized with U(-5, 5), this is done to make sure that the latent variables have more influence in generating data. We generated a dataset with 20,000 examples (train/val/test is 16000/2000/2000) each of length 10 from a vocabulary of size 1000.

In the synthetic experiment we use a LSTM encoder and LSTM decoder, both of which have 50 hidden units and 50 latent embeddings. This LSTM decoder has less capacity than the one used for creating the dataset since in the real world model capacity is usually insufficient to exactly model the empirical distribution. Parameters of LSTM decoders are initialized with U(-0.01, 0.01), except for the embedding weight matrix which is initialized with U(-0.1, 0.1). Dropout layers with probability 0.5 are applied to both input embeddings and output hidden embeddings of decoder. We use the SGD optimizer and start with a learning rate of 1.0 and decay it by a factor of 2 if the validation loss has not improved in 2 epochs and terminate training once the learning rate has decayed a total of 5 times.

B.2 TEXT:

Following Kim et al. (2018) , we use a single-layer LSTM with 1024 hidden units and 512dimensional word embeddings as the encoder and decoder for all of text models. The LSTM parameters are initialized from U(-0.01, 0.01), and embedding parameters are initialized from U(-0.1, 0.1). We use the final hidden state of the encoder to predict (via a linear transformation) the latent variable. We use the SGD optimizer and start with a learning rate of 1.0 and decay it by a factor of 2 if the validation loss has not improved in 2 epochs and terminate training once the learning rate has decayed a total of 5 times. We don't perform any text preprocessing and use the datasets as provided. We follow Kim et al. (2018) and use dropout of 0.5 on the decoder for both the input embeddings of the decoder and on the output of the decoder before the linear transformation to vocabulary space.

B.3 IMAGES:

We use the same train/val/test splits as provided by Kim et al. (2018) . We use the Adam optimizer and start with a learning rate of 0.001 and decay it by a factor of 2 if the validation loss has not improved in 20 epochs. We terminate training once the learning rate has decayed a total of 5 times. Inputs were dynamically binarized throughout training by viewing the input as Bernoulli random variables that are sampled from pixel values. We validate and test on a fixed binarization and our decoder uses binary likelihood. Our ResNet is the same as used by Chen et al. (2017) . Our 13-layer PixelCNN architecture is a larger variant based on what was used in Kim et al. (2018) and described in their Appendix B.3 section. The PixelCNN has five 7 x 7 layers, followed by, four 5 x 5 layers, and then four 3 x 3 layers. Each layer has 64 feature maps. We use batch normalization followed by an ELU activation before our final 1 x 1 convolutional layer and sigmoid nonlinearity.

C ADDITIONAL RESULTS CONTAINING ELBO 328.6 (0.0) 328.8 (0.0) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) 357.9 (0.1) 358.1 (0.1) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) β-VAE (β = 0.2) 332.2 (0.6) 335.9 (0.8) 19.1 (1.5) 3.3 (0.1) 20.4 (6.8) 360.7 (0.7) 363.2 (1.1) 11.7 (2.4) 3.0 (0.5) 10.0 (5.9) β-VAE (β = 0.4) 328.7 (0.1) 330.2 (0.4) 6.3 (1.7) 2.8 (0.6) 8.0 (5.2) 358.2 (0.3) 359.1 (0.3) 4.2 (0.4) 2.0 (0.3) 4.2 (3.8) β-VAE (β = 0.6) 328.5 (0.1) 328.9 (0.0) 0.3 (0.2) 0.2 (0.1) 1.0 (0.7) 357.9 (0.1) 358.2 (0.1) 0.2 (0.2) 0.1 (0.1) 3.8 (2.9) β-VAE (β = 0.8) 328.8 (0.1) 329.0 (0.1) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) 358.1 (0.2) 358.3 (0.2) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) SA-VAE + anneal 327.2 (0.2) 327.8 (0.2) 5.2 (1.4) 2.7 (0.5) 9.8 (1.3) 355.9 (0.1) 356.2 (0.1) 2.8 (0.5) 1.7 (0.3) 8.4 (0.9) Ours + anneal 326.7 (0.1) 328.4 (0.2) 5.7 (0.7) 2.9 (0.2) 15.0 (3.5) 355.9 (0.1) 357.2 (0.1) 3.8 (0.2) 2.4 (0.1) 11. 

D UNCERTAINTY OF EVALUATION:

To measure the uncertainty in the evaluation stage caused by random Monte Carlo samples, we load pre-trained VAE models trained by our approach and basic VAE training, and repeat our evaluation process with 10 different random seeds. We report the mean and variance values in Table 7 and Table 8 . 1 and Table 2 , we implement annealing as increasing KL weight linearly from 0.1 to 1.0 in the first 10 epochs following (Kim et al., 2018) , and observed posterior collapse for KL-annealing method. However, this annealing strategy may not be the optimal. In this section, we explore different KL-annealing schedules. Specifically, we increase KL weight linearly from 0.0 to 1.0 in the first s iterations, and s is varied as 30k, 50k, 100k, and 120k. We report results on three datasets in Table 9 . The results indicate that KL-annealing does not experience posterior collapse if the annealing procedure is sufficiently slow, but it does not produce superior predictive log likelihood to our approach, which is expected because a very slow annealing schedule resembles β-VAE training in the first many epochs, and β-VAE encourages learning latent representations but might sacrifice generative modeling performance, as we already showed in Table 1 and Table 2 . Also, the optimal KL annealing schedule varies with different datasets and model architectures, so that it requires careful tuning for the task at hand. The lagging behavior of inference networks observed in Section 3 might be caused by different magnitude of gradients of encoder and decoder9 , thus another simpler possible solution to this problem is to use separate learning rates for the encoder and decoder optimization. Here we report the results of our trial by using separate learning rates. We experiment with the Yelp dataset, and keep the decoder optimization the same as discussed before, but vary the encoder learning rates to be 1x, 2x, 4x, 6x, 8x, 10x, 30x, 50x of the decoder learning rate. We notice that training becomes very unstable

:

Code and data are available at https://github.com/jxhe/vae-lagging-encoder.

µ x,θ can be approximated through discretization of the model posterior, which we show in Appendix A.

Note that the converse is not true: the setting where all points are located at the origin may not be a local optimum. For example when a model is initialized at the origin as we show in Section 3.2.

This LSTM decoder has less capacity than the one used for creating the dataset since in real world model capacity is usually insufficient to exactly model the empirical distribution.

In Appendix G we also study the setting where the points are not initialized at origin.

We measure the uncertainty in the evaluation caused by the Monte Carlo estimates in Appendix D. The variance of our NLL estimates for a trained VAE model is smaller than 10 -3 on all datasets.

While annealing reverts back to ELBO objective after the warm-up period, we consider part of "Modified VAE objectives" since it might produce undesired behavior in the warm-up period, as we will discuss soon.

Note that in practice, we never reach exact convergence, thus here we aim to show how close to convergence is required in the inner loop update.

70 inner iterations take less time than 50 because the aggressive training of it stops earlier in practice.

In the experiments, we did observe that the gradients of decoder is much larger than the gradients of encoder.

Figures and Tables:

j Z r A W 2 h b V R B d X S A j t E Z a q A m o u g B P a F n 9 O I 8 O q / O u / M x k k 4 5 4 5 l N N B H O 5 z d M W 6 s w < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " W b l 6 / d 1 + N S e 4 D 7 I 5 N

Figure2: The projections of 500 data samples on the posterior mean space. From left to right we show the change of these points at four different time stamps in the course of training. The top row is from the basic VAE training, the bottom row is from our approach. [Can you make the font slightly bigger? -GN] [Can you also add labels indicating the four training stages ("Stage 1", "Stage 2", ..., or even better "X Epochs", "X Epochs", ...). -GN] [Add labels to the top of "Vanilla VAE" and "Our Approach" (or whatever alternative name you come up with.) -GN]

Figure 5: Training behavior on Yelp. Left: VAE + annealing. Middle: Our method. Right: β-VAE (β = 0.2).

e x i t > Figure 2: The projections of 500 data samples from a synthetic dataset on the posterior mean space over the course of training. "iter" denotes the number of updates of generators. The top row is from the basic VAE training, the bottom row is from our aggressive inference network training. The results show that while the approximate posterior is lagging far behind the true model posterior in basic VAE training, our aggressive training approach successfully moves the points onto the diagonal line and away from inference collapse.

Table content: <table/>

Results on Yahoo and Yelp datasets. We report mean values across 5 different random restarts, and standard deviation is given in parentheses when available. For LSTM-LM * we report the exact negative log likelihood.

Table content: <table><tr><td>Yahoo</td><td>Yelp</td></tr></table>

Results on OMNIGLOT dataset. We report mean values across 5 different random restarts, and standard deviation is given in parentheses when available. For PixelCNN * we report the exact negative log likelihood.

Table content: <table><tr><td>Model</td><td>NLL</td><td>KL</td><td>MI</td><td>AU</td></tr><tr><td/><td>Previous Reports</td><td/><td/><td/></tr><tr><td>VLAE (Chen et al., 2017)</td><td>89.83</td><td>-</td><td>-</td><td>-</td></tr><tr><td>VampPrior (Tomczak &amp; Welling, 2018)</td><td>89.76</td><td>-</td><td>-</td><td>-</td></tr><tr><td colspan="3">Modified VAE Objective</td><td/><td/></tr><tr><td>VAE + anneal</td><td>89.21 (0.04)</td><td>1.97 (0.12)</td><td colspan="2">1.79 (0.11) 5.3 (1.0)</td></tr><tr><td>β-VAE (β = 0.2)</td><td colspan="4">105.96 (0.38) 69.62 (2.16) 3.89 (0.03) 32.0 (0.0)</td></tr><tr><td>β-VAE (β = 0.4)</td><td colspan="4">96.09 (0.36) 44.93 (12.17) 3.91 (0.03) 32.0 (0.0)</td></tr><tr><td>β-VAE (β = 0.6)</td><td>92.14 (0.12)</td><td colspan="3">25.43 (9.12) 3.93 (0.03) 32.0 (0.0)</td></tr><tr><td>β-VAE (β = 0.8)</td><td>89.15 (0.04)</td><td>9.98 (0.20)</td><td colspan="2">3.84 (0.03) 13.0 (0.7)</td></tr><tr><td>SA-VAE + anneal</td><td>89.07 (0.06)</td><td>3.32 (0.08)</td><td colspan="2">2.63 (0.04) 8.6 (0.5)</td></tr><tr><td>Ours + anneal</td><td>89.11 (0.04)</td><td>2.36 (0.15)</td><td colspan="2">2.02 (0.12) 7.2 (1.3)</td></tr><tr><td colspan="3">Standard VAE Objective</td><td/><td/></tr><tr><td>PixelCNN  *</td><td>89.73 (0.04)</td><td>-</td><td>-</td><td>-</td></tr><tr><td>VAE</td><td>89.41 (0.04)</td><td>1.51 (0.05)</td><td colspan="2">1.43 (0.07) 3.0 (0.0)</td></tr><tr><td>SA-VAE</td><td>89.29 (0.02)</td><td>2.55 (0.05)</td><td colspan="2">2.20 (0.03) 4.0 (0.0)</td></tr><tr><td>Ours</td><td>89.05 (0.05)</td><td>2.51 (0.14)</td><td colspan="2">2.19 (0.08) 5.6 (0.5)</td></tr></table>

Comparison of total training time, in terms of relative speed and absolute hours.

Table content: <table><tr><td/><td>Yahoo</td><td/><td colspan="2">Yelp15</td><td colspan="2">OMNIGLOT</td></tr><tr><td/><td colspan="6">Relative Hours Relative Hours Relative Hours</td></tr><tr><td>VAE</td><td>1.00</td><td>5.35</td><td>1.00</td><td>5.75</td><td>1.00</td><td>4.30</td></tr><tr><td>SA-VAE</td><td>9.91</td><td>52.99</td><td>10.33</td><td>59.37</td><td>15.15</td><td>65.07</td></tr><tr><td>Ours</td><td>2.20</td><td>11.76</td><td>3.73</td><td>21.44</td><td>2.19</td><td>9.42</td></tr></table>

Results on Yelp dataset using a fixed budget of inner encoder updates

Table content: <table><tr><td colspan="4"># Inner Iterations NLL KL MI AU Hours</td></tr><tr><td>10</td><td>357.9 1.1 1.0</td><td>3</td><td>11.97</td></tr><tr><td>30</td><td>357.1 3.6 2.5</td><td>8</td><td>22.31</td></tr><tr><td>50</td><td>356.9 4.2 2.8</td><td>9</td><td>29.58</td></tr><tr><td>70</td><td colspan="2">357.1 4.4 2.7 10</td><td>24.18</td></tr><tr><td>convergence</td><td>357.0 3.8 2.6</td><td>8</td><td>21.44</td></tr><tr><td>6.4 ANALYSIS OF BASELINES</td><td/><td/><td/></tr></table>

Results on Yahoo and Yelp datasets. We report mean values across 5 different random restarts, and standard deviation is given in parentheses when available. For LSTM-LM * we report the exact negative log likelihood.

Table content: <table><tr><td/><td/><td>Yahoo</td><td/><td/><td/><td>Yelp</td><td/><td/><td/><td/></tr><tr><td>Model</td><td>IW</td><td>-ELBO</td><td>KL</td><td>MI</td><td>AU</td><td>IW</td><td>-ELBO</td><td>KL</td><td>MI</td><td>AU</td></tr><tr><td/><td/><td/><td colspan="2">Previous Reports</td><td/><td/><td/><td/><td/><td/></tr><tr><td>CNN-VAE (Yang et al., 2017)</td><td>-</td><td>332.1</td><td>10.0</td><td>-</td><td>-</td><td>-</td><td>359.1</td><td>7.6</td><td>-</td><td>-</td></tr><tr><td>SA-VAE + anneal (Kim et al., 2018)</td><td>-</td><td>327.5</td><td>7.19</td><td>-</td><td>-</td><td>-</td><td>-</td><td>-</td><td>-</td><td>-</td></tr><tr><td/><td/><td/><td colspan="3">Modified VAE Objective</td><td/><td/><td/><td/><td/></tr><tr><td>VAE + anneal</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/></tr></table>

Results on OMNIGLOT dataset. We report mean values across 5 different random restarts, and standard deviation is given in parentheses when available. For PixelCNN * we report the exact negative log likelihood.

Table content: <table><tr><td>Model</td><td>IW</td><td>-ELBO</td><td>KL</td><td>MI</td><td>AU</td></tr><tr><td/><td colspan="2">Previous Reports</td><td/><td/><td/></tr><tr><td>VLAE (Chen et al., 2017)</td><td>89.83</td><td>-</td><td>-</td><td>-</td><td/></tr><tr><td>VampPrior (Tomczak &amp; Welling, 2018)</td><td>89.76</td><td>-</td><td>-</td><td>-</td><td/></tr><tr><td/><td colspan="2">Modified VAE Objective</td><td/><td/><td/></tr><tr><td>VAE + anneal</td><td>89.21 (0.04)</td><td>89.55 (0.04)</td><td>1.97 (0.12)</td><td colspan="2">1.79 (0.11) 5.3 (1.0)</td></tr><tr><td>β-VAE (β = 0.2)</td><td colspan="5">105.96 (0.38) 113.24 (0.40) 69.62 (2.16) 3.89 (0.03) 32.0 (0.0)</td></tr><tr><td>β-VAE (β = 0.4)</td><td colspan="5">96.09 (0.36) 101.16 (0.66) 44.93 (12.17) 3.91 (0.03) 32.0 (0.0)</td></tr><tr><td>β-VAE (β = 0.6)</td><td>92.14 (0.12)</td><td>94.92 (0.47)</td><td colspan="3">25.43 (9.12) 3.93 (0.03) 32.0 (0.0)</td></tr><tr><td>β-VAE (β = 0.8)</td><td>89.15 (0.04)</td><td>90.17 (0.06)</td><td>9.98 (0.20)</td><td colspan="2">3.84 (0.03) 13.0 (0.7)</td></tr><tr><td>SA-VAE + anneal</td><td>89.07 (0.06)</td><td>89.42 (0.06)</td><td>3.32 (0.08)</td><td colspan="2">2.63 (0.04) 8.6 (0.5)</td></tr><tr><td>Ours + anneal</td><td>89.11 (0.04)</td><td>89.62 (0.16)</td><td>2.36 (0.15)</td><td colspan="2">2.02 (0.12) 7.2 (1.3)</td></tr><tr><td/><td colspan="2">Standard VAE Objective</td><td/><td/><td/></tr><tr><td>PixelCNN  *</td><td>89.73 (0.04)</td><td>-</td><td>-</td><td>-</td><td>-</td></tr><tr><td>VAE</td><td>89.41 (0.04)</td><td>89.67 (0.06)</td><td>1.51 (0.05)</td><td colspan="2">1.43 (0.07) 3.0 (0.0)</td></tr><tr><td>SA-VAE</td><td>89.29 (0.02)</td><td>89.54 (0.03)</td><td>2.55 (0.05)</td><td colspan="2">2.20 (0.03) 4.0 (0.0)</td></tr><tr><td>Ours</td><td>89.05 (0.05)</td><td>89.52 (0.03)</td><td>2.51 (0.14)</td><td colspan="2">2.19 (0.08) 5.6 (0.5)</td></tr></table>

Evaluation of a trained VAE model trained by our approach across 10 different random seeds. Mean values are reported and variance is given in parentheses. IW denotes the approximation to NLL we used in Section 6.

Table content: <table><tr><td>Dataset</td><td>IW</td><td>-ELBO</td><td>KL</td><td>MI</td><td>AU</td></tr><tr><td>Yahoo Yelp OMNIGLOT</td><td colspan="5">5.35 (0) 357.03 (10 -5 ) 358.25 (2 × 10 -4 ) 3.82 (10 -5 ) 2.61 (0.003) 8 (0) 3.01 (0.002) 8 (0) 327.98 (10 -5 ) 329.54 (5 × 10 -4 ) 89.03 (0) 89.53 (3 × 10 -4 ) 2.54 (0) 2.21 (0.001) 6 (0)</td></tr></table>

Evaluation of a trained VAE model trained by basic VAE training across 10 different random seeds.Mean values are reported and variance is given in parentheses. IW denotes the approximation to NLL we used in Section 6.

Table content: <table><tr><td>Dataset</td><td>IW</td><td>-ELBO</td><td>KL</td><td>MI</td><td>AU</td></tr><tr><td colspan="6">Yahoo Yelp OMNIGLOT 89.41 (0) 328.85 (0) 329.54 (1 × 10 -5 ) 0.00 (0) 358.17 (0) 358.38 (3 × 10 -5 ) 0.00 (0) 89.66 (2 × 10 -4 ) 1.48 (0) 1.39 (6 × 10 -4 ) 3 (0) 0.00 (0) 0 (0) 0.00 (0) 0 (0)</td></tr><tr><td colspan="5">E COMPARISON WITH DIFFERENT KL-ANNEALING SCHEDULES</td><td/></tr><tr><td colspan="2">For the annealing baseline in Table</td><td/><td/><td/><td/></tr></table>

Results on Yahoo and Yelp datasets, with different annealing schedules. Starred entries represent original annealing strategy.

Table content: <table><tr><td/><td>Yahoo</td><td/><td>Yelp</td><td/><td colspan="2">OMNIGLOT</td></tr><tr><td>Model</td><td colspan="5">NLL KL MI AU NLL KL MI AU NLL KL</td><td colspan="2">MI AU</td></tr><tr><td>VAE + anneal (30k)</td><td>328.4 0.0 0.0</td><td>0</td><td>357.9 0.2 0.2</td><td>1</td><td colspan="3">89.18 2.54 2.19 10</td></tr><tr><td>VAE + anneal (50k)</td><td>328.3 0.7 0.7</td><td>4</td><td>357.7 0.3 0.3</td><td>1</td><td colspan="3">89.15 3.18 2.58 10</td></tr><tr><td colspan="4">VAE + anneal (100k) 327.5 4.3 2.6 12 356.8 1.9 1.2</td><td>5</td><td colspan="3">89.27 4.04 2.97 16</td></tr><tr><td colspan="4">VAE + anneal (120k) 327.5 7.8 3.2 18 356.9 2.7 1.8</td><td>6</td><td colspan="3">89.32 4.12 3.00 15</td></tr><tr><td>VAE + anneal  *</td><td>328.6 0.0 0.0</td><td>0</td><td>358.0 0.0 0.0</td><td>0</td><td colspan="2">89.20 2.11 1.89</td><td>5</td></tr><tr><td>Ours + anneal  *</td><td colspan="6">326.6 6.7 3.2 15 355.9 3.7 2.3 10 89.13 2.53 2.16</td><td>8</td></tr><tr><td>Ours</td><td>328.0 5.4 3.0</td><td>8</td><td>357.0 3.8 2.6</td><td>8</td><td colspan="2">89.03 2.54 2.20</td><td>6</td></tr><tr><td colspan="8">F SEPARATE LEARNING RATES OF INFERENCE NETWORK AND GENERATOR</td></tr></table>

Peer review (iclr2019):
 Response to Authors ------------- I*ve read all other reviews and the author responses. Most responses to my issues seem to be "we will run more experiments", so my review scores haven*t changed. I*m glad the authors are planning many revised experiments, and I understand that these take time. It*s too bad revised results won*t be available before the review revision deadline (tomorrow 11/26). I guess I*m willing to take the author*s promises to update in good faith. Thus, I think this is an "accept", but only if the authors really do follow through on promises to add uncertainty quantification and include some complete comparisons to KL annealing strategies. Review Summary -------------- Overall, I think the paper offers a reasonable story for why its proposed innovation -- an alternative scheduling of parameter-specific updates where encoder parameters are always trained to convergence during early iterations -- might offer a reliable way to avoid posterior collapse that is far faster and easier-to-implement than other options that require some per-example iterations (e.g. semi-amortized VAE). My biggest concerns are that relative performance gains (in bound quality) over alternatives are not too large and hard to judge as significant because no uncertainty in these estimates is quantified. Additionally, I*d like to see more careful evaluation of the KL annealing baseline and more attention to within-model comparisons (do you really need to update until convergence?). Given the method*s simplicity and speed, I think with a satisfactory rebuttal and plan for revision I would lean towards acceptance. Paper Summary ------------- The paper investigates a common problem known as "posterior collapse" observed when training generative models such as VAEs (Kingma & Welling 2014) with high-capacity neural networks. Posterior collapse occurs when the encoder distribution q(z|x) (parameterized by a NN) becomes indistinguishable from the generative prior on codes p(z), which is often a local optima of the VI ELBO objective. While other better fixed points exist, once this one is reached during optimization it is hard to escape using the typical local gradient steps for VAEs that jointly update the parameters of an encoder and a decoder with each gradient step. The proposed solution (presented in Alg. 1) is to avoid joint gradient updates early in training, and instead use an alternating update scheme where after each single-gradient-step decoder parameter update, the encoder is updated with as many gradient steps as are needed to reach convergence. This proposed scheme, which the paper terms "aggressive updates", forces the encoder to better approximate the true posterior p(z|x) at each step. Experiments study a synthetic task where visualizing the evolution of true posterior mean of p(z|x) side-by-side with approximate q(z|x) is possible in 2D, as well as benchmark comparisons to several other methods that address posterior collapse on text modeling (Yahoo, Yelp15) and image modeling (Omniglot). Studied baselines include annealing the KL term in the VI objective, the .beta VAE (which keeps the KL term fixed with a weight .beta), and semi-amortized VAEs (SA-VAEs, Kim et al. 2018). The presented approach is said to reach better values of the log likelihood while also being ~10x faster to train than the Kim et al. approach on large datasets. Significance and Originality ---------------------------- There exists strong interest in deploying amortized VI to fit sophisticated models efficiently while avoiding posterior collapse, so the topic is definitely relevant to ICLR. Certainly solutions to this issue are welcome, though I worry with the crowded field that performance is starting to saturate and it is becoming hard to identify significant vs. marginal contributions. Thus it*s important to interpret results across multiple axes (e.g. speed and heldout likelihood). The paper does a nice job of highlighting related work on this problem, and I*d rate its methodological contributions as clearly distinct from prior work, even though the eventual procedure is simple. The closest related works in my view are: * Krishnan et al. AISTATS 2018, where VAE joint-training algorithms for nonlinear factor analysis problems are shown to be improved by an algorithm that uses the encoder NN as an *initialization* and then doing several standard SVI updates to refine per-example parameters. Encoder parameters are updated via gradient updates, *after* the decoder parameters are updated (not jointly). * SA-VAEs (Kim et al. ICML 2018) which studies VAEs for deep text models and develops an algorithm that at each a new batch uses the encoder to initialize per-example parameters, updates these via several iterations of SVI, then *backpropagates* through those updates to compute a gradient update of the encoder NN. Compared to these, the detailed algorithm presented in this work is both distinct and simpler. It does not require any per-example parameter updates, instead it only requires a different scheduling of when encoder and decoder NN updates occur. Concerns about Technical Quality (prioritized) ---------------------------------------------- ## C1: Without error bars in Table 1 and 3, hard to know which gaps are significant Are 500 Monte Carlo samples enough to be sure that the numbers reported in Table 1 are precise estimates and not too noisy? How much error is there in the estimation of various quantities like the NLL or the KL if we repeated 500-MC samples 5x or 10x or 25x? My experience is that even with 100 or more samples, evaluations of the ELBO bound for classic VAEs can differ non-trivally. I*d like to see evidence that these quantities are estimated with certainty, or (even better) some direct reporting of the uncertainties across several estimates. ## C2: Baseline comparison to KL annealing needs to be more thorough The current paper dismisses the strategy that annealing the KL term as ineffective in addressing posterior collapse (e.g. VAE + anneal has a 0.0 KL term in Table 1). However, it*s not clear that a reasonable annealing schedule was used, or even that any reasonable effort was made to try more than one schedule. For example, if we set the KL term to exactly 0.0 weight, the optimization has no incentive to push q towards the prior, and thus posterior collapse *cannot* occur. It may be that this leads to other problems, but it*s unclear to me why a schedule that keeps the KL term weight exactly at 0 for a few updates and then gradually increases the weight should lead to collapse. To me, the KL annealing story is much simpler than the presented approach and I think as a community we should invest in giving it a fair shot. If the answer is that annealing takes too long or the schedule is tough to tune, that*s sensible, but I think the claim that annealing still leads to collapse just means the schedule probably wasn*t set right. Notice that "Ours" is improved by "Ours+Annealing" for 2 datasets in Table 1. So annealing *can* be effective. Krishnan et al. 2018*s Supplementary Fig. 10 suggests that if annealing is slow enough (unfolding over 100000 updates instead of 10000 updates), then KL annealing will get close to pure SVI in effective, non-collapsed posterior approximation. The present paper*s Sec. B.3 indicates that the attempted annealing schedule was 0.1 to 1.0 linearly over 10 epochs with batch size 32 and train set size 100k, which sounds like only 30k updates of annealing were performed. I*d suggest comparing against KL annealing that both starts with a smaller weight (perhaps exactly at 0.0) and grows much slower. ## C3: Results do not analyze variability due to random initialization or random minibatch traversal Many factors can impact the final performance values of a model trained via VI, including the random initialization of its parameters and the random order of minibatches used during gradient updates. Due to local optima, often best practice is to take the best of many separate initializations (see several figures in Bishop*s PRML textbook). The present paper doesn*t make clear whether it*s reporting single runs or the best of many runs. I suggest a revision is needed to clarify. Quantifying robustness to initialization is important. ## C4: Results do not analyze relative sensitivity of encoder and decoder to using the same learning rate One possible explanation for "lagging" might be that the gradient vectors of the encoder and the decoder have different magnitudes, and thus using the same fixed learning rate for both (as seems to be done from a skim of Sec. B) might not be optimal. Perhaps a quick experiment that separately tunes learning rates of encoder and decoder is necessary? If the learning rate for encoder is too small, this could easily explain the lagging when using joint updates. ## C5: Is it necessary to update until convergence? Or would a fixed budget of 25 or 100 updates to the encoder suffice? In Alg. 1, during the "aggressive" phase the encoder is updated until convergence. I*d like to see some coverage of how long this typically takes (10 updates? 100 updates?). I*d also like to know if there are significant time-savings to be had by not going *all* the way to convergence. It*s concerning that in Fig. 1 convergence on a toy dataset takes more than 2000 iterations. ## C6: Sensitivity to the initialization of the encoder is not discussed and could matter In the synthetic example figure, it seems the encoder is initialized so that across many examples, the typical encoding will be near the origin and thus favored under the prior. Thus, the *initialization* is in some ways setting optimization up for posterior collapse. I wonder if some more diverse initialization might avoid the problem. Presentation comments --------------------- Overall the paper reads reasonably. I*d suggest mentioning the KL annealing comparison a bit earlier, but otherwise I have few complaints. I*m not sure I like the chosen terminology of "aggressive" update. The procedure is more accurately a "repeat-until-convergence" update. There*s nothing aggressive about it, it*s just repeated. Line-by-line Detailed comments ------------------------------ Citations for "traditional" VI with per-example parameters should go much further back than 2013. For example, Matthew Beal*s thesis, work by Blei in 2003 on LDA, or work by MacKay or M.I. Jordan or others even further back. Alg 1 Line 12: This update should be to .theta (model parameters), not .phi (approx posterior parameters). Alg 1: Might consider using notation like g_.theta to denote the grad. of specific parameters, rather than have the same symbol "g" overloaded as the gradient of .theta, .phi, and both in the same Algo. Fig. 3: This is interesting, but I think it*s missing something as a visualization of the algorithm. There*s nothing obvious visually that indicates the encoder update involves *many* steps, but the decoder update is only one step. I*d suggest at least turning each vertical arrow into *many* short arrows stacked end-to-end, indicating many steps. Also use a different color (not green for both). Fig. 4: Shows various quantities like KL(q, prior) traced over optimization. This figure would be more illuminating if it also showed the complete ELBO objective and the expected log likelihood term. Then it would be clear why annealing is failing to avoid posterior collapse. Table 1: How exactly is the negative log likelihood (NLL) computed? Is it the expected value of the data likelihood: -1 * E_q[log p(x|z)]? Or is it the variational lower bound on marginal likelihood?