Optimal smoothing for pathwise adjoints
Jonathan Hüser, Shih-Te Yang and Uwe Naumann
Informatik 12: Software and Tools for Computational Engineering
RWTH Aachen UniversityDecember 8, 2017

1 / 22

Binary Stochastic Neural Network1

1Techniques for Learning Binary Stochastic Feedforward Neural Networks, Raiko et al., 2015
2 / 22

Binary Stochastic Neural Network: GoalI

Generative model for binarized MNIST

3 / 22

Binary Stochastic Neural Network: DAG
hzylx

VW

  
1h=σ V,xzi ∼ Bernoulli(hi ),  
1y=σ Wz
4 / 22

Binary Stochastic Neural Network: LossI

Minimize negative loglikelihood of lower pixels
X
min Ex,y − log
P (y | W, z i )P (z i | V, x)
V,Wi

5 / 22

Binary Stochastic Neural Network: LossI

Minimize negative loglikelihood of lower pixels
X
min Ex,y − log
P (y | W, z i )P (z i | V, x)
V,WIi

MC estimator of loss
X− log
P (y | W, z i )P (z i | V, x) = − log Ez P (y | W, z)iN≈ − log

1 X
P (y | W, z j ) = L(V, W ) = l
N j=1

5 / 22Binary Stochastic Neural Network: LossI

Minimize negative loglikelihood of lower pixels
X
min Ex,y − log
P (y | W, z i )P (z i | V, x)
V,WIi

MC estimator of loss
X− log
P (y | W, z i )P (z i | V, x) = − log Ez P (y | W, z)iN≈ − logI

1 X
P (y | W, z j ) = L(V, W ) = l
N j=1

Pathwise adjoint does not give sensitivity
V(1) += ∇V L(V, W )T l(1) = 05 / 22

Score Function Trick(Likelihood Ratio Method / REINFORCE2 )2Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning,Williams, 1992
6 / 22

Score Function Trick: Idea
I

Bridge sampling gap: Use derivatives of the probability
X∇V Ez∼P (g(V )) f (z) =f (z i ) ∇V P (z i | g(V ))i=Xf (z i ) ∇V log P (z i | g(V )) P (z i | g(V ))

iN

≈

1 X
f (z j ) ∇V log P (z j | g(V ))N j=1

7 / 22

Score Function Trick: Idea
I

Bridge sampling gap: Use derivatives of the probability
X∇V Ez∼P (g(V )) f (z) =f (z i ) ∇V P (z i | g(V ))i=Xf (z i ) ∇V log P (z i | g(V )) P (z i | g(V ))

iN

≈

I

1 Xf (z j ) ∇V log P (z j | g(V ))N j=1

Seed adjoint g(1) to plug into AD computation
V(1) += ∇V g(V )T∂log P (z j | g) f (z j )∂g= f (z j ) ∇V log P (z j | g(V ))7 / 22

Score Function Trick: Variance Problem
I

Score function trick results in a high variance gradient estimator
min Ez (z − 0.49)2 + c,θz ∼ Bernoulli(θ)

1000

≈

1 X j
(z − 0.49)2 + c1000 j=18 / 22

Score Function Trick: Variance Problem
I

Score function trick results in a high variance gradient estimator
min Ez (z − 0.49)2 + c,z ∼ Bernoulli(θ)

θ1000

≈

SGD using score function adjoint
0.8
c=0
c=1
c=5c = 100.6
θ

I

1 X j
(z − 0.49)2 + c1000 j=10.4
0.2
00500

1,000

1,500

2,000

2,500

iteration
8 / 22

Reparametrization Trick9 / 22

Reparametrization Trick: IdeaI

Bridge sampling gap: Make random variable deterministic function ofparameter and nonparametric noisez ∼ Bernoulli(g(θ))→u ∼ Uniform(0, 1)
z = H(g(θ) − u)

10 / 22

Reparametrization Trick: Idea
I

Bridge sampling gap: Make random variable deterministic function ofparameter and nonparametric noisez ∼ Bernoulli(g(θ))→u ∼ Uniform(0, 1)
z = H(g(θ) − u)I

