Perception Aware Policy Optimization

Perception-Aware Policy Optimization: Fixing Multimodal Reasoning at the Perception Layer

Large multimodal models (LMMs) excel at describing images but fail at reasoning tasks requiring visual understanding—67% of errors come from perception failures, not reasoning capability. Traditional reward-based RL treats vision and reasoning as coupled, applying uniform penalties when answers are wrong without distinguishing whether the model misunderstood the image or made an inference mistake. PAPO solves this by directly measuring how much the model's outputs depend on visual input (by masking 60% of image patches) and using this divergence as an optimization signal, forcing the model to ground reasoning in actual visual content rather than hallucinating.

When fine-tuning vision-language models on multimodal benchmarks—visual math problems, charts, diagrams, spatial reasoning—the perception layer becomes the bottleneck. Standard supervised fine-tuning cannot fix hallucination; reward-based RL trains global policies without identifying which errors stem from insufficient visual grounding. Perception-aware policy optimization directly penalizes outputs that don't change when vision is corrupted, forcing the model to prove it understands what it sees.

Core Concept

PAPO measures perception quality through implicit perception loss—a KL divergence comparing the model's policy distribution over responses when given full images versus corrupted images (60% patch masking). If outputs barely change when vision is masked, the model isn't using visual information; the KL divergence remains low and the loss signal is strong. If outputs change significantly, the model grounds reasoning in vision; the divergence is high and the loss is weak. Combined with double entropy regularization to prevent collapse and asymmetric clipping in the policy gradient step, PAPO encourages models to exploit visual information while maintaining stability. The technique integrates directly into policy optimization (GRPO/DAPO) rather than being a separate reward component.

Architecture Overview

• Patch Masking Module: Randomly masks 60% of image patches to create corrupted visual inputs
• Implicit Perception Loss: Computes KL divergence between policy outputs on full vs. masked images
• Double Entropy Regularization: Constrains entropy of both original and corrupted policies to prevent degenerate solutions
• Asymmetric Clipping (Clip-Higher): Policy gradient modification encouraging exploration above symmetric PPO bounds
• Perception-Augmented Reward: Combines traditional answer correctness rewards with perception divergence signals
• Vision-Language Backbone: Standard LMM architecture (e.g., LLaVA, Qwen-VL) without structural changes

Implementation

This example demonstrates the core perception loss computation using KL divergence to measure how much model outputs depend on visual input.

# Perception-aware policy optimization for multimodal models
import torch
import torch.nn.functional as F
from torch.distributions import Categorical

class PerceptionAwarePolicyOptimizer:
    def __init__(self, model, learning_rate=1e-5, perception_weight=0.5):
        self.model = model
        self.learning_rate = learning_rate
        self.perception_weight = perception_weight
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

    def mask_image_patches(self, images, mask_ratio=0.6):
        """Randomly mask image patches to create perception-degraded inputs.
        images: [batch, channels, height, width]"""

        batch, c, h, w = images.shape

        # Patch size (assume 16x16 patches for ViT-style encoders)
        patch_size = 16
        num_h_patches = h // patch_size
        num_w_patches = w // patch_size
        total_patches = num_h_patches * num_w_patches

        # Random mask per sample
        masked_images = images.clone()

        for b in range(batch):
            # Sample which patches to mask
            num_mask = int(total_patches * mask_ratio)
            mask_indices = torch.randperm(total_patches)[:num_mask]

            # Apply masking
            for patch_idx in mask_indices:
                h_idx = patch_idx // num_w_patches
                w_idx = patch_idx % num_w_patches

                h_start = h_idx * patch_size
                h_end = h_start + patch_size
                w_start = w_idx * patch_size
                w_end = w_start + patch_size

                masked_images[b, :, h_start:h_end, w_start:w_end] = 0  # Zero out patch

        return masked_images

    def compute_implicit_perception_loss(self, images, text_prompts):
        """Compute KL divergence between policy distributions on full vs masked images.
        Lower divergence = model not using vision = higher loss signal."""

        # Forward pass with original images
        with torch.enable_grad():
            outputs_full = self.model(images, text_prompts)
            logits_full = outputs_full['logits']  # [batch, vocab_size]
            probs_full = F.softmax(logits_full, dim=-1)

        # Forward pass with masked images
        masked_images = self.mask_image_patches(images, mask_ratio=0.6)
        with torch.no_grad():
            outputs_masked = self.model(masked_images, text_prompts)
            logits_masked = outputs_masked['logits']
            probs_masked = F.softmax(logits_masked, dim=-1)

        # KL divergence: KL(p_full || p_masked)
        # High KL means outputs change significantly when vision is corrupted (good)
        # Low KL means outputs barely change (bad perception grounding)
        kl_divergence = F.kl_div(
            torch.log(probs_masked + 1e-10),
            probs_full,
            reduction='batchmean'
        )

        # Loss is negative KL (maximize divergence, minimize negative divergence)
        perception_loss = -kl_divergence

        return perception_loss, kl_divergence.item()

    def double_entropy_regularization(self, images, text_prompts):
        """Prevent policy collapse by constraining entropy of both distributions."""

        outputs_full = self.model(images, text_prompts)
        logits_full = outputs_full['logits']
        probs_full = F.softmax(logits_full, dim=-1)
        entropy_full = -torch.sum(probs_full * torch.log(probs_full + 1e-10), dim=-1)

        masked_images = self.mask_image_patches(images, mask_ratio=0.6)
        outputs_masked = self.model(masked_images, text_prompts)
        logits_masked = outputs_masked['logits']
        probs_masked = F.softmax(logits_masked, dim=-1)
        entropy_masked = -torch.sum(probs_masked * torch.log(probs_masked + 1e-10), dim=-1)

        # Both distributions should maintain reasonable entropy
        # Too low entropy = model gives up; too high = random
        min_entropy = 0.1
        max_entropy = 2.0

        entropy_loss_full = torch.relu(min_entropy - entropy_full).mean() + \
                           torch.relu(entropy_full - max_entropy).mean()
        entropy_loss_masked = torch.relu(min_entropy - entropy_masked).mean() + \
                             torch.relu(entropy_masked - max_entropy).mean()

        return entropy_loss_full + entropy_loss_masked

