Reparameterization Gradient for Non-differentiable Models Wonyeol - - PowerPoint PPT Presentation

Reparameterization Gradient for Non-differentiable Models

Wonyeol Lee Hangyeol Yu Hongseok Yang KAIST

Published at NeurIPS 2018

Reparameterization Gradient for Non-differentiable Models

Wonyeol Lee Hangyeol Yu Hongseok Yang KAIST

Published at NeurIPS 2018

Reparameterization Gradient for Non-differentiable Models

Wonyeol Lee Hangyeol Yu Hongseok Yang KAIST

Published at NeurIPS 2018

Reparameterization Gradient for Non-differentiable Models

Wonyeol Lee Hangyeol Yu Hongseok Yang KAIST

Published at NeurIPS 2018

Backgrounds

Posterior inference

• Latent variable z n.
• Observed variable x m.
• Joint density p(x,z).
• Want to infer posterior p(z|x0) given a

particular value x0 of x.

Variational inference

• 1. Fix a family of variational distr. {qθ(z)}θ.
• 2. Find qθ(z) that approximates p(z|x0) well.
• Typically, by solving

argmaxθ(ELBOθ) where ELBOθ = qθ(z)[ log( p(x0,z)/qθ(z) ) ].

Variational inference

• 1. Fix a family of variational distr. {qθ(z)}θ.
• 2. Find qθ(z) that approximates p(z|x0) well.
• Typically, by solving

argmaxθ(ELBOθ) where ELBOθ = qθ(z)[ log( p(x0,z)/qθ(z) ) ].

differentiable & easy-to-sample

Variational inference

• 1. Fix a family of variational distr. {qθ(z)}θ.
• 2. Find qθ(z) that approximates p(z|x0) well.
• Typically, by solving

argmaxθ(ELBOθ) where ELBOθ = qθ(z)[ log( p(x0,z)/qθ(z) ) ].

Typically, by solving argmaxθ(ELBOθ) where ELBOθ = qθ(z)[ log( p(x0,z)/qθ(z) ) ].

differentiable & easy-to-sample

Variational inference

• 1. Fix a family of variational distr. {qθ(z)}θ.
• 2. Find qθ(z) that approximates p(z|x0) well.
• Typically, by solving

argmaxθ(ELBOθ) where ELBOθ = qθ(z)[ log( p(x0,z)/qθ(z) ) ].

Typically, by solving argmaxθ(ELBOθ) where ELBOθ = qθ(z)[ log( p(x0,z)/qθ(z) ) ].

.. z .. z ..

differentiable & easy-to-sample

Gradient ascent

θn+1 = θn + η × θELBOθ=θn

• Difficult to compute θELBOθ.
• Use an estimated gradient instead.

Gradient ascent

θn+1 = θn + η × θELBOθ=θn

• Difficult to compute θELBOθ.
• Use an estimated gradient instead.

Gradient ascent

θn+1 = θn + η × θELBOθ=θn

• Difficult to compute θELBOθ.
• Use an estimated gradient instead.

Reparameterization estimator

• Works if p(x0,z) is differentiable wrt. z.
• Need distr. q(ε) & smooth function fθ(ε) s.t.

fθ(ε) for ε ~ q(ε) has the distr. qθ(z).

• Derived from the equation:

θELBOθ = q(ε)[ θ(.. fθ(ε) .. fθ(ε) ..) ]

Reparameterization estimator

• Works if p(x0,z) is differentiable wrt. z.
• Need distr. q(ε) & smooth function fθ(ε) s.t.

fθ(ε) for ε ~ q(ε) has the distr. qθ(z).

• Derived from the equation:

θELBOθ = q(ε)[ θ(.. fθ(ε) .. fθ(ε) ..) ]

Reparameterization estimator

• Works if p(x0,z) is differentiable wrt. z.
• Need distr. q(ε) & smooth function fθ(ε) s.t.

fθ(ε) for ε ~ q(ε) has the distr. qθ(z).

• Derived from the equation:

