求解PaLM论文中稠密自注意力单Token FLOPS公式6LH(2QT)的推导逻辑

首先明确几个核心前提：

• FLOPS统计规则：矩阵乘法中，m×k矩阵乘k×n矩阵的总FLOPS≈2mkn（乘法mkn次+加法≈mk*n次，合计近似为2倍乘法次数）。
• PaLM的MFU计算包含前向传播+反向传播的总运算量，针对的是生成式任务中每个新增Token的计算量。
• 符号定义：L=层数，H=头数，Q=单头维度，T=当前序列长度。

1. 单头单层前向传播（单个新增Token） #

稠密自注意力的核心矩阵乘法有两步：

• Query-Key乘积：新增Token的Query是1×Q的向量，已有的Key矩阵转置后是Q×T。运算量为 2*1*Q*T = 2QT。
• 注意力分数-Value乘积：得到的1×T注意力分数，乘以T×Q的Value矩阵，得到最终的1×Q输出向量。运算量同样为 2*1*T*Q = 2QT。
  单头单层前向总FLOPS：2QT + 2QT = 4QT = 2*(2QT)。

2. 单头单层反向传播（单个新增Token） #

反向传播的运算量通常是对应前向传播的2倍（需同时计算输入梯度和参数梯度）：
单头单层反向总FLOPS：4QT * 2 = 8QT = 4*(2QT)。

3. 总运算量（所有层+所有头） #

把单头单层的前向+反向运算量汇总，再乘以层数L和头数H：
总FLOPS = (2*(2QT) + 4*(2QT)) * L * H = 6LH(2QT)，完全匹配论文中的公式。

推导偏差的原因 #

你之前的推导漏了前向传播中注意力分数与Value矩阵相乘这一步的运算量，且错误估算了反向传播的倍数：反向传播的运算量是前向的2倍，而非你认为的4倍，最终导致结果不符。

内容的提问来源于stack exchange，提问作者cangozpi