🎧 HeadRouter: Dynamic Head-Weight Routing for Task-Adaptive Audio Token Pruning in Large Audio Language Models

Peize He^1,2, Yaodi Luo^1,2, Xiaoqian Liu^1,3, Xuyang Liu^1,4, Jiahang Deng¹,
Yaosong Du², Li Bangyu², Xiyan Gui^1,5, Yuxuan Chen¹, Linfeng Zhang^1✉

¹ EPIC Lab, Shanghai Jiao Tong University · ² DAIL Tech · ³ Northeastern University
⁴ Sichuan University · ⁵ Huazhong University of Science and Technology

⚡ Training-free, task-adaptive audio token pruning through dynamic attention-head routing.

🔥 News

• 2026.07.17 🤗 We release the curated HeadRouter code, calibrated routing profiles, and Qwen2.5-Omni integration.
• 2026.07.10 🎉 HeadRouter was accepted by ACM MM 2026!
• 2026.04.30 📄 The HeadRouter paper was released on arXiv.

💡 Motivation

🎯 Highlights

• Task-adaptive routing: Softly mixes semantic, uniform, and acoustic head-weight profiles for every input.
• Position-agnostic scoring: Uses a no-RoPE text-to-audio QK probe to reduce positional bias.
• Training free: Requires no additional model training or parameter update.
• Strong compression quality: Retaining 70% of audio tokens reaches 101.8% and 103.0% of the vanilla average on Qwen2.5-Omni-3B and Qwen2.5-Omni-7B, respectively.

✨ Overview

TL;DR: HeadRouter measures per-head audio selectivity, routes each sample among calibrated task profiles with a Gaussian kernel, and uses the resulting head weights to retain task-relevant audio tokens.

💥 Core Code

Component	Path
Gaussian soft routing and token selection	headrouter/routing.py
Profile loading and model configuration	headrouter/config.py
Qwen2.5-Omni integration	headrouter/models/modeling_qwen2_5_omni.py
Calibrated Qwen2.5-Omni-3B profiles	configs/qwen2_5_omni_3b/
Single-audio inference example	examples/qwen2_5_omni_inference.py

The experimental code used the internal name HAWK. Public APIs and documentation use HeadRouter; a few hawk_* configuration names remain inside the model snapshot for behavior-compatible loading.

🛠 Installation

git clone https://github.com/DabDans/HeadRouter.git
cd HeadRouter

conda create -n headrouter python=3.10 -y
conda activate headrouter
pip install --upgrade pip
pip install -e .

Flash Attention 2 is recommended for reproducing the reported efficiency. Install a build compatible with your CUDA and PyTorch versions before running with --attn-implementation flash_attention_2.

🚀 Quick Start

python examples/qwen2_5_omni_inference.py \
  --model Qwen/Qwen2.5-Omni-3B \
  --audio /path/to/example.wav \
  --prompt "Describe the audio in detail." \
  --pruning-ratio 0.6 \
  --output outputs/example.json

The example uses layer 1 for the no-RoPE probe and prunes at layer 2 by default. Set --profile-dir and --routing-config to evaluate another calibrated profile set.

Minimal model configuration

from headrouter import configure_headrouter

configure_headrouter(
    model,
    pruning_ratio=0.6,
    pruning_layer=2,
    profile_dir="configs/qwen2_5_omni_3b/layer1",
    routing_config="configs/qwen2_5_omni_3b/routing.json",
)

📊 Evaluation

The paper evaluates HeadRouter on AudioMarathon and MMAU-Pro. Use the public example as the integration template and follow each benchmark's official data preparation and metric implementation.

The repository intentionally excludes checkpoints, datasets, Slurm logs, raw per-sample predictions, and internal absolute paths. The released profiles are small calibration artifacts required by the method.

📌 Citation

ACM MM 2026

@inproceedings{he2026headrouter,
  title     = {HeadRouter: Dynamic Head-Weight Routing for Task-Adaptive Audio Token Pruning in Large Audio Language Models},
  author    = {He, Peize and Luo, Yaodi and Liu, Xiaoqian and Liu, Xuyang and Deng, Jiahang and Du, Yaosong and Li, Bangyu and Gui, Xiyan and Chen, Yuxuan and Zhang, Linfeng},
  booktitle = {Proceedings of the 34th ACM International Conference on Multimedia},
  year      = {2026}
}

arXiv

@misc{he2026headrouterarxiv,
  title         = {HeadRouter: Dynamic Head-Weight Routing for Task-Adaptive Audio Token Pruning in Large Audio Language Models},
  author        = {He, Peize and Luo, Yaodi and Liu, Xiaoqian and Liu, Xuyang and Deng, Jiahang and Du, Yaosong and Li, Bangyu and Gui, Xiyan and Chen, Yuxuan and Zhang, Linfeng},
  year          = {2026},
  eprint        = {2604.23717},
  archivePrefix = {arXiv},
  primaryClass  = {cs.SD},
  url           = {https://arxiv.org/abs/2604.23717}
}

👍 Acknowledgment

We thank the open-source communities behind Qwen2.5-Omni, AudioMarathon, and MMAU-Pro.

📩 Contact

For questions about the paper or code, please email 2023300904027@std.uestc.edu.cn or zhanglinfeng@sjtu.edu.cn.