기술 요약: 단일 세포 전사체에 대한 마스크 재구성 트랜스포머의 스케일링 법칙

문제 정의

신경망 스케일링 법칙(손실, 모델 크기, 데이터 양 사이의 멱함수 관계)은 언어 및 시각 트랜스포머에서는 잘 문서화되어 있으나, 단일 세포 유전체학에서의 존재 여부와 그 거동은 여전히 미개척 영역으로 남아 있다. 단일 세포 커뮤니티는 대규모 파운데이션 모델(예: scVI, scGPT, Geneformer)을 개발해 왔으나, 이들은 주로 다운스트림 작업(주석 달기, 배치 교정)을 통해 평가될 뿐, 사전 학습 손실과 모델 규모 사이의 근본적인 관계에 대해서는 탐구되지 않았다. 확립된 스케일링 법칙 없이는, 성능 병목 현상이 불충분한 모델 용량 때문인지, 데이터 양 때문인지, 혹은 내재적인 생물학적/기술적 노이즈 때문인지 실무자가 판단할 수 없다. 또한, 스케일링 법칙의 점근적 손실 하한(asymptotic loss floor)은 줄일 수 있는 오차에 대한 경험적 상한선을 제공하며, 이는 전사체 데이터의 내재적 엔트로피와 연관될 수 있다.

방법론

저자들은 단일 세포 RNA 시퀀싱(scRNA-seq) 데이터로 학습된 마스크 재구성 트랜스포머의 스케일링 거동에 대한 최초의 체계적인 연구를 제시한다.

• 데이터 소스: 데이터는 CELLxGENE Census에서 추출되었다. 주요 실험은 200,000개의 세포와 512개의 고변동 유전자(HVG)를 사용하는 표준화된 서브셋을 사용한다. 추가 분석에는 다양한 세포 수($D \in \{25k, 50k, 100k, 200k\}$)를 가진 고정된 어휘집(fixed-vocabulary) 서브셋이 사용된다.
• 전처리: 파이프라인은 라이브러리 크기 정규화(목표 합계 $10^4$), $\log(1+x)$ 변환, 그리고 훈련(90%), 검증(5%), 테스트(5%) 세트로의 층화 분할을 포함한다.
• 아키텍처: 유전자 순서가 생물학적 의미를 갖지 않으므로 위치 인코딩을 생략한 순열 불변(permutation-invariant) 트랜스포머 인코더를 채택하였다. 입력은 유전자 식별 임베딩, 관측된 발현 값(또는 학습된 마스크 토큰), 그리고 이진 마스크로 구성된다. 모델은 Pre-LN(pre-layer normalization)과 마스킹된 유전자 값을 재구성하기 위한 선형 예측 헤드를 사용한다.
• 학습 목표: 모델은 마스킹된 위치의 실제 발현 값과 예측된 값 사이의 평균 제곱 오차(MSE)를 최소화하는 마스크 재구성 목표를 사용하여 학습된다(15% 마스킹 비율).
• 실험 설계:
  • 파라미터 스윕(Parameter Sweep): 표준화된 $V=512, D=200k$ 데이터셋에 대해 533개에서 $1.0 \times 10^8$개의 파라미터를 갖는 6가지 모델 크기(XXS에서 L까지)를 학습시킨다. 각 크기는 3개의 무작위 시드에 대해 60,000 스텝 동안 실행된다.
  • 데이터 크기 스윕(Data-Size Sweep): 총 훈련 스텝을 일정하게 유지하면서 데이터셋 크기($D$)를 변화시키는 고정 모델($M$ 크기) 분석을 수행한다.
  • 진단(Diagnostics): 교차 유전자 잔차 이질성(cross-gene residual heterogeneity) 확인 및 소규모 직접 가우시안 음의 로그 가능도(NLL) 보정(calibration)을 수행한다.
• 스케일링 법칙 공식화: 검증 MSE($L$)와 파라미터 수($P$) 사이의 관계를 $L(P) = aP^{-\alpha} + c$로 모델링하며, 여기서 $c$는 줄일 수 없는 손실 하한을 나타낸다. 적합(fit)은 $c$에 대한 그리드 탐색과 선형화된 형태에 대한 최소제곱법을 통해 수행된다.

주요 기여

