BaLeNAS: 베이지안 학습 규칙을 통한 미분 가능한 아키텍처 탐색

최근 몇 년간, 가중치 공유와 연속적 근사(continuous relaxation)를 통해 계산 비용을 크게 줄일 수 있다는 점에서, 미분 가능한 아키텍처 탐색(Differentiable Architecture Search, DARTS)은 큰 주목을 받았다. 그러나 최근 연구들은 기존의 미분 가능한 NAS 기법들이 단순한 베이스라인을 넘어서지 못하고, 탐색이 진행될수록 성능이 저하되는 아키텍처를 도출한다는 문제점을 지적하고 있다. 본 논문은 아키텍처 파라미터를 직접 최적화하는 대신, 아키텍처 가중치를 가우시안 분포로 연속적으로 확장함으로써 신경망 아키텍처 탐색 문제를 분포 학습 문제로 재정의한다. 이를 통해 기존 코드베이스를 활용하여 메모리 및 계산 부담을 추가로 증가시키지 않고도 자연 경사 변분 추론(Natural-Gradient Variational Inference, NGVI)을 기반으로 아키텍처 분포를 효과적으로 최적화할 수 있다. 우리는 미분 가능한 NAS가 베이지안 원리로부터 어떻게 이점을 얻는지, 특히 탐색의 다양성과 안정성을 향상시키는지에 대해 설명한다. NAS-Bench-201 및 NAS-Bench-1shot1 벤치마크 데이터셋에서의 실험 결과는 제안하는 프레임워크가 상당한 성능 향상을 이끌어낼 수 있음을 입증한다. 또한, 단순히 학습된 파라미터에 대해 argmax를 적용하는 방식이 아닌, 최근 제안된 학습 불필요한 대체 지표(training-free proxies)를 활용하여 최적화된 분포에서 추출한 아키텍처 집합 중에서 최적의 아키텍처를 선택함으로써, NAS-Bench-201 및 NAS-Bench-1shot1 벤치마크에서 최신 기술 수준(SOTA)의 성능을 달성한다. 특히 DARTS 탐색 공간에서 찾은 최적 아키텍처는 CIFAR-10, CIFAR-100, ImageNet 데이터셋에서 각각 2.37%, 15.72%, 24.2%의 경쟁력 있는 테스트 오차를 기록하며 뛰어난 성능을 보였다.