MAPO

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

MAPO mitigates late-stage modality collapse in audio-reasoning LLMs using cross-modal differential entropy to identify audio-critical tokens and apply targeted policy gradients. It sustains cross-modal grounding throughout reasoning without domain-specific biases, relying on native statistical signals.

llm
multimodal
reasoning
rlhf
audio-driven

Demos

These demos illustrate how MAPO mitigates late-stage modality collapse in audio reasoning by sustaining cross-modal grounding during extended chain-of-thought generation. Evaluate whether the attention branch maintains elevated audio attention and entropy signals deep into reasoning traces, and how the dual mechanisms synergistically allocate optimization budget to audio-critical tokens while enforcing sustained grounding. Look for evidence that full MAPO prevents hallucinations compared to baselines that collapse partway through reasoning.

MAPO dual-branch architecture: modality relevance mask and attention loss branch working synergistically to prevent modality collapse

Ablation study showing sustained audio attention and differential entropy across reasoning trajectory; full MAPO maintains grounding vs. collapsed baselines

Authors: Cihan Xiao, Yiwen Shao, Chenxing Li, Xiang He, Zhenwen Liang, Steve Yves, Sanjeev Khudanpur, Liefeng Bo

Categories: cs.CL

Published 2026-05-26 · Updated 2026-05-26

Abstract

Audio and omni-modal large language models exhibit impressive cross-modal reasoning capabilities. However, applying standard reinforcement learning post-training algorithms to these models exposes a critical structural vulnerability: methods like GRPO apply uniform policy gradients across all tokens, ignoring their unequal dependence on the non-text source modality. This exacerbates late-stage modality collapse during extended chain-of-thought generation, where models progressively abandon the primary source signal in favor of compressed textual priors, leading to confident but ungrounded hallucinations. To address this, we introduce Modality-Aware Policy Optimization (MAPO), a novel dual-branch reinforcement learning framework. First, MAPO dynamically concentrates the policy gradient on modality-critical tokens using a modality relevance mask, which is derived from the cross-modal differential entropy between an audio-ablated reference and the multimodal policy. Second, it integrates an auxiliary attention loss branch that applies a targeted, temporally scaled penalty to the model's internal attention distributions. This ensures the model actively sustains cross-modal grounding deep into the reasoning trace. Evaluations on complex audio reasoning benchmarks demonstrate that MAPO substantially improves long-horizon reasoning fidelity and multimodal instruction following, achieving highly competitive performance and setting new state-of-the-art results on several key benchmarks among open-weight models. By relying strictly on native statistical signals rather than domain-specific inductive biases, MAPO offers a promising foundation for mitigating epistemic collapse across diverse multimodal systems.

Introduction

Audio and omni-modal large language models (LLMs) have demonstrated remarkable cross-modal reasoning capabilities, integrating non-text modalities such as speech, music, and environmental sounds into complex chain-of-thought (CoT) reasoning pipelines. However, a critical structural vulnerability emerges when standard reinforcement learning (RL) post-training algorithms—such as GRPO, PPO, or DPO—are applied to these models: all generated tokens are treated as equally important for the policy gradient update. This assumption is fundamentally flawed in the cross-modal setting.

In cross-modal generation, tokens fall into two qualitatively different categories. The vast majority serve as pure linguistic scaffolding that the model's internal language prior can accurately predict from text context alone. A smaller but critical subset of tokens requires genuine cross-modal perception—these tokens depend directly on the non-text source signal (e.g., raw acoustic features) to capture essential attributes such as sound events, speaker characteristics, or musical properties. Applying a uniform policy gradient squanders the optimization budget on modality-agnostic tokens while severely diluting the learning signal for tokens that require active audio grounding.

Beyond inefficient gradient allocation, this uniform weighting drives a severe and previously under-characterized failure mode the authors call late-stage modality collapse. During extended CoT generation, models exhibit a pronounced temporal decay in cross-modal attention: as the reasoning trace lengthens, the model progressively abandons the uncompressed primary audio source in favor of its own highly compressed textual summaries of what it heard in earlier steps. Once a model begins relying on these intermediate text-based representations rather than the raw audio, its reasoning drifts into confident but ungrounded hallucinations. Crucially, as the language prior takes hold, the model becomes pathologically confident—standard entropy-based signals vanish precisely when external grounding is most needed, masking the true epistemic state and making naive statistical interventions ineffective.

