SUPRA：无须额外训练，将Transformer变为高效RNN，推理速度倍增

Transformers 已经确立了自己作为首要模型架构的地位，特别是因为它们在各种任务中的出色表现。但是Transformers 的内存密集型性质和随着词元数量的指数扩展推理成本带来了重大挑战。为了解决这些问题，论文“Linearizing Large Language Models”引入了一种创新的方法，称为UPtraining for Recurrent Attention (SUPRA)。该方法利用预训练的Transformers 并将其转换为递归神经网络，在保持预训练的好处的同时实现有效的推理。

SUPRA方法旨在将预训练的大型语言模型（LLMs）转化为RNNs，具体步骤包括替换softmax归一化为GroupNorm，使用一个小型MLP投影queries和keys。这种方法不仅降低了训练成本（仅为原来的5%），还利用了现有预训练模型的强大性能和数据。

Transformers 面临着与序列长度线性增长的高推理成本。相比之下，rnn提供固定成本推理，因为它们能够保持恒定大小的隐藏状态，这使得它们对于需要高效和可扩展推理的任务具有吸引力。

引入线性Transformers的概念，减轻了标准softmax注意机制的计算开销。线性Transformers用线性相似函数代替softmax，该函数可以重新表述为像RNN一样工作。但在许多基准测试中，线性Transformers的表现仍然不如softmax，主要是由于稳定性问题和预训练的复杂性。

向上训练方法（Uptraining ）

向上训练指的是用最少的额外训练使预训练的模型适应新架构的过程，而不是微调，微调通常涉及在不同的数据集上重新训练模型。而SUPRA特别关注通过修改其注意机制将预训练的Transformers转换为rnn。

SUPRA的核心在于将Transformers的注意力机制转化为一种循环形式。softmax关注被一个允许循环更新的线性函数所取代。这种转换是至关重要的，因为它使模型能够增量地更新其状态，类似于RNN。

使用转换后的键(k)和值(v)在每个时间步更新循环状态(s)和归一化因子(z)。这些更新以允许模型以一次处理一个词元序列的方式执行，从而显著降低了推理成本。

phi(x)定义为:

1、初始化循环状态和归一化因子为零。

2、对于输入序列中的每个标记，使用学习到的权重矩阵计算查询和键。对这些向量应用MLP和旋转位置嵌入(RoPE)，以确保它们有效地处理顺序数据。

3、使用转换后的键和值更新状态和归一化因子。MLP应用的转换适当地调整键和值。

4、通过归一化转换后的查询和状态的点积来计算每个时间步骤的注意力输出。对输出应用GroupNorm，取代传统的softmax归一化，提高输出稳定性。

对序列中的每个标记重复上述步骤，确保模型以循环的方式处理序列。

对模型进行微调

在使用上述步骤将Transformers转换为RNN之后，在较小的数据集上对模型进行微调。这个微调过程调整新组件(MLP、GroupNorm等)的权重，优化性能。与从头开始预训练相比，这一步只需要一小部分计算资源。

结果展示

研究人员通过提高1B到7B参数的一系列模型(包括Llama2和Mistral模型)来测试SUPRA方法。这些模型在标准语言理解基准和长上下文评估上进行了测试，以评估它们的性能和局限性。

使用Eleuther评估工具对标准自然语言理解(NLU)任务进行评估。Llama2-7B和Mistral-7B等初始化的supra处理模型在大多数基准测试中保持了高性能，优于RWKV-5等其他线性模型。

由于需要在扩展序列上保留信息，长上下文任务对许多模型来说是具有挑战性的。SUPRA模型在各种上下文长度下对来自SCROLLS基准的任务(如Qasper和NarrativeQA)进行评估。将这些模型的性能与它们的训练上下文长度进行比较。

SUPRA模型显示出超出其训练上下文长度保持性能的能力，这是与循环模型相关的典型特征。但是未经修改的Transformers在其最大训练上下文长度上通常优于SUPRA模型。这表明尽管SUPRA模型是有效的，但在处理非常长的上下文方面仍有改进的余地。

主要论点详解

1、SUPRA方法在多大程度上能解决长上下文任务中的性能不足？

性能不足的原因：使用了固定衰减向量γ，虽然在短上下文任务中表现良好，但在长上下文任务中，衰减效应导致有效上下文长度缩短，影响了性能。虽然SUPRA方法通过简化线性注意力计算，减少了推理成本，但仍然继承了线性模型在长上下文任务中的固有不足。例如，线性模型在长上下文建模中不如基于softmax注意力的Transformer模型表现好。

论文提出了可以探索更复杂的循环状态更新规则，如门控策略（gating strategies）和高阶线性注意力（higher order linear attention），以提高长上下文任务的性能。或增加专门用于长上下文任务的训练数据，可能有助于提升模型在这些任务中的表现。

2、替换softmax归一化为GroupNorm对模型稳定性和性能的具体影响是什么？

softmax归一化在大规模模型的训练过程中容易引起数值不稳定性，尤其是当模型处理长序列数据时。这种不稳定性在高维空间中特别明显。将softmax归一化替换为GroupNorm，可以有效缓解这种不稳定性。GroupNorm通过对组内的小批量数据进行归一化，减少了模型训练时的数值波动，增加了训练过程的稳定性。论文中的实验表明，使用GroupNorm的SUPRA模型在大规模模型（例如7B参数模型）的训练过程中表现更加稳定，减少了由于数值不稳定性导致的训练中断和性能下降。

将softmax归一化替换为GroupNorm对模型的稳定性和性能有显著正面影响。虽然在某些极长上下文任务中性能有所下降，但总体而言，这种替换在提高训练稳定性和保持性能方面表现良好。未来可以进一步优化GroupNorm的使用，以期在更多任务中提升模型性能。

3、是否可以通过进一步的优化策略（如门控机制）来弥补SUPRA方法在上下文学习中的不足？

a、通过引入门控机制，可以更好地控制信息流动，保持重要信息，丢弃不必要的信息，增强模型对长上下文的建模能力。门控机制可以平滑状态更新过程，减少长序列数据中的梯度消失和爆炸问题，提升模型训练的稳定性。

b、高阶线性注意力通过在计算注意力时引入更复杂的核函数，能够捕捉到更丰富的特征交互。通过使用更复杂的核函数，可以捕捉到更复杂的特征交互，提升模型的表达能力和性能。高阶线性注意力能够更有效地处理长上下文信息，提高模型在长上下文任务中的表现。

c、动态位置编码方法（如RoPE）可以更好地处理长序列中的位置信息。动态位置编码能够增强模型对长序列中位置信息的敏感性，提高长上下文建模能力。动态位置编码能够在不同长度的序列中保持稳定的性能，提升模型的泛化能力。

总结

SUPRA方法通过替换softmax归一化为GroupNorm，显著提升了模型的稳定性和性能。然而，在处理长上下文任务时仍存在一定不足。通过引入门控机制、高阶线性注意力、专门的训练数据和任务优化、动态位置编码以及结合其他高效注意力机制，可以进一步优化SUPRA方法，弥补其在上下文学习中的不足，提升模型的整体性能。

未来的研究可以继续探索这些优化策略的组合应用，以进一步提升模型在长上下文任务中的表现，为自然语言处理领域的发展提供更多可能性。

论文地址：

https://arxiv.org/abs/2405.06640

源代码

https://github.com/tri-ml/linear_open_lm

转换后的模型：

https://huggingface.co/TRI-ML/mistral-supra