θELBOθ = q(ε)[ θ(.. fθ(ε) .. fθ(ε) ..) ] θELBOθ = θqθ(z)[.. z .. z ..] = θq(ε)[.. fθ(ε) .. fθ(ε) ..] = q(ε)[θ(.. fθ(ε) .. fθ(ε) ..)]

θ θ θ qθ(z) θ q(ε) θ θ q(ε) θ θ θ

Reparameterization estimator

• Works if p(x0,z) is differentiable wrt. z.
• Need distr. q(ε) & smooth function fθ(ε) s.t.

fθ(ε) for ε ~ q(ε) has the distr. qθ(z).

• Derived from the equation:

θELBOθ = q(ε)[ θ(.. fθ(ε) .. fθ(ε) ..) ] θELBOθ = θqθ(z)[.. z .. z ..] = θq(ε)[.. fθ(ε) .. fθ(ε) ..] = q(ε)[θ(.. fθ(ε) .. fθ(ε) ..)]

θ θ θ qθ(z) θ q(ε) θ θ q(ε) θ θ θ

Reparameterization estimator

θELBOθ = θqθ(z)[.. z .. z ..] = θq(ε)[.. fθ(ε) .. fθ(ε) ..] = q(ε)[θ(.. fθ(ε) .. fθ(ε) ..)]

• Works if p(x0,z) is differentiable wrt. z.
• Need distr. q(ε) & smooth function fθ(ε) s.t.

fθ(ε) for ε ~ q(ε) has the distr. qθ(z).

• Derived from the equation:

θELBOθ = q(ε)[ θ(.. fθ(ε) .. fθ(ε) ..) ]

Reparameterization estimator

• Works if p(x0,z) is differentiable wrt. z.
• Need distr. q(ε) & smooth function fθ(ε) s.t.

fθ(ε) for ε ~ q(ε) has the distr. qθ(z).

• Derived from the equation:

θELBOθ = q(ε)[ θ(.. fθ(ε) .. fθ(ε) ..) ] θELBOθ = θqθ(z)[.. z .. z ..] = θq(ε)[.. fθ(ε) .. fθ(ε) ..] = q(ε)[θ(.. fθ(ε) .. fθ(ε) ..)]

θ θ θ qθ(z) θ q(ε) θ θ q(ε) θ θ θ

