摘要

本文介绍了一种新型神经网络架构——关联循环记忆转换器（Associative Recurrent Memory Transformer, ARMT），旨在处理极长序列数据，并在每个时间步保持恒定的处理时间。ARMT结合了转换器的自注意力机制和段级循环存储，特别适用于关联检索任务。在BABILong多任务长上下文基准测试中，ARMT在处理超过5000万token的单事实问题时，达到了79.9%的准确率，超越了现有模型。

原理

ARMT扩展了段级循环模型RMT，通过添加层级关联记忆Al s来处理分段输入Xl s。在每个输入段s和每层l中，前一段生成的记忆标记Ml+1 s−1被添加到Al s中，用于更新输入序列和记忆嵌入。关联块的机制类似于线性转换器，但仅关注特殊记忆标记，并采用不同的计算方式。记忆标记通过线性映射转换为键和值，并存储在准线性键值记忆中。ARMT通过更新关联矩阵Al s来整合前一段的信息，并通过查询当前段嵌入来检索关联。

流程

ARMT的工作流程包括以下步骤：首先，输入序列被分段处理，每个段通过转换器块进行处理。然后，每个段的记忆标记被转换为键和值，并存储在关联记忆中。在处理下一个段时，模型会检索并更新这些记忆标记。具体来说，模型首先通过转换器块处理当前段，然后将生成的记忆标记添加到关联记忆中。在处理后续段时，模型会使用这些记忆标记来更新输入序列和记忆嵌入。通过这种方式，ARMT能够在处理长序列时保持高效的记忆检索和更新能力。

应用

ARMT在处理长上下文任务方面展现出巨大的潜力，特别是在需要大量记忆检索和更新的复杂任务中。其高效的记忆机制和处理能力使其在问答系统、语言模型和复杂推理任务中具有广泛的应用前景。随着进一步的研究和优化，ARMT有望成为处理大规模语言模型和长序列数据的重要工具。