08. High-Resolution Image Synthesis with Visual Tokens

GitHub에서 수정

The product is a button that returns a fresh, photorealistic picture of a mountain valley or a rainy city street, at 1024x1024 or 2048x2048, in about five seconds. There is no prompt in version one, but text prompts are on the roadmap. The hard part is arithmetic. A 2048px RGB image holds about 12.6 million numbers, and any model that treats each of them as a decision either runs out of memory or runs for minutes. Everything in this design follows from one move: compress the picture into a short grid of discrete symbols, learn to write that grid with a Transformer, and let a separate decoder paint the pixels back.

The core tension is how hard to compress. Every halving of the token grid side cuts generation work by four or more, but whatever detail the tokenizer drops at encode time can never come back. A good answer in an interview spends most of its time on that dial and on the system pieces that let you turn it without rebuilding everything.

The notes follow the book’s seven-step framework, then add follow-ups, interview questions, a 45-minute walkthrough and review questions.

1. Clarifying requirements

Functional: scenery only (forgiving, since foliage and water hide small errors and no identity must be preserved), unconditional now with text prompts later, two output sizes a factor of 4 apart in area, and about 5 million curated scenes plus unlabeled public images. Non-functional: about five seconds per image with a user watching a spinner, low GPU cost per image for a consumer feature, real diversity (ten presses should give ten different scenes, so mode collapse is a product bug), and safety and provenance, since even scenery models can memorize copyrighted photos.

2. Framing the problem as an ML task

Input: a request, optionally carrying a seed and a size. Output: an RGB image. The ML category is unconditional generative modeling: learn \( p(x) \) over scene images and sample from it.

Four model families are on the table. The useful question for each is what it costs per image at 1024px and whether its latent stays meaningful at that size.

The VAE failure is posterior collapse: the KL term pulls \( q(z \mid x) \) toward the prior, and once the decoder is strong enough to model images alone, the cheapest move is to make the latent carry nothing. High resolution needs a big decoder, so it bites exactly where the latent matters most. Latent diffusion (Stable Diffusion style) removes most of the pixel-diffusion cost and deserves a benchmark slot, but the token route wins here because it reuses the language-model stack end to end and text conditioning later is just a longer prefix.

So the system is two separately trained models:

1. An image tokenizer (VQ-VAE with VQGAN losses) that turns an image into a grid of integer codes and back.
2. An image generator (decoder-only Transformer) that models the probability of a code sequence.

    flowchart LR
  accTitle: End-to-end view of offline training and online generation
  accDescr: Offline, curated images train the tokenizer, the tokenizer encodes the dataset into cached token sequences, and those sequences train the Transformer generator. Online, a request goes to the generator, the token grid is decoded to a 1024 pixel image, optionally upscaled to 2048, filtered and returned.
  subgraph Offline
    D[Curated scene images] --> TT[Train tokenizer]
    TT --> ENC[Encode dataset once]
    ENC --> CACHE[(Token cache)]
    CACHE --> TG[Train generator]
  end
  subgraph Online
    REQ[User request] --> GEN[Generator service]
    GEN --> DEC[Decoder service]
    DEC --> SR{"2048 requested?"}
    SR -->|yes| UP[Super-resolution]
    SR -->|no| SAFE[Safety and watermark]
    UP --> SAFE
    SAFE --> OUT[Image to user]
  end
  TT -. decoder weights .-> DEC
  TG -. generator weights .-> GEN
  

How many tokens is an image

The token count is \( L = (H/f) \times (W/f) \) for downsampling factor \( f \). At 1024px, \( f = 16 \) gives a 64 x 64 grid (4,096 tokens), \( f = 32 \) gives 1,024 tokens and \( f = 64 \) gives 256; each count quadruples at 2048px. VQGAN used \( f = 16 \) at 256px. Holding 256 tokens at 1024px means each code stands for a 64x64 block, about 12,000 numbers in one integer, which only works because a GAN-trained decoder invents plausible texture. Fine for scenery, bad for text-in-image.