Non-differentiable models from probabilistic programming

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z) z p(z,x=0)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z) (let [ε (sample (normal 0 1)) z (+ ε θ)] z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

≈

(let [ε (sample (normal 0 1)) z (+ ε θ)] z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

≈

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

(let [ε (sample (normal 0 1)) z (+ ε θ)] z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

≈

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

(let [ε (sample (normal 0 1)) z (+ ε θ)] z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

≈

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

≈

q(ε) = (ε|0,1) z = ε+θ

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z) (let [ε (sample (normal 0 1)) z (+ ε θ)] z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

≈

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z) (let [ε (sample (normal 0 1)) z (+ ε θ)] z)

How to find a good θ?

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

≈

How to find a good θ? By gradient ascent on ELBOθ. θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z) (let [ε (sample (normal 0 1)) z (+ ε θ)] z)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

≈

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

(let [z (sample (normal 0 1))] (if (> z 0) (observe (normal 3 1) 0) (observe (normal -2 1) 0)) z) (let [ε (sample (normal 0 1)) z (+ ε θ)] z)

How to find a good θ? By gradient ascent on ELBOθ.

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[θ]

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

How to find a good θ? By gradient ascent on ELBOθ.

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[θ]

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

How to find a good θ? By gradient ascent on ELBOθ.

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[θ]

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

How to find a good θ? By gradient ascent on ELBOθ.

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[θ]

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

How to find a good θ? By gradient ascent on ELBOθ.

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[θ]

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

q(ε) = (ε|0,1) z = ε+θ

How to find a good θ? By gradient ascent on ELBOθ.

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[ -θ-ε ] = -θ

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

How to find a good θ? By gradient ascent on ELBOθ.

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ q(ε) = (ε|0,1) z = ε+θ

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[ -θ-ε ] = -θ

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

How to find a good θ? By gradient ascent on ELBOθ.

How to find a good θ? By gradient ascent on ELBOθ.

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ q(ε) = (ε|0,1) z = ε+θ

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[ -θ-ε ] = -θ

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

z

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] = q(ε)[ -θ-ε ] = -θ

θ θ θ q(ε) 1 2 q(ε) θ 1 θ 2 q(ε)

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

How to find a good θ? By gradient ascent on ELBOθ.

r1(z) = (z|0,1)(x=0|3,1) r2(z) = (z|0,1)(x=0|-2,1) p(z,x=0) = [z>0]r1(z) + [z≤0]r2(z)

θn+1 ← θn + η × θ ELBOθ=θn

q(ε) = (ε|0,1) z = ε+θ

θELBOθ = θq(ε)[[ε>-θ]log(r1(ε+θ)) + [ε≤-θ]log(r2(ε+θ))] = q(ε)[[ε>-θ]θlog(r1(ε+θ)) + [ε≤-θ]θlog(r2(ε+θ))] + Correction Term How to find a good θ? By gradient ascent on ELBOθ.

Why doesn’t it work?

• Careful when exchanging gradient and

integration.

• May fail unexpectedly.
• May hold unexpectedly, but with correction.

Why doesn’t it work?

• Careful when exchanging gradient and

integration.

• May fail unexpectedly.
• May hold unexpectedly, but with correction.

Why doesn’t it work?

• Careful when exchanging gradient and

integration.

• May fail unexpectedly.
• May hold unexpectedly, but with correction.

+ CorrectionTerm

Our results formally

Non-differentiable models

Non-differentiable models

Non-differentiable models

• forms a partition of .

is differentiable.

Non-differentiable models

• forms a partition of .

is differentiable.

has Lebesgue measure zero.

Wishful thinking

Wishful thinking

Wishful thinking

Wishful thinking

Wishful thinking

Wishful thinking

Wishful thinking

Wishful thinking

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over

Correction

• surface integral over
• Accounts for the impact of moving the boundaries.

Can be estimated by manifold sampling.

Correction

• surface integral over
• Accounts for the impact of moving the boundaries.

Can be estimated by manifold sampling.

Two ingredients

• Differentiation under moving domain:

Two ingredients

• Differentiation under moving domain:
• Divergence theorem:

Two ingredients

Surface integral over

• is a normal vector of
• .

Correction term

Surface integral over

• Correction term

Requires manifold sampling Hard to estimate in general cases

Surface integral over

• Easy to estimate if
• is a hyperplane.

Correction term

Surface integral over

• Easy to estimate if
• is a hyperplane.
• Assume the branch condition of each if-

statement is linear in .

Correction term

Subsampling

• surface integral over
• For computational efficiency,

we subsample surface integrals.

Subsampling

• surface integral over
• For computational efficiency,

we subsample surface integrals.

Experiments

Implementation

• Implemented a black-box variational

inference engine for a simple probabilistic programming language

• Supports sample, observe, if, ...
• Written in Python, using autograd package.

Benchmarks

textmsg

• Models #’s of per-day SNS msg’s, where SNS-

usage pattern changes on some day.

• Non-differentiable part: the day of change in

SNS-usage pattern.

• Given #’s of per-day SNS msg’s over 2 months,

infer the day when the pattern changes.

Benchmarks

temperature

• Models random dynamics of a controller that

tries to keep room temp. stable.

• Non-differentiable part: on/off of air conditioner,
• n which evolution of room temp. depends.
• Given noisy observations of temp. at each step,

infer on/off status of the controller at each step.

ELBO

{dotted, dashed, solid} lines: {N = 1, N = 8, N = 16}

ELBO

{dotted, dashed, solid} lines: {N = 1, N = 8, N = 16}

Computation time

High-level message

• Careful when exchanging gradient and integration.

High-level message

• Careful when exchanging gradient and integration.
• May fail unexpectedly.

High-level message

• Careful when exchanging gradient and integration.
• May fail unexpectedly.
• May hold unexpectedly, but with correction.

Any questions?