©PaperWeekly 原创 · 作者 | 苏剑林
单位 | 科学空间
研究方向 | NLP、神经网络
近年来,线性 RNN 由于其可并行训练以及常数推理成本等特性,吸引了一定研究人员的关注(例如笔者之前写的《Google新作试图“复活”RNN:RNN能否再次辉煌?》),这让 RNN 在 Transformer 遍地开花的潮流中仍有“一席之地”。然而,目前看来这“一席之地”只属于线性 RNN,因为非线性 RNN 无法高效地并行训练,所以在架构之争中是“心有余而力不足”。
不过,一篇名为《Parallelizing Non-Linear Sequential Models over the Sequence Length》的论文有不同的看法,它提出了一种迭代算法,宣传可以实现非线性 RNN 的并行训练!真有如此神奇?接下来我们一探究竟。
论文标题:
Parallelizing Non-Linear Sequential Models over the Sequence Length
https://arxiv.org/pdf/2309.12252.pdf
求不动点
原论文对其方法做了非常一般的介绍,而且其侧重点是 PDE 和 ODE,这里我们直接从 RNN 入手。考虑常见的简单非线性 RNN:
由于 的存在,它只能串行计算。现在我们在两边都减去 :
从给定 出发,反复迭代上式,理想情况下,它会收敛于一个不动点 ,这就是原来非线性 RNN 的计算结果。当然,理论上通过式(3)迭代的总计算量是比直接通过式(1)递归计算要大的,但由于每一步迭代都是可并行的线性 RNN,并且如果收敛速度比较快时迭代步数不需要太多,那么总的耗时通常都会快于直接非线性 RNN 递归(尤其是序列长度很大时)。
事实上,非线性 RNN 之所以慢,无法并行计算还是次要的,最关键是它包含了大量的非 element-wise 运算,比如式(1)的 里边的矩阵运算 ;而线性 RNN 之所以快,除了它允许并行训练之外,更关键的是它能通过对角化来将矩阵乘法变换为 element-wise 的乘法——对于 element-wise 乘法来说,即便是串行计算也不会太慢。当我们通过式(3)将非线性 RNN 转为线性 RNN 的迭代之后,同样享受线性 RNN 可对角化的“待遇”,从而提高计算速度。具体来说,在复数域中将 A 对角化为 ,那么式(3)变为由于 RNN 之后一般都还要接个投影层,所以 的 P 原则上可以合并到外接的投影层里边,也就是说,上式理论上具备跟原来的(1)具备同等的表达能力,但由于 是对角阵,递归的计算量会明显降低。上式还出现了逆矩阵 ,不单计算量大,而且不利于优化,所以我们可以干脆将 和 换成两个不相关的参数矩阵:
摄动思想
假定 ,那么式(3)其实就是将原本的非线性 RNN 就分解为一系列线性 RNN:
这正好是(8)中的第一个方程,因此如果假设成立,那么 或许已经足够接近理想的 ,后面的每一步迭代都在快速逼近它。从这里我们可以看出,“两边同时减去 ”是关键之处,这使得(3)的第一步迭代就接近于原本非线性 RNN 的一阶线性近似,这可以提高收敛速度,也是数学物理中的经典操作,名曰“摄动”。根据摄动法的思想,提高收敛速度的关键就是提高近似展开的精度,比如较为简单的改进是只假设 是小量,那么根据一阶泰勒展开有(将 视为列向量,这里的 是 Hadamard 积分)其中 。更精细的改进是在每一步迭代时,都在前一步迭代结果的基础上进行展开:理论上来说,(11)、(13)两个改进确实能提高收敛速度,然而它们使得每一步线性递归的矩阵 A 变得跟 t 甚至 n 有关了,这其实会大大增加并行的复杂度,也不能利用“简化形式”一节的对角化技巧来加速。另一方面,如果保持(3)这样的迭代格式,虽然有诸多效率上的好处,但收敛方面确实无法得到很好的保障。难道这两者的矛盾就无法调和了吗?事实上,按照笔者的观点,最直接的做法是“别去管它”——借助非线性 RNN 导出了(3)后,就忘记原本的非线性 RNN,将式(3)作为基本模型。也就是说,何必忧虑式(3)会不会收敛到原来的非线性 RNN?直接将它作为新的出发点不好吗?梯度下降学到怎样的结果就是怎样的结果,如果梯度下降学到的结果是不收敛到原来的非线性 RNN,那么就意味着不收敛到原来的 RNN 是更适合的。抛开这一层思维束缚后,其实很多问题会变得豁然开朗起来。首先,即便是式(13)在理论上拥有非常好的收敛速度,但也是有条件的,而且在深度学习的背景下,要保证这些条件会显得很奢侈。换言之,即便是式(13)的收敛性也没有绝对保证,所以何必“五十步笑百步”去苛责式(3)?其次,将式(3)视为新的出发点后,我们可以将它单纯地理解为线性 RNN 的一种新用法,或者说解决线性 RNN 缺陷(比如线性 RNN 不是图灵完备的)的一个思路,这样操作性更强。总的来说,不去管它的收敛性,似乎更能打破思维僵局,探索更一般的结果。一般情形
前面的“长篇大论”,都只围绕着简单的非线性 RNN 也就是式(1)进行讨论,对于更常用的 LSTM、GRU,结果又如何呢?初始阶段,所有门控都可以近似视为 ,那么模仿式(9)有总的来说,这种将非线性 RNN 变为线性 RNN 迭代的转换,从实践的角度来看,就是以非线性 RNN 为引,导出一种多层线性 RNN 的参数共享和组合方法,迭代了几次,那么就有几层线性 RNN 的计算量。这样自然而言就引发了一个思考:除非可以证明 GRU、LSTM 等非线性 RNN 有绝对的优势,否则直接叠加几层“线性 RNN+MLP”不好吗?文章小结
本文简单探讨了非线性 RNN 的并行计算问题——通过数学物理中的“摄动”思想,我们可以将非线性 RNN 转化为线性 RNN 的迭代,从而利用线性 RNN 的可并行性来实现非线性 RNN 的并行。
如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。
总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。
📝 稿件基本要求:
• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注
• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题
• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道:
• 投稿邮箱:[email protected]
• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿
△长按添加PaperWeekly小编
🔍
现在,在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