💬 NLP

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

The paper introduces TEXAS, a novel adaptation method for Mixture-of-Experts language models that identifies task-relevant experts by comparing their activations on successful versus failed instances and leverages this discovery to dynamically upweight supervision on answer tokens during fine-tuning, thereby achieving state-of-the-art performance across multiple benchmarks without restricting expert subsets.

Original authors: Guanzhi Deng, Haibo Wang, Kuan Wu, Xiangru Jian, Shing Yin Wong, Sichun Luo, Zhuoran Wang, Linqi Song

Published 2026-08-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Guanzhi Deng, Haibo Wang, Kuan Wu, Xiangru Jian, Shing Yin Wong, Sichun Luo, Zhuoran Wang, Linqi Song

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, super-smart robot brain made of thousands of tiny, specialized workers. This is how modern "Mixture-of-Experts" (MoE) AI models work. Instead of one giant brain doing everything, the model has a "router" that acts like a traffic cop, sending each word of a sentence to just a few of the best-suited workers (the "experts") to handle it. Some experts are great at math, others at coding, and some at telling jokes. The big question researchers are asking is: when we want to teach this robot a new skill, like solving complex math problems, how do we know which workers are actually the geniuses for that job?

Usually, people try to find the best workers by counting how often they are used. It's like assuming the most popular employee in a company is the best at the new project. But what if that employee is just busy doing something else, or is actually terrible at the new task? This paper, titled TEXAS, suggests that counting popularity isn't enough. Instead, we should look at who shows up when the robot gets the answer right versus when it gets it wrong. The authors propose a new way to train these models that listens to the "winning" workers more closely, helping the robot learn faster and better without needing to rebuild its entire brain.

The Problem: Counting Popularity vs. Finding Genius

The researchers started with a hunch: the old way of finding "task experts" (the specific workers good at a specific job) was flawed. Most previous methods looked at aggregate statistics—basically, they counted how often a specific expert was chosen across the whole dataset. They assumed that if an expert was used a lot, they must be the right one for the job.

However, the authors found a mismatch. They looked at a model trying to solve math problems and found that the "popular" experts were often just busy, not necessarily helpful. In fact, some of the most frequently used experts were actually more active when the model failed to solve a problem than when it succeeded! It's like hiring a chef because they are always in the kitchen, only to realize they are actually just burning the toast every time. The paper argues that aggregate usage statistics are an incomplete and sometimes misleading guide for finding the true experts.

The Solution: TEXAS (Task-Expert-Aware Supervision)

To fix this, the team introduced a framework called TEXAS. Instead of just counting how often an expert is used, TEXAS plays a game of "Spot the Difference" between success and failure.

Here is how it works, step-by-step:

The Trial Run: First, the model tries to solve a bunch of problems on its own. The researchers split these attempts into two piles: the "Success Pile" (where the model got the right answer) and the "Failure Pile" (where it messed up).
The Detective Work: They then look at which experts were active during the "Success Pile" versus the "Failure Pile." They are looking for the experts that show up significantly more when the model gets it right. These are the true "Task Experts."
The Training Boost: Now comes the clever part. When they start training the model on the "Failure Pile" (the mistakes), they don't treat all mistakes equally. If a mistake happens to activate one of those "True Task Experts," TEXAS says, "Hey, this expert is usually good at this! Let's pay extra attention to this specific moment." They increase the "weight" or importance of that specific part of the lesson.

Think of it like a music teacher. If a student plays a song wrong, the teacher might say, "You played the notes right, but your rhythm was off." But with TEXAS, if the student makes a mistake but their hand is in the exact right position for a difficult chord, the teacher says, "Great hand position! Let's focus all our energy on fixing the rhythm for this specific chord because your hand knows what it's doing." The model learns to trust the pathways that lead to success, even when it's currently failing.

What They Found

The team tested TEXAS on three different large AI models and six different tasks, ranging from solving math problems (GSM8K, MATH500) to writing code (HumanEval, MBPP) and following complex instructions (IFEval).

