Perception-Aware Policy Optimization: マルチモーダル推論を知覚層で修正

大規模マルチモーダルモデル(LMM)は画像の説明には優れていますが、視覚的理解が必要な推論タスクでは失敗します。エラーの67%は推論能力ではなく知覚の失敗に起因しています。従来の報酬ベースRLは視覚と推論を結合されたものとして扱い、モデルが画像を誤解したのか推論の誤りを犯したのかを区別することなく、答えが間違っているときに一律のペナルティを適用します。PAPOはこれを解決するため、モデルの出力がどの程度視覚入力に依存しているかを直接測定し(画像パッチの60%をマスキングして)、この発散を最適化信号として使用し、モデルに幻覚を見るのではなく実際の視覚コンテンツに推論を根拠付けるよう強制します。

マルチモーダルベンチマーク上で視覚言語モデルをファインチューニングする場合(視覚的数学問題、グラフ、図、空間推論)、知覚層がボトルネックになります。標準的な教師あり学習によるファインチューニングは幻覚を修正できません。報酬ベースRLはグローバルなポリシーを学習しますが、どのエラーが不十分な視覚的根拠付けに起因しているかを特定しません。知覚対応ポリシー最適化は、視覚が破損されると出力が変わらない出力に直接ペナルティを与え、モデルが見ているものを理解していることを証明するよう強制します。

核心概念

PAPOは暗黙的知覚損失を通じて知覚品質を測定します。これは、完全な画像と破損した画像(60%のパッチマスキング)が与えられたときのモデルの応答に対するポリシー分布を比較するKL発散です。視覚がマスキングされたときに出力がほとんど変わらない場合、モデルは視覚情報を使用していません。KL発散は低いままで損失信号は強いです。出力が大きく変わる場合、モデルは視覚に推論を根拠付けます。発散は高く損失は弱いです。崩壊を防ぐ二重エントロピー正則化と、ポリシー勾配ステップの非対称クリッピングと組み合わせることで、PAPOはモデルが視覚情報を活用しながら安定性を維持するよう奨励します。このテクニックはポリシー最適化(GRPO/DAPO)に直接統合され、別の報酬コンポーネントではありません。

アーキテクチャ概要

• パッチマスキングモジュール: 画像パッチの60%をランダムにマスキングして破損した視覚入力を作成
• 暗黙的知覚損失: 完全な画像と マスキングされた画像に対するポリシー出力間のKL発散を計算
• 二重エントロピー正則化: 元のポリシーとマスキングされたポリシーの両方のエントロピーを制約して、退化した解を防止
• 非対称クリッピング(Clip-Higher): 対称的なPPOの境界を超えて探索を奨励するポリシー勾配の修正
• 知覚拡張報酬: 従来の答えの正確性報酬と知覚発散信号を組み合わせる
• 視覚言語バックボーン: 構造的変更なしの標準的なLMMアーキテクチャ(例: LLaVA、Qwen-VL)

実装

このサンプルは、KL発散を使用してモデル出力が視覚入力にどの程度依存しているかを測定する、核心的な知覚損失計算を示しています。

# 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

このサンプルは、探索を奨励するために上限を削除する非対称クリッピングポリシー勾配を示しています。

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

このサンプルは、知覚損失をポリシー勾配学習に統合する完全なPAPO最適化ループを示しています。

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()
        }

実践的なガイダンス

ハイパーパラメータ	推奨値	目的
パッチマスキング比率	0.6(60%)	知覚低下と情報保持のバランス
知覚損失の重み	0.5	知覚と答え報酬のバランス
エントロピー最小閾値	0.1	退化したゼロエントロピーポリシーを防止
エントロピー最大閾値	2.0	ランダムな高エントロピーポリシーを防止
非対称クリッピングイプシロン	0.2	標準的なPPO範囲と上側非対称性
二重エントロピーの重み	0.01	正則化の強度(小さい値)
勾配クリッピング最大ノルム	1.0	学習の不安定性を防止
視覚言語バックボーン	LLaVA、Qwen-VL	固定または軽度のチューニング

使用する場合: 知覚エラーが優勢なベンチマーク上で視覚言語モデルをファインチューニングする場合、PAPOを適用してください(視覚QA、グラフ理解、空間推論、マルチモーダル数学)。答えの正確性を測定でき、視覚的根拠付けを改善したい場合に使用します。モデルが画像に相談することなく答えを幻覚する(捏造する)タスクに特に有効です。

使用しない場合: 検索ベースのタスク(例:「このオブジェクトの名前を付けてください」)は知覚が単純明快なため、PAPOを避けてください。グラウンドトゥルースの答えラベルがない場合はスキップしてください。このメソッドは教師あり報酬が必要です。計算予算が非常に限られている場合は使用しないでください。マスキングと二重の順向きパスにより2倍のオーバーヘッドが追加されます。視覚入力が小さい場合(256x256ピクセル未満)は60%のパッチマスキングがすべての情報を破壊するためスキップしてください。

よくある落とし穴: マスキング比率が高すぎる(70%を超える)と視覚コンテンツが完全に破壊され、KL発散は無意味になります。低すぎる(40%未満)とモデルはコンテキストだけから正しく推測できます。二重エントロピー正則化を適用しないとKL発散ハッキングが発生します。モデルは根拠付けを改善する代わりにエントロピーを操作します。対称的なPPOクリッピングを使用すると非対称バリアントの探索メリットが無効になります。知覚損失が負であることを忘れると(KLを最小化するのではなく最大化)最適化が反転します。ファインチューニング前に視覚エンコーダーをフリーズし忘れると視覚特性の破滅的忘却が引き起こされます。

参考文献

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