Normalising Flow Log Likelihood

Compute the exact log likelihood of data under a simple normalising flow.

A normalising flow turns a simple base distribution into a complex one through an invertible function, and it can still give the exact density through the change of variables formula. If x = f(z) with z ~ N(0, I), then

log p(x) = log N(f^-1(x); 0, I) + log |det d f^-1 / d x|

Use the elementwise affine flow x = z * exp(log_scale) + shift. Its inverse is z = (x - shift) * exp(-log_scale), and the Jacobian of the inverse is diagonal, so its log determinant is -sum(log_scale). With the standard normal log density -0.5 * z^2 - 0.5 * log(2 pi) per dimension:

log p(x) = sum_j (-0.5 * z_j^2 - 0.5 * log(2 pi)) - sum_j log_scale_j

Real flows, such as RealNVP and Glow, stack many layers of this kind, with the scale and shift computed by neural networks.

Function signature

def affine_flow_log_prob(x, log_scale, shift):

Example

Input

x         = np.array([[1.0, 2.0]])
log_scale = np.array([0.0, np.log(2.0)])
shift     = np.array([1.0, 0.0])
affine_flow_log_prob(x, log_scale, shift)

Output

array([-3.03102425])

Reasoning

Inverting gives z = ((1 - 1) * 1, (2 - 0) / 2) = (0, 1). The base log density is -0.9189 for 0 and -1.4189 for 1. The second dimension was stretched by 2, which spreads the density thinner, so subtract log 2 = 0.6931. The total is -3.0310.

Constraints and notes

• x has shape (n, d). log_scale and shift have shape (d,).
• Return one log density per row, shape (n,).