% teach:
% gap: 读者知道 MoE 的 top-k 路由但不知道 LatentMoE 的 latent 接口和 SiTU-GLU
% takeaway: LatentMoE 通过 latent 接口把 routed 专家限制在 ℓ=d/2 上算, SiTU-GLU 用软上限防溢出
% jump: 为什么 routed 专家在 latent 空间而 shared 在全宽？省参数
% omit: load balancing loss

\section{SiTU-GLU 与 Stable LatentMoE}
\splabel{C5}

\subsection{SiTU-GLU：带软上限的激活}

SwiGLU 在低精度（fp16/bf16）训练时可能溢出：$\mathrm{silu}(x) \cdot x$ 没有上限。
SiTU-GLU 用 $\tanh$ 给门控和上投影加软上限：

\[
  \mathrm{SiTU}(x) = W_o \big[\underbrace{\beta_1 \tanh\!\left(\frac{W_g x}{\beta_1}\right) \cdot \sigma(W_g x)}_{\text{gate}} \;\cdot\; \underbrace{\beta_2 \tanh\!\left(\frac{W_u x}{\beta_2}\right)}_{\text{up}}\big]
\]

\begin{center}
\begin{tabular}{lp{8cm}}
\toprule
性质 & 说明 \\
\midrule
输出上限 & $\|\mathrm{SiTU}\|_\infty \leq \beta_1 \cdot \beta_2 = 4 \times 25 = 100$ \\
原点附近 & $\tanh(x/\beta) \approx x/\beta$，所以 $\beta \cdot \tanh(x/\beta) \approx x$，退化为 SwiGLU \\
远端 & 软饱和，防 fp16 溢出 \\
\bottomrule
\end{tabular}
\end{center}

\begin{codemathtop}{layers/latent\_moe.py — SiTU}
\begin{lstlisting}
class SiTU(nn.Module):
    def __init__(self, dim_in, dim_ff, beta1=4.0, beta2=25.0):
        self.w_g = nn.Linear(dim_in, dim_ff, bias=False)
        self.w_u = nn.Linear(dim_in, dim_ff, bias=False)
        self.w_o = nn.Linear(dim_ff, dim_in, bias=False)

    def forward(self, x):                    # [*, dim_in]
        wg = self.w_g(x)
        g = self.beta1 * tanh(wg / self.beta1) * sigmoid(wg)   # gate
        u = self.beta2 * tanh(self.w_u(x) / self.beta2)        # up
        return self.w_o(g * u)               # [*, dim_in]
\end{lstlisting}
\end{codemathtop}

\subsection{LatentMoE 架构}

\begin{center}
\begin{tabular}{rl}
\toprule
组件 & 说明 \\
\midrule
\textbf{Shared 专家} & $n_{\mathrm{shared}}$ 个 SiTU，全宽 $d \to d$，所有 token 都经过 \\
\textbf{Routed 专家} & $n_{\mathrm{routed}}$ 个 SiTU，半宽 $\ell \to \ell$（$\ell = d/2$） \\
\textbf{Latent 接口} & $W_\downarrow: d \to \ell$, $W_\uparrow: \ell \to d$（压缩/还原） \\
\textbf{Router} & $W_r: d \to n_{\mathrm{routed}}$，Top-k 选择 + softmax 归一化 \\
\bottomrule
\end{tabular}
\end{center}

\subsection{计算流（五步）}

\begin{enumerate}[leftmargin=2em]
\item \textbf{Latent 压缩}：
\[
  z = W_\downarrow \cdot x \qquad \shape{B, T, \ell}
\]

\item \textbf{Routing}：
\[
  \mathrm{logits} = W_r \cdot x \qquad \shape{B, T, n_{\mathrm{routed}}}
\]
\[
  \mathrm{ids}, \mathrm{probs} = \mathrm{TopK}(\mathrm{logits}, k)
  \qquad \mathrm{ids}: \shape{B, T, k}, \;\; \mathrm{probs}: \shape{B, T, k}
\]

\item \textbf{Routed 专家}（在 latent 空间 $\ell$ 上）：
\[
  u = \sum_{i \in \mathrm{Top\text{-}k}} p_i \cdot E_i^{\mathrm{rt}}(z)
  \qquad \shape{B, T, \ell}
\]

\item \textbf{Shared 专家}（全宽 $d$）：
\[
  s = \sum_j E_j^{\mathrm{sh}}(x) \qquad \shape{B, T, d}
\]

\item \textbf{合并}：
\[
  y = s + W_\uparrow \cdot \mathrm{RMSNorm}(u) \qquad \shape{B, T, d}
\]
\end{enumerate}

\begin{importantbox}{如果你只记一件事}
Routed 专家只在 $\ell = d/2$ 的 latent 空间操作，
参数量是全宽专家的 $1/4$（$\ell^2$ vs $d^2$）。
Shared 专家保持全宽 $d$，提供基础表达能力。
\end{importantbox}

\subsection{代码对照}

\begin{codemathtop}{layers/latent\_moe.py — LatentMoE.forward}
\begin{lstlisting}
def forward(self, x):                       # [B, T, d]
    z = self.down(x)                        # [B, T, ell]

    logits = self.router(x)                 # [B, T, n_routed]
    topk = torch.topk(logits, self.top_k, dim=-1)
    ids = topk.indices                      # [B, T, k]
    probs = F.softmax(topk.values, dim=-1)  # [B, T, k]

    # All expert outputs (vectorized)
    all_out = stack([e(z) for e in self.experts])  # [R, B, T, ell]
    # Gather top-k and weighted sum
    u = zeros(B, T, ell)
    for i in range(self.top_k):
        idx = ids[:,:,i].reshape(B*T)
        sel = all_out[arange, idx]
        u += probs[:,:,i:i+1] * sel.reshape(B, T, ell)

    shared_out = stack([e(x) for e in self.shared]).sum(0)  # [B, T, d]
    return shared_out + self.up(self.norm(u))                # [B, T, d]
\end{lstlisting}
\end{codemathtop}

\begin{warningbox}{为什么 router 用 $x$（全宽）而不是 $z$（latent）？}
路由需要看到 token 的完整表示才能做好选择。
如果用 $z$ 路由，压缩过程可能丢失路由需要的信息。
K3 论文里也是用全宽 $x$ 做 Top-k，然后在 $\ell$ 空间计算。
\end{warningbox}

\subsection{形状总览}

\begin{center}
\begin{tabular}{llll}
\toprule
变量 & 形状 & 说明 \\
\midrule
$x$ & \shape{B, T, d} & 输入 \\
$z$ & \shape{B, T, \ell} & latent（$\ell = d/2$）\\
logits & \shape{B, T, n_r} & router 输出 \\
ids & \shape{B, T, k} & Top-k 专家索引 \\
probs & \shape{B, T, k} & Top-k softmax 权重 \\
\texttt{all\_out} & \shape{n_r, B, T, \ell} & 所有 routed 专家输出 \\
$u$ & \shape{B, T, \ell} & 加权求和后的 routed 输出 \\
\texttt{shared\_out} & \shape{B, T, d} & shared 专家求和 \\
$y$ & \shape{B, T, d} & 最终输出 \\
\bottomrule
\end{tabular}
\end{center}

\subsection{本章小结}

LatentMoE 把 routed 专家限制在 $\ell = d/2$ 的 latent 空间，省参数。
SiTU-GLU 给 gate 和 up 加 $\tanh$ 软上限（$\beta_1=4, \beta_2=25$），
防止低精度溢出。Shared 专家全宽，提供基础能力；routed 专家通过 Top-k 路由提供专业化能力。