3. Data preparation

Sources

The curated 5M set (known licenses) trains the generator, because it defines what outputs should look like. The tokenizer learns reconstruction, not taste, so it benefits from breadth: a web-scale pool such as LAION-400M gives a richer codebook, after license review and removal of known-bad content. The generator then learns the target distribution without inheriting web noise like memes, watermarks and screenshots.

Filtering and cleaning

• Resolution floor. Drop images whose short side is below the native size; upsampled inputs teach blur.
• Quality. Variance-of-Laplacian blur score, JPEG artifact detection, and a CLIP-based aesthetic regressor.
• Domain. Zero-shot CLIP against prompts like “a scenery photo” and “a city street” removes screenshots and product shots.
• Watermarks and text. Otherwise the generator learns to draw stock-photo logos.
• Deduplication by perceptual hash or embedding search, which cuts memorization and stops famous spots from dominating.
• Safety and privacy. NSFW classifier, blur bystander faces and license plates.
• Balance. Cluster embeddings and resample by scene type, season, time of day and region, or the model paints only sunsets.

Modality-specific preprocessing

Scale pixels to \([-1, 1]\), resize the short side to 1024 and crop square (random crops for the tokenizer). Horizontal flips are safe for scenery; vertical flips and strong color jitter are not, since skies are blue and up. Fix EXIF orientation before cropping.

Storage

Raw images go to object storage as sharded tar or WebDataset files so training reads large sequential blobs. After the tokenizer is trained, each image is encoded once and its code grid stored as compact integer arrays. The generator never sees pixels during training.

    flowchart TB
  accTitle: Data pipeline from raw images to cached token sequences
  accDescr: Raw images pass through resolution, quality, domain, safety and deduplication filters, are normalized and cropped, then stored as shards. After tokenizer training the shards are encoded to integer code grids and cached for generator training.
  RAW[Raw image sources] --> RES[Resolution floor]
  RES --> QUAL[Blur, artifact, aesthetic filters]
  QUAL --> DOM[CLIP domain filter]
  DOM --> SAFE[NSFW, face and plate handling]
  SAFE --> DEDUP[Near-duplicate removal]
  DEDUP --> PREP[Normalize, resize, crop]
  PREP --> SHARDS[(Image shards)]
  SHARDS --> TOKTRAIN[Tokenizer training]
  TOKTRAIN --> ENCODE[Encode each image once]
  SHARDS --> ENCODE
  ENCODE --> TOKCACHE[(Code grids as uint16)]
  

Worked example: storage for the token cache. With 256 codes per image and a codebook of 8,192 entries, each code fits in a 16-bit integer. 5,000,000 images x 256 codes x 2 bytes = 2.56 GB. The raw JPEGs at about 500 KB each total about 2.5 TB. The token cache is a thousand times smaller and fits in host RAM on one training node, which removes the data loader as a bottleneck. At \( f = 32 \) (1,024 codes) it is still only 10 GB.

Pitfalls

• Ordering trap. The token cache depends on the tokenizer weights; retraining the tokenizer invalidates every cached sequence and generator checkpoint. Version them together.
• Crop skew. If training used center crops and the FID reference set uses full frames, FID shifts for reasons unrelated to the model.

4. Model development

Architecture

Tokenizer: encoder, quantizer, decoder

The encoder is a convolutional network (residual blocks with downsampling, often with one self-attention block at the lowest resolution) that maps an image \( x \in \mathbb{R}^{H \times W \times 3} \) to a grid of vectors \( z_e \in \mathbb{R}^{h \times w \times d} \), with \( h = H/f \).

The quantizer holds a codebook \( \mathcal{Z} = \{e_k\}_{k=1}^{K} \) of \( K \) learned vectors in \( \mathbb{R}^{d} \). Each grid vector is snapped to its nearest entry:

Here \( z_e(x)_{ij} \) is the encoder vector at grid position \( (i, j) \), \( e_k \) is codebook entry \( k \), and \( k^{*}_{ij} \) is the integer token stored for that position. The quantizer has no weights other than the codebook itself.

Discretization buys two things: the decoder sees only one of \( K \) vectors per position, so it cannot route around the bottleneck, and the generator gets a \( K \)-way classification target instead of regressing continuous vectors toward a blurry mean.

The decoder mirrors the encoder with upsampling blocks and is the only tokenizer part deployed for serving.

Generator: decoder-only Transformer

The generator reads codes in raster order and predicts the next one: code embedding, projection to \( d_{\text{model}} \), position information, \( N \) pre-norm causal blocks of attention plus MLP, and a linear head over \( K \) codes. A 1D learned position works for a fixed grid, but row-plus-column embeddings or 2D rotary embeddings tell the model that token \( t \) and token \( t + w \) are vertical neighbors. Each attention head computes

where \( Q, K, V \) are the query, key and value projections of the token states, \( d_k \) is the per-head width, and \( M \) is the causal mask with \( -\infty \) above the diagonal so position \( t \) sees only positions before it.

    flowchart LR
  accTitle: Tokenizer and generator internals
  accDescr: On the left the tokenizer encoder downsamples an image to a feature grid, the quantizer snaps each vector to a codebook entry and emits integer codes, and the decoder upsamples codebook vectors back to pixels. On the right the generator embeds previous codes, adds 2D positions, runs N causal Transformer blocks and outputs a softmax over K codes.
  subgraph Tokenizer
    IMG[Image 1024x1024x3] --> E["Conv encoder, downsample by f"]
    E --> ZE["Feature grid h x w x d"]
    ZE --> Q["Nearest codebook entry"]
    CB[(Codebook K x d)] --> Q
    Q --> IDS["Code grid of integers"]
    IDS --> LK[Codebook lookup]
    CB --> LK
    LK --> DCD["Conv decoder, upsample by f"]
    DCD --> REC[Reconstructed image]
  end
  subgraph Generator
    PREV[Previous codes] --> EMB[Embedding and projection]
    EMB --> POS[Add 2D positions]
    POS --> BLK["N causal blocks: attention plus MLP"]
    BLK --> HEAD["Linear head, softmax over K"]
    HEAD --> NEXT[Next code]
  end
  NEXT -. full grid .-> LK
  

Design knobs

Raising \( f \) cuts tokens but makes the decoder hallucinate more. A larger codebook \( K \) is more expressive but leaves more dead entries and a harder softmax. Lower code dimension \( d \) (ViT-VQGAN uses 32 or less) often improves usage. A deeper generator fixes global layout first, at a per-step latency cost, and a bigger decoder sharpens texture from the same codes. Rule of thumb: take the largest \( f \) whose reconstruction FID (rFID) is well below the target generation FID, since the generator can never beat its decoder.

Training

Training runs in two stages that share nothing but the code format.

    flowchart TB
  accTitle: Two-stage training flow
  accDescr: Stage one trains encoder, codebook and decoder with reconstruction, codebook, commitment, perceptual and patch adversarial losses, monitoring code usage. A frozen tokenizer then encodes the curated dataset. Stage two trains the Transformer with next-code cross-entropy, evaluated by FID on decoded samples.
  subgraph S1["Stage 1: tokenizer"]
    X[Image batch] --> ENC1[Encoder]
    ENC1 --> VQ[Quantizer with straight-through gradient]
    VQ --> DEC1[Decoder]
    DEC1 --> L1["L1 or MSE plus perceptual"]
    DEC1 --> DISC[Patch discriminator]
    VQ --> LQ[Codebook and commitment terms]
    VQ --> USE[Code usage monitor]
  end
  S1 --> FREEZE[Freeze tokenizer, encode curated set]
  FREEZE --> S2
  subgraph S2["Stage 2: generator"]
    SEQ[Cached code sequences] --> TR[Causal Transformer]
    TR --> CE[Next-code cross-entropy]
    TR --> SAMP[Periodic sampling and FID]
  end
  

