论文
并行因果关联字段:用于长上下文语言建模的门控稀疏内存
Parallel Causal Associative Fields: Gated Sparse Memory for Long-Context Language Modeling
摘要
Transformer 通过提供直接的词元到词元通信路径来实现强大的语言建模性能,但因果自注意力随上下文长度呈二次方缩放。循环模型和状态空间模型降低了这种成本,但将历史记录压缩为顺序更新的固定大小状态。本文研究了第三个原语:因果后继记录上的并行内容寻址存储器。提出的并行因果关联字段(PCAF)将上下文窗口中的本地记录写入哈希桶,检索当前查询的有界候选集,在后继词元上形成稀疏缓存分布,并通过学习门将该缓存与参数本地语言模型混合。由此产生的模型保持稀疏的长上下文访问,同时避免单个固定的循环状态瓶颈。我们使用分布式 Google Cloud TPU v4-32 pod 在 WikiText-103 和 PG-19 上的完全自回归预训练下评估 PCAF。在 303M 参数和上下文长度 T = 2048 时,PCAF 语义在 WikiText-103 上达到 36.31 困惑度,在 PG-19 上达到 52.45 困惑度,而匹配的密集 Transformer 的困惑度为 47.49 和 53.84。 PCAF-semantic 在 TPU pod 上同时处理 0.61-0.62M 词元/秒,而密集和局部注意力基线的处理速度为 0.43M 词元/秒。支持 41M 参数多种子扫描和单 GPU 组件消融表明关联缓存、检索容量和学习门对速度与质量的权衡有重大影响。