Nvidia MLP 维度与 PyTorch 广播：从矩阵形状到参数量

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 训练例设输入 X 为 [B,D]、隐藏宽度 H、输出宽度 O、两层全连接和逐元素激活；题目未明确网络结构。
• 数学记号用 W1[D,H]、W2[H,O] 右乘；PyTorch `nn.Linear(in,out)` 的实际 weight 存储形状为 [out,in]，前向执行 x @ weight.T + bias。
• bias 形状分别为 [H]、[O]，按末维广播到每个 batch 行；不把广播当作每行独立可学习参数。

解题路线

1. 从最后一维开始对齐：X[B,D] @ W1[D,H] 得 [B,H]，加 b1[H] 仍为 [B,H]，激活不改形状；再 @ W2[H,O]、加 b2[O] 得 [B,O]。
2. 参数量为 DH+H+HO+O。一次前向的主要乘加数约 B(DH+HO)；若训练还要反向传播与保留中间激活，计算和内存会增加。
3. PyTorch 广播从尾维向前比：维度相等、任一为 1 或缺省才可兼容。[B,H]+[H] 正确；[B,H]+[B] 通常不兼容，若 B=H 却恰好兼容则可能沿错误的轴静默相加，应显式 reshape 成 [B,1] 才是每样本偏置。
4. 验证时先打印/断言每一步形状，再用 B≠H 的小输入测试，避免方形维度偶然掩盖轴错；还要确认 `nn.Linear` 权重存储方向与手写矩阵式的转置关系。

手动走一遍

训练例 B=2、D=3、H=4、O=2：X[2,3]→隐藏层[2,4]→输出[2,2]；参数量 3×4+4+4×2+2=26，前向主乘加数约 2×(12+8)=40。b1[4] 会复制语义地加到两行，而不是产生 8 个独立 bias 参数。

规模与取舍

前向主运算 O(BDH+BHO)，权重 O(DH+HO)、bias O(H+O)，训练需另存中间激活约 O(BH+BO) 及反向所需状态。实际 PyTorch 中广播通常以视图/stride 语义实现，不必显式复制 bias 张量；计算仍对每个输出元素执行加法。

易错点与追问

容易失分的边界

• 把 PyTorch Linear 的 weight 写成 [in,out]，却又在代码里按 `x @ weight.T` 使用，造成维度不匹配。
• 把 batch 维 B 与特征维 H 混淆；B=H 的测试可能让错误广播静默成功。
• 把 bias 广播理解成每个样本都有独立可训练偏置，误算参数量为 BH。
• 只报前向 FLOPs，不说明反向、激活缓存和大 batch 的显存需求。

继续深挖

• 若输入是 [B,T,D]，Linear 作用于最后一维，输出变 [B,T,H]；bias [H] 沿 B、T 广播。
• 若隐藏层加 residual，必须先核对相加两路的末维是否相同或是否需要 projection。
• 若改成门控 MLP，两支投影的参数量和激活形状都要重新算，不能沿用两层普通 MLP 公式。