1. 최초의 체계적 연구: 본 연구는 scRNA-seq에 대한 마스크 재구성 사전 학습의 신경 스케일링 법칙에 대한 최초의 체계적인 조사를 제공한다.
2. 스케일링 입증: 표준화된 $V=512, D=200k$ 스윕에서 멱함수 스케일링이 나타남을 입증하였으며, 홀드아웃 테스트 평가 역시 검증과 동일한 질적 경향을 따름을 보여준다.
3. 포괄적 분석: 고정된 $V$ 분석을 포함하여, 데이터 크기 스윕, 교차 유전자 잔차 이질성 진단, 소규모 직접 가우시안-NLL 보정 등을 포함한 포괄적인 연구를 수행하였다.
4. 예비 엔트로피 추정: 특정 가정하에 정보 이론 단위로 해석되는 점근적 불확실성 하한에 대한 예비 추정치를 제공한다.
5. 엄격성을 위한 로드맵: 손실-대-파라미터 곡선에서 엄격한 전사체 엔트로피 추정치로 전환하기 위해 필요한 필수 실험 요소(매칭된 $V/D$ 스윕, 가능도 기반 목표, 컴퓨트 계정 등)를 식별한다.

결과

• 파라미터 스케일링: 표준화된 스윕에서, 검증 MSE는 XXS에서 M 모델로 갈수록 꾸준히 감소한다. L 모델($\sim 10^8$ 파라미터)은 M 모델($\sim 1.9 \times 10^7$ 파라미터)보다 개선되지 않았으며, 이는 손실 곡선이 평탄해지기 시작했음을 나타낸다.
• 스케일링 지수: 적합된 멱함수 지수는 검증 시 $\alpha \approx 0.266$, 테스트 시 $\alpha \approx 0.268$이며, $R^2$는 약 0.86이다. 이는 자기회귀 언어 모델에서 보고된 지수($\alpha \approx 0.076$)보다 상당히 크지만, 멱함수 스케일링의 존재 측면에서는 질적으로 일치한다.
• 줄일 수 없는 하한: 추정된 줄일 수 없는 손실 하한 $c$는 검증 MSE 기준 $\approx 1.44$, 테스트 MSE 기준 $\approx 1.46$이다.
• 데이터 크기 효과: 고정된 스텝에서 $D$를 변화시키는 고정-$V$ 데이터 크기 스윕은 거의 평탄한 검증 적합($\beta \approx -0.002$)을 보였다. 이는 고정된 훈련 스텝에서 더 큰 데이터셋이 더 적은 유효 에포크(epoch)로 처리되어, 이 특정 설정에서는 깨끗한 $L(D)$ 법칙이 나타나는 것을 방해했음을 시사한다.
• 잔차 이질성: 잔차 오차는 유전자에 따라 강하게 의존적이다. 평균 발현량과 유전자별 MSE 사이의 상관관계는 0.892이며, 유전자 간 변동 계수는 0.730으로, 동분산(homoscedastic) 가우시안 가정이 불완전함을 나타낸다.
• 엔트로피 추정: 동분산 가우시안 근사 하에서, 점근적 하한은 마스킹된 각 유전자 위치당 약 **2.30 비트(bits)**의 줄일 수 없는 불확실성에 해당한다. 이 추정치는 MSE 하한과 변환된 가우시안 NLL 하한 모두로부터 도출되었으나, 저자들은 이들이 사후적 도출로 인해 완전히 독립적이지 않을 수 있음을 언급하였다.

의의 및 주장

본 논문은 단일 세포 파운데이션 모델의 스케일링 거동을 추론하기 위한 정량적 프레임워크를 구축한다고 주장한다. 주요 의의는 다음과 같다:

• 실무적 가이드: $\sim 200,000$개의 세포 데이터셋의 경우, 수천만 개의 파라미터를 가진 모델이 학습 가능한 구조를 포착하기에 충분하며, 현재 설정에서는 $\sim 20M$ 파라미터를 넘어서면 수익 체감(diminishing returns)이 발생함을 시사한다.
• 방법론적 명확성: 데이터 볼륨, 최적화 예산, 손실 해석이 복잡하게 상호작용함을 강조한다. 즉, 깨끗한 데이터 법칙 분석을 위해서는 데이터 크기에 따른 매칭된 최적화 예산이 필요하지만, 고정 스텝 스윕에서는 이것이 달성되지 않았다.
• 생물학적 통찰: 추정된 줄일 수 없는 하한($\sim 2.3$ 비트)은 보편적인 생물학적 상수가 아니라, 생물학적 노이즈, 기술적 노이즈, 그리고 특정 전처리 파이프라인(log-transform, HVG 선택)으로부터의 정보 손실이 결합된 작업 특이적(task-specific) 양임을 나타낸다.
• 향후 방향: 엄격한 전사체 엔트로피 추정을 위해서는 향후 매칭된 $V/D/C$ 실험 행렬, 네이티브 가능도 기반 목표($\mu$와 $\sigma^2$를 예측), 그리고 포괄적인 컴퓨트 계정이 필요함을 역설한다.

저자들은 2.3비트라는 수치가 특정 가정하의 예비 벤치마크이며, 현재의 연구가 데이터 제한적(data-limited) 상황과 모델 제한적(model-limited) 상황을 결정적으로 분리하지 못했음을 강조하며 신중한 태도를 유지한다.