[논문 리뷰] REGAL: Transfer Learning For Fast Optimization of Computation Graphs.

Aditya Paliwal, Felix Gimeno외 5명|arXiv (Cornell University)|2019. 05. 07.

Parallel Computing and Optimization Techniques참고 문헌 20인용 수 12

한 줄 요약

REGAL은 실행 시간과 최대 메모리 사용량을 줄이기 위해 신경망 계산 그래프를 최적화하는 데 깊이 있는 강화학습을 활용한 전이 학습 방법을 제안한다. 이는 새로운 그래프에 대해 재학습이 필요 없이 사전에 학습된 최적화기로 다양한 그래프에 일반화하여 초당 높은 품질의 최적화를 달성한다—기존 방법이 소요하는 수시간이 아닌 수초 내로 끝나며, 전통적 및 학습 기반 기준보다 뛰어난 성능을 보인다.

ABSTRACT

We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to an extensive set of baselines, our approach achieves significant improvements over classical and other learning-based methods on these two tasks.

연구 동기 및 목표

• 각 새로운 그래프에 대해 재학습이 필요한 기존 최적화 방법의 비효율성 문제를 해결하기 위해.
• 미리 보지 못한 모델에 대해 피지컬 튜닝 없이도 빠르고 고품질의 계산 그래프 최적화를 가능하게 하기 위해.
• 신경망 계산 그래프에서 실행 시간과 최대 메모리 사용량을 최소화하기 위해.
• 한 번 학습된 최적화기로 다양한 실세계의 텐서플로우 그래프에 일반화하기 위해.

제안 방법

• 다양한 계산 그래프에서 사전에 학습된 깊이 강화학습 에이전트를 통해 최적화 정책을 학습한다.
• 에이전트는 연산자 융합, 레이아웃 변환, 메모리 할당과 같은 결정을 내리는 데 학습한다.
• 전이 학습을 적용하여 동일한 정책 네트워크가 추가 학습 없이도 이전에 본 적 없는 그래프로 일반화할 수 있도록 한다.
• 최적화 결정을 이끄는 데 실행 시간과 메모리 사용량을 조합한 보상 함수를 사용한다.
• 표본 효율성을 향상시키기 위해 경험 재생을 활용한 비정책 기반 강화학습을 통해 정책 네트워크를 학습한다.
• 최적화 컴파일러 파이프라인에 통합되어 추론 시점에 결정을 적용한다.

실험 결과

연구 질문

• RQ1 한 번 학습된 강화학습 정책이 새로운 계산 그래프에 대해 일반화되어 최적화할 수 있는가?
• RQ2 제안된 방법은 실행 시간과 메모리 사용량 측면에서 전통적인 최적화 기법보다 어떻게 비교되는가?
• RQ3 전이 학습을 통해 새로운 모델 최적화 시 재학습이 얼마나 줄어들 수 있는가?
• RQ4 최적화 시간을 크게 단축하면서도 최적화 품질을 유지하거나 향상시킬 수 있는가?

주요 결과

• REGAL은 실행 시간과 최대 메모리 사용량 측면에서 전통적 및 학습 기반 기준과 비교해 유사하거나 뛰어난 최적화 품질을 달성한다.
• 실세계의 텐서플로우 그래프에서 최적화 시간을 수시간에서 수초로 줄였다.
• 피지컬 튜닝 없이도 새로운 그래프에 효과적으로 일반화되어 강력한 전이 능력을 입증했다.
• 다양한 모델에서 실행 시간과 메모리 프로파일을 최소화하는 데 있어 강력한 기준 기반 방법을 능가했다.
• 보상 기반 강화학습 프레임워크가 속도와 메모리 효율성 간의 트레이드오프를 성공적으로 균형 잡았다.