[论文解读] Transformers learn through gradual rank increase

Enric Boix-Adserà, Etai Littwin+3 more|arXiv (Cornell University)|Jun 12, 2023

Neural Networks and Applications被引用 5

一句话总结

本文展示了视觉与语言Transformer在训练过程中表现出渐进式学习动态，其中训练后与初始权重之差的秩以离散阶段逐步增加。理论上，在小初始化和对角权重假设下，梯度流训练导致每阶段最多增加一秩；实验上，即使在无简化假设的情况下，该现象在真实世界视觉与语言模型中依然存在。

ABSTRACT

We identify incremental learning dynamics in transformers, where the difference between trained and initial weights progressively increases in rank. We rigorously prove this occurs under the simplifying assumptions of diagonal weight matrices and small initialization. Our experiments support the theory and also show that phenomenon can occur in practice without the simplifying assumptions.

研究动机与目标

• 探究Transformer是否表现出与简单神经网络类似的渐进式学习动态。
• 在小初始化和对角权重矩阵的简化假设下，为Transformer中的渐进式学习建立理论基础。
• 在无简化假设的真实设置中，通过实验验证理论发现。
• 探索所观察到的动力学与高效微调方法（如LoRA）之间的联系。

提出的方法

• 理论分析针对一种简化Transformer变体，其注意力头权重为对角矩阵，将问题简化为依赖于两个权重向量逐元素乘积的网络。
• 证明了一个适用于形式为 $ f_{\mathsf{NN}}(\mathbf{x};\mathbf{u},\mathbf{v}) = h(\mathbf{x}; \mathbf{u} \odot \mathbf{v}) $ 的网络的一般定理，表明在小初始化下，训练以离散阶段进行，每阶段最多在 $ \mathbf{u} \odot \mathbf{v} $ 中新增一个非可忽略的条目。
• 通过将乘积 $ \mathbf{W}_K \mathbf{W}_Q^\top $ 表示为对角矩阵乘积，将理论结果应用于注意力头，从而分析有效权重矩阵中的秩增加。
• 实验比较了初始化和训练后 $ \mathbf{W}_K \mathbf{W}_Q^\top $ 的谱特性，通过测量权重差的秩来追踪渐进变化。
• 实验在CIFAR-10和ImageNet上训练的ViT模型上进行，分析注意力头权重，不假设其为对角或小初始化。
• 分析包括对训练迭代过程中归一化谱的可视化，以说明分阶段的秩增长。

实验结果

研究问题

• RQ1 Transformer是否表现出渐进式学习动态，即所学函数的复杂度以阶段方式逐步增加？
• RQ2 在小初始化和对角权重矩阵等简化假设下，Transformer中的渐进式学习动态能否被严格证明？
• RQ3 当理论假设被违反时，权重差异中的渐进秩增加现象是否仍在实践中持续存在？
• RQ4 所观察到的渐进动力学与高效微调方法（如LoRA）之间存在何种关系？

主要发现

• 在小初始化和对角权重假设下，Transformer中的梯度流训练以离散阶段进行，每个阶段持续 $ \Theta(\log(1/\alpha)) $ 时间，每阶段最多在有效权重向量 $ \mathbf{u} \odot \mathbf{v} $ 中新增一个非可忽略条目。
• 理论分析证明，由于鞍点逃逸动力学，训练后与初始权重之差的秩以每阶段最多一秩的速度渐进增加。
• 在ImageNet上训练的ViT模型实验中，$ \mathbf{W}_K \mathbf{W}_Q^\top $ 矩阵的训练后与初始值之差的秩随训练逐渐增加，谱图中清晰显示出分阶段行为。
• 即使在无理论假设的情况下（如标准训练中非对角权重和较大初始化），在视觉与语言模型中仍观察到渐进秩增加。
• 该现象在ImageNet等更大数据集上比在CIFAR-10等较小数据集上更显著，表明数据复杂度影响分阶段动力学。
• 结果提示与LoRA微调存在关联，因为两者均涉及低秩更新，所观察到的渐进动力学或可解释此类方法为何有效。