Scalable Single-Cell Gene Expression Generation with Latent Diffusion Models

Motivation

• Known: VAE(scVI 등), GAN(scGAN), diffusion 기반 모델(scDiffusion), 그리고 CFGen과 같이 scVI와 latent space에서의 flow matching을 결합한 방법 등 다양한 단일세포 유전자 발현 생성 모델이 존재한다.
• Gap: 기존 방법들은 유전자에 인위적인 순서를 부여하거나 highly variable genes(HGV)의 제한된 부분집합만 다루어, 유전자 발현이 본질적으로 순서 없는 exchangeable set이라는 생물학적 사실과 모순되며, 조직/종에 따라 유전자 vocabulary가 달라질 때 재학습이 필요하다. 또한 GAN 기반 방법은 학습 불안정성과 mode collapse 위험을 갖는다.
• Why: 유전자 순서에 의존하지 않는 permutation-invariant/equivariant 구조는 다양한 조직·종에 걸쳐 확장 가능하고 재학습 없이 유전자 vocabulary를 유연하게 다룰 수 있어, 대규모·이질적 single-cell atlas 및 perturbation 데이터에 대한 확장 가능한 생성 모델링을 가능하게 한다는 점에서 중요하다.
• Approach: 고정 크기 latent variable을 사용하는 Multi-head Cross-Attention Block(MCAB) 기반 VAE를 인코더(permutation-invariant pooling)와 디코더(permutation-equivariant unpooling)에 동일하게 재사용하고, 이 latent space 위에 Diffusion Transformer(DiT)와 linear interpolants를 활용한 latent diffusion model을 학습시켜 Gaussian prior를 대체한다.

Achievement

1. 순서 불변 VAE 아키텍처 제안: 단일 세트의 고정 크기 permutation-invariant latent variable을 사용하는 fully transformer 기반 VAE를 설계하고, 인코더/디코더 모두에서 동일한 MCAB 모듈을 재사용해 별도의 pooling/unpooling 구조(SetTransformer의 PMA+ISAB, SetVAE 등) 없이 exchangeable 확률분포를 파라미터화한다.
2. latent diffusion prior 도입: Gaussian prior 대신 flow matching loss와 linear interpolants(SiT)로 학습된 DiT 기반 latent diffusion model을 도입해 세포 상태의 복잡한 분포를 더 정교하게 모델링하고 classifier-free guidance를 통한 제어 가능한 생성을 지원한다.
3. 다중 조건 classifier-free guidance: 여러 속성(attribute) 조합에 대한 joint classifier-free guidance 공식을 제안해, CFGen에서 사용하는 additive multi-attribute guidance보다 perturbation 벤치마크에서 우수한 성능을 보임을 입증한다.
4. 관측/섭동(perturbational) 데이터 및 다운스트림 과제에서 성능 우위 입증: 재구성 지표 및 세포 수준 분류와 같은 다운스트림 예측 과제에서 fully transformer 기반 오토인코더의 강점을 보인다.

How

• 인코더: gene ID와 count를 embedding layer로 임베딩한 뒤 IPA(induced set attention) block과 transformer block을 거쳐, MCAB(Multi-head Cross-Attention Block)을 통해 inducing point를 query로 사용해 고정 크기의 permutation-invariant latent token(Gaussian head 파라미터)을 산출.
• 디코더: 동일한 MCAB 모듈을 재사용하되, query로 공유 embedding matrix에서 인덱싱된 gene-specific embedding을 사용하여 요청된 gene ID에 대해 permutation-equivariant한 conditional likelihood(count) 출력을 생성.
• 2단계 학습: 1단계에서 VAE를 학습해 고정 크기 latent token 표현을 얻고, 2단계에서 이 latent token들에 대해 Diffusion Transformer(DiT)를 denoiser로 사용하는 latent diffusion model을 linear interpolants와 flow matching loss로 학습.
• 샘플링은 Scalable Interpolant Transformers(SiT) 라이브러리를 이용해 수행하며, 여러 속성 조건에 대해 joint classifier-free guidance를 적용해 다중 조건부 생성을 지원.
• HLCA 등 observational 데이터셋과 perturbational 데이터셋에서 조건부 생성, 재구성 지표, 세포 수준 분류 등 다운스트림 과제로 평가.

Originality

• 유전자 발현 데이터의 exchangeability 속성을 근본적으로 존중하는 fully transformer 기반 아키텍처로, 기존 방법들의 고정 gene ordering이나 HVG 제한 문제를 구조적으로 해결.
• 단일 MCAB 모듈을 인코더의 permutation-invariant pooling과 디코더의 permutation-equivariant unpooling에 동시에 재사용하는 통합적 설계로, SetTransformer(PMA+ISAB)나 SetVAE와 같이 별도의 아키텍처 구성요소가 필요했던 기존 접근보다 단순화됨.
• VAE의 Gaussian prior를 DiT 기반 latent diffusion model(SiT, linear interpolants)로 대체하여 단백질/분자 도메인에서 성공적이었던 latent diffusion 패러다임을 단일세포 전사체학(transcriptomics)에 최초로 확장 적용.
• 다중 속성 조건에 대한 joint classifier-free guidance 공식을 제안하여 기존 CFGen의 additive guidance 대비 우수함을 실증.

Limitation & Further Study

• Abstract 및 발췌 본문만으로는 정량적 성능 향상 폭, 사용된 구체적 벤치마크(perturbational 데이터셋 종류) 및 baseline과의 상세 비교 수치가 명확히 드러나지 않아 실제 개선 정도를 판단하기 어려움.
• 계산 비용 및 확장성(scalability) 관련 구체적 트레이드오프(예: DiT 및 MCAB의 파라미터 수, 학습 시간)에 대한 정보가 부족함.
• 향후 연구로는 더 다양한 조직·종 간 cross-dataset generalization 검증, 더 복잡한 perturbation combinatorics에 대한 평가, 그리고 생물학적 해석 가능성(biological interpretability) 강화가 필요해 보임.
• gene ID embedding에 기반한 vocabulary 확장성 주장이 실제 매우 큰 유전자 집합이나 다종 간(cross-species) 적용에서도 유지되는지에 대한 추가 검증이 필요.

Evaluation

총평: 단일세포 유전자 발현 데이터의 순서 불변성을 구조적으로 존중하면서 latent diffusion model의 표현력을 결합한 scLDM은 기존 방법들의 근본적 한계를 창의적으로 해결한 견고한 기여로 보이며, 생성 품질과 확장성 측면에서 해당 분야에 의미 있는 진전을 제공한다.