X$^3$-OPD: Distilling Reasoning into Large Audio-Language Models via On-Policy Alignment

Technical Breakdown of X$^3$-OPD§

Core Methodology§

X$^3$-OPD addresses the reasoning deficit in large audio-language models (LALMs) by distilling reasoning from a text-based LLM teacher ($T$) into an audio-language student ($S$). The key idea is on-policy distillation: the student generates reasoning trajectories conditioned on its own acoustic perception, while the teacher provides token-level supervision using matched textual inputs.

Three-Tier Symmetric Corpus§

To cover diverse reasoning scenarios, the authors construct a corpus with three tiers: 1. Textual Reasoning Rendered into Speech: TTS-generated audio from text reasoning data (e.g., GSM8K, MATH). 2. Audio-Event Reasoning: Complex acoustic scenes with event-based questions (e.g., "How many footsteps?" after a sound clip). 3. Spoken-Dialogue Reasoning: Conversational audio with paralinguistic cues (e.g., emotion, prosody) requiring inference.

Each example consists of audio $x$, a textual prompt $p$, and a ground-truth answer $a$. For the teacher, the corresponding text input $x_{text}$ is available (e.g., transcript for TTS, event captions for tier 2, dialogue transcripts for tier 3).

Distillation Objective§

The training loss combines two terms: 1. On-Policy Distillation Loss: At each decoding step $t$, the student's logits $\mathbf{z}_t^{(S)}$ are matched to the teacher's logits $\mathbf{z}_t^{(T)}$ using KL divergence: $$\mathcal{L}_{\text{distill}} = \sum_{t} D_{\text{KL}}(\pi_T(\cdot \mid x_{text}, p, y_{<t}) \, \| \, \pi_S(\cdot \mid x, p, y_{<t}))$$ where $\pi$ denotes output probability distributions.

2. Reinforcement Learning on Verified Answers: To encourage correct final answers, a reinforcement learning term (PPO variant) is added: $$\mathcal{L}_{\text{RL}} = -\mathbb{E}_{y \sim \pi_S}[r(y, a)]$$ where $r(y,a)=1$ if the answer matches $a$, else $0$. The overall loss is $\mathcal{L} = \mathcal{L}_{\text{distill}} + \lambda \mathcal{L}_{\text{RL}}$.

Implementation Details§

• Teacher: A text-only LLM (e.g., Mistral-7B) finetuned on reasoning data.
• Student: An LALM (e.g., SALMONN or Qwen-Audio) with an audio encoder and LLM backbone.
• On-Policy Sampling: The student generates reasoning chains $y$ by sampling from its current policy $\pi_S$. These chains are used to query the teacher for token-level guidance.
• Training: The student is updated via gradient descent with the mixed loss. To preserve existing capabilities, a small replay buffer of non-reasoning tasks is used.

Code Snippet (PyTorch-like Pseudocode)§

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoProcessor

# Initialize teacher and student
teacher = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
student = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-Audio")

# Freeze teacher
for param in teacher.parameters():
    param.requires_grad = False

optimizer = torch.optim.AdamW(student.parameters(), lr=1e-5)

for batch in dataloader:
    audio, text_prompt, transcript, answer = batch
    
    # Student generates reasoning trajectory (on-policy)
    with torch.no_grad():
        student_outputs = student.generate(
            input_ids=tokenize_prompt(text_prompt),
            audio_inputs=audio,
            max_new_tokens=512,
            output_scores=True,
            return_dict_in_generate=True
        )
    student_logits = torch.stack(student_outputs.scores, dim=1)  # [B, T, V]
    student_tokens = student_outputs.sequences[:, prompt_len:]
    
    # Teacher logits on transcript
    teacher_inputs = tokenizer(transcript, return_tensors="pt").to(device)
    with torch.no_grad():
        teacher_outputs = teacher(**teacher_inputs, output_hidden_states=False)
        teacher_logits = teacher_outputs.logits[:, -student_logits.size(1)-1:-1, :]  # align length
    
    # KL divergence loss
    loss_kl = F.kl_div(
        F.log_softmax(student_logits, dim=-1),
        F.softmax(teacher_logits, dim=-1),
        reduction="batchmean"
    )
    
    # RL loss (simplified: binary reward for answer correctness)
    generated_answer = extract_answer(student_outputs.sequences)
    reward = (generated_answer == answer).float()
    log_probs = F.log_softmax(student_logits, dim=-1).gather(-1, student_tokens.unsqueeze(-1)).squeeze(-1)
    loss_rl = -(log_probs.mean(dim=1) * reward).mean()
    
    loss = loss_kl + 0.5 * loss_rl
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Results§

Experiments on MMSU, MMAU, BIG-Bench Audio, and MMAR show significant gains over baselines (e.g., +10% accuracy on MMAU) while maintaining performance on non-reasoning tasks. The on-policy nature and tiered corpus are critical for cross-modal transfer.