PixelSNAIL Architecture

• PixelSNAIL is an autoregressive generative model that employs a raster-scan causal factorization to model pixel dependencies in images.
• It interleaves dilated causal convolutions with self-attention and gated residual blocks to effectively capture both local and global spatial patterns.
• The architecture achieves state-of-the-art log-likelihoods on benchmarks like CIFAR-10 and ImageNet 32×32, demonstrating improved density estimation performance.

PixelSNAIL is an autoregressive generative model architecture designed for high-dimensional data such as images, where it models the joint distribution over all input pixels using a raster-scan causal factorization. It advances previous approaches by interleaving causal convolutions with self-attention modules, while leveraging gated residual blocks, positional embeddings, and normalization to improve upon the density estimation performance of models such as PixelCNN. PixelSNAIL attains state-of-the-art log-likelihoods on benchmarks including CIFAR-10 and ImageNet 32×32 images, substantiating the efficacy of its hybrid design principles (Chen et al., 2017).

1. Autoregressive Factorization and Causal Structure

PixelSNAIL models joint distributions $p(x)$ over all pixels $x = (x_1, \dots, x_N)$, where $N = H \times W \times C$, using a raster-scan autoregressive factorization:

$p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$

Each conditional $p(x_i \mid x_{<i})$ is estimated through a deep network that constructs a large receptive field over the context $x_{<i}$ via a stack of components:

• Causal (masked) convolutions with increasing dilation for local-to-global dependency coverage,
• Gated residual blocks for stabilization and expressive nonlinearity,
• Interleaved self-attention modules to directly capture long-range dependencies in the spatial grid.

This structured approach overcomes limitations of conventional RNN or basic convolutional factorization, allowing efficient and powerful modeling of local and global spatial patterns (Chen et al., 2017).

2. Core Architectural Components

A. Causal Convolutional Layers

Causal convolutions use a $3 \times 3$ filter, masked so that the convolution never accesses "future" pixels. Dilation is applied in powers of two (e.g., $\{1,2,4,8,16,32\}$). Stacking $L$ such layers with dilations $(d_1, \dotsc, d_L)$ yields a receptive field of $x = (x_1, \dots, x_N)$0 per spatial axis. For six layers with the aforementioned dilations, the receptive field covers $x = (x_1, \dots, x_N)$1 pixels, which exceeds the full spatial extent for $x = (x_1, \dots, x_N)$2 images.

B. Gated Residual Blocks

The gated block processes input $x = (x_1, \dots, x_N)$3:

1. Applies a masked convolution producing $x = (x_1, \dots, x_N)$4 channels, split as $x = (x_1, \dots, x_N)$5:

$x = (x_1, \dots, x_N)$6

2. Computes the gated linear unit (GLU):

$x = (x_1, \dots, x_N)$7

where $x = (x_1, \dots, x_N)$8 denotes the sigmoid.

3. Adds residual skip connection and (optionally) applies layer normalization:

$x = (x_1, \dots, x_N)$9

C. Self-Attention Modules

Self-attention is periodically applied throughout the stack:

1. The reshaped state $N = H \times W \times C$0 is linearly projected to queries $N = H \times W \times C$1, keys $N = H \times W \times C$2, and values $N = H \times W \times C$3.
2. Causal attention is enforced by applying a mask $N = H \times W \times C$4 that prevents index $N = H \times W \times C$5 from attending to $N = H \times W \times C$6:

$N = H \times W \times C$7

1. The output is reshaped and combined with the residual, followed by layer normalization. Multi-head attention is typically applied with $N = H \times W \times C$8 or $N = H \times W \times C$9 heads.

D. Pixel-Coordinate Embeddings

For spatial awareness, fixed or learned row and column embeddings $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$0 are added to the feature map before processing. This mechanism supports the model's ability to distinguish different locations, mitigating confusion between spatially distant but semantically distinct regions.

3. Network Shape, Channelization, and Flow

PixelSNAIL for $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$1 images typically operates as follows:

• Inputs are one-hot encoded ($p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$2).
• Embedding via masked $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$3 convolution to $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$4 channels: $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$5.
• $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$6 residual blocks are stacked, with $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$7 attention modules interleaved, without modifying spatial or channel dimensions.
• The output is projected via a masked $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$8 convolution to logits per-pixel and per-channel, yielding $p(x) = \prod_{i=1}^N p(x_i \mid x_{<i})$9 logits. This represents the distribution over possible intensity values in each channel.

As the stack grows to cover the entire spatial field (e.g., $p(x_i \mid x_{<i})$0 dilated layers for $p(x_i \mid x_{<i})$1), each pixel’s conditional distribution leverages the full context of preceding pixels.

4. Normalization and Stabilization

Each residual block—whether convolutional or attential—applies LayerNorm after the GLU, ensuring stable activations throughout the deep stack. Weight normalization is applied (in the style of Salimans & Kingma, 2016) to convolutional filters and Q/K/V projections. This approach accelerates convergence and partially protects against gradient-related instabilities in deep, recurrently unrolled autoregressive models (Chen et al., 2017).

5. Computational Considerations and Parallelization

Causal convolutions are fully parallelizable over the spatial grid, running at $p(x_i \mid x_{<i})$2 per layer, where $p(x_i \mid x_{<i})$3 for $p(x_i \mid x_{<i})$4 filters. Causal masking modifies only the convolution kernel support. Self-attention incurs $p(x_i \mid x_{<i})$5 time and memory per module; for $p(x_i \mid x_{<i})$6 images, this is $p(x_i \mid x_{<i})$7. To manage the quadratic scaling, only a few attention blocks are typically inserted.

A summary of per-layer complexities:

Component	Complexity per Layer	Typical Use
Causal conv	$p(x_i \mid x_{<i})$8	Many
Self-attention	$p(x_i \mid x_{<i})$9	Few blocks
1×1 conv (embed,f)	$x_{<i}$0	Begin/end

6. Empirical Results and Design Ablations

On CIFAR-10 ($x_{<i}$1), PixelSNAIL achieves bit-per-dimension scores in the range $x_{<i}$2–$x_{<i}$3, surpassing strong PixelCNN baselines (≈$x_{<i}$4 bpd). On ImageNet $x_{<i}$5, results around $x_{<i}$6 bpd are reported.

Ablation studies detail the importance of architectural features:

• Removing self-attention increases CIFAR-10 bpd by approximately $x_{<i}$7–$x_{<i}$8.
• Reducing kernel size from $x_{<i}$9 to $3 \times 3$0 increases bpd by around $3 \times 3$1.
• Omitting weight normalization results in slower convergence and a degraded bpd by about $3 \times 3$2.

Interleaving dilated convolution (rapidly expanding receptive fields) with full-image self-attention grants PixelSNAIL efficient local and global context modeling. The stabilization mechanisms (gated residuals, coordinate embeddings, normalization) collectively yield robust training dynamics and improved likelihood estimation, underpinning PixelSNAIL's state-of-the-art performance on small images (Chen et al., 2017).