---
title: PixelSNAIL Architecture
url: https://www.emergentmind.com/topics/pixelsnail-architecture
type: topic
---

# PixelSNAIL Architecture

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 [1712.09763].

## 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 [1712.09763].

## 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 $1 + 2\sum_{l=1}^L d_l$ per spatial axis. For six layers with the aforementioned dilations, the receptive field covers $127 \times 127$ pixels, which exceeds the full spatial extent for $32 \times 32$ images.

### B. Gated Residual Blocks

The gated block processes input $h_\mathrm{in} \in \mathbb{R}^{H \times W \times D}$:
1. Applies a masked convolution producing $2D$ channels, split as $(A, B)$:
   \[
   [A, B] = \mathrm{Conv}_{\mathrm{mask}}(h_\mathrm{in})
   \]
2. Computes the gated linear unit (GLU):
   \[
   \mathrm{GLU}(A, B) = A \odot \sigma(B)
   \]
   where $\sigma$ denotes the sigmoid.
3. Adds residual skip connection and (optionally) applies layer normalization:
   \[
   h_\mathrm{out} = h_\mathrm{in} + \mathrm{LayerNorm}(\mathrm{GLU}(A, B))
   \]

### C. Self-Attention Modules

Self-attention is periodically applied throughout the stack:
1. The reshaped state $h \in \mathbb{R}^{HW \times D}$ is linearly projected to queries $Q$, keys $K$, and values $V$.
2. Causal attention is enforced by applying a mask $M$ that prevents index $i$ from attending to $j \geq i$:
   \[
   \mathrm{Attn}(Q, K, V) = \mathrm{softmax}\Bigl(\frac{QK^\top}{\sqrt{d}} + M\Bigr) V
   \]
3. The output is reshaped and combined with the residual, followed by layer normalization. Multi-head attention is typically applied with $4$ or $8$ heads.

### D. Pixel-Coordinate Embeddings

For spatial awareness, fixed or learned row and column embeddings $e_{r,c} = E_{\mathrm{row}}(r) + E_{\mathrm{col}}(c)$ 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 $32 \times 32$ images typically operates as follows:
- Inputs are one-hot encoded ($x \in \{0, \ldots, 255\}^3$).
- Embedding via masked $1 \times 1$ convolution to $D=128$ channels: $32 \times 32 \times 128$.
- $B$ residual blocks are stacked, with $\lfloor B / K \rfloor$ attention modules interleaved, without modifying spatial or channel dimensions.
- The output is projected via a masked $1 \times 1$ convolution to logits per-pixel and per-channel, yielding $C \times 256$ logits. This represents the distribution over possible intensity values in each channel.

As the stack grows to cover the entire spatial field (e.g., $L$ dilated layers for $32 \times 32$), 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 [1712.09763].

## 5. Computational Considerations and Parallelization

Causal convolutions are fully parallelizable over the spatial grid, running at $O(H W D F^2)$ per layer, where $F=3$ for $3 \times 3$ filters. Causal masking modifies only the convolution kernel support. Self-attention incurs $O((HW)^2 d)$ time and memory per module; for $32 \times 32$ images, this is $O(10^6 D)$. 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     | $O(H W D F^2)$             | Many       |
| Self-attention  | $O((HW)^2 d)$              | Few blocks |
| 1×1 conv (embed,f) | $O(H W D^2)$            | Begin/end  |

## 6. Empirical Results and Design Ablations

On CIFAR-10 ($32 \times 32 \times 3$), PixelSNAIL achieves bit-per-dimension scores in the range $2.85$–$2.87$, surpassing strong PixelCNN baselines (≈$3.00$ bpd). On ImageNet $32 \times 32$, results around $3.42$ bpd are reported.

Ablation studies detail the importance of architectural features:
- Removing self-attention increases CIFAR-10 bpd by approximately $0.03$–$0.05$.
- Reducing kernel size from $3\times3$ to $1\times1$ increases bpd by around $0.05$.
- Omitting weight normalization results in slower convergence and a degraded bpd by about $0.02$.

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 [1712.09763].

Source: https://www.emergentmind.com/topics/pixelsnail-architecture