MedFeat: Model-Aware and Explainability-Driven Feature Engineering with LLMs for Clinical Tabular Prediction

이 논문은 LLM 의 도메인 지식과 SHAP 기반의 피드백을 결합하여 임상 표형 예측에서 기존 모델의 특성을 고려한 설명 가능한 특징 공학 프레임워크 'MedFeat'를 제안하고, 다양한 임상 작업에서 안정적인 성능 향상과 분포 변화에 대한 강건성을 입증합니다.

핵심

🕵️‍♂️ 비유: 수사팀과 AI 형사

의료 데이터를 분석하는 것은 마치 복잡한 사건을 해결하는 수사팀과 같습니다.

• 데이터 (환자 기록): 사건 현장에 흩어진 수많은 증거들 (혈압, 나이, 과거 병력 등).
• 기존 AI 모델 (XGBoost 등): 뛰어난 감식 능력을 가진 초능력 형사. 이 형사는 이미 많은 증거를 보고 패턴을 찾아냅니다. 하지만 모든 증거를 다 볼 수는 없거나, 숨겨진 연결고리를 놓칠 때가 있습니다.
• 문제점: 기존에는 이 형사가 스스로 모든 증거를 찾아내려 하거나, 연구자들이 일일이 "이 증거와 저 증거를 섞어보자"라고 시도했습니다. 하지만 증거가 너무 많고, 연구자들의 시도도 무작위라 비효율적이었습니다.

💡 MedFeat 의 등장: "AI 를 위한 명탐정 코난"

MedFeat 는 이 **초능력 형사 (AI) 가 놓친 단서를 찾아주는 '지식인 조력자 (LLM)'**입니다. 하지만 단순히 무작위로 단서를 만드는 게 아니라, 형사의 약점을 정확히 파악하고 **가장 중요한 단서 (SHAP 값)**에 집중합니다.

1. "형사가 무엇을 못 보는지 알아요" (Model-Awareness)

• 상황: 어떤 형사는 '직선'으로만 생각할 수 있고, 어떤 형사는 '복잡한 곡선'도 잘 봅니다.
• MedFeat 의 역할: "이 형사는 복잡한 곡선을 못 보니까, 우리가 미리 그 곡선을 만들어서 줘야 해!"라고 생각합니다.
• 예시: 만약 AI 가 '나이가 많고 가난하면 위험하다'는 복잡한 관계를 직접 찾기 어렵다면, MedFeat 는 "나이와 가난함을 곱한 새로운 점수"라는 **새로운 증거 (특징)**를 만들어 AI 에게 건네줍니다. AI 는 이 새로운 증거를 보면 바로 "아, 위험하다!"고 깨닫습니다.

2. "가장 중요한 단서부터 찾아봐요" (Explainability-Driven)

• 상황: 사건 현장에 증거가 100 개 있는데, 그중 90 개는 쓸모없는 쓰레기입니다.
• MedFeat 의 역할: AI 가 "이 증거 (혈압) 가 가장 중요해!"라고 말하면 (SHAP 값), MedFeat 는 그 가장 중요한 증거 3 개만 골라 "이 세 가지를 섞어보면 어떨까?"라고 제안합니다.
• 효과: 모든 증거를 다 섞어보는 게 아니라, 가장 유력한 용의자 (중요한 데이터) 만 집중적으로 조사하므로 시간과 에너지를 아끼고 더 정확한 단서를 찾습니다.

3. "실패한 실수는 잊지 않아요" (Feedback Loop)

• 상황: 한 번 만든 증거가 사건 해결에 도움이 안 되면?
• MedFeat 의 역할: "이 조합은 실패했어. 다음엔 다른 조합을 시도하자"라고 **기억장 (Memory Bank)**에 기록합니다. 그리고 다음 단계에서는 성공했던 조합을 더 많이 시도합니다.
• 효과: 같은 실수를 반복하지 않고, 점점 더 똑똑한 증거들을 만들어냅니다.

🏥 실제 효과: 왜 이것이 중요한가요?

이 연구는 실제 의료 데이터 (환자 사망률 예측, 입원 기간 예측 등) 에서 실험을 했습니다.

1. 더 정확한 진단: MedFeat 를 사용하면 기존 AI 모델보다 환자의 위험을 더 정확하게 예측할 수 있었습니다. (특히 데이터가 부족하거나 불균형한 상황에서 효과가 큽니다.)
2. 시간이 지나도 변하지 않는 단서: 병원마다, 시기에 따라 데이터가 달라져도 (예: ICU 환자와 일반 환자, 2020 년과 2024 년) MedFeat 가 만든 새로운 증거들은 여전히 유효했습니다. 이는 AI 모델이 자주 다시 훈련하지 않아도 된다는 뜻입니다.
3. 비밀 유지: 환자의 이름이나 구체적인 기록을 AI 에게 보내지 않고, 통계적인 중요도만 보내기 때문에 개인정보 보호 문제도 해결했습니다.

