Budget-Preserving Image Captioning

Why People Wanted It

Caption dropout removes the caption from some text-to-image training examples. The retained-caption examples teach image generation from text; the dropped-caption examples teach image generation without text conditioning. Neither format explicitly trains the model to describe an image.

The problem is the allocation. Adding image-to-text training as a separate task could spend training capacity that otherwise goes to image generation. So the real question was never just whether a model could learn both directions—it was whether captioning could use the caption-dropped slots without reducing either category of image-token training.

The Core Move

Keep caption-conditioned image generation, but replace its caption-dropped examples with image-first sequences that continue into captions.

Let $c=(c_1,\ldots,c_L)$ be caption tokens and $x=(x_1,\ldots,x_N)$ be image tokens. Tokens are discrete symbols from a shared vocabulary. Define the image span as

$$P(x)=[SOV]\ x_1\ldots x_N\ [EOV].$$

Here, $[SOV]$ and $[EOV]$ mark the start and end of the image; $[EOS]$ marks the end of an example. The complete exchange is:

Training slot	Baseline format	Replacement format
Caption retained	$c\ P(x)\ [EOS]$	$c\ P(x)\ [EOS]$
Caption dropped	$P(x)\ [EOS]$	$P(x)\ c\ [EOS]$

The first row is essential: the model must still practice generating an image after reading its caption. Only the second row changes.

Quiz

Why can an image-to-text example replace a caption-dropped example without removing its unconditional image-token training?

The later caption helps predict the preceding image through causal attention.

All caption-conditioned image-generation examples are replaced.

The replacement removes image tokens from the training loss.

The image tokens keep the same preceding context; the caption becomes an additional continuation.

$ check

Why the Shared Prefix Preserves Image Training

Here is the part that gets technical: causal next-token prediction makes the image-token loss depend on tokens before the image, not on a caption appended afterward.

For a caption-conditioned sequence, the image-token contribution to the loss is

$$\mathcal{L}_{I\mid T} =-\sum_{i=1}^{N}\log p_\theta(x_i\mid c,[SOV],x_{<i}).$$

For either a caption-dropped example or its image-to-text replacement, it is

$$\mathcal{L}_{I\mid\varnothing} =-\sum_{i=1}^{N}\log p_\theta(x_i\mid [SOV],x_{<i}).$$

Here, $x_{<i}$ denotes earlier image tokens and $\theta$ denotes model parameters. The replacement caption comes later, so causal attention cannot use it to predict these image tokens. “Unconditional” means no preceding caption; earlier image tokens remain available.

The replacement additionally trains

$$\mathcal{L}_{T\mid I} =-\sum_{j=1}^{L}\log p_\theta(c_j\mid P(x),c_{<j}).$$

These expressions isolate image and caption targets; complete sequence training also includes boundary-token targets. The image-first formats share their image prediction contexts, but differ afterward: one terminates, while the other predicts caption tokens and then terminates. How should those two formats be allocated?

Why Replacement Is Different from Adding a Task

Now the deeper accounting distinction: replacement must match the original dropout allocation and retain the caption-conditioned branch, rather than mix captioning with only unconditional image examples.

Let $p_{drop}$ be the baseline caption-drop rate and $\eta$ the replacement share assigned to image-to-text. The recipe sets

$$\eta=p_{drop},\qquad p_{drop}^{\mathrm{new}}=0.$$

The remaining share, $1-\eta$, uses caption-first text-to-image examples with no caption dropout. For a matched collection of pairs, write $B_I$ for its total image-token count. The allocation is then

$$B_{I\mid T}=(1-p_{drop})B_I, \qquad B_{I\mid\varnothing}=p_{drop}B_I.$$

Both baseline and replacement retain this split: image-to-text supplies the image tokens without preceding text, and caption-first text-to-image supplies those with preceding text. Matching only example fractions in an arbitrary variable-length sampler is not itself a guarantee of exact token counts; the intended allocation must be maintained at the image-token level.

