🤖 machine learning

From Hazard Functions to Language Space: Cox-Supervised Distillation of Survival Risk into a Large Language Model

이 논문은 콕스 모델(Cox model) 예측을 타겟으로 하여 텍스트로 변환된 임상 데이터로 미세 조정된 Qwen 기반 대규모 언어 모델이 생존 위험 구조를 효과적으로 내재화하여 시간-사건 추론(time-to-event reasoning)에 대해 경쟁력 있는 판별력과 교정 성능을 달달성할 수 있음을 입증한다.

원저자: Nicholas I-Hsien Kuo, Blanca Gallego, Louisa Jorm

게시일 2026-06-09

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Nicholas I-Hsien Kuo, Blanca Gallego, Louisa Jorm

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

핵심 아이디어: 이야기꾼에게 미래를 예측하는 법 가르치기

당신에게 환자가 심각한 질병을 앓은 후 얼마나 오래 생존할지 예측하려는 두 명의 전문가가 있다고 상상해 보세요.

과거의 전문가 (Cox 모델): 이 전문가는 전통적이고 매우 존경받는 통계학자입니다. 환자의 데이터(나이, 혈압, 병력 등)를 살펴보고 엄격한 수학 공식을 사용하여 위험 점수를 계산합니다. 이들은 매우 뛰어나지만, 오직 "수학"으로만 대화합니다. 왜 그렇게 생각하는지 설명할 수 없으며, 당신과 대화를 나눌 수도 없습니다.
새로운 학생 (대규모 언어 모델): 이 전문가는 인터넷의 거의 모든 것을 읽은 아주 똑똑한 사서처럼, 이야기를 읽고 쓰는 데 특화된 강력한 AI입니다. 자연어를 이해하지만, 본질적으로 복잡한 의학 수학을 수행하는 법은 모릅니다.

실험: 연구진은 간단한 질문을 던졌습니다. 우리가 "수학 전문가"가 정답을 적게 하고, "이야기꾼"이 그것을 그대로 따라 하게 함으로써 "이야기꾼"에게 "수학 전문가"처럼 행동하도록 가르칠 수 있을까?

방법론: "스승과 제자" 게임

AI에게 직접 복잡한 수학 방정식을 학습시키는 대신(언어 모델에게는 어려운 일입니다), 그들은 **지식 증류(Knowledge Distillation)**라는 영리한 기술을 사용했습니다. 이것은 마치 마스터 셰프가 부셰프(보조 요리사)를 가르치는 것과 같습니다.

스승 (Cox 모델): 연구진은 먼저 전통적인 수학 전문가가 수천 명의 환자 기록을 살펴보며 각 개인의 정확한 "1년 생존 위험도"를 계산하게 했습니다.
번역: 그들은 환자의 가공되지 않은 데이터(나이: 69세, 심박수: 80 등과 같은 숫자)를 의사의 진료 기록과 같은 간단한 문장으로 변환했습니다. 예: "환자는 69세이며, 심박수는 80이고, 심부전 병력이 있음."
수업: 그들은 AI에게 이 문장을 보여주며 말했습니다. "여기 환자가 있다. 정답은 0.2374이다."
연습: AI(학생)는 이 진료 기록을 읽고 숫자를 적는 연습을 했습니다. AI는 수학 공식을 배운 것이 아니라, 환자의 삶이라는 이야기가 특정 숫자와 어떻게 연결되는지에 대한 패턴을 학습했습니다.

연구진은 이를 유방암, HIV, 그리고 심근경색 환자라는 세 가지 서로 다른 환자 집단에 대해 테스트했습니다.

결과: 학생이 해냈다

결과는 놀라울 정도로 좋았습니다. AI는 단순히 "수학 문제를 푸는" 것이 아니라 (문장을 완성하듯) "다음 단어를 쓰는" 방식으로 훈련되었음에도 불구하고, 전통적인 수학 전문가만큼이나 생존 위험을 잘 예측해 냈습니다.

정확도: 환자의 순위(예: "이 환자가 저 환자보다 위험도가 높은가?")를 매기는 능력을 확인했을 때, AI는 전문가와 거의 대등한 정확도를 보였습니다.
교정(Calibration): AI가 제시한 숫자 또한 현실적이었습니다. 만약 AI가 어떤 환자의 위험도를 20%라고 말했다면, 실제로 유사한 환자들 중 약 20%가 사건(사망 등)을 경험했습니다.

숨겨진 마법: "위험 지형(Risk Landscape)"

연구진은 AI의 "뇌"(은닉 상태) 내부를 들여다보며 AI가 데이터를 어떻게 이해하고 있는지 확인했습니다.

