查看原文
其他

田渊栋新作:打开1层Transformer黑盒,注意力机制没那么神秘

新智元 新智元 2023-09-02



  新智元报道  

编辑:LRS
【新智元导读】AI理论再进一步,破解ChatGPT指日可待?


Transformer架构已经横扫了包括自然语言处理、计算机视觉、语音、多模态等多个领域,不过目前只是实验效果非常惊艳,对Transformer工作原理的相关研究仍然十分有限。


其中最大谜团在于,Transformer为什么仅依靠一个「简单的预测损失」就能从梯度训练动态(gradient training dynamics)中涌现出高效的表征?


最近田渊栋博士公布了团队的最新研究成果,以数学严格方式,分析了1层Transformer(一个自注意力层加一个解码器层)在下一个token预测任务上的SGD训练动态。


论文链接:https://arxiv.org/abs/2305.16380


这篇论文打开了自注意力层如何组合输入token动态过程的黑盒子,并揭示了潜在的归纳偏见的性质。


具体来说,在没有位置编码、长输入序列、以及解码器层比自注意力层学习更快的假设下,研究人员证明了自注意力就是一个判别式扫描算法(discriminative scanning algorithm):


从均匀分布的注意力(uniform attention)开始,对于要预测的特定下一个token,模型逐渐关注不同的key token,而较少关注那些出现在多个next token窗口中的常见token


对于不同的token,模型会逐渐降低注意力权重,遵循训练集中的key token和query token之间从低到高共现的顺序。


有趣的是,这个过程不会导致赢家通吃,而是由两层学习率控制的相变而减速,最后变成(几乎)固定的token组合,在合成和真实世界的数据上也验证了这种动态。


田渊栋博士是Meta人工智能研究院研究员、研究经理,围棋AI项目负责人,其研究方向为深度增强学习及其在游戏中的应用,以及深度学习模型的理论分析。先后于2005年及2008年获得上海交通大学本硕学位,2013年获得美国卡耐基梅隆大学机器人研究所博士学位。



曾获得2013年国际计算机视觉大会(ICCV)马尔奖提名(Marr Prize Honorable Mentions),ICML2021杰出论文荣誉提名奖。


曾在博士毕业后发布《博士五年总结》系列,从研究方向选择、阅读积累、时间管理、工作态度、收入和可持续的职业发展等方面对博士生涯总结心得和体会。


揭秘1层Transformer


基于Transformer架构的预训练模型通常只包括非常简单的监督任务,比如预测下一个单词、填空等,但却可以为下游任务提供非常丰富的表征,实在是令人费解。


之前的工作虽然已经证明了Transformer本质上就是一个通用近似器(universal approximator),但之前常用的机器学习模型,比如kNN、核SVM、多层感知机等其实也是通用近似器,这种理论无法解释这两类模型在性能上的巨大差距。



研究人员认为,了解Transformer的训练动态(training dynamics)是很重要的,也就是说,在训练过程中,可学习参数是如何随时间变化的。


文章首先以严谨数学定义的方式,形式化描述了1层无位置编码Transformer的SGD在下一个token预测(GPT系列模型常用的训练范式)上的训练动态。


1层的Transformer包含一个softmax自注意力层和预测下一个token的解码器层。



在假设序列很长,而且解码器的学习速度比自注意力层快的情况下,证明了训练期间自注意力的动态行为:


1. 频率偏差Frequency Bias


模型会逐渐关注那些与query token大量共现的key token,而对那些共现较少的token降低注意力。


2. 判别偏差Discrimitive Bias


模型更关注那些在下一个要预测的token中唯一出现的独特token,而对那些在多个下一个token中出现的通用token失去兴趣。


这两个特性表明,自注意力隐式地运行着一种判别式扫描(discriminative scanning)的算法,并存在归纳偏差(inductive bias),即偏向于经常与query token共同出现的独特的key token


此外,虽然自注意力层在训练过程中趋向于变得更加稀疏,但正如频率偏差所暗示的,模型因为训练动态中的相变(phase transition),所以不会崩溃为独热(one hot)。



学习的最后阶段并没有收敛到任何梯度为零的鞍点,而是进入了一个注意力变化缓慢的区域(即随时间变化的对数),并出现参数冻结和学会(learned)。


研究结果进一步表明,相变的开始是由学习率控制的:大的学习率会产生稀疏的注意力模式,而在固定的自注意力学习率下,大的解码器学习率会导致更快的相变和密集的注意力模式。


研究人员将工作中发现的SGD动态命名为扫描(scan)和snap:


扫描阶段自注意力集中在key tokens上,即不同的、经常与下一个预测token同时出现的token;其他所有token的注意力都下降。


snap阶段注意力全中几乎冻结,token组合固定。



这一现象在简单的真实世界数据实验中也得到验证,使用SGD在WikiText上训练的1层和3层Transformer的最低自注意力层进行观察,可以发现即使在整个训练过程中学习率保持不变,注意力也会在训练过程中的某一时刻冻结,并变得稀疏。

参考资料:
https://arxiv.org/abs/2305.16380




您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存