RNNS ARE NOT TRANSFORMERS (YET)

2024-03-03

Word count: 2.4k | Reading time≈ 8 min

论文：RNNS ARE NOT TRANSFORMERS (YET) -THE KEY BOTTLENECK ON IN-CONTEXT RETRIEVAL

本文研究了递归神经网络(rnn)和transformer在解决算法问题方面的表示能力差距。理论分析表明，CoT改善了rnn，但不足以缩小与transformer的差距。

我们证明，采用技术来增强 RNN 的上下文检索能力，包括检索增强生成（RAG）和添加单个 Transformer 层，可以使 RNN 能够通过 CoT 解决所有多项式时间可解决的问题，从而缩小了与transformer的代表性差距。

介绍

现代的RNN得到了一系列的发展，包括RWKV、RetNet和Mamba。

过去有许多研究证明了RNN不如Transformrt。比如2024年的《Repeat after me: Transformers are better than state space models at copying》证明常量记忆 RNN 没有足够的表示能力来解决对给定的输入向量子集进行平均的任务（q稀疏平均）和重复输入序列（复制），而存在可以解决这些任务的浅层 Transformer。

然而，上述结果并不排除通过额外的提示技术或微小的架构改变来增强 RNN 可能缩小与 Transformer 差距的可能性。事实上，Transformer 本身也不完美，可能需要在推理时使用额外的技术才能在某些任务上表现良好。作为一个著名的例子，思想链（CoT）是一种提示技术，要求模型在给出最终答案之前生成一系列中间标记，众所周知，它对于 Transformers 在需要数学或算法推理的任务上表现良好至关重要。部分研究者从表示能力的角度解释了这一点：变压器本身没有足够的表示能力来解决超出特定电路复杂性类别的问题（$TC^0$)，但通过 CoT，他们甚至可以模拟任何多项式时间图灵机。

CoT 在 Transformers 上的有效性自然会引发以下问题：

类似的增强（例如采用 CoT）能否改进 RNN，使其与 Transformer 相媲美？

本文贡献

本文通过研究各种方法来缩小 RNN 和 Transformer 在算法问题上的表示能力上的差距，从理论上回答了上述问题。通过一系列下限和上限结果，我们表明 CoT 提高了 RNN 的表示能力，但为了缩小与 Transformer 的差距，仅 CoT 不足以克服 RNN 的一个关键瓶颈：它们无法从上下文中检索信息，我们简称为上下文检索。

我们进一步说明，解决上下文检索瓶颈足以弥补这一差距：如果采用增强上下文检索能力的技术，包括涉及检索增强生成（RAG）和附加， RNN 可以解决所有多项式时间可解决的问题单个 Transformer 层。我们的主要贡献如下：

1.CoT 改进了 RNN，但无法缩小与 Transformer 的表示差距。

从积极的一面来看，CoT 使 RNN 严格具有更强的表达能力。
从消极的一面来看，我们表明采用 CoT 不足以缩小 RNN 和Transformer 之间的表示差距：RNN 的内存效率从根本上限制了它们执行上下文检索的能力，即使使用 CoT 也是如此。通过证明带有 CoT 的 RNN 无法解决一组直接要求上下文检索（包括联想召回）的基本算法问题，这一点得到了具体体现。通过证明 RNN 无法解决确定图是否是树的经典问题，我们进一步举例说明了在看似不相关的任务中隐含地需要上下文检索。
另一方面，我们证明 Transformers 具有轻松解决上述许多任务的表征能力，包括判断是否是树。此外，具有 CoT 的 Transformers 甚至可以有效地模拟具有 CoT 的 RNN，而参数数量只需要很小的乘法因子。

增强 RNN 的上下文检索能力可以缩小表示差距

我们证明，允许 RNN 调用函数调用来执行特定的上下文检索原语，足以提高其表示能力，以解决 CoT 的所有多项式时间可解问题，从而缩小RNN 和 Transformer 之间的表示差距。
另外，由于一层 Transformer 足以执行许多上下文检索操作，因此我们证明，通过在架构末尾添加一个 Transformer 层来隐式增强 RNN 的上下文检索能力也足以缩小差距

CoT 能否提高 RNN 的表示能力？

在本节中，我们的目标是了解带有 CoT 的 RNN 的表示能力。

