FFVB: Control variate

This section describes a control variate technique for variance reduction.

See: Variational Bayes Introduction, Fixed Form Variational Bayes

Control variate

As is typical of stochastic optimization algorithms, the performance of Algorithm 3 depends greatly on the variance of the noisy gradient. Variance reduction for the noisy gradient is a key ingredient in FFVB algorithms.

Let $θ_{s} \sim q_{λ} (θ)$ , $s = 1, \dots, S$ , be $S$ samples from the variational distribution $q_{λ} (θ)$ . A naive estimator of the $i$ th element of the vector $\nabla_{λ} LB (λ)$ is $\hat{\nabla_{λ_{i}} LB} (λ)^{naive} = \frac{1}{S} \sum_{s = 1}^{S} \nabla_{λ_{i}} [\log q_{λ} (θ_{s})] \times h_{λ} (θ_{s}),$ whose variance is often too large to be useful.

For any number $c_{i}$ , consider

\begin{matrix} (21) & \hat{\nabla_{λ_{i}} LB} (λ) = \frac{1}{S} \sum_{s = 1}^{S} \nabla_{λ_{i}} [\log q_{λ} (θ_{s})] (h_{λ} (θ_{s}) - c_{i}), \end{matrix}

which is still an unbiased estimator of $\nabla_{λ_{i}} LB (λ)$ since $E (\nabla_{λ} [\log q_{λ} (θ)]) = 0$ , whose variance can be greatly reduced by an appropriate choice of control variate $c_{i}$ .

The variance of $\hat{\nabla_{λ_{i}} LB} (λ)$ is

\frac{1}{S} V (\nabla_{λ_{i}} [\log q_{λ} (θ)] h_{λ} (θ)) + \frac{c_{i}^{2}}{S} V (\nabla_{λ_{i}} [\log q_{λ} (θ)]) - \frac{2 c_{i}}{S} cov (\nabla_{λ_{i}} [\log q_{λ} (θ)] h_{λ} (θ), \nabla_{λ_{i}} [\log q_{λ} (θ)]) .

The optimal $c_{i}$ that minimizes this variance is

\begin{matrix} (22) & c_{i} = cov (\nabla_{λ_{i}} [\log q_{λ} (θ)] h_{λ} (θ), \nabla_{λ_{i}} [\log q_{λ} (θ)]) / V (\nabla_{λ_{i}} [\log q_{λ} (θ)]) . \end{matrix}

Then

V (\hat{\nabla_{λ_{i}} LB} (λ)) = V (\hat{\nabla_{λ_{i}} LB} (λ)^{naive}) (1 - ρ_{i}^{2}) \leq V (\hat{\nabla_{λ_{i}} LB} (λ)^{naive}),

where $ρ_{i}$ is the correlation between $\nabla_{λ_{i}} [\log q_{λ} (θ)] h_{λ} (θ)$ and $\nabla_{λ_{i}} [\log q_{λ} (θ)]$ . Often, $ρ_{i}^{2}$ is very close to 1, which leads to a large variance reduction.

One can estimate the numbers $c_{i}$ in $(22)$ using samples $θ_{s} \sim q_{λ} (θ)$ . In order to ensure the unbiasedness of the gradient estimator, the samples used to estimate $c_{i}$ must be independent of the samples used to estimate the gradient.

In practice, the $c_{i}$ can be updated sequentially as follows. At iteration $t$ , we use the $c_{i}$ computed in the previous iteration $t - 1$ , i.e. based on the samples from $q_{λ^{(t - 1)}} (θ)$ , to estimate the gradient $\hat{\nabla_{λ} LB} (λ^{(t)})$ , which is computed using new samples from $q_{λ^{(t)}} (θ)$ .

We then update the $c_{i}$ using this new set of samples. By doing so, the unbiasedness is guaranteed while no extra samples are needed in updating the control variates $c_{i}$ .

Algorithm 4 provides a detailed pseudo-code implementation of the FFVB approach that uses the control variate for variance reduction and moving average adaptive learning, and Algorithm 5 implements the FFVB approach that uses the control variate and natural gradient.

Algorithm 4: FFVB with control variates and adaptive learning

Input: Initial $λ^{(0)}$ , adaptive learning weights $β_{1}, β_{2} \in (0, 1)$ , fixed learning rate $ϵ_{0}$ , threshold $τ$ , rolling window size $t_{W}$ and maximum patience $P$ . Model-specific requirement: function $h (θ) := \log (p (θ) p (y ∣ θ))$ .
Initialization
- Generate $θ_{s} \sim q_{λ^{(0)}} (θ)$ , $s = 1, \dots, S$ .
- Compute the unbiased estimate of the LB gradient
  $\hat{\nabla_{λ} LB} (λ^{(0)}) := \frac{1}{S} \sum_{s = 1}^{S} \nabla_{λ} \log q_{λ} (θ_{s}) \times h_{λ} (θ_{s}) |_{λ = λ^{(0)}} .$
- Set $g_{0} := \hat{\nabla_{λ} LB} (λ^{(0)})$ , $v_{0} := (g_{0})^{2}$ , $\bar{g} := g_{0}$ , $\bar{v} := v_{0}$ .
- Estimate the vector of control variates $c$ as in $(22)$ using the samples $θ_{s}, s = 1, \dots, S$ .
- Set $t = 0$ , $patience = 0$ and $stop=false$ .
While $stop=false$ :
- Generate $θ_{s} \sim q_{λ^{(t)}} (θ)$ , $s = 1, \dots, S$ .
- Compute the unbiased estimate of the LB gradient
  $g_{t} := \hat{\nabla_{λ} LB} (λ^{(t)}) = \frac{1}{S} \sum_{s = 1}^{S} \nabla_{λ} \log q_{λ} (θ_{s}) \circ (h_{λ} (θ_{s}) - c) |_{λ = λ^{(t)}} .$
