Skip to content

Token-Sparse Medical Multimodal Reasoning via Dual-Stream Reinforcement Learning

🕒 Published (v1): 2026-06-30 12:47 UTC · Source: Arxiv · link

Why this paper was selected

Dual-stream RL for token-sparse medical VLM reasoning; novel sparse-evidence approach

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

ViToS is a dual-stream reinforcement learning framework that jointly trains visual token pruning (VTP) and sparse-evidence reasoning in a single policy model for medical VLMs. By reformulating VTP as a policy-guided evidence selection process and introducing cross-feedback sequential optimization, it reduces visual token sequences to 77% of their original length while improving accuracy. On seven medical VQA benchmarks it achieves 108.27% relative performance on Lingshu-7B and ~Ă—4.5 inference speedup over heuristic VTP competitors.

Problem

Medical VLMs process images via dense, uniform visual token encoders, yet clinically relevant evidence occupies only a sparse subset of the image. Existing VTP methods (FastV, SparseVLM, VisionZip, etc.) apply heuristic, non-differentiable pruning rules that are disconnected from end-to-end task training, causing degraded or unstable performance. Existing medical RL methods (Med-R1, ViTAR) each optimize only one ability—either grounding or reasoning—rather than jointly internalizing both. No unified RL scheme for simultaneously learning evidence localization and sparse-token reasoning existed.

Method

Dual-stream RL (DS-RL): A single shared policy model \(\pi_\theta\) is trained with two cascaded branches operating under GRPO.

  1. Localization branch: Given input \(X = \{x_v, x_t\}\), the policy generates \(G\) rollouts \(\{o_1, \ldots, o_G\}\) each containing a predicted foreground bounding box \(b_{\text{pred}}\).

  2. Token-sparse reasoning branch (grounding-aware VTP): The visual encoder produces \(E \in \mathbb{R}^{N \times D}\). Conditioned on \(b_{\text{pred}}^{(g)}\), tokens are split into foreground \(E_{\text{kept}}^{(g)}\) and background \(E_{\text{others}}^{(g)}\). Each background token is fused to its most semantically similar foreground token via normalized cosine similarity and weighted aggregation, yielding compact representation \(\hat{E}_{\text{kept}}^{(g)}\). Positional embeddings of the pruned set are reassigned with RoPE to preserve relative spatial order. The decoder then generates \(G\) reasoning trajectories per pruned representation, producing an expanded rollout set \(\{o'_{g,j}\}_{g,j=1}^{G \times G}\).

Cross-feedback sequential optimization: To avoid gradient conflict from the coupled policy, training proceeds in two stages. First, only the localization branch is optimized via reward \(R^L_{\text{total}} = 3(1-\lambda)R^L_{\text{format}} + \lambda R_{\text{IoU}} + \lambda R^S_{\text{acc}}\), where \(R^S_{\text{acc}}\) (downstream reasoning accuracy) serves as a cross-feedback signal. Once localization converges and is frozen, the reasoning branch is optimized via \(R^S_{\text{total}} = 2(1-\lambda)R^S_{\text{format}} + \lambda R^S_{\text{acc}} + \lambda R^L_{\text{acc}}\). The GRPO objective clips the policy ratio and applies KL regularization toward the reference model \(\pi_{\text{ref}}\).

Training runs on 8Ă— NVIDIA H200 GPUs for 8 hours with veRL; inference uses the vLLM engine.

Key Contributions

  • Dual-stream RL framework with a unified policy model that jointly internalizes spatial grounding and token-sparse reasoning, framing VTP as policy-guided evidence selection rather than a static heuristic.
  • Cross-feedback sequential optimization that decouples gradients between branches, using each branch's accuracy signal as a reward for the other to resolve the coupled policy learning problem.
  • Trainable grounding-aware VTP: foreground tokens identified by bounding box prediction, background tokens fused via cosine-similarity weighting, with RoPE re-indexing to restore positional coherence.
  • Full compatibility with vLLM inference, enabling practical inference speedup alongside improved accuracy.

Results

  • Visual token sequence reduced to 77% of original length on average.
  • DS-RL + GTP on Lingshu-7B: average accuracy 68.95% vs. 63.68% baseline (+5.27 pp); 108.27% relative performance vs. baseline.
  • DS-RL + GTP on HuatuoGPT-Vision-7B: average accuracy 59.17% vs. 56.81% baseline; 104.16% relative performance.
  • Lingshu-32B: DS-RL + GTP average 70.21% vs. 66.46% baseline; SLAKE peaks at 90.62%, PathVQA at 85.37%.
  • Outperforms all heuristic VTP baselines (VisionZip: 98.82%, VScan: 99.44%, FastV: 94.65%, PDrop: 99.35% relative on Lingshu-7B vs. ViToS at 108.27%).
  • Outperforms medical RL baselines: Med-R1, MedVLM-R1, MedCCO, MedGVLThinker, ViTAR (best average performance).
  • Inference speedup: ~Ă—4.5 on Lingshu-7B and ~Ă—3.4 on Lingshu-32B vs. competing VTP methods; even without vLLM, completes in 31 min (7B) and 81 min (32B), faster than competitors.
  • Benchmarks: PathVQA, SLAKE, VQA-RAD, OmniMedVQA, PMC-VQA, MMMU-Med, MedXpertQA.

Limitations

  • Localization branch requires ground-truth bounding box annotations for IoU reward; applicability to datasets lacking spatial labels is unaddressed.
  • Two-stage sequential training increases training pipeline complexity and requires a convergence criterion between stages.
  • Cross-feedback reward weighting \(\lambda\) is a hyperparameter requiring tuning; sensitivity analysis is in the appendix but not fully excerpted here.
  • Evaluated only on 7B and 32B scales; performance at sub-7B or other model families is unexplored.
  • Grounding-aware VTP assumes a single dominant foreground region per query; pathologies spanning diffuse or multi-focal regions may not be well-captured by a single bounding box.
  • Generalization beyond medical imaging (natural image VQA) is not assessed; domain specificity of the approach is unclear.

Relevance to Vision-Language Models

This paper directly advances the VLM efficiency literature by demonstrating that visual token pruning can be learned end-to-end through RL rather than imposed as a post-hoc heuristic—a finding relevant to any VLM deployment where redundant visual context degrades reasoning. The cross-feedback sequential optimization strategy addresses gradient conflict in multi-objective RL fine-tuning, a broadly applicable challenge for VLMs with coupled sub-tasks (e.g., grounding + captioning, detection + VQA). The work sits at the intersection of GRPO-based reasoning (DeepSeek-R1 lineage) and visual token compression research (SparseVLM, LLaVA-PruMerge), showing these threads can be unified rather than traded off. For researchers tracking VLMs, ViToS establishes a template for domain-specific VLM training where sparse visual evidence is the norm.