You Only Search Once (YOSO)

작성일 | In |
AutoML 관련 논문 리뷰입니다. 기존 NAS(Neural Architecture Search)들이 강화학습이라면 여기선 NAS에 다른 방식의 optimization 방법(DSO)을 결합한 새로운 AutoML 접근법을 제안하고 있습니다. AutoML 분야에서는 최초로 이미지넷 데이터에 적용했다고 합니다.

You Only Search Once: Single Shot Neural Architecture Search via Direct Sparse Optimization

Introduction

AutoML은 최적의 네트워크 구조를 스스로 학습하는 개념입니다. 2016년 구글이 처음으로 NAS(Neural Architecture Search)를 제안하였습니다. 2017년에는 계산량 관점에서 더욱 효율적인 ENAS를 제안하기도 했습니다. 실제로 지금은 상용화시켜서 기업에 클라우드 AutoML 이라는 이름으로 제품을 팔고 있습니다(구글에 AutoML을 검색하면 가장 위에 나옵니다).

이번에 소개해드릴 논문에서는 NAS와 ENAS처럼 강화학습 기반 방법론이 아닙니다. DSO(Direct Sparse Optimization)과 NAS를 합친 DSO-NAS 방법을 제안하고 있습니다. 단일 모델 하나를 학습시키며, 필요 없는 부분을 쳐내는 pruning 과정이 있습니다. 이 방법을 사용하면 훨씬 간단하게 네트워크 최적화가 가능하다고 합니다. 또한 이전 AutoML 방법론들은 CIFAR 데이터와 같이 비교적 용량이 적고 간단한 데이터에만 적용이 가능했으나, YOSO는 이미지넷 처럼 훨씬 큰 데이터셋에 적용을 할 수 있으며 실제 실험 결과도 제시하고 있습니다.

  • 단일 모델만 학습하면 되는 sparse optimization 기반의 모델 pruning 방법 제안
  • DSO를 위한 효과적이고 효율적인 최적화 기법 제안
  • 다른 NAS 모델에 비해 훨씬 간단하면서도 학습이 빠름을 증명

Network architecture

DSO-NAS의 전체적인 architecture space는 완전 연결된 DAG(Directed Acyclic Graph)입니다. 여기서 노드 사이를 연결한 엣지들을 pruning 하면 우리가 원하는 최적화 된 네트워크 구조가 완성됩니다. 네트워크의 전체 구조는 여러 개의 블록으로 이루어져 있으며, 6개의 노드로 이루어진 하나의 블록은 아래 그림과 같습니다.

$T$개의 노드로 이루어진 DAG에서, 노드 $i$의 output을 나타내면 아래 식과 같습니다. 여기서 ${h}^{(j)}$($j < i$)는 각 노드의 output입니다.

\[{ h }^{ (i) }={\mathcal{O}}^{ (i) }\left( \sum _{ j=1 }^{ i-1 }{ { h }^{ (j) } } \right)\]

여기서 모델 pruning을 위해서는 각 노드끼리 연결된 엣지를 쳐내야 합니다. 이를 위하여 논문에서는 아래 식 처럼 각 엣지마다 scaling factors $\lambda$를 설정하였습니다.

\[{ h }^{ (i) }={ \mathcal{O} }^{ (i) }\left( \sum _{ j=1 }^{ i-1 }{ { \lambda }_{ (j) }^{ (i) }{ h }^{ (j) } } \right)\]
${ \lambda }_{ (j) }^{ (i) }$는 노드 $j$에서 $i$로 가는 엣지의 scaling factor를 나타냅니다. 직관적으로 봤을 때 ${ \lambda }_{ (j) }^{ (i) }$값을 0으로 바꿔주면 자연스럽게 엣지를 잘라낼 수 있습니다. 아래에서 따로 설명 드릴 예정이지만, 이렇게 엣지는 일반적인 네트워크 학습과는 다른 방법으로 sparse optimization을 적용하여 학습을 진행합니다.

블록 학습 과정을 그림으로 나타내면 위와 같습니다. 네트워크 전체는 여러 개의 블록이 차례대로 연결되어 있습니다. 블록 내부에는 $M$개의 레벨(level)과 각 레벨마다 $N$개의 operation이 있습니다. 레벨 내부의 operation은 각자 이전 레벨들의 operation과 input이 입력으로 들어 오고, operation의 output은 다음 레벨들의 operation 그리고 output과 연결되어 있습니다. 다만 같은 레벨의 operation 끼리는 서로 연결되어 있지 않습니다. 그림 (b)처럼 연결된 edge들은 0과 1사이의 값을 가집니다. $b$번째 블록의 $i$번째 레이어의 $j$번째 operation의 output을 수식으로 나타내면 아래와 같습니다.

