단일 세포 생성을 위한 자기 회귀 트랜스포머 모델의 확장 연구

Scaling an Autoregressive Transformer for Single-Cell Generation

본 연구에서는 단일 세포 유전자 발현 벡터를 생성하는 자기 지도 학습 방법을 다룹니다. 특정 세포 유형에서 얻은 벡터 집합을 기반으로, 해당 세포 유형의 추가적인 유전자 발현 벡터를 생성하는 것을 목표로 합니다. 이 작업에서 우리는 생성된 유전자 발현 벡터의 생물학적 정확성과 사전 훈련 손실의 확장 특성을 모두 분석합니다. 제안하는 모델은 인과 관계 트랜스포머와 학습된 양자화 VAE 토크나이저를 결합하여 사용하며, 교차 엔트로피 손실을 통해 훈련됩니다. 모델의 성능을 평가하기 위해, 특정 세포 유형의 보류된 유전자 발현 벡터를 조건으로 사용하여 유전자 발현 벡터를 생성하고, 생성된 유전자 발현 벡터 분포와 해당 세포 유형의 실제 분포를 비교합니다. 제안하는 아키텍처의 확장 특성을 분석하기 위해 훈련된 파라미터 수와 훈련 데이터 양을 다양하게 조절했습니다. 현재까지 알려진 바로는, 본 연구에서 단일 세포 기반 모델에 대한 최초의 공동 최적화된 두 지수 스케일링 법칙과 계산적으로 최적화된 경계를 제시합니다. 마지막으로, 사전 훈련된 이 모델이 어떻게 교란 반응 예측을 위해 미세 조정될 수 있는지 논의합니다.

We study a self-supervised generation task for single-cell gene expression vectors: given a set of vectors from a cell type, we aim to generate additional gene expression vectors of that cell type. For this task we characterize both the biological fidelity of the generated gene expression vectors and the scaling behavior of the pretraining loss. The model is a causal transformer paired with a learned quantized VAE tokenizer, trained with a cross-entropy loss. To evaluate the model, we condition it on held-out gene expression vectors of a cell type and generate vectors of gene expression, comparing the resulting distribution over gene expression vectors to the ground truth distribution of that cell type. We study the scaling properties of the proposed architecture by varying the number of trained parameters and the amount of training data. To our knowledge, we find the first jointly-fit two-exponent scaling law and compute-optimal frontier for a single-cell foundation model. Finally, we discuss how this pretrained model could be finetuned for perturbation response prediction.