H is the Heaviside step function(1H(x) =0if x ≥ 0otherwise

10 / 22

Reparametrization Trick: Smoothing
I

Approximate H(x) ≈ σ(x/p) to obtain a differentiable surrogate model

3
The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables, Maddison
et al., 201711 / 22

Reparametrization Trick: Smoothing
I
I

Approximate H(x) ≈ σ(x/p) to obtain a differentiable surrogate model
Choice of p introduces bias variance tradeoff→ optimal smoothing

3
The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables, Maddison
et al., 201711 / 22

Reparametrization Trick: Smoothing

IApproximate H(x) ≈ σ(x/p) to obtain a differentiable surrogate model
Choice of p introduces bias variance tradeoff→ optimal smoothing

I

Resulting z is a continuous relaxation

I3
The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables, Maddison
et al., 201711 / 22

Reparametrization Trick: Smoothing

I

Approximate H(x) ≈ σ(x/p) to obtain a differentiable surrogate model
Choice of p introduces bias variance tradeoff→ optimal smoothing

I

Resulting z is a continuous relaxation

I

CONCRETE3 relaxation z = σ((log(θ/(1 − θ)) + log(1 − u) − log(u))/p)I

3

The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables, Maddison
et al., 201711 / 22

Reparametrization Trick: Bias Problem
I

Smoothing reparametrization trick results in biased gradient estimator
min Ez (z − 0.49)2 + 10,
θz ∼ Bernoulli(θ)

12 / 22

Reparametrization Trick: Bias Problem
I

Smoothing reparametrization trick results in biased gradient estimator
min Ez (z − 0.49)2 + 10,θ

SGD using CONCRETE relaxation adjoint with p = 0.1
0.8
score function
concrete0.6
θIz ∼ Bernoulli(θ)

0.4
0.200

1,000

2,000

3,000

4,000

5,000

iteration
12 / 22

Control Variates

13 / 22

Control Variates: Idea
I

Reduce variance using correlated function with “known” meanNEz f (z) ≈1 X(f (z j ) − f˜(z j )) + Ez f˜(z)N j=1

4
REBAR: Low-Variance, Unbiased Gradient Estimates for Discrete Latent Variable Models,
Tucker et al., 201714 / 22

Control Variates: Idea
I

Reduce variance using correlated function with “known” meanNEz f (z) ≈I

1 X(f (z j ) − f˜(z j )) + Ez f˜(z)N j=1

REBAR4 : Use smoothing as control variate for score function→ low variance unbiased estimator

4REBAR: Low-Variance, Unbiased Gradient Estimates for Discrete Latent Variable Models,Tucker et al., 201714 / 22

Control Variates: Idea
IReduce variance using correlated function with “known” meanNEz f (z) ≈II

1 X(f (z j ) − f˜(z j )) + Ez f˜(z)N j=1

REBAR4 : Use smoothing as control variate for score function→ low variance unbiased estimator
Roughly as follows (details more involved due to marginalization)N∇θ Ez L(z) ≈N

1 X
1 X(L(z j ) − L̃(z j )) ∇θ log P (z j | θ) +∇θ L̃(z j )
N j=1
N i=1
|{z}
{z}
|
score function adjointsmoothed pathwise adjoint

4REBAR: Low-Variance, Unbiased Gradient Estimates for Discrete Latent Variable Models,
Tucker et al., 201714 / 22

Control Variates: No Problem
I

REBAR control variate results in low variance unbiased gradient estimator
min Ez (z − 0.49)2 + 10,
θz ∼ Bernoulli(θ)

15 / 22

Control Variates: No ProblemIREBAR control variate results in low variance unbiased gradient estimator
min Ez (z − 0.49)2 + 10,
θSGD using REBAR with p = 0.10.8score function
concreterebar

0.6
θIz ∼ Bernoulli(θ)

0.4
0.2
00

1,000

2,000

3,000

4,000

5,000

iteration

15 / 22

Control Variates: No Problem
I

REBAR control variate results in low variance unbiased gradient estimator
min Ez (z − 0.49)2 + 10,
θIz ∼ Bernoulli(θ)

SGD using REBAR with p = 0.10.8score function
concrete
rebar

θ

0.6
0.4
0.2
001,0002,000

3,000

4,000

5,000

iteration
IOptimal smoothing: Parameter p still “unkown”
15 / 22

Control Variates: Adaptive Optimization
ILet ∇θ L(θ, p) be the REBAR estimator, minimize variancemin kE (∇θ L(θ, p)2 ) − (E ∇θ L(θ, p))2 k22 = G(p) = gp

16 / 22

Control Variates: Adaptive Optimization
ILet ∇θ L(θ, p) be the REBAR estimator, minimize variancemin kE (∇θ L(θ, p)2 ) − (E ∇θ L(θ, p))2 k22 = G(p) = gp

ITo optimize via SGD we can use a one sample gradient estimator∂∂
E (2∇θ L(θ, p)∇θ L(θ, p) ) − 2(E ∇θ L(θ, p))
E ∇θ L(θ, p)∂p∂p||{z}{z}second-order adjoint

zero

16 / 22

Control Variates: Adaptive Optimization
ILet ∇θ L(θ, p) be the REBAR estimator, minimize variancemin kE (∇θ L(θ, p)2 ) − (E ∇θ L(θ, p))2 k22 = G(p) = gp

To optimize via SGD we can use a one sample gradient estimator∂∂
E (2∇θ L(θ, p)∇θ L(θ, p) ) − 2(E ∇θ L(θ, p))
E ∇θ L(θ, p)∂p∂p||
{z}
{z}
zero

second-order adjoint

0.8
score function
concreterebar
adaptive

0.6
θ

I

0.4
0.2
0

0

1,000

2,000

3,000

4,000

5,000

iteration
16 / 22

Second-order Adjoints56

5Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, Griewank
and Walther, 2008
6
The Art of Differentiating Computer Programs, Naumann, 201117 / 22

Second-order Adjoints: Reverse-over-reverse 1
I

7

Pathwise adjoint of l = L(θ, p) isθ(1)∇θ L(θ, p)T+=l(1)Tp(1)∇p L(θ, p)

dco/c++: Derivative Code by Overloading in C++, Leppkes et al., under review
18 / 22

Second-order Adjoints: Reverse-over-reverse 1IPathwise adjoint of l = L(θ, p) is



θ(1)
∇θ L(θ, p)T+=l(1)Tp(1)∇p L(θ, p)I

Adjoint of hyper parameter objective g = G(θ, p) isp(2) += (∂p G(θ, p, θ(1) )T + ∂p θ(1) (θ, p)T ∂θ(1) G(θ, p, θ(1) )T ) g(2)7dco/c++: Derivative Code by Overloading in C++, Leppkes et al., under review
18 / 22

Second-order Adjoints: Reverse-over-reverse 1I

Pathwise adjoint of l = L(θ, p) isθ(1)∇θ L(θ, p)T+=l(1)Tp(1)∇p L(θ, p)I

Adjoint of hyper parameter objective g = G(θ, p) isp(2) += (∂p G(θ, p, θ(1) )T + ∂p θ(1) (θ, p)T ∂θ(1) G(θ, p, θ(1) )T ) g(2)I

Implementation via type generic overloading AD (dco/c++7 )dco :: ga1s < float >:: type x_a1s ;dco :: ga1s < dco :: ga1s < float >:: type >:: type x_a1s_a2s ;7dco/c++: Derivative Code by Overloading in C++, Leppkes et al., under review
18 / 22

Second-order Adjoints: Reverse-over-reverse 2I

Adjoint of hyper parameter objective g = G(θ, p, θ(1) (θ, p)) isp(2) += (∂p G(θ, p, θ(1) )T + ∂p θ(1) (θ, p)T ∂θ(1) G(θ, p, θ(1) )T ) g(2)
g

θ(1)

lθp

19 / 22

Second-order Adjoints: Forward-over-reverse 1
ISave memory for second-order adjoints by local replacement with tangent

20 / 22

Second-order Adjoints: Forward-over-reverse 1II

Save memory for second-order adjoints by local replacement with tangent
Pathwise adjoint of l = L(θ, p) is


θ(1)
∇θ L(θ, p)T+=l = ∇L(θ, p)T l(1)p(1)∇p L(θ, p)T (1)20 / 22

Second-order Adjoints: Forward-over-reverse 1II

I

Save memory for second-order adjoints by local replacement with tangent
Pathwise adjoint of l = L(θ, p) isθ(1)∇θ L(θ, p)T+=l = ∇L(θ, p)T l(1)p(1)∇p L(θ, p)T (1)Tangent of pathwise adjoint is" (2) # (2) θ(1)θT
2+=l∇L(θ,p)
(1)
(2)
p(2)
p(1)
∂2L
∂θ1 ∂θ1


..

 2.
 ∂ L
 ∂θn ∂θ12
θ∇ L(θ, p) =  ∂2L
 ∂p1 ∂θ1


..
 2.

∂ L
∂pnp ∂θ1...

...
...

...

∂2L
∂θ1 ∂θnθ∂2L∂θ1 ∂p1..
.

..
.

∂2L
∂θnθ ∂θnθ
∂2L∂p1 ∂θnθ∂2L∂θnθ ∂p1∂2L∂p1 ∂p1..
.

..
.

∂2L
∂pnp ∂θnθ

∂2L∂pnp ∂p1...

∂2L∂θ1 ∂pnp

..
.

...
...

...


∂2L
∂θnθ ∂pnp ∂2L∂p1 ∂pnp 
..
.
2∂ L
∂pnp ∂pnp

20 / 22

Second-order Adjoints: Forward-over-reverse 2
I

Adjoint of hyper parameter objective g = G(θ, p, θ(1) (θ, p)) isp(2) += (∂p G(θ, p, θ(1) )T + ∂p θ(1) (θ, p)T ∂θ(1) G(θ, p, θ(1) )T ) g(2)I8

Tangent of pathwise adjoint is" (2) # (2) θ(1)θT
2+=l∇L(θ,p)(1)(2)p(2)p(1)dco/c++: Derivative Code by Overloading in C++, Leppkes et al., under review
21 / 22

Second-order Adjoints: Forward-over-reverse 2I

Adjoint of hyper parameter objective g = G(θ, p, θ(1) (θ, p)) isp(2) += (∂p G(θ, p, θ(1) )T + ∂p θ(1) (θ, p)T ∂θ(1) G(θ, p, θ(1) )T ) g(2)II8Tangent of pathwise adjoint is" (2) # (2) θ(1)θT
2+=l∇L(θ,p)(1)(2)p(2)p(1)(2)

We seed θ(2) = ∂θ(1) G(θ, p, θ(1) )T to get p(1) = ∂p θ(1) (θ, p)T θ(2)dco/c++: Derivative Code by Overloading in C++, Leppkes et al., under review
21 / 22

Second-order Adjoints: Forward-over-reverse 2
I

Adjoint of hyper parameter objective g = G(θ, p, θ(1) (θ, p)) isp(2) += (∂p G(θ, p, θ(1) )T + ∂p θ(1) (θ, p)T ∂θ(1) G(θ, p, θ(1) )T ) g(2)III8Tangent of pathwise adjoint is" (2) # (2) θ(1)θT2+=l∇L(θ,p)
(1)
(2)
p(2)
p(1)
(2)

We seed θ(2) = ∂θ(1) G(θ, p, θ(1) )T to get p(1) = ∂p θ(1) (θ, p)T θ(2)Local replacement of reverse-over-reverse by forward-over-reverse is an
Adjoint Code Design Pattern (see Poster)

dco/c++: Derivative Code by Overloading in C++, Leppkes et al., under review
21 / 22

Second-order Adjoints: Forward-over-reverse 2I

Adjoint of hyper parameter objective g = G(θ, p, θ(1) (θ, p)) isp(2) += (∂p G(θ, p, θ(1) )T + ∂p θ(1) (θ, p)T ∂θ(1) G(θ, p, θ(1) )T ) g(2)IIII

Tangent of pathwise adjoint is" (2) # (2) θ(1)θT
2+=l∇L(θ,p)
(1)
(2)
p(2)
p(1)
(2)

We seed θ(2) = ∂θ(1) G(θ, p, θ(1) )T to get p(1) = ∂p θ(1) (θ, p)T θ(2)Local replacement of reverse-over-reverse by forward-over-reverse is an
Adjoint Code Design Pattern (see Poster)
Implementation via type generic overloading AD (dco/c++8 )dco :: ga1s < float >:: type x_a1s ;dco :: ga1s < dco :: gt1s < float >:: type >:: type x_a1s_t2s ;8dco/c++: Derivative Code by Overloading in C++, Leppkes et al., under review
21 / 22

Thank you

22 / 22