Autoregressive models: PixelRNN and PixelCNN

The families of this course deal with the intractable likelihood in different ways: a VAE bounds it, a GAN avoids it, a flow computes it through an invertible map, diffusion bounds it again. Autoregressive models remove the difficulty at its source. They write the density of an image as a product of one-dimensional conditional distributions, each a softmax over a few hundred values and therefore normalized by construction. The likelihood is exact, training is a cross-entropy loss, and the price is paid at sampling time, one pixel after another.

Notation. An image is a vector of $n$ values $x = (x_1, \dots, x_n)$ in a fixed order, usually row by row from the top left, and for colour images red, green and blue within each pixel. $x_{<i} = (x_1, \dots, x_{i-1})$ are the values before position $i$. Values are integers in $\{0, \dots, 255\}$ unless stated otherwise.

The chain rule as a model#

Any joint distribution can be written as

$$p(x) = \prod_{i=1}^{n} p(x_i \mid x_{<i}), \qquad \log p(x) = \sum_{i=1}^{n} \log p(x_i \mid x_{<i}).$$

This is the chain rule of probability, not an assumption: it holds for every distribution and every ordering of the coordinates. A model becomes restrictive only through the way the conditionals are parameterized. Maximum likelihood, which the KL note derives from minimizing $D_{\mathrm{KL}}(p_{\text{data}} \,\|\, p_\theta)$, becomes a sum of $n$ classification losses: at every position, predict the next value from the previous ones, and pay the cross-entropy of the true value.

The whole difficulty is in representing the conditionals compactly. Take binary images of $n$ pixels. A table of the joint distribution has $2^n - 1$ free numbers, and so do the tables of all conditionals together, since the conditional of pixel $i$ needs one number for each of the $2^{i-1}$ histories. The fully visible sigmoid belief network replaces each table by a logistic regression,

$$p(x_i = 1 \mid x_{<i}) = \sigma\Big(\alpha_i + \sum_{j<i} \alpha_{ij} x_j\Big),$$

with $i$ parameters for pixel $i$ and $n(n+1)/2$ in total. For the $28 \times 28$ images of MNIST, $n = 784$, that is 307,720 parameters instead of $2^{784} - 1$. Later models share more: NADE (Larochelle and Murray, 2011) computes all conditionals with one hidden layer whose weights are reused across positions, and MADE (Germain et al., 2015) gets them all from one pass of an autoencoder whose weights are masked so that output $i$ depends only on inputs before $i$; the same masks build the autoregressive flows. PixelRNN and PixelCNN go further and reuse one recurrent or convolutional network at every position, so the number of parameters no longer depends on the size of the image.

Sampling follows the product: draw $x_1$, then $x_2$ given $x_1$, and so on. Each draw is the inverse-CDF method for a discrete variable: take $u \sim \mathcal U(0, 1)$ and return the smallest value whose cumulative probability reaches $u$. An autoregressive model is still noise plus a learned transformation, with $n$ uniform numbers as the noise.

Pixels as discrete values#

PixelRNN (van den Oord, Kalchbrenner and Kavukcuoglu, 2016) treats each colour value as a category and predicts it with a 256-way softmax. A discrete distribution needs no assumption about its shape, so it can be multimodal or skewed, and it assigns probabilities rather than densities: $\log p(x) \le 0$ and no dequantization is needed, unlike the continuous densities of flows.

Likelihoods of image models are reported in bits per dimension:

$$\text{bits/dim} = -\frac{\log_2 p(x)}{n} = -\frac{\ln p(x)}{n \ln 2},$$

averaged over test images. A model that spreads probability uniformly over 256 values gets exactly 8 bits/dim. The number has a direct meaning: with arithmetic coding, a model with $b$ bits/dim compresses images losslessly to about $b \cdot n$ bits. On CIFAR-10, $n = 32 \cdot 32 \cdot 3 = 3072$, and the best PixelRNN reaches 3.00 bits/dim: $3.00 \cdot 3072 = 9216$ bits, about 1.1 KiB per image, or 37.5% of the 24,576 bits of the raw image.

A model that predicts fewer levels solves a coarser problem. With 10 levels of grey a uniform guess costs $\ln 10 \approx 2.30$ nats or $\log_2 10 \approx 3.32$ bits per pixel, and such numbers cannot be compared with bits/dim over 256 values.

