[논문 리뷰] Joint Distributions for TensorFlow Probability

Dan Piponi, Dave Moore외 1명|arXiv (Cornell University)|2020. 01. 22.

Bayesian Modeling and Causal Inference참고 문헌 16인용 수 11

한 줄 요약

이 논문은 텐서플로 확률 프로브리티(TensorFlow Probability)에서 JointDistributions를 소개하며, 다수의 선언적 및 명령형 인터페이스를 통해 확률적 그래픽 모델을 유연하고 조합 가능한 방식으로 지정할 수 있는 프레임워크를 제공한다. 이는 텐서플로 분포를 확장하여 구조적, 배치 처리, 미분 가능한 공동 분포를 지원함으로써 모델 지정과 추론을 통합하며, 다양한 모델 아키텍처에서 하드웨어 가속을 통한 효율적인 추론을 가능하게 한다. 이는 벡터화된 샘플링과 기울기 계산을 통해 실현된다.

ABSTRACT

A central tenet of probabilistic programming is that a model is specified exactly once in a canonical representation which is usable by inference algorithms. We describe JointDistributions, a family of declarative representations of directed graphical models in TensorFlow Probability.

연구 동기 및 목표

• 텐서플로 확률 프로브리티에서 복잡한 확률 모델을 지정하기 위한 통합적이고 확장 가능한 인터페이스를 제공하는 것.
• 일관된 분포 계약 아래에서 선언적, 명령형, 코루틴 기반의 다양한 모델 지정 방식을 지원하는 것.
• 텐서플로의 배치 및 형태 의미 체계를 구조적 공동 분포에 확장하여 효율적인 벡터화된 추론을 가능하게 하는 것.
• 표준화된 인터페이스를 통해 모델 지정과 추론 알고리즘을 분리함으로써, 미분 가능성과 기울기 지원을 유지하는 것.
• 자동 벡터화를 통해 다중 체인 MCMC와 다중 샘플 변분 경계와 같은 확장 가능한 추론 패러다임을 지원하는 것.

제안 방법

• JointDistributions는 텐서플로 분포 API를 확장하여 텐서의 구조적 출력(튜플, 딕셔너리, 리스트)을 지원하며, 이벤트, 배치, 샘플 형태 의미 체계를 유지한다.
• 세 가지 주요 유형이 구현되어 있다: JointDistributionNamed(딕셔너리 기반), JointDistributionSequential(리스트 기반), JointDistributionCoroutine(제너레이터 기반)으로, 각각 다른 프로그래밍 스타일을 지원한다.
• 샘플 형태의 전파를 통해 모델의 자동 벡터화를 지원하여 효율적인 배치 샘플링과 로그확률 계산을 가능하게 한다.
• sample_distributions를 도입하여 조건부 분포와 그 샘플된 값을 동시에 반환함으로써, 로컬 로그 밀도 평가와 같은 세밀한 확률 계산을 가능하게 한다.
• 시스템은 텐서플로의 자동 미분을 활용하여 모델 파라미터에 대한 log_prob 및 샘플의 기울기를 자동으로 계산하며, HMC와 블랙박스 변분 추론과 같은 기울기 기반 추론을 가능하게 한다.
• 모델 파라미터는 tf.Variable과 TransformedVariable을 통해 관리되어, 미분 가능성과 제약 조건(예: 음수 아님 스케일)을 유지하며, tfp.math.minimize를 통한 엔드 투 엔드 최적화를 가능하게 한다.

실험 결과

연구 질문

• RQ1 확률 모델을 인간이 읽기 쉬운 방식으로 지정하면서도 하드웨어 가속 기반 추론과 천연적으로 호환되는 방식은 무엇인가?
• RQ2 다양한 지정 방식(예: 명령형 대비 선언형)이 하나의 통합 인터페이스 아래에서 상호 운용 가능하도록 하기 위한 설계 패턴은 무엇인가?
• RQ3 복잡한 중첩 모델에서 구조적 공동 분포가 이벤트, 배치, 샘플 형태 의미 체계를 일관되게 유지할 수 있는 방법은 무엇인가?
• RQ4 공동 분포가 자동 벡터화 및 배치 처리를 통해 어느 정도 확장 가능한 추론을 가능하게 할 수 있는가?
• RQ5 모델 파라미터는 어떻게 효율적으로 최적화하면서도 미분 가능성과 제약 조건을 유지할 수 있는가?

주요 결과

• JointDistributions는 상호 운용성에 손상을 주지 않으면서도 딕셔너리 기반, 리스트 기반, 코루틴 기반의 다양한 지정 스타일을 지원하는 통합적이고 확장 가능한 인터페이스를 제공한다.
• 복잡하고 중첩된 공동 분포에 걸쳐 효율적인 벡터화된 샘플링과 로그확률 계산이 가능하며, 하드웨어 가속과 자동 배치 처리 덕분에 성능이 향상된다.
• sample_distributions 메서드를 통해 사용자는 중간 조건부 분포와 그 값에 접근할 수 있어, 구성 요소별 로그 밀도 계산과 같은 세밀한 분석이 가능하다.
• log_prob 및 샘플의 기울기는 역방향 자동 미분을 통해 자동으로 계산되며, 해밀턴 몬테카를로 및 블랙박스 변분 추론과의 원활한 통합이 가능하다.
• tf.Variable 또는 TransformedVariable로 표현된 모델 파라미터는 자동으로 추적되고, 표준 텐서플로 최적화기로 엔드 투 엔드 최적화가 가능하다.
• JointDistribution 객체의 중첩은 네이티브로 지원되어, 구성 요소 자체가 공동 분포인 모듈러이고 계층적인 모델 구축이 가능하다.