그들은 시각화 도구를 사용하여 AI의 생각을 지도화했습니다. AI가 환자를 별개의 고립된 상자(예: "저위험" 대 "고위험")로 취급하는 대신, 매끄럽고 연속적인 지형을 그려내는 것을 발견했습니다.

비유: 언덕을 상상해 보세요. 언덕 아래에는 위험도가 낮은 건강한 환자들이 있습니다. 언덕을 올라갈수록 위험도는 점점 더 높아집니다. AI는 뚜렷한 단계가 있는 계단을 만든 것이 아니라, 매끄러운 경사로를 만들었습니다. AI는 위험도가 단순히 켜지고 꺼지는 스위치가 아니라, 미끄러지듯 변하는 연속적인 척도라는 것을 이해하고 있었습니다. 이는 AI가 단순히 숫자를 암기한 것이 아니라, 생존 위험이라는 개념을 진정으로 "이해"했음을 시사합니다.

한계점: 아직 완벽한 대체재는 아니다

이 논문은 다음과 같은 한계점을 솔직하게 밝히고 있습니다:

천장 효과: AI는 스승보다 똑똑해질 수 없습니다. AI는 단지 수학 전문가의 답을 복사하는 과정이었기에, 수학 전문가가 놓친 것을 새로 발견할 수는 없습니다. 즉, 스승의 품질에 제한을 받습니다.
비용: 이 AI를 훈련하는 것은 단순한 수학 전문가를 실행하는 것보다 훨씬 더 비싸고 많은 계산 자원이 필요합니다.
데이터 문제: AI는 데이터가 깨끗할 때는 잘 작동하지만, 연구진은 실제 의료 현장의 기록은 무질서하고 일관성이 없다는 점을 언급했습니다. 이 AI가 전통적인 방식만큼 실제 세계의 "노이즈(잡음)"를 잘 처리할 수 있는지는 아직 증명되지 않았습니다.

결론

이 논문은 대규모 언어 모델이 환자에 대한 텍text 설명을 읽고 전통적인 모델의 답을 복사하는 것만으로도 사건 발생 시간 위험도(생존과 같은)를 이해할 수 있음을 입증했습니다. 이는 이러한 AI 모델이 복잡한 의학적 패턴을 내면화하고 이를 자연어를 통해 환자의 결과를 추론하는 도구로 활용할 수 있는 가능성을 열어주며, 향후 자연어로 환자의 예후를 다루는 도구의 토대를 마련했습니다.

기술 요약: 위험 함수에서 언어 공간으로

문제 정의

생존 분석은 중도 절단(censoring)이 포함된 시간-사건(time-to-event) 결과에 대한 근본적인 임상 예측 작업입니다. Cox 비례 위험 모델과 같은 고전적 방법론은 해석 가능성과 견고함을 유지하고 있지만, 현대의 신경망 생존 모델은 일반적으로 언어 기반의 감독(supervision)보다는 명시적인 예측 헤드와 특화된 생존 손실 함수(예: Cox 부분 가능도)에 의존합니다. 반면, 최근 표 형식의 대규모 언어 모델(LLM) 연구는 구조화된 데이터를 분류 작업을 위한 자연어 프롬프트로 처리할 수 있음을 입증했습니다. 그러나 생성형 LLM이 생존 특화 손실 함수 없이 표준 텍exam 생성 목적 함수만을 통해 학습될 때, 통계적 생존 모델이 추정한 연속적인 시간-사건 위험 구조를 내재화할 수 있는지 여부는 여전히 불분명합니다.

방법론

저자들은 통계적 교사(teacher) 모델로부터 생성형 LLM 학생(student) 모델로 생존 위험 정보를 전달하는 Cox 감독 증류(Cox-supervised distillation) 프레임워크를 제안합니다. 이 파이프라인은 네 가지 주요 단계로 구성됩니다:

Cox 교사 학습: 구조화된 임상 공변량( $x_i$ ), 추적 관찰 시간( $t_i$ ), 그리고 이벤트 지표( $\delta_i$ )에 대해 표준 Cox 비례 위험 모델을 적합(fit)시킵니다. 이 모델은 환자별 1년 생존 위험 추정치( $r_i = 1 - S_i(\tau)$ )를 생성하며, 이는 정답(ground-truth) 학습 타겟 역할을 합니다.
프롬프트 구성: 구조화된 환자 기록을 자연어 프롬프트로 변환합니다. 각 프롬프트에는 작업 지침, 구조화된 임상 공변량 목록(예: "Age: 69", "BMI: 26.4"), 그리고 큐(cue)인 "Risk:"가 포함됩니다.
LLM 미세 조정 (학생): Qwen2.5-1.5B-Instruct 모델을 **저순위 적응(LoRA)**을 사용하여 미세 조정합니다. 학습 목적 함수는 표준적인 자기회귀적 지도 미세 조정(SFT)이며, **완성 부분 전용 손실(completion-only loss)**을 사용합니다. 모델은 텍스트 프롬프트를 네 자리 소수점 형태의 문자열(예: "0.2374")로 된 스칼라 위험 값으로 매핑하는 법을 배웁니다. 결정적으로, 이 모델은 생존 특화 손실(논문의 식 1과 같은)을 통해 최적화되지 않으며, 오직 Cox에서 유도된 위험 문자열을 타겟으로 하는 순수한 텍스트 생성 작업으로서 학습됩니다.
추론 및 평가: 미세 조정된 모델은 추론 과정에서 자기회귀적으로 위험 문자열을 생성합니다. 이 문자열들은 정규 표현식을 통해 숫자로 파싱됩니다. 성능은 표준 생존 지표인 Harrell의 C-index(변별력), 교정 오차( $D_{21}$ ), 그리고 백분위 기반 재분류 개선도(NRI)를 사용하여 평가됩니다.

