Tag handelling and new post.

Author: localhost433

posts/entries/011-MoE-1.md

@@ -1,7 +1,7 @@
 ---
 title: Auxiliary-loss Load Balancing in MoEs (1)
 date: 2025-07-06
-tags: [cs, finance, ai]
+tags: [cs, ai, notes]
 author: R
 location: Above Illulissat, Greenland while on a plane from New York to Hong Kong
 ---
@@ -17,9 +17,9 @@ Let
 - $K$ the number of experts selected per token,  
 - $T$ the total number of tokens in the batch,  
 - $\mathbf{u}_t\in\R^d$ the input for token $t$,  
-- $\mathrm{FFN}_i: \R^d\to\R^d$ the $i$-th expert network,  
+- $\text{FFN}_i: \R^d\to\R^d$ the $i$-th expert network,  
 - $e_i\in\R^d$ the centroid (parameter) of expert $i$, and  
-- $G\colon\R\to\R_{>0}$ a positive gating function (e.g. $\exp$, $\mathrm{sigmoid}$, or $\mathrm{softmax}$).
+- $G\colon\R\to\R_{>0}$ a positive gating function (e.g. $\exp$, $\text{sigmoid}$, or $\text{softmax}$).
 
 Compute for each token $t$ and expert $i$:
 $$
@@ -30,7 +30,7 @@ $$
                  &          (\text{if $s_{i,t}$ is among the top-$K$ scores})
         \\
         0, & \text{otherwise}
-        \end{cases}\\
+        \end{cases} \\
     s_{i,t} &= G(\textbf{u}_t^\top e_i)
 \end{align*}
 $$
@@ -40,54 +40,118 @@ $$
 \textbf{h}_t = \textbf{u}_t + \sum^N_{i=1} g_{i,t} \text{FFN}_i (\textbf{u}_t) \\
 $$
 
-So, here $G$ could be something that is $\R \to \R_{>0}$, some conventional ones could be a $\exp$, softmax or sigmoid (TBH I have to search these up to see what they actually are). In this paper they have use the latter 2.
+So, here $G$ could be something that is $\R \to \R_{>0}$, some conventional ones could be a $\exp$, softmax or sigmoid (TBH I have to search these two up to see what they are exactly). In this paper they have use the latter 2.
 
 And there is the Expert consulted following the gating function.
 
-### Problem: Inbalanced routing
+## Problem: Inbalanced routing
 But one problem MoEs often experience is inbalanced routing (a few number of experts recieve the most token), thus *a risk of routing collapse (Shazeer et al., 2017), where the model consistently selects only a few experts, hindering sufficient training of the other experts*; or, a *computational bottleneck by load inbalance*.
 
 I was wondering how it could cause a computational bottleneck, but then I realized the way I thought about it that it could easily scale through parallelism or some other ways is not easily achievable. Since there are different machines hosting each model, it depends more on the load given to a certain expert.
 
 Plus, the training loop should undergo a substantial redesign for it to use the idle computational power to catch up. Even if I create replicas for the "hot" experts on more hosting devices, they need to be in sync, therefore creating a lot of cost by itself. Merging gradients across replicas requires collective operations every step, at that point it will just recreate the original problem trying to overcome if 1 of these slowed down...
 
-#### Solution: Auxiliary-loss
-To address this issue, there is auxiliary-loss encourage balanced load thus avoids inbalanced routing in training MoEs. To do this, it penalized the use of only a few number of agents. Its mostly within the process of the gating function. Defined as such:
+### Solution: Auxiliary-loss
+To address this issue, there is auxiliary-loss encourage balanced load thus avoids inbalanced routing in training MoEs. To do this, it penalized the use of only a few number of agents. Its mostly within the process of the gating function. 
 
+**Key variables:**
+- $N$: number of experts in the MoE layer
+- $K$: number of experts selected per token (top-K)  
+- $T$: total number of tokens in the batch
+- $\mathbb{1}$: indicator function (equals 1 if condition is true, 0 otherwise)
+- $\alpha$: balancing‐loss weight (manually set hyperparameter)
 
