论文

用于对数线性注意力的自适应记忆衰减

Adaptive Memory Decay for Log-Linear Attention

模型架构Transformer

摘要

序列模型面临着内存容量和计算效率之间的基本权衡。 Transformer 以二次成本实现表达性上下文建模,而线性注意力和状态空间模型通过将上下文压缩为固定大小的隐藏状态来在线性时间内运行,本质上限制了召回率。对数线性注意力通过在 Fenwick 树层次结构中组织内存,以对数线性计算成本随序列长度对数增长其隐藏状态来实现这种权衡。然而,它的记忆衰减参数 λ 是固定的并且与输入无关,无论内容如何,​​在所有层次结构级别上分配统一的权重,这引入了不必要的刚性。我们建议通过轻量级两层 MLP 直接从输入中学习 λ,产生适应内容而不是位置的每个标记、每个级别的衰减。 softplus 激活让每个 Fenwick 树级别独立扩展,避免了 softmax 引入的层间竞争。这种修改准确地保留了对数线性复杂性,并增加了可忽略不计的参数开销。我们对联想回忆、选择性复制和语言建模进行了评估,发现依赖于输入的衰减始终优于基线,在基线 λ 退化或完全崩溃的长程记忆设置中收益最大。