Towards end-to-end LLM-based censoring-aware survival analysis

기술 요약: LLMSurvival

문제 제기
생존 분석은 의료 예측의 초석으로, 사건의 발생과 그 발생까지의 시간을 모두 모델링하는 것을 목표로 합니다. 그러나 대규모 언어 모델 (LLM) 을 엔드 투 엔드 생존 모델로 적용하는 것은 **중도절단 (censoring)**이라는 도전 과제에 의해 방해받아 왔습니다. 중도절단 데이터에서는 대상자가 탈퇴하거나 관찰 기간이 사건 발생 전에 종료될 수 있으므로 사건의 정확한 시점을 알 수 없습니다. 이러한 명시적인 정답 레이블의 부재는 확정적인 목표를 의존하는 표준 지도형 미세 조정 (SFT) 및 강화 학습 프레임워크의 직접적인 사용을 불가능하게 합니다. 결과적으로 기존 연구는 LLM 을 특징 추출기로만 국한시키거나 (예: LLM 위에 신경망을 추가하는 등) LLM 을 주요 예측 엔진으로 활용하기보다는 상당한 구조적 수정을 요구해 왔습니다.

방법론: LLMSurvival
저자들은 구조화된 임상 표본 데이터를 직접 처리하는 수정되지 않은 LLM 을 사용하여 중도절단을 인식하는 생존 분석을 가능하게 하는 LLMSurvival 프레임워크를 제안합니다. 핵심 혁신은 절대적인 시간 레이블의 필요성을 우회하기 위해 사건 발생 시간 예측을 **쌍별 순위 문제 (pairwise ranking problem)**로 재구성한 것입니다.

1. 데이터 변환: 구조화된 표본 변수 (예: 나이, 생체 징후, 검사 결과) 는 템플릿 (예: "나이는 79 세; 성별은 여성") 을 사용하여 자연어 문장으로 변환됩니다.
2. 쌍별 공식화: 특정 생존 시간을 예측하는 대신, 모델은 두 개의 비교 가능한 대상 중 누가 먼저 사건을 경험할지 판단하도록 훈련됩니다.
   • 훈련 쌍: 대상자 $i$가 사건을 경험 ($e_i=1$) 했고 대상자 $j$의 추적 관찰 시간보다 일찍 ($t_i < t_j$) 발생했을 때 유효한 훈련 쌍 $(i, j)$이 구성됩니다. 이는 다른 대상자의 사건 발생 시간보다 더 오래 관찰된 경우라면 중도절단된 대상자도 포함합니다.
   • 목적: 모델은 지도 학습을 통해 상대적 순서를 예측하도록 미세 조정됩니다. 손실 함수는 유효한 비교 가능한 쌍에 대한 이진 교차 엔트로피 목적 함수이며, 사건을 더 일찍 경험한 대상자에게 더 높은 위험 점수를 할당할 확률을 최적화합니다.
   • 샘플링: 모든 가능한 쌍을 생성하는 계산 비용을 관리하기 위해, 각 사건 사례에 대해 고정된 수의 대조군 (추적 관찰 기간이 더 긴 대상자) 을 선택하는 중첩 사례 - 대조군 샘플링 전략이 사용됩니다.
3. 앵커 비교를 통한 추론: 테스트 중에는 대상자의 위험 점수가 직접 예측되지 않습니다. 대신, 테스트 대상자는 훈련 코호트에서 샘플링된 고정된 $K$개의 "앵커" 개인과 비교됩니다.
   • 모델은 테스트 대상자가 각 앵커보다 위험이 높은지 여부를 예측합니다.
   • 최종 위험 점수는 테스트 대상자가 더 높은 위험으로 순위 매겨진 앵커의 비율입니다.
4. 모델: 이 프레임워크는 두 가지 지시형 튜닝 로컬 배포 가능 LLM 인 Llama 3.1-8B-Instruct와 Qwen 2.5-7B-Instruct를 사용하여 평가되었습니다.

주요 기여