Stage 1: tokenizer objectives

The loss is a weighted sum of four parts.

Reconstruction. Pixel MSE or L1 between the image and its reconstruction:

with \( x_i \) and \( \hat{x}_i \) the original and reconstructed pixel values and \( n \) the number of pixel values. L1 is often preferred because MSE punishes outliers and pushes toward blur.

Codebook and commitment.

The first term moves the selected codebook entries toward the encoder outputs, with \( \text{sg}[\cdot] \) (stop-gradient) freezing the encoder for that term; the second, weighted by \( \beta \) (about 0.25), keeps encoder outputs committed to their codes. An exponential moving average codebook update is a more stable replacement for the first term.

Straight-through gradient. The argmin has no useful gradient, so the backward pass treats quantization as the identity:

The value equals \( z_q \), but its gradient with respect to \( z_e \) is the identity, so the decoder’s gradient reaches the encoder. Only selected entries are updated, which is why dead codes appear.

Perceptual loss.

where \( \phi_l \) is the layer-\( l \) activation map of a frozen network such as VGG and \( C_l H_l W_l \) its size (LPIPS adds learned channel weights). A one-pixel shift costs a lot of MSE but almost nothing here, so this term rewards right texture over exact alignment.

Patch adversarial loss. A convolutional discriminator \( D \) scores overlapping patches as real or reconstructed, and the tokenizer minimizes

where \( \hat{x} \) is the reconstruction. Judging patches focuses the critic on local texture, exactly what the other losses leave soft.

Total and balancing.

The \( \lambda \) values are tuned. VQGAN sets \( \lambda_{\text{adv}} \) adaptively from the ratio of reconstruction and adversarial gradient norms at the decoder’s last layer, and the discriminator is usually switched on only after tens of thousands of steps, since an early critic wins too easily.

Stage 2: generator objective

The image distribution is factored over the raster-ordered code sequence \( s = (s_1, \dots, s_L) \):

where \( s_t \) is the code at position \( t \), \( s_{<t} \) all earlier codes, and \( \theta \) the Transformer weights. Teacher forcing scores all \( L \) positions in one forward pass. Prepend a learned start token so the first code is modeled too, and once text arrives, drop the caption about 10 percent of the time to enable classifier-free guidance.

Efficiency techniques

• bf16 mixed precision, but keep codebook distances in fp32, where underflow flips assignments.
• FSDP or ZeRO sharding for a 0.5 to 3B generator; tensor parallelism only above that.
• Activation checkpointing in the tokenizer, whose full-resolution feature maps dominate memory.
• FlashAttention in the generator once sequences reach 1,024 tokens.

Worked example: generator training compute. Take a 1B-parameter generator and 256 tokens per image. One pass over 5M images is \( 1.28 \times 10^{9} \) tokens. Train for 20 epochs: \( N = 10^{9} \), \( D = 2.56 \times 10^{10} \) tokens. Using the standard estimate

An H100 at about 990 TFLOP/s dense bf16 peak and 40 percent utilization delivers about \( 4 \times 10^{14} \) FLOP/s, so this is \( 3.8 \times 10^{5} \) GPU-seconds, about 105 GPU-hours, or under a day on 8 GPUs. The tokenizer is the heavier job, since its convolutions run at full resolution over possibly hundreds of millions of web images. Twenty epochs is also a lot of repetition, so watch validation loss for memorization.

Failure modes and fixes

• Codebook collapse (few codes used, flat perplexity): EMA updates, restart dead codes from encoder outputs, normalized low-dimensional codes.
• Smeared reconstructions: perceptual and patch adversarial losses were missing or too weak.
• Checkerboard texture or loss spikes once the critic starts: adaptive weight, later critic start, upsample plus conv instead of transposed conv.
• Broken global layout: generator too shallow or missing 2D positions.
• Copied training images: duplicates or too many epochs; deduplicate, stop early, audit with nearest-neighbor search.

