DASH: A Single-GPU, Minute-Level Hybrid Attention Architecture Search Framework

DASH (Differentiable Architecture Search for Hybrid Attention) is a differentiable search framework designed for hybrid attention architectures, focusing on solving the challenge of selecting optimal attention operators for each layer. Through three key innovations—continuous architecture relaxation, teacher-aligned candidates, and pure architecture search with frozen weights—it achieves a 12.3 million token, ~20-minute single-GPU search, reducing search costs by 99.994% compared to Jet-Nemotron while maintaining performance advantages.

Hybrid attention architectures are an important paradigm for improving large model inference efficiency, balancing quality and efficiency via local/global/sparse/linear attention. Existing methods have limitations: manual design relies on experience and is hard to optimize; proxy signal selectors deviate from final performance; NAS methods like Jet-Nemotron consume 200 billion tokens in the PostNAS phase, leading to extremely high costs.

1. Continuous Architecture Relaxation: Convert discrete operator assignment into continuous architectural logic, supporting gradient optimization to avoid combinatorial explosion;
2. Teacher-Aligned Candidates: Pre-train linear candidates aligned with the teacher model’s behavior to ensure search starting point quality;
3. Pure Architecture Search with Frozen Weights: Only update architectural logic without repeated model training, improving efficiency and stability.

Performance Comparison: Outperforms all selector baselines on Qwen2.5-3B-Instruct, surpasses Jet-Nemotron on the RULER long-context benchmark, and maintains competitiveness on short-context/general benchmarks. Efficiency Data:

Metric	DASH	Jet-Nemotron	Savings Ratio
Search Token Count	12.3 million	200 billion	99.994%
Search Time	~20 minutes	Several days	99%+
GPU Requirement	Single RTX Pro6000	Multi-card cluster	-

Differentiable Selection Mechanism: Convert architectural logic into probabilities via softmax, forward pass uses weighted outputs of candidate operators, backward pass propagates gradients to update logic; Architectural Regularization: Introduce sparsity regularization, continuity penalty, and computational cost constraints to prevent architectural complexity; Post-Search Processing: Convert continuous logic to discrete configurations via Top-K selection/threshold truncation, which can be lightly fine-tuned for optimization.

1. Rapid Prototype Validation: Explore hybrid architecture configurations in minutes to accelerate iteration;
2. Model Customization: Search optimal configurations for scenarios like long-document processing, code generation, and edge deployment;
3. Architecture Research: Understand layer sensitivity to attention types, task preference patterns, and combination methods.

Limitations: Search space is limited to predefined candidates; may overfit to the search task; efficiency evaluation is based on specific GPUs; Future Directions: Expand the search space to include attention variants; multi-task generalized architectures; dynamic adaptive architectures; joint optimization of architecture and quantization precision.

DASH enables minute-level hybrid attention architecture search through efficient design, reducing costs by over 99% while delivering excellent performance. Its success proves efficiency and quality can coexist, turning architecture search from an expert privilege into a daily tool. It aligns with trends in AI model compression, efficient training, and inference optimization, pointing the way for NAS research.