LSTM之父Jürgen Schmidhuber新作:一种方法,超越线性Transformers
Transformers 已成为处理文本数据的最流行的神经网络之一。它在神经机器翻译上的成功迅速转移到自然语言处理的其他任务上,如语言建模或问答。最近,它也广泛活跃在其他领域,包括图像处理或数学问题求解等等。
尽管 Transformers 突破了 RNN 不能并行化的限制,其不容忽视的缺点仍是在时间复杂度和空间复杂度上是序列长度的平方级。此外,在自回归版本中,状态大小( state size )随序列长度呈线性增加的趋势。这使得 Transformers 不适用于处理非常长或潜在无限序列的自动回归设置,从而迫使研究者截断时间上下文( temporal contexts )并忽略超出固定大小的时间窗口的长期依赖关系。
最近的工作正试图解决这个问题,但提及在部分可观察环境中的 RL 时, Transformers 依旧极具有挑战性,这仍然是诸如 LSTM 之类的处理序列数据的RNN所主导的。
为了将 Transformers 扩展到更长的序列,近期的工作已提出将 self-attention 计算中的 softmax 线性化,并以顺序的方式重新组织后者。
由“ LSTM ”之父 Jürgen Schmidhuber 、瑞士人工智能实验室 IDSIA 撰写的论文 Going Beyond Linear Transformers with Recurrent Fast Weight Programmers ,介绍了采用“快速加权编程器”( FWPs , Fast Weight Programmers )的思想,改进线性 Transformer 变体的方法。
最近的调查研究表明,线性 Transformer 本质上相当于 90 年代基于外积的 FWPs 。团队成员在本文将进一步探索这种联系,以描述更强大的 FWPs 。
团队成员通过在语言建模数据集 Wikitext-103 上的实验证明, RFWPs 与常规 Transformers 相比具有明显的竞争力。
其次,团队在代码执行和顺序 ListOps 两个合成算法任务上研究了所提出模型的各种性质。最后,证实模型完全可以替代 LSTMs ,直接应用到 RL 问题中。
除此之外,文章提出的 RFWPs 在许多 Atari 2600 2D 游戏环境中也达到了比 LSTM 基线更大的改进。
简而言之,此工作主要有以下两个贡献:
(1)从 FWPs 的角度出发,研究了用于序列处理的新型强大的 FWPs ,证明了神经网络可以很容易地学会控制比单个前馈层更复杂的神经网络;
(2)从 Transformer 模型的角度出发, RFWPs 通过循环的方式扩展线性 Transformers ,解决现有自回归 Transformer 模型的一般限制。
FWPs 背景知识
我们简单回顾下 FWPs 的一般概念及其两个具体实例:线性 Transformers 和 Delta Net 。
FWPs公式化表示如下:
其中,从先前的迭代中获取快速权值 ,以产生新的基于 条件下的快速权值 。从本质上看,更新规则为FWP使用的可微的基本编程指令。
FWP 的具体实例:一般情况下,快速权值 的维数太大,神经网络无法方便的对其进行参数化。相反地,它在 1991 年被提出通过慢网产生的两个向量的外积来执行秩1( rank-one )更新。近期的两个模型直接对应于这种基于外积的 FWPs :线性 Transformers 和 Delta Net 。
线性 Transformers :“ linear Transformers ”是 Transformers 家族中的一种,其注意力中的 softmax 是线性化的。这通过用核函数(kernel function)替换softmax来实现,然后self-attention可被重写( rewritten)为一个基本的基于外积的FWP。先前的研究工作主要聚焦在不同的带有属性的 映射上,如增加容量或在极限中保证softmax的近似值。对目标而言,的特定选择是无关紧要的,需要做的是设定 :。用 ,代替 , ,以简化下面的公式。对于新输入 ,输出 为:
其中, 。
Delta Net :Delta Net 是通过将线性 Transformer 中的纯加法编程指令替换为类似于 Delta 规则的指令来获得的。公式如下:
其中 是权值为 、激活函数为sigmoid的慢网生成的更新规则的快速参数:
而 是由之前的快速权值 和键 构造而来:
引入该更新规则是为了解决使用纯加法更新规则影响线性 Transformers 的内存容量问题。对应的 , 公式可以表示为:
其中,
因此,与线性 Transformers 不同, Delta Net 中的慢网将之前的快速权值 考虑在内,以生成新的快速权值更新。
快速加权编程器
在本节中,研究团队提出了快网和慢网两种条件下的 RFWPs 。具体内容如下:
快网扩展( Fast Network Extensions )
原则上,任何神经网络架构都可以快速构建。它的快速权值版本是通过另一个慢网参数化的快速权值代替网络权值而得到的。比如,两个权重矩阵 和 的常规RNN层:
用 和 替换 和 可以得到快速权值版本。
虽然这个观点说明了 FWP 的普遍性,但团队成员所采用模型的角度略有不同:引入循环作为一种增强现有线性 Transformers 的方法。
Delta RNN | 通过在线性 Transformers 的前馈快速网络中增加一个循环项,得到一种称为 Delta RNN 的快速权值 RNN 。
其中, 作为引入循环连接的一种快速权值矩阵。与 类似,它也是由慢网使用 Delta 更新规则生成的,只不过有附加的慢速权值。
替代前馈快网 | 虽然这项工作的重点是 RNN ,但在前面公式中也有一些有趣的快速前馈模型,这可能会导致更强的前馈基线。例如,考虑使用 K 层的深度网络以替代公式中的单层快网。
另一种有趣的方法是使用 Delta Net 自身作为快网,即令 Delta 网络中慢速权重变快(得到 Delta Net )。从原则上讲,这样的模型可以学习适应根据上下文快速生成权值的方法。
慢网扩展( Slow Network Extensions )
RDN 是通过修改键、值和查询向量的生成以及 Delta Net 的学习率获得的。研究团队为循环连接添加了额外了慢速权重,以对先前的快网输出进行有效连接。
其中, 。
与常规 RNNs 相比,所有这些快速权值变体都保持标准的一维状态和二维快速权值状态,这两种状态都作为额外的内存。在后面的实验部分证明了慢网可以学习控制这些相当复杂的变体的权值,而在快网中循环的加入便引入了一个关键特性,而这正是常规 Transformers 中所缺少的。
实验及分析
研究团队选择在四种不同的环境下进行实验。
首先是评估模型在语言建模任务中的表现,以获得其性能并讨论计算成本。语言建模是评估序列模型的一个很好的任务。
其次,为了突出模型的功能特性,团队考虑在算法任务上同样对模型进行评估。
事实上,众所周知, RNNs 的实际能力因架构而异。团队成员对此非常感兴趣,考虑到这一目标,团队在两个合成算法任务上进行实验,即代码执行和顺序 ListOps ,旨在比较模型的基本序列处理能力。最后,团队将模型应用于 2D 环境中的RL,替代 LSTM 。
( 1 )语言建模:
选择 Wikitext-103 数据集,遵循小型模型设置,在语言任务中从性能概述和计算效率等方面评估所有模型。
在评价指标 1 ——困惑度( Perplexity )上,具体结果如表 1 所示,该表概述了不同的模型、不同的慢网&快网架构以及不同的更新规则选择。
虽然语言建模作为模型的合理性验证是有用的(例如,除了 Delta Net 之外,所有模型都实现了合理化性能),但这项任务过于一般化,无法识别模型的某些重要特征。
在评价指标 2 ——计算效率( Computational efficiency )上,前文描述的扩展修改究竟会给线性 Transformers / FWPs 带来哪些额外的计算成本呢?
首先声明,线性 Transformers 的核心复杂度不会受影响:它们都具有恒定的空间复杂度和线性的时间复杂度。然而,每个时间步的计算成本在不同的模型之间却有明显差异。它们都具有常数空间和线性时间复杂度,即序列长度。然而,每个时间步的计算成本在不同的模型之间有很大的不同,这里用具体的实现以量化模型的训练速度。所有模型都使用自定义 CUDA 内核实现,除了使用常规 PyTorch 代码的基线 Transformer 。表1中 LT 和 Delta Net 的训练速度分别为每秒 66k 和 63k words ( vs. 33k 的基线 Transformer )。成本最昂贵的模型是 Delta LSTM ,这个快速 LSTM 网络有 6 个权重矩阵,每个都由单独的Delta规则操作,其相应速度是每秒 14k words ,对于扩展到更多的实验来说太慢了。相比之下, Delta RNN 和 RDN 的速度更合理一些:分别为每秒 41 K 和 35 K words 。因此,考虑到性能和计算成本两方面的前景和实用性,后续的实验将集中在这两种循环体系结构上。
( 2 )代码执行任务:学习维护和更新变量状态。
在该任务中,训练模型依次读取 word-level文本提供的输入代码,并预测相应地代码执行结果,任务设置为一个条件和三个基本陈述。
表 2 的左侧显示了结果。团队成员通过修改变量数量( 3 或 5 )来控制任务难度。需要注意的是,该模型体系结构是固定的:LSTM 只有一个具有 256 个节点的层,所有 Transformer 变体具有相同的架构, 4 个层,隐藏层大小为 256 ,使用 16 heads ,内部前馈层大小为1024。
从表 2 中很容易注意到 LSTM 在这两个难度级别上均取得最佳性能,而且随着变量数量的增加,性能下降幅度也最小。与前面的声明相反,LSTM 显然能够将多个变量的值存储在单个向量中。
有了这三个变量,常规 Transformer 的性能已经大大低于其他具有可变内存的模型:Delta Net 、 Delta RNN 和 RDN 。
而线性 Transformers 在此任务中完全失败,可能是由于内存容量问题。通过将变量的数量增加到 5 个,基线 Transformers 、 Delta Net 和 RDN 变得不稳定。与 Delta Net 相比, RDN 中引入的循环连接的优势更加明显( 76.3 vs 61.4% )。相比之下,Delta RNN性能保持稳定,并在Transformer变体中达到最佳性能( 85.1% ),显示出了循环的优势之处,尤其是在常规的 RNN 结构中。然而,为了匹配 LSTM 在这个任务上的性能,这些模型需要更多的层进行实验,具体结果请参见原文附录。
( 3 )顺序 ListOps :学习层次结构和计算。
( 4 ) 2D 游戏环境中的RL。
团队成员在 20 个环境中充分评估模型,重复评估 5 次,以报告带有标准偏差的平均分数。图 1 和 2 分别报告了 50 M 和 200 M env.steps 的结果。
图 1 展示了 50 M 交互后 RDN 对于 LSTM 的相对改进情况。总的来说,在 6 个游戏中 LSTM 在 200 M env.steps 上的表现明显优于 RDN ,而在 9 个游戏中结果却恰恰相反。
鉴于上述结果,一个很自然的问题是,在 LSTM 占主导地位的游戏中,较大型的模型能否改善 RDN 呢?
对此,团队专注于以下四款游戏开展实验:Berzerk 、Gopher 和 Seaquest ,具体结果如表 3 所示。很显然,RDN larger 模型减小了与 LSTM 的差距( Berzerk 除外)。这表明,进一步扩展 RDN 很可能与在其他领域扩展常规 Transformers 一样拥有广阔的发展前景。
论文地址:https://arxiv.org/abs/2106.06295