Overview of the MAPO framework showing late-stage modality collapse and the dual-branch architecture — Overview of the MAPO framework. **(a)** Late-stage modality collapse, where attention shifts from the source audio to the text prior during CoT reasoning. **(b)** MAPO mitigates this via a dual-branch architecture. A *modality relevance mask* uses cross-modal differential entropy ($\Delta h_t$) to focus the policy gradient on audio-critical tokens. Simultaneously, an *attention loss branch* applies a temporally scaled penalty to substantive tokens with low audio attention, enforcing sustained cross-modal grounding.

To address these intertwined challenges, the paper introduces Modality-Aware Policy Optimization (MAPO), a novel RL framework built atop GRPO that enforces sustained cross-modal grounding through two complementary mechanisms:

A Modality Relevance Mask that dynamically reweights the policy gradient by computing cross-modal differential entropy between a text-only reference and the full multimodal policy, concentrating gradient updates on tokens where the audio signal genuinely shifts the predictive distribution.
An auxiliary Attention Loss Branch that directly penalizes neglect of the source modality at linguistically substantive tokens during late-stage generation, using inverse relevance weighting and temporally increasing penalties to enforce persistent cross-modal grounding.

The paper is submitted for consideration to NeurIPS 2026 and demonstrates state-of-the-art performance among open-weight models on several complex audio reasoning benchmarks spanning speech, music, and environmental sound events.

Related Work

Omni-modal Large Language Models

Modern omni-modal LLMs process text, vision, and audio in a unified, native manner. The base model used in MAPO experiments is Qwen3-Omni-Thinking, which is optimized specifically for extended cross-modal CoT generation (as opposed to Qwen3-Omni-Instruct, which targets instruction following). The broader ecosystem includes proprietary systems such as the Gemini series, and open-weight models including Audio Flamingo 3, Step-Audio-2, MiMo-Audio, and Covo-Audio. All of these systems, despite their architectural innovations, are trained via standard text-based post-training objectives that uniformly weight all generated tokens, leaving them structurally susceptible to language-prior dominance and late-stage modality collapse.

Group Relative Policy Optimization (GRPO)

GRPO is an efficient RL alignment method that estimates advantages over a sampled group of completions rather than using a separate value model. For a prompt $x$ and $G$ completions $\{y^{(g)}\}_{g=1}^{G}$, the normalized advantage is:

$$\hat{A}^{(g)} = \frac{r(x, y^{(g)}) - \operatorname{mean}(\{r(x, y^{(g')})\})}{\operatorname{std}(\{r(x, y^{(g')})\}) + \varepsilon}$$

