Source-linked AI summary

Rethinking Token-Level Policy Optimization for Multimodal Chain-of-Thought

Yunheng Li, Hangyi Kuang, Hengrui Zhang, Jiangxia Cao, Zhaojie Liu, Qibin Hou, Ming-Ming Cheng

arXiv:2603.22847v1cs.CV

TL;DR

LVLM RLVR methods commonly optimize multimodal CoT at sequence level, overlooking how visual grounding and exploratory uncertainty vary across tokens. PEPO analyzes these token dynamics, combines hidden-state-based perception priors with entropy through smooth gating, and consistently improves RL baselines across multimodal benchmarks while maintaining stable training. The method is implemented within GRPO and DAPO without auxiliary branches or additional supervision.

  • Problem

    Existing RLVR methods apply coarse or uniform supervision to multimodal CoT, while entropy-based advantages capture textual uncertainty but weakly reflect visual semantics.

  • Method

    PEPO derives a perception prior from response–visual hidden-state similarity and smoothly fuses it with token entropy to produce token-level advantage weights within GRPO or DAPO.

  • Results

    PEPO consistently outperforms GRPO and DAPO across geometry reasoning, visual puzzles, visual grounding, and few-shot classification while maintaining training stability.

  • Takeaways & Limitations

    Combining perceptual grounding with entropy-guided exploration provides a supported approach for fine-grained multimodal policy optimization.

  • Takeaways & Limitations

    Experiments cover recent 2B or 3B LVLMs and a curated benchmark set, not larger backbones, longer contexts, video understanding, or tool-augmented reasoning.

Abstract

from arXiv · show

Multimodal Chain-of-Thought (CoT) reasoning requires large vision-language models to construct reasoning trajectories that interleave perceptual grounding with multi-step inference. However, existing Reinforcement Learning with Verifiable Rewards (RLVR) methods typically optimize reasoning at a coarse granularity, treating CoT uniformly without distinguishing their varying degrees of visual grounding. In this work, we conduct a token-level analysis of multimodal reasoning trajectories and show that successful reasoning is characterized by structured token dynamics reflecting both perceptual grounding and exploratory inference. Building upon this analysis, we propose Perception-Exploration Policy Optimization (PEPO), which derives a perception prior from hidden state similarity and integrates it with token entropy through a smooth gating mechanism to produce token-level advantages. PEPO integrates seamlessly with existing RLVR frameworks such as GRPO and DAPO, requiring neither additional supervision nor auxiliary branches. Extensive experiments across diverse multimodal benchmarks demonstrate consistent and robust improvements over strong RL baselines, spanning geometry reasoning, visual grounding, visual puzzle solving, and few-shot classification, while maintaining stable training dynamics. Code: https://github.com/xzxxntxdy/PEPO

1 Introduction

Existing LVLM RLVR methods optimize CoT too coarsely, while successful multimodal reasoning combines perceptual grounding with exploratory inference. PEPO addresses this by assigning token-level weights from visual similarity and entropy, improving performance across diverse benchmarks.

  • Motivation: Sequence-level RLVR assigns uniform advantages across response tokens, failing to distinguish intermediate CoT contributions.Entropy-based token advantages encourage exploration but primarily capture textual uncertainty and weakly reflect visual semantics.
  • Motivation: Correct multimodal reasoning depends on compact subsets of visually aligned tokens that anchor the CoT process.Hidden-state similarity between response and visual tokens provides a modality-specific indicator of linguistic-perceptual alignment.
  • PEPO: PEPO converts hidden-state similarity into a perception prior and smoothly fuses it with token entropy to produce normalized token weights.The framework requires neither auxiliary branches nor additional supervision and refines sequence-level advantage estimation.
  • Contribution: The method is motivated by complementary roles: visual grounding anchors reasoning, while high-entropy tokens identify steps for exploring alternative paths.This token-level perception-exploration view is presented as a distinct perspective on multimodal reasoning.
  • Results: PEPO integrates with GRPO and DAPO and yields consistent gains across geometry, math/logic, visual puzzles, visual grounding, and few-shot classification.The reported improvements span multiple multimodal reasoning benchmarks and model backbones.

2 Related Work

Related work improves LVLM reasoning through RLVR, data construction, reward design, chain-of-thought supervision, and step-wise instruction tuning. These approaches advance multimodal reasoning but leave adaptive feedback and further framework development as active concerns.

  • RLVR for LVLMs: RLVR has become an effective approach for training LVLMs, with GRPO widely used for stable, critic-free policy optimization.Related work develops both data-centric training strategies and reward-centric methods for multimodal tasks.
  • Reasoning in LVLMs: Chain-of-thought supervision and step-wise instruction tuning encourage structured LVLM inference but remain limited by static supervision and lack adaptive feedback.Reinforcement learning is introduced as a way to refine reasoning beyond these constraints.

3 Methodology

