RADLADS: Rapid Attention Distillation to Linear Attention Decoders at Scale
原作者: Daniel Goldstein, Eric Alcaide, Janna Lu, Eugene Cheah
原作者: Daniel Goldstein, Eric Alcaide, Janna Lu, Eugene Cheah
原始论文采用 CC BY 4.0 许可(http://creativecommons.org/licenses/by/4.0/)。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性,请参阅原始论文。 阅读完整免责声明
技术摘要:RADLADS
问题陈述
虽然线性注意力变体(如 RWKV 和 Mamba)相比传统的 Softmax 注意力 Transformer 具有显著优势——具体表现为推理时间为 O(1) 且通过避免键值(KV)缓存实现了常数级显存占用——但从头开始训练大规模线性模型仍然极其昂贵。最先进的(SoTA)Transformer 模型通常需要在超过十万亿个 token 上进行训练,这对大多数机构来说是无法承受的成本。现有的将预训练 Softmax Transformer 转换为线性或循环模型的方法(例如 T2R、SUPRA、LoLCats、MOHAWK)在历史上往往需要海量的 token 量(范围从 200 亿到 1000 亿+ 个 token),或者导致下游性能较差,特别是在 MMLU 等基准测试上。此外,许多转换尝试依赖于混合架构(保留部分 Softmax 注意力),或者无法达到原始教师模型的质量水平。
方法论
作者提出了 RADLADS(大规模快速注意力蒸馏至线性注意力解码器),这是一种旨在以极低计算成本将 Softmax 注意力 Transformer 转换为线性注意力解码器模型的协议。该过程包含三个主要阶段和特定的架构适配:
1. RADLADS 协议
该转换是一个三步过程,仅需 3.5 亿至 7 亿个 token(不到教师模型预训练数据量的 0.005%):
- 设置(注意力权重转移): 将教师模型的注意力相关权重(Wq,Wk,Wv,Wo)直接转移到学生模型的序列混合层中。其他权重通过标准的预训练方法进行初始化,或设置为模仿教师的行为但不产生即时影响。
- 步骤 1:注意力隐藏状态对齐: 训练学生模型的序列混合层以逼近对应教师注意力层的隐藏状态输出。这是通过 L2 距离(或 MSE)目标函数实现的。作者发现,使用门控线性注意力内核(移除“错位一位”的衰减项和奖励项)比标准 RWKV-6 能更紧密地拟合教师的隐藏状态。
- 步骤 2:知识蒸馏: 整个学生模型通过 KL 散度损失来逼近教师模型的输出 Logits。作者假设事实性知识主要存在于教师模型的 MLP 和嵌入层中;因此,保持这些组件的学习率较低或固定,以防止灾后遗忘,而对序列混合器进行更积极的训练。
- 步骤 3:上下文长度扩展(可选): 使用标准的交叉熵损失在更长的序列(最高 16k tokens)上对模型进行微调,以增强长上下文能力。或者,作者提出了 步骤 2a,即在蒸馏阶段本身将序列长度扩展到 4096,这可能使得无需单独的步骤 3。
2. 新型架构
作者发现标准的 RWKV 架构并不总是最优的转换选择。他们引入了两种新变体:
- RAD-RWKV6 ("RADFinch"): 对 RWKV-6 的一种修改,使用门控线性注意力内核和状态平衡技术来提高稳定性和拟合度。
- RAD-RWKV7 ("RADGoose"): 对 RWKV-7 的一种修改,移除了在本文语境下并无益处的“token 偏移(tokenshift)”机制,并直接应用旋转位置编码(RoPE)。与未修改的 RWKV-7 相比,该架构展现出了更快的收敛速度和更低的蒸馏损失。
3. 超参数与数据
- 数据集: 作者在所有转换步骤中最终选择了 DCLM (DataComp-LM),发现其在转换 Qwen 模型方面优于 FineWeb 或 FineWeb-Edu。
- 学习率: 步骤 1 使用余弦退火调度,从高学习率(10−3)开始以对齐隐藏状态,最后接近教师模型最终预训练的学习率(10−5)。步骤 2 和步骤 3 使用恒定学习率。
核心贡献
- RADLADS 蒸馏配方: 一个详细的、分步骤的协议,包括特定的超参数、token 数量和数据集选择,能够以极小的数据量实现高质量转换。
- 新型架构: 引入了 RAD-RWKV6 和 RAD-RWKV7,它们针对转换过程进行了优化,比其未经修改的前身具有更快的推理速度和更好的教师模型对齐效果。
- 大规模模型: 成功将流行的开源 Qwen2.5 模型转换为参数规模为 7B、32B 和 72B 的线性注意力变体。
- 开源发布: 在 Apache 2.0 许可协议下(对于 72B 模型受 Qwen 许可限制)发布代码和转换后的模型(QRWKV6/7-7B/32B/72B),允许他人复现该过程。
结果
转换后的模型在同规模的纯循环模型中达到了最先进的性能:
- 效率: 转换一个 72B 模型成本不到 2,000 美元,且仅需约 7 亿个 token。
- 性能: 在标准基准测试(Lambada, MMLU, ARC 等)上,RADLADS 模型持续优于其他转换方法(如 SUPRA, LoLCats, MOHAWK, ARWKV)。
- QRWKV7-7B-Instruct 的相对 MMLU 得分为教师模型的 92.4%,显著优于其他 7B 规模的转换模型。
- QRWKV6-72B-Instruct 的相对 MMLU 得分为 89.9%,创下了该规模下纯 RNN 语言模型的新纪录。
- 推理速度: 由于采用了线性注意力机制,转换后的模型随着上下文长度的增加展现出显著的速度提升。例如,在输入 8k tokens、输出 256 tokens 的情况下,32B 的 QRWKV7 比 Qwen3-32B 教师模型快 1.61 倍;在输入 6k、输出 2k 的情况下,加速比达到 3.41 倍。
局限性与声明
论文谦逊地承认了若干局限性:
- 推理与长上下文: 虽然在标准基准测试上表现强劲,但模型在复杂推理任务(如 Minerva Math)和极长上下文任务(如 RULER)中表现出局限性,即它们在增加输出 token 时,性能提升不如教师模型明显。
- 架构敏感性: 每种新的架构设计都需要细致的测试以确保与 RADLADS 协议兼容。例如,某些线性注意力变体中常见的 GroupNorm/LayerNorm 在 14B+ 规模下会导致训练不稳定,因此必须替换为预缩放(pre-scaling)技术。
- 数据集对齐: 作者指出,推理数据集可能需要与教师模型的分布更加紧密对齐,以防止出现重复循环行为,作者通过将 DCLM 与 OpenThoughts 进行混合部分解决了这一挑战。
重要意义
论文声称 RADLADS 提供了一条具有成本效益的路径,使大规模线性注意力模型的使用变得民主化。通过将训练所需的 token 量从万亿级降低到亿级,它使研究人员和规模较小的机构能够在无需承担极端预训练成本的情况下,测试、训练并部署新类型的 RNN 架构。作者将其定位为加速下一代压缩状态注意力变体发展的工具。
您所在领域的论文太多了?
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要,使用您的语言。
每周获取最佳 AI 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。