[Paper 리뷰] HELP: Hardware-Adaptive Efficient Latency Prediction for NAS via Meta-Learning

HELP: Hardware-Adaptive Efficient Latency Prediction for NAS via Meta-Learning

---

• 배포와 device별 제약 조건을 만족하기 위해 Neural Architecture Search (NAS)는 hardware를 고려해야 함
• 기존에는 Lookup table / Latency estimator를 통해 target device의 latency sample을 수집
  -> 서로 다른 사양의 hardware를 많이 요구하기 때문에 비실용적임
• Hardware-adaptive Efficient Latency Predictor (HELP)
  • Device별 latency 추정 문제를 meta-learning 문제로 공식화
  • Latency를 출력하는 black-box를 고려해 모든 device를 embedding 하는 hardware embedding 방식을 제안
  • Hardware embedding을 사용해 Hardware-adaptive Latency Predictor를 meta-learning
• 논문 (NeurIPS 2021) : Paper Link

---

1. Introduction

• NAS가 실제 시나리오에서 적용되기 위해서는 hardware를 고려해야 함
• 메모리, latency, 에너지 소비 등 다양한 device 제약 조건을 만족하는 architecture를 찾아야 함
• Latency가 efficieny 제약으로 가장 많이 고려됨
  
  • Latency를 측정하기 위해서는 target device의 architecture-latency pair를 수집해야 하므로 비용이 많이 소모됨
  • OFA는 고성능의 supernet을 활용해 모델의 훈련 시간을 줄이지만, 각 device별로 latency predictor를 구축해야 함
  • BigNAS는 FLOP을 기준으로 efficieny를 측정하지만, 정확한 efficieny 기준이라고 볼 수 없음
  • Device의 종류와 task에 따라 여러 시간이 소요되므로 latency-constrained NAS의 병목으로 작용함

-> 그래서 효율적인 latency predictor인 HELP를 제안

 

• 알려진 device를 통해 학습한 knowledge를 새로운 device로 transfer 하여 모르는 device에서 latency 추정 시 높은 효율성을 달성
  
  • Latency prediction 문제를 few-shot 회귀 문제로 고려해 주어진 architecture-device pair에서 latency를 추정
  • 새로운 hardware embedding 방법 제안
    - 각 device에서 참조 architecture의 latency를 활용해 모든 device를 embedding
  • 제안된 hardware embedding을 활용해 여러 device에 걸쳐 latency predictor를 학습
    - amortized meta-learning, gradient-based meta-learning을 결합
• Device에 구애받지 않는 hardware embedding을 통해 모든 hardware platform 및 architecture search space에 적용가능
• NAS framework와 HELP를 결합하여 latencty-constrained NAS의 컴퓨팅 병목 현상을 줄일 수 있음
  - MetaD2A, OFA, HAT와 결합했을 때 효과적임

< Overall of HELP (Hardware-adaptive Efficient Latency Predictor) >

• Latency 추정을 few-shot 회귀로 공식화해 주어진 architecture-device pair로부터 latency를 출력
• 참조 architecture의 latency에 따라 device를 embedding 하는 device-agnostic hardware embedding을 제안
• 새로운 device의 latency를 추정하기 위해 few-shot 회귀 모델을 meta 학습하는 HELP를 제안
• HELP를 기존 NAS framework와 결합해 latency-constrained NAS의 병목 현상을 해소

Concept of HELP

2. Method

- Problem Definition

• (목표) Device에서 적은 수의 sample만을 사용하여 architecture-device pair의 latency를 예측하는 모델을 설계하는 것
• Task specification $\tau =\{h^{\tau },{\mathbf{\text{X}}}^{\tau },{\mathbf{\text{Y}}}^{\tau }\}$가 주어졌을때, 주어진 hardware device $h$에 대한 neural architecture $x$의 latency $y$를 추정하는 회귀 모델 $f(x;\theta ):\mathbf{\text{X}}\to \mathbb{R}$를 학습
  • $h^{\tau }\in \mathbf{\text{H}}$ : Hardware device
  • $X^{\tau }\subset \mathbf{\text{X}}$ : Neural architecture 집합
  • $Y^{\tau }\subset \mathbf{\text{Y}}$ : ${\mathbf{\text{X}}}^{\tau }$의 latency 집합
  • 예측값 $f({\mathbf{\text{X}}}^{\tau };\theta )$와 실제 측정값 ${\mathbf{\text{Y}}}^{\tau }$에 대한 Loss $L$을 최소화 : $mi{n}_{\theta }L(f^{\tau }({\mathbf{\text{X}}}^{\tau };\theta ),{\mathbf{\text{Y}}}^{\tau })$
