EMNLP 2025 Main Conference

Router-Tuning: A Simple and Effective Approach for Dynamic-Depth in Transformers

Fine-tune only lightweight router heads to empower existing Transformers with dynamic-depth token skipping — capturing the efficiency of Mixture of Depths at a fraction (<0.001%) of full-model retraining cost.

Shwai He¹ Tao Ge² Guoheng Sun¹ Bowei Tian¹ Xiaoyang Wang² Dong Yu²

¹ University of Maryland, College Park ² Tencent AI Lab, Bellevue, WA

arXiv:2410.13184 ACL Anthology Code & Checkpoints Live Simulator

Figure: Router-Tuning vs. Traditional Mixture-of-Depths Adaptation Backbone Frozen · Only Linear Heads Updated

<0.001%

Tuned Parameters

100% Backbone Frozen

Up to 48%

FLOPs Reduction

Layer-wise Token Skipping

1.45×

Inference Acceleration

Measurable real-world latency

>99.3%

Accuracy Maintained

<0.5% drop on GSM8K/MMLU

Interactive Demo

Dynamic-Depth & Token-Routing Simulator

Adjust the token compute budget and routing granularity to observe real-time estimated FLOPs savings, speedup, and token-skipping heatmaps across Transformer layers!

Model & Router Controls

Backbone Architecture

Routing Granularity

Target Token Capacity ($C_{\text{target}}$) 50%

20% (Max Skip) 50% (Recommended) 100% (Dense)

Active Router Layers 16 Layers

4 deep layers 16 middle/deep All layers

Simulation Prompt

"The swift brown fox analyzes quantum fluctuations in de Sitter space with minimal compute budget."

FLOPs Saving

47.8%

vs Dense Baseline

Speedup

1.44×

Theoretical Latency

KV & Mem Saved

42.1%

Memory Bandwidth

Accuracy Est.

99.2%

>70.6 pts (GSM/MMLU)

Token Routing Heatmap Across Depth

Green = Computed Layer • Dotted Grey = Skipped (Residual Pass)

Executed Skipped

Routing Insight: Semantic tokens (e.g. "quantum", "fluctuations", "de Sitter") maintain high execution density in deep layers, whereas grammatical glue words (e.g. "the", "in", "with") are safely bypassed after initial representation forming.

Architectural Principles

How Router-Tuning Operates

Instead of updating billions of LLM parameters, Router-Tuning optimizes ultra-compact linear routing heads with Straight-Through Estimators (STE) to enable end-to-end discrete token routing.

Router-Only Optimization

The transformer backbone weights $\mathbf{W}_{\text{backbone}}$ are 100% frozen. We only train a lightweight router projection $\mathbf{W}_r \in \mathbb{R}^{d \times 1}$ per layer, drastically eliminating catastrophic forgetting and reducing tuning VRAM.

$$\mathbf{s}_l = \sigma(\mathbf{W}_r \mathbf{h}_l)$$

Discrete STE Backprop

Binary routing masks $\mathbf{m} \in \{0, 1\}$ are formed via thresholding. The Straight-Through Estimator copies upstream gradients directly to routing logits, allowing gradient descent across discrete skip decisions.

$$\mathbf{m}_l = \text{STE}(\mathbf{s}_l - \tau)$$

Capacity Regularization

An auxiliary capacity loss forces the average activation ratio to obey the user-specified compute budget $C_{\text{target}}$, preventing trivial all-skip or all-compute collapse.

$$\mathcal{L}_{\text{cap}} = \lambda \cdot \text{ReLU}(\bar{\mathbf{m}} - C_{\text{target}})$$

Experimental Results

Pareto Frontier: Accuracy vs. Acceleration

Router-Tuning outperforms static pruning and matches full-model dynamic tuning across extensive benchmarks, all with orders-of-magnitude fewer trainable parameters.

Benchmark Quality-Speedup Tradeoff

Interactive comparison across LLM efficiency paradigms

Figure 1

Main Benchmark Tradeoffs

Router-Tuning establishes an optimal Pareto frontier across diverse compute budgets compared to full-parameter adaptation.

Figure 2

Expert Routing Analysis

Routing specialization emerges naturally: early layers focus on general representations while deeper layers selectively engage.

Figure 3

LoRA Compatibility

Router-Tuning effortlessly composes with parameter-efficient fine-tuning (LoRA) adapters with zero interference.

Developer Guide

Quickstart in 3 Simple Steps

Get started with environment setup, dataset preparation, distributed router fine-tuning, and downstream dynamic evaluation.

bash terminal

# 1. Create and activate a fresh python environment
conda create -n router-tuning python=3.10 -y
conda activate router-tuning

# 2. Clone repository & install dependencies
git clone https://github.com/CASE-Lab-UMD/Router-Tuning-Mixture-of-Depths.git
cd Router-Tuning-Mixture-of-Depths
pip install -r requirements.txt

# 3. (Optional) Install FlashAttention-2 for faster execution
pip install --no-build-isolation flash-attn==2.6.3

# Reformat raw instruction datasets into unified message schemas
python entrypoints/data/reformat_datasets.py \
  --raw_data_root ./data/raw \
  --save_path ./data/reformatted

# (Optional) Mix multiple datasets for balanced instruction tuning
python entrypoints/data/mix_datasets.py \
  --reformatted_dir ./data/reformatted \
  --save_path ./data/mixed

# Launch distributed router-tuning with DeepSpeed + Accelerate
MODEL_NAME_OR_PATH="Qwen/Qwen2.5-7B" \
DATA_TYPE="alpaca" \
GRANULARITY="attn_sequence" \
ROUTER_LAYERS=16 \
TARGET_CAPACITY=0.5 \
NUM_PROCESSES=4 \
bash scripts/finetune_router_tuning.sh

# Evaluate fine-tuned checkpoint on standard reasoning suites via lm-eval
lm_eval --model hf \
  --model_args pretrained=./trained_models/qwen2.5-7b/alpaca/1000/attn_sequence_epoch1_router_layers16_lambda0.0_lr1e-05_wd0.0 \
  --tasks gsm8k,mmlu,arc_challenge,hellaswag \
  --batch_size 16

BibTeX Citation

Please cite our work if you find this project helpful in your research

@inproceedings{he2025routertuning,
  title     = {Router-Tuning: A Simple and Effective Approach for Enabling Dynamic-Depth in Transformers},
  author    = {He, Shwai and Ge, Tao and Sun, Guoheng and Tian, Bowei and Wang, Xiaoyang and Yu, Dong},
  booktitle = {Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing (EMNLP 2025)},
  year      = {2025},
  url       = {https://arxiv.org/abs/2410.13184}
}