Learned Interpolated Posterior Variance

1. Background Knowledge

In DDPM (Denoising Diffusion Probabilistic Models), the reverse process $q(x_{t-1} | x_t, x_0)$ is a Gaussian whose mean and variance are derived from the forward diffusion schedule. The variance of this posterior, often denoted $\tilde{\beta}_t$, is given by:

where $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$ and $\alpha_t = 1 - \beta_t$. This $\tilde{\beta}_t$ is the theoretically optimal variance for the reverse step. However, Improved DDPM (Nichol & Dhariwal, 2021) observes that letting the network learn the variance can improve sample quality. Instead of fixing the variance to $\tilde{\beta}_t$ or $\beta_t$, the model outputs a scalar $v \in [-1, 1]$ per dimension, which is mapped to an interpolation factor $f \in [0, 1]$.

The learned variance is computed in log-space to ensure positivity and numerical stability:

This is a geometric interpolation between the two bounds. When $f = 0$ (i.e., $v = -1$), $\Sigma_t = \tilde{\beta}_t$; when $f = 1$ (i.e., $v = +1$), $\Sigma_t = \beta_t$. The mapping $f = \frac{v + 1}{2}$ linearly transforms the network output from $[-1, 1]$ to $[0, 1]$.

2. Algorithm Approach

The problem is a direct element-wise computation with no iterative or search-based logic. For each dimension:

1. Compute the posterior variance $\tilde{\beta}_t$ from the given schedule parameters.
2. Map the raw network output $v$ to the interpolation factor $f$.
3. Compute the log-interpolated variance and exponentiate to get $\Sigma_t$.
4. Round to 6 decimal places.

This is a vectorizable operation: all dimensions share the same $\beta_t$, $\bar{\alpha}_t$, and $\bar{\alpha}_{t-1}$, so $\tilde{\beta}_t$ is computed once, and the per-dimension work is just the $f$-mapping and log-interpolation.

3. Step-by-Step Strategy

• Compute $\tilde{\beta}_t$: Use the formula $\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t$. This is a single scalar value shared across all dimensions.
• Precompute logs: Calculate $\log \beta_t$ and $\log \tilde{\beta}_t$ once. These are constants for the entire list.
• Loop over each $v_i$:
• Compute $f_i = \frac{v_i + 1}{2}$.
• Compute $\log \Sigma_i = f_i \cdot \log \beta_t + (1 - f_i) \cdot \log \tilde{\beta}_t$.
• Compute $\Sigma_i = \exp(\log \Sigma_i)$.
• Round $\Sigma_i$ to 6 decimal places.
• Collect results into a list and return.

import math

def learned_variance(v, beta_t, alpha_bar_t, alpha_bar_prev):
 beta_tilde = (1 - alpha_bar_prev) / (1 - alpha_bar_t) * beta_t
 log_beta = math.log(beta_t)
 log_beta_tilde = math.log(beta_tilde)
 result = []
 for vi in v:
 f = (vi + 1) / 2.0
 log_sigma = f * log_beta + (1 - f) * log_beta_tilde
 sigma = math.exp(log_sigma)
 result.append(round(sigma, 6))
 return result

4. Common Pitfalls

• Confusing $\bar{\alpha}_t$ and $\bar{\alpha}_{t-1}$: The posterior formula uses $\bar{\alpha}_{t-1}$ in the numerator and $\bar{\alpha}_t$ in the denominator. Swapping them gives an incorrect $\tilde{\beta}_t$.
• Forgetting to exponentiate: The interpolation is done in log-space. If you forget the final $\exp(\cdot)$, you return log-variances instead of variances.
• Incorrect $f$ mapping: The mapping is $f = \frac{v+1}{2}$, not $f = v$ or $f = \frac{v-1}{2}$. Verify with the boundary conditions: $v = -1 \Rightarrow f = 0$, $v = +1 \Rightarrow f = 1$.
• Rounding too early: Round only the final variance, not intermediate values like $f$ or $\log \Sigma_t$. Premature rounding can accumulate error.
• Division by zero: If $\bar{\alpha}_t = 1$, the denominator $1 - \bar{\alpha}_t$ is zero. In practice, $\bar{\alpha}_t < 1$ for all finite $t$, but be aware of this edge case.

5. Time & Space Complexity

• Time: $O(n)$ where $n = \text{len}(v)$. Each dimension requires a constant number of arithmetic operations (one division, one multiplication, one addition, one exponential). The precomputation of $\tilde{\beta}_t$ and its log is $O(1)$.
• Space: $O(n)$ for the output list. No additional data structures are needed beyond the input and output.