Sampling

Start from a start token, then for \( t = 1 \dots L \) run the Transformer on the prefix, adjust the logits, sample one code and append it. Reshape the result to \( h \times w \) and run the decoder once. Temperature rescales logits \( z \):

where \( \tau < 1 \) sharpens toward likely codes and \( \tau > 1 \) flattens. Top-k keeps the \( k \) most likely codes; top-p keeps the smallest set reaching cumulative probability \( p \), so the candidate set is narrow for sky and wide for foliage. Greedy and beam search pick maximum-likelihood codes, which in images means flat, repetitive regions.

Worked example: latency budget at 1024px. Use a 1B-parameter generator in bf16 (2 GB of weights) on one H100 with about 3.35 TB/s of memory bandwidth. Decoding one token at batch size 1 must read all weights, so the floor is about 2 GB / 3.35 TB/s = 0.6 ms per step. Real kernels with launch overhead, sampling and the KV cache run closer to 5 ms per step.

The 256-token grid fits the 5-second target with headroom for queueing; the 1,024-token grid is borderline; native 2048 generation is out. This is the quantitative argument for “generate at 1024, upscale to 2048.”

Worked example: KV cache size. For a model with 24 layers and width 2048, each token stores a key and a value per layer: \( 2 \times 24 \times 2048 \times 2 \) bytes = 196,608 bytes, about 0.19 MB. A 256-token image needs about 50 MB of cache; a 1,024-token image about 200 MB. On an 80 GB GPU with 2 GB of weights you could hold a few hundred concurrent 1,024-token sequences, so batch size is limited by latency targets, not memory.

5. Evaluation

Offline metrics

Fréchet Inception Distance fits a Gaussian to Inception-v3 features of real and generated images and measures the distance between them:

where \( \mu_r, \Sigma_r \) and \( \mu_g, \Sigma_g \) are the means and covariances of real and generated features. The mean term catches systematic shifts (everything too dark), the covariance term lost variety. Use 10,000 to 50,000 samples, since FID is biased upward on small sets, and compute it on delivered 2048 outputs too, because Inception downsizes to 299px and hides fine-detail problems.

Inception Score uses only generated images:

where \( p(y \mid x) \) is the classifier’s label distribution for one image and \( p(y) \) its average over samples. It rewards confident, varied classes (“valley”, “lakeside” and “streetcar” exist in ImageNet) but never compares against real data. Precision and recall for generative models split fidelity from coverage, the clearest view of a temperature change.

Tokenizer metrics: rFID, PSNR and LPIPS on held-out reconstructions, plus code perplexity:

where \( \bar{p}_k \) is the share of validation positions assigned to code \( k \); a value far below \( K \) signals collapse. Add a memorization audit: nearest-neighbor search of generations against the training set.

Human evaluation

Raters see randomized side-by-side pairs and pick the more photographic and the more pleasing image. Report win rates with confidence intervals over a few hundred pairs, mix in real photos as a sanity check, and show grids of eight samples to judge diversity.

Online metrics

Regeneration rate (an immediate second press signals disappointment), downloads and shares, thumbs up or down, subscription conversion, churn, and p50 and p95 latency. They lag and are noisy, so they confirm launches rather than diagnose regressions.

6. Overall ML system design

