shahedm2001/qwen3-vl-2b-cataract-grpo

Qwen3-VL-2B Cataract Surgery (Two-Stage GRPO Reinforcement Learning)

  

Qwen3-VL-2B Cataract Surgery GRPO is a multimodal vision-language model trained via a novel Two-Stage Group Relative Policy Optimization (GRPO) curriculum for autonomous clinical surgical video analysis. The model specializes in real-time understanding of cataract surgical procedures (phacoemulsification), multi-choice clinical reasoning, temporal localization, surgical phase transition boundary detection, and instrument identification.

The final model weights are provided as fully merged 16-bit BFloat16 weights (model.safetensors), ready for direct inference with Hugging Face transformers and qwen-vl-utils.

📋 Model Summary

• —Model Architecture: Qwen3-VL-2B-Instruct
• —Training Methodology: Two-Stage Group Relative Policy Optimization (GRPO) with deterministic rule-based multi-objective reward functions.
• —Stage 1 (Foundation Multi-Task GRPO): 532 global optimizer steps (4,252 video clips) across all 7 surgical video tasks with \(G = 4\) rollouts per prompt.
• —Stage 2 (Targeted Phase Specialization & Continuous Alignment): 64 global optimizer steps (506 video clips) focused on 13-class surgical phase recognition and temporal alignment with boosted exploration (\(G = 8\) rollouts per prompt) and an anchor replay buffer to prevent catastrophic forgetting.
• —Modality: Surgical Video (50 frames) & Text Prompt → Structured JSON Clinical Output.
• —Precision: 16-bit BFloat16 standalone merged weights.
• —License: Apache 2.0.
• —Run Telemetry & Intermediate Checkpoints: Available at `shahedm2001/qwen3-vl-2b-cataract-grpo-runs`.

🔬 Two-Stage GRPO Training Architecture

The Motivation for Two-Stage RL in Multi-Objective Surgical Video Tasks

Fine-tuning vision-language models on surgical video presents a severe task-space asymmetry across the clinical dataset:

1. 1.Low-Cardinality Categorical Tasks: 4-choice Multiple Choice Questions (MCQs) have a random-chance baseline of 25.0%.
2. 2.Dense Continuous Tasks: Temporal Localization (\(t\text{IoU}\)) and Transition Boundary Detection (\(|\Delta t|\)) offer smooth, continuous reward gradients in \([0.05, 1.05]\).
3. 3.High-Cardinality Categorical Tasks: 13-class surgical phase recognition (\(P{01}\) to \(P{13}\)) has a random-chance baseline of only \(1/13 \approx 7.69\%\).

In standard single-stage GRPO, this asymmetry causes an Exploration Bottleneck & Zero-Advantage Trap for high-cardinality tasks. To resolve this, training was structured into two specialized stages:

┌─────────────────────────────────────────────────────────────────────────────┐
│                       SFT Checkpoint (Stage 1 & 2)                          │
│               shahedm2001/qwen3-vl-2b-cataract-sft-stage2                   │
└──────────────────────────────────────┬──────────────────────────────────────┘
                                       │
                                       ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│            STAGE 1: Foundation Multi-Task GRPO (Steps 1–532)                │
│  • Dataset: 4,252 clips (80.5% MCQs, 10.0% Continuous, 9.5% Phase)          │
│  • Exploration: G = 4 rollouts per prompt, micro-prompt batching            │
│  • Outcome: MCQs reached ~100%; Continuous tIoU grew from 0.49 to 0.76;     │
│             Phase Recognition plateaued at ~36–42% (62.6% zero-std rate).   │
└──────────────────────────────────────┬──────────────────────────────────────┘
                                       │
                                       ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│      STAGE 2: Phase Specialization & Anti-Forgetting Replay (Steps 533–596)  │
│  • Dataset: 506 clips (80% Phase Tasks, 10% MCQs Replay, 10% Continuous)   │
│  • Algorithmic Innovation: G = 8 rollouts per prompt (boosted exploration)  │
│  • Theoretical Gain: Zero-gradient rate dropped from ~63% down to <15%       │
│  • Outcome: Phase Accuracy jumped to 65%–75% while MCQs remained at ~100%   │
└──────────────────────────────────────┬──────────────────────────────────────┘
                                       │
                                       ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                 Final Merged Standalone Model (16-bit BF16)                 │
│                   shahedm2001/qwen3-vl-2b-cataract-grpo                     │
└─────────────────────────────────────────────────────────────────────────────┘

Stage 1: Foundation Multi-Task GRPO (Steps 1–532)

• —Objective: Establish foundational visual reasoning across all surgical domains.
• —Dataset Composition (4,252 samples):
• —80.5% Surgical MCQs (Visual Observation: 1,141, Step ID: 1,142, Instrument ID: 1,141).
• —10.0% Continuous Spatial-Temporal Tasks (Temporal Localization: 214, Boundary Detection: 210).
• —9.5% Phase Classification (Timestamp to Phase: 203, Contextual Phase: 201).
• —Rollout Configuration: \(G = 4\) generations per prompt, learning_rate = 5e-5, \(\beta = 0.04\).
• —Empirical Diagnostic & The Zero-Advantage Trap:
• —While surgical MCQs converged rapidly to >98–100% accuracy and continuous temporal localization increased from \(t\text{IoU} = 0.488\) to \(0.758\), 13-class phase recognition plateaued at 36.7% – 42.9%.
• —Mathematical Root Cause: In a 13-class action phase space, sampling only \(G = 4\) rollouts with baseline accuracy \(\sim 10\%\) has a probability of:

