• Add stochastic layers \(z \sim p(z)\)
• The task then becomes $$ \mathbb{E}_{p(z)} L(f(x, z|\phi), y) \to \min_\phi$$
• Monte Carlo estimate the gradient if don't want to learn \(z\)'s parameters $$ \nabla_\phi L(f(x, z|\phi), y), z \sim p(z) $$

This talk: What if \(z \sim p(z|\theta)\)?

(and we'd like to find optimal \(\theta\) as well)

• General idea: we just need to figure out how to backpropagate through stochasticity
• Our task is $$\mathbb{E}_{p(z|\theta)} F_\theta(z) \to \min_\theta$$
• Let \(\mathcal{F}(\theta, \phi) = \mathbb{E}_{p(z|\theta)} F_\phi(z)\), then $$ \nabla_\theta \mathcal{F}(\theta, \theta) = \nabla_\theta F(\theta, \phi) \Bigr|_{\phi=\theta} + \nabla_\phi \mathcal{F}(\theta, \phi)\Bigr|_{\phi=\theta} $$
• Backpropagation through \(F_\theta(z)\) is an easy problem, just backprop as usual with noise fixed
• How to differentiate w.r.t. expectation's parameters?
• See "Gradient Estimation Using Stochastic Computation Graphs" by J. Schulman et al.

Let \(F(z) = L(f(x, z|\phi), y)\), then

$$ \nabla_\theta \mathbb{E}_{p(z|\theta)} F(z) = \nabla_\theta \int F(z) p(z|\theta) dz = \int F(z) \nabla_\theta p(z|\theta) dz $$ $$= \int F(z) \nabla_\theta \log p(z|\theta) p(z|\theta) dz = \mathbb{E}_{p(z|\theta)} \nabla_\theta \log p(z|\theta) F(z) $$

Intuition: push probabilities of good samples (as measured by \(F(z)\)) up.

Pros: very general, does not require differentiable \(F\).

Cons: known to have large variance, sensitive to values of \(F\).

We'll get back to this estimator later in the talk.

• Can actually differentiate a sample w.r.t. its parameters!
• For a univariate r.v. \(X_\theta\) with CDF \(\mathbb{F}(x|\theta)\) we have $$ \mathbb{F}(X_\theta|\theta) \sim U[0, 1] $$
• Hence $$ 0 = \frac{\partial}{\partial \theta} \mathbb{F}(X_\theta |\theta) = \frac{\partial}{\partial \theta} \mathbb{F}(x|\theta)|_{x=X_\theta} + p(X_\theta) \frac{\partial}{\partial \theta} X_\theta $$ $$ \frac{\partial X_\theta}{\partial \theta} = - \frac{1}{p(X_\theta)} \int_{-\infty}^{X_\theta} \frac{\partial}{\partial \theta} p(x|\theta) dx$$
• Extends to multivariate case with conditional marginal CDFs

Pros: backprop through sample

Cons: requires ability to differentiate CDF w.r.t. \(\theta\)

• Stochastic Computation Graph with Gaussian hidden nodes: $$ z|x \sim \mathcal{N}(z \mid \mu(x|\theta), \Sigma(x|\theta)) $$
• \(z|x\) has the same distribution as \(\mu(x|\theta) + \Sigma(x|\theta)^{1/2} \varepsilon \)

• If \(z|\theta\) has the same distribution as \(\mathcal{T}_\theta(\varepsilon) \) for some differentiable \(\mathcal{T}_\theta\)
• Then $$ \mathbb{E}_{p(z|\theta)} F(z) = \mathbb{E}_{p(\varepsilon)} F(\mathcal{T}_\theta(\varepsilon)) $$ $$ \nabla_\theta \mathbb{E}_{p(z|\theta)} F(z) = \mathbb{E}_{p(\varepsilon)} \nabla_\theta F(\mathcal{T}_\theta(\varepsilon)) $$

• Pros: obtained estimator is very efficient in practise
• Cons:
  • not every distribution has a useful reparametrisation (e.g. Dirichlet distribution)
  • requires differentiable \(F\)

• Idea: reparametrisation removes dependence on parameters completely. What if we remove it just partially?
• For example, whiten first moments: $$ \varepsilon = \mathcal{T}_\theta^{-1}(z) = \frac{z - \mu(\theta)}{\sigma(\theta)} \quad\quad z = \mathcal{T}_\theta(\varepsilon) $$
• \(\varepsilon\) still depends on \(\theta\): \(\varepsilon \sim p(\varepsilon|\theta)\)

$$ \nabla_\theta \mathbb{E}_{p(\varepsilon|\theta)} F(z) = \mathbb{E}_{p(\varepsilon|\theta)} \nabla_\theta F(\mathcal{T}_\theta(\varepsilon)) + \mathbb{E}_{p(\varepsilon|\theta)} \nabla_\theta \log p(\varepsilon|\theta) F(\mathcal{T}_\theta(\varepsilon))$$

• The first term is reparametrised gradient, the second is the correction to make the gradient unbiased.
• Can also be seen from the rejection sampling perspective *
  • Don't need the \(\mathcal{T}^{-1}_\theta\) then

The formula $$ \nabla_\theta \mathbb{E}_{p(\varepsilon|\theta)} F(z) = \mathbb{E}_{p(\varepsilon|\theta)} \nabla_\theta F(\mathcal{T}_\theta(\varepsilon)) + \mathbb{E}_{p(\varepsilon|\theta)} \nabla_\theta \log p(\varepsilon|\theta) F(\mathcal{T}_\theta(\varepsilon))$$ requires us to sample \( \varepsilon|\theta\). With a bit of algebra we can rewrite these addends in terms of samples \(z|\theta\): $$ \mathbb{E}_{p(z|\theta)} \nabla_z F(z) \nabla_\theta h_\theta(\mathcal{T}^{-1}_\theta(z)) $$ $$\mathbb{E}_{p(z|\theta)} F(z) \left[ \nabla_\theta \log p(z|\theta) + \nabla_z\log p(z|\theta) h_\theta(\mathcal{T}^{-1}_\theta(z)) + u_\theta(\mathcal{T}^{-1}_\theta(z))\right]$$

Where

• \(h_\theta(\varepsilon) = \nabla_\theta \mathcal{T}_\theta(\varepsilon)\) – Jacobian of \(\mathcal{T}_\theta\) w.r.t. \(\theta\)
• \(u_\theta(\varepsilon) = \nabla_\theta \log \Bigl|\text{det} \nabla_\varepsilon \mathcal{T}_\theta(\varepsilon)\Bigr| \)

Pros: interpolates between reparametrisation and REINFORCE

Cons: need to come up with differentiable \(\mathcal{T}_\theta\)

• Suppose you have binary dropout and want to adjust each neuron's dropout rate.
• Consider a pair of layers \(h_1\) and \(h_2\) such that \( h_2 = \sigma(W h_1 + b) \)
• Dropout:  \( h_2 = \sigma(W (h_1 \circ z) + b), \quad z \sim \text{Bernoulli}(p) \)
• Dropout masks only enter sums with many addends
• We might hope for CLT to hold
• Hopefully preactivations of \(h_2\) are approximately normal:
  • Mean \( W(h_1 \circ \mathbb{E} z) + b = W (h_1 \circ p) + b\)
  • Covariance \( W \mathbb{D} z W^T = W (p \circ (1-p)) W^T\)
• Equivalent to multiplicative \(\mathcal{N}(p, p(1-p))\) noise

• Pros: benefits of Gaussian Reparametrisation
• Cons: very limited applicability scope, depends on how good the CLT approximation is

• Can use any continuous r.v. \(X\) with CDF \(\mathbb{F}(x)\) to generate \(\text{Bernoulli}(p)\) r.v. by $$ z = [X> \mathbb{F}^{-1}(1-p)] = H(X- \mathbb{F}^{-1}(1-p))$$
• Approximate the step function using a sigmoid function with temperature: $$ H(x-a) = \lim_{\tau \to 0} \sigma(\tfrac{x-a}{\tau}) $$

• Which \(X\) to choose?
  • Sample values close to 0 and 1 more often
    • ⇒ need modes there
    • ⇒ \(X\) should have all \(\mathbb{R}\) as a support
  • Want to be able to relax general categorical r.v.

• Gumbel-Max Trick will help

• Gumbel random variables \(\{\gamma_k\}_{k=1}^K\) have a nice property: $$ \mathbb{P}(k = \text{argmax}_j [\gamma_j + \log p_j]) = \frac{p_k}{\sum_j p_j} $$
• Gumbel-Max Trick: Sample categorical r.v. via maximisation of independently perturbed log-probabilities

• Suppose argmax returns a one-hot vector
• Approximate it with a softmax with temperature $$ \text{argmax}(x) = \lim_{\tau \to 0} \text{softmax}(\tfrac{x}{\tau}) $$ (temperature makes outputs more contrast)

$$ z = \text{argmax}_k \left[\gamma_k + \log p_k\right] $$

$$ \zeta = \text{softmax}_\tau \left(\gamma_k + \log p_k\right) $$

Many other estimators only relax the backward pass

• Straight Through Estimator: backpropagate through hard thresholding as if it was an identity function: $$\tfrac{\partial}{\partial \theta} H(x - \sigma(-\theta)) = 1$$
• Or one might keep the derivative of \(\sigma\): $$H'(x) = 1 \Rightarrow \tfrac{\partial}{\partial \theta} H(x - \sigma(-\theta)) = \sigma(\theta) \sigma(-\theta)$$
• Gumbel Straight Through: backward pass as in the Gumbel-Softmax relaxation

Pros: don't see any

Cons: mathematically unsound ¯\_(ツ)_/¯

• REINFORCE gradient estimate: \( \nabla_\theta \log p(z|\theta) F(z) \)
• We've seen how offsetting \(F(z)\) can break the estimator

• Easy to fix by just centring \(F(z)\)
• Leads to a baseline: \( b \nabla_\theta \log p(z|\theta) \)
• Optimal \(b\) can be found analytically, and estimated using moving averages

• Pros: Simple, sparse gradients
• Cons: Not efficient enough

• We've seen some methods and experiments, time to ask ourselves "What could make REINFORCE more efficient?"
• Control Variates: $$ \mathbb{E}_{p(z|\theta)} \left[F(z) \nabla_\theta \log p(z|\theta) + B(z)\right]- \mathbb{E}_{p(z|\theta)} B(z)$$ if \(B(z)\) and \(F(z) \nabla_\theta \log p(z|\theta)\) are negatively correlated, resulting estimator has lower variance.

How to design a control variate?

• Using prior information (dependency structure) often helps
• Gradient information \(\nabla_z F(z)\) should help
• Good baselines "integrate out" as much as possible:
  • Ideally we'd like to decompose REINFORCE estimate into easy to work with parts
  • Or at least reduce its influence
  • Thus likely to use samples \(z\)

• Neural Variational Inference and Learning in Belief Networks
• Very similar to what RL people do

1. Centring learning signal: $$ \mathbb{E}_{p(z|x, \theta)} [F(x, z) - \color{red}{b(x) - C}] \nabla_\theta \log p(z|x, \theta)$$ (\(b(x)\) does not have to be scalar)
2. Variance Normalisation: divide by estimate of standard deviation.
3. Local learning signals: exploit inner structure of \(p(z|x, \theta)\) and \(F(x,z)\)
   • Akin to using cumulative future reward in RL

• VIMCO: extension to multi-sample objectives

Pros: more efficient, still easy to implement

Cons: requires training an extra model \(b(x)\) (unless VIMCO), does not use \(z\) in the baseline, doesn't use gradient \(\nabla_z F(z)\)

• Idea: use first-order Taylor approximation as a baseline $$ F(z) \approx F(\zeta) + \nabla F(\zeta)^T (z - \zeta) $$
• Easy to fix the bias: $$ \mathbb{E}_{p(z|\theta)} \nabla F(\zeta)^T z \nabla_\theta \log p(z|\theta) = \nabla_\theta \mathbb{E}_{p(z|\theta)} \nabla F(\zeta)^T z = \nabla F(\zeta)^T \nabla_\theta \mathbb{E}_{p(z|\theta)} z $$
• \(\zeta\) is chosen as a "mean-field" point

Pros: uses gradient information \(\nabla_z F(z)\)

Cons:

• Need to compute \(\mathbb{E}_{p(z|\theta)} z\) analytically to backprop through it
• Only a first order approximation

• Idea: use Gumbel-Softmax relaxation as a baseline
• Turns out can compensate for the bias

• Assume \( z|\theta = H(X|\theta)\), then $$ \mathbb{E}_{p(X|\theta)} F(z) \nabla_\theta \log p(z|\theta) = \mathbb{E}_{p(X|\theta)}F(H(X|\theta)) \nabla_\theta \log p(X|\theta)$$ RHS has high variance due to non-marginalised sampling
• Control variate: $$ \mathbb{E}_{p(X|\theta)} F(\sigma_\tau(X)) \nabla_\theta \log p(X|\theta) = \nabla_\theta \mathbb{E}_{p(X|\theta)} F(\sigma_\tau(X))$$
• Can conditionally marginalise the estimator

\( \nabla_\theta \mathbb{E}_{p(X|\theta)} F(\sigma_\tau(X)) = \nabla_\theta \mathbb{E}_{p(X, z|\theta)} F(\sigma_\tau(X)) = \)

$$ \mathbb{E}_{p(z|\theta)} \left[\nabla_\theta \mathbb{E}_{p(X|z, \theta)} F(\sigma_\tau(X|z))\right] + \mathbb{E}_{p(z|\theta)} \mathbb{E}_{p(X|z,\theta)} [F(\sigma_\tau(X|z))] \nabla_\theta \log p(z|\theta) $$ 

$$\nabla_\theta \mathbb{E}_{p(z|\theta)} F(z) = \mathbb{E}_{u, v} \left[ \left(F(z) - \eta F(\zeta|z) \right) \nabla_\theta \log p(z|\theta) + \eta \nabla_\theta F(\zeta) - \eta \nabla_\theta F(\zeta|z) \right]$$ 

Where \(z=H(X)\), \(\zeta = \sigma_\tau(X)\), \(\zeta|z = \sigma_\tau(X|z)\)

\( X = \log \frac{u}{1-u} + \log \tfrac{\mu(\theta)}{1-\mu(\theta)} \)

\(\eta\) and \(\tau\) are tuneable parameters, optimised to reduce the variance

Pros:

• uses \(F\)'s gradient information
• quality of approximation is controlled by temperature
• essentially Gumbel-Softmax with bias removed

Cons:

• quite involved
• requires several evaluations of \(F\)

What if \(F(z)\) is not differentiable or we don't know its gradients (like in RL)?

• If \(z\) is continuous and reparametrisable, we can learn a baseline \(c(z)\): $$ \hat g = (F(z) - \tilde{F}(z)) \nabla_\theta \log p(z|\theta) + \nabla_\theta \tilde{F}(z), \quad z = \mathcal{T}(\varepsilon|\theta), \varepsilon \sim p(\varepsilon)$$We reparametrised the last term; \(\tilde{F}\) is optimised to minimise the variance of this gradient estimate
• Discrete case: relax \(z\) just as in REBAR, gives us $$\hat g = (F(z) - \eta \tilde{F}(\zeta|z)) \nabla_\theta \log p(z|\theta) + \eta \nabla_\theta F(\zeta) - \eta \nabla_\theta F(\zeta|z) $$

The \(\tilde{F}\) is optimized to minimize the variance $$\text{Var} \hat{g}_i = \mathbb{E} \hat{g}_i^2 - \left(\mathbb{E} \hat{g}_i\right)^2$$ \(\hat{g}_i\) is unbiased, hence the second term does not depend on \(\tilde{F}\)