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

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

Generative Adversarial Networks and Image Synthesis참고 문헌 32인용 수 36

한 줄 요약

이 논문은 정규화 흐름과 확산을 활용하여 범주형 데이터를 모델링하기 위한 Argmax Flows와 Multinomial Diffusion을 제안하고, 텍스트 및 세분화 작업에서 더 나은 로그-가능도(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.

연구 동기 및 목표

• Autoregressive 모델을 넘어서는 고차원 범주 분포 학습에 대한 동기 부여.
• 확률적 역 Argmax를 통해 이산 데이터와 연속 밀도 모델 간의 다리를 놓기.
• Argmax Flows와 Multinomial Diffusion 두 프레임워크를 통해 이산 데이터의 직접 학습 가능하게 하기.
• 언어 및 세분화 작업에서 Dequantization 벤치마크보다 로그-가능도가 개선되었음을 입증하기.

제안 방법

• Argmax Flows 도입: 연속 밀도 p(v)와 이를 이산 x로 매핑하는 Argmax 매핑을 결합; Argmax 제약을 만족하는 q(v|x)라는 확률적 역을 학습.
• P(x)와의 내적 불가능한 적분을 통해 변분 하한 ELBO를 정의하고, Argmax 제약 하에서 q(v|x)를 사용하여 최적화.
• 세 가지 q(v|x) 구성: 임계값 기반, 멀티노믹 기반, 그리고 임계값 기반-Gumbel, 각 구성의 로그 밀도 계산.
• Argmax Flows의 Cartesian 곱을 허용하여 서로 다른 기본 표현(K-진수 x를 D 차원에 걸쳐) 다루기.
• Multinomial Diffusion 도입: 고정된 마르코프 체인 q(x_t|x_{t-1})과 노이즈 제거 생성 모델 p(x_{t-1}|x_t)를 가지고 직접 원-핫 범주 벡터에서 정의된 확산 과정.
• KL 기반 변분 목적과 폐쇄형 후사분포 q(x_{t-1}|x_t,x_0)를 통한 학습.

실험 결과

연구 질문

• RQ1 Argmax Flows가 이산 범주 데이터를 연속 공간으로 효과적으로 올려density 모델링을 할 수 있는가?
• RQ2 Argmax의 확률적 역이 이산 데이터의 학습을 가능한가하게 만들고 dequantization보다 더 나은 로그-가능도를 제공하는가?
• RQ3 확산이 이산 데이터에서 직접 정의되고 이산-흐름 벤치마크보다 경쟁력 있거나 우수하게 작용하는가?
• RQ4 Argmax Flows와 Multinomial Diffusion이 언어 및 세분화 작업에서 자기회귀 및 비자기회귀 벤치마크와 어떻게 비교되는가?
• RQ5 표현 형식(Cartesian 곱)이 성능과 효율성에 어떤 영향을 주는가?

주요 결과

• Argmax Flows가 텍스트8과 enwik8에서 로그-가능도 측면에서 균일(dequantization) 및 변분(dequantization)보다 우수하다.
• 자기회귀 Argmax Flows가 언어 데이터에서 여러 VAE 및 비-Argmax 이산 흐름 벤치마크보다 더 나은 성능을 달성한다.
• Multinomial Diffusion은 이산 텍스트 데이터에서 강력한 결과를 내고 세분화 맵에서 ELBO/IWBO에 경쟁력 있는 성능을 보여주며, 한 설정에서 과적합이 나타나 보강이 필요하다고 보고된다.
• 비자기회귀 다항 확산은 일부 설정에서 Argmax 결합 흐름보다 우수할 수 있으며, 확산 경로와 모델의 강인성 간의 trade-off를 강조한다.
• 카르테시안 곱 표현은 대칭성을 차원 수와 교환해 서로 다른 기본 표현(K-진수)에 적응하도록 한다.
• 샘플은 합리적인 텍스트 및 세분화 출력 가능성을 시사하며, 확산 모델의 부수 효과로 철자 교정이 시연된다.