The old’ 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)\) (“which module \(m\) should we assign, given token \(x\)”), 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}

Its MOEReview Li: Branch-Train-Merge but MoEs now. Each layer is combined by standard moe routing with a weight that is tuned.

A single kernel to scatter the residuals and then run forward pass at the same time instead of copying and grouping first.

Split \(Q\) projection and attention out projection into experts, with one router coordinating them.

Better than MHA performanec.

Bellman-Ford Algorithm is a dynamic programming problem that solves single-source shortest path.

d[v] = \infty
d[s] = 0
for i = 0, ..., n-2:
    for v in v:
        d_prev = d
        for u in v.in_neigbors:
            d[v] = min(d_prev[v], d_prev[u] + w(u,v))

Notice you need \(O\qty(n)\) space (2 \(d\) rounds, the previous round and the next round), and runtime is \(O\qty(nm)\) (outer \(n\) loop, inner is an iteration between \(deg\qty(v)*|v| = |e| = m\); that is, we have an minimum over the degree of each \(v\) for every \(v\), which adds up to the total number of edges.)