• BUT, 회귀 모델을 학습하는 것은 간단하지 않음
  1. 여러 device에 걸쳐 일반화할 수 없기 때문에 $N$개의 device에 대해 $N$개의 predictor $\{f^{\tau }(\cdot ;{\theta }^{\tau }){\}}_{\tau =1}^N$을 개별적으로 학습하고 sample을 수집해야 함
  2. 회귀 모델을 overfit 하지 않고 신뢰할만한 예측 성능을 얻으려면 각 device 별로 많은 수의 architecture-latency pair를 필요로 함
  3. Device와 architecture에 대한 일반화가 부족하기 때문에 NAS framework는 새로운 device가 제공될 때마다 sample 수집을 반복해야 하므로 시간 소모적임
• Single predictor $f(\cdot ;\theta )$의 사용
  • 적은 수의 architecture-latency 만을 수집하여 ( $X^{\tau }\ll {\mathbf{\text{X}}}^{\tau },Y^{\tau }\ll {\mathbf{\text{Y}}}^{\tau }$ ) 새로운 target device와 architecture에 빠르게 일반화 가능 
  • Meta-learning framework를 통해 device와 architrcture pool $p(\tau )$에서 얻은 knowledge를 transfer 함

Overview of HELP

- Hardware-adaptive Latency Prediction with Device Embedding

• 측정된 latency $y\leftarrow (x,h)$는 device type $h$와 architecture $x$에 따라 달라짐
  • 기존의 latency predictor는 device별로 개별적으로 학습되기 때문에 device 제약조건을 무시하고 $f(x;\theta )$ 형식을 취함
    - 새로운 device를 포함하는 환경에서 Single latency predictor를 학습하면 성능이 저하됨
  • Hardware-conditioned prediction model이 필요 : $f(x,h;\theta ):\mathbf{\text{X}}\times \mathbf{\text{H}}\to \mathbb{R}$
• Hardware-conditioned prediction model
  • 동일한 architecture $x$에 대해서도 device type에 따라 latency를 다르게 예측 가능
  • Platform type에 관계없이 모든 device에 대해 hardware device $h$를 표현하는 것이 중요
    - Hardware device의 물리적 architectrue가 다를 수 있기 때문 (CPU, FPGA..)
  • Hardware device를 주어진 architecture에서 추론된 latency을 출력하는 black-box function으로 고려
• 고정된 참조 neural architecture 집합에 대한 device의 latency : $V_h$
  • $E$ : 참조 neural architecture $\{x_1,x_2,...,x_d\}\subset \mathbf{\text{X}}$의 집합
    - meta-training, meta-test 작업 전반에서 고정되어 사용됨
  • $d$ : 참조 architecture의 수
  • $y_i^{\ast }(x_i,h)=\{y_i(x_i,h)-min(V_h^{(0)})\}/\{max(V_h^{(0)})-min(V_h^{(0)})\}$ : 표준화된 latency 값 (0~1)
    - 이때, $V_h^{(0)}=\{y_1(x_1,h),y_2(x_2,h),...,y_d(x_d,h)\}$
  • 참조 device 집합은 대표성이 있어야 하므로 다양하고 이질적인 것으로 선택
    - 참조 architecture는 search space에서 random sampling

고정된 참조 architecture 집합에 대한 latency

• Hardware device에 대한 black-box 처리 및 latency-based hardware embedding을 통해 자세한 hardware specification을 고려하지 않고 새로운 device를 embedding 가능

- Meta-Learning the Hardware-adaptive Latency Predictor

• 수집된 device 및 architecture pool $p(\tau )$를 활용하여 여러 device에 대한 few-shot 회귀 문제를 해결해야 함
  • Meta-learning을 활용한 hardware-adaptive latency predictor 
    - Task distribution $p(\tau )$에서 $f(x,h;\theta )$를 meta-learning
    - Predictor $f(x,h;{\theta }^{\tau })$를 새로운 task specification $\tau =\{h^{\tau },V_h,X^{\tau },Y^{\tau }\}$가 주어졌을 때 neural architecture $x$에 빠르게 적응시킴
• Meta-training 단계에서는 Episodic Training Strategy 활용
  • Device-architecture pool $p(\tau )$에서 task $\tau$를 random sampling 하여 각 반복마다 few-shot regression 수행
    - Training set $D=\{h^{\tau },X^{\tau },Y^{\tau }\}$ / Test set $\tilde{D}=\{h^{\tau },{\tilde{X}}^{\tau },{\tilde{Y}}^{\tau }\}$
    - $X^{\tau }\subset \mathbf{\text{X}}$, ${\tilde{X}}^{\tau }\subset \mathbf{\text{X}}$ : neural architecture sample의 집합
    - $X^{\tau }$ : few-shot sample의 집합 ($|X^{\tau }|\ll |\mathbf{\text{X}}|$)
    -> 이 집합 간에는 교집합이 없음 ($X^{\tau }\bigcap {\tilde{X}}^{\tau }=\mathrm{\varnothing }$)
    - $Y^{\tau }\subset \mathbf{\text{Y}}$, ${\tilde{Y}}^{\tau }\subset \mathbf{\text{Y}}$ : neural architecture $X^{\tau }$, ${\tilde{X}}^{\tau }$에 대응하는 latency 값의 집합
  • Hardware-adaptive 예측 모델 : $f(X,V^{\tau };{\theta }^{\tau })$

• Latency predictor를 Meta-training 하기 위한 test loss $L(\cdot ;{\tilde{D}}^{\tau })$ 
  
  • Task embedding $V_h^{\tau }$를 사용해서 task-conditioned latency predictor를 얻을 수 있음
  • Amortized meta-learning framework의 사용 : dataset-adaptive performance predictor와 NAS framework를 meta-learn 하는 것이 목표

Meta-training을 위한 test loss

• Target device에서 수집된 latency를 사용해서 few-shot adaptation을 추가로 수행
  • $t$ : $t$번째 inner gradient step
  • $T$ : inner gradient step의 총 수
  • $\alpha$ : multi-dimensional global learning rate vector
  • Meta-learning에서 사용되는 device에 대한 knowledge를 사용하여 새로운 device에 빠르게 적응 가능
  • Meta-device pool과 관련성이 없는 완전히 새로운 device를 접하게 되면 ${\theta }_{(t)}^{\tau }$에 의해 기존에 캡처된 meta-knowledge에서 벗어나게 됨  

Few-shot adaptation을 위한 Inner gradient update

• Hardware-adaptive modulator : $Z^{\tau }=g(V_h^{\tau };\varphi ):{\mathbb{R}}^d\to {\mathbb{R}}^{d_{\theta }}$
  • Initial parameter를 ${\theta }_{(0)}=\theta \ast z^{\tau }$로 modulate 하기 위함
  • ${\theta }_{(0)}$ : hardware $h^{\tau }$에 대한 새로운 initialization
  • Weight 설정 : ${\theta }_{(0)}\leftarrow \theta \circ z^{\tau }$ ($\circ$ : element-wise multiplication operator)
  • Bias 설정 : ${\theta }_{(0)}\leftarrow \theta +z^{\tau }$

Hardware-adaptive modulator를 반영한 Inner gradient update rule

• 최종적인 Meta-learning objective function
  • Hardware-adaptive latency predictor와 shared initial parameter에 대한 modulator를 모두 meta-learn

Meta-learning objective

- Few-shot Adaptation to Unseen Devices (Meta-Test)