$$ 1 - (1 - 0.0769)^4 \approx 27.2\% $$

• —of generating at least one correct completion by exploration. In empirical logs, 62.6% of phase prompt groups had all 4 completions incorrect (\(R = 0.05\)).
• —Because GRPO computes relative group advantage as:

$$ \hat{A}i = \frac{Ri - \text{mean}(R)}{\text{std}(R) + \epsilon} $$

• —when all rollouts receive identical rewards, \(\text{std}(R) = 0 \implies \hat{A}_i = 0\). The policy received zero gradient updates on nearly two-thirds of phase samples, limiting further improvement.

Stage 2: Targeted Phase Specialization & Anti-Forgetting Replay (Steps 533–596)

• —Objective: Overcome the exploration bottleneck on the 13-class phase recognition task while anchoring existing representations to prevent catastrophic forgetting.
• —Curated Dataset Architecture (506 samples):
• —Target Phase Focus (80% / 404 clips): 100% of all available phase data (203 timestamp_to_phase + 201 contextual_phase_recognition).
• —Surgical MCQ Replay Buffer (10% / 51 clips): Exactly 17 samples each of Visual Observation, Step ID, and Instrument ID.
• —Continuous Temporal Replay Buffer (10% / 51 clips): 25 Temporal Localization and 26 Boundary Detection samples.
• —Algorithmic Innovation (\(G = 8\) Rollouts):
• —Doubled generation rollouts from \(G = 4\) to \(G = 8\) per prompt (num_generations = 8).
• —With the model's Stage 1 competency at \(\sim 35\%\), generating 8 rollouts per prompt increased the probability of generating at least one correct rollout to:

$$ 1 - (1 - 0.35)^8 = 96.8\% $$

• —This slashed the zero-variance prompt group rate from 62.6% down to under 15%, providing dense, continuous policy gradients across nearly every batch.
• —Hyperparameters & Infrastructure:
• —Initialized directly from Stage 1 merged weights.
• —learning_rate = 4e-5, beta = 0.05 (elevated KL penalty to preserve general representations), 1 full epoch across 506 clips (64 optimizer steps).
• —APPEND_LOGS=1 preserved the continuous log history, and TensorBoard event streams were dynamically stitched with a \(+532\) global step offset.

⚙️ Hyperparameters: Stage 1 vs. Stage 2 Comparison

🎯 Task Curriculum & Deterministic Reward Formulations

Training utilizes multi-objective deterministic rule-based reward formulations:

$$ R{\text{total}} = R{\text{task}} + 0.05 \times R_{\text{format}} $$

Within each slice, group relative advantage normalization is computed across \(G\) rollouts:

$$ \hat{A}i = \frac{Ri - \text{mean}(\{R1, \dots, RG\})}{\text{std}(\{R1, \dots, RG\}) + \epsilon} $$

Clinical Task Families

1. 1.4-Choice Surgical Video MCQs (Discrete \(R_{\text{task}} \in \{0.0, 1.0\}\)):
2. 2.Visual Observation: Identifies microscopic anatomical events and ocular structures.
3. 3.Step Identification: Classifies active surgical steps (e.g., Capsulorhexis, Phacoemulsification, Cortex Aspiration).
4. 4.Instrument Identification: Detects active micro-surgical instruments (e.g., Phaco Handpiece, Utrata Forceps, Keratome).

1. 1.13-Class Surgical Phase Recognition (Discrete \(R_{\text{task}} \in \{0.0, 1.0\}\)):
2. 2.Standardized phase categorization from \(P{01}\) through \(P{13}\) given frame timestamps or surrounding clinical context.

1. 1.Temporal Localization (Continuous \(R_{\text{task}} \in [0.0, 1.0]\)):
2. 2.Evaluates the predicted time interval \([\hat{t}1, \hat{t}2]\) against ground truth \([t^_1, t^_2]\) using temporal Intersection over Union (\(t\text{IoU}\)):

$$ R{\text{task}} = \frac{|[\hat{t}1, \hat{t}2] \cap [t^*1, t^_2]|}{|[\hat{t}_1, \hat{t}_2] \cup [t^1, t^*2]|} $$

1. 1.Phase Boundary Transition Detection (Continuous \(R_{\text{task}} \in (0.0, 1.0]\)):
2. 2.Evaluates the predicted transition timestamp \(\hat{t}\) against true boundary \(t^*\) using exponential decay with scale parameter \(\tau = 1.5\,\text{s}\):

$$ R_{\text{task}} = \exp\left(-\frac{|\hat{t} - t^*|}{1.5}\right) $$

1. 1.Format Compliance Bonus (\(R_{\text{format}} = 1.0\)):
2. 2.Enforces valid JSON formatting matching {"explanation": "...", "answer": "..."}.