The results were quite strong. In 17 out of 18 different test scenarios, TEXAS performed the best or tied for the best. On average, it improved the model's performance by 1.3 to 1.5 points compared to the strongest existing methods. For example, on the GSM8K math benchmark, TEXAS boosted the DeepSeek model's score from 59.9 to 62.5.

The researchers also ran "ablation studies" (which are like taking the engine apart to see which part makes it run) to prove their ideas. They found that:

Correctness matters: If they just picked experts based on popularity (frequency) instead of success, the model didn't do as well.
Focus matters: If they boosted the learning signal for all mistakes, rather than just the ones involving the "True Task Experts," the improvement was smaller.
The experts are real: When they "masked" (turned off) the experts found by TEXAS, the model's performance dropped much more than when they turned off experts found by other methods. This suggests TEXAS really did find the functional geniuses.

Why It Works (and What It Doesn't Do)

The paper suggests that TEXAS works because it reinforces the "pathways" in the AI's brain that are associated with getting things right. By giving extra weight to the moments where the right experts are active during a failure, the model learns to strengthen those specific connections.

Importantly, the authors clarify what TEXAS is not. It does not force the model to use a fixed set of experts, nor does it try to change the model's internal "traffic cop" (the router) to force it to pick specific workers. It simply listens to the natural behavior of the model and amplifies the learning signal where it matters most.

The paper also notes a small exception: on the MMLU benchmark (a general knowledge test), TEXAS performed comparably to other top methods but didn't show the same massive jump. The authors suggest this might be because general knowledge is so broad and scattered that it doesn't rely on a single, consistent set of "genius" experts in the same way math or coding does.

In short, TEXAS suggests that when teaching a complex AI, we shouldn't just look at who is busy; we should look at who is winning. By paying extra attention to the "winning" workers even when the team is losing, we can help the whole team learn faster and smarter.

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

Problem Statement

Mixture-of-Experts (MoE) language models utilize sparse routing to activate a subset of experts for each token, enabling scalability and expert specialization. While recent studies suggest that routing patterns encode signals useful for identifying "task experts" (experts relevant to specific downstream capabilities), existing adaptation methods suffer from two primary limitations:

Ineffective Task Expert Identification: Current approaches typically identify task experts using aggregate routing statistics (e.g., average gate scores or selection frequencies) over task data. However, the paper demonstrates that high aggregate usage does not reliably correlate with successful task completion. Experts selected by frequency often exhibit mixed or negative associations with success, failing to distinguish between experts that drive correct outputs and those that do not.
Underutilized Supervision Signals: Existing methods primarily leverage identified experts for selective fine-tuning (updating only specific experts) or routing optimization (forcing specific routing distributions). They largely overlook the potential of task-expert activations as dynamic signals for allocating stronger training supervision at the token level.

Methodology: TEXAS

The authors propose Task-Expert-Aware Supervision (TEXAS), a framework that combines correctness-conditioned task expert discovery with token-level supervision allocation. TEXAS operates in two tightly coupled stages:

1. Correctness-Conditioned Task Expert Discovery

Instead of relying on aggregate statistics, TEXAS identifies task experts by comparing their activation patterns on instances the base model solves successfully versus those it fails.

Data Partitioning: The training set $D$ is split into a successful subset $D^+$ (where the base model $M_0$ meets the success criterion) and a failed subset $D^-$ .
Activation Analysis: For each layer-expert pair, the method calculates the answer-level activation rate on both subsets.
Statistical Selection: A one-sided Welch's t-test is performed to identify experts with significantly higher activation rates on successful instances ( $H_1: \mu^+ > \mu^-$ ). False discovery rates are controlled using the Benjamini–Hochberg procedure.
Candidate Filtering: To exclude statistically significant but rarely activated experts, the top- $K$ experts (based on mean activation rates on $D^+$ ) are retained. The final task-expert set $E_{task}$ is the intersection of statistically significant experts and these high-activation candidates.

2. Task-Expert-Aware Supervision

During fine-tuning, TEXAS dynamically adjusts the supervision weight of answer tokens based on whether they activate the discovered task experts.

Weighting Strategy: For instances in the failed subset $D^-$ , answer tokens that activate at least one discovered task expert receive an increased cross-entropy weight $\alpha$ (where $\alpha > 1$ ). Tokens in successful instances or those in failed instances that do not activate task experts retain a standard weight of 1.
Objective Function: The model optimizes a mean weighted cross-entropy loss:
$L_{TEXAS} = \frac{1}{N_B} \sum_{i \in B} \sum_{t \in A_i} w_{i,t} CE_{i,t}$
where $w_{i,t} = \alpha$ if the token activates a task expert in a failed instance, and $1$ otherwise.
Key Distinction: Unlike methods that restrict training to a fixed expert subset or impose a target routing distribution, TEXAS allows all experts to be updated but amplifies the learning signal for tokens lying on task-relevant computation paths identified by the base model's failures.

Key Contributions

Identification of Limitations: The paper highlights that aggregate routing statistics are an incomplete proxy for identifying experts associated with successful task completion and that task-expert activations remain underexplored for supervision allocation.
The TEXAS Framework: Introduces a method that discovers task experts via correctness-conditioned activation comparisons and utilizes these discoveries to dynamically upweight supervision for specific tokens in failed instances.
Empirical Validation: Extensive experiments across three MoE models (DeepSeek-V2-Lite, OLMoE, Qwen1.5-MoE-A2.7B) and six benchmarks (covering math, code, knowledge, and instruction following) demonstrate the framework's effectiveness.

Experimental Results

TEXAS was evaluated against baselines including standard Supervised Fine-Tuning (SFT), Expert-Selective Fine-Tuning (ESFT), and Routing Optimization (RoMA).

Performance: TEXAS achieved the best or tied-best performance in 17 out of 18 model-task settings.
Improvements: On average, TEXAS improved performance over the strongest baseline by 1.3–1.5 points. Compared to standard SFT, it yielded average gains of 2.4 to 3.0 points across the three models.
Ablation Studies:
- Replacing correctness-conditioned discovery with frequency-based selection (TEXAS-Freq) resulted in lower gains, validating the importance of the correctness-conditioned approach.
- Replacing loss upweighting with an auxiliary routing objective (TEXAS-Route) yielded smaller improvements, suggesting direct supervision amplification is more effective.
- Applying weighting to all instances or all tokens in failed instances (TEXAS-AllInst, TEXAS-AllTok) performed worse than the selective token-level approach, confirming the value of targeting specific task-relevant tokens in failed cases.
Analysis of Mechanism:
- SuccessGap: Experts discovered by TEXAS showed a significantly higher "SuccessGap" (activation difference between success and failure) on unseen benchmarks compared to those selected by ESFT or random selection.
- Functional Importance: Masking TEXAS-discovered experts caused greater performance degradation than masking ESFT-selected or random experts, indicating they are functionally critical.
- Pathway Reinforcement: TEXAS was shown to strengthen the activation of task-expert pathways, particularly on instances where the base model failed but the adapted model succeeded.
- Token Relevance: Tokens upweighted by TEXAS were found to be more enriched with task-relevant content (e.g., mathematical operators, code syntax) compared to high-loss tokens or random tokens.

Significance and Claims

The paper posits that internal computation pathways, specifically the activation of task-relevant experts, provide a valuable signal for allocating learning resources during adaptation. TEXAS demonstrates that:

Correctness-conditioned discovery is superior to aggregate statistics for identifying experts that genuinely contribute to task success.
Token-level supervision allocation based on these experts is a viable and effective strategy, avoiding the constraints of fixed expert subsets or explicit routing targets.
The method improves performance broadly across diverse tasks (math, code, instruction following) without requiring architectural changes or significant computational overhead during the fine-tuning phase (the discovery cost is a one-time offline step).

The authors conclude that TEXAS offers a new perspective on MoE adaptation: rather than just selecting which experts to train or how to route, one can use expert activations to determine where (at the token level) stronger supervision should be applied to correct specific failure modes.