Learned Interpolated Posterior Variance

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:

beta_tilde=(1-0.64)/(1-0.5)0.2=0.144. v=1 -> f=1 -> beta_t=0.2; v=-1 -> f=0 -> 0.144; v=0 -> f=0.5 -> exp(0.5ln0.2+0.5ln0.144)=sqrt(0.20.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.