[Paper Review] DAAS: Differentiable Architecture and Augmentation Policy Search

Xiaoxing Wang, Xiangxiang Chu+2 more|arXiv (Cornell University)|Sep 30, 2021

Advanced Neural Network Applications32 references4 citations

TL;DR

DAAS proposes a differentiable joint search framework for neural architecture and data augmentation policy optimization using bi-level optimization with Gumbel-softmax for architecture search and policy gradient-based second-order approximation for augmentation search. It achieves 97.91% accuracy on CIFAR-10 and 76.6% top-1 accuracy on ImageNet in just 1 GPU-day, outperforming prior methods by up to 0.5% on CIFAR-10 and 1.0% on ImageNet through joint optimization.

ABSTRACT

Neural architecture search (NAS) has been an active direction of automatic machine learning (Auto-ML), aiming to explore efficient network structures. The searched architecture is evaluated by training on datasets with fixed data augmentation policies. However, recent works on auto-augmentation show that the suited augmentation policies can vary over different structures. Therefore, this work considers the possible coupling between neural architectures and data augmentation and proposes an effective algorithm jointly searching for them. Specifically, 1) for the NAS task, we adopt a single-path based differentiable method with Gumbel-softmax reparameterization strategy due to its memory efficiency; 2) for the auto-augmentation task, we introduce a novel search method based on policy gradient algorithm, which can significantly reduce the computation complexity. Our approach achieves 97.91% accuracy on CIFAR-10 and 76.6% Top-1 accuracy on ImageNet dataset, showing the outstanding performance of our search algorithm.

Motivation & Objective

• To address the limitation of NAS and AA being optimized independently, where fixed augmentation policies may bias architecture evaluation.
• To identify the coupling between neural architectures and data augmentation policies, showing that optimal policies vary with architecture.
• To propose a joint search framework that simultaneously optimizes architecture and augmentation policies for improved performance.
• To overcome the biased gradient approximation in prior differentiable auto-augmentation methods by introducing a policy gradient-based second-order approximation.
• To demonstrate that joint search significantly improves model accuracy compared to independent NAS and AA optimization.

Proposed method

• Formulates a bi-level optimization problem that alternately updates architecture parameters (via Gumbel-softmax reparameterization) and augmentation parameters (via policy gradient-based second-order approximation).
• Introduces a novel policy gradient-based method to estimate gradients for non-differentiable augmentation parameters, avoiding the biased gradient approximation of prior methods.
• Employs a memory-efficient NAS strategy by sampling a subset of supernet edges at each training iteration to reduce computational cost.
• Uses a differentiable search space that combines DARTS-based architecture search with a rich augmentation policy space including operation types, magnitudes, and application probabilities.
• Applies second-order approximation to the inner optimization loop for both NAS and AA, enabling end-to-end gradient-based optimization.
• Employs a two-stage training process: first, joint search for architecture and augmentation policies; second, full training of discovered models under both default and searched policies to validate performance.

Experimental results

Research questions

• RQ1 Does joint optimization of neural architecture and data augmentation policy lead to better model performance than independent optimization?
• RQ2 Can a policy gradient-based second-order approximation improve the accuracy and efficiency of auto-augmentation compared to biased gradient approximations in prior methods?
• RQ3 How does the performance of NAS models change when evaluated under their own searched augmentation policies versus fixed default policies?
• RQ4 To what extent does the optimal data augmentation policy depend on the specific neural architecture?
• RQ5 Can a joint search framework achieve state-of-the-art accuracy with minimal computational cost (e.g., 1 GPU-day)?

Key findings

• DAAS achieves 97.91% top-1 accuracy on CIFAR-10 and 76.6% top-1 accuracy on ImageNet, surpassing prior NAS methods by up to 0.5% and 1.0% respectively.
• The joint search framework outperforms independent search by 0.2% on average (2.24% vs. 2.38% top-1 error) on CIFAR-10.
• Architecture performance improves by up to 0.5% when trained with the searched augmentation policy instead of default policies, demonstrating the importance of policy-architecture coupling.
• The policy gradient-based method for auto-augmentation outperforms prior differentiable AA methods and even the original AA (Cubuk et al., 2019) on 4 benchmarks, with a search cost of less than 1 GPU-day.
• Ablation studies show that mixed combinations of architectures and policies discovered in isolation perform worse than their matched pairs, confirming the value of joint discovery.
• The proposed method reduces computational cost by simplifying second-order partial derivatives through policy gradient, avoiding the biased gradient approximation of prior approaches.