DASH用分钟级搜索找到高效混合注意力结构,比现有方法快3万倍。
DASH: Fast Differentiable Architecture Search for Hybrid Attention in Minutes on a Single GPU

- 将离散的注意力模块选择转为连续可微优化,实现快速搜索。
- 单次搜索仅需1230万词元,20分钟内完成,效率提升超3万倍。
- 适合追求高效推理的LLM开发者,尤其关注轻量化部署场景。
混合注意力架构正成为提升大语言模型推理效率并保持模型质量的关键范式,因此其设计成为核心问题。现有方法多依赖人工经验或代理信号进行层间算子分配。近期基于神经架构搜索(NAS)的系统如Jet-Nemotron展示了自动化设计的潜力,但其后置搜索阶段耗时高达2000亿词元,难以作为常规方法使用。本文提出DASH,一种快速可微分混合注意力架构搜索框架:通过将离散的层间注意力算子分配松弛为连续架构逻辑值,构建可复用的教师对齐线性候选集,并在冻结模型与算子权重的前提下仅搜索架构本身,极大提升效率。在Qwen2.5-3B-Instruct上,DASH持续优于各类现有选择式混合注意力基线,证明直接可微搜索能发现更优架构。此外,其在RULER评测中表现超越已发布的Jet-Nemotron模型,同时在短上下文和通用基准上保持竞争力。值得注意的是,每次DASH搜索仅消耗1230万词元,约20分钟即可在单张RTX Pro 6000 GPU上完成,仅为Jet-Nemotron后置搜索所用词元量的0.006%。结果表明,高质量混合注意力架构可在分钟级可微搜索中获得,为混合架构设计提供了新方向。
原文摘要 · Abstract (English)
Hybrid attention architectures are becoming an increasingly important paradigm for improving LLM inference efficiency while preserving model quality, making hybrid architecture design a central problem. Existing designs often rely on manual empirical rules or proxy-based selector signals for layer-wise operator allocation. Recent NAS-style systems such as Jet-Nemotron demonstrate the promise of automated hybrid architecture search. However, Jet-Nemotron's PostNAS search stages alone use 200B tokens, making such search pipelines difficult to use as routine methods for hybrid architecture design. We introduce DASH, a fast differentiable search framework for hybrid attention architecture design, which relaxes discrete layer-wise attention operator placement into continuous architecture logits, prepares reusable teacher-aligned linear candidates, and performs architecture-only search with model and operator weights frozen to significantly enhance search efficiency. On Qwen2.5-3B-Instruct, DASH consistently outperforms a comprehensive suite of existing selector-style hybrid attention design baselines, showing that direct differentiable search can discover stronger hybrid architectures. Moreover, DASH achieves stronger RULER performance than released Jet-Nemotron models while remaining competitive on overlapping short-context and general benchmarks. Notably, each DASH search run uses only 12.3M tokens and takes about 20 minutes on a single RTX Pro 6000 GPU, corresponding to merely 0.006% of the PostNAS search tokens reported by Jet-Nemotron. These results suggest that high-quality hybrid attention architectures can be obtained through minutes-level differentiable search, providing a promising direction for hybrid architecture design.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。