PEPO addresses the coarse, uniform token supervision in GRPO by analyzing visual grounding and entropy as complementary signals, then using them to modulate token-level advantages. Its framework derives visual similarity from hidden-state correlations and combines normalized perception and exploration scores through smooth gating.

  • 3.1 Background and Motivation: GRPO applies each response's sequence-level advantage uniformly across all tokens, limiting optimization granularity for multimodal reasoning.This can create imbalance because visual grounding primarily determines correctness while textual reasoning receives more extensive gradient updates.
  • 3.2 Token-Level Analysis of Multimodal Reasoning: Visual similarity measures visual dependency through mean cosine similarity between response-token and vision-token hidden states across model layers.The analysis uses layer-wise hidden states and defines Mglob, Mhigh, and Mlow by aggregating token-level scores across responses.
  • 3.2 Token-Level Analysis of Multimodal Reasoning: Correct responses show rightward shifts in global and top-K visual similarity, while bottom-K similarity separates minimally.The pattern indicates that successful reasoning relies on a compact subset of visually aligned tokens rather than uniformly grounded responses.
  • 3.2 Token-Level Analysis of Multimodal Reasoning: Removing images causes substantially larger hidden-state shifts for high-visual-similarity tokens than for high-entropy tokens.This controlled perturbation supports visual similarity as a perceptual-grounding signal, while entropy captures a different aspect of reasoning dynamics.
  • 3.3 Perception-Exploration Policy Optimization: PEPO computes per-token visual similarity and entropy from response and vision hidden states and policy-model logits, respectively.Visual similarity supplies the perception prior, while entropy models uncertainty and exploratory transitions.
  • 3.3 Perception-Exploration Policy Optimization: PEPO min-max normalizes both scores, applies a smooth gate to their centered sum, and uses the resulting weights to redistribute sequence-level advantage across tokens.The gate is multiplied by visual similarity so entropy-driven modulation remains conditioned on visually grounded tokens; normalization preserves the overall advantage scale.

4 Experiments

PEPO is evaluated across multiple multimodal reasoning settings and consistently improves performance and training behavior over established RL baselines. Analyses indicate that combining visual grounding with exploration, calibrated token weighting, and multi-layer signals supports robust optimization.

  • Experiment Setup: PEPO is compared with GRPO, DAPO, and High-Entropy RL using Qwen2.5-VL-3B-Instruct and InternVL3-2B-Instruct across five multimodal task categories.The evaluation covers geometry reasoning, visual grounding, few-shot classification, visual puzzles, and scalability.
  • Main Results: PEPO improves visual grounding under domain shift, few-shot classification across 1-, 2-, and 4-shot settings, and both in-domain and out-of-domain visual puzzle benchmarks.Reported few-shot average gains over GRPO are +5.32 points on FGVC Aircraft and +1.46 points on Flower102.
  • Efficiency and training dynamics: PEPO maintains low computational overhead, with ρ below 1% across benchmarks, while shorter responses can preserve comparable or improved effective throughput.Training dynamics also show higher reward, stronger plateaus, increasing visual similarity, and more gradual response-length reduction for PEPOG.
  • Component analysis: The full visual-entropy formulation outperforms perception-only and exploration-only variants, whose limitations are reduced reasoning diversity and unstable optimization, respectively.The component analysis identifies complementary roles for visual similarity and entropy.
  • Ablations: Progressive scheduling, per-sample min-max normalization, gated weighting, and aggregation across all layers are important for performance and stable token weighting.Single shallow, intermediate, or deep layer configurations underperform the all-layer default.

5 Conclusions

The paper concludes that PEPO provides a principled token-level approach for multimodal reasoning by jointly modeling visual grounding and exploratory inference. Experiments across architectures and task types show consistent gains with stable and scalable training.

  • 5 Conclusions: PEPO couples visual similarity and token entropy through smooth gating to estimate token-level advantages for LVLM reinforcement learning.The framework is presented as a perception-exploration policy optimization method.
  • 5 Conclusions: Across geometry reasoning, visual puzzles, visual grounding, and few-shot classification, PEPO consistently outperforms GRPO and DAPO while maintaining training stability and scalability.The conclusion attributes the approach’s effectiveness to integrating perception and exploration.
  • 5 Conclusions: The results support integrating perception and exploration as an effective approach to advancing multimodal reasoning in LVLMs.

A Implementation Details

The implementation uses five task settings within a shared RLVR framework, covering reasoning, grounding, classification, puzzles, and scalability. Experiments use common optimization infrastructure while varying data sources and hyperparameters by task.

  • Task settings: The five task settings cover geometry and logic reasoning, visual grounding, few-shot classification, visual puzzle reasoning, and scalability analysis.The corresponding datasets include Geometry3K, RefCOCO, FGVC Aircraft, PuzzleVQA, and ViRL39K.
  • Training setup: All settings use a unified RLVR setup with AdamW, full-parameter bfloat16 fine-tuning, gradient checkpointing, eight sampled responses, and DeepSpeed ZeRO-2.The experiments are implemented within the Swift framework.
  • Prompts and rewards: Prompt formats and programmatically verifiable rewards are designed separately across the five task settings.The reward design governs output formatting and deterministic evaluation.

