All in the Exponential Family: Bregman Duality in Thermodynamic - - PowerPoint PPT Presentation

All in the Exponential Family: Bregman Duality in Thermodynamic Variational Inference

Rob Brekelmans Vaden Masrani Frank Wood Greg Ver Steeg Aram Galstyan

1

Thermodynamic Variational Objective

• TVO is a recent objective for training deep generative models
• Generalizes and tightens the ELBO
• qφ(z|x)

pθ(z|x) πβ

1) Masrani et. al. “The Thermodynamic Variational Objective”. NeurIPS 2019

πβ(z|x) ∝ qφ(z|x)1−βpθ(x, z)β

log pθ(x) =

1

Z Eπβ ⇥ log pθ(x, z) qφ(z|x) ⇤ dβ

1

Eπβ[log pθ(x, z) qφ(z|x) ]

βk−1

βk

TVOL(θ, φ, x)

β

≥

K−1

X

k=0

Eπβk ⇥ log pθ(x, z) qφ(z|x) ⇤ · ∆βk

2

Thermodynamic Variational Objective

• TVO is a recent objective for training deep generative models
• Generalizes and tightens the ELBO
• qφ(z|x)

pθ(z|x) πβ

1) Masrani et. al. “The Thermodynamic Variational Objective”. NeurIPS 2019

πβ(z|x) ∝ qφ(z|x)1−βpθ(x, z)β

log pθ(x) =

1

Z Eπβ ⇥ log pθ(x, z) qφ(z|x) ⇤ dβ

1

Eπβ[log pθ(x, z) qφ(z|x) ]

βk−1

βk

TVOL(θ, φ, x)

β

≥

K−1

X

k=0

Eπβk ⇥ log pθ(x, z) qφ(z|x) ⇤ · ∆βk

3 Eπβ[log pθ(x, z) qφ(z|x) ]

1

ELBO(θ, φ, x)

β

Problems with TVO

• Gap in TVO bounds previously unknown
• Choosing intermediate
• Log-uniform spacing, static across epochs
• Required grid search over

1

ELBO(θ, φ, x)

DKL[qφ(z|x)||pθ(z|x)]

{ βk }K

k=1

LELBO = log pθ(x) − DKL[qφ(z|x)||pθ(z|x)]

β1

1

TVOL(θ, φ, x)

βk−1

βk

???

LTVO = log pθ(x) − ???

4

Exponential Family Interpretation

πβ(z|x) = qφ(z|x)1−βpθ(x, z)β / Zβ(x)

πβ(z|x) = qφ(z|x) | {z }

Base (β=0)

exp{ β · log pθ(x, z) qφ(z|x) | {z }

Sufficient Statistics

− log Zβ(x)}

5

Our Contributions

• TVO Bound Gaps using Bregman Divergences
• Adaptively select using dual

parameters of exponential family

• Single β can notably improve upon ELBO

2) Grosse et. al. “Annealing between Distributions by Averaging Moments”. NeurIPS 2013

LTVO = log pθ(x) −

K

X

k=1

DKL[πβk−1(z|x)||πβk(z|x)]

1

TVOL(θ, φ, x)

βk

βk−1

DKL[πβk−1||πβk]

6

{ βk }K

k=1

Rest of this Talk

• Path Exponential Family
• Bregman Divergence intuition
• Moments Scheduling
• Results
• Future Directions

7

Path Exponential Family

8

Path Exponential Family

• Strictly convex in β

∝ qφ(z|x)⇠⇠⇠ ⇠ exp log ✓pθ(x, z) qφ(z|x) ◆β

∝ qφ(z|x)1−β pθ(x, z)β

πβ(z|x) = qφ(z|x) exp{ β · log pθ(x, z) qφ(z|x) − log Zβ(x)}

1

β2

β1

β3

log Zβ(x)

log Zβ(x) = log Z qφ(z|x)1−β pθ(x, z)β dz

= ) rβ log Zβ(x)

strictly increasing

9

Mean Parameters

1

β2

β1

β3

η1

η2

η3

log Zβ(x)

as a dual parameterization

= ) ηk := rβ log Zβk(x)

πβ(z|x) = qφ(z|x) exp { β · log pθ(x, z) qφ(z|x) | {z }

Sufficient Statistics

− log Zβ}

ηk := rβ log Zβk(x) = Eπβk ⇥ log pθ(x, z) qφ(z|x) ⇤

10

Thermodynamic Integration

• Consider endpoints
• Thermodynamic Integration / Fundamental Theorem of Calculus
• π1(z|x) ∝ pθ(x, z)

log Z1 = log pθ(x)

• π0(z|x) = qφ(z|x)

• log Z0 = 0

• πβ(z|x) = qφ(z|x)1−βpθ(x, z)β

Zβ(x)

log Z1 log Z0 =

1

Z rβ log Zβ(x) dβ

11

Thermodynamic Integration

• Consider endpoints
• Thermodynamic Integration / Fundamental Theorem of Calculus
• π1(z|x) ∝ pθ(x, z)

log Z1 = log pθ(x)

• π0(z|x) = qφ(z|x)

• log Z0 = 0

• πβ(z|x) = qφ(z|x)1−βpθ(x, z)β

Zβ(x)

log pθ(x) 0 =

1

Z rβ log Zβ(x) dβ

12

Thermodynamic Integration

• Consider endpoints
• Thermodynamic Integration / Fundamental Theorem of Calculus
• π1(z|x) ∝ pθ(x, z)