이 프레임워크는 세 가지 공개 데이터셋인 GBSG2(유방암), ACTG320(HIV/AIDS), WHAS500(급성 심근경색)에 대해 평가되었습니다.

주요 결과

세 가지 데이터셋 모두에서 Qwen 기반 LLM은 텍스트 생성 작업으로 학습되었음에도 불구하고 Cox 교사 모델과 비교하여 경쟁력 있는 성능을 보였습니다:

변별력(Discrimination): WHAS500 데이터셋에서 LLM은 0.7563의 C-index를 달려 Cox 모델의 0.7658에 근접한 성능을 보였습니다. GBSG2(0.6801 vs. 0.6780)와 ACTG320(0.7066 vs. 0.7038)에서도 유사한 결과가 관찰되었습니다.
교정(Calibration): LLM은 완만한 교정 능력을 유지했으나, WHAS500의 경우 Cox 모델(0.0583)보다 약간 높은 오차(0.0711)를 보였습니다. 일부 하위 그룹(예: GBSG2의 Age > 60y)에서는 LLM이 Cox 모델보다 낮은 교정 오차를 달성하기도 했습니다.
잠재 공간 기하학(Latent Space Geometry): 모델의 은닉 상태(hidden states)에 대한 t-SNE 시각화 결과, 잠재 공간 내에서 **매끄러운 위험 경사(smooth risk gradients)**가 나타났습니다. 예측 위험도가 유사한 환자들이 함께 클러스터링되는 것을 통해, 모델이 생존 위험을 이산적인 범주가 아닌 연속적인 구조로 표현하고 있음을 시사합니다.
재분류(Reclassification): 백분위 기반 NRI 분석에 따르면, LLM은 전반적으로 위험 순위를 보존했지만 Cox 모델의 순위를 완전히 보존하지는 못했으며, 특히 이벤트 발생 사례에서 성능 저하를 보였습니다 (WHAS500의 경우 NRI cases = -0.2067).

의의 및 주장

본 논문은 대규모 언어 모델이 명시적인 생존 목적 함수 없이도 자기회귀적 지도 미세 조정만으로 생존 위험 구조를 내재화할 수 있음을 주장합니다. 핵심 기여는 다음과 같습니다:

텍스트 기반 생존 모델링의 타당성: 본 연구는 구조화된 공변량을 텍스트로 변환하고 LLM이 위험 점수를 생성하도록 학습시키는 것이 전통적인 신경망 생존 아키텍처의 실행 가능한 대안임을 입증합니다.
연속적 표현: 잠재 공간에서의 매끄러운 위험 경사 관찰은 LLM이 시간-사건 위험을 고립된 분류 빈(bin)으로 취급하는 대신, 임상적으로 의미 있는 연속적 표현으로 학습할 수 있음을 시사합니다.
이질적 데이터에 대한 유연성: 저자들은 텍스트 기반 프레임워크가 엄격한 수치적 특징 벡터 대신 자연어 설명을 조건으로 하기 때문에, 이질적인 임상 데이터나 호환되지 않는 피처 스키마를 가진 환경에서 실질적인 이점을 제공할 수 있다고 제안합니다.

저자들은 Cox 모델이 이 증류 접근법의 "예측 천장(predictive ceiling)"을 정의하며, LLM이 더 높은 계산 비용을 초래한다는 점을 인정하며 신중한 입장을 유지합니다. 또한, 이 프레임워크가 현재로서 Cox 모델을 능가하지는 못하지만, 시간-사건 추론을 언어 모델에 통합하는 새로운 경로를 제시한다고 언급했습니다. 향후 연구로는 정보성 결측(informative missingness)에 대한 강건성과 기관 간 분포 변화(distribution shifts) 문제를 해결하는 것이 제안되었습니다.