shahedm2001

qwen3-vl-2b-cataract-grpo

README

License: apache-2.0

📋 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 LLM_MARKDOWN_PROTECTED_1.

🔬 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:

Low-Cardinality Categorical Tasks: 4-choice Multiple Choice Questions (MCQs) have a random-chance baseline of 25.0%.
Dense Continuous Tasks: Temporal Localization ( $t IoU$ ) and Transition Boundary Detection ( $∣Δ t ∣$ ) offer smooth, continuous reward gradients in $[0.05, 1.05]$ .
High-Cardinality Categorical Tasks: 13-class surgical phase recognition ( $P 01 $ to $P$ ) has a random-chance baseline of only .

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:

text

┌─────────────────────────────────────────────────────────────────────────────┐│                       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, $β = 0.04$ .
Empirical Diagnostic & The Zero-Advantage Trap:

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:

A^i ​ = std (R) + ϵ R i ​ - mean (R) ​

when all rollouts receive identical rewards, $std (R) = 0 ⟹ 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 to per prompt ().

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

Table with columns: Hyperparameter, Stage 1 (Foundation GRPO), Stage 2 (Phase Specialization), Rationale
Hyperparameter	Stage 1 (Foundation GRPO)	Stage 2 (Phase Specialization)	Rationale
Base Model Initialization	SFT Stage 2 Checkpoint	Stage 1 Merged 16-bit Weights	Progressive policy refinement
Training Samples	4,252 video clips	506 curated clips	80% Phase focus + 20% Replay anchors
Optimizer Steps	532 steps (1 epoch)	64 steps (1 epoch)	Short targeted specialization
Rollouts per Prompt ()

🎯 Task Curriculum & Deterministic Reward Formulations

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

R total ​ = R task ​ + 0.05 \times R format ​

Within each slice, group relative advantage normalization is computed across $G$ rollouts:

A^i ​ = std ({R 1 ​, \dots, R G ​})

Clinical Task Families

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

R task ​ = ∣ [t^1 ​, t^

Phase Boundary Transition Detection (Continuous $R task \in (0.0, 1.0]$ ):
- Evaluates the predicted transition timestamp $t^$ against true boundary $t *$ using exponential decay with scale parameter :

R task ​ = exp (- 1.5 ∣ t^- t * ∣ ​)

Format Compliance Bonus ( $R format = 1.0$ ):
- Enforces valid JSON formatting matching {"explanation": "...", "answer": "..."}.

📊 Empirical Evaluation & Learning Dynamics by Data Group

1. Group 1: 4-Choice Surgical Video MCQs

Group 1 Surgical MCQs

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 $std (R) \to 0$ as policy outputs saturated the optimal choice.

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

Group 2 Phase Recognition

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 $std (R) = 0 ⟹ A$ and zero policy gradients.

3. Group 3: Continuous Fine-Grained Temporal Tasks

Group 3 Continuous Temporal Tasks

Temporal Localization & Boundary Transition Detection (475 prompt slices):
- Temporal Localization ( $t 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 IoU > 0.85$ .
- Rolling median boundary error decreased from (Early) down to (Late) and reached sub-second precision () during Stage 2.

4. Stage 2 Specialization Dashboard & Unified Policy Dynamics

Stage 2 Specialization Dashboard Unified Optimization Dynamics

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 KL $ ) stayed tightly contained within [0.015, 0.035].
Exploration Coverage: Active gradient groups ( $σ > 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

bash

pip install torch transformers accelerate qwen-vl-utils decord

Video Inference Example

python

import jsonimport torchfrom transformers import AutoProcessor, AutoModelForImageTextToTextfrom qwen_vl_utils import process_vision_info model_id = "shahedm2001/qwen3-vl-2b-cataract-grpo" # 1. Load model and processormodel = AutoModelForImageTextToText.from_pretrained(    model_id,    torch_dtype=torch.bfloat16,    device_map="auto")processor = AutoProcessor.from_pretrained(model_id) # 2. Formulate video analysis querymessages = [    {        "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 inputsprompt = 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 predictionwith 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: LLM_MARKDOWN_PROTECTED_36
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.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.

Learn more

Model Details

Model Provider