Lingshu-Cell: 가상 세포 구현을 위한 전사체 모델링용 생성형 세포 월드 모델

세포의 상태와 외부 자극에 대한 반응을 컴퓨터로 정밀하게 시뮬레이션하는 '가상 세포' 구현의 기술적 토대를 마련했다. 기존 모델들이 해결하지 못한 단일 세포 데이터의 희소성과 이산적 특성을 Masked Discrete Diffusion 구조로 극복하여 질병 연구 및 신약 스크리닝의 효율성을 극대화할 수 있다.

기존의 단일 세포 AI 모델은 주로 세포를 분류하거나 특징을 추출하는 '정적 표현' 학습에 머물러 있었다. 하지만 실제 생물학적 연구에서는 특정 약물을 투여하거나 유전자를 조절했을 때 세포가 어떻게 변할지 예측하는 '동적 시뮬레이션'이 훨씬 중요하다. 기존 생성 모델은 데이터를 연속적인 수치로 가정하고 가우시안 노이즈를 섞는 방식을 썼는데, 이는 0이 매우 많고 정수 단위로 측정되는 실제 세포 데이터(UMI count)의 특성과 맞지 않아 정확도가 떨어지는 한계가 있었다.

Lingshu-Cell은 이 문제를 해결하기 위해 언어 모델이 문장의 빈칸을 채우는 방식과 유사한 '마스크 기반 이산 확산 모델'을 도입했다. 유전자 발현 데이터의 일부를 무작위로 가리고(Masking), 모델이 주변 유전자들의 정보를 바탕으로 원래 어떤 값이 있었는지 맞히게 함으로써 유전자 간의 복잡한 조절 관계를 학습한다. 이는 데이터의 이산적 성질을 그대로 유지하면서도 수만 개의 유전자가 서로 어떻게 영향을 주고받는지 파악할 수 있게 한다.

또한, 유전자의 순서가 고정되지 않은 생물학적 데이터의 특성을 반영하여 순서에 구애받지 않는 양방향 Transformer 구조를 채택했다. 결과적으로 세포의 종류나 외부 자극 조건을 입력하면, 모델은 학습한 유전자 네트워크 지식을 바탕으로 해당 조건에서 세포가 가질 수 있는 가장 현실적인 전사체 상태를 시뮬레이션해낸다.

Lingshu-Cell은 크게 데이터 양자화, 마스크 확산 프로세스, 시퀀스 압축의 세 단계로 작동한다. 먼저 UMI count 데이터를 281개의 이산 토큰으로 변환하는 양자화 함수 $q(x)$를 사용한다. [0~99 사이의 정수 입력 → 그대로 유지, 100 이상의 값 → 로그 스케일 기반 구간화 → 인덱스 토큰 출력] 과정을 통해 데이터의 넓은 동적 범위를 효율적인 토큰 공간으로 매핑한다.

핵심 아키텍처는 마스크 이산 확산 프로세스를 따르는 Transformer 기반 마스크 예측기 $p_\theta(x_0|x_t)$이다. 순방향에서는 데이터를 점진적으로 마스크 토큰으로 교체하고, 역방향에서는 [마스킹된 시퀀스 입력 → 양방향 Self-Attention 연산 → 각 위치별 원래 토큰의 확률 분포 출력] 과정을 거쳐 데이터를 복원한다. 이를 통해 모델은 특정 유전자 발현이 다른 유전자들에 미치는 조건부 확률 분포를 학습한다.

약 18,000개의 유전자를 동시에 처리하기 위해 임베딩 공간 시퀀스 압축(Embedding-Space Sequence Compression)을 도입했다. [유전자 임베딩 벡터 입력 → 무작위 순열 및 $S$개 단위 그룹화 → 선형 투영($W_{down}$) → 압축된 벡터 출력] 과정을 통해 Transformer가 처리해야 할 시퀀스 길이를 획기적으로 줄인다. 이는 연산 효율성을 높일 뿐만 아니라, 여러 유전자의 신호를 혼합하여 노이즈에 강한 표현력을 갖게 한다.

조건부 생성을 위해 세포 타입, 기증자 정보, 섭동 타겟 등을 특수 토큰으로 정의하여 시퀀스 앞에 추가한다. 추론 시에는 Classifier-Free Guidance(CFG)를 적용하여 [조건부 로짓과 무조건부 로짓의 차이 계산 → 가중치 $w$ 곱셈 → 최종 확률 분포 생성] 과정을 수행함으로써, 섭동에 의한 미세한 발현 변화 신호를 증폭시켜 예측 정확도를 높인다.

Virtual Cell Challenge H1 유전적 섭동 벤치마크에서 26개 팀 중 평균 순위 8.7위로 종합 1위를 차지했다. 특히 MAE(0.052)와 Pearson-Delta(0.306) 지표에서 최고점을 기록하며, 유전자 변이에 따른 전체 전사체의 변화 방향과 크기를 가장 정확하게 예측함을 증명했다.

인간 PBMC 데이터셋을 활용한 사이토카인 자극 반응 예측 실험에서도 scGPT, scVI 등 기존 모델 대비 압도적인 성능을 보였다. 90가지 자극 조건 중 70%를 학습에서 제외한 제로샷(Zero-shot) 환경에서도 PDS(Perturbation Discrimination Score)와 Pearson-Delta 상관관계에서 SOTA 성능을 기록하여 미지의 자극에 대한 일반화 능력을 확인했다.

무조건부 생성 실험에서는 9가지 조직과 5가지 종(인간, 마우스, 원숭이, 제브라피시, 초파리)에 대해 실제 데이터와 통계적으로 거의 일치하는 세포 상태를 생성했다. MMD(Maximum Mean Discrepancy) 지표에서 0.0088을 기록하여 scDiffusion(0.0178)이나 scVI(0.0343)보다 실제 세포 분포를 훨씬 더 정밀하게 재현하는 것으로 나타났다.

Lingshu-Cell의 백본은 LLaMA 스타일의 아키텍처를 따르며, 13개의 Transformer 블록으로 구성된다. 각 블록은 RMSNorm과 SwiGLU 활성화 함수를 사용하며, Rotary Position Embedding(RoPE)을 통해 위치 정보를 주입한다. 모든 선형 레이어는 편향(Bias)을 제거하여 학습 안정성을 높였다.

학습 목적 함수는 마스킹된 토큰 위치에서만 계산되는 Cross-Entropy Loss를 사용하며, 이는 이론적으로 데이터의 Negative Log-Likelihood에 대한 Variational Upper Bound를 최소화하는 것과 같다. 학습 시에는 bfloat16 혼합 정밀도와 AdamW 옵티마이저를 사용하며, 모델 가중치의 지수 이동 평균(EMA)을 유지하여 생성 품질을 안정화했다.

추론 시에는 '생물학적 사전 지식 주입(Biological Prior Injection)' 기법을 활용한다. 외부 데이터셋에서 확인된 섭동 관련 하향 조절(Downregulated) 유전자 정보를 초기 마스크 상태의 일부에 고정값으로 입력함으로써, 모델이 생물학적으로 타당한 섭동 경로를 더 빠르게 탐색하도록 유도하는 전략을 취한다.

현재 모델은 전사체(RNA) 데이터만을 기반으로 하며, 후성유전체(Epigenomic), 단백질체(Proteomic), 공간적 정보(Spatial modality) 등은 통합되지 않았다. 또한, 생성된 결과가 통계적으로 실제와 유사하더라도 그것이 반드시 생물학적 인과관계를 의미하는 것은 아니므로 실험적 검증이 병행되어야 한다.