📝 한 줄 요약

MedFeat 는 "AI 가 놓친 중요한 단서를, 의사의 지식과 AI 의 피드백을 결합해 자동으로 찾아내어, 더 정확하고 안전한 의료 진단을 가능하게 하는 스마트 조력자"입니다.

이 도구는 의료진이 매일 쏟아지는 복잡한 환자 데이터 속에서, 진짜 중요한 신호를 찾아내어 생명을 구하는 데 도움을 줄 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

의료 분야 표형 데이터 (Tabular Data) 예측에서 기존 딥러닝 모델보다 전통적인 머신러닝 모델 (예: 트리 기반 모델) 이 더 우수한 성능을 보이는 경우가 많습니다. 그러나 이러한 모델의 성능을 극대화하기 위해서는 **수동적인 특징 공학 (Feature Engineering)**이 필수적이지만, 이는 시간과 비용이 많이 들고 전문 지식이 요구됩니다.

최근 대규모 언어 모델 (LLM) 을 활용한 자동 특징 공학 연구가 등장했으나, 기존 방법론들은 다음과 같은 한계를 가지고 있습니다:

• 모델 무관성 (Model-Agnostic Design): 하류 모델 (Downstream Model) 의 특성을 고려하지 않습니다. 예를 들어, 로지스틱 회귀는 비선형성이 필요하지만 트리 모델은 이를 이미 학습할 수 있어, 모델에 따라 필요한 특징이 다름에도 불구하고 동일한 방식으로 특징을 생성합니다.
• 중요도 신호 무시: 모든 특징을 동등하게 취급하여, 예측에 중요한 신호를 놓치거나 불필요한 탐색을 수행합니다.
• 프롬프트 길이 문제: 모든 특징의 설명을 LLM 에 전달하면 프롬프트 길이가 길어지고 생성 품질이 저하됩니다.
• 개인정보 보호 문제: 일부 방법은 환자 수준의 원본 데이터를 LLM 에 전송하여 프라이버시 문제가 발생할 수 있습니다.

2. 방법론 (Methodology: MedFeat)

저자들은 MedFeat이라는 새로운 프레임워크를 제안합니다. 이는 LLM 의 추론 능력과 도메인 지식을 활용하되, **모델 인식 (Model-Aware)**과 **설명 가능성 (Explainability)**을 기반으로 한 피드백 루프를 통해 특징을 생성합니다.

핵심 구성 요소

1. 모델 인식 특징 생성 (Model-Aware Feature Generation):

   • 하류 모델의 특성에 맞춰 특징을 제안합니다.
   • 로지스틱 회귀 (LR) 인 경우: 명시적인 비선형 변환과 상호작용 항 (Interaction Terms) 을 생성하도록 유도합니다.
   • XGBoost 인 경우: 트리 모델이 이미 학습 가능한 단순한 패턴은 피하고, 복잡한 시간적 패턴, 전역 통계, 문맥 기반 상호작용 등 모델이 학습하기 어려운 구조를 포착하는 특징을 생성하도록 유도합니다.

2. 설명 가능성 기반 피드백 (Explainability-Driven Feedback):

   • SHAP 값 활용: 하류 모델의 예측 성능이 아닌, SHAP (SHapley Additive exPlanations) 값을 기반으로 특징의 중요도를 산출합니다.
   • 이 중요도 점수는 특징 생성의 우선순위를 정하고, LLM 프롬프트에 구조화된 중요도 정보를 제공하여 생성의 방향성을 잡는 데 사용됩니다.

3. 중요도 가중 섬 샘플링 (Importance-Weighted Island Sampling):

   • 모든 특징을 한 번에 LLM 에 전달하는 대신, SHAP 점수에 따라 중요한 특징들의 부분집합을 **'섬 (Island)'**으로 구성합니다.
   • 각 섬은 작은 특징 집합으로, LLM 프롬프트 길이를 제한하고 토큰 사용을 줄이며, 특정 지역 (Feature Space) 에 집중하여 생성 품질을 높입니다.

4. 피드백 메모리 및 반복 최적화:

   • 성공한 특징과 실패한 특징을 메모리에 저장하여 다음 반복에서 유사한 저품질 제안은 피하고, 성능을 향상시킨 패턴은 강화합니다.
   • 프라이버시 보호: 환자 수준의 원본 데이터는 LLM 에 절대 전송되지 않으며, 메타데이터와 SHAP 점수 (집계된 피드백) 만을 전송합니다.

