Discrete Flow Matching

Despite Flow Matching and diffusion models having emerged as powerful generative paradigms for continuous variables such as images and videos, their application to high-dimensional discrete data, such as language, is still limited. In this work, we present Discrete Flow Matching, a novel discrete flow paradigm designed specifically for generating discrete data. Discrete Flow Matching offers several key contributions:(i) it works with a general family of probability paths interpolating between source and target distributions; (ii) it allows for a generic formula for sampling from these probability paths using learned posteriors such as the probability denoiser ($x$-prediction) and noise-prediction ($\epsilon$-prediction); (iii) practically, focusing on specific probability paths defined with different schedulers improves generative perplexity compared to previous discrete diffusion and flow models; and (iv) by scaling Discrete Flow Matching models up to 1.7B parameters, we reach 6.7% Pass@1 and 13.4% Pass@10 on HumanEval and 6.7% Pass@1 and 20.6% Pass@10 on 1-shot MBPP coding benchmarks. Our approach is capable of generating high-quality discrete data in a non-autoregressive fashion, significantly closing the gap between autoregressive models and discrete flow models.

한 줄 요약

Discrete Flow Matching은 비자동회귀 방식으로 높은 품질의 이산 데이터를 생성하는 새로운 이산 흐름 모델링 기법이다.

핵심 기여도

일반적인 확률 경로 가설을 사용하여 소스와 타겟 분포 간의 보간을 가능하게 함.
확률 제거기($x$-prediction)와 노이즈 예측기($\epsilon$-prediction)를 활용한 샘플링 공식 제시.
다양한 스케줄러를 사용한 확률 경로 최적화로 이전 이산 확산 및 흐름 모델 대비 생성 퍼플렉시티 개선.
1.7B 파라미터 규모 모델로 HumanEval에서 6.7% Pass@1, 13.4% Pass@10, 1-shot MBPP에서 6.7% Pass@1, 20.6% Pass@10 달성.

핵심 아이디어

기존 이산 확산 및 흐름 모델은 연속 공간에 이산 데이터를 임베딩하거나 복잡한 상태 공간 확산을 사용하는 경향이 있다. Discrete Flow Matching은 이산 데이터 자체에서 확률 흐름을 정의하는 새로운 이론적 프레임워크를 제시한다. 이 모델은 소스(노이즈) 분포와 타겟(데이터) 분포 간의 일반적인 확률 경로를 사용하며, 이 경로는 시간에 따라 변화하는 스케줄러에 의해 정의된다. 특히, 샘플링 알고리즘에서 사용되는 'generating probability velocity'는 연속 Flow Matching과 동일한 형태를 가지며, 이는 이산 공간에서도 유의미한 성능을 보장한다. 이 접근법은 이산 흐름 모델의 이론적 기반을 확장하고, 비자동회귀 생성 방식의 성능을 대폭 향상시킨다.

기술적 접근법

**확률 경로**: 소스와 타겟 분포 간의 보간을 위한 일반적인 확률 경로 가설 사용.
**샘플링 알고리즘**: 학습된 posterior(예: $x$-prediction, $\epsilon$-prediction)를 기반으로 샘플링.
**스케줄러**: 확률 경로와 corrector 단계를 정의하는 스케줄러 $\kappa_t$, $\alpha_t$ 사용.
**모델 규모**: 1.7B 파라미터 규모의 모델을 Llama-2와 CodeLlama과 동일한 데이터셋으로 학습.
**학습 손실**: 확률 제거기의 손실 함수 (식 28)를 사용.
**샘플링 알고리즘**: generating velocity (식 24)와 Algorithm 1을 사용.

주요 결과

**HumanEval**: 6.7% Pass@1, 13.4% Pass@10 (기존 이산 비자동회귀 모델 대비 향상).
**1-shot MBPP**: 6.7% Pass@1, 20.6% Pass@10.
**Generative Perplexity**: Llama-3 8B 기준 9.7 (1.7B 자동회귀 모델 22.3 대비 훨씬 낮음).
**이미지 생성**: CIFAR10에서 완전히 이산적인 생성 성능을 보임.

의의 및 한계

Discrete Flow Matching은 이산 데이터 생성에서 비자동회귀 모델의 한계를 극복하는 중요한 발전을 의미한다. 특히, 이산 흐름 모델의 이론적 기반을 확장하고, 생성 품질을 대폭 향상시킴으로써 자동회귀 모델과의 성능 격차를 줄이는 데 기여한다. 그러나, 샘플링 효율성은 여전히 연속 흐름 모델보다 낮아, 향후 연구 주제로 제시된다. 또한, 확률 경로의 다양한 형태를 탐색하는 것도 중요한 발전 방향으로 언급된다.

실용적 활용

이 모델은 프로그래밍 코드 생성, 자연어 생성, 이미지 생성 등 다양한 이산 시퀀스 생성 작업에 적용 가능하다. 특히, 대규모 언어 모델과 유사한 성능을 비자동회귀 방식으로 달성함으로써, 추론 속도와 에너지 효율성 측면에서 실용적 이점을 제공한다.