This example shows the asymmetric clipping policy gradient that encourages exploration by removing the upper bound on advantage.

def asymmetric_clipping_ppo_loss(advantages, probs_new, probs_old, eps=0.2):
    """Clip-Higher: asymmetric PPO clipping that encourages exploration.
    Only clip downward to prevent excessive conservatism."""

    # Probability ratio
    ratio = probs_new / (probs_old + 1e-10)

    # Standard PPO clipping (both sides)
    surr1 = ratio * advantages

    # Asymmetric clipping: only clip upward side
    surr2 = torch.clamp(ratio, min=1-eps, max=1+eps) * advantages
    # But allow higher values than 1+eps if they improve advantage
    surr2_upper = torch.max(torch.tensor(1+eps), ratio) * advantages

    # Take minimum for lower side, but allow higher values
    ppo_loss = torch.min(surr1, surr2_upper).mean()

    return -ppo_loss  # Policy gradient: minimize negative advantage

This example demonstrates the complete PAPO optimization loop integrating perception loss into policy gradient training.

class PAPOTrainer:
    def __init__(self, model, learning_rate=1e-5):
        self.optimizer = PerceptionAwarePolicyOptimizer(model, learning_rate=learning_rate)

    def training_step(self, images, text_prompts, answers, ground_truth_answers):
        """Complete PAPO training step: perception + entropy + policy gradient."""

        # Reward 1: Answer correctness
        correct = torch.tensor(
            [ans == gt for ans, gt in zip(answers, ground_truth_answers)]
        ).float()
        reward_answer = correct

        # Perception loss: encourage visual grounding
        perception_loss, kl_div = self.optimizer.compute_implicit_perception_loss(
            images, text_prompts
        )

        # Entropy regularization: prevent collapse
        entropy_loss = self.optimizer.double_entropy_regularization(images, text_prompts)

        # Policy gradient with asymmetric clipping
        outputs = self.optimizer.model(images, text_prompts)
        logits = outputs['logits']
        probs_new = F.softmax(logits, dim=-1)

        # Compute advantages
        advantages = reward_answer - reward_answer.mean()
        advantages = advantages / (advantages.std() + 1e-10)

        # Asymmetric PPO loss
        ppo_loss = asymmetric_clipping_ppo_loss(advantages, probs_new, probs_new)

        # Combined loss
        total_loss = (
            ppo_loss +
            self.optimizer.perception_weight * perception_loss +
            0.01 * entropy_loss
        )

        self.optimizer.optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.optimizer.model.parameters(), max_norm=1.0)
        self.optimizer.optimizer.step()

        return {
            'answer_reward': reward_answer.mean().item(),
            'perception_loss': perception_loss.item(),
            'kl_divergence': kl_div,
            'entropy_loss': entropy_loss.item(),
            'ppo_loss': ppo_loss.item(),
            'total_loss': total_loss.item()
        }

Practical Guidance

When to use: Apply PAPO when fine-tuning vision-language models on benchmarks where perception errors dominate (visual QA, chart understanding, spatial reasoning, multimodal math). Use when you can measure answer correctness and want to improve visual grounding. Particularly effective for tasks where models hallucinate (confabulate) answers without consulting the image.

When NOT to use: Avoid PAPO for retrieval-based tasks (e.g., "name this object") where perception is straightforward. Skip if you have no ground truth answer labels—the method requires supervised rewards. Don't use if computational budget is severely limited, as masking and dual forward passes add 2× overhead. Skip if your visual inputs are small (< 256x256 pixels) where 60% patch masking destroys all information.

Common pitfalls: Masking ratio too high (>70%) destroys visual content entirely, making KL divergence meaningless. Too low (<40%) allows models to guess correctly from context alone. Not applying double entropy regularization causes KL divergence hacking—models manipulate entropy rather than improving grounding. Using symmetric PPO clipping negates the exploration benefits of asymmetric variants. Forgetting that perception loss is negative (maximize KL, not minimize) inverts the optimization. Not freezing the vision encoder before fine-tuning causes catastrophic forgetting of visual features.

Reference

Team PAPO. (2025). Perception-Aware Policy Optimization for Multimodal Reasoning. arXiv preprint arXiv:2507.06448. https://arxiv.org/abs/2507.06448