+Defined as such:
 
-- **Normalized load**  
+- **Normalized load**   
+  $f_i$:= the fraction of tokens routed to expert $i$:
   $$
-    f_i = \frac{N}{KT} \sum_{t=1}^T \mathbb{1} (i \in \mathrm{Topk} \mid \mathbf{u}_t )
-    \quad (\text{fraction of tokens routed to expert }i ).
+    f_i = \frac{N}{KT} \sum_{t=1}^T \mathbb{1} (i \in \text{Topk} \mid \mathbf{u}_t )
   $$
-- **Average gating weight**
+
+- **Average gating weight**  
+  $P_i$:= the mean score assigned by the gate to expert $i$:
   $$
-    P_i = \frac{1}{T}\sum_{t=1}^T s_{i,t}
-    \quad (\text{mean score assigned by the gate to expert }i ).
+    P_i = \frac{1}{T} \sum_{t=1}^T s_{i,t}
   $$
-`
-##### Balance loss
 
 Combine these into a single penalty term:
+  $$
+    \mathcal{L}_{\text{balance}} = \alpha \sum_{i=1}^N f_i P_i.
+  $$
+
+
+**Regularization terms:**  
+Introduce two small-weight penalties on the imbalance of $\{P_i\}$ and $\{f_i\}$:
+$$
+\begin{align*}
+  \mathcal{L}_{P} &= \lambda_{P} \text{CV}^2 (\{P_i\} )\\
+  \mathcal{L}_{f} &= \lambda_{f} \text{CV}^2 (\{f_i\} )
+\end{align*}
+$$
+where typically $\lambda_{P} \approx\lambda_{f} \sim10^{-2}$.
+
+> This is actually optional, for simpler just use $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \mathcal{L}_{\text{balance}}$.  
+> I write it this way just to follow the [original MoE auxiliary-loss formulation paper (Shazeer et al. (2017))](https://arxiv.org/pdf/1701.06538).
+
+
+**Imbalance metric: coefficient of variation squared**  
+For any set of scalars $\{z_i\}_{i=1}^N$, define
+$$
+  \text{CV}^2(\{z_i\}) =
+  \frac{\frac{1}{N} \sum_{i=1}^N z_i^2 - (\frac{1}{N} \sum_{i=1}^N z_i )^2}
+       {(\frac{1}{N} \sum_{i=1}^N z_i )^2},
+$$
+which satisfies $\text{CV}^2=0$ exactly when all $z_i$ are equal.
+
+> Btw, this looks like the variance so much.  
+> Write $\mu = \tfrac1N \sum_i z_i$ and $\nu = \tfrac1N \sum_i z_i^2$. Then
+> $$
+> \text{CV}^2 = \frac{\nu - \mu^2}{\mu^2} = \frac{\text{Var}}{(\text{Mean})^2}
+> $$
+> Its partial derivative w. one coordinate $z_k$ is
+> $$
+> \frac{\partial \text{CV}^2}{\partial z_k}
+> = \frac{2}{N}\Bigl(\frac{z_k}{\mu^2} - \frac{\nu}{\mu^3}\Bigr).
+> $$
+> > Details:
+> > $$
+> > \begin{align*}
+> >   \frac{\partial}{\partial z_k} (\tfrac{\nu - \mu^2}{\mu^2})
+> >   &= \frac{1}{\mu^2} \frac{\partial\nu}{\partial z_k}
+> >    - \frac{\nu - \mu^2}{\mu^4} 2\mu \frac{\partial\mu}{\partial z_k}\\
+> >   &= \frac{1}{\mu^2} \frac{2z_k}{N}
+> >    - \frac{\nu - \mu^2}{\mu^4} \frac{2\mu}{N}\\
+> >   &= \frac{2}{N}\Bigl(\frac{z_k}{\mu^2} - \frac{\nu}{\mu^3}\Bigr).
+> > \end{align*}
+> > $$
+> 
+> Because $\nu/\mu^3$ is the same constant for all $k$, this gradient pushes down any $z_k > \mu$ (overloaded expert) and pushes up any $z_k < \mu$ (underloaded expert). In other words, the derivative of a variance term normalized by $\mu^2$.
+
+
+**Total training objective**  
+Combine with the primary task loss $L_{\text{task}}$:
 $$
-  \mathcal{L}_{\mathrm{balance}}
-   = 
-  \alpha \sum_{i=1}^N f_i P_i.
+  \mathcal{L}_{\text{total}}
+  = \mathcal{L}_{\text{task}}
+  + \mathcal{L}_{P}
+  + \mathcal{L}_{f}.
 $$
 
-##### Why this encourages balanced routing
+#### Intuition
+- The penalty grows, as either $f_i$ or $P_i$ grows (since it's a product). Then the routing distribution is driven toward uniformity by the penalties. Backpropagation through the parameters plays a role in this process.
+- Minimizing $\text{CV}^2$ drives the variance of $\{\text{Imp}_i\}$ or $\{\text{Load}_i\}$ toward zero relative to their mean. (see derivation of $\partial \text{CV}^2/\partial z_k$ above)
+- Any expert $i$ with above-average usage raises its own $\text{Imp}_i$ or $\text{Load}_i$, increasing the penalty.
 
-- $f_i$ captures how heavily expert $i$ is used, while $P_i$ captures its average gate score.  
-- If an expert is over-selected ($f_i$ large), the product $f_iP_i$ grows, increasing the penalty.  
-- Gradients then adjust the gating parameters to **decrease** routing to over-used experts and **increase** routing to under-used ones, driving the distribution toward uniformity.
 
-##### Why this balances load
+#### Drawbacks
+The ICLR 2025 mentioned that auxiliary loss might introduce unwanted gradients, as the MoE models performs worse on some metrics.
 
-- Minimizing $\mathrm{CV}^2$ drives the variance of $\{\mathrm{Imp}_i\}$ or $\{\mathrm{Load}_i\}$ toward zero *relative* to their mean.  
-- Any expert $i$ with above-average usage raises its own $\mathrm{Imp}_i$ or $\mathrm{Load}_i$, increasing the penalty.  
-- Backpropagation through the gating parameters encourages **reduced** routing to overloaded experts and **increased** routing to underutilized ones, leading to a more uniform expert selection distribution.  
+However, I wasn't really convinced about this reasoning. Because the performance was not improved that significantly (I was expecting a larger gap) for the Validation Perplexity. The load balance one sound ok, and that's the main point of the paper, so it's good.
 
-Basically, divide by the variance.
+The true drawback, in my opinion comes with the act of rebalancing through auxiliary-loss itself.
+- The idea of MoE is having a lot of highly specialized experts, auxiliary-loss fights any concentration of weight, even if that concentration was beneficial for modeling those tokens.
+- The balancing gradient for expert involves all experts' totals. So updating the logit for one expert now depends on every other expert’s load. It's obvious it can drown out more specialized signals.
+- Naturally, experts that are good at certain tokens are expected to get those; trying to make the router equalize loads regardless of quality can route a token to a weaker expert, simply because the "best" expert is already slightly busier.
 
 (TBC)

posts/js/blog.js

@@ -97,8 +97,14 @@ searchBox.addEventListener("input", e => {
 // Handle tag button clicks
 function handleTagClick(tag) {
   const params = new URLSearchParams(location.search);
-  params.set("tag", tag);
-  history.pushState({ tag }, "", `?${params.toString()}`);
+  if (tag) {
+    params.set("tag", tag);
+  } else {
+    params.delete("tag");
+  }
+  const newUrl = window.location.pathname +
+    (params.toString() ? `?${params.toString()}` : "");
+  history.pushState({ tag }, "", newUrl);
   applyTagFilter(tag);
 }