Mission 1
概念练习
先让每个 token 给所有专家分配概率
🎯 先猜一猜
router_logits.shape=[8,8]，其中第 0 维是 token，第 1 维是专家。
softmax 应该沿哪个维度计算，才能让每个 token 独立给所有专家分配概率？
先猜一个答案，再对照下面的讲解。可以随时修改选择，猜错会得到解释。
先补的知识
- 输入 hidden_states 的 shape 是 [batch_size, seq_len, hidden_size]。B 是句子数，S 是每句 token 数，H 是每个 token 的特征数。
- view(-1, hidden_size) 把 B 和 S 合并成 token 总数 T=B×S；它不改变数据，只把 [B,S,H] 看成 [T,H]。
- gate 是 nn.Linear(hidden_size, num_experts, bias=False)，所以每个 token 的 H 个特征会变成 E=num_experts 个打分。
- logit 是未归一化分数，可以为负，也不要求总和为 1；softmax 才把一行分数变成概率分布。
图解原理
把每个 token 想成一张问诊单，E 位专家都先给它一个匹配分。Router 要沿专家这一维把分数变成概率，因此同一 token 的 E 个概率相加为 1。Notebook 还先把 logits 转成 FP32，因为 softmax 里的指数运算对低精度更敏感。
[B,S,H]句子中的 token 特征
[T,H]T=B×S，排成 token 队列
[T,E]gate 给每位专家打分
[T,E]每行变成概率分布
先认清最后一维
hidden_states[2,4,16]，共 8 个 token
router_logits[8,8]，每行对应 8 位专家
dim=-1在每行的专家列表内归一化
输出 dtypeTODO 先用 FP32，返回前再恢复输入 dtype
一行只属于一个 token
logits
[-1, 2, 0, 1] 转 FP32 softmax
全部为正 一行总和 = 1
[-1, 2, 0, 1] 转 FP32 softmax
全部为正 一行总和 = 1
不能沿 token 维做 softmax，否则不同 token 会互相争概率，Router 就失去“每个 token 自己选专家”的含义。
| 变量 | shape | 最后一维表示什么 | 本步是否改变 shape |
|---|---|---|---|
hidden_states | [T,H] | 隐藏特征 | 已在前面展平 |
router_logits | [T,E] | 专家原始分数 | gate 把 H 改为 E |
routing_probs | [T,E] | 专家全局概率 | softmax 不改 shape |
这里“全局 Softmax”是指对一个 token 的全部 E 位专家归一化，不是把 batch 中所有 token 混在一起归一化。
语法热身：给每位学生的所有兴趣方向分配概率
scores = torch.tensor([
[1.0, 3.0, -1.0],
[2.0, 0.0, 1.0],
])
# 每一行是一位学生，最后一维是 3 个方向
probabilities = F.softmax(scores.float(), dim=-1)
print(probabilities.shape) # [2, 3]
print(probabilities.sum(dim=-1)) # [1, 1]独立例子如何迁移回 Notebook
scores对应 gate 产出的router_logits。3 个方向对应num_experts位专家。probabilities对应 TODO 1 要创建的routing_probs。- 保留
.float()和dim=-1：前者提高 softmax 稳定性，后者明确沿专家维计算。
巩固一下
输入是 [2,4,16]，num_experts=8。展平、gate、softmax 后 routing_probs 的 shape 是什么？
学完这一段，试着做
用一个小动作确认自己理解了；最后再进入官方题目。