Parallel training, sequential sampling#

During training the whole image is known. If the network is built so that its output at position $i$ depends only on the inputs $x_{<i}$, a single forward pass on the real image yields all $n$ conditionals at once, and the loss is their sum. Feeding the true previous values rather than the model’s own samples is called teacher forcing, and it makes training as parallel as for any convolutional network.

Sampling cannot be parallel: $x_{i+1}$ is computed from $x_i$, which must be sampled first. A model that predicts one value at a time needs $n$ network passes for a sample, 3072 for a CIFAR-10 image and $89 \cdot 60 = 5340$ for a grey image of $89 \times 60$ pixels. A model that predicts a whole pixel per pass, such as PixelCNN++ below, needs one pass per pixel, 1024 for CIFAR-10, and draws the three values of the pixel from the predicted parameters; the positions are still sequential. A plain implementation also recomputes the whole image at every pass, although only one output is used. This asymmetry, cheap density and expensive sampling, is the same as for the masked autoregressive flow.

PixelRNN and PixelCNN#

The two networks of van den Oord, Kalchbrenner and Kavukcuoglu (2016) differ in how the output at position $i$ collects information from all of $x_{<i}$ while seeing nothing of $x_i$ and later values.

Recurrent layers#

Row LSTM processes the image row by row from top to bottom and computes the state of a whole row at once. The input-to-state part is a $k \times 1$ convolution along the row ($k \ge 3$), and the state-to-state part is a one-dimensional convolution over the previous row of states. Each step looks a little to the left and right in the row above, so the context of a pixel is a triangle above it and misses part of the available past.

Diagonal BiLSTM captures the whole past. It skews the input map, shifting each row one position to the right of the previous one, so that the diagonals of the image become columns of an $n \times (2n-1)$ map. The state-to-state convolution is $2 \times 1$ along these columns, and two LSTMs scan in opposite directions; the output of the right-to-left one is shifted down a row before the two are added, so that no future pixel leaks in. The models use up to 12 layers with residual connections between them. The recurrences run sequentially over rows or diagonals even during training, which makes both models slow to train.

Masked convolutions#

PixelCNN replaces the recurrences with convolutions whose kernels are multiplied by a mask. For a $k \times k$ kernel with centre $c = \lfloor k/2 \rfloor$,

$$M_{uv} = \begin{cases} 1, & u < c \text{ or } (u = c \text{ and } v < c), \\ 0, & \text{otherwise}, \end{cases}$$

so a position sees the rows above it and the pixels to its left. The first layer uses mask A, which excludes the centre, because its input at the centre is the value to be predicted. The later layers use mask B, which includes it: after the first layer, the feature at position $i$ has been computed from $x_{<i}$ only, and passing it on is safe. The PixelCNN of the paper has fifteen layers, a $7 \times 7$ first layer with mask A followed by residual blocks with $3 \times 3$ convolutions. Training is fully parallel; sampling is not.

Figure

The 5×5 convolution masks of PixelCNN

Colour channels#

With colour, each pixel holds three values ordered R, G, B, and the factorization goes over all $3 \times$ (number of pixels) of them:

$$p(x_i \mid x_{<i}) = p(x_{i,R} \mid x_{<i})\; p(x_{i,G} \mid x_{<i}, x_{i,R})\; p(x_{i,B} \mid x_{<i}, x_{i,R}, x_{i,G}),$$

where $x_i$ is now the whole pixel. The feature maps of every layer are split into three groups, one per colour. Mask A lets the red group see no value of the current pixel, the green group see red, and the blue group see red and green; mask B in addition connects each group to itself.

The simplest alternative models the channels independently, for example with three grey-level models, one per channel. Each channel is then a plausible image on its own, but nothing ties the three together: in a sample, the red, green and blue values of a pixel, and the edges they draw, need not agree. A milder version conditions every channel on all earlier pixels of all channels but drops the dependence among the three values of the current pixel. The chain rule needs that dependence, so this version is an assumption rather than a factorization. A model that keeps the dependence either follows the channel order inside the masks, as above, or, as PixelCNN++ does below, predicts the three values of a pixel jointly with a simple dependence between them.

The blind spot#

