Question about expert choice rouiting

[lionel-me]

Hello, thank you for your work. I'd like to ask some questions about expert choice routing:

1. While looking for code implementations of this part, I found that Google's implementation mentions in a comment that expert choice routing is not suitable for decoder-only architectures (https://github.com/google/flaxformer/blob/main/flaxformer/architectures/moe/routing.py#L655). I'm curious about how you handled this issue in your experiments.
2. I'd like to know how you perform expert choice routing calculations during inference. With kv-cache, the ffn layer typically needs to compute only one token at a time. Or, does the expert choice routing model not utilize kv-cache during inference?

[Muennighoff]

1. Hm I don't recall doing something to avoid the model cheating during EC & I don't think it is very easy for the model to learn to cheat as they mention in their comment?
2. We didn't have to deal with that as we only trained & evaled on tasks during training that do not require generation; I think you can do something like this paper to deal with it: https://arxiv.org/abs/2404.02258