[논문 리뷰] Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs

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

Metaheuristic Optimization Algorithms Research참고 문헌 42인용 수 25

한 줄 요약

이 논문은 재훈련 없이도 실세계 텐서플로우 그래프에서 실행 시간과 최대 메모리 사용량을 최소화하는 데 있어 Biased Random-Key Genetic Algorithm(BRKGA)를 안내하기 위해 그래프 신경망을 학습시키는 강화학습 기반 유전 알고리즘인 REGAL을 제안한다. BRKGA에 대한 입력 기반 제안 분포를 학습함으로써 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.

연구 동기 및 목표

• 재훈련 없이 다양한 실세계 모델에서 작동하는 빠르고 일반화 가능한 신경망 계산 그래프 최적화기 개발
• 정적 컴파일러에서 배치 및 스케줄링 결정을 동시에 최적화하여 실행 비용(실행 시간 또는 최대 메모리 사용량)을 최소화
• 기존의 학습 기반 최적화기의 한계를 극복하기 위해 각 그래프당 수 시간의 훈련이 필요하거나 새로운 아키텍처로의 일반화 능력이 떨어지는 문제 해결
• 대규모 그래프에서 몇 초 내에 고품질의 해를 달성하여 최적화 컴파일러의 생산 환경 구현 가능화

제안 방법

• BRKGA의 돌연변이 생성 단계에서 각 노드의 비균일한 베타 분포 파라미터를 예측하기 위해 REINFORCE를 사용한 문맥 기반 밴딧 설정으로 그래프 신경망(GNN)을 훈련한다.
• BRKGA 알고리즘은 균일한 기본값 대신 학습된 입력 의존적 제안 분포를 사용하여 더 높은 품질의 해로 향하는 탐색을 안내한다.
• GNN은 계산 그래프의 구조, 노드 특징(예: 연산 유형, 텐서 크기), 엣지 종속성을 처리하여 배치 및 스케줄링에 대한 노드별 행동 확률을 산출한다.
• 정책은 명시적 지도 학습 없이 목적 함수 값(예: 실행 시간 또는 메모리 사용량)을 희박한 보상 신호로 사용하여 훈련된다.
• 후보 해의 실행 비용을 빠르게 평가하기 위해 훈련 중에 GNN의 피드백을 제공하는 조잡한 정적 비용 모델이 사용된다.
• 학습된 이동 가능한 결정 패턴을 통해 다양한 그래프 구조, 특히 실세계 텐서플로우 모델에 걸쳐 일반화된다.

실험 결과

연구 질문

• RQ1 재훈련 없이도 학습된 정책이 새로운 계산 그래프에서 더 나은 해를 찾을 수 있는가?
• RQ2 GNN 기반 정책 학습이 정적 컴파일러에서 동시 배치 및 스케줄링 최적화의 해 품질을 향상시키는가?
• RQ3 이 방법은 수 시간이 아닌 몇 초 내에 고품질의 결과를 도출할 수 있는가? 이는 실시간으로 생산 컴파일러에 적용 가능한가?
• RQ4 이 모델은 다양한 신경망 아키텍처와 워크로드에 대해 어떻게 일반화되는가?

주요 결과

• REGAL은 실세계 텐서플로우 그래프에서 실행 시간과 최대 메모리 사용량을 모두 최소화하는 데 있어 기존의 전통적 및 학습 기반 기준보다 뚜렷한 개선을 보였다.
• 기존의 학습 기반 접근 방식(예: Mirhoseini et al., 2017, 2018)이 수 시간이 소요되는 데 비해, REGAL은 몇 초 내에 고품질의 해를 생성한다.
• GNN 정책은 데이터 증강을 통해 생성된 새로운 그래프 구조에도 효과적으로 일반화되며, 다양한 아키텍처 간 성능 변동이 매우 낮다.
• 단순 히ュ리스틱을 초월해 데이터에 따라 변하는 전략—예를 들어 고메모리 노드에 낮은 스케줄링 우선순위를 할당하고 무거운 노드에 대해 장치 편향을 피하는 것—을 학습한다.
• TF Runtime 테스트 세트에서 REGAL은 95%의 경우 BRKGA 기준보다 우수한 성능을 보였으며, 보상 값이 -1보다 큰 경우 더 나은 해를 의미한다.