Align before Fuse: Vision and Language Representation Learning with Momentum Distillation

1 Introduction . Vision-and-Language Pre-training ( VLP ) aims to learn multimodal representations from large-scale image-text pairs that can improve downstream Vision-and-Language ( V+L ) tasks . Most existing VLP methods ( e.g . LXMERT [ 1 ] , UNITER [ 2 ] , OSCAR [ 3 ] ) rely on pre-trained object detectors to extract region-based image features , and employ a multimodal encoder to fuse the image features with word tokens . The multimodal encoder is trained to solve tasks that require joint understanding of image and text , such as masked language modeling ( MLM ) and image-text matching ( ITM ) . While effective , this VLP framework suffers from several key limitations : ( 1 ) The image features and the word token embeddings reside in their own spaces , which makes it challenging for the multimodal encoder to learn to model their interactions ; ( 2 ) The object detector is both annotation-expensive and compute-expensive , because it requires bounding box annotations during pre-training , and highresolution ( e.g . 600⇥1000 ) images during inference ; ( 3 ) The widely used image-text datasets [ 4 , 5 ] are collected from the web and are inherently noisy , and existing pre-training objectives such as MLM may overfit to the noisy text and degrade the model ’ s generalization performance . We propose ALign BEfore Fuse ( ALBEF ) , a new VLP framework to address these limitations . We first encode the image and text independently with a detector-free image encoder and a text encoder . Then we use a multimodal encoder to fuse the image features with the text features through crossmodal attention . We introduce an intermediate image-text contrastive ( ITC ) loss on representations from the unimodal encoders , which serves three purposes : ( 1 ) it aligns the image features and the text features , making it easier for the multimodal encoder to perform cross-modal learning ; ( 2 ) it 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . improves the unimodal encoders to better understand the semantic meaning of images and texts ; ( 3 ) it learns a common low-dimensional space to embed images and texts , which enables the image-text matching objective to find more informative samples through our contrastive hard negative mining . To improve learning under noisy supervision , we propose Momentum Distillation ( MoD ) , a simple method which enables the model to leverage a larger uncurated web dataset . During training , we keep a momentum version of the model by taking the moving-average of its parameters , and use the momentum model to generate pseudo-targets as additional supervision . With MoD , the model is not penalized for producing other reasonable outputs that are different from the web annotation . We show that MoD not only improves pre-training , but also downstream tasks with clean annotations . We provide theoretical justifications on ALBEF from the perspective of mutual information maximization . Specifically , we show that ITC and MLM maximize a lower bound on the mutual information between different views of an image-text pair , where the views are generated by taking partial information from each pair . From this perspective , our momentum distillation can be interpreted as generating new views with semantically similar samples . Therefore , ALBEF learns vision-language representations that are invariant to semantic-preserving transformations . We demonstrate the effectiveness of ALBEF on various downstream V+L tasks including image-text retrieval , visual question answering , visual reasoning , visual entailment , and weakly-supervised visual grounding . ALBEF achieves substantial improvements over existing state-of-the-art methods . On image-text retrieval , it outperforms methods that are pre-trained on orders of magnitude larger datasets ( CLIP [ 6 ] and ALIGN [ 7 ] ) . On VQA and NLVR2 , it achieves absolute improvements of 2.37 % and 3.84 % compared to the state-of-the-art method VILLA [ 8 ] , while enjoying much faster inference speed . We also provide quantitative and qualitative analysis on ALBEF using Grad-CAM [ 9 ] , which reveals its ability to perform accurate object , attribute and relationship grounding implicitly . 2 Related Work . 2.1 Vision-Language Representation Learning . Most existing work on vision-language representation learning fall into two categories . The first category focuses on modelling the interactions between image and text features with transformerbased multimodal encoders [ 10 , 11 , 12 , 13 , 1 , 14 , 15 , 2 , 3 , 16 , 8 , 17 , 18 ] . Methods in this category achieve superior performance on downstream V+L tasks that require complex reasoning over image and text ( e.g . NLVR2 [ 19 ] , VQA [ 20 ] ) , but most of them require high-resolution input images and pre-trained object detectors . A recent method [ 21 ] improves inference speed by removing the object detector , but results in lower performance . The second category focuses on learning separate unimodal encoders for image and text [ 22 , 23 , 6 , 7 ] . The recent CLIP [ 6 ] and ALIGN [ 7 ] perform pre-training on massive noisy web data using a contrastive loss , one of the most effective loss for representation learning [ 24 , 25 , 26 , 27 ] . They achieve remarkable performance on image-text retrieval tasks , but lack the ability to model more complex interactions between image and text for other V+L tasks [ 21 ] . ALBEF unifies the two categories , leading to strong unimodal and multimodal representations with superior performance on both retrieval and reasoning tasks . Furthermore , ALBEF does not require object detectors , a major computation bottleneck for many existing methods [ 1 , 2 , 3 , 8 , 17 ] . 2.2 Knowledge Distillation . Knowledge distillation [ 28 ] aims to improve a student model ’ s performance by distilling knowledge from a teacher model , usually through matching the student ’ s prediction with the teacher ’ s . While most methods focus on distilling knowledge from a pre-trained teacher model [ 28 , 29 , 30 , 31 , 32 ] , online distillation [ 33 , 34 ] simultaneously trains multiple models and use their ensemble as the teacher . Our momentum distillation can be interpreted as a form of online self-distillation , where a temporal ensemble of the student model is used as the teacher . Similar ideas have been explored in semi-supervised learning [ 35 ] , label noise learning [ 36 ] , and very recently in contrastive learning [ 37 ] . Different from existing studies , we theoretically and experimentally show that momentum distillation is a generic learning algorithm that can improve the model ’ s performance on many V+L tasks . 3 ALBEF Pre-training . In this section , we first introduce the model architecture ( Section 3.1 ) . Then we delineate the pretraining objectives ( Section 3.2 ) , followed by the proposed momentum distillation ( Section 3.3 ) . Lastly we describe the pre-training datasets ( Section 3.4 ) and implementation details ( Section 3.5 ) . 3.1 Model Architecture . As illustrated in Figure 1 , ALBEF contains an image encoder , a text encoder , and a multimodal encoder . We use a 12-layer visual transformer ViT-B/16 [ 38 ] as the image encoder , and initialize it with weights pre-trained on ImageNet-1k from [ 31 ] . An input image I is encoded into a sequence of embeddings : { vcls , v1 , ... , vN } , where vcls is the embedding of the [ CLS ] token . We use a 6-layer transformer [ 39 ] for both the text encoder and the multimodal encoder . The text encoder is initialized using the first 6 layers of the BERTbase [ 40 ] model , and the multimodal encoder is initialized using the last 6 layers of the BERTbase . The text encoder transforms an input text T into a sequence of embeddings { wcls , w1 , ... , wN } , which is fed to the multimodal encoder . The image features are fused with the text features through cross attention at each layer of the multimodal encoder . 3.2 Pre-training Objectives . We pre-train ALBEF with three objectives : image-text contrastive learning ( ITC ) on the unimodal encoders , masked language modeling ( MLM ) and image-text matching ( ITM ) on the multimodal encoder . We improve ITM with online contrastive hard negative mining . Image-Text Contrastive Learning aims to learn better unimodal representations before fusion . It learns a similarity function s = gv ( vcls ) > gw ( wcls ) , such that parallel image-text pairs have higher similarity scores . gv and gw are linear transformations that map the [ CLS ] embeddings to normalized lower-dimensional ( 256-d ) representations . Inspired by MoCo [ 24 ] , we maintain two queues to store the most recent M image-text representations from the momentum unimodal encoders . The normalized features from the momentum encoders are denoted as g0v ( v0cls ) and g 0 w ( w 0 cls ) . We define s ( I , T ) = gv ( vcls ) > g0w ( w 0 cls ) and s ( T , I ) = gw ( wcls ) > g0v ( v 0 cls ) . For each image and text , we calculate the softmax-normalized image-to-text and text-to-image similarity as : pi2tm ( I ) = exp ( s ( I , Tm ) /⌧ ) PM m=1 exp ( s ( I , Tm ) /⌧ ) , pt2im ( T ) = exp ( s ( T , Im ) /⌧ ) PM m=1 exp ( s ( T , Im ) /⌧ ) ( 1 ) where ⌧ is a learnable temperature parameter . Let yi2t ( I ) and yt2i ( T ) denote the ground-truth one-hot similarity , where negative pairs have a probability of 0 and the positive pair has a probability of 1 . The image-text contrastive loss is defined as the cross-entropy H between p and y : Litc = 1 2 E ( I , T ) ⇠D ⇥ H ( yi2t ( I ) , pi2t ( I ) ) + H ( yt2i ( T ) , pt2i ( T ) ) ⇤ ( 2 ) Masked Language Modeling utilizes both the image and the contextual text to predict the masked words . We randomly mask out the input tokens with a probability of 15 % and replace them with the special token [ MASK ] 1 . Let T̂ denote a masked text , and pmsk ( I , T̂ ) denote the model ’ s predicted probability for a masked token . MLM minimizes a cross-entropy loss : Lmlm = E ( I , T̂ ) ⇠DH ( y msk , pmsk ( I , T̂ ) ) ( 3 ) where ymsk is a one-hot vocabulary distribution where the ground-truth token has a probability of 1 . Image-Text Matching predicts whether a pair of image and text is positive ( matched ) or negative ( not matched ) . We use the multimodal encoder ’ s output embedding of the [ CLS ] token as the joint representation of the image-text pair , and append a fully-connected ( FC ) layer followed by softmax to predict a two-class probability pitm . The ITM loss is : Litm = E ( I , T ) ⇠DH ( yitm , pitm ( I , T ) ) ( 4 ) where yitm is a 2-dimensional one-hot vector representing the ground-truth label . We propose a strategy to sample hard negatives for the ITM task with zero computational overhead . A negative image-text pair is hard if they share similar semantics but differ in fine-grained details . We use the contrastive similarity from Equation 1 to find in-batch hard negatives . For each image in a mini-batch , we sample one negative text from the same batch following the contrastive similarity distribution , where texts that are more similar to the image have a higher chance to be sampled . Likewise , we also sample one hard negative image for each text . The full pre-training objective of ALBEF is : L = Litc + Lmlm + Litm ( 5 )

Summary: The paper proposes to learn joint vision and language representations by addressing several limitations in the existing methods. Firstly, they learn separate image and text representations using unimodal encoders without the need for bounding box annotations. Second, they fuse these two representations using a multimodal encoder based on cross modal attention and, third, they address the noise in large scale web based image-text datasets using a Momentum Distillation approach that generates pseudo-targets as additional supervision. They demonstrate the effectiveness of their approach on various downstream V+L tasks and gain substantial improvement over SOTA methods. Their main contributions are: removing the need for pre-trained object detector and high resolution images and combining it with the contrastive loss function for learning effective multi-modal representations, image-text contrastive learning loss (ITC). They also generate pseudo-targets for the ITC loss and the masked language modeling (MLM) loss using the momentum encoder model to address the weak correlations in the noisy image-text web data. This achieves high performance on both reasoning and retrieval tasks unlike other methods. The paper  is very well written and easy to follow and understand.