Affine Coupling Flow

Build a RealNVP style normalising flow and train it by maximum likelihood on two moons.

A flow is an invertible map from data x to a latent z that has a standard normal distribution. The change of variables formula gives an exact log likelihood:

log p(x) = log N(z; 0, I) + log |det dz/dx|

An affine coupling layer makes the determinant cheap. A fixed binary mask picks the coordinates that pass through unchanged. A network looks only at those coordinates and outputs a log scale s and a shift t for the others:

x_a = x * mask
s, t = net(x_a)                      # both zero where mask == 1
y = x_a + (1 - mask) * (x * exp(s) + t)
log_det = sum of s over the coordinates

The Jacobian is triangular, so its log determinant is the sum of s. The inverse uses the same network, because x_a = y * mask:

x = y_a + (1 - mask) * (y - t) * exp(-s)

Stack layers with alternating masks into a Flow. Then train it in train_flow by minimising -log_prob(X).mean().

Function signature

class AffineCoupling(nn.Module):
    def __init__(self, mask, hidden=64): ...
    def forward(self, x): ...     # -> (y, log_det), log_det of shape (n,)
    def inverse(self, y): ...     # -> x


class Flow(nn.Module):
    def __init__(self, dim=2, n_layers=6, hidden=64): ...
    def forward(self, x): ...     # data to latent -> (z, log_det)
    def inverse(self, z): ...     # latent to data
    def log_prob(self, x): ...    # -> (n,)
    def sample(self, n): ...      # -> (n, dim)


def train_flow(X):

Example

Input

mask = torch.tensor([1.0, 0.0])
x    = torch.tensor([[2.0, 3.0]])
# suppose net(x * mask) gives s = 0.5 and t = 1.0 for coordinate 1

Output

y       = tensor([[2.0000, 5.9462]])
log_det = tensor([0.5000])

Reasoning

Coordinate 0 has mask 1, so it passes through as 2.0. Coordinate 1 is scaled and shifted: 3 * exp(0.5) + 1 = 5.9462. Only one coordinate is scaled, so the log determinant is that coordinate's s, which is 0.5.

Constraints and notes

• mask is a float tensor of zeros and ones with shape (dim,). The coordinates where it is 1 must come out unchanged.
• The tests compare your log_det with the log determinant of the true Jacobian, which they get with torch.autograd.functional.jacobian, after they set your weights to random values. Your layers must work in float64.
• inverse(forward(x)) must give back x, for one layer and for the flow.
• Use alternating masks in Flow, so every coordinate gets transformed.
• train_flow gets a float tensor of shape (n, 2). On held-out points from the same distribution, the mean log likelihood must be above -2.3. The best single Gaussian gets about -2.7.
• Keep training under about 10 seconds on a CPU.