Skip to content

VQ-VAE + PixelCNN Prior — Visual Guide

The VQ-VAE turns each ECG into a sequence of discrete code indices (stage 1), and a separate autoregressive model learns which index sequences are likely (stage 2). Generation samples indices from the prior and decodes them. Code: src/ecgen/families/vqvae/, losses: src/ecgen/families/vqvae/losses.py. Written reference: VQ-VAE + Prior.

Two stages at a glance

Stage 1VQVAELightning — learn the codebook, encoder and decoder Encoder VQ Decoder MSE + VQ loss Stage 2PriorLightning — learn p(codes) with the stage-1 model frozen Encoder + VQ (frozen) codes PixelCNN cross-entropy SamplePriorLightning.sample() PixelCNN codes Decoder

Stage 1: forward pass

input x (B, 12, 5000) Encoder logvar ✕ z VQcodebook512 × 64 z_q Decoder x̂ (B, 12, 5000) 37 5 5 212 88 401 37 … indices (B, 625), values in [0, 512) forward() returns (recon, vq_loss, indices)
Hover over a block for details.

The encoder and decoder are the shared backbone with latent_channels = 64:

Hover over a block for the exact layer. The faded logvar half is computed but not used.

Vector quantization

Each latent vector zt is replaced by the nearest of the 512 codebook vectors ek. The animation uses a 2-D toy space with 12 codes and 8 time steps:

1. encoder outputs zt (circles) · 2. find the nearest codebook vector (dashed) · 3. snap to it and record its index.

Stage 1: how the loss is computed

recon, vq_loss, indices = vqvae(x)
total, recon_loss, vq_loss = vqvae_loss(recon, x, vq_loss)   # total = recon_loss + vq_loss
\[ \mathcal{L} = \underbrace{\operatorname{MSE}(\hat{x}, x)}_{\texttt{recon\_loss}} + \underbrace{\operatorname{MSE}(\operatorname{sg}[z_q],\, z)}_{\text{trains the encoder}} + \beta \cdot \underbrace{\operatorname{MSE}(z_q,\, \operatorname{sg}[z])}_{\text{trains the codebook}}, \qquad \beta = \texttt{commitment\_cost} = 0.25 \]
∂ recon_loss: decoder → z_q straight-through copy → z Encoder z argmin ‖z − e‖+ lookup codebook z_q Decoder MSE(x̂, x) MSE(sg[z_q], z)→ encoder only β · MSE(z_q, sg[z])→ codebook only + ℒ purple = gradient path · sg[·] = .detach() · the argmin itself passes no gradient
Hover over the loss terms and the straight-through label.

β is on the codebook term here

The original VQ-VAE paper puts β on the encoder (commitment) term. In VectorQuantizer.forward, β multiplies the term that updates the codebook, even though that variable is named commitment_loss. The diagram shows what the code does.

Stage 2: the PixelCNN prior

PriorLightning loads the stage-1 checkpoint, freezes it, and turns every training ECG into its 625 code indices. PixelCNNPrior is then trained to predict each index from the ones before it.

VQ-VAEencode_to_indices c (B, 625) Embedding shift → 1 10 × CausalGatedBlock 1 2 4 8 512 Σ skips head (1×1 convs) logits cross-entropytarget = c (not shifted) targets
Hover over a block. Use PriorLightning.compute_loss(indices); never shift the targets yourself — the shift is inside the model.

Each gated block

x left pad d(causal) conv k=2, d tanh σ ⊙ h 1×1 res 1×1 skip + x′ skip
tanh(Wf ∗ x) ⊙ σ(Wg ∗ x): the sigmoid gate decides how much of the filter output passes. Hover over the padding and the conv.

Why logit t never sees code t

Hover over (or tap) a top node to see which codes it depends on. This toy stack has three layers with dilations 1, 2, 4; the real prior has ten (1 … 512).

Sampling

PixelCNNPrior.sample() fills the 625 positions left to right. At step t it runs the network on the codes so far, takes the logits at t, divides by temperature, and draws one code. Then the frozen VQ-VAE decodes all indices at once. Lower temperatures make the most likely code win more often.

The toy “prior” here prefers codes close to the previous one, so you can see the effect of temperature. Sampling 625 codes takes 625 network passes, which makes this the slowest generator of the three.