log Z1 = log pθ(x)

• π0(z|x) = qφ(z|x)

• log Z0 = 0

• πβ(z|x) = qφ(z|x)1−βpθ(x, z)β

Zβ(x)

13

log pθ(x) =

1

Z ηβ dβ

Thermodynamic Integration

• Consider endpoints
• Thermodynamic Integration / Fundamental Theorem of Calculus
• π1(z|x) ∝ pθ(x, z)

log Z1 = log pθ(x)

• π0(z|x) = qφ(z|x)

• log Z0 = 0

• πβ(z|x) = qφ(z|x)1−βpθ(x, z)β

Zβ(x)

14

log pθ(x) =

1

Z Eπβ ⇥ log pθ(x, z) qφ(z|x) ⇤ dβ

Thermodynamic Variational Objective

• Lower or Upper bounds via Riemann Sum approximations
• Since integrand is increasing,

ηβ = rβ log Zβ

Left Riemann Sum

1

βk−1

βk

TVOL(θ, φ, x)

ηβ

1

βk−1

βk

ηβ

log pθ(x)

log pθ(x)

≤

Right Riemann Sum

1

βk−1

βk

ηβ

TVOU(θ, φ, x)

≤

15

Thermodynamic Variational Objective

• Log likelihood as an integral
• Left Riemann Lower Bound
• Self-normalized importance sampling for each

1

βk−1

βk

TVOL(θ, φ, x)

ηβ

log pθ(x) =

1

Z Eπβ ⇥ log pθ(x, z) qφ(z|x) ⇤ | {z } dβ

log pθ(x) ≥

K−1

X

k=0

(βk+1 − βk) Eπβk ⇥ log pθ(x, z) qφ(z|x) ⇤

Eπβk ⇥ · ⇤

16

Bregman Divergence

17

• Let , which is convex
• We can then define the Bregman Divergence

Bregman Divergence

ψ(β) := log Zβ(x)

18

Dψ[βk+1 : βk] = ψ(βk+1) ψ(βk) (βk+1 βk) · rψ(βk)

• Let , which is convex
• We can then define the Bregman Divergence
• For , Bregman divergence = KL divergence

Bregman Divergence

ψ(β) := log Zβ(x)

Dψ[βk+1 : βk] = DKL[πβk||πβk+1]

log Zβ(x)

(w/arguments reversed)

19

1

ψ(β)

βk

βk+1

Dψ[βk+1 : βk]

Dψ[βk+1 : βk] = ψ(βk+1) ψ(βk) (βk+1 βk) · rβψ(βk) | {z }

First order Taylor Approximation

Gap in TVO Lower Bound

20

ψ(βk+1) ψ(βk) | {z }

• (βk+1 βk) · rβ ψ(βk)

| {z } = Dψ[βk+1 : βk] | {z }

Gap in TVO Lower Bound

Area under curve

βk+1

Z

βk

rβ log Zβ dβ

1

ηβ

βk

βk+1

21

log Zβk+1 log Zβk | {z }

• (βk+1 βk) · rβ log Zβk

| {z } = Dψ[βk+1 : βk] | {z }

DKL[πβk||πβk+1]

=

1

βk

ηβ

DKL[πβk||πβk+1]

βk+1

Left Riemann term

1

ηβ

βk

βk+1

Width

·

Height

−

(βk+1 − βk) · ηβk

Bregman Divergences in TVO

• TVO Lower Bound (Left Riemann):
• TVO Upper Bound (Right Riemann):
• Additional analysis:

LTVOU = log pθ(x) +

K−1

X

k=0

DKL[πβk+1(z|x)||πβk(z|x)]

LTVOL = log pθ(x) −

K−1

X

k=0

DKL[πβk(z|x)||πβk+1(z|x)]

22

• Symmetrized KL divergence
• Taylor Remainder theorem
• Renyi Divergence VI (Li and Turner 2016)

Moments Scheduling

23

Choosing β “Schedules”

• Budget of K intermediate points as a hyper parameter
• Goal: Assign more intermediate where integrand is changing quickly
• Shape of the TVO integrand related to posterior mismatch
• Adaptive choice of based on training progress

βk

{ βk }K

k=1

24

Moment Spacing Schedule

• Find to yield equal spacing in the mean parameters 2
• Legendre transform
• Efficient with SNIS, binary search
• Corresponds to Lebesgue integration

βk

ηk

2) Grosse et. al. “Annealing between Distributions by Averaging Moments”. NeurIPS 2013

1

= η0

rβψ(β)

ηk

βk

= ηK

non-uniform in uniform in

η

β

25

Results

26

Results

• TVO with single intermediate
• (2 term Riemann approx)
• Compare with:
• Grid search over
• ELBO (single term )

β1

β1

β0 = 0

27

Results

• TVO with REINFORCE
• Moments schedule
• Deteriorating performance with

higher Κ (i.e. # β)

28

Results

• TVO with REINFORCE
• Moments schedule
• Deteriorating performance with

higher Κ (i.e. # β)

• TVO with reparameterization (Ours)
• Schedules perform similarly

29

Results

• Comparing TVO with reparamerization, original TVO with IWAE
• By number of importance samples S

30

Summary

• Path exponential family
• Clarifies analysis of TVO
• General construction with rich geometric structure
• Future Directions in “Thermodynamic Variational Inference”
• Improve SNIS sampler with MCMC methods
• Explore deeper connections with thermodynamics

31