DRAW: teaching a network to draw step by step
Does attention help a network draw, and does it matter more as the images get harder?
Summary
We trained DRAW, a network that draws an image over several steps, with and without its attention window on three handwriting datasets of rising difficulty. Attention did not help on digits and only paid off on the hardest dataset, where it first blurred thin strokes until we made its window finer.
The scaled-up model drawing five Omniglot characters, one step per column (T = 20).
Key results
- 0.206 vs 0.213MNIST validation loss without vs with attention: attention does not help on digits
- 0.199 vs 0.201Omniglot validation loss with vs without attention: the only dataset where attention wins
- 0.112 → 0.099Omniglot reconstruction loss when the attention grid grew from 9×9 to 18×18
The problem
DRAW (Gregor et al., 2015) is a generative model that does not produce an image in one go. It draws it over a fixed number of steps, and at each step an attention window decides where on the canvas to read and where to write. The original paper showed that attention helps, but not how much it matters for different kinds of images. We asked whether attention improves the generated images, and whether its effect grows with the difficulty of the data, using three handwriting datasets: MNIST (digits), EMNIST Letters and Omniglot (1,623 characters from 50 alphabets, with thin, detailed strokes).
What we did
- Implemented DRAW in PyTorch: an LSTM encoder and decoder, a latent variable at every step, a canvas that is updated over T steps, and read and write attention through an N×N grid of Gaussian filters. A switch turns attention off, so the model reads and writes the whole canvas instead.
- Trained both variants on all three datasets (28×28 pixels, binarized), and scaled the model with the difficulty: from 10 steps, a 5×5 grid and a latent size of 100 on MNIST up to 20 steps, a 9×9 grid and a latent size of 256 on Omniglot. Training used Adam with batch size 128 for 30 to 50 epochs on Google Colab.
- Compared the variants on the validation loss (ELBO), split into reconstruction loss and KL divergence, and by looking at the reconstructed and generated images.
- Looked into why the attention model reconstructed Omniglot worse. Visualizing the attention window showed that a 9×9 grid has to stretch over the whole canvas, which blurs thin strokes.
- Trained a high-resolution variant with an 18×18 grid and 30 steps to test that explanation.
Results
On MNIST, the model without attention is slightly better (validation loss 0.206 vs 0.213) and reconstructs digits more accurately. On EMNIST Letters the two variants tie: the difference is 0.0003. On Omniglot, attention wins on the overall loss (0.199 vs 0.201), because it organizes the latent space better (lower KL), even though it reconstructs the images worse.
That worse reconstruction came from the attention grid being too small for thin strokes. With an 18×18 grid and 30 steps, the reconstruction loss dropped from 0.112 to 0.099, level with the model without attention (0.098), while keeping the better overall loss.
| Dataset | Loss, with attention | Loss, without | Reconstruction, with | Reconstruction, without |
|---|---|---|---|---|
| MNIST | 0.213 | 0.206 | 0.164 | 0.147 |
| EMNIST Letters | 0.192 | 0.192 | 0.115 | 0.114 |
| Omniglot (9×9 grid) | 0.199 | 0.201 | 0.112 | 0.098 |
| Omniglot (18×18 grid) | 0.199 | 0.201 | 0.099 | 0.098 |
What I learned
- An ablation answers “when”, not just “whether”. Attention made no difference on digits and only helped on the hardest dataset.
- Look inside the model when a number surprises you. Visualizing the attention window showed it stretching across the whole image, which explained the blurry strokes and pointed to the fix.
- Small differences need careful reading. On EMNIST the two variants differ by 0.0003, and each was trained once, so that is a tie, not a winner.