夕小瑶科技说 原创
作者 | 付奶茶、python近期，Transformer再度成为学术界的热门话题！

Meta的一项最新研究带来了一个新的发现：在某些情况下，"Transformer实际上可以被视为多状态的RNN（MSRNN）。

自2017年问世以来，Transformer以其在多个自然语言处理任务中的卓越表现迅速成为了这一领域的宠儿，同时也使得原本在文本序列处理中占据主导地位的RNN相形见绌。

Meta的这一发现指出，尽管Transformer在理念上与RNN有所区别，但Transformer实际上可以看作是一种具有无限隐藏状态容量的RNN变体。通过限制处理的标记数量，预训练Transformer可以被转换为有限的MSRNN。此外，论文中提出了一种新的策略“TOVA”，该策略基于注意力分数选择保留哪些标记，相较于现有的缓存压缩技术更为简单。接下来，让我们一起深入探索这篇论文～

论文标题:
Transformers are Multi-State RNNs

论文链接:
https://arxiv.org/pdf/2401.06104.pdf

为什么说Transformer是MSRNN

RNNs的核心在于递归式，在时间接收t时刻的表示和上一时刻的隐藏状态，输出两个更新后的Token和新的隐藏状态：

而Transformer采用非递归方式，每一层接受一系列Token表示作为输入，返回一个经过变换的表示:

Transformer每一层由自注意力（SelfAttn）和前馈网络组成（FF）。自注意力机制将输入投影到三个维度为 的矩阵：Q（查询），K（键），和 V（值），并据此进行计算。

本篇论文的核心观点是将原本仅用的解码器（decoder-only）的Transformer重新概念化为无限多状态的多状态递归神经网络MSRNN。

MSRNN的表达式为：

在MSRNN中代表第第l层在时间点t的隐藏状态矩阵，其中每一行代表一个独立的单状态，整个矩阵构成一个多状态系统。当函数在所拥有时间点t上取均值为1时，，MSRNN退化为传统的单状态RNN，表明每个时间点只有一个活跃状态。在特殊的情况下，若，则每个时间步t的活跃的单状态数量与该时间步的输入Token数相等。在这种的设置下，Transformer可以被概念化为MSRNN的一种形式，其中第l层在时间点t的隐藏状态H_t^l由键值对组成，每层间计算遵循以下方式

其中, , 是的自注意力投影,将以上元素整合，可以得到描述Transformer作为MSRNN时的完整方程式：

研究团队提出将预训练的Transformer视为具有无限状态数的MSRNN，并探讨了这些模型是否真正利用了这种无限容量。为了创建一个有限的MSRNN，他们定义了一个函数，其中是一个常数。

虽然将Transformer视为具有无限状态的多状态递归神经网络（MSRNN）在理论上具有吸引力的，但在实际操作中面临几个关键限制:

• 计算资源的限制：在实际应用中，无限状态的模型会要求无限的内存和计算资源，这在现实中是不可行的。每增加一个状态，都需要更多的内存来存储和更多的计算资源来处理这些状态。

• 效率问题：随着状态数的增加，模型的效率会显著下降。

• 长期依赖问题：尽管理论上无限状态能够更好地处理长期依赖问题，但在实际中，长时间累积的状态会导致信息的混淆和性能下降。

因此，研究团队提出当时间点超过时，应用压缩策略以符合内存限制。这种策略可以是预先定义的，或在训练期间学习的。

无限转有限的策略

在过去其他人的研究工作中，已经提出了几种压缩策略，这些策略可以用来将无限的MSRNN转换为有限的MSRNN。

窗口策略

采用先进先出（FIFO）的方式，当状态容量达到限制时，丢弃最旧的状态，确保只保留最新的状态。

窗口+i策略

在窗口策略上进行了扩展，保留前i个Token，其中i是一个常数。先前的研究显示，窗口+i策略使用极少数早期Token，显著优于纯窗口策略。

H2O策略

该方法同样保留了最近的Token窗口和额外的早期Token，与窗口+i不同，它通过整个序列动态选择非窗口Token，根据注意力分数保留得分最高的Token。

但是，现有方法中没有一个考虑到策略与RNN之间的关联，因此研究团队提出了一种新的策略TOVA：

TOVA策略

TOVA比之前提出的方法更简单，根据最后一个Token的注意力权重保留顶部状态。在每个解码步骤中，考虑当前查询对所有多状态中Token以及当前Token的softmax标准化注意力分数，得分最低的Token被丢弃。

TOVA不像之前的策略那样固定最近Token的窗口或偏好早期Token。尽管存在轻微的近期偏见，但许多最近的Token仍然被丢弃，一些初始Token在数千个解码步骤中被保留，表明它们对于成功解码非常重要。实验表明，TOVA策略在多个任务中表现优于其他策略，并且与相应的无限MSRNN（常规预训练的Transformer）的结果接近。

实验

实验部分，研究的目标是检验不同压缩策略（有限的MSRNN）是否能够匹配完整的、无限的MSRNN模型的性能。实验分为三个主要部分：语言建模、长范围理解和长文本生成。

语言建模：使用PG-19测试集来报告困惑度，这是评估长期语言模型性能的广泛使用的基准测试。

长范围理解：利用ZeroSCROLLS基准测试的两个测试集，分别关注长范围摘要和长范围问答（QA）。

长文本生成：模型以生成长篇故事的提示为输入，从每个模型版本中抽取100个独特故事。鉴于比较两个故事的复杂性，使用GPT-4作为评估器。

实验使用了三个最先进的Transformer解码器LLM家族：LLaMA-2、Mistral和Yi，每个家族提供了约7B参数的版本进行评估。对于不同的任务（语言建模、长期理解和长文本生成）进行了评估。

在长期理解任务中，针对SQuALITY和QASPER测试集对经过指令调优的LLMs进行了评估。长期摘要结果显示，TOVA策略在所有多状态尺寸和模型上一致优于基线，使用TOVA策略与完整上下文的顶线模型的结果相差不到一个点。长期QA的结果表明，TOVA策略与基线之间的差距在某些情况下超过了5个F1分数点。然而，TOVA策略需要半数的完整多状态尺寸才能在一个F1分数点范围内达到顶线模型的性能。

在长文本生成任务中，限制多状态尺寸会导致生成的文本更短。全模型的平均故事长度为1566个Token，而在多状态尺寸为1024时保持不变，但在512个Token时降至1503个，256个Token时降至1361个。GPT-4评估结果表明，使用较小的多状态尺寸（256个Token），TOVA策略在47%的情况下输给顶线模型，但在其余情况下胜出或持平。

小结

整体而言，该论文提出了一种新颖的视角，将Transformer模型视为具备多状态的循环神经网络（RNN）。研究团队引入了一种名为TOVA（Token Omission Via Attention）的策略，有效地在保持模型性能的同时显著减少了计算资源的需求。这一发现对理解和优化Transformer模型具有重大意义。