Flows and Discrete VAEs Instructor: John Thickstun Discussion Board: - - PowerPoint PPT Presentation

Discussion Board: Available on Ed Zoom Link: Available on Canvas Instructor Contact: thickstn@cs.washington.edu Course Webpage: https://courses.cs.washington.edu/courses/cse599i/20au/ CSE 599I: Generative Models University of Washington Lecture 9 Instructor: John Thickstun

Flows and Discrete VAEs

Variational Autoencoders

CSE 599I: Generative Models University of Washington Lecture 9

• Generative model . Learn , where
• Estimate the MLE using the ELBO:
• Modeling choices: prior , likelihood , and proposal .
• Construct an expressive family proposals so that the ELBO is tight.

ˆ θmle ≡ arg max

θ

E

x∼p log pθ(x) = arg max θ

sup

qϕ

E

x∼p z∼qϕ(·|x)

log pθ(x, z) qϕ(z|x) .

pθ(x, z) = pθ(x|z)r(z)

pθ(x) ≈ p(x)

pθ(x) = Ez∼r[pθ(x|z)] = Z

Z

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

r(z)

pθ(x|z)

qϕ(z|x)

Normalizing Flows

CSE 599I: Generative Models University of Washington Lecture 9

• Build a series of transformations of our initial proposal
• Let and define
• The log-density of the pushforward distribution on is given by
• Choose functions so that is easy to calculate.

gs : Z → Z

z0 ∼ qφ(·|x).

zt = gt · · · g1(z0).

zt

gs

logdet ✓∂gs(zs−1) ∂zs−1 ◆

log qt(zt) = log q0(z0) −

t

X

s=1

logdet ✓∂gs(zs−1) ∂zs−1 ◆ .

Autoregressive Model?

CSE 599I: Generative Models University of Washington Lecture 9

• The density of the pushforward distribution on is directly parameterized.
• No need to accumulate densities of pushfoward distributions.
• But sampling is slow: O(p) where p is the dimensionality of the latent space.

zt

Inverse Autoregressive Model

CSE 599I: Generative Models University of Washington Lecture 9

• Autoregressive model: serially

transform into sample

• Inverse autoregressive model:

parallel transform into samples

• IAF whitens AR samples.
• Figure credit: Brian Keng.

x ∼ N(0, I)

y ∼ q.

y ∼ q

x ∼ N(0, I).

Inverse Autoregressive Flow

CSE 599I: Generative Models University of Washington Lecture 9

• Inverse autoregressive transformation:
• Claim: Proof:

zt = zt−1 − µt(zt−1) σt(zt−1) .

∂zt ∂zt−1 =      

∂zt,0 ∂zt−1,0

. . .

∂zt,0 ∂zt−1,1 ∂zt,1 ∂zt−1,1

. . . . . . ... ... . . .

∂zt,0 ∂zt−1,p

. . .

∂zt,p−1 ∂zt−1,p ∂zt,p ∂zt−1,p

      =      

1 σt,0

. . .

∂zt,0 ∂zt−1,1 1 σt,1

. . . . . . ... ... . . .

∂zt,0 ∂zt−1,p

. . .

∂zt,p−1 ∂zt−1,p 1 σt,p

      .

logdet ✓ ∂zt ∂zt−1 ◆ = −

p

X

k=1

log σt,k(zt−1,<k).

• Re-write the inverse autoregressive transformation:
• Don’t parameterize mean and deviation; instead parameterize :
• Then the IAF looks like:

Re-parameterizing IAF

CSE 599I: Generative Models University of Washington Lecture 9

zt = zt−1 − µt(zt−1) σt(zt−1) = zt−1 σt(zt−1) − µt(zt−1) σt(zt−1).

zt = vt(zt−1)zt−1 + (1 vt(zt−1)) mt(zt−1).

mt, st : Z → Z

vt = sigmoid(st), σt = 1/vt, µt = σt (1 vt) mt.

IAF Posterior vs AR Prior

CSE 599I: Generative Models University of Washington Lecture 9 IAF Posterior Autoregressive Prior

z0

zt

x

z0

zt

x

≈ ε

≈ ε

(IAF) (AR) Encoder Decoder

5-Minute Break

CSE 599I: Generative Models University of Washington Lecture 9

Variational Autoencoders

CSE 599I: Generative Models University of Washington Lecture 9

• Generative model . Learn , where
• Estimate the MLE using the ELBO:
• Modeling choices: prior , likelihood , and proposal .
• What if is finite, but large?

ˆ θmle ≡ arg max

θ

E

x∼p log pθ(x) = arg max θ

