Learned Interpolated Posterior Variance

hardFREEParameterizations and the Reverse PosteriorDiffusion ModelsDiffusion Models 1

Problem Statement

Improved-DDPM lets the network learn the reverse variance by interpolating, in log space, between the two natural bounds beta_t and the posterior beta_tilde_t. Implement that interpolation from the network's raw output.

Background

The model emits a value v per dimension in [-1, 1] (after a tanh, say). The reverse variance is

$\Sigma_t = \exp\big(f\, \log \beta_t + (1 - f)\, \log \tilde{\beta}_t\big), \qquad f = \frac{v + 1}{2}$

where f maps v from [-1, 1] to [0, 1]. The posterior lower bound is beta_tilde_t = (1 - alpha_bar_prev)/(1 - alpha_bar_t) * beta_t. At v = -1 the variance is exactly beta_tilde_t; at v = +1 it is beta_t.

Your Task

Implement:

def learned_variance(v, beta_t, alpha_bar_t, alpha_bar_prev):

v: list of per-dimension model outputs in [-1, 1].

Return the list of variances rounded to 6 decimals.

Input Format

v: list of floats in [-1, 1].
beta_t, alpha_bar_t, alpha_bar_prev (float).

Output Format

A list of floats rounded to 6 decimals.

Sample

print(learned_variance([1.0, -1.0, 0.0], 0.2, 0.5, 0.64))

Output:

[0.2, 0.144, 0.169706]

Example

Input

print(learned_variance([1.0, -1.0, 0.0], 0.2, 0.5, 0.64))

Output

[0.2, 0.144, 0.169706]

Reasoning

Compute the posterior lower bound $\tilde{\beta}_t$ using the given parameters: $\tilde{\beta}_t = \frac{1 - 0.64}{1 - 0.5} \times 0.2 = \frac{0.36}{0.5} \times 0.2 = 0.72 \times 0.2 = 0.144$ .
Determine the interpolation factor $f$ for each dimension in $v = [1.0, -1.0, 0.0]$ using $f = \frac{v+1}{2}$ , yielding $f = [1.0, 0.0, 0.5]$ .
Calculate the log-variance for each dimension via $\log \Sigma_t = f \log \beta_t + (1-f) \log \tilde{\beta}_t$ $lo g Σ_{t} = f lo g β_{t} + (1 - f) lo g \tilde{β}_{t}$ :
- For $v=1.0$ ( $f=1$ ): $\log \Sigma = 1 \cdot \log(0.2) + 0 \cdot \log(0.144) = \log(0.2)$ .
- For $v=-1.0$ ( $f=0$ ): $\log \Sigma = 0 \cdot \log(0.2) + 1 \cdot \log(0.144) = \log(0.144)$ .
- For $v=0.0$ ( $f=0.5$ ): $\log \Sigma = 0.5 \cdot \log(0.2) + 0.5 \cdot \log(0.144) = \log(\sqrt{0.2 \times 0.144})$ .
Exponentiate the log-variances to obtain the final variances:
- $\Sigma_1 = \exp(\log(0.2)) = 0.2$ .
- $\Sigma_2 = \exp(\log(0.144)) = 0.144$ .
- $\Sigma_3 = \sqrt{0.2 \times 0.144} = \sqrt{0.0288} \approx 0.1697056$ .
The final output is [0.2, 0.144, 0.169706]

Constraints

v[i] in [-1, 1]; schedule values in (0, 1).
beta_tilde = (1 - alpha_bar_prev)/(1 - alpha_bar_t)*beta_t.
Interpolate in log space with f = (v+1)/2; round to 6 decimals.

Learn this in

Sampling & AccelerationDiffusion Study Plan

More from Diffusion Models 1

Classifier-Free Guidance: Uncond/Cond Mixingeasy DDPM Posterior Mean and Variance from x0medium Adaptive CFG Scheduling over the Trajectoryhard Classifier Gradient Guidance Stepmedium Negative Prompt as the Unconditional Branchmedium

Concept primer

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:

\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t

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:

\log \Sigma_t = f \log \beta_t + (1 - f) \log \tilde{\beta}_t

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:

Compute the posterior variance $\tilde{\beta}_t$ from the given schedule parameters.
Map the raw network output $v$ to the interpolation factor $f$ .
Compute the log-interpolated variance and exponentiate to get $\Sigma_t$ .
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.

solution.py

Python 3.13

Test results