- Estimate the new control variate vector $c$ as in $(22)$ using the samples $θ_{s}, s = 1, \dots, S$ .
- Compute $v_{t} = (g_{t})^{2}$ and
  $\bar{g} = β_{1} \bar{g} + (1 - β_{1}) g_{t}, \bar{v} = β_{2} \bar{v} + (1 - β_{2}) v_{t} .$
- Compute $α_{t} = min (ϵ_{0}, ϵ_{0} \frac{τ}{t})$ and update
  $λ^{(t + 1)} = λ^{(t)} + α_{t} \bar{g} / \sqrt{\bar{v}}$
- Compute the lower bound estimate
  
  $\hat{LB} (λ^{(t)}) := \frac{1}{S} \sum_{s = 1}^{S} h_{λ^{(t)}} (θ_{s}) .$ \item If $t \geq t_{W}$ : compute the moving averaged lower bound ${\overset{―}{LB}}_{t - t_{W} + 1} = \frac{1}{t_{W}} \sum_{k = 1}^{t_{W}} \hat{LB} (λ^{(t - k + 1)}),$
  
  and if ${\overset{―}{LB}}_{t - t_{W} + 1} \geq max (\overset{―}{LB})$ patience = 0; else $patience := patience + 1$ .
- If $patience \geq P$ , $stop=true$ .
- Set $t := t + 1$ .

Note: The term $\nabla_{λ} \log q_{λ} (θ_{s}) \circ (h_{λ} (θ_{s}) - c)$ should be understood component-wise, i.e. it is the vector whose $i$ th element is $\nabla_{λ_{i}} \log q_{λ} (θ_{s}) \times (h_{λ} (θ_{s}) - c_{i})$ .

Algorithm 5: FFVB with control variates and natural gradient

Input: Initial $λ^{(0)}$ , momentum weight $α_{m}$ , fixed learning rate $ϵ_{0}$ , threshold $τ$ , rolling window size $t_{W}$ and maximum patience $P$ . Model-specific requirement: function function $h (θ) := \log (p (θ) p (y ∣ θ))$ .
Initialization
- Generate $θ_{s} \sim q_{λ^{(0)}} (θ)$ , $s = 1, \dots, S$ .
- Compute the unbiased estimate of the LB gradient
  $\hat{\nabla_{λ} LB} (λ^{(0)}) := \frac{1}{S} \sum_{s = 1}^{S} \nabla_{λ} \log q_{λ} (θ_{s}) \times h_{λ} (θ_{s}) |_{λ = λ^{(0)}} .$
  and the natural gradient
  $\hat{\nabla_{λ} LB} (λ^{(0)})^{nat} := I_{F}^{- 1} (λ^{(0)}) \hat{\nabla_{λ} LB} (λ^{(0)}) .$
- Set momentum gradient
$\overset{―}{\nabla_{λ} LB} := \hat{\nabla_{λ} LB} (λ^{(0)})^{nat} .$
- Estimate control variate vector $c$ as in $(22)$ using the samples $θ_{s}, s = 1, \dots, S$ .
- Set $t = 0$ , $patience = 0$ and $stop=false$ .
While $stop=false$ :
- Generate $θ_{s} \sim q_{λ^{(t)}} (θ)$ , $s = 1, \dots, S$ .
- Compute the unbiased estimate of the LB gradient
  $\hat{\nabla_{λ} LB} (λ^{(t)}) = \frac{1}{S} \sum_{s = 1}^{S} \nabla_{λ} \log q_{λ} (θ_{s}) \circ (h_{λ} (θ_{s}) - c) |_{λ = λ^{(t)}}$
  and the natural gradient
  $\hat{\nabla_{λ} LB} (λ^{(t)})^{nat} = I_{F}^{- 1} (λ^{(t)}) \hat{\nabla_{λ} LB} (λ^{(t)}) .$
- Estimate the new control variate vector $c$ as in $(22)$ using the samples $θ_{s}, s = 1, \dots, S$ .
- Compute the momentum gradient
  $\overset{―}{\nabla_{λ} LB} = α_{m} \overset{―}{\nabla_{λ} LB} + (1 - α_{m}) \hat{\nabla_{λ} LB} (λ^{(t)})^{nat} .$
- Compute $α_{t} = min (ϵ_{0}, ϵ_{0} \frac{τ}{t})$ and update
  $λ^{(t + 1)} = λ^{(t)} + α_{t} \overset{―}{\nabla_{λ} LB} .$
- Compute the lower bound estimate
  $\hat{LB} (λ^{(t)}) := \frac{1}{S} \sum_{s = 1}^{S} h_{λ^{(t)}} (θ_{s}) .$
- If $t \geq t_{W}$ : compute the moving average lower bound
  ${\overset{―}{LB}}_{t - t_{W} + 1} = \frac{1}{t_{W}} \sum_{k = 1}^{t_{W}} \hat{LB} (λ^{(t - k + 1)}),$
  and if ${\overset{―}{LB}}_{t - t_{W} + 1} \geq max (\overset{―}{LB})$ patience = 0; else $patience := patience + 1$ .
- If $patience \geq P$ , $stop=true$ .
- Set $t := t + 1$ .

Next: FFVB with reparameterization trick