[Paper Review] Efficient Architecture Search for Diverse Tasks

Jun-Hong Shen, Mikhail Khodak+1 more|arXiv (Cornell University)|Apr 15, 2022

Machine Learning and Data Classification5 citations

TL;DR

This paper proposes Dash, a differentiable neural architecture search method that efficiently searches over diverse kernel sizes and dilations in convolutional networks to improve performance across multiple domains. By leveraging Fourier diagonalization and kernel mixing, Dash achieves up to 10x faster search than baseline methods while outperforming state-of-the-art AutoML on seven of ten diverse tasks, including protein folding and PDE solving.

ABSTRACT

While neural architecture search (NAS) has enabled automated machine learning (AutoML) for well-researched areas, its application to tasks beyond computer vision is still under-explored. As less-studied domains are precisely those where we expect AutoML to have the greatest impact, in this work we study NAS for efficiently solving diverse problems. Seeking an approach that is fast, simple, and broadly applicable, we fix a standard convolutional network (CNN) topology and propose to search for the right kernel sizes and dilations its operations should take on. This dramatically expands the model's capacity to extract features at multiple resolutions for different types of data while only requiring search over the operation space. To overcome the efficiency challenges of naive weight-sharing in this search space, we introduce DASH, a differentiable NAS algorithm that computes the mixture-of-operations using the Fourier diagonalization of convolution, achieving both a better asymptotic complexity and an up-to-10x search time speedup in practice. We evaluate DASH on ten tasks spanning a variety of application domains such as PDE solving, protein folding, and heart disease detection. DASH outperforms state-of-the-art AutoML methods in aggregate, attaining the best-known automated performance on seven tasks. Meanwhile, on six of the ten tasks, the combined search and retraining time is less than 2x slower than simply training a CNN backbone that is far less accurate.

Motivation & Objective

• Address the lack of generalization in existing NAS methods beyond computer vision and natural language processing.
• Develop an efficient, scalable NAS approach that maintains high accuracy across diverse, under-explored domains such as biology, healthcare, and physical sciences.
• Overcome the computational bottleneck of searching over large kernel sizes and dilations in standard CNNs without sacrificing speed or model efficiency.
• Design a method that is both faster than existing differentiable NAS and more expressive than discrete search spaces, enabling broader applicability.
• Demonstrate that automated architecture search can surpass hand-crafted models in performance while remaining practical for real-world deployment.

Proposed method

• Fix a standard CNN backbone and search over a cross-scale operator space defined by varying kernel sizes (up to 15) and dilation rates (up to 127), enabling multi-resolution feature extraction.
• Introduce a kernel-mixing technique that combines multiple convolutions into a single convolution with a composite kernel, reducing complexity from O(|K||D|) to O(1).
• Apply Fourier diagonalization of convolutions to reduce computational cost from O(nk) to O(n log n), enabling efficient computation with large kernels.
• Use Kronecker products of undilated kernels and sparse matrices to accelerate dilated convolution computation on GPUs, achieving additional speedup.
• Employ a differentiable search framework with a shared supernet, followed by discrete architecture selection and retraining, ensuring practical deployment.
• Integrate the learned architecture into existing backbones (e.g., Wide ResNet, ConvNeXt, TCN) for backbone-agnostic performance gains.

Experimental results

Research questions

• RQ1 Can a differentiable NAS method efficiently search over large kernel sizes and high dilation rates to improve performance on diverse, non-vision tasks?
• RQ2 Does using Fourier-based convolution diagonalization enable scalable search over large kernels without incurring prohibitive computational costs?
• RQ3 Can a single architecture search method achieve performance comparable to or better than hand-designed models across multiple domains, including healthcare, physics, and biology?
• RQ4 Is it possible to maintain high search efficiency while significantly expanding the search space beyond standard small-kernel NAS methods?
• RQ5 To what extent can the learned architecture patterns reveal intrinsic data properties or task-specific inductive biases?

Key findings

• Dash achieves the best-known automated performance on seven out of ten diverse tasks, including heart disease detection, protein folding, and PDE solving.
• On six of the ten tasks, the combined search and retraining time for Dash is less than twice the time required to train a standard Wide ResNet, demonstrating high practical efficiency.
• Dash reduces search time by up to 10x compared to baseline methods like DARTS, while maintaining or improving accuracy.
• Dash improves top-1 accuracy on ImageNet by 1.4% over the ConvNeXt backbone, surpassing even human-tuned models in this setting.
• The method is backbone-agnostic and successfully enhances performance on TCN and ConvNeXt, showing compatibility with various architectures.
• The aggregated convolution mechanism enables efficient, differentiable search over large kernel spaces without sacrificing model efficiency or accuracy.