Summary (Overview)

• Proposes Adaptive Reward Routing (ARR): A novel framework for multi-reward reinforcement learning (RL) post-training of joint audio-video diffusion models that dynamically adapts both where reward-driven updates should act (update routing) and how competing rewards should be coordinated (reward weighting).
• Cross-Modal Influence-Guided Routing: Uses bidirectional cross-attention response norms (pre-gate A2V/V2A outputs) as a computationally efficient proxy to dynamically reweight token-level losses and scale gradients across cross-modal layers, without requiring additional model interventions.
• Preference-Preserving Modality-Aware Reweighting: Combines user-defined reward weights (as preference priors) with branch-specific, gradient-geometry-based residual corrections (via MARBLE-style optimization) after a warm-up phase, preventing dominant rewards from suppressing weak but essential objectives.
• Empirical validation: Demonstrates consistent improvements across modality quality, semantic consistency, and audio-video synchronization on the JavisBench benchmark using two backbones (LTX-2 19B and LTX-2.3 22B), outperforming strong baselines (GDPO, MARBLE, OmniNFT).
• Mechanism analysis: Validates that the cross-attention response proxy accurately identifies functionally important layers (Spearman correlations of 0.98 for A2V, 0.97 for V2A) and tokens (top-scoring tokens cause 1.62×–1.74× larger prediction changes when blocked), while showing that static routing maps become stale during training.

Introduction and Theoretical Foundation

Background

Joint audio-video diffusion models (e.g., LTX-2) generate synchronized visual and audio content from text prompts. High-quality generation requires simultaneously satisfying:

• Modality-specific visual and audio quality
• Cross-modal semantic alignment
• Temporal synchronization

These objectives are difficult to capture with a single supervised objective, motivating reward-guided diffusion reinforcement learning.

Key Challenges

The paper identifies two coupled dimensions of dynamic multi-reward optimization:

1. Dynamic Reward Routing (Where to optimize): Joint audio-video models contain modality-specific branches coupled through cross-attention. Static routing (as in OmniNFT) becomes stale because cross-modal functions and gradient flows evolve during fine-tuning. The authors demonstrate this via:

   • Forward KV Ablation (Fig. 1a): Shifting synchronization-critical layers between base and trained models
   • Backward Gradient Analysis (Fig. 1b): Evolving cross-modal gradient paths

2. Dynamic Reward Coordination (How to balance): Rewards frequently disagree on the same sample (Fig. 1c). Existing approaches fail:

   • GDPO: Normalizes rewards but uses fixed weights, leaving conflicts unadapted
   • MARBLE: Adjusts weights via gradient geometry, but coefficients reflect gradient compatibility rather than user preferences (Fig. 1d shows vanishing AV-DeSync weight)

Theoretical Foundation

The method builds on DiffusionNFT (Zheng et al., 2025), a forward-process RL framework. The core idea is constructing implicit positive and negative policies:

$$v^+_\theta = (1-\beta)v^{\text{updated}} + \beta v_\theta, \quad v^-_\theta = (1+\beta)v^{\text{updated}} - \beta v_\theta$$

The reward of sample $n$ is converted to a group-relative advantage:

$$A(n) = \frac{R(n) - \mu_R}{\sigma_R + \varepsilon}, \quad r(n) = \frac{1}{2} + \frac{1}{2}\text{clip}\left(\frac{A(n)}{A_{\max}}, -1, 1\right)$$

The resulting objective is:

$$\mathcal{L}_{\text{NFT}} = \mathbb{E}_{n,t}\left[r(n)\|v^+_\theta(x^{(n)}_t, c, t) - u^{(n)}\|^2_2 + (1-r(n))\|v^-_\theta(x^{(n)}_t, c, t) - u^{(n)}\|^2_2\right]$$

Methodology

Framework Overview

The complete optimization flow is:

