MR2-Molmo2-4B-RM

MR²-Molmo2-4B-RM

MR²-Molmo2-4B-RM is a multi-response multimodal reward model built on top of Molmo2-4B. Unlike conventional reward models that score each response independently, our model scores all N candidate responses in a single forward pass by concatenating them into one sequence, enabling direct comparative reasoning across candidates while achieving up to N× inference speedup.

This model is described in our paper: You Only Judge Once: Multi-response Reward Modeling in a Single Forward Pass (arXiv).

Quick links:

• 📃 Paper
• 💾 Benchmark Datasets
• 💻 Code

Key Features

• Single-pass multi-response scoring: All N candidates processed in one forward pass
• Cross-response comparative reasoning: Under causal attention, each response attends to all preceding responses
• Up to N× efficiency gain: Replaces N independent forward passes (discriminative) or C(N,2) pairwise comparisons (generative)
• Supports both image and video inputs
• State-of-the-art on six multimodal reward benchmarks with only 4B parameters

Results

Proprietary models (as judge)

Open-source VLMs (as judge)

Generative reward models

Discriminative reward models (including ours)

VL-RB: VL-RewardBench (macro pairwise acc.); MM-RB: Multimodal RewardBench (pairwise acc.); MMRLHF: MM-RLHF RewardBench (pairwise acc.); MR²B-I: MR²Bench-Image (best-of-4 acc.); VRB: VideoRewardBench (macro pairwise acc.); MR²B-V: MR²Bench-Video (best-of-4 acc.).

Quick Start

Setup

git clone https://github.com/yinuoyang01/multi-response-rm.git
cd multi-response-rm
pip install -r requirements.txt

Basic Usage: Score N Candidate Responses

import json
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText
from huggingface_hub import hf_hub_download
from PIL import Image

from mr2rm.models.reward_model import MultiResponseRewardModel
from mr2rm.data.dataset import add_resp_sep_token, RESP_SEP_TOKEN

model_id = "yinuoy/MR2-Molmo2-4B-RM"

# 1. Processor (register the <|resp_sep|> special token)
processor = AutoProcessor.from_pretrained(
    model_id, trust_remote_code=True, dtype="auto", device_map="auto",
)
add_resp_sep_token(processor.tokenizer)

# 2. Base model + reward model (value-head config from reward_model_config.json)
base_model = AutoModelForImageTextToText.from_pretrained(
    model_id, trust_remote_code=True, dtype="auto", device_map="auto",
)
config_path = hf_hub_download(repo_id=model_id, filename="reward_model_config.json")
with open(config_path) as f:
    rm_config = json.load(f)

reward_model = MultiResponseRewardModel(
    base_model=base_model,
    value_head_type=rm_config["value_head_type"],
    value_head_hidden_dim=rm_config["value_head_hidden_dim"],
    value_head_activation=rm_config["value_head_activation"],
    resp_repr_mode=rm_config["resp_repr_mode"],
)

# 3. Load value-head weights
vh_path = hf_hub_download(repo_id=model_id, filename="value_head.pt")
reward_model.value_head.load_state_dict(torch.load(vh_path, map_location="cpu"))
device = next(reward_model.base_model.parameters()).device
dtype = next(reward_model.base_model.parameters()).dtype
reward_model.value_head = reward_model.value_head.to(device=device, dtype=dtype)
reward_model.eval()

# 4. Build input: user=(prompt+image), assistant=(responses joined by <|resp_sep|>)
image = Image.open("example.jpg").convert("RGB")
prompt = "Describe this image."
responses = [
    "A golden retriever sitting on grass.",
    "A dog in a park on a sunny day.",
    "There is an animal outside.",
    "I don't know.",
]
sep = f"\n\n{RESP_SEP_TOKEN}\n\n"
assistant_text = sep.join(responses)
messages = [
    {"role": "user", "content": [dict(type="image", image=image), dict(type="text", text=prompt)]},
    {"role": "assistant", "content": [dict(type="text", text=assistant_text)]},
]
inputs = processor.apply_chat_template(
    messages, tokenize=True, return_tensors="pt", return_dict=True,
)
inputs = {k: v.to(device) for k, v in inputs.items()}

# 5. Locate the last-token position of each response (just before each <|resp_sep|>, plus end-of-sequence for the final one)
sep_token_id = processor.tokenizer.convert_tokens_to_ids(RESP_SEP_TOKEN)
input_ids = inputs["input_ids"][0]
sep_positions = (input_ids == sep_token_id).nonzero(as_tuple=True)[0].tolist()
end_positions = [p - 1 for p in sep_positions] + [input_ids.size(0) - 1]
resp_indices = torch.tensor([end_positions], device=device)

# 6. One forward pass — scores for all N responses
with torch.inference_mode():
    (scores,) = reward_model(
        input_ids=inputs["input_ids"],
        attention_mask=inputs.get("attention_mask"),
        resp_indices=resp_indices,
        **{k: v for k, v in inputs.items() if k not in ["input_ids", "attention_mask"]},
    )
rewards = scores[0].tolist()

print("Scores:", rewards)
print("Best response:", responses[rewards.index(max(rewards))])

Using for Best-of-N Selection

# Score N candidates and pick the best
best_idx = scores.argmax().item()
best_response = responses[best_idx]

Using as GRPO Reward

# For each rollout group of N responses, compute scores in one forward pass
# Use scores as rewards for GRPO advantage computation

Model Details

• Base model: Molmo2-4B
• Architecture: Molmo2-4B backbone + 2-layer MLP value head (hidden dim 1024, SiLU); last-token hidden state before <|resp_sep|> as the per-response representation
• Objective: Cross-entropy over N response scores (single forward pass)
• Training data: 436K preference samples drawn from 10 public datasets — MM-RLHF, LLaVA-Critic-113k, RLAIF-V, VLFeedback, POVID, WildVision, Tulu, Skywork-Reward, Nectar, PKU-SafeRLHF
• Training: LoRA rank 128, 3 epochs, effective batch size 64 — see the paper for full hyperparameters

License

This model is licensed under Apache 2.0, consistent with the Molmo2 base model license. It is intended for research and educational use. This model is trained on third-party datasets that are subject to their respective licenses; please review the sources listed above to determine if this model is appropriate for your use case.

Citation

@misc{yang2026judgeoncemultiresponsereward,
      title={You Only Judge Once: Multi-response Reward Modeling in a Single Forward Pass},
      author={Yinuo Yang and Zixian Ma and Manasi Ganti and Jieyu Zhang and Ranjay Krishna},
      year={2026},
      eprint={2604.10966},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2604.10966},
}