This is preservation of image-token accounting under the fixed pretraining budget, not a claim that caption targets cost no computation or that every sequence has the same targets. Does matching the allocation also determine learning behavior?

Why Interleaving Matters for Caption Learning

Now the version for readers who want more of the scheduling machinery: equal aggregate token counts do not specify when the model receives each training objective.

Staging puts image-to-text training before text-to-image training. Mixing interleaves both throughout training. In the separate scheduling comparison, these schedules use identical token counts at each image-to-text allocation; ordinary mixed training there still applies caption dropout to text-to-image examples, unlike the replacement recipe.

At the smaller tested allocation, both schedules preserve comparable text-to-image generation. At the larger allocation, staging achieves substantially lower gFD, an image-generation distribution metric for which lower is better. Only mixing substantially reduces image-conditioned text negative log-likelihood, the caption prediction loss. These findings distinguish scheduling from allocation; they do not make arbitrary mixtures equivalent to dropout replacement. How much captioning practice must remain present throughout training?

Minimal Causal Code

This plain PyTorch loop exposes both branches and the parameter update. The supplied task schedule assigns the original dropout allocation to i2t and the remainder to t2i; it must enforce the intended token budget. There is no additional caption dropout. The model returns causal next-token logits over the joint vocabulary.

An initial EOS supplies an explicit preceding sequence boundary for this standalone illustration. Image tokenization, batching, padding, and budget construction are omitted; padded implementations must exclude padding from the loss.

import torch
import torch.nn.functional as F


def make_sequence(image_tokens, caption_tokens, task):
    image_span = [SOV, *image_tokens, EOV]
    if task == "t2i":
        # Keep caption-conditioned image generation.
        payload = [*caption_tokens, *image_span, EOS]
    elif task == "i2t":
        # Replace the dropped-caption slot:
        # image prediction first, caption prediction afterward.
        payload = [*image_span, *caption_tokens, EOS]
    else:
        raise ValueError(task)
    return [EOS, *payload]


model.train()
for image_tokens, caption_tokens, task in scheduled_examples:
    # task is "t2i" or "i2t"; the schedule matches the allocation.
    sequence = make_sequence(image_tokens, caption_tokens, task)
    tokens = torch.tensor(sequence, dtype=torch.long, device=device)
    inputs = tokens[:-1].unsqueeze(0)
    targets = tokens[1:].unsqueeze(0)

    optimizer.zero_grad()
    logits = model(inputs)  # causal: no access to future targets
    loss = F.cross_entropy(
        logits.reshape(-1, logits.size(-1)),
        targets.reshape(-1),
    )
    loss.backward()
    optimizer.step()  # update the shared model on this sequence

On t2i examples, image targets can attend to the caption. On i2t examples, image targets cannot attend to the later caption, but caption targets can attend to the image. Training covers the sequence, not just its final continuation; masking out image targets in i2t would break the image-training accounting.

What You Gain, and What You Lose

You gain caption supervision under the same pretraining budget while preserving conditional and unconditional image-token counts in the replacement setup. Across the tested model sizes, captioning negative log-likelihood decreases substantially while the reported text-to-image metrics remain comparable to baseline.

You do not preserve every target or obtain a universally interchangeable objective. Replacement adds caption predictions after image prediction, changes the post-image termination target, and requires retaining the caption-conditioned generation branch. Equal image-token counts also do not guarantee equal outcomes across schedules: the separate staged-versus-mixed comparison exposes a captioning–image-generation tradeoff at the larger allocation.

Quiz

What does the replacement recipe preserve, and what changes?

It preserves the budget by eliminating caption-conditioned image generation.

It guarantees identical generation quality for every schedule.

It preserves conditional and unconditional image-token counts but changes post-image targets to include captions.

It preserves every prediction target and only changes sequence order.

$ check