The policy gradient loss uses a clipped surrogate objective. With importance ratio $\rho_t^{(g)} = \frac{\pi_\theta(y_t^{(g)} \mid y_{ $$\mathcal{L}_{\text{PG}}^{\text{GRPO}} = -\frac{1}{T_g}\sum_{t=1}^{T_g} \min\!\Big(\rho_t^{(g)} \hat{A}^{(g)},\; \text{clip}(\rho_t^{(g)}, 1{-}\epsilon, 1{+}\epsilon)\,\hat{A}^{(g)}\Big)$$

The KL regularization term uses the unbiased $k_3$ estimator:

$$\mathcal{L}_{\text{KL}} = \frac{1}{T_g}\sum_{t=1}^{T_g}\Big[\exp\!\big(\log\pi_{\text{ref}}(y_t^{(g)}) - \log\pi_\theta(y_t^{(g)})\big) - \big(\log\pi_{\text{ref}}(y_t^{(g)}) - \log\pi_\theta(y_t^{(g)})\big) - 1\Big]$$

The total GRPO objective is $\mathcal{L}_{\text{GRPO}} = \mathcal{L}_{\text{PG}}^{\text{GRPO}} + \beta \mathcal{L}_{\text{KL}}$. The key structural problem MAPO addresses is the inherent $1/T_g$ uniform averaging—all tokens receive an equal share of optimization budget regardless of their dependence on the audio source.

Cross-Modal Alignment and Grounding

Prior works inject modality-aware feedback via external rewards from large teacher models (e.g., AudioThinker, AudioDeepThinker), which provide auxiliary sequence-level signals. While effective, these extrinsic methods still apply feedback uniformly at the sequence level, ignoring token-by-token variance in modality dependence. MAPO is orthogonal and complementary to these approaches.

At a finer granularity, PAPO and concurrent vision-language works have recognized the limitations of uniform token weighting. Spotlight isolates perceptually pivotal visual tokens to focus policy updates; VGPO combats temporal visual forgetting via progressive attention compensation. These methods, however, rely on domain-specific visual similarity metrics. An entropy-based token selection approach has also been proposed, but as the paper argues, standard predictive entropy fails during late-stage modality collapse—the model becomes overconfident in its hallucinations, causing the entropy signal to vanish. MAPO addresses this directly by anchoring its metric against a text-only reference.

Method: Modality-Aware Policy Optimization (MAPO)

Problem Formulation

The framework considers a cross-modal generation setting where an omni-modal LLM with $L$ transformer layers and $H$ attention heads per layer acts as policy $\pi_\theta$. The model receives a multimodal prompt $x$ comprising a text instruction sequence $x_{\text{text}}$ and non-text source modality $x_{\text{audio}}$ (continuous acoustic features mapped to the model's input space). The policy autoregressively generates a text completion $y = (y_1, \dots, y_T)$ where each token $y_t \in \mathcal{V}$, with likelihood $\prod_{t=1}^T \pi_\theta(y_t \mid y_{

MAPO challenges the uniform $1/T$ token weighting of standard RL objectives. The goal is to design a dynamic, token-level weighting mechanism that distinguishes tokens driven purely by the language prior from those requiring active cross-modal perception.

Audio Attention Mass and Late-Stage Modality Collapse

MAPO introduces the audio attention mass as a dynamic quantitative measure of how much the model grounds its reasoning in the source audio at any given generation step. This metric captures the aggregate probability mass within the model's self-attention distribution that is routed to the cross-modal input tokens.

Let $\alpha_{t,j}^{(l,h)}$ denote the attention weight that query position $t$ assigns to key position $j$ in layer $l$ and head $h$. Let $\mathcal{S}_{\text{audio}}$ be the index set of audio input tokens. The raw audio attention mass for a specific head is:

$$m_t^{(l,h)} = \sum_{j \in \mathcal{S}_{\text{audio}}} \alpha_{t,j}^{(l,h)}$$

Both max-head and mean-head reductions are defined:

$$m_t^{(l,\text{max})} = \max_{h}\, m_t^{(l,h)}, \qquad m_t^{(l,\text{mean})} = \frac{1}{H}\sum_{h=1}^{H} m_t^{(l,h)}$$

The final audio attention mass $a_t$ for token position $t$ averages the max-head signal across a targeted set of deep transformer layers $L_{\text{tgt}}$:

$$a_t = \frac{1}{|L_{\text{tgt}}|}\sum_{l \in L_{\text{tgt}}} m_t^{(l,\text{max})}$$

Using the max-head reduction captures the strongest cross-modal grounding signal. The paper demonstrates empirically that in the early stages of CoT generation $a_t$ is relatively high, but it exhibits a pronounced temporal decay as generation lengthens. In the late stages, $a_t$ collapses near zero—the model has effectively stopped attending to the raw audio and is instead auto-regressively conditioned on its own textual context.

Late-stage modality collapse and key token extraction. Plots display the temporal decay of audio attention mass (**solid:** max-head reduction, *dashed:* mean-head reduction) alongside uncertainty metrics. **(a)** Text-only entropy fails to track this loss of source grounding. **(b)** Cross-modal differential entropy effectively isolates true "key tokens"; this signal strongly correlates with the attention decay.

Cross-Modal Differential Entropy

To counteract modality collapse, MAPO must identify which tokens are genuine "key tokens" requiring audio perception. Standard predictive entropy fails here: during late-stage collapse, the model is confidently wrong (low entropy but incorrect), so entropy provides no useful signal.

MAPO introduces cross-modal differential entropy ($\Delta h_t$) to expose hidden modality dependence. This is computed as the token-level entropy difference between a frozen text-only reference model and the full multimodal policy:

$$\Delta h_t = H\!\left(\pi_{\text{ref}}^{\text{text}}(y_t \mid y_{where $H(\cdot)$ denotes the predictive entropy over the vocabulary. A large positive $\Delta h_t$ indicates that conditioning on the audio signal dramatically sharpens the model's prediction—i.e., the token genuinely depends on acoustic information. A near-zero or negative value indicates the token is predictable from text context alone. This cross-modal differential entropy circumvents the collapse problem: even when $\pi_\theta$ is falsely confident (low entropy), if $\pi_{\text{ref}}^{\text{text}}$ is equally uncertain, $\Delta h_t$ remains informative about the token's true audio dependence.

Part-of-Speech Gating

MAPO also incorporates part-of-speech (POS) gating to further focus the attention loss on linguistically substantive tokens. By analyzing cross-modal differential entropy stratified by POS category, the paper shows that content words (nouns, verbs, adjectives, adverbs) exhibit systematically higher audio dependence than function words and punctuation. A POS gate $g_t \in \{0, 1\}$ selects substantive tokens, preventing gradient updates and attention penalties from being diluted by tokens that are trivially predictable regardless of the audio content.

POS entropy summary across cross-modal reasoning trajectories. The top panel displays average text-only predictive entropy, isolating the uncertainty of the language prior. The bottom panel shows average cross-modal differential entropy, indicating the degree of acoustic dependence. Content words show consistently higher audio dependence than function words.

Modality Relevance Mask

The first branch of MAPO is the Modality Relevance Mask. Using $\Delta h_t$ as a per-token score, MAPO constructs a soft weighting mask $w_t$ that dynamically concentrates the GRPO policy gradient on audio-critical tokens. Tokens where $\Delta h_t$ is large receive amplified gradient weight; tokens that are textually predictable receive downweighted gradient contributions.

Formally, for the $g$-th completion, the masked policy gradient replaces the uniform $1/T_g$ average with:

$$\mathcal{L}_{\text{PG}}^{\text{MAPO}} = -\frac{\sum_{t=1}^{T_g} w_t \cdot \min\!\Big(\rho_t^{(g)} \hat{A}^{(g)},\; \text{clip}(\rho_t^{(g)}, 1{-}\epsilon, 1{+}\epsilon)\,\hat{A}^{(g)}\Big)}{\sum_{t=1}^{T_g} w_t}$$

where $w_t$ is a normalized function of $\Delta h_t^{(g)}$, derived from the per-token differential entropy at position $t$ within that specific completion. This ensures that audio-critical tokens contribute disproportionately to the learning signal while modality-agnostic scaffolding tokens are effectively down-weighted. The normalization maintains numerical stability and prevents the loss from collapsing when modality dependence is globally low.

Attention Loss Branch

The second branch is the Attention Loss Branch, which directly combats the late-stage collapse mechanism by penalizing insufficient audio attention at substantive generation steps. This branch is guided by the complement of the relevance mask (i.e., an inverse weighting) combined with temporally increasing weights that grow as the generation trace lengthens—precisely targeting the late stages where collapse is most pronounced.

The attention loss penalizes low audio attention mass $a_t$ at positions that are (1) linguistically substantive (pass the POS gate $g_t = 1$), (2) at positions where $\Delta h_t$ suggests audio dependence, and (3) in completions where the model is failing (gated by a soft task-failure signal to avoid penalizing already-correct completions where collapse is less harmful). The temporal scaling ensures the penalty grows with token index $t$, making late-stage grounding violations more costly than early-stage ones.

Concretely, the attention loss for a single completion is:

$$\mathcal{L}_{\text{Attn}} = \frac{1}{Z}\sum_{t=1}^{T_g} \phi_t \cdot g_t \cdot \mathcal{F} \cdot \max(0,\, \tau - a_t)$$

where $\phi_t$ is the temporally scaled weight (increasing with $t$), $g_t$ is the POS gate, $\mathcal{F}$ is a soft task-failure gate derived from the completion's reward signal, $\tau$ is a target threshold for audio attention mass, and $Z$ is a normalization constant. This formulation ensures that the penalty is applied only where it is both meaningful (substantive tokens, audio-dependent positions) and necessary (incomplete or failed reasoning).

Combined MAPO Objective

The final MAPO training objective combines the masked policy gradient, the KL regularization term inherited from GRPO, and the attention loss branch:

$$\mathcal{L}_{\text{MAPO}} = \mathcal{L}_{\text{PG}}^{\text{MAPO}} + \beta \mathcal{L}_{\text{KL}} + \lambda \mathcal{L}_{\text{Attn}}$$

where $\beta$ controls KL regularization strength and $\lambda$ controls the attention loss weight. The dual-branch design is key: the modality relevance mask addresses the gradient allocation problem (making RL learning more efficient), while the attention loss branch directly combats the mechanistic cause of collapse (forcing sustained audio attention). Together, they form a mutually reinforcing system.

Analysis of Modality Collapse

The paper provides a detailed analysis of internal model dynamics before and after MAPO training, comparing three conditions: the unmodified Baseline (standard GRPO), an ablated variant without the attention loss (MAPO-No-Attn-Loss), and the full MAPO with both branches.

Impact of the attention loss branch on internal attention dynamics. Plots display cross-modal differential entropy $\Delta h_t$ (blue bars) against max-head (solid orange) and mean-head (dashed orange) audio attention mass. **Bottom & Middle:** The baseline and partially ablated (No-Attn-Loss) models exhibit severe late-stage modality collapse; audio attention decays precipitously midway through generation, accompanied by a weakened cross-modal differential entropy signal. **Top:** The full MAPO framework mitigates this temporal decay, maintaining both sustained audio attention and stronger distributional deviations from the text-only reference during late-stage reasoning.

In the baseline model, audio attention mass $a_t$ starts high at the beginning of reasoning but decays sharply as the CoT trace extends. The cross-modal differential entropy $\Delta h_t$ likewise weakens in the late stages, confirming that the model has effectively internalized the audio content into compressed text representations and no longer needs to re-consult the raw signal. The ablated MAPO-No-Attn-Loss model shows some improvement in early-stage grounding (due to the modality relevance mask improving gradient efficiency) but still collapses late in generation. The full MAPO model maintains significantly higher and more stable audio attention throughout the reasoning trace, with the temporal attention penalty directly counteracting the collapse mechanism.

Qualitative Analysis: Token-Level Grounding

The paper provides rich qualitative analyses through token-level visualizations of audio attention mass across three model conditions on example audio reasoning questions. Two example tasks are analyzed:

Example 1: Mechanical Sound Identification

Given an audio clip, the task is to identify the source of mechanical sounds (choices: Factory machinery, Wind turbine, Car engine, Airplane). The ground truth is "Factory machinery".

The Baseline model predicts "Wind turbine" (incorrect). Its token stream shows high audio attention early but severe decay in the late reasoning phase, where the model rationalizes "whooshing" and "thwop" sounds consistent with wind turbines—features that appear to be confabulated from language priors rather than grounded in the actual audio signal.
MAPO-No-Attn-Loss predicts "Factory machinery" (correct). The modality relevance mask improves gradient allocation, but some mid-stream attention decay is still visible.
The full MAPO+Attn-Loss model predicts "Factory machinery" (correct) and exhibits consistently green (high) audio attention mass throughout its entire reasoning trace, including during the option evaluation and final conclusion phases. The model correctly identifies the layered mechanical sounds, whirring/whining, and industrial-scale characteristics as consistent with factory machinery rather than wind turbines.

Example 2: Transportation Mode Identification

Given an audio clip, the task is to identify the transportation mode (choices: Automobile, Train, Aeroplane, Horse-drawn wagon). The ground truth is "Horse-drawn wagon".

Both the Baseline and MAPO-No-Attn-Loss models incorrectly predict "Train," reasoning that the rhythmic metallic clattering and "clack-clack" pattern is characteristic of steel wheels on rail joints. Despite relatively high audio attention in some segments, both models succumb to the language prior that strongly associates rhythmic metallic sounds with trains.
The full MAPO+Attn-Loss model correctly identifies "Horse-drawn wagon." It maintains sustained audio attention and correctly discerns that the sound has an organic "thudding" quality characteristic of hooves, combined with wheel creaking and harness sounds, ruling out the purely metallic signature of train wheels on steel tracks. This demonstrates that sustained cross-modal grounding enables finer perceptual discrimination that collapses without the attention enforcement mechanism.

Training Dynamics

Training accuracy curves — Training accuracy across training steps for MAPO compared to baselines.

Completion length during training — Completion length evolution during training, showing how MAPO affects the length of generated reasoning traces.

Mean audio attention mass over training steps for different attention loss weights — Mean audio attention mass over 800 training steps for different attention loss weights ($\lambda$). Higher values of $\lambda$ structurally enforce a significantly higher baseline of cross-modal grounding throughout training.

Training dynamics reveal that increasing the attention loss weight $\lambda$ monotonically increases the average audio attention mass maintained during generation. This demonstrates that the attention loss branch has a direct and structural effect on cross-modal grounding, not merely an indirect one via improved task performance. The paper reports that there exists an optimal range of $\lambda$: too small and the effect is negligible; too large and the model may over-attend to audio even for truly language-predictable tokens, degrading overall performance.

Instruct model training accuracy comparison — Training accuracy comparison between GRPO and MAPO when applied to the instruction-following model variant.

Instruct model completion length during training — Completion length evolution during training for the instruction-following model variant under GRPO and MAPO.

Experiments

Model and Base Architecture

MAPO is evaluated on top of Qwen3-Omni-Thinking, an omni-modal LLM specifically optimized for extended cross-modal chain-of-thought generation. The authors also test MAPO applied to Qwen3-Omni-Instruct, the instruction-following variant, to assess generalizability across model training regimes.

Benchmarks

Evaluations are conducted on complex audio reasoning benchmarks covering three audio domains:

Speech: Tasks requiring understanding of spoken language content, speaker attributes, and paralinguistic features.
Music: Tasks involving musical attribute recognition, genre classification, and compositional reasoning.
Sound Events: Tasks involving identification and reasoning about environmental and mechanical sound sources.

The paper reports results on several key benchmarks where MAPO achieves state-of-the-art performance among open-weight models, substantially improving long-horizon reasoning fidelity and multimodal instruction following.

Baselines and Comparisons

MAPO is compared against:

The base Qwen3-Omni-Thinking model (unmodified)
Standard GRPO fine-tuning without modality-aware modifications
MAPO-No-Attn-Loss (ablation: modality relevance mask only, no attention loss branch)
Other published open-weight audio reasoning models including Audio Flamingo 3, Step-Audio-2, MiMo-Audio, and Covo-Audio
Proprietary systems such as the Gemini series

Ablation Studies

The paper conducts systematic ablations to isolate the contribution of each MAPO component:

Modality Relevance Mask Alone (MAPO-No-Attn-Loss)

Applying only the modality relevance mask without the attention loss branch improves task accuracy over standard GRPO in many conditions. This confirms that reweighting the gradient toward audio-critical tokens is beneficial in itself. However, as shown in the qualitative analysis (e.g., the horse-drawn wagon example), the model still suffers from late-stage collapse because the mask acts only indirectly on attention dynamics through gradient optimization pressure, without a direct mechanistic constraint.

Attention Loss Branch Alone

The attention loss branch in isolation shows that directly penalizing audio attention neglect is a strong intervention for maintaining grounding. However, without the modality relevance mask, gradient efficiency is suboptimal and some performance gains are left on the table.

Full MAPO (Both Branches)

The combined framework consistently outperforms both ablations, confirming that the two branches are complementary: the mask addresses gradient allocation efficiency while the attention loss provides a direct mechanistic constraint against collapse.

Attention Loss Weight ($\lambda$)

As shown in the audio attention mass ablation figure, the choice of $\lambda$ has a pronounced effect on cross-modal grounding. The paper identifies an optimal range that balances the structural enforcement of audio attention against the flexibility needed for the model to reason efficiently.

Temporal Scaling

Experiments confirm the importance of the temporally increasing penalty in the attention loss branch. Without temporal scaling (uniform penalty across all positions), the model receives equal pressure at early and late positions, which is less effective than concentrating the penalty where collapse actually occurs.

POS Gating

Part-of-speech gating is shown to improve both efficiency and performance by avoiding wasted penalty mass on function words and punctuation that are inherently text-predictable regardless of the audio content.

Key Results

MAPO achieves highly competitive performance and sets new state-of-the-art results on several audio reasoning benchmarks among open-weight models. Key empirical findings include:

Long-horizon reasoning fidelity: MAPO shows the largest gains over GRPO on examples requiring long CoT traces, exactly where late-stage collapse is most severe. Short-answer tasks show smaller but consistent improvements.
Multimodal instruction following: MAPO improves the model's ability to follow audio-grounded instructions, reducing instances where the model ignores the audio content and answers based purely on question text.
Robustness to language-prior hallucinations: On tasks where language priors are misleadingly strong (e.g., rhythmic metallic sounds that statistically correlate with trains but are actually horse-drawn wagons), MAPO's sustained audio attention enables finer acoustic discrimination that the baseline and GRPO models fail to achieve.
Generality: The method works on both the Thinking and Instruct variants of the base model, suggesting broad applicability across model training regimes.

Design Principles and Novelty

MAPO's design is guided by three key principles that distinguish it from prior work:

Native statistical signals, no domain-specific inductive biases: Unlike methods that rely on visual similarity metrics or domain-specific audio features, MAPO derives its key signal (cross-modal differential entropy) from the model's own predictive distributions, making it applicable in principle to any multimodal system regardless of the specific non-text modality.
Mechanistic targeting of collapse: Rather than relying solely on reward shaping or external teacher signals, MAPO directly intervenes in the model's internal attention mechanism through the attention loss branch, addressing the proximate cause of late-stage collapse rather than its downstream effects.
Dual-branch synergy: The two branches address distinct failure modes (gradient allocation inefficiency and attention decay) that are related but mechanistically distinct, and their combination produces effects that neither branch achieves alone.

The paper also argues that MAPO is fully orthogonal to and compatible with sequence-level reward engineering approaches. Practitioners could combine MAPO with, for example, external teacher-model reward signals to benefit from both fine-grained token-level optimization and high-quality sequence-level supervision.

Limitations and Future Directions

The paper acknowledges several limitations and open questions:

Computational overhead of the text-only reference: Computing cross-modal differential entropy requires a forward pass through a text-only reference model at each training step, which adds computational cost. Efficient approximations or cached reference distributions could reduce this overhead.
Hyperparameter sensitivity: The attention loss weight $\lambda$, temporal scaling schedule, and attention threshold $\tau$ introduce additional hyperparameters that require tuning. The paper provides guidance but does not derive principled automatic selection methods.
Scope of evaluation: While MAPO is motivated as a general framework for multimodal systems, empirical validation is focused on audio reasoning. Extension to vision-language models or video-language models would further validate the generality claim.
Task-failure gate calibration: The soft task-failure gate that conditions the attention penalty on completion quality requires careful calibration to avoid suppressing grounding enforcement on near-correct completions.
Interaction with very long CoT traces: While MAPO demonstrably improves late-stage grounding, extremely long reasoning traces (e.g., thousands of tokens) may present challenges that the current temporal scaling schedule is not optimally designed for.

Future work directions suggested by the paper include extending MAPO to vision-language and video-language settings, developing automatic methods for threshold and weight selection, investigating whether the cross-modal differential entropy signal can be computed more efficiently via sampling or distillation, and combining MAPO with model architecture modifications (e.g., cross-modal memory mechanisms) that could structurally preserve audio access across long reasoning traces.

Conclusion

MAPO introduces a principled, dual-branch reinforcement learning framework that addresses two intertwined problems in cross-modal RL post-training: the inefficiency of uniform token weighting and the mechanistic failure mode of late-stage modality collapse. By deriving a modality relevance mask from cross-modal differential entropy and coupling it with a temporally scaled attention loss branch, MAPO dynamically concentrates the policy gradient on audio-critical tokens while directly enforcing sustained cross-modal grounding throughout extended reasoning traces. Evaluated on complex audio reasoning benchmarks spanning speech, music, and environmental sound events, MAPO achieves state-of-the-art performance among open-weight models and demonstrates substantially improved long-horizon reasoning fidelity. The framework's reliance on native statistical signals rather than domain-specific heuristics positions it as a broadly applicable foundation for mitigating epistemic collapse across diverse multimodal systems.

Code & Implementation

The repository implements MAPO as an extension to the Megatron-Core distributed training framework, integrated within the MS-Swift RL framework. The codebase is organized around three core components:

Core Training Logic

MAPO Trainer (libs/ms-swift/swift/megatron/trainers/mapo_trainer.py): Extends MegatronGRPOTrainer to implement the dual-branch architecture. The trainer manages the modality relevance mask computation (derived from cross-modal differential entropy between an audio-ablated reference model and the policy) and orchestrates the auxiliary attention loss branch across distributed training nodes.
Attention Collector (libs/ms-swift/swift/megatron/trainers/mapo_attention_collector.py): Dynamically extracts cross-modal attention distributions from the transformer's multi-head attention mechanism. Detects audio tokens, constructs audio token masks, and reconstructs attention probabilities from query/key tensors for Qwen-style thinker architectures.
POS Gating Utilities (libs/ms-swift/swift/megatron/trainers/mapo_pos_utils.py): Implements optional part-of-speech (POS) tagging to further refine token-level importance gating, allowing fine-grained control over which tokens receive the auxiliary attention penalty.

Evaluation & Reward Functions

Audio QA Rewards (src/rewards/audio_qa_rewards.py): Defines reward extraction and scoring for audio reasoning tasks, including answer parsing (tag-based, post-think, plain-text variants), counting-aware MSE scoring, and MCQ handling.
Benchmark Scripts (src/mmau/, src/mmar/, src/mmsu/, src/mmau-pro/): Inference and evaluation pipelines for the four audio-reasoning benchmarks (Multi-modal Audio Understanding, Deep Audio Reasoning, Spoken Language Understanding, Instruction Following).

Training & Inference Orchestration

Asynchronous Rollout Server (src/mapo/start_rollout.sh): Launches a vLLM-based inference server on dedicated GPU nodes, decoupled from training nodes to maximize throughput and GPU utilization during distributed RL training.
Consistency Checker Server (src/mapo/start_checker.sh): Optional auxiliary service for multi-turn consistency verification, supporting more complex reward signals.
Training Launcher (src/mapo/start_train.sh): Entry point that configures Megatron-Core distributed training, manages multi-node synchronization, and integrates the MAPO trainer with the reference model rollout.

Paper–Repo Mapping

Paper Component	Repository Implementation
Modality Relevance Mask (ω̃)	`mapo_trainer.py`: entropy-based mask computation from reference vs. policy logits
Attention Loss Branch (ℒ_attn)	`mapo_trainer.py`: temporal scaling penalty on audio-attention heads
Cross-modal Attention Extraction	`mapo_attention_collector.py`: audio token detection and multi-head reconstruction
Distributed RL Training	Megatron-Core integration via `MegatronGRPOTrainer` base class
Rollout (Inference)	vLLM server decoupled from training; `start_rollout.sh` orchestration
Reward Functions & Eval	`src/rewards/` and per-benchmark `scripts/evaluation.py`

Quick Start

Training requires Python 3.11, CUDA 12.8, PyTorch 2.9.1, Megatron-Core 0.15.4, and vLLM 0.14.0. The three-step training procedure (start vLLM rollout server, optional consistency checker, then launch distributed training nodes) is coordinated via shell scripts that configure environment variables and GPU topology.