Stacking masked convolutions does not give the full past. A $3 \times 3$ mask reaches one row up and one column to each side, or one column to the left in the current row. After $L$ layers, the row $a$ steps above the target is covered from $L$ columns to the left up to only $a$ columns to the right, whatever $L$ is. The pixels further right in the rows just above stay out of reach at any depth. Within the window of $L$ rows and $\pm L$ columns, this blind spot holds $L(L-1)/2$ of the $2L(L+1)$ earlier pixels, 6 of 40 for $L = 4$; its share grows toward a quarter, the figure van den Oord et al. (2016b) give. The Code section below finds the same 6 pixels by differentiating a randomly initialized network.

Figure

The blind spot of PixelCNN

Gated PixelCNN#

van den Oord et al. (2016b) close the gap between PixelCNN and PixelRNN with two changes. The first fixes the blind spot with two stacks of convolutions. A vertical stack is not masked within its window and sees all rows above the current one, so its receptive field grows as a rectangle. A horizontal stack sees the current row up to the current pixel. After each layer the vertical stack feeds the horizontal one, never the other way round: information from the horizontal stack would carry pixels to the right of or below the target into the vertical stack.

The second change replaces the ReLUs between layers with multiplicative units, a stand-in for the gates of an LSTM:

$$y = \tanh(W_{k,f} * x) \odot \sigma(W_{k,g} * x),$$

where $*$ is a convolution, $k$ the layer and $\odot$ the element-wise product. A condition enters as an extra term inside both nonlinearities: for a vector $h$, such as a one-hot class label, $y = \tanh(W_{k,f} * x + V_{k,f}^\top h) \odot \sigma(W_{k,g} * x + V_{k,g}^\top h)$, and for a spatial map $s$ the products $V^\top h$ become convolutions $V * s$. The conditional model is $p(x \mid h) = \prod_i p(x_i \mid x_{<i}, h)$.

On CIFAR-10 the Gated PixelCNN reaches 3.03 bits/dim, against 3.14 for PixelCNN and 3.00 for PixelRNN. On ImageNet it is slightly better than PixelRNN, 3.83 against 3.86 bits/dim at $32 \times 32$ and 3.57 against 3.63 at $64 \times 64$, and it matches the Row LSTM in less than half the training time (60 hours on 32 GPUs). The ImageNet model has 20 layers of 384 units with $5 \times 5$ filters.

PixelCNN++#

A 256-way softmax knows nothing about the order of its classes: 127 and 128 are as unrelated as 0 and 255 until the network learns otherwise, a value never seen in training gets almost no probability, and the output layer has $3 \cdot 256 = 768$ numbers per pixel. PixelCNN++ (Salimans et al., 2017) assumes a continuous intensity $\nu$ with a mixture of $K$ logistic distributions and models the observed value as $\nu$ rounded to an integer:

$$P(x \mid \pi, \mu, s) = \sum_{k=1}^{K} \pi_k \Big[\sigma\Big(\frac{x + 0.5 - \mu_k}{s_k}\Big) - \sigma\Big(\frac{x - 0.5 - \mu_k}{s_k}\Big)\Big],$$

where $\sigma$ is the logistic sigmoid, the CDF of the logistic distribution. For $x = 0$ the term $x - 0.5$ is replaced by $-\infty$, and for $x = 255$ the term $x + 0.5$ by $+\infty$, so the probabilities sum to one over $\{0, \dots, 255\}$ and the edges collect the mass of everything beyond them. With one component, $\mu = 127.3$ and $s = 2$, the values 126, 127, 128 and 129 get probabilities 0.112, 0.124, 0.121 and 0.105: neighbours are close automatically. With $\mu = 254$ and $s = 2$, the edge value 255 gets 0.438 and 254 gets 0.124, which suits the excess of saturated pixels in real photographs.

The other changes:

• Whole pixels. The network conditions on whole pixels above and to the left and predicts the three channels of a pixel together. The mixture indicator is shared by the channels, and inside a component the means depend linearly on the channels already drawn: $\mu_G + \alpha R$ and $\mu_B + \beta R + \gamma G$, with $\alpha, \beta, \gamma$ produced by the network. The output is $10K$ numbers per pixel ($K$ weights, a mean and a scale for each channel and component, and three coefficients per component), 50 for $K = 5$, which the paper finds sufficient.
• Downsampling with stride-2 convolutions instead of dilated convolutions to see large regions at several resolutions, with short-cut connections that carry the full-resolution features past the downsampled layers; without them the model fails to train.
• Dropout on the residual path. Without it the training loss falls below 2.0 bits per sub-pixel and the test loss rises above 6.0.