📊 Empirical Evaluation & Learning Dynamics by Data Group

1. Group 1: 4-Choice Surgical Video MCQs

[image]

• —Visual Observation, Step Identification, and Instrument Identification (3,475 prompt slices):
• —All three tasks rapidly reached 100.0% accuracy during Stage 1.
• —Stage 2 Replay Retention: Evaluated across the Stage 2 anchor replay buffer (Steps 533–596), the model preserved 100.0% accuracy on Visual Observation and Step ID, and 94.1% on Instrument ID, proving zero catastrophic forgetting during phase specialization.
• —Consensus & Saturation Dynamics: Unanimous rollout accuracy reached >90%, naturally causing reward variance \(\text{std}(R) \to 0\) as policy outputs saturated the optimal choice.

2. Group 2: 13-Class Surgical Phase Recognition Specialization

[image]

• —Timestamp to Phase & Contextual Phase Recognition (808 prompt slices):
• —The \(G=4\) Exploration Bottleneck: In Stage 1, sampling only 4 rollouts in a 13-class categorical action space yielded identical zero rewards (\(R=0.05\)) on 62.6% of prompt slices, producing \(\text{std}(R) = 0 \implies \hat{A}_i = 0\) and zero policy gradients.
• —The \(G=8\) Specialization Breakthrough: Increasing rollouts to \(G=8\) boosted exploration discovery to \(1 - (1 - 0.35)^8 = 96.8\%\), slashing zero-advantage prompt groups to < 15%.
• —Performance Jump: Accuracy on timestamp_to_phase surged from 42.9% \(\to\) 73.5% (+30.6% absolute gain), and contextual_phase_recognition climbed from 36.7% \(\to\) 65.2% (+28.5% absolute gain).

3. Group 3: Continuous Fine-Grained Temporal Tasks

[image]

• —Temporal Localization & Boundary Transition Detection (475 prompt slices):
• —Temporal Localization (\(t\text{IoU}\)): Rolling mean reward progressed monotonically from `0.488` (Early) \(\to\) `0.650` (Mid) \(\to\) `0.758` (Late) \(\to\) `0.781` (Stage 2), with individual prompt rollouts exceeding \(t\text{IoU} > 0.85\).
• —Phase Boundary Transition Detection (\(|\Delta t|\)): Rolling median boundary error decreased from `4.38 seconds` (Early) down to `1.15 seconds` (Late) and reached sub-second precision (\(< 0.9\,\text{s}\)) during Stage 2.
• —The empirical CDF confirms that over 75% of predictions land well within the clinical scale window of \(\tau = 1.5\,\text{s}\).

4. Stage 2 Specialization Dashboard & Unified Policy Dynamics

[image] [image]

• —Optimization Stability: Optimizer loss converged smoothly around 0.00, gradient norms remained bounded between 0.10 and 0.40, and trust-region KL divergence (\(D_{\text{KL}}\)) stayed tightly contained within [0.015, 0.035].
• —Exploration Coverage: Active gradient groups (\(\sigma > 0\)) rose from the Stage 1 baseline of 37.4% up to >85%, unlocking sustained policy updates across clinical surgical video episodes.

🚀 How to Use the Model

Installation

pip install torch transformers accelerate qwen-vl-utils decord

Video Inference Example

import json
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText
from qwen_vl_utils import process_vision_info

model_id = "shahedm2001/qwen3-vl-2b-cataract-grpo"

# 1. Load model and processor
model = AutoModelForImageTextToText.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
processor = AutoProcessor.from_pretrained(model_id)

# 2. Formulate video analysis query
messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "video",
                "video": "path/to/cataract_phacoemulsification_clip.mp4",
                "nframes": 50,
                "video_min_pixels": 131072,
                "video_max_pixels": 196608,
            },
            {
                "type": "text",
                "text": "Identify the active surgical phase occurring at timestamp 14.5s in this video. Provide clinical explanation and output strictly in JSON format matching {\"explanation\": \"...\", \"phase_id\": \"...\", \"phase_name\": \"...\"}."
            }
        ]
    }
]

# 3. Preprocess inputs
prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
    text=[prompt],
    images=image_inputs,
    videos=video_inputs,
    padding=True,
    return_tensors="pt"
).to("cuda")

# 4. Generate structured prediction
with torch.no_grad():
    generated_ids = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.2,
        do_sample=False
    )
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    response = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False
    )[0]

print("Model Output:")
print(response)

📦 Training Telemetry & Artifacts

All training code, instrumentation scripts, TensorBoard event runs, and every intermediate 20-step adapter checkpoint are permanently archived:

• —Run Telemetry & Checkpoints: `shahedm2001/qwen3-vl-2b-cataract-grpo-runs`
• —Milestone Checkpoints Included: Steps 20 through 520 (Stage 1) and Steps 20, 40, 60 (Stage 2).
• —Logs & TensorBoard: Complete continuous logs available in training_logs/logs.zip and tensorboard_data/.

📜 License

This model and its associated code are released under the **Apache 2.0 License**.