我们首先展示了积极的结果，即带有 CoT 的 RNN 可以解决在没有 CoT 固定状态大小的情况下 RNN 无法完成的任务。然后我们继续了解 CoT 是否可以使 RNN 具有像 Transformer 一样的表现力。

我们表明，即使使用 CoT，RNN 仍然难以解决明确需要上下文检索的问题，并且这种表示差距会传播到看似与检索无关的推理任务，例如 IsTree。最后，我们表明这种差距确实是单方面的：只存在 Transformer 需要的参数比 RNN 指数少的任务，但反之则不然。

CoT 严格改进 RNN

CoT 无法缩小与 Transformer 的代表性差距

上下文检索的简单问题

首先，我们证明 RNN 在解决几个直接测试上下文检索能力的简单算法问题时与 Transformer 存在显着的表示差距。

对于上限，我们证明 Transformer 可以通过利用称为“注意力机制”的注意力机制来解决问题匹配，它接受查询令牌并关注与某些预定义坐标上的查询令牌相匹配的先前键。这种机制允许Transformer像键值字典一样读取其上下文窗口，从而可以完美解决问题。对于计数问题，我们另外使用数数注意力机制，通过均匀地关注查询标记的每次出现来计算查询标记的出现次数。

了解 RNN 超越简单上下文检索问题的表示能力

一个自然的问题是，如果算法问题不直接测试上下文检索能力，我们是否可以希望 RNN 具有解决它的表示能力？

在这种情况下，RNN 和 Transformer 是否具有相同的表示能力？

我们表明，RNN 中有限的内存大小仍然可能成为解决算法问题的瓶颈。即使检索能力没有在算法问题中显式测试，在推理答案时仍然可能隐式地需要它。

我们通过一个最小的算法问题示例（称为IsTree)：

给定一个无向图G的n节点，判断是否是一棵树，即每一对节点是否都由一条简单路径连接。 IsTree 的经典解决方案是运行深度优先搜索 (DFS)，它需要O(n)时间。

Transformer 明显比 RNN 更具表现力

上述定理表明，存在一些任务，其中 Transformer 所需的内存比 RNN 少得多。

然而，他们并没有排除是否存在相应任务，这种任务其中 Transformer 会比 RNN 更加冗余并且需要指数级更多的参数。

然而，以下定理证实了常规 RNN 不存在这样的任务。

增强上下文检索能力缩小表征差距

在前面中，我们表明 RNN 在上下文检索方面存在缺陷，因此导致与Transformer 的表示存在显着差距。

在本节中，我们的目标是了解：如果我们增强RNN 的上下文检索能力，RNN 与 Transformer 是否仍然存在表示差距？

我们通过考虑显式和隐式方法来增强上下文检索能力来回答这个问题，并表明这两种方法都可以缩小 RNN 和 Transformer 在解决算法问题时的表示差距。

Explicit Retrieval Through Regular Expression

首先，我们探索 RNN 与检索增强生成 (RAG) 的强大功能，它使 LM 能够检索相关信息以协助生成。在我们的上下文中，我们特别感兴趣的是允许 LM 调用函数从其上下文中检索信息，我们将其称为In-context Retrieval Augmented Generation (In-context RAG).

通过仅附加一层 Transformer 层进行隐式检索

我们在本节中正式表明，这种形式的隐式检索可以缩小 RNN 和 Transformer 在解决算法问题时的表示差距。我们考虑以下混合架构，它通过将单个 Transformer 层附加到 RNN 输出来组合 RNN 和 Transformer。

实验

训练了三种不同的架构：

（1）LLaMa

（2）Mamba

（3） Mamba with one additional layer of LLaMA block representing hybrid architectures.

图中可观察到：

1.CoT 提高了 Transformer 和 RNN 的性能。然而，随着图尺寸的增加，RNN 的性能急剧下降，并且 Transformer 的性能始终优于 RNN。这与我们的理论一致，即 CoT 可以提高 RNN 模型的表达能力，但表达能力仍然不足以解决 IsTree 任务

2.通过正则表达式的检索增强使所有模型达到近乎完美的准确性。这与我们的理论一致，即通过正则表达式进行检索增强可以提高 RNN 模型解决算法任务的表达能力

3.混合模型显示了所有模型中最好的性能，即使没有 CoT 或显式检索增强也能达到近乎完美的精度，这也反映在理论结果中。