텐서 컴파일러를 위한 구성 크로스-어텐션을 사용한 그래프 신경망

최근 신경망의 인기가 높아짐에 따라 추론 작업 부하를 효율적으로 제공하는 필요성이 대두되고 있습니다. 신경망 추론 작업 부하는 연산자로 구성된 계산 그래프로 표현할 수 있으며, 이 그래프의 노드는 다차원 텐서를 변환합니다. 텐서는 조합적으로 매우 많은 방법으로 전치(transpose)되거나 타일(tiled)될 수 있으며, 일부 구성을 통해 추론 속도가 가속화될 수 있습니다. 우리는 이러한 목표 계산 그래프의 빠른 구성을 선별할 수 있는 신경 그래프 아키텍처인 TGraph를 제안합니다. 이는 기존의 휴리스틱 기반 컴파일러와 달리 인공지능(AI) 텐서 컴파일러를 나타냅니다. 제안된 해결책은 TpuGraphs의 레이아웃 컬렉션에서 Kendall's $\tau$ 평균을 안정적인 베이스라인의 29.8%에서 TGraph의 67.4%로 개선했습니다. 우리의 연구와 관련된 잠재적 CO$_2$ 배출량 감소를 AI 중심 데이터 센터가 위치한 지역의 총 가구 배출량의 50% 이상에 해당한다고 추정됩니다.