3. 주요 기여 (Key Contributions)

1. 최초의 모델 인식 LLM 특징 공학 프레임워크: 하류 학습자의 표현 한계 (Representational Limits) 에 맞춰 특징 제안을 적응화하여, 모델이 이미 학습 가능한 특징에 대한 평가 예산을 낭비하지 않고 모델이 학습하기 어려운 신호를 생성합니다.
2. SHAP 기반 생성 신호: 특징 생성 시 SHAP 점수와 랭킹을 활용하여 예측 신호가 이미 입증된 특징을 우선시하고, 생성 프롬프트에 구조화된 중요도 정보를 조건부 (Conditioning) 로 제공합니다.
3. 효율적인 섬 탐색 전략: 중요도 기반의 섬 샘플링을 통해 토큰 사용을 제한하고, 약한 예측 변수에 대한 집중되지 않은 변환을 방지하여 특징 품질과 효율성을 동시에 개선합니다.
4. 실제 배포 가능성 검증: 분포 변화 (Distribution Shift) 와 시간적 드리프트 (Temporal Drift) 하에서도 안정적인 성능을 보여주며, ICU 코호트에서 일반 입원 환자 집단으로의 전이 학습에서도 효과성을 입증했습니다.

4. 실험 결과 (Results)

저자들은 IORD, MIMIC-IV, HRS 등 다양한 임상 데이터셋과 24 시간 입원 사망률, 10 년 사망률, 심부전 예측 등 여러 태스크에서 MedFeat 를 평가했습니다.

• 성능 향상:
  • XGBoost: 기본 설정 (Untuned) 에서 모든 태스크에서 AUC 가 개선되었으며, 특히 24 시간 입원 사망률 예측에서 AUC 가 0.686 에서 0.740 으로 7.87% 향상되었습니다.
  • 로지스틱 회귀: 선형 모델의 표현 능력을 확장하는 특징을 생성하여 AUC 와 F1 점수 모두에서 일관된 개선을 보였습니다.
  • 기존 방법 대비: AutoFeat, OpenFE, CAAFE, FeatLLM, OCTree 등 기존 자동 특징 공학 방법론들보다 안정적이고 일관된 성능 향상을 보였습니다. 특히 기존 방법들은 한 지표는 개선되지만 다른 지표는 악화시키는 불안정성이 있었으나, MedFeat 는 AUC 와 F1 을 동시에 개선하는 경향이 있었습니다.
• 하이퍼파라미터 최적화 (HPO) 하의 강건성:
  • XGBoost 의 경우 HPO 를 수행하면 특징 공학의 이득이 줄어들지만, MedFeat 는 여전히 F1 점수에서 유의미한 개선을 보였습니다.
  • 로지스틱 회귀의 경우 HPO 후에도 특징 공학의 이득이 지속적으로 유지되었습니다.
• 일반화 및 드리프트 강건성:
  • 전이 학습: MIMIC-IV(중환자실) 에서 생성된 특징을 IORD(일반 입원) 데이터에 적용했을 때, 성능이 향상되고 변동성이 감소했습니다.
  • 시간적 드리프트: 2020 년 이전 데이터로 학습된 모델을 2020~2024 년 데이터에 적용할 때, MedFeat 특징을 사용한 모델은 재학습 (Retraining) 한 베이스라인보다 훨씬 안정적인 성능을 유지했습니다.

5. 의의 및 결론 (Significance)

• 실용적 배포 경로 제시: MedFeat 는 의료 현장의 제약 (데이터 불균형, 결측치, 프라이버시, 해석 가능성 요구) 을 고려하여, LLM 을 활용한 확장 가능하고 해석 가능한 특징 공학의 실현 가능한 경로를 제시합니다.
• 모델 중심 설계의 중요성: 특징 공학이 단순히 검증 점수 (Validation Score) 를 높이는 탐색이 아니라, 하류 모델의 인덕티브 바이어스 (Inductive Bias) 를 보완하는 구조를 주입하는 과정임을 입증했습니다.
• 프라이버시 보호: 환자 원본 데이터를 LLM 에 보내지 않고 메타데이터와 설명 가능성 신호만 활용함으로써, 의료 데이터의 민감한 프라이버시 문제를 해결하면서도 LLM 의 지식을 활용할 수 있음을 보여줍니다.

결론적으로 MedFeat 는 의료 표형 데이터 예측에서 전통적인 머신러닝 모델의 성능을 극대화하기 위한 모델 인식적이고 설명 가능성에 기반한 자동 특징 공학 프레임워크로서, 실제 임상 환경에서의 배포와 안정성을 크게 향상시킬 수 있는 잠재력을 가지고 있습니다.