Transformer朴素采样时，生成 T 个 token 需要的 FLOPs 计算推导过程

Transformer 的朴素采样

在生成每个token时，都需要将整个历史序列输入到Transformer中。

假设你有一个 transformer 模型，输入 prompt（never gonna give you），模型会给出下一个 token 在词汇表上的 logits 分布，然后从中采样。一旦得到 token（up,），就把它加到 prompt 后面。再输入给 transformer，接着模型给出分布，采样得到 token（never）。依此类推。

缺点：时间复杂度很高。 生成T个token需要O(T³)的浮点运算次数(FLOPs)（因为，每次前馈传递的复杂度是O(T²)，T个 token 要进行 T 次前馈传递，复杂度之和接近O(T³)）

所以，这种方法效率低下，因为每次生成新token时都会重复计算历史token的表示。

推理过程

要理解朴素采样中“生成T个token需要O(T³) FLOPs”的计算逻辑，需从Transformer的前馈传递复杂度和生成过程的序列长度变化两方面拆解，核心是自注意力机制的计算量随序列长度的增长规律。

一、前提：生成过程的序列长度变化

在朴素采样中，生成第$k$个token时，输入Transformer的序列长度为$k$（假设初始prompt长度为0，仅考虑生成的token；若包含初始prompt，逻辑相同，只是基数增加）：

• 生成第1个token：输入序列长度$n=1$（空序列或初始prompt，简化为1）；
• 生成第2个token：输入序列长度$n=2$（第1个token+新增的输入）；
• …
• 生成第$T$个token：输入序列长度$n=T$（前$T-1$个token+新增输入）。

二、关键：Transformer单次前馈传递的FLOPs（以自注意力为主）

Transformer的前馈传递中，自注意力机制是计算量最大的部分，其复杂度主导了整体FLOPs。对于长度为$n$的序列，自注意力的核心计算步骤及复杂度如下：

1. 自注意力的核心计算（简化版）

自注意力的核心是计算“每个位置对所有位置的注意力权重”，具体步骤：

• 生成Query（Q）、Key（K）、Value（V）：每个矩阵维度为$n\times d$（$n$是序列长度，$d$是隐藏维度，如768），由输入序列通过线性层得到，复杂度为$O(n{d}^2)$（次要项，$d$是固定常数）；
• 计算注意力分数（$Q{K}^T$）：Q是$n\times d$，K是$n\times d$，$Q{K}^T$的结果是$n\times n$的矩阵，每个元素需要$d$次乘法和$d-1$次加法，总复杂度为$O(n^{2}d)$（主导项）；
• 计算softmax和与V相乘：softmax对$n\times n$矩阵操作，复杂度$O(n^2)$；与V（$n\times d$）相乘的结果是$n\times d$，复杂度$O(n^{2}d)$（主导项）。

2. 单次前馈传递的主导复杂度

忽略常数项（如$d$固定，可视为常数），自注意力的主导复杂度为$O(n^2)$。
其他模块（如前馈网络FFN）的复杂度为$O(n{d}^2)$，因$d$是固定值（如768），当$n$（序列长度）增大时，$n^2$的增长远快于$n{d}^2$，因此单次前馈传递的整体复杂度可近似为$O(n^2)$。

三、总FLOPs：T步生成的累加复杂度

生成$T$个token的过程中，每一步的序列长度为$n=1,2,...,T$，每步的前馈传递复杂度为$O(n^2)$。因此总FLOPs是各步复杂度的总和：

$$\text{总FLOPs}\approx \sum _{n=1}^{T}O(n^2)$$

根据数学公式，平方和的近似结果为：
$$\sum _{n=1}^T{n}^2=\frac{T(T+1)(2T+1)}{6}\approx O(T^3)$$

总结：为什么是O(T³)？

• 生成第$k$个token时，序列长度为$k$，单次前馈传递复杂度为$O(k^2)$（自注意力主导）；
• 生成$T$个token的总复杂度是$1^2+2^2+...+T^2$的累加，数学上近似为$O(T^3)$。

这就是朴素采样中“生成T个token需要O(T³) FLOPs”的核心逻辑——序列长度随生成步骤线性增长，而每步的计算量随序列长度的平方增长，最终导致总复杂度为三次方级。