Adaptive Reward Routing: Dynamic Multi-Reward Optimization for Joint Audio-Video Diffusion via Forward-Process RL

AI summary

Overview

Research area: Computer vision and generative modeling, specifically reinforcement-learning-based post-training of joint audio-video diffusion models (text-to-audio-video generation).

Technical level: Advanced. The paper assumes familiarity with diffusion/flow-matching models, cross-attention, GRPO-style advantage estimation, and multi-objective gradient methods.

Scope: The paper proposes Adaptive Reward Routing, a forward-process (DiffusionNFT) reinforcement learning framework that dynamically decides where reward-driven gradient updates should act and how competing rewards should be balanced during training of joint audio-video diffusion models.

What This Paper Is About

Joint audio-video diffusion models must satisfy several goals at once: visual quality, audio quality, text alignment, and tight audio-video synchronization. Reward-guided RL is used to optimize these goals, but two things change as training proceeds: which layers and tokens actually carry cross-modal influence, and how much the five reward signals disagree with each other. Existing methods fix both the layer/token routing (OmniNFT) and the reward weights (GDPO) up front, so their routing and weighting become stale as the model evolves. This paper's goal is to adapt both quantities on the fly while still respecting user-defined reward priorities.

Key Contributions

1. Formulation of joint audio-video diffusion RL as dynamic multi-reward optimization over two coupled dimensions — where modality-conditioned reward updates should act, and how multiple rewards should be coordinated — together with empirical evidence (probing base versus trained checkpoints, reward-conflict analysis, and gradient-based weight drift) that static solutions are insufficient.
2. Cross-Modal Influence-Guided Routing, which uses bidirectional cross-attention responses as a cheap proxy for cross-modal influence: layer aggregation produces token weights, and token aggregation produces layer scales that preserve backward gradient flow through influential cross-modal pathways without extra model interventions.
3. Preference-Preserving Modality-Aware Reweighting, which estimates reward conflicts within the branch responsible for each objective and, after a warm-up period, applies those coefficients as residual corrections to predefined reward weights rather than replacing them, preventing dominant rewards from suppressing weak but essential objectives.
4. Comprehensive comparisons, ablations, and mechanism analyses on two backbones showing consistent gains, complementary benefits of routing and reweighting, and direct-intervention validation that the cross-attention proxy identifies functionally important layers and tokens.

Main Findings

