Technical Summary: Escape the Language Prior: Mitigating Late-Stage Modality Collapse in Audio Reasoning via Modality-Aware Policy Optimization

1. Problem Statement

Audio and omni-modal large language models (LLMs) have demonstrated impressive cross-modal reasoning capabilities. However, applying standard reinforcement learning (RL) post-training algorithms (e.g., PPO, DPO, GRPO) to these models exposes a critical structural vulnerability: late-stage modality collapse.

Standard RL methods apply uniform policy gradients across all generated tokens, assuming each token contributes equally to the learning signal. In cross-modal generation, this assumption is flawed. Many tokens serve as linguistic scaffolding predictable by the model's internal language prior, while a critical subset relies entirely on the non-text source modality (e.g., raw audio) to capture essential acoustic attributes. Uniform weighting squanders the optimization budget on modality-agnostic text and dilutes the learning signal for tokens requiring active cross-modal grounding.

This inefficiency drives a severe failure mode during extended Chain-of-Thought (CoT) generation: as the reasoning trace lengthens, models progressively abandon the uncompressed primary source signal in favor of compressed textual priors. This leads to "confident but ungrounded hallucinations," where the model becomes epistemically disconnected from the source audio. Crucially, as the model succumbs to its language prior, it becomes pathologically confident, causing standard entropy-based interventions to fail in recovering lost grounding.

2. Methodology: Modality-Aware Policy Optimization (MAPO)

To address these challenges, the authors propose Modality-Aware Policy Optimization (MAPO), a novel dual-branch reinforcement learning framework built upon Group Relative Policy Optimization (GRPO). MAPO modifies the learning objective through two complementary mechanisms designed to break the language prior and enforce sustained cross-modal grounding.

2.1 Core Metrics

Before defining the optimization objective, MAPO introduces two quantitative metrics:

• Audio Attention Mass ($a_t$): A measure of the aggregate probability mass routed to the cross-modal input sequence within the model's self-attention distribution. It serves as a proxy for active perceptual bandwidth. Empirical analysis shows this mass decays significantly in late-stage generation, indicating the model stops "listening."
• Cross-Modal Differential Entropy ($\Delta h_t$): The token-level entropy difference between a frozen text-only reference model ($\pi_{text-ref}$) and the live multimodal policy ($\pi_\theta$).
  $$\Delta h_t = H(\pi_{text-ref}(\cdot | y_{<t}, x_{text})) - H(\pi_\theta(\cdot | y_{<t}, x))$$
  This metric isolates the information gain provided by the non-text modality. High $\Delta h_t$ indicates tokens where the audio source significantly alters the predictive distribution, identifying "key tokens" that require cross-modal perception.

2.2 Dual-Branch Optimization Objective

The total MAPO objective is defined as:
$$\mathcal{L}_{MAPO} = \mathcal{L}_{PG}(\tilde{\omega}) + \beta \mathcal{L}_{KL} + \eta \mathcal{L}_{attn}$$

A. Modality Relevance Mask and Reweighted Policy Gradient ($\mathcal{L}_{PG}$)

Standard GRPO averages token losses uniformly ($1/T$). MAPO replaces this with a dynamic, token-level Modality Relevance Mask ($\tilde{\omega}$) derived from the cross-modal differential entropy.

• Mechanism: The mask reweights the policy gradient to disproportionately allocate the optimization budget toward tokens where the source audio strongly influences the predictive distribution.
• Implementation: The mask is constructed by applying a softmax over the absolute differential entropy $|\Delta h_t|$, scaled by a length-dependent temperature and capped to prevent extreme amplification. The final loss normalizes by the sum of weights to maintain a stable learning rate.
• Effect: This dynamically upweights the gradient contribution of modality-critical tokens while downweighting linguistic scaffolding.

B. Auxiliary Attention Loss Branch ($\mathcal{L}_{attn}$)

While the mask corrects the gradient budget, it relies on the model retaining some latent awareness of the signal. To directly combat late-stage collapse, MAPO introduces an auxiliary loss branch that penalizes the neglect of the source signal deep within the reasoning chain.