C Evaluation

The evaluation uses standard splits and deterministic, task-specific metrics across reasoning, classification, puzzle, grounding, and scaling benchmarks. Each setting defines correctness through programmatic rules or exact matching.

  • Evaluation protocol: Standard dataset splits and deterministic programmatic evaluation are used throughout the experiments.
  • Geometry and logic reasoning: Geometry and logic reasoning uses avg@8 accuracy from eight responses per image-question pair, excluding most free-form-answer instances.
  • Few-shot classification: Few-shot classification evaluates exact decoded class-name matches under 1-, 2-, and 4-shot settings.
  • Visual puzzle reasoning: Visual puzzle reasoning reports top-1 accuracy with predefined options, greedy decoding, and malformed outputs counted as incorrect.
  • Visual grounding: Visual grounding reports IoU@50 for greedily predicted bounding boxes, requiring intersection-over-union of at least 0.5 with a ground-truth annotation.
  • Scaling evaluation: Scaling benchmarks use avg@8 evaluation with public PAPO scripts and exclude LLM-as-a-judge instances except Geometry3K.

D Policy Gradient View of PEPO

PEPO replaces uniformly applied sequence-level advantages with token-wise weights that redistribute credit toward visually grounded or high-entropy tokens while preserving the overall policy-gradient scale.

  • D Policy Gradient View of PEPO: PEPO introduces token-wise weights into GRPO-style optimization to refine sequence-level advantages at the token level.The algorithm samples grouped responses and applies the resulting token-level advantage in a GRPO/PPO-style objective.
  • D Policy Gradient View of PEPO: The unit-mean constraint on token weights preserves the sequence-level policy-gradient scale and only redistributes credit among tokens.This keeps the total advantage mass unchanged while changing how credit is allocated within the response.
  • D Policy Gradient View of PEPO: The perception module computes visual similarity for response tokens, while the exploration module computes token-level entropy.These signals are normalized and fused through a perception–exploration module before producing token-level advantages.
  • D Policy Gradient View of PEPO: PEPO therefore preferentially allocates credit to visually grounded or high-entropy tokens without introducing a global scaling change in the policy gradient.The modification affects the distribution of credit inside the token summation rather than the overall gradient magnitude.

E Additional Ablation Results

Additional ablations find that PEPO is robust to gate strength and that cosine similarity is the strongest tested perception-prior measure, while excessive entropy modulation slightly reduces performance.

  • E Additional Ablation Results: PEPO remains robust across a range of gate-strength α values, with best average results at α = 0.05 or α = 0.10 depending on the setting.The perception-only variant at α = 0 already outperforms GRPO, while overly small or large settings can reduce performance slightly.
  • E Additional Ablation Results: A perception-only variant at α = 0 outperforms GRPO, showing that visual similarity alone provides a strong training signal.Adding a moderate entropy gate further improves performance.
  • E Additional Ablation Results: Cosine similarity consistently outperforms L1 and L2 distances across geometry reasoning and few-shot classification benchmarks.Replacing cosine similarity with either distance-based alternative causes noticeable performance degradation.
  • E Additional Ablation Results: Cosine similarity is used as the default perception prior because it is scale-invariant and effective for vision–language alignment.The ablation attributes its stronger performance to a more stable and semantically aligned perception signal.

F Additional Details for the Hidden-state Shift Analysis

The hidden-state shift analysis measures how strongly response tokens depend on image input, bins tokens by visual similarity or entropy, and summarizes lexical patterns for the two token groups.

  • F Additional Details for the Hidden-state Shift Analysis: Representational shift D_t is defined from hidden states with and without the image input for each token and layer.The analysis compares h_with and h_without representations to quantify image-associated token changes.
  • F Additional Details for the Hidden-state Shift Analysis: Tokens are assigned to percentile bins by visual similarity or entropy, and the average D_t is reported within each bin.Both visual similarity and entropy are computed under the image-present condition.
  • F Additional Details for the Hidden-state Shift Analysis: The word-cloud analysis ranks Geometry3K tokens separately by mean entropy or mean visual similarity.It selects the top 100 tokens from each ranking, weights them by aggregated frequency, and removes rare, special, and non-semantic tokens.

G Limitations

The experiments are limited to 2B or 3B LVLM backbones and a curated set of multimodal reasoning and grounding benchmarks, leaving larger models, longer contexts, and additional task types for future study.

  • G Limitations: PEPO is evaluated on recent 2B or 3B LVLMs, not on larger backbones such as 7B or above or on longer-context configurations.The authors attribute this scope boundary to computational and memory limitations.
  • G Limitations: The evaluation is confined to a curated set of multimodal reasoning and grounding benchmarks.Applying PEPO to stronger base models and tasks such as video understanding and tool-augmented reasoning is identified as future research.
Loading 2603.22847v1…