OpenAI, Week 9-10 // Fangirling Over DALL-E

OpenAI rang in the new year by unveiling two exciting projects, DALL-E and CLIP. Both are awesome and deserving of separate blog posts, so I'll focus on DALL-E for today. The page about DALL-E on the OpenAI blog does an amazing job showing off its capabilities and I highly encourage everyone to check it out. It's a bit light on the technical details - probably because the paper has yet to be released. It does briefly discuss the general idea of the model, which I've sketched below:

My scribbled high-level diagram of DALL-E.

Basically, DALL-E is a giant Transformer decoder that takes a combined text-image input and outputs a generated image. It's trained on image-caption datasets and tokenizes each separately; the captions are tokenized using byte-pair encoding (subword data compression for text), while the images are fed into the encoder of something like a VQ-VAE (which discretizes the latent space). Some or potentially all of the image is masked, starting from the lower right corner. 

After concatenating the two types of tokens, you have the input for DALL-E, which tries to predict more tokens representing the masked part of the image. Each image token is able to attend to all text tokens. The image tokens are allowed to attend to previous image tokens based on certain patterns (see OpenAI's Sparse Transformer paper), with different patterns in different layers. 

Once the new tokens are generated, the image part of the sequence is fed into the decoder of the VQ-VAE type thing to convert back into an image, which has either been filled in or generated from scratch, depending on how much was masked in the beginning.

Quick note: I don't know exactly what the VQ-VAE type thing is here. There is a technical footnote on the blog post about this: "Similar to VQVAE, each image is compressed to a 32x32 grid of discrete latent codes using a discrete VAE that we pretrained using a continuous relaxation. We found that training using the relaxation obviates the need for an explicit codebook, EMA loss, or tricks like dead code revival, and can scale up to large vocabulary sizes." Hence my diagram may not be strictly accurate because I don't fully understand how their discrete VAE works.

At first, I thought the model was a Transformer-on-VAE sandwich, with all the parameters trained simultaneously. But no, it seems the VAE is pre-trained. I don't know exactly what their loss function is, but I guess it's purely based on the tokens, not on comparing the input vs. output image. That makes it even more impressive, I think; it shows how good the Transformer is at generating these tokens, that they don't introduce too much noise when converted into actual images from a pre-trained VAE decoder. I'm also curious about how easy it would be to flip the model i.e. turn it into a caption generator. I suspect you' have to use something fancier than a byte-pair encoding/decoding in this case. Anyway, very cool stuff, and I look forward to reading the paper when it hits the arXiv.