• 엔드 투 엔드 LLM 생존 모델링: 이 논문은 구조적 수정 없이 LLM 을 직접 생존 모델로 훈련시키는 개념 증명 (proof of concept) 을 제시하여, 특징 추출기라는 전통적인 역할을 넘어설 수 있음을 보여줍니다.
• 중도절단 해결: 문제를 순위 작업으로 재구성함으로써, 이 방법은 사건 발생 시간의 부재라는 근본적인 장벽을 해결하며 중도절단된 관찰을 훈련 과정에 자연스럽게 통합합니다.
• 이식성: 이 프레임워크는 급성 ICU 사망률과 만성 골절 위험 등 다양한 임상 맥락에서 작업별 구조 변경 없이 이식 가능함이 입증되었습니다.
• 로컬 배포: 컴팩트하고 공개적으로 이용 가능한 베이스 모델 (7B/8B 파라미터) 을 사용함으로써 외부 API 호출과 관련된 개인정보 보호 문제를 완화하는 로컬 배포가 가능합니다.

결과
이 프레임워크는 두 가지 다른 임상 작업에서 평가되었습니다:

1. ICU 사망률 예측 (MIMIC-IV):
   • LLMSurvival 은 전체 C-지수에서 0.776(Llama) 및 0.784(Qwen)를 달성했습니다.
   • 이는 콕스 비례 위험 모델 (0.753) 대비 3.1% 향상이며, 확립된 임상 점수 SAPS-II(0.746) 대비 3.0% 향상입니다.
   • 평균적으로 세 가지 딥러닝 베이스라인 (Cox-nnet, DeepHit, DeepSurv) 보다 우수한 성능을 보였습니다.
2. 취약성 골절 예측 (NYP/WCM 코호트):
   • LLMSurvival 은 전체 C-지수에서 0.742(Llama) 및 0.741(Qwen)를 달성했습니다.
   • 이는 콕스 모델 대비 0.5% 향상이며, 임상 점수 FRAX(0.660) 대비 8.2% 향상을 보였습니다.
   • 성능은 다양한 시간 범위 (2 년, 3 년, 5 년) 에서 딥러닝 베이스라인과 경쟁하거나 더 우수했습니다.

두 작업 모두에서 카플란 - 마이어 분석은 LLMSurvival 이 환자를 명확한 생존 궤적을 가진 고위험군과 저위험군으로 효과적으로 계층화했음을 확인시켰습니다 (ICU 위험비 4.80, 골절 위험비 4.02). 모델 출력에 대한 정성적 분석은 LLM 이 입원 중증도, 산소화 수준, 동반 질환 부담과 같은 관련 위험 요인을 인용하여 쌍별 비교에 대해 임상적으로 일관된 근거를 생성했음을 나타냈습니다.

의의 및 주장
저자들은 LLMSurvival 이 비교 기반 재구성을 통해 중도절단이 포함된 생존 모델링이 LLM 미세 조정과 호환될 수 있음을 입증했다고 주장합니다. 이 연구는 LLM 이 하이브리드 아키텍처에 의존하기보다 복잡한 통계 문제를 독립적으로 해결할 수 있는 능력을 확장할 수 있음을 보여주는 개념 증명입니다. 전문가가 선별한 임상 점수 (SAPS-II 및 FRAX) 를 능가하는 프레임워크의 능력은 LLM 이 데이터에서 직접 의미 있는 위험 계층화를 학습할 수 있음을 시사합니다. 또한, 이 작업은 통계적으로 어려운 문제를 LLM 의 능력 (예: 비교적 추론) 과 일치하는 작업으로 재구성함으로써 생물의학 기계 학습에서 언어 모델의 적용 범위를 확장할 수 있다는 더 넓은 원리를 보여줍니다.

한계점
저자들은 몇 가지 한계를 인정합니다:

• 훈련은 샘플링 구성에 따라 일부 변이를 보였으며, 이는 최적화 안정성 개선의 여지가 있음을 시사합니다.
• 쌍 선택은 대조군의 무작위 샘플링에 의존했으며, 구조화된 전략이 효율성을 개선할 수 있습니다.
• 추론은 앵커에 대한 여러 쌍별 비교를 필요로 하므로 계산 오버헤드가 발생합니다.
• 현재 구현은 구조화된 표본 입력만 사용하며, 임상 기록 및 의료 영상과 같은 다중 모달 신호를 통합하기 위한 향후 연구가 필요합니다.