Three stages with different hardware profiles sit behind one API. The generation service runs the Transformer: sequential, memory-bandwidth bound, latency linear in token count, throughput from batching. The decoding service runs one compute-bound convolutional pass in a fraction of a second; the encoder is never deployed. The super-resolution service runs a 2x upscaler (Real-ESRGAN style or a small diffusion upscaler) only for 2048 requests. Separate services pick their own GPU types, autoscale on their own queues and ship independently. The hop between them is cheap: a code grid is 512 bytes.

    sequenceDiagram
  accTitle: Request flow for a 2048 pixel image
  accDescr: The client calls the API gateway, which checks quota and enqueues the job. The generation service samples a code grid with KV cache and continuous batching, the decoding service turns it into a 1024 image, the super-resolution service upscales to 2048, a safety service filters and watermarks, and the client receives a URL from storage.
  participant C as Client
  participant G as API gateway
  participant Q as Job queue
  participant GEN as Generation service
  participant DEC as Decoding service
  participant SR as Super-resolution
  participant S as Safety and storage
  C->>G: generate size=2048 seed=optional
  G->>Q: enqueue job after quota check
  Q->>GEN: batch of jobs
  GEN->>GEN: sample 256 codes with KV cache
  GEN->>DEC: code grid 16x16
  DEC->>SR: 1024 image
  SR->>S: 2048 image
  S->>S: NSFW and memorization check, watermark
  S-->>C: signed image URL
  

Latency and cost optimizations

Because there is no prompt, any image answers any request, so off-peak capacity can fill a pool that serves most traffic instantly, with each pooled image consumed once.

Worked example: fleet sizing. 2 million images a day is 2,000,000 / 86,400 = 23 per second, about 70 at a 3x peak. One H100 batching 32 sequences of 256 tokens at 6 ms per step finishes 32 images every \( 256 \times 6 \) ms = 1.5 s, about 21 per second, so peak needs 4 generator GPUs (plan 6 to 8). The decoder at 32 images per 0.5 s batch needs 2 GPUs. If 30 percent ask for 2048, the upscaler sees 21 per second; at about 20 per second per GPU that is 2 to 3 GPUs. Roughly a dozen GPUs, dominated by the generator.

7. Deployment and monitoring

Safety and monitoring

• Output NSFW and violence classifiers before return; log and drop failures.
• Memorization check against an embedding index of the training set for a sampled fraction of outputs.
• Invisible watermark plus C2PA content credentials so images are labeled as AI-generated.
• Dashboards per stage: queue depth, p95 latency, GPU utilization, error rate. Model health: rolling FID on a daily sample, code-usage histogram of generated grids (a shift signals a bad deploy), safety-filter hit rate.

Rollout and feedback

Ship the tokenizer, its token cache and the generator as one versioned bundle, and route each request to a single bundle end to end, since codes from one tokenizer mean nothing to another decoder. A new bundle first runs in shadow on a copy of traffic to compare latency and the code-usage histogram, then takes a small canary share. Promotion requires rolling FID on daily samples and the regeneration rate to hold within a set margin of the current version. The previous bundle stays warm for rollback, and the pre-generated pool is refilled from the new bundle only after promotion.

Regenerate presses, downloads and thumbs feed a preference log. Frequently downloaded images can seed a small curated fine-tuning set, and repeated regenerations on one kind of scene point to gaps in the curated data. A retrain is triggered when rolling FID drifts upward for several days or when the curated set has grown enough to matter.

Follow-ups and extensions

• Text-to-image. Prefix caption tokens to the image codes and train on pairs, as Parti and the first DALL-E did; drop captions sometimes to enable classifier-free guidance.
• Aspect ratios. Variable grids with 2D positions and a size token make 16:9 just a 12 x 21 grid.
• Inpainting and outpainting. Generate only the missing codes; holes in the middle suit masked bidirectional models better than raster order.
• Better tokenizers. Residual quantization, lookup-free or finite scalar quantization without dead codes, ViT-based encoders.
• Faster generation. MaskGIT iterative parallel decoding, or VAR-style next-scale prediction from coarse to fine grids.
• Interactive refinement and personalization. Lock a region and resample the rest; bias seeds or adapters from regenerate and download signals.

Interview questions

Clarifying

