[논문 리뷰] Neural Topological Ordering for Computation Graphs

Mukul Gagrani, Corrado Rainone외 8명|arXiv (Cornell University)|2022. 07. 13.

Advanced Graph Neural Networks인용 수 4

한 줄 요약

이 논문은 컴파일러 워크로드에서 피크 메모리 사용량을 최소화하기 위해 방향 무사이클 그래프(DAG)에서 최적의 순서를 학습하기 위한 새로운 엔드 투 엔드 딥러닝 프레임워크인 Topoformer를 제안한다. 상위 순서에 민감한 메시지 전달 방식을 사용하는 비자기적(non-autoregressive), 어텐션 기반 그래프 신경망을 통해, 기존의 기준선들보다 훨씬 빠른 추론 속도를 확보하면서도 최신 기술 수준 또는 경쟁 가능한 성능을 달성한다. 이는 최대 4000배 이상 더 많은 런타임이 필요한 기준선들과 비교할 때 뚜렷한 성능 향상을 보인다.

ABSTRACT

Recent works on machine learning for combinatorial optimization have shown that learning based approaches can outperform heuristic methods in terms of speed and performance. In this paper, we consider the problem of finding an optimal topological order on a directed acyclic graph with focus on the memory minimization problem which arises in compilers. We propose an end-to-end machine learning based approach for topological ordering using an encoder-decoder framework. Our encoder is a novel attention based graph neural network architecture called \emph{Topoformer} which uses different topological transforms of a DAG for message passing. The node embeddings produced by the encoder are converted into node priorities which are used by the decoder to generate a probability distribution over topological orders. We train our model on a dataset of synthetically generated graphs called layered graphs. We show that our model outperforms, or is on-par, with several topological ordering baselines while being significantly faster on synthetic graphs with up to 2k nodes. We also train and test our model on a set of real-world computation graphs, showing performance improvements.

연구 동기 및 목표

• 컴파일러 파이pline에서 피크 메모리 사용량을 최소화하기 위한 DAG 내 최적의 순서를 찾는 조합 최적화 문제를 해결하기 위해.
• 휴리스틱 방법을 학습된 정책으로 대체할 수 있는 확장 가능한 엔드 투 엔드 딥러닝 프레임워크를 개발하기 위해.
• DAG에서 상위 순서에 민감한 전역 메시지 전달을 가능하게 하는 새로운 그래프 신경망 아키텍처인 Topoformer를 도입하기 위해.
• 재현 가능한 벤치마킹을 위한 층화된, 신경망 유사 계산 그래프로 구성된 공개 가능한 합성 데이터셋을 구축하기 위해.
• 실제 산업 컴파일러 워크로드에 일반화 가능할 뿐 아니라, 합성 및 실제 계산 그래프 모두에서 성능 향상을 입증하기 위해.

제안 방법

• 비자기적(NAR) 디코더를 사용하여 순서 집합 위의 확률 분포를 효율적으로 매개변수화함으로써, 추론 속도를 높이는 데 기여한다.
• 에코더는 Topoformer이며, 이는 노드 쌍 간의 상위 순서 관계(예: 조상, 후손, 관련 없음)에 따라 다른 학습 가능한 파라미터를 사용해 메시지 전달을 수행하는 새로운 어텐션 기반 그래프 신경망이다.
• Topoformer는 정방향 및 역방향 간선 전파를 포함한 일곱 가지의 서로 다른 상위 순서 변환을 통해 DAG를 처리함으로써 전역적인 구조적 종속성을 포착한다.
• Topoformer에서 생성된 노드 임베딩은 디코더의 순서 생성을 이끄는 노드 우선순위로 변환된다.
• 모델은 실제 신경망 워크로드를 모방하고 스케일러블한 데이터 생성을 가능하게 하는 층화된 계산 그래프로 구성된 합성 데이터셋에서 훈련된다.
• 훈련은 정확한 동적 프로그래밍(DP) 솔루션을 타겟으로 삼는 지도 학습을 사용하며, 이는 정책 네트워크를 감독한다.

실험 결과

연구 질문

• RQ1 학습 가능한 엔드 투 엔드 딥러닝 프레임워크가 기존 휴리스틱 방법보다 DAG에서 피크 메모리 사용량을 최소화하는 데 더 나은 성능을 보일 수 있는가?
• RQ2 지역적 또는 간선 전용 메시지 전달과 비교했을 때, GNN 아키텍처에서 상위 순서에 민감한 전역 메시지 전달은 DAG에서 얼마나 효과적인가?
• RQ3 비자기적 디코딩 기법이 자기적 기준선들보다 훨씬 빠른 추론 시간을 확보하면서도 경쟁 가능한 성능을 달성할 수 있는가?
• RQ4 층화된 DAG로 구성된 공개 가능한 합성 데이터셋은 기계 학습을 통한 조합 최적화 분야에서 재현 가능한 벤치마킹을 얼마나 잘 가능하게 하는가?
• RQ5 제안된 방법이 합성 그래프를 넘어서 실제 컴파일러 워크로드에도 일반화 가능한가?

주요 결과

• 500노드 합성 그래프에서, 전체 Topoformer 모델은 빔 서치를 사용해 약 3.08 ± 0.51%의 갭을 기록하며, GAT(정방향+역방향) 및 기타 기준선들을 능가한다.
• 1000노드 그래프에서, 빔 서치를 사용한 Topoformer는 약 0.05 ± 0.36%의 갭을 기록하며, GAT(정방향+역방향) 및 기타 모델들을 크게 능가한다.
• 전체 Topoformer 모델은 정방향+역방향 버전 대비 약 1.5~2.5%p의 갭 감소를 기록하여, 전역 메시지 전달의 유용성을 입증한다.
• Topoformer는 500노드 그래프에서 빔 서치 시 평균 추론 시간 1.62초를 기록하며 최신 기술 수준의 성능을 달성했고, 기존 기준선들보다 최대 4000배 이상 더 많은 시간이 소요되는 것을 고려할 때 뚜렷한 성능 향상을 보였다.
• 실제 계산 그래프에서, 모델은 강력한 기준선들보다 일관된 성능 향상을 보이며, 산업적 컴파일러 워크플로우에 실질적인 관련성을 입증했다.
• 제안된 합성 그래프 생성 알고리즘은 확장 가능한 재현 가능한 벤치마킹을 가능하게 하며, 연구 공동체가 공개적으로 이용할 수 있도록 제공된다.