기술 요약: OCOO-T – 단순하고 확장 가능한 가상 세포 모델

문제 정의

유전적, 화학적, 사이토카인 섭동(perturbation)에 대한 단일 세포 전사체 반응을 예측하는 것은 계산 생물학의 핵심 과제이며, 이는 신약 개발 및 유전자 조절 네트워크 규명에 중요한 함의를 갖는다. 이 분야의 근본적인 어려움은 섭동 전후의 진정한 세포 수준 대응 관계가 부재하다는 점이다. 단일 세포 RNA 시퀀싱은 파괴적인 방식이므로, 모델이 결정론적인 세포 간 매핑을 학습할 수 없다. 대신, 이 작업은 본질적으로 분포적(distributional)이며, 인구 집단 수준의 유전자 발현 변화를 포착해야 한다.

기존 접근 방식들은 다음과 같은 상당한 트레이드오프에 직착한다:

• 점별 예측 모델 (예: scGPT): 집단 수준의 분포 변화를 포착하는 데 실패하는 경우가 많다.
• 잠재 공간 모델 (예: scGen, CPA, Squidiff): 고차원 프로파일을 압축하기 위해 보조 변이 오토인코더(VAE)에 의존하며, 이는 구조적 복잡성, 잠재적인 정보 병목 현상, 그리고 원래의 발현 공간에서의 엔드 투 엔드(end-to-end) 최적화 방해를 초래한다.
• 분포 매칭 모델 (예: STATE, PerturbDiff): 데이터의 비쌍생성(unpaired) 특성을 다루지만, 확장성 측면에서 어려움을 겪는 경우가 많다. 예를 들어, PerturbDiff는 재생 커널 힐베르트 공간(RKHS) 내의 커널 평균 임베딩에 의존하는데, 이는 수천 개의 유전자로 구성된 긴 전사체 프로파일로 확장될 때 통계적으로 저하된다.
• 롱 프로파일(Long-profile)의 한계: 대부분의 모델은 셀프 어텐션(self-attention)의 이차 복잡도 때문에 고변이 유전자(HVGs, 약 2,000~5,000개)로 입력을 제한하며, 이로 인해 중요한 저발현 유전자 및 경로가 배제된다.

방법론: OCOO-T

OCOO-T는 보조 인코더나 특화된 잠재 공간을 거치지 않고, 연속적인 유전자 발현 공간에서 직접 전사체 섭동 반응을 예측하도록 설계된 미니멀리즘 기반의 흐름 매칭(flow-matching) 가상 세포(AIVC) 모델이다.

핵심 아키텍처

OCOO-T는 연속적인 유전자 발현 프로파일 상에서 작동하는 **바닐라 트랜스포머 스택(vanilla Transformer stack)**을 활용한다. 주요 아키텍처 구성 요소는 다음과 같다:

• 연속 값 인코딩 (Continuous Value Encoding): 이산적 빈닝(binning)을 사용하는 모델과 달리, OCOO-T는 노이즈가 섞인 연속 발현 값($z_t$)을 MLP를 통해 히든 차원으로 직접 투영하여 수치적 정밀도를 보존한다.
• 유전자 정체성 임베딩 (Gene Identity Embeddings): 사전 학습된 단백질 언어 모델 임베딩(예: ESM2)을 유전자 토큰의 의미론적 사전 지식(semantic priors)으로 사용하며, 이를 발현 값과 함께 주입한다.
• 조건화 메커니즘 (Conditioning Mechanisms):
  • 적응형 레이어 정규화 (AdaLN): 섭동 정체성, 용량, 세포주 및 배치 정보를 하나의 조건 벡터로 집계하여 AdaLN을 통해 트랜스포머 블록을 조절한다.
  • 인-컨텍스트 토큰 (In-Context Tokens): 중간 레이어에서 시퀀스 앞에 학습 가능한 토큰을 추가하여 조건별 정보를 명시적으로 전달함으로써, 복잡한 이중 조건 시나리오를 처리하는 능력을 향상시킨다.
  • 암시적 컨텍스트 주입 (Implicit Context Injection): 타겟 세포주가 관찰되지 않은 경우, 모델은 제어 세포 집단의 평균 프로파일을 입력 컨텍스트로 받아 학습된 범주형 임베딩 없이도 분포적 특성을 요약하여 수용할 수 있다.
• 패칭 및 디패칭 (Patching and Depatching): 긴 전사체 프로파일(예: 약 18,000개 유전자)을 처리하기 위해, OCOO-T는 트랜스포머 처리 전 유전자 토큰을 연속적인 패치로 그룹화한 뒤 다시 유전자 수준으로 디패치하는 전략을 사용한다. 이는 백본을 변경하지 않고도 유효 어텐션 길이를 $O(G)$에서 $O(G/s)$로 줄여준다.

생성 프레임워크

OCOO-T는 **정방향 흐름(Rectified Flow)**을 사용하여 섭동 반응 예측을 조건부 연속 시간 디노이징 프로세스로 공식화한다:

• 학습: 순수 가우시안 노이즈와 타겟 섭동 발현 프로파일 사이의 선형 보간(linear interpolation)을 정의한다. 모델은 이 궤적의 속도장(velocity field, $v = x - \epsilon$)을 예측하도록 학습된다.
• 손실 함수: 속도 예측 손실($v$-loss)로 학습되며, 이는 깨끗한 샘플에 가까운 큰-$t$ 타임스텝의 기여도를 가중함으로써 경험적으로 더 나은 생성 품질을 제공한다.
• 추론: 예측된 속도장을 ODE 솔버(Euler 또는 Heun)를 통해 적분하여 노이즈로부터 섭동 프로파일을 생성한다.
• 가이던스 (Guidance): 추론 중에는 섭동과 세포주 조건 모두에 대한 준수도를 균형 있게 조절하기 위해 이중 조건 분류기 없는 가이던스(Classifier-Free Guidance, CFG) 전략이 채택된다.

주요 기여

1. 미니멀리즘 설계: OCOO-T는 효과적인 섭동 반응 모델링이 복잡한 보조 분포 정렬 모듈이나 사전 학습된 오토인코더 없이도 달성될 수 있음을 보여준다. 이는 조건부 공변량(covariates)에 조건화된 평범한 트랜스포머 백본만으로 충분하다.
2. 직접적인 연속 모델링: 연속적인 유전자 발현 프로파일에서 직접 작동함으로써, 잠재 공간 압축과 관련된 정보 병목 현상 및 최적화 불일치를 피한다.
3. 롱 프로파일로의 확장성: 패칭/디패칭 메커니즘을 통해 전체 전사체 입력(예: 18,080개 유전자)을 효율적으로 모델링할 수 있으며, 이는 표준 셀프 어텐션의 이차 복잡도 장벽을 극복한다.
4. 유연한 조건화: 모델은 명시적 공변량 임베딩(관찰된 세포주의 경우)과 평균 제어 프로파일을 통한 암시적 집단 수준 조건화(관찰되지 않았거나 미세한 컨텍스트의 경우)를 모두 지원한다.

실험 결과

OCOO-T는 세 가지 주요 벤치마크인 Tahoe100M(화학적 섭동), Replogle-Nadir(유전적 섭동), PBMC(사이토카인 자극)를 통해 평가되었다.

• 전반적인 성능: OCOO-T는 발현 수준 재구성(MSE, MAE), 섭동 식별(PDS), 차등 발현 회복(DEOver, DEPrec, DirAgr)을 포함한 다양한 지표에서 최첨단(SOTA) 성능을 달성했다.
  • Tahoe100M에서 모든 베이스라인 중 가장 높은 PDCorr(0.952)과 가장 낮은 재구성 오차를 기록했다.
  • Replogle-Nadir에서 PDCorr에 대해 STATE(강력한 집단 수준 베이스라인)와 대등하거나 능가하는 성능을 보였으며, 가장 강력한 차등 발현 회복 점수를 달성했다.
  • PBMC에서 사이토카인 유도 전사 방향성과 차등 반응 구조를 회복하는 데 있어 우수한 능력을 입증했다.
• 롱 프로파일 확장성: 18,080개 유전자를 사용한 Replogle-Nadir 벤치마크 실험에서, 패칭(크기 8, 16, 32)을 사용해도 숏 패널(2,000개 HVG) 모델과 유사한 성능을 유지함을 보여, 패칭 전략의 효율성을 검증했다.
• 조건화 전략: 학습된 세포주 임베딩이 관찰된 컨텍스트에서 최고 성능을 냈으나, 평균 제어 세포 프로파일을 암시적 컨텍스트로 사용하는 방식은 관찰되지 않은 컨텍스트에 대해서도 견고한 예측을 제공하여, 적용 가능성과 벤치마크 점수 사이의 실질적인 트레이드오프를 보여주었다.

의의 및 주장

본 논문은 디노이징 목적 함수와 조건화 체계가 적절히 설계된다면, 아키텍처의 단순함만으로도 높은 충실도의 섭동 반응 예측이 가능하다는 점을 시사한다. 단일 세포 오믹스에 대해 트랜스포머 기반 디노이징의 단순함을 활용함으로써, OCOO-T는 효과적이고 확장 가능한 in-silico 세포 시뮬레이션 프레임워크를 제공한다.

저자들은 자신들의 접근 방식이 커널 기반 방법의 고차원 통계적 장벽과 2단계 VAE 파이프라인의 구조적 복잡성을 피한다는 점을 강조한다. 그러나 현재의 평가 방식이 실제 생물학적 유용성(예: 경로 일관성)을 완전히 포착하지 못할 수 있다는 점을 인정하며, 미지의 개입, 실험실, 프로토콜 등 분포 외(out-of-distribution) 설정에서의 일반화 능력을 테스트하기 위한 추가 연구가 필요함을 겸허히 밝히고 있다.