CoCa: Contrastive Captioners are Image-Text Foundation Models

TL;DR

• CoCa trains a single image-text model jointly with two losses — a contrastive loss (like CLIP) and an autoregressive captioning loss — in one forward pass, requiring no separate training stages.
• The text decoder is split into a unimodal lower half (self-attention only, produces the text embedding for contrastive alignment) and a multimodal upper half (adds cross-attention to image features, produces captioning logits).
• Task-specific attentional poolers convert image token sequences into either a single global vector (for contrastive use) or multiple fine-grained vectors (for generative use), with learned compression rather than a guarantee that every visual detail survives.
• The resulting foundation model achieves state-of-the-art results on image classification, image-text retrieval, and image captioning — all from the same pretrained weights.

One model, two objectives

Most vision-language models commit to one training paradigm: contrastive models (like CLIP) learn to align images and text in a shared embedding space, while generative models learn to produce descriptive captions token by token. Each approach excels at different tasks — contrastive models shine at retrieval and zero-shot classification; generative models enable open-ended captioning and VQA. CoCa asks: why choose?

The key insight is that both objectives can share the same forward pass. The image encoder produces a sequence of patch tokens once. The text decoder runs once, but its lower layers (unimodal, self-attention only) produce a text embedding that feeds the contrastive loss, while its upper layers (multimodal, with cross-attention to image features) produce token predictions that feed the captioning loss. The total training loss is a weighted sum: L_total = λ_con · L_con + λ_cap · L_cap, with λ_cap = 2 in the paper.

Computing both losses in a single forward pass shares the image encoder and lower text decoder computation. In the paper's Table 8b ablation, combined training costs 1.18 times the contrastive-only baseline, close to captioning-only at 1.17 times. Shared computation does not make captioning free or establish identical cost to CLIP. There is no alternation between contrastive and generative batches, no two-stage pretraining schedule, and no separate model to maintain for each downstream task type.

Decoupled decoder

The architecture that makes the dual objective possible is the decoupled text decoder. A standard captioning decoder would apply cross-attention to image features at every layer, which would prevent the text embedding from being independent of the image — exactly what the contrastive loss requires. CoCa resolves this cleanly: the lower half of the decoder uses only causal self-attention (masked attention), making it a pure text representation. The upper half stacks cross-attention on top, fusing image information to generate tokens conditioned on what the image contains.

At training time, a learned [CLS] token is appended at the end of the sentence so causal attention can summarize preceding text. The unimodal half’s trailing [CLS] representation is projected and compared against the image embedding under the contrastive loss. The multimodal half’s outputs are compared against ground-truth tokens under the standard cross-entropy captioning loss. Both paths share the same decoder weights up to the split point, so language representations are learned jointly across both objectives.

The split depth is a hyperparameter. In the paper, roughly half the decoder layers are unimodal and half are multimodal. Shallower splits give the contrastive branch fewer layers to refine a clean text embedding; deeper splits leave fewer layers for cross-modal fusion. The chosen split balances alignment quality against generation quality on downstream benchmarks.

Attentional pooling

The image encoder (a ViT) outputs a sequence of patch tokens — rich spatial information, but not a single vector the contrastive loss can compare against a text CLS token. CoCa introduces task-specific attentional poolers to supply different representation lengths. Pooling is learned compression and can lose visual detail.

An attentional pooler is a small cross-attention module with a fixed set of learned query vectors. The queries attend over the supplied key/value tokens, and the pooler outputs one vector per query. For the generative (captioning) objective, many queries attend to the encoder tokens and produce a sequence of visual features for the multimodal decoder. In the original default cascade design, a single contrastive query then attends to those generative pooler outputs to produce a global summary vector. Figure 2 shows a parallel alternative, but §4.3 and Table 8f state that cascade is the default in all CoCa models.

Using separate poolers lets each objective request a different representation length. The contrastive pooler can specialise in producing compact, discriminative global representations. The generative pooler can retain spatial granularity. In the default cascade, the generative pooler reads encoder tokens and the contrastive pooler reads its outputs. Both objectives still share one image encoder evaluation; they do not both attend directly to the encoder sequence.

Implementation

The pooling design now connects to the decoder split: one visual sequence supplies both objectives, while the contrastive text representation stays independent of images. This reduced forward makes those routes, the next-token alignment, and their combined gradients executable. Its random weights demonstrate mechanics rather than caption meaning.

class ImageEncoder(nn.Module):
    """Maps 4x4 patches to width-D tokens; deliberately smaller than a ViT."""

    def __init__(self, width):
        super().__init__()
        self.patches = nn.Conv2d(3, width, kernel_size=4, stride=4)
        self.calls = 0

    def forward(self, images):
        """Evaluates the shared image encoder once per joint model forward."""
        self.calls += 1
        return self.patches(images).flatten(2).transpose(1, 2)

class Pooler(nn.Module):
    """Pools keys/values with learned queries, one output per query."""

    def __init__(self, queries, width, heads):
        super().__init__()
        self.queries = nn.Parameter(torch.randn(queries, width) * 0.02)
        self.attention = nn.MultiheadAttention(width, heads, batch_first=True)

    def forward(self, tokens):
        """Returns [B, queries, D] rather than claiming lossless compression."""
        queries = self.queries.unsqueeze(0).expand(tokens.shape[0], -1, -1)
        return self.attention(queries, tokens, tokens, need_weights=False)[0]


def cascade_poolers(model, visual):
    """Follows the original default in section 4.3/Table 8f, not Figure 2."""
    caption_tokens = model.caption_pooler(visual)
    contrastive_token = model.contrastive_pooler(caption_tokens)
    return caption_tokens, contrastive_token

def encode_text(self, tokens):
    """Appends CLS so a causal query can summarize all preceding valid text."""
    hidden = self.embedding(tokens)
    cls = self.cls.expand(tokens.shape[0], -1, -1)
    hidden = torch.cat((hidden, cls), dim=1)
    hidden = hidden + self.positions[: hidden.shape[1]]
    padding = torch.cat(
        (
            tokens.eq(PAD),
            torch.zeros(
                tokens.shape[0], 1, dtype=torch.bool, device=tokens.device
            ),
        ),
        dim=1,
    )
    # No image tensor enters this function; pads are blocked as keys.
    return self.unimodal(hidden, padding)

def contrastive(self, image_token, unimodal):
    """Uses trailing CLS, unit embeddings and a trainable positive temperature."""
    image = F.normalize(self.image_projection(image_token[:, 0]), dim=-1)
    text = F.normalize(self.text_projection(unimodal[:, -1]), dim=-1)
    logits = image @ text.T / self.log_temperature.exp()
    labels = torch.arange(image.shape[0], device=image.device)
    # Eq. 2 adds directional batch means: deliberately no factor of 1/2.
    loss = F.cross_entropy(logits, labels) + F.cross_entropy(
        logits.T, labels
    )
    return image, text, logits, loss

def caption_logits(self, unimodal, tokens, caption_tokens):
    """Fuses images after the unimodal split; excludes CLS from caption states."""
    text_states = unimodal[:, :-1]
    conditioned = self.multimodal(
        text_states, tokens.eq(PAD), caption_tokens
    )
    return self.caption_head(conditioned)

def caption_loss(logits, tokens):
    """Sums next-token NLL per example, then averages examples, ignoring PAD."""
    # Position i predicts input token i+1; BOS and appended CLS are never targets.
    targets = tokens[:, 1:]
    predictions = logits[:, :-1]
    losses = F.cross_entropy(
        predictions.transpose(1, 2), targets, ignore_index=PAD, reduction="none"
    )
    return losses.sum(dim=1).mean()

def joint_loss(output, tokens, contrastive_weight=1.0, caption_weight=2.0):
    """Weights explicitly declared reductions; these weights follow section 4.1."""
    weights = (contrastive_weight, caption_weight)
    if any(not math.isfinite(w) or w < 0 for w in weights) or not any(weights):
        raise ValueError(
            "Loss weights must be finite, nonnegative, and not all zero"
        )
    contrastive = output["contrastive_loss"]
    caption = caption_loss(output["caption_logits"], tokens)
    total = contrastive_weight * contrastive + caption_weight * caption
    return total, contrastive, caption


def train_step(model, images, tokens, optimizer):
    """Runs one joint forward, one backward, then one optimizer update."""
    optimizer.zero_grad(set_to_none=True)
    output = model(images, tokens)
    total, contrastive, caption = joint_loss(output, tokens)
    total.backward()
    optimizer.step()
    return output, (total.detach(), contrastive.detach(), caption.detach())

python -m pip install -r requirements-coca.txt
python coca.py --smoke --seed 7

Why it mattered

Before CoCa, achieving top results across retrieval, classification, and captioning typically required separate models or multi-stage pipelines: first pretrain a contrastive model, then attach and fine-tune a captioning head. CoCa showed that a single pretraining recipe — one encoder, one decoupled decoder, two poolers, two losses, one forward pass — produces a foundation model competitive with or better than specialized systems on all three task families simultaneously.

This unification matters in practice. A single set of pretrained weights can be fine-tuned for image captioning (adapt the multimodal decoder), for image-text retrieval (use the contrastive embeddings directly), or for image classification (linear probe on the contrastive image embedding). No architectural changes are needed between tasks; only the fine-tuning head or objective changes.

CoCa also demonstrated that the generative captioning loss acts as a useful regulariser for the contrastive objective and vice versa: learning to generate captions forces the model to retain compositional image content, while the alignment loss anchors text representations to visual semantics. The combination yields stronger representations than either objective alone at equivalent model scale.

Related Reading

• CLIP — the contrastive pretraining paradigm that CoCa extends with a joint captioning objective
• SigLIP — a successor contrastive approach that replaces the softmax loss with a sigmoid loss for improved scaling
• BLIP-2 — bridges frozen image encoders and large language models with a lightweight querying transformer, approaching the dual-objective idea from the fine-tuning angle