Research · Master's thesis
Adversarial Text-to-Image Generation using Contrastive Learning
Caltech-UCSD Birds-200-2011 (CUB-200-2011) · 200 species · 11,788 images · 10 captions per image
May 2022MSc Data Science & Artificial Intelligence, Saarland University
About the project
Text-to-image generation with a Generative Adversarial Network (GAN) had a blind spot: people describe the same bird in very different words, and the model drew different birds for them. I extended the Semantic-Spatial Aware Generative Adversarial Network (SSA-GAN) with contrastive learning (CL) so that captions of the same image are pulled together and captions of different images are pushed apart, both when the text and image encoders are pre-trained and while the generator trains.
The model is written in Python with PyTorch: a bi-directional Long Short-Term Memory (LSTM) text encoder and an Inception-v3 image encoder, together forming the Deep Attentional Multimodal Similarity Model (DAMSM), pre-trained with a contrastive loss. An SSA-GAN generator uses mask predictors to decide where each part of a caption applies. It was trained for 600 epochs on the Caltech-UCSD Birds-200-2011 dataset (CUB-200-2011) and evaluated with Fréchet Inception Distance (FID), Inception Score and R-precision.
The result had the lowest FID of all models I compared, 15.10, and raised R-precision from 73.55 for SSA-GAN as I reproduced it to 84.08: the images became more realistic and matched their captions more closely. The code and the trained models are public.
What makes it unique
The most realistic birds
Its images had the lowest FID of all the models I compared: 15.10, where lower means closer to real photos.
Pictures that match the words
R-precision, which measures whether an image fits its caption, rose from 73.55 for SSA-GAN, as I reproduced it, to 84.08.
A first, as of May 2022
To our knowledge, the first model to learn locally weighted image and text features together with contrastive learning.
Caption
"small bird, red head, black wings"
Text encoder
trained jointly
SSA-GAN
+ semantic mask predictor
Image
matches the caption
contrastive loss pulls matching caption ↔ image pairs together and pushes mismatched ones apart, so different wordings of the same scene give the same picture.
Three problems it tackles
a.Local semantics ignored
Conditional batch norm treated the whole feature map equally. SSA-GAN adds a weakly-supervised mask predictor.
b.Frozen text encoder
Trained end-to-end, so the text encoder learns joint text-image representations with the generator.
c.Caption disparity
People describe the same image in different words. Contrastive learning maximises mutual information between text and image.
Data and training
- Dataset
- Caltech-UCSD Birds-200-2011 (CUB-200-2011): 200 bird species, 11,788 images, 10 written descriptions per image
- Image size
- 256 × 256 pixels
- Text encoder
- Bi-directional Long Short-Term Memory (LSTM) with 256-dimensional embeddings
- Image encoder
- Inception-v3, pre-trained on ImageNet; together, both encoders form the Deep Attentional Multimodal Similarity Model (DAMSM)
- Contrastive loss
- Normalized Temperature-scaled Cross Entropy (NT-Xent), temperature 0.5, other pairs in the batch as negatives
- Adversarial loss
- Hinge loss
- Optimiser
- Adam: generator 0.0001, discriminator 0.0004, betas (0.0, 0.9)
- Training
- Batch size 24, 600 epochs on one GPU; the model from epoch 550 is evaluated
- Metrics
- Fréchet Inception Distance (FID; realism, lower is better), Inception Score (quality and variety), R-precision (image matches its caption)
Results on CUB birds
My model against published text-to-image models. Fréchet Inception Distance (FID) measures realism (lower is better), Inception Score quality and variety, and R-precision how well an image matches its caption.
| Model | Inception Score ↑ | FID ↓ | R-precision ↑ |
|---|---|---|---|
| StackGAN++ | 4.04 ± 0.06 | 15.30 | – |
| AttnGAN | 4.36 ± 0.03 | 23.98 | 67.82 ± 4.43 |
| AttnGAN + CL | 4.42 ± 0.05 | 16.34 | 69.64 ± 0.63 |
| ControlGAN | 4.58 ± 0.09 | – | – |
| SD-GAN | 4.67 ± 0.09 | – | – |
| DM-GAN | 4.75 ± 0.07 | 16.09 | – |
| DM-GAN + CL | 4.77 ± 0.05 | 14.38 | 78.99 ± 0.66 |
| DF-GAN | 4.86 ± 0.04 | 19.24 | – |
| SSA-GAN * | 5.17 ± 0.08 | 15.61 | 86.70 ± 0.91 |
| SSA-GAN # | 5.13 ± 0.06 | 15.60 | 73.55 ± 0.68 |
| SSA-GAN + CL (mine) | 4.85 ± 0.08 | 15.10 | 84.08 ± 0.73 |
* as published in the SSA-GAN paper. # SSA-GAN as I reproduced it with the same setup. My model has the lowest FID in the table and raises R-precision by 10.5 points over the reproduced SSA-GAN; its Inception Score is lower than SSA-GAN's.
Screenshots
Figures from the thesis: birds generated from captions of the CUB test set, and the results table.
Generated birds
Four captions, the real photo, and the birds generated by AttnGAN + CL, SSA-GAN and my model.
Results table from the thesis
Inception Score, FID and R-precision against published models on CUB birds.
Challenges and learning
Challenge: People describe the same bird in very different words, and the model drew a different bird for each wording.
Solution: A contrastive loss pulls captions of the same image together and pushes captions of different images apart, first when pre-training the text and image encoders and then while the generator trains.
Learned: Metrics pull in different directions: realism (FID) and caption match (R-precision) improved, while the Inception Score dropped slightly. Reporting all three, and reproducing the baseline myself instead of only quoting its paper, kept the comparison honest.
Technologies used
Framework ·Python and PyTorch
- The whole model: generator, discriminator, text and image encoders, custom losses and training loops.
Model ·Semantic-Spatial Aware Generative Adversarial Network (SSA-GAN)
- Semantic-aware blocks with a mask predictor at each stage that decides where in the image each part of the caption applies.
- Generates 256 × 256 images.
Model ·Deep Attentional Multimodal Similarity Model (DAMSM)
- A bi-directional Long Short-Term Memory (LSTM) for the caption and Inception-v3, pre-trained on ImageNet, for the image.
- Pre-trained again with a contrastive loss so that captions of the same image sit close together.
Learning ·Contrastive learning (CL), Normalized Temperature-scaled Cross Entropy (NT-Xent)
- A SimCLR-style loss with temperature 0.5 and the other pairs in the batch as negatives.
- Used twice: when pre-training the encoders and while training the generator.
Learning ·Adversarial training
- Hinge loss; Adam with a learning rate of 0.0001 for the generator and 0.0004 for the discriminator.
Data ·Caltech-UCSD Birds-200-2011 (CUB-200-2011)
- 200 bird species, 11,788 images, 10 written descriptions per image.
Evaluation ·Fréchet Inception Distance (FID), Inception Score and R-precision
- FID for realism, Inception Score for quality and variety, R-precision for how well an image matches its caption.
Tools ·CUDA and tensorboardX
- GPU training for 600 epochs with batch size 24; losses and samples tracked in TensorBoard.
Want to know more?
I'm happy to walk you through the product and engineering decisions on a call. Let's connect.