论文

ProxyFormer:用于超长上下文和高分辨率生成的双流代理架构

ProxyFormer: A Dual-Stream Proxy Architecture for Ultra-Long Context and High-Resolution Generation

模型架构Transformer

摘要

注意力计算和键值(KV)缓存相对于序列长度的二次增长是超长上下文语言模型和高分辨率生成模型的中心瓶颈。我们提出 ProxyFormer,一种基于代理词元构建的通用双流架构。在每一层中,细粒度的局部特征被自下而上地压缩为一小组代理状态;昂贵的全局交互仅在压缩代理空间中执行;然后,全局上下文代理被解压缩并自上而下注入回本地流中。由于本地流跨层持续存在,因此一个压缩步骤未捕获的细粒度信息仍然可以访问以供以后细化,从而减轻了传统一次性压缩的不可逆信息丢失。我们进一步介绍了分解的多级压缩/解压缩、逐层动态压缩比、非对称双嵌入和仅代理的 KV 缓存推理方案。在批量大小为 1 的 16GB GPU 上,标准的仅解码器模型只能训练大约 20K 个 token 的序列,而压缩比为 64 的 ProxyFormer 将可训练序列长度扩展到大约 0.7M。使用 64K 窗口训练的模型在具有 1,048,576 个标记的多针检索任务中保留了 92%-95% 的检索准确率,而使用 8K 窗口训练的模型在外推到 256K 标记时的准确率超过 94%。初步图像生成实验证明了 ProxyFormer 对于像素空间和潜在空间流匹配的可行性。