Scaling an Autoregressive Transformer for Single-Cell Generation

기술 요약: 단일 세포 생성을 위한 자기회귀 트랜스포머의 스케일링

문제 정의
본 논문은 특정 세포 유형에 대한 현실적인 단일 세포 유전자 발현 벡터를 생성하는 문제를 다룹니다. 기존 방법들은 세포주를 범주형 이름, 단일 예시 벡터, 또는 벡터의 시퀀스로 표현하지만, 희소한 측정치를 보강하고 in-silico 실험을 가능하게 할 수 있는 확장 가능한 자기지도 학습 프레임워크가 필요합니다. 또한, 단일 세포 파운데이션 모델의 사전 학습 손실(pretraining loss)이 모델 크기($N$)와 데이터 크기($D$)에 따라 어떻게 스케일링되는지에 대한 명확한 이해가 부족하며, 이 모델들이 예측 가능한 멱함수 법칙(power-law)을 따르는지 아니면 데이터 다양성 제약으로 인해 정체(plateau)되는지에 대해 상충하는 보고들이 존재합니다.

방법론
저자들은 학습된 양자화 VAE(Variational Autoencoder) 토크나이저와 결합된 자기회귀 트랜스포머 아키텍처를 제안합니다. 이 접근 방식은 순차적으로 훈련되는 두 가지 주요 구성 요소로 이루어집니다:

1. 토크나이저 (RQ-VAE): 잔차 양자화 VAE(Residual Quantized VAE)는 연속적인 단일 세포 유전자 발현 벡터(18,080개 유전자)를 8개의 이산 토큰 시퀀스로 매핑합니다. 인코더는 입력을 256차원 잠재 벡터로 압축하며, 이는 8개의 잔차 코드북을 통해 거친 단계에서 미세한 단계(coarse-to-fine)로 양자화됩니다. 이 토크나이저는 scBaseCount 아틀라스에서 한 번 훈련된 후 고정되어 트랜스포머의 이산 표현 계층 역할을 수행합니다.
2. 트랜스포머 모델: 인과적 트랜스포머(LLaMA 아키텍처 기반)는 이러한 이산 토큰의 시퀀스를 모델링합니다. "문장(sentence)"은 동일한 세포 유형을 가진 32개 세포의 토큰 블록을 학습된 구분자(separator) 토큰으로 연결하여 구성됩니다. 모델은 다음 토큰 예측(next-token prediction)을 통해 크로스 엔트로피 손실을 사용하여 훈련됩니다.
3. 스케일링 그리드: 스케일링 법칙을 조사하기 위해 저자들은 모델 크기($N$)와 데이터 크기($D$)를 독립적으로 변화시키는 통제된 그리드를 구축했습니다.
   • 모델 크기: 고정된 너비 대 깊이 비율을 유지하면서 1.31M에서 83.9M 파라미터에 이르는 5가지 구성을 사용했습니다.
   • 데이터 크기: 약 2억 개의 세포로 구성된 scBaseCount 아틀라스의 고정된 부분 집합에서 자연스러운 세포 유형 구성을 보존하기 위해 비례 샘플링을 적용하여, 38.1M에서 362M 세포까지의 훈련 예산을 설정했습니다.
   • 총계: 22개의 뚜렷한 $(N, D)$ 지점이 평가되었습니다.

주요 기여

1. 태스크 정식화: 이산 토큰화를 활용하여 세포 라인 컨텍스트가 주어졌을 때 합성 유전자 발현 벡터를 생성하는 자기지도 학습 태스크를 정의했습니다.
2. 평가 지표: 평균 발현의 피어슨 상관계수, 평균 절대 오차(MAE), 세포 유형 판별력 점수를 포함하여 생물학적 충실도를 위한 다운스트림 지표를 확립했습니다.
3. 손실 상관관계: 사전 학습 크로스 엔트로피 손실이 다운스트림 생성 품질 지표와 강하게 상관되어 있음을 입증했습니다.
4. 결합 스케일링 법칙: 사전 학습 손실에 대해 결합된 두 지수 함수 형태의 법칙 $L(N, D) = E + A N^{-\alpha} + B D^{-\beta}$를 피팅했습니다. 저자들이 아는 바로는, 이는 독립적으로 변화된 $N \times D$ 그리드 상에서 단일 세포 트랜스포머에 대해 수행된 첫 번째 피팅입니다.
5. 컴퓨팅 최적 프런티어: 피팅된 지수를 바탕으로 컴퓨팅 최적 할당 전략을 도출하였으며, 현재 규모에서는 모델 파라미터를 늘리는 것보다 데이터 볼륨을 늘리는 데 계산 자원을 투입하는 것이 더 효율적임을 시사했습니다.

