Skip to content

Add image-text contrastive loss to Blip2ForImageTextRetrieval - #48583

Open
xquantize wants to merge 1 commit into
huggingface:mainfrom
xquantize:blip2-itc-loss
Open

Add image-text contrastive loss to Blip2ForImageTextRetrieval#48583
xquantize wants to merge 1 commit into
huggingface:mainfrom
xquantize:blip2-itc-loss

Conversation

@xquantize

@xquantize xquantize commented Sep 6, 2026

Copy link
Copy Markdown

CPU CI GPU run-slow

What does this PR do?

Blip2ForImageTextRetrieval already computes ITC and ITM logits, but stops there — there's no loss path, so the class can't be used for training or fine-tuning. This adds the contrastive (ITC) half of that, as discussed in #34019.

Two things are needed for a usable ITC loss:

  1. A learnable temperature. The original LAVIS checkpoints carry a temp parameter, but the conversion script drops it. convert_blip_2_original_to_pytorch.py L223 asserts it away as an unexpected key, so converted models have no temperature to scale similarities with. This PR adds logit_scale as a learnable parameter on the class, constant-initialised from a new Blip2Config.logit_scale_init_value (2.6592, matching CLIP). It does not read the LAVIS temp value; mapping that through conversion would be a reasonable follow-up.

  2. The loss itself. A new return_loss argument on forward computes symmetric cross entropy over logits_per_image and logits_per_text, each scaled by logit_scale.exp(). The two directions are averaged and returned in Blip2ImageTextMatchingModelOutput.loss (and in the tuple path when return_dict=False).

Backwards compatible. return_loss defaults to False, and the returned logits stay unscaled either way, so existing inference behaviour is unchanged. There's a test for this.

Scope. ITC only. return_loss=True with use_image_text_matching_head=True raises a ValueError rather than silently doing the wrong thing. This is deliberately not the full QFormer pretraining recipe from #34019: no hard-negative ITM sampling, no ITG. @qgallouedec noted that's a training recipe rather than a forward(labels=...) addition. Partial progress on #34019.

Test plan

New tests in tests/models/blip_2/test_modeling_blip_2.py:

  • loss is None when return_loss is not passed
  • return_loss=True gives a finite scalar loss
  • logits are identical with and without the loss computation
  • backward reaches logit_scale.grad
  • ITM head + return_loss=True raises ValueError

Before submitting

Who can review?

@qgallouedec — you offered to look at a scoped draft on this class in #34019.

@github-actions

github-actions Bot commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: blip_2

@github-actions

github-actions Bot commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 34065895245
Result: success | Grafana metrics are not available yet.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant