[논문 리뷰] POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging

Shishir G. Patil, Paras Jain외 3명|arXiv (Cornell University)|2022. 07. 15.

Advanced Neural Network Applications인용 수 6

한 줄 요약

POET은 메모리 제약이 있는 엣지 디바이스에서 대규모 신경망을 에너지 최적화된 방식으로 훈련할 수 있도록 혼합정수선형계획법(MILP) 기반 프레임워크를 도입한다. 이는 재물리화와 페이지링을 동시에 최적화함으로써 달성된다. 독립적인 방법에 비해 최대 40% 낮은 에너지 소비를 기록하며, 32 KB RAM이 있는 Cortex-M 계열 디바이스에서 ResNet-18과 BERT의 훈련을 가능하게 하여 기존 엣지 훈련 기법들보다 에너지 효율성과 확장성 면에서 뛰어나다.

ABSTRACT

Fine-tuning models on edge devices like mobile phones would enable privacy-preserving personalization over sensitive data. However, edge training has historically been limited to relatively small models with simple architectures because training is both memory and energy intensive. We present POET, an algorithm to enable training large neural networks on memory-scarce battery-operated edge devices. POET jointly optimizes the integrated search search spaces of rematerialization and paging, two algorithms to reduce the memory consumption of backpropagation. Given a memory budget and a run-time constraint, we formulate a mixed-integer linear program (MILP) for energy-optimal training. Our approach enables training significantly larger models on embedded devices while reducing energy consumption while not modifying mathematical correctness of backpropagation. We demonstrate that it is possible to fine-tune both ResNet-18 and BERT within the memory constraints of a Cortex-M class embedded device while outperforming current edge training methods in energy efficiency. POET is an open-source project available at https://github.com/ShishirPatil/poet

연구 동기 및 목표

• 모바일폰과 마이크로컨트롤러와 같은 배터리 구동, 메모리가 부족한 엣지 디바이스에서 대규모 현대 신경망의 훈련을 가능하게 하기 위해.
• 현재는 로지스틱 회귀와 같은 작은 아키텍처로 제한되어 있는 엣지 훈련에서 백프로파게이션의 높은 메모리 및 에너지 요구량을 해결하기 위해.
• 두 가지 활성화 메모리 절감 기법인 재물리화와 페이지링을 통합된 에너지 최소화 프레임워크 내에서 동시에 최적화하기 위해.
• 실제 엣지 배포 환경에서 흔히 발생하는 하드웨어 런타임 제약(예: 유휴 기간 동안 훈련 또는 충전 중 훈련)을 충족하기 위해.
• 기존 딥러닝 모델에 아키텍처 수정 없이도 수학적으로 정확하고 값 유지 가능한 솔루션을 제공하기 위해.

제안 방법

• 메모리 및 런타임 제약 조건 하에서 재물리화와 페이지링 결정을 동시에 최적화하는 혼합정수선형계획문제(MILP)로 엣지 훈련을 수식화한다.
• 대상 하드웨어에서 모델 연산자들을 프로파일링하여 에너지, 메모리, 실행 시간에 대한 세밀한 비용 모델을 구축한다.
• MILP 솔버를 사용하여 메모리 및 시간 예산을 준수하면서 총 에너지 소비를 최소화하는 에너지 최적 스케줄을 찾는다.
• 페이지링 작업의 파ip라이닝을 재물리화와 통합하여 지연을 숨기고 엄격한 런타임 제약을 충족시킨다.
• 연산자 수준의 비용 프로파일을 일반화함으로써 선형 및 비선형 모델, 특히 ResNet-18과 BERT를 포함한 다양한 모델을 지원한다.
• 저비용 연산자(예: ReLU)의 재물리화와 고비용 연산자(예: 컨볼루션)의 페이지링을 통합한 유일한 검색 공간을 활용하여 에너지 오버헤드를 감소시킨다.

실험 결과

연구 질문

• RQ1 재물리화와 페이지링을 동시에 최적화함으로써, 각각을 별도로 사용할 경우에 비해 엣지 디바이스에서 더 낮은 에너지 소비를 달성할 수 있는가?
• RQ2 엄격한 런타임 제약 조건을 충족하면서도 최소 32 KB 메모리가 있는 엣지 디바이스에서 BERT 및 ResNet-18과 같은 복잡한 모델을 훈련시킬 수 있는가?
• RQ3 재물리화와 페이지링의 통합은 히وري스틱 또는 순차적 접근 방식에 비해 에너지 효율성과 메모리 사용에 어떤 영향을 미치는가?
• RQ4 MILP 기반 수식이 현실적인 하드웨어 제약 조건 하에서 엣지 훈련에 대해 증명 가능한 최적 스케줄을 찾을 수 있는가?
• RQ5 메모리 및 시간 예산을 변화시킬 때 에너지 효율성과 런타임 성능 간의 상호 교환 관계는 어떠한가?

주요 결과

• A72 플랫폼에서 타이트한 메모리 예산 조건 하에서 ResNet-18 훈련 시, POET은 전체 메모리 기반 기준 대비 최대 141% 낮은 에너지 오버헤드를 기록한다.
• POET은 전체 메모리 훈련과 비교해 1% 미만의 에너지 오버헤드를 유발하지만, 히وري스틱 기반 방법인 Capuchin은 동일 조건에서 73%에서 141%까지 더 높은 에너지 소비를 기록한다.
• Jetson TX2에서 POET은 선형 모델에 국한된 이전 방법인 POFO 대비 피크 메모리 소비를 8.3% 감소시키고, 처리량을 13% 향상시켰다.
• POET의 통합 솔루션은 재물리화 또는 페이지링을 별도로 사용할 경우에 비해 최대 40% 낮은 에너지 소비를 기록하며, 특히 엄격한 런타임 제약 조건(예: 1 에포크당 0.5 ms)에서 두드러진다.
• 조금의 일반 하드웨어에서도 복잡한 모델인 BERT 및 ResNet-18에 대해 MILP 솔버가 10분 이내에 최적 해를 찾는다.
• POET은 32 KB 메모리만을 갖는 Cortex-M 계열 디바이스에서 ResNet-18과 BERT의 훈련을 성공적으로 수행하여, 개인정보 보호를 위한 현장 내 개인화의 가능성을 입증한다.