Architecture-Aware Minimization (A$^2$M): How to Find Flat Minima in Neural Architecture Search

1. 문제 정의 (Problem)

신경망 구조 탐색 (Neural Architecture Search, NAS) 은 주어진 작업에 최적의 신경망 아키텍처를 자동으로 찾는 강력한 패러다임입니다. 특히 미분 가능한 NAS (Differentiable NAS, D-NAS) 방법론 (예: DARTS) 은 효율성 면에서 큰 진전을 이루었으나, 여전히 다음과 같은 근본적인 문제들이 존재합니다.

• 성능 붕괴 (Performance Collapse): 탐색 과정에서 발견된 연속적인 아키텍처 (Supernet) 와 최종 이산화된 (Discrete) 아키텍처 간의 성능 격차가 발생합니다.
• 기하학적 특성의 부재: 가중치 공간 (Weight Space) 에서는 손실 함수의 '평탄한 최소값 (Flat Minima)'이 일반화 성능과 밀접한 연관이 있음이 잘 알려져 있지만, 아키텍처 공간 (Architecture Space) 의 기하학적 구조는 아직 충분히 탐구되지 않았습니다.
• 일반화 실패: 기존 DARTS 기반 방법들은 종종 날카로운 최소값 (Sharp Minima) 에 수렴하여, 작은 아키텍처 변화에도 성능이 급격히 떨어지는 불안정한 구조를 찾곤 합니다.

2. 방법론 (Methodology)

이 논문은 아키텍처 공간의 기하학적 특성을 분석하고, 이를 바탕으로 새로운 최적화 알고리즘을 제안합니다.

가. 아키텍처 공간의 기하학적 분석 (Geometric Analysis)

저자들은 NAS-Bench-201 과 DARTS 검색 공간을 대상으로 다음과 같은 분석 도구를 도입했습니다.

• 아키텍처 이웃 (Architecture Neighborhood): 아키텍처를 시퀀스 또는 인접 행렬로 표현하고, 단일 연산 변경 (Operation Change) 또는 연결 조정 (Connection Adjustment) 을 통해 '반경 (Radius)'을 정의합니다.
• 경로 및 장벽 (Paths and Barriers): 두 개의 서로 다른 아키텍처를 연결하는 최단 경로 (Shortest Path) 를 구성하고, 경로 상의 아키텍처 정확도 변화를 분석합니다.
• 분석 결과:
  • 고성능 아키텍처의 군집화: 높은 정확도를 보이는 아키텍처들은 서로 가까이 모여 있으며 (Clustering), 주변 이웃들의 성능 변화가 작아 **평탄한 지역 (Flat Regions)**을 형성합니다.
  • 저성능 아키텍처의 고립: 성능이 낮은 아키텍처들은 고립되어 있으며, 이웃으로 이동할 때 정확도가 급격히 떨어지는 **높은 장벽 (Accuracy Barriers)**을 가집니다.

나. 제안 알고리즘: A2M (Architecture-Aware Minimization)

가중치 공간에서의 Sharpness-Aware Minimization (SAM) 개념을 아키텍처 공간으로 확장하여, 미분 가능한 NAS 의 이중 최적화 (Bi-level Optimization) 문제를 재정의했습니다.

• 핵심 아이디어: 아키텍처 파라미터 ($\alpha$) 의 경사 하강 업데이트 시, 평탄한 최소값을 찾도록 명시적으로 편향 (Bias) 을 줍니다.
• 수식적 접근:
  • 기존 DARTS 의 1 차 경사 업데이트에 SAM 의 원리를 적용하여, 아키텍처 파라미터를 $\rho_\alpha$만큼 교란 (Perturbation) 시킨 지점에서의 경사를 고려합니다.
  • 유한 차분 (Finite Difference) 기법을 사용하여 아키텍처 공간의 곡률 (Curvature) 성분을 근사적으로 계산하고, 이를 경사 업데이트 식에 통합합니다.
  • A2M 업데이트 공식 (Eq. 8):
    $$\nabla_\alpha L_{val} \approx \nabla_{\tilde{\alpha}} L_{val} + \rho_\alpha \frac{\nabla_{\alpha^+} L_{val} - \nabla_{\alpha^-} L_{val}}{2\epsilon}$$
    여기서 $\tilde{\alpha}$는 SAM 원리에 따라 이동된 아키텍처 파라미터이며, $\alpha^\pm$는 이를 기준으로 미세하게 교란된 지점들입니다.
