Autoregressive Flow

Normalizing Flows

Normalizing flows (Papamakarios et al., 2021) are the pioneers in generative models. The main idea is to transform a random vector $\mathrm{u}$ through an invertible transport map $\mathrm{T}$:

If $\mathrm{T}$ is bijective and differentiable, the density of $\mathrm{x}$ can be obtained by a change of variables

where \(\text{det}\ \mathrm{J}_{\mathrm{T}}(\mathrm{u})\) quantifies the change of volume w.r.t. u, such that \(\text{det}\ \mathrm{J}_{\mathrm{T}}(\mathrm{u}) \approx \frac{\text{Vol}(\mathrm{d} \mathrm{x})}{\text{Vol}(\mathrm{d} \mathrm{u})}\); the second equality holds since $\text{det} (\mathrm{A} \mathrm{A}^{-1})=1=\text{det} (\mathrm{A}) \text{det}(\mathrm{A}^{-1})$.

Composition Property

To model complex distributions, it is standard to conduct multiple transformations $\mathrm{T}_1$, $\mathrm{T}_2, \cdots$, $\mathrm{T}_K$ such that $\mathrm{T}=\mathrm{T}_K \circ \cdots \circ \mathrm{T}_2 \circ \mathrm{T}_1$. If \(\{\mathrm{T}_k\}_{k=1}^K\) are also bijective and invertible, we have

KL Divergence

Denote by $\phi$ and $\psi$ the parameters of $\mathrm{T}$ and $\mathrm{p}_{\mathrm{u}}(\mathrm{u})$, respectively. The forward KL divergence follows that

The forward KL divergence is well-suited to generate the data distribution if we can model the transport maps $\mathrm{T}$ efficiently as shown in the masked autoencoder below, but is less efficient in computing the model density $\mathrm{p}^{\star}_{\mathrm{X}}(\mathrm{x})$ due to the iterative maps in Eq.\eqref{iterative_maps}.

Alternatively, when we are interested to evaluate model density \(\mathrm{p}_{\mathrm{X}^{\star}}(\mathrm{x})\), we can try to optimize the reverse KL divergence

In empirical training, we resrot to Monte Carlo samples to approximate the expectation.

Autogressive Flows

Autoregressive flows represent one of the earliest developments in flow-based models. We can map the data distribution $\mathrm{p}_\mathrm{X}(\mathrm{x})$ into a uniform distribution in $(0, 1)^\mathrm{D}$ using a transport map with a triangular Jacobian.

To achieve the goal of universal representation, we leverage conditional probability and decompose $\mathrm{p}_\mathrm{X}(\mathrm{x})$ into a product of conditional probabilities as follows

Define the transformation $\mathrm{F}$ to be the CDF function of the conditional density:

The invertibility of $\mathrm{F}$ leads to

Notably, the Jacobian of $\mathrm{F}$ is a lower trianagular matrix and the determinant of $\mathrm{F}$ is equal to the product of diagonal elemants.

If there is another map $\mathrm{G}$ that transforms a general prior distribution to the uniform distribution $(0, 1)^\mathrm{D}$, the flow with transformation $\mathrm{T}=\mathrm{F}^{-1}\circ \mathrm{G}$ can map a general prior to the data distribution.

Modeling

Masked Autoencoder

To model the sequential structure \eqref{seq_structure} via an efficient parallelism, (Germain et al., 2015) abandoned the RNN encoder and proposed a mind-blowing idea by using an masked autoencoder. The code snippet is from (Karpathy, 2019).

class MaskedLinear(nn.Linear):
    """ apply masks to a standard MLP """
    def __init__(self, in_features, out_features, bias=True):
        super().__init__(in_features, out_features, bias)        
        self.register_buffer('mask', torch.ones(out_features, in_features))
        
    def set_mask(self, mask):
        self.mask.data.copy_(torch.from_numpy(mask.astype(np.uint8).T))
        
    def forward(self, input):
        return F.linear(input, self.mask * self.weight, self.bias)