• Mechanism: This branch applies a targeted, temporally scaled penalty to the model's internal attention distributions. It penalizes tokens with low audio attention mass ($a_t$) specifically at linguistically substantive positions (isolated via a Part-of-Speech gate).
• Temporal Scaling: The penalty is weighted by an inverse relevance mask (targeting tokens with suspiciously low differential entropy) and a non-decreasing temporal weight $(t/T)^\kappa$, concentrating the penalty at the tail end of the sequence where collapse typically occurs.
• Gradient Flow: Gradients from this branch flow strictly through the audio attention mass directly into the transformer's multi-head attention parameters (Query and Key projections), leaving the language modeling head unperturbed. This forces the model to actively sustain cross-modal grounding deep into the reasoning trace.

3. Key Contributions

1. Formalization of Late-Stage Modality Collapse: The authors identify and formalize the limitations of uniform token weighting in cross-modal RL, highlighting how the language prior masks true model uncertainty during long-horizon reasoning, leading to confident hallucinations.
2. MAPO Framework: They propose a dual-branch optimization framework that leverages cross-modal differential entropy to prioritize modality-critical tokens and introduces a targeted attention penalty to enforce sustained cross-modal grounding.
3. Empirical Validation: Through extensive evaluation on complex audio reasoning tasks (speech, music, sound events), they demonstrate that MAPO improves both the fidelity of long-form reasoning and robustness against language-prior hallucinations, setting new state-of-the-art results among open-weights models.

4. Experimental Results

The authors evaluated MAPO using the Qwen3-Omni-30B-A3B-Thinking model as the base policy, trained on a diverse mixture of audio QA datasets (including AVQA, AudioSet, LibriTTS-R, and MusicCaps).

4.1 Ablation Studies

• Component Impact: Standard GRPO on the base model slightly degraded performance (Avg: 70.21 $\to$ 69.51). Adding the modality relevance mask recovered performance (70.53). Integrating the auxiliary attention loss yielded a substantial leap (71.51), validating that structurally enforcing sustained source grounding mitigates collapse.
• Full Fine-Tuning: MAPO maintained a clear advantage over standard GRPO even under full parameter updates, achieving an overall average of 73.34 on the final Phase 2 benchmark.

4.2 Benchmark Performance

• Modality-Centric Benchmarks (MMAU, MMAR): MAPO achieved a new state-of-the-art for open-weights models. On MMAU, it scored 77.80, matching the leading proprietary model (Nova 2 Omni) and substantially outperforming Gemini 2.5 Pro (71.60) and GPT-4o Audio (62.50). Compared to its Qwen3-Omni-Thinking baseline, MAPO improved MMAU by +2.80 and MMAR by +4.00.
• Capability and Instruction Following (MMSU, MMAU-Pro): MAPO achieved the top overall score on MMSU (79.36), with simultaneous improvements in Perception and Reasoning. Notably, on MMAU-Pro, it elevated the Audio Instruction Following (AIF) score from 62.07 to 95.40, surpassing highly optimized models like Gemini 2.5 Flash (95.10).

4.3 Qualitative Analysis

Case studies revealed that without the attention loss branch, models often identified initial acoustic cues correctly but then drifted into incorrect textual priors (e.g., misidentifying a "Horse-drawn wagon" as a "Train" due to rhythmic clattering). The full MAPO framework maintained audio attention deep into the reasoning trace, allowing the model to refine ambiguous cues and reject misleading priors.

5. Significance and Claims

The paper claims that MAPO offers a promising foundation for mitigating epistemic collapse across diverse multimodal systems. Its significance lies in three areas:

1. Structural Correction: It addresses a fundamental structural inefficiency in standard RL post-training (uniform token weighting) that causes models to abandon source modalities during extended reasoning.
2. Native Signal Reliance: Unlike prior methods that rely on external teacher models or domain-specific inductive biases (e.g., visual similarity metrics), MAPO relies strictly on native statistical signals (attention weights and predictive entropy) derived from the model itself.
3. Generalizability: While validated on audio, the authors posit that the framework's token-reweighting and attention-steering mechanisms are modality-agnostic, suggesting applicability to other continuous non-text modalities like vision and video.

The authors explicitly note that MAPO is an alignment and regularization framework, not a capability-injection mechanism. It ensures models fully leverage existing perceptual and reasoning capacities but does not synthesize capabilities absent from the base pre-training. Furthermore, they observe that applying MAPO to instruction-tuned models lacking intrinsic CoT capabilities yields limited success, as the underlying prior of brevity constrains the optimization potential.