[Paper Review] Three Towers: Flexible Contrastive Learning with Pretrained Image Models

Jannik Kossen, Mark Collier+8 more|arXiv (Cornell University)|May 26, 2023

Multimodal Machine Learning Applications4 citations

TL;DR

Three Towers (3T) introduces a flexible contrastive learning framework that enhances vision-language models by integrating frozen pretrained image classifier embeddings through a third tower, while allowing the main image and text towers to be jointly trained via contrastive loss. This approach consistently outperforms both LiT and from-scratch baselines in retrieval and classification, especially on diverse pretraining datasets like ImageNet-21k and Places365.

ABSTRACT

We introduce Three Towers (3T), a flexible method to improve the contrastive learning of vision-language models by incorporating pretrained image classifiers. While contrastive models are usually trained from scratch, LiT (Zhai et al., 2022) has recently shown performance gains from using pretrained classifier embeddings. However, LiT directly replaces the image tower with the frozen embeddings, excluding any potential benefits from training the image tower contrastively. With 3T, we propose a more flexible strategy that allows the image tower to benefit from both pretrained embeddings and contrastive training. To achieve this, we introduce a third tower that contains the frozen pretrained embeddings, and we encourage alignment between this third tower and the main image-text towers. Empirically, 3T consistently improves over LiT and the CLIP-style from-scratch baseline for retrieval tasks. For classification, 3T reliably improves over the from-scratch baseline, and while it underperforms relative to LiT for JFT-pretrained models, it outperforms LiT for ImageNet-21k and Places365 pretraining.

Motivation & Objective

• To address the limitation of LiT, which locks pretrained image embeddings and prevents the image tower from benefiting from contrastive training.
• To develop a flexible method that enables joint learning from both pretrained image representations and contrastive signals in vision-language models.
• To improve robustness and performance across diverse downstream tasks, especially when pretrained models are suboptimal or mismatched to downstream data.
• To enable trade-offs between prior knowledge from pretrained models and inductive bias from contrastive learning via a modular three-tower architecture.
• To empirically validate that combining pretrained embeddings with contrastive training yields better generalization than relying solely on frozen embeddings or training from scratch.

Proposed method

• Introduce a third tower containing frozen, pretrained image embeddings from a classifier (e.g., ImageNet-21k or Places365), distinct from the main image and text towers.
• Train the main image and text towers from scratch using contrastive learning on web-scraped image-caption pairs.
• Apply additional contrastive loss terms to align the main image tower with the third tower (pretrained embeddings), enabling knowledge transfer.
• Use multihead attention pooling to aggregate patch embeddings into a single representation for both the main image tower and the third tower.
• Optimize the full model using Adafactor with learning rate warmup and cosine decay, gradient clipping, and weight decay.
• Maintain inference efficiency by discarding the third tower at test time, using only the main image and text towers for downstream tasks.

Experimental results

Research questions

• RQ1 Can we improve contrastive vision-language learning by combining frozen pretrained image embeddings with trainable image towers?
• RQ2 Does the proposed Three Towers framework outperform LiT and from-scratch baselines in retrieval and zero-shot classification?
• RQ3 How does 3T perform when the pretrained model is mismatched to downstream tasks, such as with Places365 pretraining?
• RQ4 Does the method remain robust when the pretrained model lacks coverage for certain image categories or domains?
• RQ5 Can 3T achieve better performance than LiT on diverse pretraining datasets like ImageNet-21k and JFT?

Key findings

• 3T improves top-1 recall on Flickr30k and COCO retrieval benchmarks over both LiT and from-scratch baselines, achieving 87.3% (img2txt) and 48.5% (txt2img) respectively.
• For ImageNet-1k zero-shot classification, 3T outperforms the from-scratch baseline and matches or exceeds LiT when using ImageNet-21k-pretrained models.
• On Places365-pretrained models, 3T significantly outperforms LiT, which fails to generalize to downstream tasks like ImageNet-1k and CIFAR-100.
• 3T shows improved robustness to pretraining dataset limitations, particularly in scenarios where the pretrained model's labels or distribution do not align with downstream tasks.
• The method benefits more from larger model scales and longer training, with convergence observed in ~3 days on 256 TPU chips for L-scale models.
• For g-scale models, 3T runs converge in about 5 days on 512 TPU chips, with training stability improved by adjusting $eta_2$ to 0.95.