Normalizing Flows: Exact Likelihood with Expressive Distributions

Module: Generative AI | Difficulty: Advanced

Change of Variables

RealNVP Coupling Layers

Jacobian: — computed in .

Optimal Transport Connection

Brenier's theorem: optimal transport map for convex .

Residual Flows + Hutchinson Trace

import torch, torch.nn as nn

class AffineCoupling(nn.Module):
    def __init__(self, dim, h=256):
        super().__init__()
        self.dim = dim
        self.s = nn.Sequential(nn.Linear(dim//2,h),nn.ReLU(),nn.Linear(h,dim//2),nn.Tanh())
        self.t = nn.Sequential(nn.Linear(dim//2,h),nn.ReLU(),nn.Linear(h,dim//2))
    def forward(self, x, rev=False):
        x1, x2 = x[:,:self.dim//2], x[:,self.dim//2:]
        s, t = self.s(x1), self.t(x1)
        if not rev: return torch.cat([x1, x2*torch.exp(s)+t],1), s.sum(1)
        return torch.cat([x1, (x2-t)*torch.exp(-s)],1), -s.sum(1)

| Method | Exact LL | Sampling Speed | FID | |--------|---------|----------------|-----| | RealNVP | Yes | Fast | 28.7 | | Glow | Yes | Fast | 25.3 | | FFJORD | Approx | Slow | 23.8 | | Diffusion | Approx | Slow | 2.9 |