PixelCNN++ reaches 2.92 bits/dim on CIFAR-10. The same model trained on images with added uniform noise, a continuous density instead of the rounding, gives a bound of 3.11.

How the likelihoods compare#

Test negative log-likelihood on CIFAR-10, in bits per dimension (van den Oord et al., 2016a, 2016b; Salimans et al., 2017):

Model	Bits/dim
Uniform distribution	8.00
Multivariate Gaussian	4.70
NICE (flow)	4.48
Real NVP (flow)	3.49
PixelCNN	3.14
Row LSTM	3.07
Gated PixelCNN	3.03
Diagonal BiLSTM (PixelRNN)	3.00
PixelCNN++	2.92

For several years after 2016 autoregressive models held the best likelihoods on these benchmarks. A likelihood measures how well the model compresses the test set, which is not the same as how good its samples look: the evaluation notes measure that separately, and a model can win on one and lose on the other.

Strengths and limits#

• Exact likelihood and simple training. The loss is a sum of cross-entropies with no bound, no adversary and no Jacobian, and it can be monitored on held-out data.
• Slow sampling. One network evaluation per position: $n$ for one image with a sub-pixel model, one per pixel ($n/3$ for colour) when the network predicts whole pixels as PixelCNN++ does, with the channel values then drawn from its output. Either way the cost grows with resolution. Diffusion also needs many passes, but their number is set by the noise schedule, not by the number of pixels.
• No latent code. Nothing summarizes an image in a few numbers, and the ordering is arbitrary: the model must learn global structure from left-to-right, top-to-bottom prediction. PixelVAE (Gulrajani et al., 2017) combines the two ideas: a VAE whose decoder is a PixelCNN. It needs far fewer autoregressive layers than a PixelCNN alone, and its codes are more compressed than those of a plain VAE while still capturing most of the structure of the image.
• Strong priors. Over a small grid of discrete codes, an autoregressive model is an excellent prior. This is its role in VQ-VAE, which fits a PixelCNN over the code indices, and the starting point of the next note.

Notebook#

The PixelCNN notebook of the HSE-LAMBDA Deep Generative Models repository trains a small PixelCNN on movie posters (cell numbers count every cell from 1). It has no saved outputs, so it documents the procedure rather than results.

• Cells 1–3: setup. Posters of $268 \times 182$ pixels are downscaled by 3 to $89 \times 60$; batches of 32.
• Cells 4–5, task 1: bucketize maps 0–255 to 10 levels by integer division by 26 and clamping at 9. The buckets are not quite equal, as the docstring says: levels 0–8 hold 26 values each and level 9 the 22 values 234–255. The task text gives the shape as batch × channels × width × height; the tensors are batch × channels × height × width.
• Cells 6–7, task 2: resize, convert to grey, ToTensor, multiply by 255 and bucketize. The product returns the exact integers here, so no pixel lands in a wrong bucket.
• Cells 9–13: CustomImageDataset lists the image files in ../data/movie-genre-from-its-poster/posters/, the folder name of the Kaggle Movie Genre from its Poster dataset of IMDB posters. It opens every file once with PIL’s verify() and skips those it cannot read, converts each image to RGB before the transforms, and returns a dummy label 0 with every image so that the default loader works. The last 100 files in sorted order are the validation set. The validation loader drops the last incomplete batch, so 96 of them are used.
• Cells 15–19, task 3: make_mask and MaskedCNN, which multiplies weight.data by the mask in place on every forward pass. This works: the masked weights still receive gradients and the optimizer moves them between steps, but the next forward pass zeroes them again, the same effect as the PixelRNN paper’s zeroing after each update. Multiplying in the forward pass without .data, as in the Code section, keeps their gradients at zero.
• Cells 20–21, task 4: a model of one mask-A layer, one mask-B layer and a $1 \times 1$ output to 10 logits per pixel. Two masked layers see little: each prediction reaches at most $2 \lfloor k/2 \rfloor$ rows up, 6 of the 89 rows for $7 \times 7$ kernels. The model receives the level indices 0–9 as floats without rescaling.
• Cells 23–24: Adam, cross-entropy over the 10 levels and 50 epochs. After every epoch two validation posters are completed from row 10 and column 10: the region below and to the right is zeroed and refilled pixel by pixel, $79 \cdot 50 = 3950$ network passes for the region, while the pixels left of column 10 in the lower rows stay real. The loss is in nats per pixel, with the uniform guess at $\ln 10 \approx 2.30$; it measures the 10-level problem and is not comparable with bits/dim over 256 values.
• Cell 25: eight completions from row 30 and column 30, $59 \cdot 30 = 1770$ network passes for the batch.
• Cell 26: the question about colour images, answered in Colour channels above.

