[논문 리뷰] Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic Circuits

Robert Peharz, Steven Lang외 7명|arXiv (Cornell University)|2020. 04. 13.

Machine Learning and Data Classification참고 문헌 34인용 수 27

한 줄 요약

이 논문은 희소한 계산 그래프 대신 단일 einsum 연산을 사용하는 새로운 기법인 Einsum Networks (EiNets)를 소개한다. 이는 희소한 계산 그래프를 대체하여 학습과 추론을 가속화한다. 효율적인 텐서 수축과 자동 미분을 활용함으로써 EiNets는 이전 방법 대비 최대 두 배수의 속도 향상과 메모리 절감을 달성하며, SVHN 및 CelebA와 같은 대규모 데이터셋에서 확장 가능한 학습이 가능해졌고, 복잡한 조건부 질의에 대해 정확하고 효율적인 추론을 지원한다.

ABSTRACT

Probabilistic circuits (PCs) are a promising avenue for probabilistic modeling, as they permit a wide range of exact and efficient inference routines. Recent ``deep-learning-style'' implementations of PCs strive for a better scalability, but are still difficult to train on real-world data, due to their sparsely connected computational graphs. In this paper, we propose Einsum Networks (EiNets), a novel implementation design for PCs, improving prior art in several regards. At their core, EiNets combine a large number of arithmetic operations in a single monolithic einsum-operation, leading to speedups and memory savings of up to two orders of magnitude, in comparison to previous implementations. As an algorithmic contribution, we show that the implementation of Expectation-Maximization (EM) can be simplified for PCs, by leveraging automatic differentiation. Furthermore, we demonstrate that EiNets scale well to datasets which were previously out of reach, such as SVHN and CelebA, and that they can be used as faithful generative image models.

연구 동기 및 목표

• 기존 확률적 회로(PC) 구현 방식이 희소하게 연결된 로그 도메인 계산 그래프를 사용함으로써 발생하는 낮은 확장성과 높은 계산 비용 문제를 해결하기 위해.
• 이전에는 PC 모델에 비해 비현실적이었지만, SVHN 및 CelebA와 같은 대규모 실세계 데이터셋에 대해 효율적인 학습과 정확한 추론을 가능하게 하기 위해.
• 자동 미분을 활용하여 EM 기반 학습을 단순화하여 구현 복잡도를 감소시키기 위해.
• 효율적인 텐서 연산을 통해 곱셈과 덧셈 연산을 통합하는 새로운 PC 아키텍처를 설계하여, 속도와 메모리 효율성을 모두 향상시키기 위해.

제안 방법

• EiNets는 여러 개의 덧셈 및 곱셈 연산을 하나의 고도로 최적화된 텐서 수축으로 통합하는 단일 einsum 연산의 계층적 구조를 사용하여 확률적 회로를 표현한다.
• 학습 중 수치적 안정성을 유지하기 위해 표준 log-sum-exp 기법을 einsum 환경에 확장한 새로운 'log-einsum-exp' 기법을 사용한다.
• 자식 노드가 여러 개인 덧셈 노드의 경우, EiNets는 두 단계 프로세스로 분해한다: 첫 번째로 einsum를 통해 개별 덧셈 연산을 계산하고, 두 번째로 전용 '믹싱 레이어'를 사용해 요소별 볼록 혼합을 적용한다.
• 다중 자식 노드를 캐스케이딩된 단순 덧셈 노드로 과잉 매개변수화한 후 믹싱 레이어를 추가함으로써, 비이진 트리 포함한 임의의 PC 구조를 지원한다.
• 학습은 자동 미분을 통해 가속화되며, 로그우도의 기울기를 직접 계산함으로써 기대값-최대화(EM)의 구현을 단순화한다.
• 추론과 학습은 로그 도메인 계산에서 발생하는 성능 저하를 피하기 위해 자연 도메인에서 수행된다.

실험 결과

연구 질문

• RQ1 단일 einsum 연산을 사용함으로써 확률적 회로의 추론 및 학습 속도와 메모리 효율성이 크게 향상되는가?
• RQ2 EiNets는 이전에는 타당한 확률적 모델이 불가능했던 대규모 이미지 데이터셋인 SVHN 및 CelebA에 확장 가능한가?
• RQ3 기존 방법과 비교해 자동 미분을 사용함으로써 확률적 회로의 EM 학습 구현이 단순화되는가?
• RQ4 EiNets의 성능은 LibSPN 및 SPFlow와 같은 기존 PC 프레임워크와 비교해 학습 시간, 추론 속도, 메모리 소비 측면에서 어떻게 다른가?

주요 결과

• EiNets는 LibSPN 및 SPFlow와 같은 이전 구현 방식 대비 최대 두 배수의 속도 향상과 메모리 절감을 달성한다. 특히 대규모 데이터셋에서 두드러진다.
• 모델은 SVHN 및 CelebA에서 성공적으로 학습되었으며, 이는 EiNets가 이러한 복잡한 데이터셋에서도 충실한 생성 이미지 모델로 기능할 수 있음을 보여준다.
• EiNets의 추론 시간은 SPFlow 대비 최대 세 배수 빠르며, 네트워크 깊이가 깊어질수록도 뚜렷한 향상이 있다.
• 레이어당 단일 einsum 연산을 사용함으로써 CPU 및 GPU 모두에서 매우 효율적인 계산이 가능하며, 믹싱 레이어의 텐서 재구성 및 제로 패딩에 의한 오버헤드가 최소화된다.
• log-einsum-exp 기법은 학습 중 수치적 안정성을 보장하여 자연 도메인에서의 신뢰할 수 있는 최적화를 가능하게 하며 성능 저하 없이 구현된다.
• 자동 미분은 수동 기울기 유도가 필요 없게 하여 EM 학습 과정을 보다 강인하고 쉽게 구현할 수 있도록 한다.