• Best overall on both backbones. On JavisBench, the method achieves the best result on nine of the ten metrics under each of LTX-2 and LTX-2.3. Under LTX-2: VQ 3.336 ± 0.0094, AQ 5.868 ± 0.0158, TV-IB 0.268, TA-IB 0.167, CLIP 0.314, CLAP 0.425, AV-IB 0.235, AVHScore 0.234, JavisScore 0.206, DeSync 0.341. Under LTX-2.3: VQ 3.599, AQ 5.979, TV-IB 0.274, TA-IB 0.177, CLIP 0.314, CLAP 0.460, AV-IB 0.267, AVHScore 0.266, JavisScore 0.236, DeSync 0.302.
• Fixed reward aggregation harms audio and synchronization. Under LTX-2, GDPO raises visual quality (VQ 2.722 versus the base model's 1.883) but worsens DeSync to 0.671, worse than the base model's 0.604. MARBLE reaches DeSync 0.618 and OmniNFT 0.390; OmniNFT* is 0.378.
• Token and layer routing are complementary. Adding token weighting over GDPO gives VQ 3.008, AQ 5.663, DeSync 0.482; adding layer scaling gives VQ 3.192, AQ 5.784, DeSync 0.366; combining both gives VQ 3.315, AQ 5.839, DeSync 0.343.
• Each reward-coordination component adds value. Isolated from MARBLE without the routing stack, branch-aware balancing gives VQ 2.612, AQ 5.290, DeSync 0.492; adding residual preference-preserving mixing gives 2.891 / 5.526 / 0.455; adding warm-up gives 2.999 / 5.791 / 0.369. The complete method (routing plus weighting) reaches 3.336 / 5.868 / 0.341 under LTX-2.
• The proxy matches direct intervention. Across 48 A2V and V2A blocks, the proxy recovers the intervention-based layer ranking with Spearman correlations of 0.98 for A2V and 0.97 for V2A.
• Fixed routing goes stale. The proxy recomputed from the current model stays above 0.96 throughout fine-tuning, whereas a proxy frozen at initialization falls to 0.56 and 0.38.
• High-scoring tokens are functionally important. Blocking the cross-modal responses of the top-scoring 10 percent of target tokens changes the final prediction 1.62 times more for A2V and 1.74 times more for V2A than blocking an equally sized random group.
• More balanced multi-reward progress. Training trajectories show the method reaches the highest average reward while the baselines make less balanced progress across audio, video, and synchronization objectives.
• Qualitative stability. Examples spanning Chinese news delivery, stylized English speech, two-speaker dialogue, and dog barking show less subject and appearance drift than LTX-2 and fewer temporal inconsistencies than OmniNFT.

Methodology in Plain English

The setting is a joint audio-video diffusion model with two streams (audio and video) connected by bidirectional cross-attention. A standard RL post-training loop (DiffusionNFT) generates a group of samples per prompt, scores them with five reward models, converts each reward into a group-relative advantage, and pushes the model toward the samples that beat their peers.

The paper changes two things in that loop.

First, it decides where the update should land. Instead of disabling cross-modal paths (expensive) it reads a number already computed during the forward pass: the magnitude of the pre-gate cross-attention response from audio to video and from video to audio. Averaged over layers and timesteps, these responses become per-token weights that emphasize cross-modally influential tokens, with audio normalized globally and video normalized per frame. Averaged over tokens instead, they become per-layer coefficients that partially detach the source key/value tensors, so strongly influential cross-modal layers keep more backward gradient while weakly coupled ones are increasingly detached. The forward output is unchanged — only the gradient routing changes.

Second, it decides how strongly each reward counts. Each reward is probed only through the branch it supervises (video and audio rewards through their own branches, cross-modal rewards through both), with token routing disabled so the measurement is unbiased. A MARBLE-style simplex solve produces conflict-aware coefficients within each branch. Before warm-up, the predefined prior weights are used as-is. After warm-up, the smoothed coefficient is blended in as a residual correction with a rescaling factor that preserves total prior weight in that branch, so the prior acts as a nonzero floor.

Finally, the weighted rewards form separate video and audio advantages, which become negative-aware token-level losses, weighted by the token scores, combined across branches, and regularized toward a fixed reference policy with a KL term.

Why This Matters

Impact on research. The paper reframes multi-reward post-training of multimodal generative models as a dynamic credit-assignment problem rather than a fixed recipe. It shows empirically that cross-modal function and gradient flow shift during fine-tuning (frozen proxy accuracy dropping to 0.56 and 0.38), which implies that any static routing or weighting scheme will degrade over training. It also demonstrates that pure gradient-geometry weighting can suppress objectives that users consider essential, motivating preference-anchored adaptation.

Real-world applications:

• Automated dubbing and localization, where lip and speech synchronization plus speech quality must both hold.
• Text-to-video advertising or short-form content pipelines that need coherent visuals and matching sound effects or dialogue.
• Virtual avatars and stylized character generation for games or interactive media, where identity and appearance drift across a clip is a visible failure.
• Accessibility tooling such as audio description and synchronized caption generation for video content.

Industry relevance. The work is a collaboration between MMLab@HKUST, Tencent Video, and The University of Hong Kong, and the evaluation is on production-scale backbones (LTX-2 at 19B parameters and LTX-2.3 at 22B) with a benchmark of 10,140 prompts. Post-training methods that improve synchronization and audio quality without retraining from scratch are directly relevant to video platforms, which have to serve multilingual, synchronized audio-video content at scale.

Future Directions

1. A unified audio-video reward. The authors state that no established unified reward jointly captures modality quality, semantic consistency, and temporal synchronization, and that human-preference models provide a complementary overall signal but do not replace fine-grained modality-specific supervision.
2. Extension to other diffusion objectives. The paper adopts DiffusionNFT for direct supervision of sampled flow-matching timesteps without reverse-process likelihood estimation; the authors argue routing can extend to other objectives through their token losses and gradient paths, while reward coordination requires only reward-wise gradients.
3. Architectural generalization. The limitations section states routing applies when modality tokens are identifiable and directional interaction responses can be isolated, and notes validation on both dual-stream LTX backbones and the unified single-stream JavisDiT++ backbone. Models with inseparable modality representations or inaccessible interaction responses are left as future work. (The main results table reports only the two LTX backbones.)
4. Beyond audio-video. The conclusion points toward multimodal learning systems in which feedback is continually reorganized as the model acquires new capabilities, rather than routed by a fixed recipe.

Target Audience

Researchers and engineers working on diffusion-based generative models, reinforcement learning for generative post-training, and multimodal audio-video generation. It is most useful to readers already comfortable with flow matching, cross-attention architectures, and multi-objective optimization who want to understand how to make reward routing and reward weighting adaptive during training. Readers looking for an introductory treatment of diffusion models will find it demanding.

Authors’ abstract

Multi-reward guided reinforcement learning (i.e., RL) offers a promising way to improve joint audio-video diffusion models along several complementary objectives, including modality-specific quality, cross-modal semantic alignment, and temporal synchronization. Its effectiveness, however, depends on two quantities that change during training: where reward-driven updates should act, and how competing rewards should be combined. Existing methods tend to rely on fixed routing and reward weights, failing to track evolving model functions. To address these limitations, we propose Adaptive Reward Routing to jointly adapt update locations and reward coordination during forward-process RL (i.e., DiffusionNFT) of joint audio-video diffusion models. Our method consists of two components. (i) Cross-Modal Influence-Guided Routing (Localizing Updates): We use bidirectional cross-attention responses as an efficient proxy for evolving cross-modal influence, dynamically reweighting token-aware losses and scaling gradients across cross-modal layers without additional model interventions. (ii) Preference-Preserving Modality-Aware Reweighting (Coordinating Rewards): We preserve predefined weights as preference priors and use branch-specific reward-gradient interactions as residual corrections after warm-up. This resolves evolving conflicts without letting dominant rewards suppress weak but essential objectives. Extensive experiments demonstrate consistent improvements in modality quality, semantic consistency, and audio-video synchronization over strong RL baselines. Ablations and mechanism analyses further validate the complementary benefits of adaptive update routing and reward coordination.