srpo-multimodal-reflection-rl

name	srpo-multimodal-reflection-rl
title	SRPO: Enhancing Multimodal LLM Reasoning via Reflection-Aware Reinforcement Learning
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2506.01713
keywords	["Multimodal","Reinforcement Learning","Reflection","Self-Correction"]
description	Teach multimodal language models to reflect on their reasoning and improve answers through structured RL training.

SRPO: Teach MLLMs to Think Out Loud and Correct Themselves

Multimodal large language models struggle with complex reasoning because they lack the ability to reflect: to pause, analyze their work, and correct mistakes. SRPO adds explicit self-reflection to MLLMs through a two-stage framework. First, generate high-quality reflection examples using an advanced model. Second, train using Group Relative Policy Optimization to reward both correct answers and meaningful, concise reflections that avoid redundancy. This teaches models to reason visually, express uncertainty, and self-correct—improving performance on tasks like visual math and detailed understanding.

Core Concept

Explicit reflection is a trainable skill. Most MLLMs generate answers directly; they don't ask themselves "Wait, let me reconsider this visual detail." SRPO makes reflection explicit and trainable by creating datasets of reflection-augmented reasoning, then optimizing a reward function that values both answer correctness and reflection quality. Good reflections are concise, identify errors, propose corrections, and avoid repeating information already stated.

Architecture Overview

• Reflection Dataset Construction: Use advanced MLLM to generate high-quality reflections on reasoning problems; reflections explain uncertainty, identify visual details, propose corrections
• Dual-Component Rewards: Separate rewards for answer correctness and reflection quality; avoid rewarding verbose or redundant reflections
• GRPO Training Loop: Group Relative Policy Optimization compares samples within groups, upweighting high-reward reflections while downweighting low-reward ones
• Baseline Multimodal Model: Standard MLLM architecture (vision transformer + language model); no architectural changes needed
• Benchmark Integration: Training on visual math (MathVista), detailed understanding (MMMU-Pro), and multimodal reasoning benchmarks

Implementation

This implementation demonstrates reflection-aware RL for multimodal models.

First, build a reflection dataset generator using an advanced model:

import torch
from transformers import CLIPVisionModel, CLIPProcessor, AutoTokenizer, AutoModelForCausalLM
from typing import List, Dict
from dataclasses import dataclass

@dataclass
class ReflectionExample:
    image_id: str
    problem_text: str
    image: torch.Tensor
    initial_answer: str
    reflection: str
    correct_answer: str
    is_correct: bool

class ReflectionDatasetGenerator:
    """Generate reflection-augmented reasoning examples."""

    def __init__(self, mllm_name: str = "openai/clip-vit-large-patch14"):
        self.processor = CLIPProcessor.from_pretrained(mllm_name)
        self.vision_model = CLIPVisionModel.from_pretrained(mllm_name)
        self.tokenizer = AutoTokenizer.from_pretrained("gpt2-large")
        self.language_model = AutoModelForCausalLM.from_pretrained("gpt2-large")

    def generate_reflection(self, problem_text: str, image: torch.Tensor,
                           initial_answer: str) -> str:
        """
        Generate high-quality reflection using advanced model.
        Reflection format:
        - Identifies what's visible in image
        - Explains reasoning
        - Identifies errors if any
        - Proposes correction
        """
        # Process image
        image_inputs = .processor(images=[image], return_tensors=)

        
        prompt = 

        
        reflection =  \
                     \
                     \
                    

         reflection

     () -> [ReflectionExample]:
        
        examples = []

         i, problem  (problems):
            
            initial_answer = ._get_initial_answer(problem)

            
            reflection = .generate_reflection(
                problem[],
                problem[],
                initial_answer
            )

            
            is_correct = (initial_answer == problem[])

            example = ReflectionExample(
                image_id=,
                problem_text=problem[],
                image=problem[],
                initial_answer=initial_answer,
                reflection=reflection,
                correct_answer=problem[],
                is_correct=is_correct
            )
            examples.append(example)

         examples

     () -> :
        
           


generator = ReflectionDatasetGenerator()
reflection_examples = generator.create_reflection_dataset(
    [
        {
            : torch.randn(, , ),
            : ,
            : 
        }
    ]
)

Implement dual-component reward model:

