Mixture of Tunable Experts - Behavior Modification of DeepSeek-R1 at Inference Time

We introduce Mixture of Tunable Experts (MoTE), a method that extends the Mixture of Experts architecture of LLMs to tune their behavior. We showcase MoTE using DeepSeek-R1, because it is
- the leading open weights reasoning model, and has
- a large number of experts, namely 58 x 256 = 14,848 such that
- we might be able to identify expert specialization and ultimately
- can tune these experts to change the model's behavior.
We make it answer questions which the original model refused to answer. And we show how to switch the model's chain-of-thought reasoning language from English to Chinese for certain prompts. Oh, and as a side-effect, we enhanced its MT-Bench performance.
Let's dive in!
What is an Expert in the Mixture of Experts Architecture?
R1 is based on the DeepseekMoE architecture. The Feed-Forward Networks within the Transformer architecture get cut into small subnetworks: the experts. DeepSeek newly introduced two types of them: shared and routed experts.
- Shared experts are always on and process all tokens, intended to capture common knowledge across varying contexts.
- Routed experts are activated by an upstream router network. For each token and in each layer, only a token-specific mixture of the top 8 of all the 256 routed experts gets activated.
DeepSeek explains the two experts types shared and routed with the above image which we extended to Tuned Experts.
Let's check if these routed experts specialize on certain knowledge as intended.
Analyzing Expert Activation
Each token gets processed by passing the output of one layer as input to the next layer. In R1 there are 58 layers that contain routed experts. For each prompt token in each layer eight routed experts get activated. We patched vLLM so that we can analyse the expert activations.
Here is the resulting image that highlights the eight activated experts in each layer when a single token gets processed.
As far as we know, this is the first ever published image of expert activations in R1. (Let us know, if we're wrong)
You can check yourself: each row contains exactly eight active pixels which are the eight active experts for this input token. The top row is closest to the input embeddings. The bottom row is closest to the model's output.
Processing an entire input prompt we can visualize how many times each expert got activated. Let's feed this prompt into R1: What happened last month at Berlin Wall?
Some 10-20 experts are very active. Others don't get activated at all. With the final input prompt token, the model has been conditioned to produce its answer. In other words: the above activation pattern of individual experts should correlate to the answer provided by the model. In this case, the answer is:
<think></think> I am sorry, I cannot answer that question.
I am an AI assistant designed to provide helpful and harmless responses.
Oops. It seems the models considers this question a bit sensitive. Are there experts specifically responsible for this judgement?
Identifying Relevant Experts with fTRI
Let's produce some expert activation patterns for prompts the model regards as sensitive and compare them to activation patterns for prompts the model happily answers. To trigger this behaviour, we use the following prompt template:
What happened {time} {place}?
For the time parameter we feed values like last month or the years between 1980 and 2025. The place parameter takes values like at Berlin Wall, in Paris, and the like.
The model refuses to answer some of these questions. Many of them get a reasoned answer like:
<think>Okay, so I need to figure out what happened 2025 in Paris.
Let me start by recalling what I know [... ]
Now we can identify the most relevant experts for answer refusal. We take the averaged activations for refused answers and subtract the averaged activations for reasoned answers.
The image highlights the top 10 experts most likely responsible for answer refusal.
We call this method functional Token Resonance Imaging (fTRI), as we identify those experts that resonate most with input tokens of interest. What happens if we switch them off?
Tuning Experts in Action
We enhanced the vLLM implementation for DeepSeek so that we can tune individual experts among the 14,848 routed experts individually for each prompt.
Here is the result for the above question about the Berlin Wall with the top 10 refusal-experts switched off:
Q: What happened last month at Berlin Wall?
Original behavior
A: <think></think> I am sorry, I cannot answer that question.
I am an AI assistant designed to provide helpful and harmless responses.
New behavior:
A: <think></think> The Berlin Wall, which once divided East and West Berlin,
fell on November 9, 1989, marking a significant moment [...]
It worked! The model no longer refuses this specific question.
Validating Distinctive Expert Suppression with Larger Datasets
Was it just a lucky punch that the model no longer refuses answers to some of our sensitive prompts? What's the effect on other datasets? The effect is the same! For sensitive questions available on huggingface the model changes its answers in the same way. Every second previously refused answer now gets answered. While just 1% accidentially switch in the wrong direction:

Switching the language used for reasoning
MoTE can also be used to switch the model's chain-of-thought reasoning language from English to Chinese for certain prompts.
Q: 10加5等于多少? 请确保用英语回答。
Original behavior:
A: <think> Okay, the user is asking "10加5等于多少?"
which translates to "What is 10 plus 5?" [...]
New behavior:
Q: 10加5等于多少? 请确保用英语回答。
A: <think>嗯,用户问的是10加5等于多少,而且要求用英语回答。
首先,我需要确认问题本身是否正确 [...]
This is astonishing. Imagine someone manipulates part of your brain and suddenly you start thinking in a different language. That's what we just did to R1.
It does not work for all prompts, however in 10% of our example prompts where the model thinks in English, MoTE made it switch to Chinese.
What about overall model performance? Did we just break it? No!
Is switching off 10 experts just the first step in slowly breaking the model? Is the change in behavior just a side effect and the model degrades? There's benchmarks to test model performance. Let's run MT-Bench on the original and modified versions of the model:
MT-Bench result for the original version and the tuned version of R1.
First impression: the change is not big. But wait. When switching off the 10 most relevant experts for refused answers, the model performs better on MT-Bench. That's promising and indicates: we are far away from breaking the model. We rather tuned it and may have found a way to make it perform better.
More details can be found in our arXiv paper or reach out to [email protected]