\[{ h }_{ (b,i,j) }={ \mathcal{O} }_{ (b,i,j) }(\sum _{ m=1 }^{ i-1 }{ \sum _{ n=1 }^{ N }{ { \lambda }_{ (b,m,n) }^{ (i,j) }{ h }_{ (b,m,n) } } } +{ \lambda }_{ (b,0,0) }^{ (i,j) }O _{ (b-1) })\]
수식이 굉장히 복잡해 보이지만 하나씩 뜯어 보면 어렵지 않습니다. $O _{ (b-1) }$는 그림의 input 노드입니다. 마지막 output 노드같은 경우 1x1 conv를 사용하여 차원 수를 줄여주는 역할을 합니다. output node ${O}_{b}$의 수식은 아래와 같습니다. 만약 모든 operation의 엣지가 잘려나가면, 자연스럽게 identity mapping이 됩니다.
\[{ O }_{ (b) }=\mathbb{R}\left( \left[ { \lambda }_{ (b,1,1) }^{ (M+1,0) }{ h }_{ (b,1,1) } \right] ,{ \lambda }_{ (b,1,2) }^{ (M+1,0) }{ h }_{ (b,1,2) },...,{ \lambda }_{ (b,m,n) }^{ (M+1,0) }{ h }_{ (b,m,n) },...,{ \lambda }_{ (b,1,2) }^{ (M+1,0) }{ h }_{ (b,M,N) } \right) \\ +{ O }_{ (b-1) },m\in \left[ 1,M \right] ,n\in \left[ 1,N \right]\]

네트워크의 전체 구조는 아래 그림과 같습니다. 총 $S$개의 스테이지가 있으며, 각 스테이지 내부에는 $B$개의 블록과 spatial size를 줄여 주는 reduction block이 있습니다. 블록 내부는 위에서 설명한 레벨과 operation으로 이루어져 있습니다. Operation은 아래 그림의 총 4가지 방법(seperable convlution, pooling)을 사용하였습니다.

Reduction block의 경우 단순히 커널 사이즈가 1x1과 3x3인 convolution layer 두 개로 되어 있으며, stride를 2로 설정하여 피쳐맵의 사이즈를 1/2로 줄여주었습니다. 두 conv layer를 거친 결과물을 더해 주었다고(elementwise add) 합니다.

아래 식은 네트워크 학습을 위한 최종 방정식과 notation입니다.

\[\min _{ W,\lambda }{ \frac { 1 }{ K } \sum _{ i=1 }^{ K }{ \mathcal{L}\left( { y }_{ i },Net\left( { x }_{ i },W,\lambda \right) \right) } +\delta { \left\| W \right\| }_{ F }^{ 2 }+\gamma { \left\| \lambda \right\| }_{ 1 } }\]
  • ${x}_{i}$ : input data
  • ${y}_{i}$ : label
  • $K$ : number of training samples
  • $W$ : network weight
  • $\lambda$ : edge weight
  • $\delta, \gamma$ : weight of regularization

Sparse Optimization

뉴럴네트워크, 특히 딥러닝처럼 덩치가 큰 네트워크들은 일반적으로,stochastic 방법들로 최적화를 시킵니다. 그러나 논문에서 제시하는 $\lambda$를 동일한 방법으로 제대로 학습시키기는 쉽지 않다고 합니다. 물론 휴리스틱한 방법으로 찾아낼 수는 있긴 하지만 여전히 결과물이 불안정하게 됩니다. 이에 여기서는 SSS(Sparse Structure Selection)를 활용한 새로운 APG(Accelerated Proximal Gradient)라는 방법을 사용하고 있습니다.

\[{ z }_{ (t) }={ \lambda }_{ (t-1) }-{ \eta }_{ (t) }\nabla \mathcal{G}\left( { \lambda }_{ (t-1) } \right) \\ { v }_{ (t) }={ S }_{ \eta (t)\gamma }\left( { z }_{ (t) } \right) -{ \lambda }_{ (t-1) }+{ \mu }_{ (t-1) }{ v }_{ (t-1) }\\ { \lambda }_{ (t) }={ S }_{ \eta (t)\gamma }\left( { z }_{ (t) } \right) +{\mu }_{ (t) }{ v }_{ (t) }\\ { S }_{ \alpha }{ \left( { z } \right) }_{ i }=sign({ z }_{ i })\left( \left| { z }_{ i } \right| -\alpha \right) _{ + }\]

파이토치 코드로 나타내면 아래와 같이 간단하게 나타낼 수 있습니다.

def apg_updater(weight, lr, grad, mom, gamma=0.01):
    z = weight - lr * grad
    def soft_thresholding(x, gamma):
        y = torch.max(torch.tensor(0.), torch.abs(x) - torch.tensor(gamma))
        return torch.sign(x) * y   
    z = soft_thresholding(z, lr * gamma)
    mom = z - weight + 0.9 * mom
    weight = z + 0.9 * mom
    return weight

위 식과 같은 최적화 방법이 APG-NAG 입니다. 여기서 NAG는 Nesterov Accelerated Gradient로 우리가 흔히 알고 있는 momentum optimizer를 살짝 변형한 optimizer 입니다. APG-NAG를 직접적으로 네트워크에 적용할 경우 과적합 되는 경향이 생길 수 있기 때문에, 테스트 셋에서 좋지 않은 결과가 나올 수 있다고 합니다. 이에 데이터셋을 $W$와 $\lambda$ 학습용 두 개로 나누었습니다.

