MoCa: Modality-aware Continual Pre-training Makes Better Bidirectional Multimodal Embeddings
version
0.0.2
engine
skillxiv-v0.0.2-claude-opus-4.6
license
MIT
url
https://arxiv.org/abs/2506.23115
keywords
["Multimodal Embeddings","Continual Pre-training","Vision Language Models","Bidirectional Embeddings","Cross-modal Retrieval"]
description
Transform pre-trained vision-language models into powerful bidirectional multimodal embeddings through modality-aware continual pre-training and heterogeneous contrastive fine-tuning. 3B model matches 7B baselines.
MoCa: Efficient Multimodal Embeddings Through Modality-Aware Pre-training
Vision-language models are excellent at understanding images and text, but they're designed for generative tasks (caption generation, visual question answering). When you extract embeddings from VLMs for retrieval tasks (find similar images to a text query), performance lags significantly behind specialized embedding models. The problem is that VLM embeddings are optimized for generation, not for forming a unified representation space where similar images and text are close in distance.
MoCa solves this by adapting pre-trained VLMs into embedding models through two stages: (1) modality-aware continual pre-training using both masked language modeling and masked autoencoding to encourage cross-modal understanding, and (2) heterogeneous contrastive fine-tuning on diverse data including long-form documents, curated pairs, and text-only data. The result achieves state-of-the-art multimodal retrieval performance with smaller models (3B matching 7B baselines).
Core Concept
The key insight is that embedding spaces have different properties than generative representations. VLMs optimize for:
Causal attention (left-to-right generation)
Autoregressive loss (predict next token)
Task-specific outputs (caption generation)
MoCa embeddings need:
Bidirectional context (both left and right matter)
Unified cross-modal space (image and text in same space)
This requires modification at two levels:
Changing the attention mechanism from causal to bidirectional
Changing the training objective from generation to contrastive learning
The approach uses two stages: (1) unlabeled continual pre-training to adapt representations without task-specific labels, and (2) contrastive fine-tuning with diverse data to build discriminative embeddings.
Architecture Overview
MoCa modifies a vision-language model at the architecture and training level:
Vision Encoder: Vision Transformer (unchanged from VLM)
Text Encoder: Modified to use bidirectional attention instead of causal
Shared Representation Space: Single unified embedding space for both modalities
Mean Pooling: Image and text embeddings are averaged (not CLS tokens) for better uniformity
Multi-stage Training: CPT (continual pre-training) followed by heterogeneous contrastive fine-tuning
Implementation
Step 1: Modify the VLM for bidirectional embeddings
Convert the causal language model to bidirectional by changing attention patterns and adding bidirectional embeddings.
import torch
import torch.nn as nn
classBidirectionalTextEncoder(nn.Module):
"""
Modify a causal LLM to use bidirectional attention for embeddings.
Key changes: replace causal mask with bidirectional, use mean pooling.
"""def__init__(self, original_model):
super().__init__()
self.embedding = original_model.embed_tokens
self.layers = original_model.model.layers
# Replace causal attention with bidirectionalself.layers = self._convert_to_bidirectional(self.layers)
self.norm = original_model.model.norm
def_convert_to_bidirectional(self, layers):
"""
Modify attention layers from causal (triangular mask) to bidirectional.
"""
modified_layers = nn.ModuleList()
for layer in layers:
# Get the attention block
attn = layer.self_attn
# Create a wrapper that removes the causal mask
original_forward = attn.forward
defbidirectional_forward(hidden_states, *args, **kwargs):
# Remove attention_mask or replace with all-ones mask# This enables bidirectional attention
kwargs['attention_mask'] = Nonereturn original_forward(hidden_states, *args, **kwargs)
attn.forward = bidirectional_forward
modified_layers.append(layer)
return modified_layers
defforward(self, input_ids, attention_mask=None):
"""
Encode text using bidirectional attention.
Return mean-pooled embeddings instead of next-token logits.
"""# Embed tokens
hidden_states = self.embedding(input_ids) # [batch, seq_length, hidden_dim]# Apply transformer layers with bidirectional attentionfor layer inself.layers:
layer_outputs = layer(hidden_states, attention_mask=attention_mask)
hidden_states = layer_outputs[0]
# Final layer norm
hidden_states = self.norm(hidden_states)
# Mean pooling: average over sequence length# This is better than CLS token for embedding uniformity
embeddings = hidden_states.mean(dim=1) # [batch, hidden_dim]return embeddings
