DASH：基于单张 GPU 以分钟级速度进行混合注意力可微架构搜索
机器学习
2026-05-21 v1 人工智能
计算与语言
摘要
混合注意力架构正在成为提高 LLM 推理效率而不牺牲模型质量的重要范式，使得混合架构设计成为核心问题。现有设计通常依赖手动经验规则或基于代理的选择器信号进行层级算子分配。最近的 NAS 类系统如 Jet-Nemotron 展示了自动化混合架构搜索的前景。然而，Jet-Nemotron 的 PostNAS 搜索阶段alone 使用了 2000 亿标记，使得此类搜索管道难以作为日常方法用于混合架构设计。我们引入 DASH，一种用于混合注意力架构设计的快速可微搜索框架，将离散的层级注意力算子放置松散为连续的架构逻辑，对齐的线性候选者进行可重用化，并通过冻结模型和算子权重进行仅架构搜索，从而显著提升搜索效率。在 Qwen2.5-3B-Instruct 上，DASH 始终优于综合套件中现有的基于选择器的混合注意力设计基线，表明直接的可微搜索可以发现更强的混合架构。此外，DASH 的 RULER 性能优于公开的 Jet-Nemotron 模型，同时在重叠的短上下文和通用基准测试中保持竞争力。值得注意的是，每次 DASH 搜索仅使用 1230 万标记，约 20 分钟即可在单张 RTX Pro 6000 GPU 上完成，相当于 Jet-Nemotron PostNAS 搜索标记的仅 0.006%。这些结果表明，通过分钟级可微搜索可以获得高质量的混合注意力架构，为混合架构设计指明了前景方向。
引用
@article{arxiv.2605.20936,
title = {DASH: Fast Differentiable Architecture Search for Hybrid Attention in Minutes on a Single GPU},
author = {Weizhe Chen and Miao Zhang and Junpeng Jiang and Yaping Li and Weili Guan and Liqiang Nie},
journal= {arXiv preprint arXiv:2605.20936},
year = {2026}
}
备注
19 pages, 7 figures