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

> Adaptive Reward Routing dynamically adjusts both update locations and reward weights during RL post-training, improving audio-video diffusion quality across all metrics on JavisBench.

- **Source:** [arXiv](https://arxiv.org/abs/2609.37200)
- **Published:** 2026-10-03
- **Permalink:** https://picx.dev/p/9mr0D3
- **Whiteboard:** https://picx.dev/p/9mr0D3/image

## Summary

## 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)

| Backbone | Method | VQ ↑ | AQ ↑ | TV-IB ↑ | TA-IB ↑ | CLIP ↑ | CLAP ↑ | AV-IB ↑ | AVHScore ↑ | JavisScore ↑ | DeSync ↓ |
|----------|--------|------|------|---------|---------|--------|--------|---------|------------|-------------|----------|
| **LTX-2** | Base | 1.883 | 5.201 | 0.265 | 0.143 | 0.312 | 0.358 | 0.180 | 0.177 | 0.153 | 0.604 |
| | +GDPO | 2.722 | 5.450 | 0.261 | 0.138 | 0.312 | 0.347 | 0.174 | 0.175 | 0.155 | 0.671 |
| | +MARBLE | 2.384 | 5.100 | 0.265 | 0.138 | 0.311 | 0.365 | 0.182 | 0.182 | 0.158 | 0.618 |
| | +OmniNFT | 3.136 | 5.614 | 0.265 | 0.145 | 0.312 | 0.416 | 0.222 | 0.219 | 0.195 | 0.390 |
| | **+Ours** | **3.336** | **5.868** | **0.268** | **0.167** | **0.314** | **0.425** | **0.235** | **0.234** | **0.206** | **0.341** |
| **LTX-2.3** | Base | 2.032 | 5.218 | 0.271 | 0.151 | 0.308 | 0.387 | 0.205 | 0.202 | 0.175 | 0.504 |
| | +GDPO | 2.929 | 5.251 | 0.272 | 0.144 | 0.309 | 0.376 | 0.218 | 0.199 | 0.178 | 0.560 |
| | +MARBLE | 2.582 | 5.176 | 0.271 | 0.147 | 0.309 | 0.394 | 0.217 | 0.207 | 0.182 | 0.496 |
| | +OmniNFT | 3.489 | 5.693 | 0.271 | 0.163 | 0.316 | 0.449 | 0.250 | 0.238 | 0.224 | 0.369 |
| | **+Ours** | **3.599** | **5.979** | **0.274** | **0.177** | **0.314** | **0.460** | **0.267** | **0.266** | **0.236** | **0.302** |

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.

---

_Markdown view of https://picx.dev/p/9mr0D3, served by PicX — AI-generated visual whiteboard summaries of research papers._