论文

DASH:在单个 GPU 上在几分钟内快速进行混合注意力的可微分架构搜索

DASH: Fast Differentiable Architecture Search for Hybrid Attention in Minutes on a Single GPU

模型架构Transformer

摘要

混合注意力架构正在成为提高 LLM 推理效率同时保持模型质量的越来越重要的范例,这使得混合架构设计成为一个中心问题。现有的设计通常依赖于手动经验规则或基于代理的选择器信号来进行分层算子分配。最近的 NAS 式系统(例如 Jet-Nemotron)展示了自动化混合架构搜索的前景。然而,Jet-Nemotron 的 PostNAS 搜索阶段仅使用 200B 词元,使得此类搜索管道难以用作混合架构设计的常规方法。我们引入了 DASH,一种用于混合注意力架构设计的快速可微搜索框架,它将离散分层注意力算子放置到连续架构 logits 中,准备可重用的教师对齐线性候选,并执行仅架构搜索,模型和算子权重被冻结,以显着提高搜索效率。在 Qwen2.5-3B-Instruct 上,DASH 始终优于一整套现有的选择器式混合注意力设计基线,这表明直接可微搜索可以发现更强大的混合架构。此外,DASH 实现了比已发布的 Jet-Nemotron 模型更强的 RULER 性能,同时在重叠的短上下文和一般基准上保持竞争力。值得注意的是,每次 DASH 搜索运行仅使用 1230 万个词元,在单个 RTX Pro 6000 GPU 上大约需要 20 分钟,仅相当于 Jet-Nemotron 报告的 PostNAS 搜索词元的 0.006%。这些结果表明,可以通过分钟级可微搜索获得高质量的混合注意力架构,为混合架构设计提供了一个有前途的方向。