Multi-modal Alignment using Representation Codebook

TL;DR

This work addresses aligning heterogeneous image and text features for vision-language pretraining. It introduces CODIS, a framework that learns a joint multimodal codebook of $K$ prototypes in a $d_c$-dimensional space (e.g., $K=4000$, $d_c=256$) and uses Optimal Transport to align modality-specific features to prototypes, with a momentum teacher system to stabilize learning. The method combines multimodal codebook learning with teacher-student distillation and standard SSL objectives (ITM and MLM) into a unified pretraining objective ${\mathcal L}_{final}= {\mathcal L}_{mlm}+{\mathcal L}_{itm}+{\mathcal L}_{ica}+{\mathcal L}_{code}$, achieving SoTA zero-shot cross-modality retrieval and strong transfer on V&L tasks. By bridging modalities at the prototype level and leveraging stable teacher guidance, CODIS effectively mitigates distribution gaps and prototype collapse, enabling efficient multimodal alignment without requiring massive data resources.

Abstract

Aligning signals from different modalities is an important step in vision-language representation learning as it affects the performance of later stages such as cross-modality fusion. Since image and text typically reside in different regions of the feature space, directly aligning them at instance level is challenging especially when features are still evolving during training. In this paper, we propose to align at a higher and more stable level using cluster representation. Specifically, we treat image and text as two "views" of the same entity, and encode them into a joint vision-language coding space spanned by a dictionary of cluster centers (codebook). We contrast positive and negative samples via their cluster assignments while simultaneously optimizing the cluster centers. To further smooth out the learning process, we adopt a teacher-student distillation paradigm, where the momentum teacher of one view guides the student learning of the other. We evaluated our approach on common vision language benchmarks and obtain new SoTA on zero-shot cross modality retrieval while being competitive on various other transfer tasks.

Figures (4)

• Figure 1: We propose to use a learnable codebook to better align the image and text modalities. The codebook serves as a "bridge" between the image and text features. Each codeword can be interpreted as a prototype, which enables contrasting image and text at the cluster level. We then solve an optimal transport ambrosio2008gradient problem to optimize the distance between each modality to the prototypes, which in turn optimizes the alignment between the two modalities. Prototype vectors are learned along with the feature encoders in our V&L framework.
• Figure 2: Overview of our framework. For simplicity, we only display a pair of teacher-student encoders (e.g., teacher for the image and student for the text) and similarly for the memory queue. The teacher is updated with an exponential moving average of the student (from the same modality). The codebook helps bridge the gap between the different modalities. The entire framework is end-to-end optimized.
• Figure 3: This is the diagram illustrating how to calculate four codebook losses. "$\rightarrow$": softmax operator. "$\rightarrow$": IPOT algorithm. "$\rightarrow$": OT loss. "$\rightarrow$": cross entropy.
• Figure 4: Grad-CAM visualization on the cross-attention maps corresponding to individual words