CountsDiff: A diffusion model on the natural numbers for generation and imputation of count-based data

Motivation

• Known: 연속형 데이터에 대한 Gaussian diffusion과 범주형(discrete) 데이터에 대한 categorical diffusion model은 각각 잘 발전되어 있으며, 이미지와 토큰화된 텍스트 등에서 널리 활용되고 있다. Blackout diffusion은 pure-death process를 이용해 자연수 상의 diffusion을 다룬 선행 연구이다.
• Gap: scRNA-seq, ATAC-seq, metagenomic read count 등 생물학적 assay 데이터는 자연수 형태의 abundance 측정값이지만, 기존 방법은 이를 실수로 relax하여 continuous diffusion model을 적용하거나(양자화 문제 발생), 각 값을 독립적인 class로 취급하는 categorical diffusion을 적용(순서 정보 손실, 계산량 급증)해야 하는 한계가 있었다. 또한 Blackout diffusion 자체도 현대적 diffusion 기법(연속시간 학습, guidance, churn 등)이 결여되어 있었다.
• Why: whole-genome sequencing, RNA-seq, ATAC-seq 등 다양한 생물학적 count 데이터에 자연스럽게 적용 가능한 원칙적인 diffusion 프레임워크를 제공함으로써, 기존 continuous/discrete adaptation의 pitfall을 피하고 순서 정보를 보존하면서도 확장 가능한 생성 모델링 및 imputation을 가능하게 한다는 점에서 중요하다.
• Approach: Blackout diffusion의 forward process를 survival probability p(t)로 일반화하여 재매개변수화하고, 여기에 continuous-time training, classifier-free guidance, churn/remasking reverse dynamics, stochastic rounding 등 현대 diffusion model의 도구들을 통합적으로 이식하는 접근을 취한다.

Achievement

1. CountsDiff forward process 정식화: pure-death process를 p-schedule p(t)로 매개변수화하여 존재성(Proposition 3.1)을 증명하고, Blackout diffusion을 특수 사례로 포함시킴.
2. Gaussian diffusion과의 SNR 대응관계 규명: p-schedule 기반 CountsDiff의 signal-to-noise ratio가 Gaussian diffusion의 ᾱt 기반 SNR과 정확히 일치함을 보여, 기존 noise schedule(cosine schedule 등)을 직접 이식 가능하게 함.
3. 현대적 diffusion 기법 통합: continuous-time training/sampling, weighted objective, classifier-free guidance, churn/remasking 역동역학, stochastic rounding을 count 기반 도메인에 최초로 도입.
4. 다중 도메인 검증: synthetic count 데이터, CIFAR-10/CelebA 자연 이미지, fetal 및 heart cell atlas의 scRNA-seq imputation에서 검증하여, 단순한 초기 구현만으로도 state-of-the-art discrete generative model 및 scRNA-seq imputation 기법과 동등하거나 우수한 성능을 달성.

How

• Blackout diffusion의 forward kernel을 inhomogeneous pure-death process R(fw)_{i,j}(t) = iμ(t)(δ_{i-1,j} - δ_{i,j})로 재정의하고, 이항분포(binomial) marginal q(xt|x0)을 유도.
• p-schedule로 cos(πt/2)^2 (Nichol & Dhariwal, 2021)을 채택하여 Gaussian diffusion의 noise schedule을 그대로 이식.
• Kolmogorov forward equation을 활용해 continuous-time formulation을 유도하고, churn/remasking을 통한 non-monotone reverse trajectory를 가능케 하는 attrition schedule 도입.
• Classifier-free guidance 및 explicit loss weighting을 통해 학습 목적함수를 재구성.
• Stochastic rounding 알고리즘을 도입하여 Blackout Diffusion의 실패 모드(failure mode)를 방지.
• CIFAR-10, CelebA에서 설계 파라미터(noise schedule, loss weighting, attrition)의 효과를 분석하고, fetal/heart cell atlas의 scRNA-seq 데이터에 대해 imputation 성능을 벤치마킹.

Originality

• Blackout diffusion을 survival probability p(t) 기반으로 일반화하여, 기존보다 명확하고 유연한 forward process 정식화를 제시.
• 자연수 도메인의 diffusion model과 Gaussian diffusion 간 SNR 동치성을 규명하여, 기존 continuous diffusion 설계 자산(noise schedule 등)을 count 도메인에 이식 가능하게 한 이론적 기여.
• churn/remasking을 통한 non-monotone reverse dynamics를 count 기반 discrete diffusion에 최초로 도입.
• scRNA-seq imputation이라는 생물학적 count 데이터에 diffusion model을 적용한 새로운 응용 사례 제시.

Limitation & Further Study

• 논문에서 스스로 명시하듯, 제안된 것은 "초기 구현(initial instantiation)"으로, 설계 선택(design choice)을 최적화하면 추가적인 성능 향상 여지가 상당히 남아있음.
• 발췌된 본문에서는 실험 섹션(4장 이후) 및 정량적 결과가 충분히 제시되지 않아, 실제 성능 향상 폭이나 baseline 대비 우위의 정도를 명확히 판단하기 어려움.
• 후속 연구로 최적화된 noise schedule, loss weighting, churn 파라미터 탐색 및 더 큰 규모의 생물학적 count 데이터셋(ATAC-seq, metagenomic 등)으로의 확장이 필요해 보임.
• 다차원(multivariate) count 데이터에서 차원 간 상관관계를 어떻게 모델링하는지에 대한 구체적 설명이 발췌 부분에는 부족함.

Evaluation

총평: Blackout diffusion을 일반화하고 현대적 diffusion 기법을 count 도메인에 체계적으로 이식한 견고한 이론적·실용적 기여이며, scRNA-seq imputation이라는 실질적 응용에서도 경쟁력 있는 성능을 보여준 점에서 인상적인 연구이다. 다만 초기 구현 단계임을 저자들도 인정하는 만큼, 후속 최적화 연구의 여지가 크다.