Introduction
- 기존의 efficient network를 디자인하는 방법은 크게 두가지, MobileNet, ShuffleNet과 같은 'Compact Model Design'과 Pruning, Quantization, Distillation을 이용한 'Model Compression'이 있다. 이 방법은 deployment scenario에 따라서 total cost가 선형적으로 증가함. 가령, 스마트폰에 모델을 탑재하려고해도 AP에 따라 specialized model이 달라질 것이다.
- 이를 해결하기 위해 diverse sub-network을 가질 수 있는 'Once for All Network (OFA Network)'를 디자인했음. 플랫폼에 탑재 시에는 ofa network의 sub-network를 선택하였음. 이때 sub-network는 재학습없이 depth, width, kernel size, resolution을 조절하여 선택한다.
- 모델 학습과 NAS를 분리시킴. sub-network들 다 최적화시키고 specialized sub-network를 찾기 때문에 training time은 ofa network를 학습시킬 때만 듬.
- OFA Network를 학습시키는 건 sub-network간 가중치의 joint optimization이 필요하기 때문에 쉽지 않은 문제임. 매번 업데이트 단계에서 gradient를 얻기 위해 모든 sub-network를 계산하는 것은 경우의 수가 너무 많아 사실상 불가능하고, sub-network를 샘플링하는 방법은 정확도가 많이 떨어질 수 있음.
- 이것의 근본적인 문제는 서로 다른 sub-network가 서로 간섭해서 전체 네트워크(ofa network)를 학습하는 과정이 비효율적인 것. 이를 해결하기 위해서 'Progressive Shrinking Algorithm'을 제안함.
- 이 방법은 maximum depth, width, kernel size를 갖는 large neural network(ofa network)를 학습하고 큰 네트워크와 가중치를 '공유'하는 smaller sub-network를 서포트하기 위해 ofa network를 점진적으로 파인튜닝 해줌.
- 여기서 small sub-network를 초기화할 때 larger sub-network에서 중요한 가중치를 뽑아서 쓸 수 있기 때문에 나은 초기화값을 제공해줄 수도 있고, larger sub-network에서 small sub-network로 distillation을 해줄 수도 있음.
- 일반적으로 Pruning은 단일 pruned network를 타겟으로 압축해 정확도 유지하는 것을 목표하지만 ofa network는 모든 sub-network의 정확도를 유지하는 것을 목표로 함.
Method
Problem Formalization
- OFA Network의 가중치를 $W_{o}$라 하고, architectural configuration을 $\{arch_{i}\}$라 하면 문제는 다음과 같이 정의할 수 있음.
- 여기서 $C(W_{o},arch_{i})$는 ofa network에서 $arch_{i}$ configuration을 갖는 sub-network를 고르는 selection scheme이다.
- 전체 학습은 ofa network의 가중치를 최적화하는 방향으로 이뤄지는데, 이는 각 서포팅되는 sub-network가 네트워크를 독립적으로 학습하는 것과 같은 수준의 정확도를 유지하게끔 만들어준다.
Architecture Space
- ofa network는 하나의 모델을 학습시키지만 cnn에서 중요한 네가지(width, depth, kernel, resolution)을 커버하는 서로 다른 크기의 sub-network를 지원함.
- CNN model의 피쳐맵 사이즈는 줄이고 채널 개수는 점진적으로 늘렸다.
- 피쳐맵 사이즈가 줄어든 경우에는 첫번째 레이어의 stride=2로 두어 유닛을 레이어의 시퀀스로 구성했다. 이후의 모든 레이어는 stride=1이 되게끔 구성함.
- 각 유닛이 가질 수 있는 레이어 개수 임의로 두었다. (elastic depth)
- 각 레이어에서 채널 개수, 커널 크기를 임의로 두었다. (elastic width, elastic kernel size)
- 모델이 임의의 인풋 이미지 사이즈를 받을 수 있게 했다. (elastic resolution)
- 예를 들어서, 인풋 이미지 범위는 128-224 with a stride 4, 각 유닛이 가질 수 있는 depth는 ${2,3,4}$, 각 레이어에서 expansion ratio는 ${4,5,6}$, 커널 크기는 {3,5,7}로 잡았다고 하자. 이때 5개의 유닛으로 생성할 수 있는 네트워크 아키텍쳐 개수는 $((3 \times 3)^2 + (3 \times 3)^3 + (3 \times 3)^4)^5 \approx 2 \times 10^{19}$가 된다.
- 그런데 얘네들은 ofa network의 가중치를 모두 ofa network로 본인들의 가중치를 공유하기 때문에 이를 저장하는 데, $5.1M$개의 파라미터가 필요했다고 한다.
Training The Once-For-All Network
Naive Approach
- ofa network를 학습시키는 것은 sub-network 하나에 대응되는 목표함수들을 최적화하는 multi-objective problem으로 볼 수 있다.
- 이 관점에서 ofa network를 전체 objective function의 정확학 gradient를 사용하여 최적화하는 방법을 생각할 수 있다.
- 근데 이건 매 스텝마다 모든 sub-network를 뽑아내야 하는데, sub-network 개수만큼 코스트가 증가할 것이다.
- 다른 접근법은 모두 뽑아내지말고 스텝마다 sub-network를 샘플링하는 방법이다. 이건 sub-network 개수에 의한 코스트 증가는 없다. 하지만 얘네가 가중치를 공유하기 때문에 sub-network끼리 간섭을 주고 이는 정확도를 떨어뜨린다고 프랙티컬하게 확인했다.
- 위의 방법이 다 좋지 않아서 'Progressive Shrinking'을 제안했다.
Progressive Shrinking (PS)
- ofa network는 large sub-network에 서로 다른 크기의 small sub-network가 중첩돼있음. sub-network간 간섭을 피하기 위해, 학습 순서를 large sub-network에서 small sub-network로 했음. 이를 PS라 부른다.
- PS는 먼저 maximum kernel size, depth, width를 갖는 large network를 학습시킴 (큰 네트워크부터)
- 다음에는 smaller sub-network를 서포팅할 수 있게 샘플링 스페이스에 sub-network들을 추가시킴. (위의 그림에서 sample하는 부분)
- 가장 큰 네트워크를 학습시킨 이후에는 각 레이어에서 ${3,5,7}$ 중에서 커널 크기를 선택할 수 있게끔 했고 width, depth는 최대값을 유지했음. resolution은 전체 학습 프로세스동안 탄력적으로 두었는데, 이건 학습의 각 배치에서 서로 다른 이미지 사이즈를 샘플링함으로써 구현됨.
- large network를 학습시킨 이후에는 distillation도 사용했음.
- naive approach와 비교했을 때, ps는 ofa network가 small sub-network를 지원하도록 파인 튜닝된 경우, large sub-network가 이미 잘 훈련돼있기 때문에 small sub-network가 large sub-network를 간섭하는 것을 방지할 수 있음.
- small sub-network의 가중치는 large sub-network와 공유함. 그래서 ps를 사용하면 잘 훈련된 large sub-network의 가장 중요한 가중치를 갖는 small sub-network를 초기화할 수 있기 때문에 학습 과정이 빨라질 수 있음.
Elastic Kernel Size
- 7x7 kernel의 중앙에는 5x5 kernel, 이녀석의 중앙에는 3x3 kernel을 두었음. 그래서 커널 사이즈는 엘라스틱해지는데, 여기서 문제점은 중앙의 서브 커널들이 여러 역할을 수행해야된다는 거. (독립적인 커널로써 역할, 큰 커널의 일부분으로써 역할) 그래서 중앙에 위치한 서브 커널들의 가중치는 역할에 따라 다른 분포는 크기를 가질 필요가 있을 수 있음.
- 그래서 커널의 가중치를 공유할 때 'kernel transformation matrix'를 도입했음. 서로 다른 레이어에 대해서는 별도의 transformation matrix를 사용했음.
Elastic Depth
- $N$개의 레이어를 갖는 유닛에서 $D$개의 레이어를 갖는 sub-network를 뽑아내기 위해서 첫번째 $D$개의 레이어는 남기고, 뒤에 $N-D$개의 레이어는 날리는 방법을 사용했음.
- 따라서 하나의 depth를 설정하는 것은 하나의 조합에만 해당하고, 첫번째 $D$개의 레이어를 갖는 네트워크의 가중치는 큰 모델과 작은 모델간에 공유됨.
Elastic Width
- Width는 채널 개수를 뜻함.
- 각 레이어에서 서로 다른 channel expansion ratio를 고를 수 있게 했음.
- PS scheme에 따라 먼저 full-width model을 학습시키고 partial width를 지원하기 위해서 channel sorting operation을 도입했음. 이는 L1 norm이 큰 건 중요한 채널로 보고 낮은 걸 날린 뒤 채널을 재구성하는 거.
Specialized Model Deployment with Once-For-All Network
- ofa network를 학습시켰으면, 주어진 deployment scenario에 대해 specialized sub-network를 뽑아야 한다. 목표는 정확도는 높고 타겟 하드웨어의 제약 조건을 만족시키는 네트워크를 찾는 것.
- ofa가 학습과 네트워크 서치를 분리시켰기 때문에 여기서는 학습 코스트가 필요하지 않음.
- 주어진 모델에 대해 퀄리티를 빠르게 피드백하기 위해 주어진 뉴럴넷의 latency와 정확도를 예측하는 'neural net twins'를 빌드함. 이는 실제 acc/latency를 측정하기 위해 네트워크를 학습시키면 시간이 너무 오래 걸리니 이를 예측 acc/latency로 대체하겠다는 것.
- 네트워크를 찾는 방법은 'Progressive neural architecture search (ECCV' 18)'의 방법을 따랐다고 함.
- accuracy predictor는 각 레이어에는 400개의 히든 유닛을 갖고, 3개의 레이어로 구성된 MLP를 사용했음. 모델의 각 레이어를 kernel size, expansion ratio에 기반해 원핫 벡터로 인코딩했음. 레이어가 날려진 경우에는 0으로 두고... (자세한 건 appendix.A에..)