결과

• 생성 품질: 모델은 경험적 복제본과 매우 유사한 세포 유형 특이적 벡터를 생성합니다. 가장 낮은 검증 손실 체크포인트에서, 생성된 모집단은 홀드아웃(held-out) 경험적 데이터에 대해 0.960의 피어슨 상관계수와 0.041의 MAE를 달 기록했으며, 이는 두 개의 서로 다른 경험적 부분 집합을 비교했을 때의 베이스라인(0.964 및 0.038)에 거의 근접한 수치입니다. 세포 유형 판별력은 0.990으로, 생성기가 단일 모드로 붕괴되지 않았음을 나타냅니다.
• 스케일링 법칙 피팅: 피팅된 스케일링 법칙은 $\alpha \approx 0.81$ (모델 크기) 및 $\beta \approx 0.57$ (데이터 크기)의 지수와 $E \approx 2.90$ nats의 불가피한 손실 하한선(irreducible loss floor)을 산출했습니다. 이 피팅은 0.977의 $R^2$를 달성했습니다.
• 예측력: 피팅된 법칙은 $N$과 $D$를 모두 변화시킨 홀드아웃 실행(run)을 실제 측정값의 0.5% 이내 오차로 예측하여, 모델의 외삽 능력을 검증했습니다.
• 컴퓨팅 최적 프런티어: $\alpha > \beta$이므로, 최적의 할당은 데이터 쪽으로 기울어집니다 ($D^* \propto C^{0.59}$ vs. $N^* \propto C^{0.41}$).

의의 및 주장
본 논문은 단일 세포 스케일링 법칙에 관한 논쟁을 해결한다고 주장합니다. 이전 연구들("다양성 우선" 클러스터)은 단일 세포 모델이 큰 규모에서 정체된다고 보고하며 이를 컨텍스트 다양성의 부족으로 돌렸으나, 저자들은 이것이 용량 병목 현상(capacity bottleneck)에 대한 오해라고 주장합니다. 저자들의 결과에 따르면, $N$과 $D$를 결합하여 변화시킬 때 명확한 멱함수 스케일링이 나타납니다. 이전 연구에서 관찰된 정체 현상은 모델 크기가 고정되었거나 통제되지 않은 상태에서 데이터를 증가시켰을 때, 데이터 항이 용량 항보다 낮아지면서 모델이 추가 데이터를 활용할 능력이 부족해질 때 발생하는 것으로 간주됩니다.

저자들은 이 연구를 단일 세포 파운데이션 모델을 위한 기초적인 단계로 위치시키며, 훈련을 위한 컴퓨팅 최적 프런티어를 제공합니다. 또한 "2단계" 응용 분야로서, 사전 학습된 모델을 섭동 반응(perturbation response) 예측을 위해 미세 조정하는 것을 논의합니다. 그러나 사전 학습 과정에서 관찰된 바에 따르면, 섭동 벤치마크에 대한 홀드아웃 손실이 메인 스케일링 손실과는 다른 "체크마크 모양(checkmark-shaped)" 궤적(감소 후 상승)을 그렸는데, 이는 일반적인 사전 학습 코퍼스와 특정 섭동 컨텍스트 사이의 구성 불일치를 시사합니다. 이는 보고된 결과라기보다 향후 연구를 위한 열린 과제로 남겨두었습니다.

결론적으로, 본 논문은 제안된 아키텍처가 높은 생물학적 충실도를 가진 유전자 발현 벡터를 성공적으로 합성하며, 단일 세포 트랜스포머의 스케일링 동작은 훈련 그리드가 충분히 통제된다면 자연어 처리와 유사하게 예측 가능한 결합 멱함수 법칙을 따른다는 것을 보여줍니다.