长上下文 Transformer 的规模极限

我们研究了固定查询和来自维度为 $n$ 的 i.i.d. 随机上下文上的 softmax 自注意力的长上下文极限, 将逆温度 $\beta_n$ 作为决定注意力是退化为均匀平均还是坍缩到最近单个 key 的缩放参数。我们证明, 选择性临界尺度由距离查询分布在零处的局部指数决定而非由上下文的全局特征决定, 且对均匀分布于 $\mathbb{S}^{d-1}$ 上的 key, 其尺度为 $\beta_n^* \asymp n^{2/(d-1)}$。此外, 我们描述了有序注意力权重和注意力输出在所有 $\beta_n$ regimes 中的极限分布: 在亚临界 regime 中, 输出简化为围绕 $q$ 的局部平均, 伴随显式的确定性偏差和高斯波动; 在临界 regime 中, 少数最近 key 保持宏观质量, 但无单个 key 崩溃; 在超临界 regime 中, 所有质量集中在最近的 key 上。值得注意的是, 在价值矩阵为 identity 的亚临界情况下, 注意力图大约实现了逆热方程。