• 구현: 기존 DARTS 기반 프레임워크에 플러그인 (Plug-in) 형태로 쉽게 통합 가능하며, 추가적인 하이퍼파라미터는 이웃 반경 $\rho_\alpha$ 하나뿐입니다.

3. 주요 기여 (Key Contributions)

1. 아키텍처 공간의 기하학적 특성 규명: 아키텍처 공간 내 '이웃'과 '경로'를 정의하여, 고성능 아키텍처가 평탄한 지역에 군집하고 저성능 아키텍처는 고립되어 있음을 정량적으로 증명했습니다.
2. A2M 알고리즘 개발: 아키텍처 공간의 평탄성 (Flatness) 을 명시적으로 최적화하는 새로운 경사 업데이트 규칙을 분석적으로 유도했습니다. 이는 가중치 공간의 SAM을 아키텍처 공간으로 처음 확장한 사례입니다.
3. 범용성 및 성능 향상: 기존 최첨단 (SotA) DARTS 기반 방법론 (DARTS, PC-DARTS, $\beta$-DARTS 등) 에 A2M 을 적용하여, 다양한 검색 공간 (NAS-Bench-201, DARTS) 과 데이터셋 (CIFAR-10/100, ImageNet-16-120) 에서 일관된 성능 개선을 달성했습니다.

4. 실험 결과 (Results)

• 데이터셋 및 환경: CIFAR-10, CIFAR-100, ImageNet-16-120 에서 NAS-Bench-201 과 DARTS 검색 공간을 대상으로 실험 수행.
• 성능 향상 (평균 Test Accuracy):
  • CIFAR-10: +3.60% 향상
  • CIFAR-100: +4.60% 향상
  • ImageNet-16-120: +3.64% 향상
  • 특히 PC-DARTS 의 경우 ImageNet-16-120 에서 정확도가 2 배 이상 (18.41% → 39.59%) 향상되는 놀라운 결과를 보였습니다.
• 기하학적 검증:
  • A2M 으로 찾은 아키텍처는 주변 이웃 (Radius-1, 2, 3) 에 대해 더 낮은 정확도 장벽 (Accuracy Barriers) 을 보여주었습니다. 이는 A2M 이 실제로 더 평탄하고 안정적인 지역을 찾았음을 의미합니다.
  • DARTS 기반 모델 대비 A2-DARTS 모델의 주변 아키텍처 성능 변동폭이 현저히 작아, 로컬 안정성이 높음을 확인했습니다.
• 계산 비용: A2M 은 아키텍처 파라미터에 대한 역전파 횟수를 1 회에서 4 회로 늘리지만, 가중치 파라미터에 대한 비용은 동일합니다. 전체 실행 시간은 약 1.1 배 ~ 1.3 배 증가하는 것으로 측정되어, 성능 향상에 비해 계산 오버헤드가 매우 낮습니다.

5. 의의 및 결론 (Significance)

• NAS 최적화의 새로운 관점: 이 연구는 NAS 문제를 단순히 아키텍처 조합을 찾는 것을 넘어, **아키텍처 공간의 기하학적 구조 (평탄성)**를 이해하고 이를 최적화 목표에 명시적으로 반영해야 함을 보여줍니다.
• 일반화 능력 증대: 평탄한 최소값을 찾는 것이 아키텍처의 일반화 성능을 결정하는 핵심 요소임을 입증했습니다.
• 실용성: A2M 은 기존 D-NAS 프레임워크에 쉽게 적용 가능한 범용 도구로, 자동화 머신러닝 (AutoML) 분야에서 더 강력하고 안정적인 모델 탐색을 가능하게 합니다.
• 미래 방향: 이 연구는 이산적 NAS (Discrete NAS) 나 다른 작업 (NLP, 적대적 공격 방어 등) 으로 확장할 수 있는 기반을 마련했으며, 계산 효율성을 높이기 위한 SAM 변형 알고리즘 적용 등의 후속 연구를 유도합니다.

요약하자면, 이 논문은 NAS 공간의 '평탄한 지역'이 고성능 아키텍처의 핵심임을 발견하고, 이를 찾기 위한 A2M 알고리즘을 제안함으로써, 기존 DARTS 기반 방법론의 한계를 극복하고 일반화 성능을 획기적으로 개선했습니다.