Which content categories and resolutions must version one support? I would ask because both set the tokenizer’s downsampling factor. Scenery tolerates large patches, so a 16 x 16 code grid at 1024px is defensible; faces or text would push me toward 32 x 32 or finer, which quadruples generation time. The two output sizes tell me whether to train two generators or one plus an upscaler. With a factor of two between them, one native size plus super-resolution is the cheaper plan.

How much training data exists, and can we use unlabeled public images? The two models want different data. The tokenizer only learns to reconstruct, so it benefits from hundreds of millions of diverse images with no labels. The generator defines the output distribution, so it should train on the curated set. If licensing forbids web data, I would train the tokenizer on the curated 5M and accept a slightly worse codebook.

Is conditioning on text coming soon? It changes the architecture choice. If it is, I want a decoder-only Transformer that can take a prefix of condition tokens now, and I want condition dropout baked into the training recipe so guidance works later. If conditioning is never coming, a GAN or a small latent diffusion model becomes more competitive. I would also ask whether five seconds is a p50 or a hard p99, since the latter forces fewer tokens or a pre-generated pool.

Design and trade-offs

Trade-off: autoregressive tokens versus latent diffusion. Both run in a compressed space, so the gap is smaller than with pixel diffusion. Autoregressive models reuse LLM infrastructure, give exact likelihoods, and add text conditioning naturally, but they decode sequentially and suffer from exposure bias. Latent diffusion uses a fixed number of steps independent of resolution, often wins on sample quality, and parallelizes over pixels, but needs more steps for quality and its guidance math is more complex. For this chapter’s constraints I would ship the token model and keep a diffusion baseline in the benchmark suite.

Trade-off: downsampling factor f. Doubling \( f \) cuts tokens by 4 and attention cost by up to 16, which is the difference between 20 seconds and 1.3 seconds at 1024px. The price is information: each code must stand for 4 times as many pixels, so the decoder invents more and fine structure like thin branches or window grids degrades. I choose by plotting rFID against \( f \) and taking the largest \( f \) whose rFID is well under the target generation FID.

Trade-off: native 2048 generation versus 1024 plus 2x upscaling. Native generation gives globally coherent fine detail but needs 4 times the tokens and roughly 4 to 16 times the compute. A super-resolution model sharpens and adds plausible texture but cannot fix composition errors at the small scale and sometimes invents repeated textures. With a 5-second budget and scenery content, upscaling wins; for print-quality products I would revisit.

Why quantize instead of predicting continuous latents? Discrete codes give the generator a categorical target where cross-entropy can express multi-modal uncertainty: the next patch might be sky or tree, and softmax can put mass on both. Regressing a continuous vector with MSE averages those modes into a muddy mix. Quantization also prevents the decoder from ignoring the latent. Recent work does model continuous tokens with diffusion heads, but that adds complexity.

Deep dives

Walk through the straight-through estimator and why codes die. Nearest-neighbor lookup has no useful gradient, so in the backward pass we treat quantization as the identity and pass the decoder’s gradient straight to the encoder output. The codebook itself only learns through the codebook term, and only for the entries actually selected. An entry that no encoder output is near never gets selected and never moves, so it stays dead. Fixes are EMA updates, reinitializing dead entries from recent encoder outputs, and normalizing codes and features so distances are comparable.

What does each tokenizer loss contribute, and what breaks without it? Pixel loss keeps color and layout right but alone produces blur. The codebook and commitment terms keep encoder outputs and codes close; without commitment, encoder outputs drift and assignments become unstable. Perceptual loss rewards matching texture and structure; without it, edges are soft. The patch adversarial loss removes the residual mush; without it, 1024px reconstructions look painted. Too much adversarial weight causes artifacts and instability, so its weight is balanced by gradient norms.

How would you add 2D awareness to a raster-order Transformer? Raster order makes the token directly above position \( t \) sit \( w \) steps back, and a 1D position embedding does not tell the model that. I would use separate row and column embeddings summed together, or 2D rotary embeddings that rotate query and key by row and column angles. Local attention windows that include the previous row help both quality and speed. Some systems also try alternative scan orders, but raster with 2D positions is the simple default.

