GSA通过门控机制提升线性注意力模型的上下文记忆能力,兼顾高效训练与推理。
Gated Slot Attention for Efficient Linear-Time Sequence Modeling
- 采用双层门控线性注意力结构,实现上下文感知的记忆读取与自适应遗忘
- 在需要上下文召回的任务中超越传统Transformer,且训练资源需求更低
- 特别适合微调预训练模型转为RNN架构的场景,减少从零训练成本
线性注意力Transformer及其门控变体虽支持并行训练和高效循环推理,但在需要高召回的任务中仍逊于传统Transformer,且从头训练需大量资源。本文提出门控槽注意力(Gated Slot Attention, GSA),通过引入受门控线性注意力(GLA)启发的门控机制,增强注意力的有界记忆控制(ABC)。GSA由两层GLA通过softmax连接,结合上下文感知的记忆读取与自适应遗忘,提升记忆容量的同时保持紧凑的递归状态。该设计通过GLA的硬件友好训练算法和减小的状态尺寸,显著提高训练与推理效率。此外,保留softmax操作在将预训练Transformer微调为RNN(T2R)设置中尤为有利,降低了对大规模从头训练的需求。大量实验验证了GSA在需要上下文召回及T2R场景中的优越性能。
原文摘要 · Abstract (English)
Linear attention Transformers and their gated variants, celebrated for enabling parallel training and efficient recurrent inference, still fall short in recall-intensive tasks compared to traditional Transformers and demand significant resources for training from scratch. This paper introduces Gated Slot Attention (GSA), which enhances Attention with Bounded-memory-Control (ABC) by incorporating a gating mechanism inspired by Gated Linear Attention (GLA). Essentially, GSA comprises a two-layer GLA linked via $\operatorname{softmax}$, utilizing context-aware memory reading and adaptive forgetting to improve memory capacity while maintaining compact recurrent state size. This design greatly enhances both training and inference efficiency through GLA's hardware-efficient training algorithm and reduced state size. Additionally, retaining the $\operatorname{softmax}$ operation is particularly beneficial in "finetuning pretrained Transformers to RNNs" (T2R) settings, reducing the need for extensive training from scratch. Extensive experiments confirm GSA's superior performance in scenarios requiring in-context recall and in T2R settings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。