$$\{A_k\} \xrightarrow{\text{reward reweighting}} \{\omega_{m,k}A_k\} \xrightarrow{\text{branch routing}} (A_v, A_a) \xrightarrow{\text{token routing}} \mathcal{L} \xrightarrow{\text{layer routing}} \nabla_\theta \mathcal{L}$$

1. Cross-Modal Influence-Guided Routing

Proxy Design: Instead of direct interventions (disabling A2V/V2A and comparing outputs), the method uses the pre-gate response of cross-attention paths already produced during the forward pass:

$$d^{l,t}_{v,i} = \|o^{l,t}_{a \to v,i}\|_2, \quad d^{l,t}_{a,i} = \|o^{l,t}_{v \to a,i}\|_2$$

These are collected over an intermediate-to-late denoising window and detached before policy optimization.

Token-Level Routing: For each target token, responses are averaged over selected timesteps and cross-modal blocks, then normalized:

$$\lambda_{m,i} = 1 + (\lambda_{\max} - 1)\text{Norm}_{99}\left(\frac{1}{|\mathcal{B}||\mathcal{T}|}\sum_{l \in \mathcal{B}}\sum_{t \in \mathcal{T}} d^{l,t}_{m,i}\right)$$

Audio responses are normalized globally; video responses are normalized within each frame.

Layer-Level Routing: Layer scores are converted to soft detachment coefficients:

$$\alpha^l_m = (1 - \tilde{\delta}^l_m)^{1/\tau}$$

For a source key/value tensor $X \in \{K, V\}$:

$$\tilde{X}^{l,t}_{\bar{m} \to m} = \alpha^l_m \text{sg}(X^{l,t}_{\bar{m}}) + (1 - \alpha^l_m)X^{l,t}_{\bar{m}}$$

This leaves the forward value unchanged but scales the backward gradient by $1 - \alpha^l_m$.

2. Preference-Preserving Modality-Aware Reweighting

Each reward is probed only through the branch it supervises (video/audio rewards use their respective branches; cross-modal rewards use both). MARBLE produces conflict-aware coefficients $\gamma_{m,k}$ within each branch (with token routing disabled during probes). After warm-up:

$$\omega_{m,k} = \begin{cases} \omega^{\text{prior}}_{m,k}, & e < e_{\text{warm}} \\ (1-\kappa)\omega^{\text{prior}}_{m,k} + \kappa C_m \bar{\gamma}_{m,k}, & e \geq e_{\text{warm}} \end{cases}$$

where $C_m$ rescales simplex coefficients to preserve total prior weight, and $\bar{\gamma}_m \leftarrow \rho\bar{\gamma}_m + (1-\rho)\gamma^*_m$ smooths estimates.

3. Training Objective

Branch advantages combine reward-specific advantages:

$$A^{(n)}_m = \sum_{k \in \mathcal{K}_m \cup \mathcal{K}_c} \omega_{m,k} A^{(n)}_k, \quad m \in \{v, a\}$$

The negative-aware token loss:

$$\ell^{(n)}_{m,i} = \frac{r^{(n)}_m \|v^+_{\theta,m,i} - u_{m,i}\|^2_2}{w^{+,(n)}_m + \varepsilon} + \frac{(1-r^{(n)}_m)\|v^-_{\theta,m,i} - u_{m,i}\|^2_2}{w^{-,(n)}_m + \varepsilon}$$

The modality loss with token routing:

$$\mathcal{L}^{\text{policy}}_m = \mathbb{E}_n\left[\frac{\sum_{i \in \mathcal{I}_m} \lambda^{(n)}_{m,i}\ell^{(n)}_{m,i}}{\sum_{i \in \mathcal{I}_m} \lambda^{(n)}_{m,i}}\right]$$

Final objective with KL regularization:

$$\mathcal{L}(\theta) = \sum_{m \in \mathcal{M}} \mathcal{L}^{\text{policy}}_m + \lambda_{\text{KL}} \sum_{m \in \mathcal{M}} \mathcal{L}_{\text{KL},m}(\theta)$$