Estimate generation latency for a 1,024-token grid on one H100. A 1B bf16 model is 2 GB, and at 3.35 TB/s reading it takes 0.6 ms, the floor for one decoding step. Realistic per-step time with overhead is around 5 ms, so 1,024 steps take about 5 seconds before decoding and upscaling. That busts a 5-second budget, which pushes me to 256 tokens, int8 weights to halve memory traffic, or parallel decoding. Batching does not reduce single-request latency, it only raises throughput.

Estimate the GPU fleet for 2 million images per day. That is 23 per second on average and about 70 at a 3x peak. One H100 batching 32 sequences of 256 tokens at about 6 ms per step finishes 32 images in about 1.5 seconds, around 21 per second, so peak needs 4 generator GPUs, 6 to 8 with headroom. The decoder and upscaler need 2 to 3 GPUs each. Roughly a dozen GPUs total, and a pre-generated pool could shave the peak further.

Operations and failure modes

After a deploy, users say images look repetitive. How do you debug? First check sampling configuration: a temperature or top-p change is the most common cause. Then compare the histogram of generated codes against the training distribution; a collapse to fewer codes points at the generator, while normal histograms with bad images point at the decoder. Compute precision and recall on a fresh sample; precision up and recall down confirms a diversity loss. Roll back if confirmed, then add a code-entropy alarm to prevent a repeat.

How do you detect and reduce memorization? Deduplicate training data aggressively, since duplicated images are the main source of verbatim copies. Limit epochs and watch the train-validation gap. At evaluation and in production, embed a sample of outputs and search the training index for near matches above a threshold, and block outputs that match. Low temperature raises copy risk, so check it at the operating point actually used.

You retrained the tokenizer. What else must change? Every cached code sequence is now meaningless because code IDs map to different vectors. The generator has to be retrained on newly encoded data, and the decoding service has to deploy the new decoder at the same moment as the new generator. I version tokenizer, token cache and generator as one bundle and route traffic by bundle version so a half-upgraded state cannot occur.

A 45-minute walkthrough

Review questions

Answer each question in your own words first, then open the answer to check.

Concepts and definitions

1. What is the downsampling factor \( f \), and how does it set the number of tokens per image?

2. What does the quantizer do, and what are its trainable parameters?

3. What is code perplexity, and what does a value far below \( K \) mean?

4. What is rFID, and why does it set a ceiling for the whole system?

5. How do top-k and top-p differ, and why does top-p suit images with both sky and foliage?

Mechanisms and math

6. Explain the two terms of \( \mathcal{L}_{\text{vq}} \) and the role of the stop-gradient.

7. Why does VQGAN switch on the discriminator late and set its weight adaptively?

8. Why can the generator score every position in one forward pass during training but needs \( L \) passes at sampling time?

9. How does temperature reshape the next-code distribution, and what happens in the limit of very low temperature?

10. What does the KV cache store, and why does it change the cost of each decoding step?

Numbers and estimates

11. How big is the token cache for 5 million images at 256 codes, and at 1,024 codes?

12. Use \( C \approx 6ND \) to estimate generator training compute for 1B parameters, 256 tokens per image, 5M images and 20 epochs.

13. How much KV cache does one image take for a model with 24 layers and width 2048 in 16-bit precision?

Design reasoning and trade-offs

14. Why run generation, decoding and super-resolution as three separate services?

15. Why can this product serve most traffic from a pre-generated pool, and what rule keeps that safe?

16. Why are horizontal flips allowed in preprocessing but vertical flips and strong color jitter are not?

17. Why compute FID on the delivered 2048px outputs as well as at 1024px?

Common mistakes

18. What goes wrong if the codebook distance computation runs in bf16?

19. What goes wrong if training used center crops but the FID reference set uses full frames?

20. What goes wrong if images below the resolution floor are upsampled and kept in training?