sup

qϕ

E

x∼p z∼qϕ(·|x)

log pθ(x, z) qϕ(z|x) .

pθ(x, z) = pθ(x|z)r(z)

pθ(x) ≈ p(x)

pθ(x) = Ez∼r[pθ(x|z)] = Z

Z

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

r(z)

pθ(x|z)

qϕ(z|x)

Z

• The general form of the ELBO optimization:
• For discrete latent spaces, use a uniform prior where
• The divergence term becomes:
• Likelihood term is unchanged. Combined variational objective is

A Discrete ELBO

CSE 599I: Generative Models University of Washington Lecture 9

ˆ θmle = arg max

θ

sup

q

E

x∼p z∼qφ(·|x)

⇥ log pθ(x|z) D(qφ(z|x) k r(z)) ⇤ .

r(z) = 1/K

K = |Z|.

D(qφ(z|x) k r(z)) = log(K) H(qφ(z|x)).

L(x, θ, φ) = E

z∼qφ(·|x) [log pθ(x|z)] + H(qφ(z|x)) − log(K).

How to Get a Gradient?

CSE 599I: Generative Models University of Washington Lecture 9

• Want to estimate a gradient with respect to of the discrete ELBO:
• Can’t just use the re-parameterization trick: gradients are zero!
• No closed form for the entropy term; we’ll need to estimate that.
• Need to estimate the gradient of the entropy too.

L(x, θ, φ) = E

z∼qφ(·|x) [log pθ(x|z)] + H(qφ(z|x)) − log(K).

φ

The Policy Gradient Theorem

CSE 599I: Generative Models University of Washington Lecture 9

• Also known as REINFORCE:

rφ E

z∼qφ(·|x)

⇥ log pθ(x|z) ⇤ = X

z∈Z

log pθ(x|z)rφqφ(z|x) = X

z∈Z

log pθ(x|z)qφ(z|x)rφ log qφ(z|x) = E

z∼qφ(·|x)

⇥ log pθ(x|z)rφ log qφ(z|x) ⇤ = E

z∼qφ(·|x)

log pθ(x|z) qφ(z|x) rφqφ(z|x)

• .

• A similar trick applies to the entropy term. But the variance is high!

Gumbel Re-parameterization

CSE 599I: Generative Models University of Washington Lecture 9

• Suppose has some factorized structure, e.g.
• Homework 1: if then where
• We can use this to re-parameterize the ELBO. More generally,
• Gradients can come inside the expectation now.
• But gradient of an argmax operation is zero…

z ∼ q

E

z∼qφ

⇥ f(z) ⇤ = E

gu∼Gumbel(0,1)

 f ✓ arg max

u

[log qφ(u) + gu] ◆ .

Z

zi = arg max

u

⇥ log qi(u) + gi,u ⇤ .

Z = {1, . . . , d}k(K = dk).

gi,1, . . . , gi,d ∼ Gumbel(0, 1)

Gumbel Softmax

CSE 599I: Generative Models University of Washington Lecture 9

• Re-parameterized expectation:
• What if we relax the argmax operation to a softmax?
• Softmax at temperature defined by:

E

z∼qφ

⇥ f(z) ⇤ = E

gu∼Gumbel(0,1)

 f ✓ arg max

u

[log qφ(u) + gu] ◆ .

softmax

u

(h(u); τ) = exp(h(u)/τ) P

v exp(h(v)/τ).

E

z∼qφ

⇥ f(z) ⇤ = E

gu∼Gumbel(0,1)

h f ⇣ softmax

u

[log qφ(u) + gu; τ] ⌘i .

τ

Gumbel Softmax Critiqued

CSE 599I: Generative Models University of Washington Lecture 9

• Gumbel-softmax distribution also called the Concrete distribution.
• It is a biased estimator of the gradient.
• Bias goes away as , but gradients blow up.
• Sometimes it works!
• Doesn’t seem to work well when is large.

τ → 0

d

Straight-Through Estimation?

CSE 599I: Generative Models University of Washington Lecture 9

• Want to take a gradient of the form
• The straight-through estimator:
• Where does this come from? Ignore sampling & define
• Use this to compute gradients of expressions with samples :
• Unclear to me why this sometimes works.

rφ E

z∼qφ

⇥ f(z) ⇤ .

˜ rφ E

z∼qφ

[f(z)] ⌘ E

z∼qφ

[(rφf(z))rφqφ(z)] .

˜ rφz ⌘ rφqφ(z).

z

˜ rφf(z) = (rφf(z)) ˜ rφz = (rφf(z))rφqφ(z).