Incorporating different budgets

\[\min _{ W,\lambda }{ \frac { 1 }{ K } \sum _{ i=1 }^{ K }{ \mathcal{L}\left( { y }_{ i },Net\left( { x }_{ i },W,\lambda \right) \right) } +\delta { \left\| W \right\| }_{ F }^{ 2 }+\gamma { \left\| \lambda \right\| }_{ 1 } }\]

네트워크 최적화를 위한 최종 수식에서 엣지와 네트워크 웨이트는 각각 서로 다른 regularization 상수를 달고 있습니다. 이 중에서 각 블록의 FLOPs를 균형있게 다뤄주기 위하여 아래 수식과 같이 $\gamma$를 유동적으로 변화 시켰습니다.

\[{ \gamma }^{ t }=\frac { { FLOPs }^{ t } }{ { FLOPs }_{ block } } \gamma\]

${ FLOPs }_{ block }$는 엣지가 모두 연결된 블록의 FLOPs이며 ${ FLOPs }^{ t }$는 엣지 웨이트 $\lambda$까지 고려한 FLOPs입니다. 적은 FLOPs를 가지는 블록에는 상대적으로 적은 L1 regularization 효과를 주게 됩니다. 이렇게 FLOPs를 고려해주면 블록마다 $\lambda$ 값을 치우치지 않게 분배하는 효과가 생기게 됩니다. 이를 Adaptive FLOPs라고 합니다.

두 번째로는 MAC(Memory Access Cost)를 고려하였습니다.$\gamma$가 $m$번째 레벨의 $n$번째 operation에 iteration $t$마다 다음 수식과 같이 적용됩니다.

\[{ { { \gamma }_{ (m,n) }^{ t } } }=\frac { { MAC_{ (m,n) } } }{ { MAC }_{ max }^{ t } } \gamma\]

${ MAC }_{ max }$는 전체 네트워크의 MAC를 말합니다. 이를 Adaptive MAC라고 하며 더 좋은 성능의 네트워크를 만드는데 도움이 된다고 합니다.

Experiments

학습은 다음과 같은 방식으로 진행되었습니다.

  1. 좋은 모델 웨이트를 위하여 완전 연결된 네트워크를 몇 에폭동안 학습
  2. 엣지 웨이트도 같이 학습시키면서 네트워크 구조 탐색
  3. 최종적으로 결정된 네트워크 구조를 밑바닥부터 학습

1, 2번의 경우 $\lambda$ 학습을 방지하기 위하여 배치 노말라이제이션의 scaling factor를 1로 고정 시켰습니다.

CIFAR-10의 실험 결과는 위와 같습니다. DSO-NAS-share는 블록 내 엣지 웨이트를 서로 공유하는 방식이고, DSO-NAS-full은 엣지 웨이트를 전부 독립적으로 취급 하였습니다. 스테이지 개수 $S$는 3, 블록 $B$는 8, 블록 내부 레벨을 정해주는 $M$은 4로 설정 하였습니다. random의 경우 step2까지 학습하고 이후 웨이트를 랜덤으로 지정해준 것으로 보입니다. 확실히 엣지 웨이트를 따로 학습시키는 결과가 더 좋게 나오는 것을 확인할 수 있습니다. 또한 ENAS와 유사한 성능을 내면서 파라미터 수는 훨씬 줄었으며, DARTS와 비교해서도 성능과 파라미터 모두 더 우수함을 확인할 수 있습니다.

이미지넷 데이터의 실험 결과 또한 SOTA 급의 성능을 보여주고 있습니다. *이 달린 부분은 CIFAR-10으로 미리 학습해둔 블록을 그대로 사용했다고 합니다. 논문에서는 NAS 계열 방법론 중에선 최초로 대형 데이터셋에 적용 했다고 강조하고 있습니다. 아래 그림은 실제로 구조 탐색을 찾아 낸 블록의 모습입니다.

아래 그림은 Adaptive FLOPs를 적용했을 때 어떤 변화가 나타나는지 나타내고 있습니다. 확실히 (b)와 (c)를 보면 adaptive FLOPs를 적용했을 때 같은 FLOPs 대비 더 좋은 결과가 나타나는 것을 확인할 수 있습니다.

마치며

Sparse optimization을 통한 AutoML을 풀어내는데 집중한 논문입니다. 파라미터 수가 상대적으로 적고, 학습 시간도 적은 좋은 결과였지만 ENAS를 모든 방면에서 완벽히 뛰어 넘지는 못한 것 같습니다. 오토케라스처럼 사용자가 쉽게 적용할 수 있는 AutoML api도 많이 나오는 상황인 만큼 발전 가능성이 높은 분야가 아닐까 합니다.

Reference

Related Posts