Empirical Validation / Results

Experimental Setup

• Backbones: LTX-2 (19B) and LTX-2.3 (22B) joint audio-video diffusion models
• Training Data: 19,487 audio-video prompts from VGGSound-derived corpus
• Rewards: VideoAlign, HPSv3 (video quality); AudioBox Aesthetics (audio quality); CLAP (text-audio alignment); DeSync→AV-DeSync (synchronization)
• Evaluation: JavisBench (10,140 prompts), four metric groups: AV-Quality, Text-Consistency, AV-Consistency, AV-Synchrony

Main Results (Table 1)

The method achieves the best result on nine of ten metrics for each backbone.

Ablation Studies (Table 2, LTX-2 backbone)

• Routing components (over GDPO): Token weighting alone → 3.008 VQ, 0.482 DeSync; Layer scaling alone → 3.192 VQ, 0.366 DeSync; Combined → 3.315 VQ, 0.343 DeSync
• Weighting components (over MARBLE): Branch-aware → 2.612 VQ; +Residual → 2.891 VQ; +Warm-Up → 2.999 VQ, 0.369 DeSync
• Complete method (Ours): 3.336 VQ, 0.341 DeSync — best overall

Proxy Validation (Section 5.4)

• Layer-level fidelity: Spearman correlations of 0.98 (A2V) and 0.97 (V2A) between proxy scores and intervention-based layer rankings
• Dynamic tracking: Current-model proxy stays above 0.96 throughout training; proxy frozen at initialization drops to 0.56 and 0.38
• Token importance: Blocking top-scoring 10% of tokens changes final prediction 1.62× (A2V) and 1.74× (V2A) more than blocking random groups

Theoretical and Practical Implications

Theoretical Contributions

1. Formalization of dynamic credit assignment: The paper reframes multi-reward post-training for joint audio-video generation as a dynamic credit-assignment problem where both update localization and reward coordination must co-adapt with the evolving model.
2. Efficient proxy for functional influence: Demonstrates that pre-gate cross-attention response norms serve as a valid, computationally free proxy for direct intervention-based measures of cross-modal influence, with near-perfect ranking agreement.
3. Preference-anchored adaptation: Shows that pure gradient-geometry-based reward weighting (MARBLE) can suppress weak objectives (e.g., AV-DeSync weight vanishing), while combining gradient residuals with preference priors preserves essential objectives.

Practical Implications

• Stable multi-reward optimization: The method prevents dominant rewards from suppressing weak objectives while still adapting to changing conflicts, leading to more balanced improvements across all quality dimensions.
• Architecture-agnostic components: Reward coordination is architecture-independent; routing applies whenever modality tokens are identifiable and directional interaction responses can be isolated.
• Training efficiency: The proxy requires no additional forward passes or interventions, making the approach computationally practical for large-scale models (19B–22B parameters).

Conclusion

Main Takeaways

• Adaptive Reward Routing jointly adapts update locations and reward coordination during forward-process RL, addressing two coupled failure modes of static approaches.
• The cross-attention response proxy accurately identifies functionally important layers and tokens, and dynamic recomputation is necessary because fixed routing becomes stale during training.
• Preference-preserving reward reweighting resolves evolving conflicts without letting dominant rewards suppress weak but essential objectives.
• Consistent gains across backbones (LTX-2, LTX-2.3), metrics, and ablations demonstrate the importance of adapting the optimization process in step with the model itself.

Future Directions

1. Unified reward models: No established reward jointly captures modality quality, semantic consistency, and temporal synchronization; human-preference models provide complementary signals but don't replace fine-grained supervision.
2. RL framework extension: The routing can extend to other diffusion objectives through their token losses and gradient paths; reward coordination only requires reward-wise gradients.
3. Architectural scope: Models with inseparable modality representations or inaccessible interaction responses remain unexplored; routing requires identifiable modality tokens and isolable directional interaction responses.