기술 요약: SHIFT – 불완전하고 이질적인 유전체 데이터를 활용한 생존 예측

1. 문제 정의

환자의 생존을 예측하기 위한 유전체 예측 모델은 **구조적 특징 결측(structural feature missingness)**으로 인해 의료 기관 간 일반화에 실패하는 경우가 많습니다. 무작위로 발생하는 결측 데이터와 달리, 유전체 데이터의 결측은 대개 체계적입니다. 즉, 각 센터마다 서로 다른 시퀀싱 패널을 사용하여 측정하는 유전자 집합이 다릅니다. 이는 광범위한 패널(예: 197개 유전자)로 학습된 모델을 좁은 범위의 패널(예: 22개 유전자)을 사용하는 기관에 직접 배포할 때, 상당한 데이터 손실이나 임퓨테이션(imputation, 결측치 보정) 없이는 불가능하다는 시나리오를 만듭니다.

기존의 접근 방식은 이러한 과제에 대해 일반적으로 세 가지 차선책 중 하나에 의존합니다:

1. 분석 범위 제한: 모든 코호트에서 공유되는 특징들의 교집합으로 분석을 제한하여, 잠재적으로 유익한 유전체 데이터를 폐기합니다.
2. 환자 제외: 불완전한 프로필을 가진 환자를 제외함으로써 코호트의 크기와 다양성을 줄이고 선택 편향을 증폭시킵니다.
3. 테스트 시점 임퓨테이션 (예: KNN, 평균 임퓨테이션): 이는 생물학적 신호를 왜곡하고, 편향을 도입하며, 특히 대규모 특징 블록이 구조적으로 부재할 때 임상적으로 유의미한 특징을 환각(hallucinate)할 수 있습니다.

본 논문은 이러한 한계가 정밀 종양학 분야에서 견고한 다기관 생존 모델의 배포를 저해한다고 주장합니다.

2. 방법론: SHIFT 모델

저자들은 임퓨테이션 없이 불완전한 유전체 입력값으로부터 직접 결과를 예측하도록 설계된 결측 인식 생존 모델인 SHIFT(Survival prediction Handling Incomplete Features using Transformer)를 제안합니다.

아키텍처

SHIFT는 표 형식의 유전체 데이터에 적응된 트랜스포머 기반 아키텍처로, 세 가지 주요 구성 요소로 이루어져 있습니다:

1. 특징 임베딩 블록 (Feature Embedding Block): 전체 유전체 벡터를 공동으로 처리하는 대신, SHIFT는 각 유전 변수를 자기 정규화 네트워크(Self-Normalizing Network, SNN) 블록(Linear $\to$ SELU $\to$ AlphaDropout)을 사용하여 독립적인 토큰 표현으로 매핑합니다. 각 특징 $x_l$은 전용 $SNN(1, D)$(여기서$D=128$)를 통해 임베딩됩니다. 누락된 특징은 플레이스홀더로서 0으로 채워지지만, 유효한 입력으로서 처리되지는 않습니다.
2. 마스크드 트랜스포머 인코더 (Masked Transformer Encoder): 2계층 트랜스포머 인코더가 특징 토큰의 시퀀스를 처리합니다. 결정적으로, 이 모델은 키-패딩 마스크(key-padding mask)로서 **결측 마스크($M$)**를 사용합니다. 이는 결측된 특징에 해당하는 토큰이 셀프 어텐션(self-attention) 메커니즘에서 제외되도록 보장합니다. 결과적으로 학습 가능한 [CLS] 토큰은 관찰된 유전 변수로부터만 정보를 집계합니다. 유전적 특징에는 자연적인 순서가 없으므로 위치 임베딩(positional embeddings)은 생략되었습니다.
3. 생존 예측 블록 (Survival Prediction Block): 최종 [CLS] 표현은 SNN과 선형 레이어를 통과하여 $K=4$개의 이산 시간 간격에 대한 로짓(logits)을 생성합니다. 시그모이드 활성화 함수는 이를 이산적 위험 확률(discrete hazard probabilities)로 변환합니다. 모델은 검열된 데이터(censored data)를 고려하는 이산 시간 음의 로그 가능도(negative log-likelihood) 손실을 사용하여 엔드투엔드로 학습됩니다.

학습 전략: 가변 비율 마스킹 (Variable-Rate Masking, VRM)

이질적인 결측 패턴에 대한 강건성을 높이기 위해, 저자들은 학습 중에 **가변 비율 마스킹(VRM)**을 도입했습니다.

• 각 학습 에포크와 환자에 대해, 무작위로 결정된 특징의 비율($k_i$)이 마스킹됩니다. 여기서 $k_i$는 최대 마스킹 비율 $f$까지의 균등 분포에서 샘플링됩니다.
• 이는 모델이 부분적으로 관찰된 입력으로부터 특징 간의 관계를 학습하도록 강제하여, 다양한 기관에서 나타나는 변동성을 시뮬레이션합니다.
• 기본 설정은 $f=0.5$를 사용하여, 학습 중에 최대 50%의 특징이 무작위로 마스킹됨을 의미합니다.
• 중요한 점은, VRM은 오직 학습 중에만 사용된다는 것입니다. 추론 시에는 모델이 결측 마스크를 통해 자연스럽게 누락된 특징을 처리하므로 별도의 임퓨테이션이 필요하지 않습니다.

3. 주요 기여

