이 논문 "Prediction-Augmented Generation for Automatic Diagnosis Tasks"는 자동 진단(Automatic Diagnosis System, ADS)과 같은 추론 기반 작업에서 대규모 언어 모델(LLM)의 한계를 극복하기 위해 예측 증강 생성(Prediction-Augmented Generation, PAG)이라는 새로운 방법론을 제안합니다.
초록 및 서론
대부분의 LLM은 이전 문맥을 기반으로 다음 단어 토큰을 예측하는 자기 회귀(autoregressive) 아키텍처를 채택합니다. 이러한 방식은 글쓰기나 요약과 같은 언어 생성 작업에는 강력하지만, 예측 및 의사 결정과 같은 높은 수준의 추론 작업에는 한계가 있습니다. LLM은 인간의 사고 방식과 덜 유사하며, 정답을 직접 출력하는 예측 기반 모델과 달리 정답을 나타내는 단어 토큰을 간접적으로 예측하여 전반적인 예측 성능이 저하될 수 있습니다. 이는 모델이 "환각(hallucinations)"을 생성하거나 입력 단어의 작은 변화에도 결과가 달라지는 내재적 불안정성을 유발합니다.
기존의 검색 증강 생성(Retrieval-Augmented Generation, RAG) 방식은 외부 지식을 제공하여 이러한 문제를 해결하려 하지만, 질병 진단, 판단, 전략적 결정과 같이 방대한 경험과 예측 능력을 요구하는 고급 작업에는 LLM이 RAG를 사용하더라도 적절히 수행하기 어렵습니다. 실제로, 예측 모델이 질병 예측 작업에서 일반적으로 더 나은 성능을 보입니다.
PAG는 LLM에 작업별 예측 모델을 외부 도구로 통합하여 이러한 한계를 극복하고자 합니다. RAG가 외부 지식으로 LLM을 보완하는 전략에서 영감을 받아, PAG는 예측 모델을 활용하여 레이블 예측을 수행한 다음, 이 결과를 LLM의 광범위한 지식과 통합하여 예측 정확도를 향상시킵니다. 단순히 예측 모델의 추론 결과를 복사하는 것이 아니라, LLM의 지식을 활용하여 추론 결과를 개선하는 것이 특징입니다.
주요 기여는 다음과 같습니다:
- 예측 증강 생성(PAG) 개념 최초 도입: 향후 연구를 위한 분석 및 기반을 제공합니다.
- 최고 수준의 성능 달성: 질병 진단 작업에서 기존 예측 모델과 LLM을 능가하는 성능을 보입니다.
- 새로운 평가 방법 도입: 언어 모델에 대한 보다 정량적이고 객관적인 평가를 가능하게 합니다.
관련 연구
LLM의 초기 연구는 기계 번역, 요약, 질의응답과 같은 언어 이해 및 생성 능력에 중점을 두었습니다. 하지만 지식 업데이트를 위한 재훈련 비용과 출력의 환각 현상과 같은 문제에 직면했습니다. 이러한 한계를 극복하기 위해 외부 도구(external tools), 예를 들어 검색(retrieval), API 호출을 통한 광범위한 외부 도구 지원, 코드 실행기를 통한 계산 지원 등이 탐구되었습니다.
LLM의 성능을 평가하는 것도 중요한 주제입니다. 초기에는 단어 수 기반의 표면적 일치를 비교하는 방식(BLEU)이나 벡터 기반 유사도 측정 방식이 사용되었으나, 더 깊은 의미를 포착하지 못했습니다. 이후 LLM 평가를 위해 객관식 질의응답 데이터셋이 개발되었지만, 전용 데이터셋 구축이 필요하며 LLM이 진정한 이해나 깊은 추론 없이 통계적 패턴에 의존할 수 있다는 한계가 있습니다. 최근에는 LLM 자체를 평가 도구로 사용하는 방식도 등장했으나, LLM이 중간 점수나 중립적인 판단을 내리는 경향이 있다는 새로운 도전 과제를 제시합니다.
예비 지식 및 PAG 방법론
자동 진단 시스템(ADS)의 과제
ADS 연구의 한 가지 과제는 진단 데이터 수집의 어려움입니다. 특히 통계적으로 덜 알려진 질병의 경우 더욱 그렇습니다. PAG는 LLM의 방대한 지식과 예측 모델을 결합하여 이 문제를 해결하는 데 기여할 것으로 기대됩니다.
문제 정의 및 PAG 프로세스
PAG의 전체 프로세스는 Figure 3에 자세히 설명되어 있습니다.
- 1단계: 특징 추출 (Feature Extraction)
- 환자의 비정형 발화($U$)에서 관련 의료 개념을 추출합니다.
- LLM이 이 추출을 수행하며, 누락된 속성을 채우고 개념을 예측 모델의 입력 스키마에 맞게 정규화합니다.
- 2단계: 특징 정렬 (Feature Alignment)
- LLM이 생성한 추출된 의료 개념($E^I$)을 미리 정의된 의료 용어 사전($T$)에 매핑합니다.
- 두 단계 매칭 방식을 사용합니다:
- 정확 일치(Exact Match): 쿼리 용어($q_n$)가 사전($T$)에 있으면 직접 할당합니다.
- 순위 기반 매칭: 정확 일치가 없으면, 어휘 기반(lexical-based) 유사도 점수 (단어 집합의 교집합 크기)와 의미 기반(semantic-based) 유사도 점수 (단어 임베딩의 코사인 유사도)를 결합하여 용어 후보를 순위화합니다.
- 상위 $k$개 후보 용어($C$)는 LLM에 구조화된 프롬프트를 통해 제공되며, LLM은 문맥상 가장 적절한 의료 용어를 선택합니다.
- 선택된 용어는 원-핫 인코딩, 단어 임베딩 또는 통계적 표현으로 변환되어 예측 모델의 입력($X$)으로 준비됩니다.
- 3단계: 예측 (Prediction)
- 준비된 입력($X$)은 예측 모델에 입력되어 $k$개의 예측 결과를 생성합니다.
- 예측된 수치 질병 레이블은 $L^N$, 해당 텍스트 레이블은 $L^T$로 표시됩니다.
- 4단계: 지식 통합 (Knowledge Integration)
- $L^T$는 LLM에 의해 추가로 처리되어 각 예측된 질병에 대한 상세 의료 지식($K_d$)을 생성합니다. 이는 진단 결과에 대한 설명 가능성을 높여줍니다.
- 5단계: 진단 (Diagnosis)
- 환자의 원본 발화($U$), 예측 결과($L^T$), 생성된 의료 지식($K_d$)이 통합되어 LLM에 제공됩니다.
- LLM은 이를 기반으로 정제된 진단 텍스트($D^T$)와 해당 질병 레이블($D^L$)을 생성합니다.
- 6단계: 집계 (Aggregation)
- 예측 모델의 출력($L^N$)과 LLM이 생성한 진단($D^L$)을 결합하기 위한 집계 메커니즘이 사용됩니다.
- 각 소스에서 가장 가능성이 높은 질병에 가장 높은 점수($k$)를 부여하고, 그다음에는 $k-1$ 등을 부여하는 방식으로 순위를 매깁니다.
- 각 질병에 대한 점수는 $\text{Vote}(d) = ∑ S_s(d)$로 계산되며, 최종 진단($D$)은 가장 높은 집계 점수를 가진 질병으로 결정됩니다. 이는 예측 모델과 LLM 모두의 통찰력을 통합한 포괄적인 결론을 나타냅니다.
실험
PAG는 단일 질병 예측을 위한 MDD, Muzhi(MZ), DXY 데이터셋과 다중 진단을 위한 PolyMed 데이터셋에서 평가되었습니다. Recall@K, Precision@K, NDCG@K, 그리고 PolyMed의 경우 WAM(Weighted Arithmetic Mean) 점수가 지표로 사용되었습니다. LDA, TabNet, XB-Net, M2DIA, MAEA와 같은 예측 모델과 GPT-4o, GPT-3.5, Gemma2, Llama3.1, Mixtral, Mistral-large와 같은 다양한 LLM이 기준 모델로 사용되었습니다.
데이터 복원 테스트
- 환자 비정형 텍스트에서 구조화된 데이터를 추출하고 미리 정의된 용어에 매핑하는 능력을 평가했습니다.
- Mistral-large 모델이 Feature alignment recall(FAR)에서 가장 우수했지만 (Table 1), GPT-4o와 LDA 조합의 PAG가 진단 성능에서 가장 큰 향상을 기록했습니다 (Table 2).
- 이는 데이터 복원 성능이 최종 예측 성능을 보장하지 않으며, PAG는 단어 매칭, 지식 생성, 예측 정제 등 다양한 작업을 포함하므로 각 단계에서 균형 잡힌 성능을 유지하는 모델이 더 적합함을 시사합니다.
단일 진단 작업
- PAG는 모든 데이터셋에서 예측 성능을 크게 향상시켰습니다 (Table 3).
- Mistral-large와 같이 예측 성능이 좋지 않은 단일 LLM에서도 PAG 적용 후 상당한 성능 향상을 보였는데, 이는 PAG가 예측 작업에서 LLM의 성능을 보완하는 강력한 외부 도구로 활용될 수 있음을 나타냅니다.
- PAG는 RAG 방식보다 전반적으로 우수한 성능을 보였습니다. 이는 질병 진단에서 추론 기반 도구가 지식 검색보다 성능에 더 중요하다는 것을 시사합니다.
다중 진단 작업
- PAG는 모든 실험에서 예측 성능을 향상시켰으며, 특히 Top-5 예측에서 현저한 향상을 보였습니다 (Table 4).
- 이는 예측 모델이 단일 레이블 출력에 훈련되는 경향이 있어 여러 후보를 생성하는 데 실패할 수 있는 반면, PAG는 LLM 지식을 활용하여 예측을 정제하고 후보 풀을 향상시켜 다중 의심 질병을 제공하는 자동 진단 시스템에 특히 유용함을 보여줍니다.
LLM 지식의 영향
- 전체 PAG 모델($K_d$ 포함)과 지식 생성 모듈이 없는 변형(w/o $K_d$)을 비교했을 때, 전체 모델은 Top-3 및 Top-5 지표에서 더 나은 성능을 보였습니다.
- 하지만 Top-1에서는 w/o Kd 변형이 유리한 결과를 보였는데, 이는 LLM 지식이 전반적인 예측 풀의 품질을 향상시킨다는 것을 시사합니다.
- 특히 학습되지 않은 질병을 추론해야 하는 PolyMed의 Unseen test(미확인 테스트)에서 PAG w/o $K_d$는 성능이 저하되었고, 전체 PAG 모델이 LLM 단독 성능을 능가했습니다. 이는 LLM 유래 지식을 통합하는 것이 예측 정확도와 설명 가능성을 모두 향상시킨다는 것을 의미합니다.
정보량의 영향
- PAG(M2DIA)와 함께 LLM에 제공되는 질병 예측 후보의 수를 변경하며 성능을 평가한 결과 (Figure 4), 성능은 초기에는 향상되지만 특정 지점(25~35개 예측)을 넘어서면 정체되거나 심지어 감소합니다.
- 이는 LLM이 과도한 정보를 처리하는 제한된 능력과 내재적인 진단 제약을 반영하며, PAG에서 예측 수는 최적화해야 할 하이퍼파라미터로 다루어져야 함을 나타냅니다.
요소 제거 연구(Ablation Study)
- Figure 5에서, 전체 PAG 모델은 LLM 단독 변형보다 Top-1 성능이 좋지만, 예측 모델 단독 또는 지식 없는 버전보다는 낮습니다. 이는 예측 모델이 단일 레이블 정확도에 최적화되어 있는 반면, LLM은 더 광범위한 추론에 적합하기 때문입니다.
- 그러나 Top-5 지표에서는 전체 PAG가 다른 모든 변형을 일관되게 능가합니다. 이는 LLM이 생성한 지식을 통합하는 것이 후보 예측의 전반적인 품질을 향상시킨다는 것을 입증합니다.
사례 연구
- PAG가 단계별 추론을 통해 질병 예측을 어떻게 정제하는지 보여주는 사례 (Table 6)에서, 예측 모델이 여러 잠재적 질환을 식별하더라도, LLM은 가슴뼈 뒤의 작열감과 복통 증상의 특이성을 바탕으로 식도염(esophagitis)을 가장 그럴듯한 진단으로 정제합니다.
- PAG는 통계적 예측과 문맥적 의료 추론을 모두 통합하며, 생성된 지식을 바탕으로 추론의 합리적 근거를 제공하여 설명 가능성을 향상시킵니다. 이는 의료 분야에서 신뢰성과 신뢰성이 중요한 실제 배포에 PAG의 잠재력을 보여줍니다.
결론 및 한계
논문은 PAG 기술이 예측 모델을 LLM의 외부 도구로 활용하여 추론 기반 작업에서 LLM을 지원하는 방법을 제안했습니다. 이 방법은 질병 진단 생성뿐만 아니라 예측 모델의 예측 성능까지 효과적으로 향상시킴을 입증했습니다. 특히, LLM의 개입이 질병 예측 풀의 품질을 효과적으로 향상시킴을 발견했습니다. 결과적으로 PAG는 예측 모델과 LLM을 효과적으로 결합하여 의료 진단의 신뢰성을 향상시키고, 실제 의료 응용 분야에서 강력한 도구로 기여할 수 있는 핵심 기술이 될 수 있습니다.
한계점은 다음과 같습니다:
- 지식 통합 (Knowledge Integration): 예측 결과와 LLM 지식의 융합을 최적화하기 위한 추가 연구가 필요합니다. PAG는 다중 질병 후보를 생성하는 데 뛰어나지만, 단일 가장 가능성 있는 진단을 정확히 찾아내는 능력은 상대적으로 약합니다.
- 제어 가능성 (Controllability): LLM은 프롬프트에 제공된 규칙과 구조를 따르지 못하는 경우가 있으며, 작은 모델에서는 오타나 부정확한 토큰 출력과 같은 오류가 발생할 수 있습니다.
- 설명 가능성 (Explainability): PAG는 명시적인 지식 생성 및 추론을 통해 의미 있는 설명 가능성을 제공하지만, 이는 모델의 내부 계산에 대한 직접적인 분석이 아니라 기존의 사후(post-hoc) 방법론을 통해 생성된 것입니다. 따라서 모든 내부 프로세스를 완전히 설명하지는 못하며, 본질적으로 설명 가능한 AI 기술의 추가 통합이 필요합니다.
PAG는 마치 경험 많은 의사가 (예측 모델) 임상 데이터를 바탕으로 일차 진단을 내리고, 그 진단을 (LLM) 방대한 의학 지식과 환자 개개인의 특수성을 결합하여 더욱 정교하고 설명 가능한 최종 진단을 도출하는 과정과 유사합니다. 이처럼 PAG는 기계 학습의 정밀한 예측 능력과 LLM의 폭넓은 지식 및 추론 능력을 융합하여, 복잡한 의료 진단과 같은 중요한 분야에서 더욱 신뢰할 수 있는 결과를 제공할 잠재력을 가지고 있습니다.