[논문 리뷰] Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

Emiel Hoogeboom, Didrik Nielsen외 3명|arXiv (Cornell University)|2021. 02. 10.

Music and Audio Processing인용 수 79

한 줄 요약

논문은 Argmax Flows와 Multinomial Diffusion을 도입하여 범주형 분포를 학습하고, 텍스트 및 세분화 맵과 같은 이산 데이터를 다룰 수 있도록 normalizing flows와 diffusion 모델을 확장하며, 표준 Dequantization 방법보다 로그 가능도(log-likelihood)가 향상됨을 보인다.

ABSTRACT

Generative flows and diffusion models have been predominantly trained on ordinal data, for example natural images. This paper introduces two extensions of flows and diffusion for categorical data such as language or image segmentation: Argmax Flows and Multinomial Diffusion. Argmax Flows are defined by a composition of a continuous distribution (such as a normalizing flow), and an argmax function. To optimize this model, we learn a probabilistic inverse for the argmax that lifts the categorical data to a continuous space. Multinomial Diffusion gradually adds categorical noise in a diffusion process, for which the generative denoising process is learned. We demonstrate that our method outperforms existing dequantization approaches on text modelling and modelling on image segmentation maps in log-likelihood.

연구 동기 및 목표

• 생성 흐름 및 확산 모델로 고차원의 범주형 데이터를 학습하도록 동기를 부여한다.
• argmax 기반 리프팅과 확률적 역함수를 통해 normalizing flows를 직접적인 범주 분포로 확장한다.
• 범주형 변수에 직접 정의된 확산 모델(Multinomial Diffusion)을 제안한다.
• Dequantization 기준선과 비교하여 언어 모델링 및 세분화 데이터셋에서 로그-가능도 향상을 시연한다.

제안 방법

• Argmax Flows는 연속 밀도 p(v)와 argmax v -> x 매핑, 그리고 argmax 제약을 만족하는 q(v|x) 역함수를 결합한다.
• q(v|x)에 대해 Thresholding, Gumbel 기반, 그리고 Gumbel Thresholding의 세 가지 확률적 역함수를 탐구하며, 각각 훈련용 로그-결정식(log-determinant terms)을 갖는다.
• 최적화는 변분 하한(ELBO)을 이용하며, log P(x) >= E_{v~q(v|x)}[log p(v) - log q(v|x)]의 형태를 가진다.
• Argmax Flows의 기본 공간으로 대칭성을 교환하기 위해 Cartesian 곱을 base-M 표현을 통해 base-K 공간으로 확장한다.
• Multinomial Diffusion은 고정된 잡음 제거 궤적과 닫힌 형태의 사후 확률 q(x_{t-1}|x_t,x_0)를 갖는 하나의 분류 벡터(원-핫) 위에서 직접 확산 과정을 정의한다.
• 모델은 v에 대해 flows를 통한 p(v) (Argmax Flows) 예측 혹은 범주 공간에서 잡음 제거 궤적 p(x_{t-1}|x_t)를 학습하는 두 가지 경로를 사용한다.

실험 결과

연구 질문

• RQ1 Argmax Flows가 이산 범주형 데이터를 연속 공간으로 효과적으로 올려 density 모델링에 사용할 수 있는가?
• RQ2 argmax 제약을 만족하는 확률적 역함수 q(v|x)를 어떻게 구성하고 학습할 수 있는가?
• RQ3 범주형 변수에 정의된 확산 과정(Multinomial Diffusion)이 이산 데이터에 대해 Dequantization 기준선보다 더 나은 로그-가능도를 제공하는가?
• RQ4 텍스트 및 세분화 맵을 모델링할 때 자기회귀(auto-regressive)와 비자기회귀(커플링) 흐름 아키텍처 간의 트레이드오프는 무엇인가?
• RQ5 Argmax Flows의 Cartesian 곱이 큰 범주 수에서 성능과 모델 용량에 어떤 영향을 미치는가?

주요 결과

• Argmax Flows는 text8 및 enwik8에서 로그-가능도(bit per character/byte) 측면에서 균일 분해(dequantization) 대비 일관되게 우수한 성능을 보인다.
• 자기회귀 Argmax Flows는 언어 데이터에서 여러 VAE 및 이산 흐름 기준선보다 더 나은 성능을 달성한다.
• 비자기회귀 Argmax Flows(커플링)는 일부 이산 흐름 기준선을 능가하지만 텍스트에서 자기회귀 변종보다 뒤처지고, Multinomial Diffusion은 특정 설정에서 일부 비자기회귀 Argmax 변종을 능가할 수 있다.
• Cityscapes와 같은 이미지 세분화 맵에서 Argmax Flows와 Multinomial Diffusion은 경쟁력 있는 ELBO와 IWBO 점수를 달성하며 전통적인 Dequantization 방법보다 우수한 성능을 보인다.
• 표준 자기회귀 모델, 연속 자기회귀 모델, 비자기회귀 연속 모델 간의 성능 간격이 존재함을 시사하며, 이산-연속 리라이프 및 자기회귀 구성요소의 향후 연구 방향을 제시한다.