[Paper Review] RepNAS: Searching for Efficient Re-parameterizing Blocks

Mingyang Zhang, Xinyi Yu+2 more|arXiv (Cornell University)|Sep 8, 2021

Advanced Neural Network Applications29 references4 citations

TL;DR

RepNAS proposes a novel one-stage NAS framework that searches for efficient, re-parameterizable diverse branch blocks (ODBBs) in a flexible search space allowing arbitrary branch counts per block. By using gradient-based architecture search with block-wise weight sharing and dynamic pruning under memory constraints, RepNAS achieves state-of-the-art accuracy on ImageNet with inference speeds comparable to single-path models, outperforming both human-designed and prior NAS models under similar latency constraints.

ABSTRACT

In the past years, significant improvements in the field of neural architecture search(NAS) have been made. However, it is still challenging to search for efficient networks due to the gap between the searched constraint and real inference time exists. To search for a high-performance network with low inference time, several previous works set a computational complexity constraint for the search algorithm. However, many factors affect the speed of inference(e.g., FLOPs, MACs). The correlation between a single indicator and the latency is not strong. Currently, some re-parameterization(Rep) techniques are proposed to convert multi-branch to single-path architecture which is inference-friendly. Nevertheless, multi-branch architectures are still human-defined and inefficient. In this work, we propose a new search space that is suitable for structural re-parameterization techniques. RepNAS, a one-stage NAS approach, is present to efficiently search the optimal diverse branch block(ODBB) for each layer under the branch number constraint. Our experimental results show the searched ODBB can easily surpass the manual diverse branch block(DBB) with efficient training.

Motivation & Objective

• Address the gap between neural architecture search (NAS) performance and real-world inference speed by designing a search space compatible with structural re-parameterization.
• Overcome the limitations of fixed, manually designed multi-branch blocks (e.g., RepVGG, ACB) that suffer from high memory usage and suboptimal branch configurations.
• Enable automatic, end-to-end search for optimal diverse branch blocks (ODBBs) without retraining, while maintaining fast inference through single-path fusion.
• Develop a memory-efficient training strategy that dynamically prunes low-importance branches during backpropagation to scale to larger models under GPU memory constraints.

Proposed method

• Propose a new search space, Rep search space, where each block independently preserves an arbitrary number of branches during training and can be fused into a single path at inference time.
• Introduce RepNAS, a one-stage gradient-based NAS method that learns the importance of entire branches (not individual edges) via architecture parameters updated in backpropagation.
• Apply block-wise weight sharing in the supernet to enable efficient training and enable accurate architecture ranking without retraining.
• Implement a memory-aware training strategy: during each forward pass, branches with low importance (determined by temperature-controlled gates) are pruned sequentially until GPU memory constraints are met.
• Use a differentiable architecture search (DARTS-style) with a temperature parameter to soft-assign branch importance, allowing gradient flow to update both network and architecture parameters simultaneously.
• After training, the final ODBB is obtained by fusing all high-importance branches into a single-path model, enabling direct deployment without retraining.

Experimental results

Research questions

• RQ1 Can a flexible, re-parameterizable search space be designed to allow arbitrary branch counts per block while enabling efficient, single-path inference?
• RQ2 Can a one-stage NAS method effectively search for optimal diverse branch blocks (ODBBs) without retraining, under memory constraints?
• RQ3 Does dynamic pruning of low-importance branches during training improve memory efficiency without degrading performance?
• RQ4 Can the proposed method outperform both manually designed blocks (e.g., RepVGG, ACB) and prior NAS models in accuracy and latency trade-off?

Key findings

• The searched ODBB-A0 achieves 31.4% COCO AP on object detection with 0.048s inference time, surpassing ResNet-18 (28.1% AP) at faster speed.
• ODBB-B3 achieves 37.3% COCO AP with 0.148s inference time, outperforming ResNet-101 (34.6% AP) by 2.7% AP while maintaining comparable latency.
• ODBB with only 50 branches achieves higher top-1 accuracy than RepVGG (66 branches), DBB (88 branches), and ACB (66 branches), demonstrating superior efficiency.
• The ODBB with 75 branches achieves performance equivalent to larger configurations, indicating diminishing returns beyond a certain branch count.
• Weight sharing in the supernet leads to faster convergence and more stable performance of sampled subnetworks compared to independent weight updates.
• The supernet trained with RepNAS effectively ranks high-performing subnetworks, as evidenced by consistently higher and more stable accuracy in sampled subnets over training epochs.