CountsDiff: 자연수를 기반으로 한 확산 모델 - 계수 기반 데이터 생성 및 보완을 위한 방법

CountsDiff: A Diffusion Model on the Natural Numbers for Generation and Imputation of Count-Based Data

확산 모델은 연속 데이터 및 토큰 기반 데이터 영역에서 생성 작업에 뛰어난 성능을 보여주었지만, 이산적인 순서 데이터에 대한 응용은 아직 발전 단계에 있습니다. 본 논문에서는 자연수 분포를 효과적으로 모델링하도록 설계된 확산 프레임워크인 CountsDiff를 소개합니다. CountsDiff는 생존 확률 스케줄과 명시적인 손실 가중치를 통해 수식을 간소화하여 Blackout 확산 프레임워크를 확장합니다. 이를 통해 기존 확산 모델링 프레임워크에서 찾을 수 있는 설계 매개변수를 사용하여 유연성을 제공합니다. CountsDiff는 기존 계수 기반 영역에서는 찾아보기 어려웠던 최신 확산 모델의 기능을 도입합니다. 여기에는 연속 시간 훈련, 분류기-프리 가이드, 그리고 단조적이지 않은 역방향 경로를 허용하는 Churn/Remasking 역동 특성이 포함됩니다. 본 논문에서는 CountsDiff의 초기 구현을 제시하고, CIFAR-10 및 CelebA와 같은 자연 이미지 데이터 세트에서 이를 검증합니다. 또한, 다양한 설계 매개변수가 복잡하고 잘 연구되었으며 해석 가능한 데이터 영역에 미치는 영향을 분석합니다. 추가적으로, 생물학적 계수 분석을 CountsDiff의 자연스러운 활용 사례로 제시하고, 배아 세포 및 심장 세포 아틀라스에서 단일 세포 RNA-seq 데이터 보완 작업에 CountsDiff를 적용하여 성능을 평가합니다. 놀랍게도, 본 논문에서 제시된 단순한 구현 방식임에도 불구하고, CountsDiff는 최첨단 이산 생성 모델 및 선도적인 RNA-seq 보완 방법의 성능과 동등하거나 그 이상을 보여주며, 향후 최적화된 설계 선택을 통해 더 큰 성능 향상을 기대할 수 있습니다.

Diffusion models have excelled at generative tasks for both continuous and token-based domains, but their application to discrete ordinal data remains underdeveloped. We present CountsDiff, a diffusion framework designed to natively model distributions on the natural numbers. CountsDiff extends the Blackout diffusion framework by simplifying its formulation through a direct parameterization in terms of a survival probability schedule and an explicit loss weighting. This introduces flexibility through design parameters with direct analogues in existing diffusion modeling frameworks. Beyond this reparameterization, CountsDiff introduces features from modern diffusion models, previously absent in counts-based domains, including continuous-time training, classifier-free guidance, and churn/remasking reverse dynamics that allow non-monotone reverse trajectories. We propose an initial instantiation of CountsDiff and validate it on natural image datasets (CIFAR-10, CelebA), exploring the effects of varying the introduced design parameters in a complex, well-studied, and interpretable data domain. We then highlight biological count assays as a natural use case, evaluating CountsDiff on single-cell RNA-seq imputation in a fetal cell and heart cell atlas. Remarkably, we find that even this simple instantiation matches or surpasses the performance of a state-of-the-art discrete generative model and leading RNA-seq imputation methods, while leaving substantial headroom for further gains through optimized design choices in future work.