Shaped Transformer:无限深度与宽度极限下的注意力模型
机器学习
2023-12-12 v2 机器学习
摘要
在深度学习理论中,表征的协方差矩阵可作为检验网络可训练性的代理量。受 Transformer 成功的启发,我们研究了在无限深度与宽度成比例极限下一种改进的基于 Softmax 的带跳跃连接的注意力模型的协方差矩阵。我们证明,在初始化时,该极限分布可由一个以深度-宽度比为指标的随机微分方程(SDE)描述。为获得定义良好的随机极限,Transformer 的注意力机制被修改为将 Softmax 输出以恒等矩阵为中心,并以依赖于宽度的温度参数对 Softmax 对数几率进行缩放。我们通过相应的 SDE 考察网络的稳定性,展示如何借助残差连接优雅地控制漂移与扩散的尺度。稳定 SDE 的存在意味着协方差结构行为良好,即便在极大深度与宽度下亦如此,从而避免了深度注意力模型中臭名昭著的秩退化问题。最后,我们通过模拟表明,该 SDE 对相应的有限规模模型提供了出奇良好的描述。我们将这些架构修改命名为 shaped Transformer。
引用
@article{arxiv.2306.17759,
title = {The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit},
author = {Lorenzo Noci and Chuning Li and Mufan Bill Li and Bobby He and Thomas Hofmann and Chris Maddison and Daniel M. Roy},
journal= {arXiv preprint arXiv:2306.17759},
year = {2023}
}