Suggest edit — MOEReview Shen: ModuleFormer

Title

Name

Note

---
title: "MOEReview Shen: ModuleFormer"
source: https://www.jemoka.com/posts/kbhmoereview_shen_moduleformer/
---

The old&rsquo; load balancing loss. Instead of training a router with explicitly labeled data for each expert, a load balancing + load concentration loss induces the modularity in data.
Insight: we want to maximize the mutual information between tokens and modules. For the router \(m \sim g\qty(\cdot \mid x)\) (&ldquo;which module \(m\) should we assign, given token \(x\)&rdquo;), we write:
\begin{equation} \ell_{MI} = \underbrace{\sum_{m=1}^{N} p\qty(m) \log p\qty(m)}_{-H\qty(m)} - \frac{1}{|X|} \sum_{x \in X}^{} \underbrace{\sum_{m=1}^{N} g\qty(m|x) \log g\qty(m|x)}_{H\qty(m|x)} \end{equation}
that is, we want to maximize the entropy of the marginal distribution (&ldquo;without knowing what tokens, we shouldn&rsquo;t be opinionated about which module&rdquo;) while minimizing the entropy of the conditional distribution (&ldquo;knowing which tokens, we should know which modules.&rdquo;)
At fine-tuning time (i.e. if we do have task specific data), we can directly minimize the entropy:
\begin{equation} \ell_{\text{concentration}} = - \sum_{m=1}^{N} p\qty(m) \log p\qty(m) \end{equation}
since we want to use one module for each specialized fine tuning mix.