Gradient-aligned routing improves multi-task text classification accuracy
Routing in Gradient Space: Balanced Usage Is Not Expert Specialization
Machine Learning
Summary
This work looks at how computer models choose which specialized parts handle different tasks when learning many things at once. The authors found that organizing these choices based on how training signals align in their gradients leads to better accuracy than just balancing workload or routing based on task loss alone. Their method, called gradient-aligned routing, groups similar training signals together, helping the model learn more effectively across multiple text classification tasks.
What this means in practice
- •For natural language processing engineers: Implement gradient-aligned routing to improve multi-task text classification performance in large language model architectures.
- •For machine learning infrastructure teams: Develop routing algorithms that balance workload while grouping tasks by gradient similarity to optimize resource usage and model training efficiency.
Authors
Yuchen Li, Mingyu Du, Zongqi Fan, Nguyen H. Tran, Ken-Tye Yong
Abstract
Sparse expert models can distribute traffic evenly while still grouping incompatible training signals within the same experts. We study routing as a gradient-partitioning problem and introduce gradient-aligned routing (GAR), whose load-normalized router objective rewards grouping observations with aligned gradients. On five multi-task text-classification mixtures, we compare GAR with task-loss-only routing, gradient-combination and gradient-conflict methods, and load-balancing losses. With a fully trainable RoBERTa backbone and classification-head experts, GAR has the highest aggregate validation accuracy, 1.07 percentage points above task-loss-only routing. With frozen DeBERTa and Qwen3-1.7B backbones and low-rank adapter experts, it again ranks first, 1.10 points above task-loss-only routing, with better-balanced expert load and higher gradient-mass purity, the share of each expert's gradient-norm mass from its dominant task; the load-balancing losses flatten load further but leave this purity near its task-loss-only level. Top-1 routing, trainable full-parameter feed-forward network (FFN) experts, and a larger backbone also show positive aggregate gains. The results distinguish expert-load balance from gradient-based routing organization and indicate the predictive value of gradient-informed routing in multi-task text classification.