Code#

A PixelCNN for grey images with levels values per pixel, in PyTorch, and a check of its receptive field:

import torch
import torch.nn as nn
import torch.nn.functional as F

def make_mask(include_center, height, width):
    mask = torch.zeros(height, width)
    mask[: height // 2] = 1                       # rows above the centre
    mask[height // 2, : width // 2] = 1           # left of the centre in its row
    if include_center:
        mask[height // 2, width // 2] = 1
    return mask

class MaskedConv2d(nn.Conv2d):
    def __init__(self, include_center, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.register_buffer('mask', make_mask(include_center, *self.kernel_size))

    def forward(self, x):
        return F.conv2d(x, self.weight * self.mask, self.bias, padding=self.padding)

class PixelCNN(nn.Module):
    def __init__(self, levels=10, hidden=64, layers=4, k=3):
        super().__init__()
        self.levels = levels
        blocks = [MaskedConv2d(False, 1, hidden, k, padding=k // 2), nn.ReLU()]
        for _ in range(layers - 1):
            blocks += [MaskedConv2d(True, hidden, hidden, k, padding=k // 2), nn.ReLU()]
        self.net = nn.Sequential(*blocks, nn.Conv2d(hidden, levels, 1))

    def forward(self, x):                         # x: [B, 1, H, W] with integer levels
        return self.net(x.float() / (self.levels - 1))   # logits: [B, levels, H, W]

    @torch.no_grad()
    def sample(self, x, start=(0, 0)):            # fill x in raster order from start
        x = x.clone()
        H, W = x.shape[-2:]
        for i in range(start[0], H):
            for j in range(start[1] if i == start[0] else 0, W):
                probs = self(x)[:, :, i, j].softmax(dim=1)
                x[:, 0, i, j] = torch.multinomial(probs, 1).squeeze(1)
        return x

# training step: loss = F.cross_entropy(model(x), x.long().squeeze(1)), in nats per pixel

torch.manual_seed(0)
model = PixelCNN()
x = torch.randint(0, 10, (1, 1, 11, 11)).float().requires_grad_(True)
model(x)[0, :, 5, 5].sum().backward()             # which inputs reach the output at (5, 5)?
print((x.grad[0, 0, 1:6, 1:10] != 0).int())

The printed window covers the four rows above position $(5, 5)$, its own row and four columns to each side:

tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 0]], dtype=torch.int32)

No input at or after $(5, 5)$ in raster order has a nonzero gradient, and the zeros in the rows above are the 6 pixels of the blind spot in the figure.

Next#

PixelCNN predicts one pixel from a fixed neighbourhood. Autoregressive transformers replace the convolutions with self-attention, which sees the whole past at every layer, and, for large images, move from pixels to the discrete codes of a VQ autoencoder: a grid of 256 to 1024 tokens describes a $256 \times 256$ image and can be continued from a text prompt.

References#

• A. van den Oord, N. Kalchbrenner, K. Kavukcuoglu. Pixel Recurrent Neural Networks. ICML, 2016. (2016a)
• A. van den Oord, N. Kalchbrenner, O. Vinyals, L. Espeholt, A. Graves, K. Kavukcuoglu. Conditional Image Generation with PixelCNN Decoders. NeurIPS, 2016. (2016b)
• T. Salimans, A. Karpathy, X. Chen, D. P. Kingma. PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture Likelihood and Other Modifications. ICLR, 2017.
• H. Larochelle, I. Murray. The Neural Autoregressive Distribution Estimator. AISTATS, 2011.
• M. Germain, K. Gregor, I. Murray, H. Larochelle. MADE: Masked Autoencoder for Distribution Estimation. ICML, 2015.
• I. Gulrajani, K. Kumar, F. Ahmed, A. A. Taiga, F. Visin, D. Vazquez, A. Courville. PixelVAE: A Latent Variable Model for Natural Images. ICLR, 2017.