TEXAS: Task-Expert-Aware Supervision for Downstream Mixture-of-Experts LLM Adaptation

2026-07-31

AI Summary

This work addresses the challenge that existing Mixture-of-Experts (MoE) language models struggle to accurately identify task-relevant experts during downstream adaptation and fail to effectively leverage expert activation signals. The authors propose a correctness-conditioned mechanism for task-specific expert discovery, which dynamically identifies critical experts by contrasting activation patterns between successful and failed inference samples. During fine-tuning, higher supervision weights are applied to tokens corresponding to incorrect answers along the activation paths of these identified experts. Notably, this approach requires no predefined expert subsets or fixed routing distributions, enabling more efficient adaptation. Evaluated across three MoE models and six benchmarks—yielding 18 experimental settings—the method achieves state-of-the-art or tied-best performance in 17 cases, outperforming the strongest baseline by an average of 1.3–1.5 points.

Abstract

Mixture-of-Experts (MoE) language models route each token through a small subset of experts, making routing patterns useful for identifying task-relevant experts during downstream adaptation. Yet current approaches have two limitations: task experts are typically identified from aggregate routing statistics that reflect usage rather than association with successful task completion, and task-expert activations remain underexplored as signals for supervision allocation. We introduce Task-Expert-Aware Supervision (TEXAS), which combines correctness-conditioned task expert discovery with token-level supervision allocation. TEXAS compares expert activations on instances that the base model solves successfully and those it fails to solve, and retains experts more strongly activated on successful instances. During fine-tuning, it upweights answer tokens in failed instances when they activate these experts. TEXAS therefore leverages existing routing behavior without restricting adaptation to a fixed expert subset or imposing an explicit target routing distribution. Across three MoE models and six benchmarks, TEXAS achieves the best or tied-best performance in 17 of 18 settings and improves over the strongest baseline by 1.3--1.5 points on average. Ablations and further analyses validate both the discovered experts and the resulting supervision strategy.

Problem

• Mixture-of-Experts
• downstream adaptation
• task-expert association
• supervision allocation
• expert routing

Innovation

• Mixture-of-Experts
• task-expert discovery
• correctness-conditioned supervision
• token-level supervision allocation
• expert activation analysis