The challenge is to design proper masks to satisfy the autoregressive (triangular) property. A good tutorial is provided below (Khapra, 2019).

To ensure the output $\tilde x_d$ only depend on the preceding inputs $x_{<d}$.

1. Assign the $k$-th hidden unit some random number $m(k)$ in $\{1, 2, \cdots, D-1\}$, which is the maximum number of input units that can be connected.

    # sample the order of the inputs and the connectivity of all neurons
    self.m[-1] = np.arange(self.nin) if self.natural_ordering else rng.permutation(self.nin)
    for l in range(L):
        self.m[l] = rng.randint(self.m[l-1].min(), self.nin-1, size=self.hidden_sizes[l])

1. For any layers except the last one, the mask value is \(1_{m(k)\geq d}\) for the $d$-th output; The sign is changed to $>$ for the last layer.

    # construct the mask matrices
    masks = [self.m[l-1][:,None] <= self.m[l][None,:] for l in range(L)]
    masks.append(self.m[L-1][:,None] < self.m[-1][None,:])

Masked Autogressive Flow (MAF)

MAF (Papamakarios et al., 2017) propose to optimize the reverse KL objective \eqref{reverse_KL}, which is efficient in the conditional density estimation while the sampling stage is slow in high dimensions because we need to iterate the dimension sequentially. To help understand the masked autoregressive flow, I simplified (Karpathy, 2019)’s code template and added some edits and comments.

from nflib.made import ARMLP

class MAF(nn.Module):
    """ Masked Autoregressive Flow that uses a MADE-style network for fast forward """
    def __init__(self, dim, parity, net_class=ARMLP, nh=24):
        super().__init__()
        self.dim = dim
        self.net = net_class(dim, dim*2, nh)
        self.parity = parity

    def forward(self, x):
        s, t = self.net(x)
        z = x * torch.exp(s) + t
        # permute elements for efficiency - pg5 https://arxiv.org/pdf/1705.07057.pdf
        z = z.flip(dims=(1,)) if self.parity else z
        log_det = torch.sum(s, dim=1)
        return z, log_det
    
    def backward(self, z):
        x = torch.zeros_like(z)
        log_det = 0 
        z = z.flip(dims=(1,)) if self.parity else z
        for i in range(self.dim): # iterative sampling from d=1 to D
            s, t = self.net(x.clone()) # clone to avoid in-place op errors if using IAF
            x[:, i] = (z[:, i] - t[:, i]) * torch.exp(-s[:, i])
            log_det += -s[:, i]
        return x, log_det

class NormalizingFlowModel(nn.Module):
    """ A Normalizing Flow Model is a (prior, flow) pair """
    def __init__(self, prior, af_flows):
        super().__init__()
        self.prior = prior
        self.af_flow = nn.ModuleList(af_flows)
        
    def iterate(self, data, direction=1):
        log_det = torch.zeros(data.shape[0])
        for af_flow in self.af_flow[::direction]:
            cur_flow = af_flow.forward if direction == 1 else af_flow.backward
            data, ld = cur_flow(data)
            log_det += ld
        return data, log_det
    
    def forward(self, x):
        z, log_det = self.iterate(x, direction=+1)
        prior_logprob = self.prior.log_prob(z).view(x.size(0), -1).sum(1)
        return z, prior_logprob, log_det

    def backward(self, z):
        return self.iterate(z, direction=-1)
    
    def sample(self, num_samples):
        z = self.prior.sample((num_samples,))
        x, _ = self.iterate(z, direction=-1)
        return x

Inverse Autogressive Flow (IAF)

Similar to MAF, IAF (Kingma et al., 2016) proposed to flip the backward and forward sampling to speed up the sampling efficiency. The density estimation, however, becomes slower. The optimization corresponds to the forward KL divergence.

class IAF(MAF):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.forward, self.backward = self.backward, self.forward