• Meta-train 된 latency prediction model $f(\cdot ;\theta )$를 통해 새로운 device $h_{\nu }$에서 architecture의 device 추정
  • Device-conditioned meta-learning을 활용하면 few latency 값 만으로도 새로운 architecture ${\tilde{x}}^{\nu }$의 latency $y^{\nu }$를 측정할 수 있음
  • For example,
    1. (Given) 새로운 device $h^{\nu },\nu =\{h^{\nu },X^{\nu }\}$에서 architecture ${\tilde{x}}^{\nu }$의 latency 예측
    2. Hardware embedding $V_h$ 계산 : 고정된 참조 architecture 집합에서 latency를 얻음
    3. 얻어진 $V_h$를 Latency predictor ${\theta }_{(T+1)}^{\nu }$의 device-optimized parameter로 활용
    4. Device-optimized latency predictor $f(\cdot ,V_{h^{\nu }};{\theta }_{(T+1)}^{\nu })$을 사용해 ${\tilde{x}}^{\nu }$의 latency 예측
  • Meta latency predictor를 NAS method와 결합하여 새로운 device에 대해 latency가 제한된 NAS를 수행 가능

- Computational Complexity of HELP

• Latency predictor의 meta-training은 한 번만 수행되고 얻어진 latency predictor를 조정해서 사용할 수 있음
• 기존 방식은 많은 양의 sample을 수집해 각 target device별로 latency predictor를 학습시켜야 하지만, HELP는 device 당 10개의 sample만 수집하면 됨
  • Time complexity를 $O(DN)$ 에서 $O(N)$으로 줄임
  • $D$ : device 수
  • $N$ : sample 수

Computational complexity 비교

3. Experiment

- Settings

• Search space : NAS-Bench-201, FBNet, MoblieNetV3, HAT
• Meta-training / Meta-test pool : 18 heterogeneous device (CPU, GPU, mobile..)
• Baseline : MAML, Meta-SGD, ANP, BRP-NAS, MetaD2A

- Efficacy of HELP on Few-shot Latency Estimation for Novel Devices

• Meta-knowledge (Meta-SGD, HELP)를 사용하는 것이 처음부터 Latency predictor를 학습하는 것보다 우수한 성능을 보임
  • 특히 HELP는 샘플 수가 적을 때 다른 방식들보다 훨씬 높은 성능 보임
  • Hardware-adaptive initial parameter ${\theta }_0$을 사용하기 때문

Latency Estimation Performance for Unseen Devices

• HELP가 다른 hardware-independent 한 방식들보다 우수한 성능을 보임
  •  Meta-training dataset에서 device의 이질성으로 인해 task conditioning이 더 중요해지기 때문

Effect of the Hardware-adaptive Meta-learning

• Meta-training pool에 있는 device 수가 10개 이상일 때, 새로운 target device에 대해 0.9 이상의 상관관계를 달성

Effect of the Size of the Device Pool

• HELP는 10개의 sample 만으로 GPU에서 0.987, CPU에서 0.989의 Spearman 순위 상관을 얻어 최고의 sample 효율성을 보임

Sample-efficiency of HELP

- End-to-end Latency-constrained NAS with HELP

• HELP가 Pixel2, Titan RTX에서 각각 90배의 샘플 효율성과 9.8배, 9.9배의 계산 효율성을 달성해 최고의 성능을 보임
  
  • HELP+MetaD2A는 125초 만에 최적의 Latency-constrained architecture를 탐색했음

여러 Latency predictor와 MetaD2A를 결합한 Latency-constrained NAS 결과

• HELP가 각 traget device에 대해 총 NAS 비용을 크게 줄여줌

ImageNet, MobileNet search space에서 Latency-constrained NAS 결과

• OFA+HELP는 target device의 총 NAS 비용을 2140배까지 줄임

총 NAS 비용 비교

• Transformer 모델에 대한 NAS를 수행했을 때, HELP는 기존의 HAT보다 200배 더 적은 sample을 사용하면서 경쟁력 있는 Transformer 모델을 얻었음 

Hardware-aware Transformer Architecture Search