AhmedSSabir/BERT-CNN-Reward-Caption-Reranker

BERT-CNN Reward Caption Reranker

A BERT-CNN caption reranker fine-tuned with a listwise expected-CIDEr reward. Given the nine beam-search captions of an image and image-context words (objects, scene labels), it picks the caption most likely to match human references. The model does not read images or generate captions.

It starts from the BERT-CNN visual semantic relatedness model of Visual Semantic Relatedness Dataset for Image Captioning and adds a second training stage. No new human annotation is used: the reward comes from existing COCO reference captions.

It also improves a recent vision-language model. Without any retraining, choosing among Google's Gemma 4 12B nine beam captions raises CIDEr from 0.9624 to 0.9783 on 5,000 COCO Karpathy test set. The gain varies by captioner and is absent on BLIP-2 OPT-2.7B COCO. Details and examples are under Evaluation.

Method	BLEU-4	METEOR	ROUGE-L	CIDEr	SPICE	BERTScore
Baseline: Gemma 4 12B, top beam	0.2441	0.2739	0.5288	0.9624	0.2248	0.9438
Gemma 4 12B + this model	0.2458	0.2766	0.5314	0.9783	0.2267	0.9448

Training

Stage I: relatedness pre-training (the paper's BERT-CNN). BERT-base with a CNN head, trained with binary cross-entropy on the visual semantic relatedness dataset (330,790 context/caption pairs, related if the similarity score is at least 0.3).

Stage II: listwise reward fine-tuning (this model).

1. A frozen BLIP captioner (Salesforce/blip-image-captioning-base) writes nine beam-search captions for 10,000 COCO training images.
2. Each candidate gets a reward r_k = CIDEr(c_k, references) against that image's COCO human captions.
3. The reranking score is s_k = b_k + w * f(v, c_k). Here f is the BERT-CNN logit, b_k is a learned bonus for beam position k, and w is a learned weight.
4. The model maximizes the exact expected reward over the nine candidates, E_{k ~ softmax(s)}[r_k], with the mean reward as baseline. This is a policy-gradient objective computed without sampling.
5. Training: 4 epochs (5,000 steps) with Adam. Learning rates: 1e-6 for BERT, 1e-5 for the CNN head, 1e-2 for b, w. It took 7 minutes on one RTX A5000.

Usage

The model reads the visual context (text labels extracted from the image) and the captions. It does not read the image itself, and this repository does not include the label extractor, so you supply the labels.

from huggingface_hub import snapshot_download
import sys

path = snapshot_download("AhmedSSabir/BERT-CNN-Reward-Caption-Reranker")
sys.path.insert(0, path)
from inference import CaptionReranker

reranker = CaptionReranker(path, device="cpu")  # "cuda" for a GPU
beams = [  # exactly nine captions, in beam-search order (best beam first)
    "A large black and gold clock stands against a blue sky.",
    "A black and gold street clock stands against a blue sky.",
    "A large black and gold clock stands under a blue sky.",
    "A large black and gold street clock stands against a blue sky.",
    "A large black and gold street clock stands under a blue sky.",
    "A black and gold street clock stands under a blue sky.",
    "A large ornate street clock stands against a blue sky.",
    "A large ornate clock stands in front of a building.",
    "A large black and gold clock stands in a city.",
]
visual_context = "wall clock analog tree"  # text labels extracted from the image
best, scores = reranker.rerank(visual_context, beams)
print(beams[best])  # -> A large ornate clock stands in front of a building.

The beam bonus means rerank needs the nine captions in the generator's beam order. For other inputs, reranker.logits(context, captions) returns the raw BERT-CNN scores.

Evaluation

Every captioner below gives nine beam-search candidates per image on the same 5,000 Karpathy test images, and is compared with its own top beam. All rows use the same scorers.

In-domain test: BLIP 9-beam candidates for the same 5,000 images.

Method	CIDEr
Best Beam	1.0632
This model	1.1470

Gemma 4 12B candidates. google/gemma-4-12B-it, prompted for a short COCO-style caption with nine beams, on the same 5,000 images. This model was not retrained on them.

Method	BLEU-4	METEOR	ROUGE-L	CIDEr	SPICE	BERTScore
Gemma 4 12B, beam 1 of 9	0.2441	0.2739	0.5288	0.9624	0.2248	0.9438
Gemma 4 12B + this model	0.2458	0.2766	0.5314	0.9783	0.2267	0.9448

Gemma 4 examples, CIDEr per caption. The model replaces Gemma's top caption on 1,298 of 5,000 images: CIDEr improves on 751, drops on 543 and ties on 4. The rows below are the first three improvements and the first drop in a random draw from those images.

Image	Context words	Gemma 4 top caption	This model picks
COCO 27585	wall clock analog tree	A large black and gold clock stands against a blue sky. (0.70)	A large ornate clock stands in front of a building. (1.25)
COCO 320370	croquet ball soccer tree	A man in a red shirt stands on a green field. (0.61)	A man in a red shirt stands on a green lawn. (0.70)
COCO 300233	mixing bowl corn food	Ingredients and a food processor are arranged on a gray surface. (0.40)	Ingredients and a food processor are laid out on a table. (0.47)
COCO 419228 (failure)	ice bear polar	A white polar bear walks on rocks in an enclosure. (2.54)	A large white polar bear walks on a rocky surface. (1.85)

Summary: CIDEr for each captioner's own top beam and with this model.

Captioner	Top beam	+ This model
BLIP-base (training captioner)	1.0632	1.1470*
Qwen3-VL-8B	0.5723	0.5927*
Gemma 4 12B	0.9624	0.9783*
BLIP-large	1.1509	1.1638*
Florence-2-large-ft	1.4099	1.4203*
BLIP-2 OPT-2.7B COCO	1.4034	1.3960

* Significantly different from the top beam. Best result in bold.

The gain shrinks as the captioner gets stronger. This model was not retrained on any captioner in this table except BLIP-base.

All metrics, nine beams each. Bold marks the better result for each captioner.

System	BLEU-4	METEOR	ROUGE-L	CIDEr	SPICE	BERTScore
Qwen3-VL-8B, Best Beam	0.0669	0.1872	0.3357	0.5723	0.1582	0.9223
Qwen3-VL-8B + This model	0.0710	0.1910	0.3404	0.5927	0.1613	0.9231
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━
Gemma 4 12B, Best Beam	0.2441	0.2739	0.5288	0.9624	0.2248	0.9438
Gemma 4 12B + This model	0.2458	0.2766	0.5314	0.9783	0.2267	0.9448
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━
BLIP-large, Best Beam	0.3004	0.2993	0.5378	1.1509	0.2381	0.9338
BLIP-large + This model	0.3035	0.2996	0.5389	1.1638	0.2395	0.9342
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━
Florence-2-large-ft, Best Beam	0.4079	0.3214	0.6121	1.4099	0.2498	0.9545
Florence-2-large-ft + This model	0.4090	0.3220	0.6143	1.4203	0.2512	0.9549
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━	━━━━━━━━
BLIP-2 OPT-2.7B COCO, Best Beam	0.3952	0.3256	0.6061	1.4034	0.2630	0.9453
BLIP-2 OPT-2.7B COCO + This model	0.3919	0.3251	0.6056	1.3960	0.2628	0.9454

Limitations

• Use it as a reranker, not a relatedness classifier. Reward training shifts the scores upward. For relatedness probabilities, use the stage I BERT-CNN.
• The beam bonus was learned for 9-beam candidates. Other beam widths may need retraining. It was trained on BLIP candidates; gains carry over to other captioners but are smaller there, and absent on BLIP-2 OPT-2.7B COCO.
• A reranker is a cheap way to improve an existing captioner, but it can only choose among its candidates, so it does not replace a stronger captioner. On the same images, COCO-tuned captioners score higher on their own (single caption, CIDEr: Florence-2-large-ft 1.445, BLIP-2 OPT-2.7B COCO 1.369). On their nine-beam candidates this model adds a small gain for Florence-2 and none for BLIP-2.
• The best validation epoch was the last epoch, so longer training was not explored.

Citation

@InProceedings{Sabir_2023_CVPR,
  author    = {Sabir, Ahmed and Moreno-Noguer, Francesc and Padr{\'o}, Llu{\'i}s},
  title     = {Visual Semantic Relatedness Dataset for Image Captioning},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops},
  month     = {June},
  year      = {2023},
  pages     = {5598--5606},
  doi       = {10.1109/CVPRW59228.2023.00592},
  url       = {https://openaccess.thecvf.com/content/CVPR2023W/O-DRUM/html/Sabir_Visual_Semantic_Relatedness_Dataset_for_Image_Captioning_CVPRW_2023_paper.html}
}