TikTok Transformer FFN：逐 Token 非线性与 Attention 的分工

题意与假设

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

• 训练传统两层FFN x→activation(xW1+b1)W2+b2，不把现代门控变体当唯一结构。
• block的pre/post norm与激活需按模型合同，不预设所有Transformer一致。

解题路线

1. attention根据上下文混合不同token的value；FFN对每个token表示单独做通道变换，参数跨位置共享。它不直接在位置之间连边，但输入已经可包含上下文。
2. 两层线性之间加入非线性，通常从dmodel扩展到dff再投影回dmodel，为表示增加非线性组合能力；若无非线性，两层可合成单个线性映射，表达作用不同。
3. FFN与residual维度要闭合：输出回到dmodel才能与输入相加。pre/post norm影响训练梯度路径，不能把norm位置随便挪却说同一模型。
4. 容量与算力取舍按dff估计，每层参数约2dmodel·dff加bias；很多模型FFN占显著参数/FLOPs，不能只优化attention便忽略它。
5. GELU/ReLU/SwiGLU等激活/门控是选择，门控常多一份投影，参数与中间张量成本也变。MoE把部分FFN换成路由专家，又引入负载均衡/通信，不是免费扩容。
6. 面试写伪代码时标明[B,T,dmodel]→[B,T,dff]→[B,T,dmodel]，验证无额外token轴混合与广播错误。理论说“逐token”不等于每个token独立于上下文，因为attention已在前层混合。

手动走一遍

两个token各自乘同一W1/W2，shape保留2个位置。注意力可以把第二个token信息带到第一个，随后FFN加工这份上下文表示；因此不是“FFN看不到其他token信息”。

规模与取舍

两层FFN计算O(BT·dmodel·dff)，参数O(dmodel·dff)，激活O(BT·dff)；门控/专家另计。

易错点与追问

容易失分的边界

• 说FFN直接混合不同位置。
• 去掉激活仍说相同表达力。
• 扩展后不投回维度却做残差。

继续深挖

• 门控FFN怎样改参数量？
• pre/post norm怎样影响训练？
• MoE路由的额外成本是什么？