class ReflectionRewardModel:
    """Score both answer correctness and reflection quality."""

    def __init__(self):
        # Answer evaluator: checks correctness
        self.answer_verifier = self._build_verifier()
        # Reflection evaluator: scores quality
        self.reflection_scorer = self._build_scorer()

    def _build_verifier(self):
        """Verifier that checks answer correctness."""
        def verify(answer: str, correct_answer: str) -> float:
            if answer.strip().lower() == correct_answer.strip().lower():
                return 1.0
            return 0.0
        return verify

    def _build_scorer(self):
        """Scorer that evaluates reflection quality."""
        def score_reflection(reflection: str, problem: str, image_analysis: str) -> float:
            # Multi-component reflection score
            score = 0.0

            # 1. Conciseness: avoid verbose reflections
            word_count = len(reflection.split())
            if word_count < :
                score += 
             word_count < :
                score += 
            

            
            visual_keywords = [, , , , , ]
             (kw  reflection.lower()  kw  visual_keywords):
                score += 

            
            error_keywords = [, , , , ]
             (kw  reflection.lower()  kw  error_keywords):
                score += 

            
               reflection.lower()    reflection.lower():
                score += 

             (, score)

         score_reflection

     () -> :
        
        
        answer_reward = .answer_verifier(
            generated_answer,
            example.correct_answer
        )

        
        reflection_reward = .reflection_scorer(
            generated_reflection,
            example.problem_text,
              
        )

        
        combined_reward =  * answer_reward +  * reflection_reward

         {
            : answer_reward,
            : reflection_reward,
            : combined_reward,
            : {
                : answer_reward > ,
                : reflection_reward
            }
        }


reward_model = ReflectionRewardModel()

test_example = reflection_examples[]
gen_answer = 
gen_reflection = 

rewards = reward_model.compute_reward(test_example, gen_answer, gen_reflection)
()
()
()

Implement GRPO training loop for reflection-augmented answers:

class ReflectionGRPOTrainer:
    """Group Relative Policy Optimization for reflection-aware reasoning."""

    def __init__(self, model, reward_model: ReflectionRewardModel,
                 learning_rate: float = 1e-5):
        self.model = model
        self.reward_model = reward_model
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

    def generate_with_reflection(self, problem_text: str, image: torch.Tensor,
                                temperature: float = 0.7):
        """
        Generate both answer and reflection from multimodal model.
        """
        prompt = f"Problem: {problem_text}\nThink step-by-step and reflect."

        # In practice: use actual MLLM generation
        answer = "placeholder_answer"
        reflection = "placeholder_reflection"

        return answer, reflection

    def compute_grpo_loss(self, examples: List[ReflectionExample],
                         group_size: int = 4) -> torch.Tensor:
        """
        Compute GRPO loss across groups of examples.
        Within each group, upweight high-reward samples.
        """
        losses = []

        for i in range(0, len(examples), group_size):
            batch = examples[i:i+group_size]

            # Generate and score all samples in group
            group_rewards = []
            group_log_probs = []

             example  batch:
                
                answer, reflection = .generate_with_reflection(
                    example.problem_text,
                    example.image
                )

                
                reward_dict = .reward_model.compute_reward(
                    example,
                    answer,
                    reflection
                )
                group_rewards.append(reward_dict[])

                
                log_prob =   
                group_log_probs.append(log_prob)

            
            rewards_tensor = torch.tensor(group_rewards, dtype=torch.)
            normalized_rewards = (rewards_tensor - rewards_tensor.mean()) / \
                                (rewards_tensor.std() + )

            
             log_prob, norm_reward  (group_log_probs, normalized_rewards):
                loss = -log_prob * norm_reward
                losses.append(loss)

         torch.stack(losses).mean()  losses  torch.tensor()

     () -> :
        
        loss = .compute_grpo_loss(examples, group_size=)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(.model.parameters(), )
        .optimizer.step()
        .optimizer.zero_grad()

         {: loss.item()}

Practical Guidance

Aspect	Details
Reflection Dataset Size	5k-20k examples sufficient; use advanced model to bootstrap, then filter
Answer/Reflection Weight	Start 70/30; visual tasks may need 60/40, pure reasoning 80/20
Group Size	4-8 samples; larger groups provide better relative comparisons
Reflection Length Target	50-150 tokens optimal; reward brevity to avoid verbosity
Training Epochs	3-5 epochs on reflection data; monitor for overfitting

When to Use:

• Multimodal reasoning tasks requiring visual understanding and self-correction
• Need explainability: reflections show model's reasoning process
• Combining visual and textual information where tradeoffs exist
• Training on benchmarks like MathVista, MMMU where visual math matters
• Want to improve reasoning without scaling model size

When NOT to Use:

• Tasks where reflection adds latency costs that outweigh benefits
• Models already achieving ceiling performance (reflection won't help)
• Domains without clear visual-semantic alignment
• Applications requiring fast inference where reflection generation is prohibitive
• Few-shot learning where reflection dataset can't be generated

Common Pitfalls:

• Reflection dataset quality low: biased initial model generates bad reflections; validate manually
• Reward model misalignment: if reflection scorer doesn't match human judgment, RL optimizes wrong objective
• Verbosity explosion: models learn to generate long reflections if not penalized; enforce length limits
• Reflection-answer mismatch: models generate good reflections but poor answers; weight answer reward heavily
• Overfitting to reflection dataset: use diverse generation strategies during data collection

Reference

SRPO: Enhancing Multimodal LLM Reasoning via Reflection-Aware Reinforcement Learning https://arxiv.org/abs/2506.01713