1. SHIFT 모델: 임퓨테이션 없이 결측된 유전 데이터로부터 생존 예측이 가능한 트랜스포머 기반 모델을 도입하였으며, 마스크드 셀프 어텐션을 사용하여 관찰된 입력에만 집중합니다.
2. 가변 비율 마스킹 (VRM): 모델을 다양한 관찰된 특징 집합에 노출시킴으로써 코호트 간 변화(cross-cohort shifts)에 대한 모델의 강건성을 향상시키는 학습 전략입니다.
3. 포괄적 모델링: 심각하게 불완전한 코호트(예: 88%의 특징이 결측된 경우)도 모델 개발에 포함하여 외부 데이터에 대한 성능을 개선할 수 있음을 입증함으로써, 모델의 포괄성을 보여주었습니다.

4. 실험 결과

저자들은 공공 데이터(TCGA, CPTAC) 및 개인 코호트를 사용하여 두 가지 암 유형인 **교모세포종(GBM)**과 **폐편평세포암(LUSC)**에 대해 SHIFT를 평가했습니다.

완전한 데이터에서의 성능

• 학습과 테스트 데이터가 동일한 특징 공간을 공유하는 설정(결측 없음)에서, SHIFT(VRM 적용 여부와 상관없이)는 표준 베이스라인(CoxPH, RSF, DeepSurv, SNN)과 대등한 성능을 보였습니다.
• VRM의 영향: VRM이 없는 SHIFT도 경쟁력이 있지만, VRM을 추가하면 기관 간 일반화 성능(예: TCGA에서 CPTAC/독일 코호트로의 이동)이 크게 향상되었습니다. 이는 VRM이 특정 특징 집합에 대한 과적합을 방지하는 정규화 도구 역할을 함을 시사합니다.

불완전한 데이터에서의 성능 (LUSC)

• 시나리오: 미국 코호트는 197개 특징 중 22개만이 사용 가능했습니다 (88.8%의 구조적 결측 발생).
• 베이스라인: 197개 전체 특징으로 학습된 표준 모델들은 테스트 시점에 KNN 또는 평균 임퓨테이션이 필요했습니다. 이러한 방식은 임퓨테이션 전략에 따라 민감하게 반응했으며, 많은 베이스라인이 성능이 크게 저하되었습니다 (예: SNN은 KNN 사용 시 C-index가 약 0.42까지 하락).
• SHIFT: 197개 전체 특징으로 학습되고, 임퓨테이션 없이 22개 특징을 가진 미국 코호트에 대해 네이티브하게 평가된 SHIFT-VRM은 가장 우수한 앙상블 성능(C-index 0.570)을 달성했습니다. 이는 가장 강력한 임퓨테이션 베이스라인(평균 임퓨테이션을 사용한 CoxPH)과 일치하며, 다른 대부분의 모델을 능가했습니다.
• 대조군: 모든 모델을 공유된 22개 특징으로 제한했을 때, SHIFT-VRM은 최상의 베이스라인(SNN)보다 뛰어난 성능을 보이지 못했습니다. 이는 SHIFT의 장점이 단순히 트랜스포머 아키텍처 자체에서 오는 것이 아니라, 풍부한 학습 데이터를 활용하면서도 배포 제약에 강건하게 대응하는 능력에서 온다는 것을 확인시켜 줍니다.

다기관 학습

• 저자들은 불완전한 미국 코호트(22개 특징)를 학습 세트(TCGA와 함께)에 포함하는 것이 외부 코호트인 CPTAT의 성능을 개선하는지 테스트했습니다.
• 결과: 불완전한 미국 코호트를 포함함으로써 CPTAP의 C-index가 TCGA 단독 모델(0.602)에서 TCGA+US 모델(0.629)로 향상되었습니다. 이는 부분적으로 관찰된 코호트가 유용한 신호를 제공하며, 배제될 필요가 없다는 주장을 뒷받침합니다.

절제 연구 (Ablation Study)

• 마스킹 비율 $f$를 변화시킨 결과, $f=0.5$가 코호트 전반에 걸쳐 가장 일관된 성능을 제공함을 확인했습니다.
• VRM을 제거했을 때($f=0.0$)는 테스트 데이터가 완전한 설정에서도 외부 일반화 성능이 지속적으로 약화되었으며, 이는 VRM이 정규화 도구로서 수행하는 역할을 확인시켜 줍니다.

5. 의의 및 주장

본 논문은 SHIFT가 유전체 데이터의 이질성에 내재된 "구조적 결측" 문제를 해결함으로써 정밀 종양학을 위한 실질적인 다기관 생존 예측 전략을 제공한다고 주장합니다.

• 단일 모델 배포: 하나의 모델을 이질적인 데이터로 학습시켜, 서로 다른 시퀀싱 패널을 사용하는 기관들에 사이트별 재학습이나 임퓨테이션 없이 배포할 수 있습니다.
• 구조적 결측에 대한 강건성: (마스킹을 통해) 관찰된 입력에 대해서만 조건화하는 것이, 측정되지 않은 대규모 특징 블록을 임퓨테이션을 통해 재구성하는 것만큼 혹은 그보다 더 효과적일 수 있습니다. 특히 결측이 무작위가 아닌 구조적인 경우 더욱 그러합니다.
• 데이터 포괄성: 본 연구는 불완전한 코호트를 제외하는 관행에 도전합니다. 결측을 네이티브하게 처리함으로써, 모델은 더 좁은 패널을 가진 센터의 데이터를 활용하여 성능 저하 없이 샘플 크기와 다양성을 높일 수 있습니다.

저자들은 두 가지 암 유형에 국한된 점, 외부 코호트 크기의 미미함, 그리고 사이트 구조적 패턴보다는 무작위 마스킹을 사용했다는 점 등의 한계를 인정합니다. 그러나 이들은 결측을 명시적으로 모델링하는 것이 더 배포 가능하고 포괄적인 유전체 생존 모델로 나아가는 실행 가능한 경로임을 결론짓습니다.