기술 요약: ATHENA-R1 – 치료 추론을 위한 AI 에이전트

문제 정의

치료 추론(Treatment reasoning)은 질병의 맥락, 환자의 특성, 동반 질환, 병용 약물, 금기 사항, 그리고 진화하는 생물 의학적 지식의 통합을 요구하는 의학의 근간이자 매우 까다로운 과업입니다. 단순한 사실 검색이나 패턴 인식과 달리, 치료 추론은 본질적으로 반복적입니다. 즉, 후보 전략을 수집하고, 여러 제약 조건에 따라 이를 평가하며, 새로운 증거가 나타남에 따라 분석을 수정해야 합니다.

현재의 대규모 언어 모델(LLM)은 주로 모델 가중치에 저장된 파라미터 지식(parametric knowledge)에 의존합니다. 이러한 모델들은 유창한 응답을 생성하고 광범위한 임상 패턴을 포착할 수는 있지만, 업데이트되고 검증된 의학 정보에 접근하는 능력이 부족하며, 특정 금기 사항이나 환자의 제약 조건을 고려하지 못한 권고를 생성할 수 있습니다. 검색 증강 생성(RAG)과 도구 증강 LLM이 외부 데이터에 대한 접근성을 제공하긴 하지만, 단순히 생물 의학 도구에 접근할 수 있다고 해서 치료 추론이 보장되는 것은 아닙니다. 모델은 어떤 증거가 필요한지, 적절한 출처를 무엇을 선택할지, 결과를 맥락에 맞게 해석할지, 그리고 증거가 불완전하거나 상충할 때 분석을 어떻게 수정할지를 자율적으로 결정할 수 있어야 합니다. 이러한 능력은 도구에 대한 접근성만으로는 가정될 수 없으며, 반드시 학습되어야 합니다.

방법론

저자들은 1939년 이후 승인된 모든 FDA 승인 약물을 대상으로 치료 추론을 수행하도록 설계된 AI 에이전트인 ATHENA-R1을 소개합니다. 이 시스템은 다단계 분석과 212개의 생물 의학 도구 유니버스에 대한 직접적인 접근을 결방하여 작동합니다.

핵심 아키텍처

ATHENA-R1은 단 한 단계로 답을 내놓지 않습니다. 대신 다음과 같은 과정을 반복합니다:

1. 필요한 정보 식별: 특정 치료 결정에 필요한 누락된 정보를 파악합니다.
2. 관련 도구 선택 및 실행: 212개의 생물 의학 리소스 라이브러리(예: 약물 상호작 작용, 금기 사항, 약리학, 질병 표현형 쿼리)에서 관련 도구를 선택하고 실행합니다.
3. 증거 검색 및 해석: FDA 처방 정보 및 생물 의학 지식 베이스를 포함한 큐레이션된 소스로부터 증거를 검색하고 해석합니다.
4. 분석 업데이트: 다음 단계로 넘어가기 전, 검색된 증거를 바탕으로 분석 내용을 업데이트합니다.

학습 프레k워크: 2단계 자기 학습(Two-Level Self-Learning)

학nom습에 필요한 방대하고 다양하며 다단계적인 추론 궤적(reasoning traces)을 생성하는 것은 인간 주석가들에게 불가능한 작업입니다. 따라서 ATHENA-R1은 2단계 자기 학습 프레임워크를 통해 학습됩니다.

1. 단계 1: 멀티 에이전트 구성을 통한 지도 미세 조정(SFT)

   • 멀티 에이전트 시스템이 생물 의학 도구, 치료 과업 및 추론 궤적을 자동으로 구축합니다.
   • 이 과정은 85,340개의 추론 궤적에서 파생된 378,027개의 인스트럭션 튜닝 샘플로 구성된 ATHENA-R1-Instruct 데이터셋을 생성합니다.
   • 이 데이터셋은 1939년 이후의 FDA 약물 라벨에 근거한 177,626개의 추론 단계와 281,695개의 도구 호출을 포함합니다.
   • 이 단계는 문제 분해, 증거 검색, 도구 사용, 증거 해석을 포함한 치료 추론의 구조를 모델에게 가르칩니다.

2. 단계 2: 과학적 피드백을 통한 강화 학습(RL)

   • 모델은 212개의 도구 라이브러리가 있는 라이브 환경에서 정교화됩니다.
   • 각 프롬프트에 대해 다회차(multi-turn) 추론 궤적을 생성합니다.
   • 보상 함수는 정답의 정확성, 출력 형식의 유효성, 증거 수집, 다단계 추론, 도구 인자 근거 설정(tool-argument grounding), 추론의 비중복성 등 6가지 차원에서 과학적 피드백을 제공합니다.
   • 그룹 상대 정책 최적화(GRPO)를 사용하여 더 높은 점수를 받은 궤적의 확률을 높임으로써, 에이전트의 증거 탐색 전략과 도구 선택 능력을 정교화합니다.

주요 기여

• ATHENA-R1 에이전트: 단순한 파라미터 지식 의존을 넘어, 실시간 증거 수집과 다단계 분석을 통합하는 혁신적인 반복적 치료 추론 가능 AI 에이전트입니다.
• 자기 학습 프레임워크: 인간이 작성한 궤적 없이도 멀티 에이전트 시스템을 통해 학습 데이터를 생성하고, 과학적 피드백이 있는 RL을 통해 추론 전략을 개선하는 확장 가능한 방법론을 제시합니다.
• 포괄적인 도구 유니버스: 적응증, 용법/용량, 안전성, 약리학 및 환자별 제약 사항을 다루는 212개의 도구 라이브러리를 통합하고 활용합니다.
• 평가 스위트: 희귀 질환 전문가, 현직 의사 및 인구 규모의 전자 건강 기록(EHR) 분석을 포함하는 엄격한 벤치마크(DrugPC 및 TreatmentPC)와 실세계 검증 프로토콜을 개발했습니다.

결과

벤치마크 성능

ATHENA-R1은 3,168개의 약물 추론 과업과 456개의 환자 치료 사례를 아우르는 5개의 데이터셋에서 평가되었습니다.

• DrugPC (약물 처방 카드): ATHENA-R1은 개방형 약물 추론에서 94.7%의 정확도를 달성하여, GPT-5(76.9%), DeepSeek-R1(68.8%), Qwen3(48.7%)를 능가했습니다.
• TreatmentPC (환자 치료 선택): ATHENA-R1은 82.9%의 정확도를 기록하여, GPT-5(72.2%)보다 10.7%포인트, DeepSeek-R1(67.5%)보다 15.4%포인트 높은 성능을 보였습니다.
• 도구 사용의 필수성: 도구에 대한 접근성은 있으나 도구 출력에 기반해 추론하는 학습된 능력이 없는 모델들은 낮은 성능을 보였습니다. 예를 들어, GPT-5에게 선택적 도구 접근권을 주었을 때, 도구를 호출한 경우는 1%에 불과했으며 정확도는 도구가 없는 기본 모델보다 낮아졌습니다. 반면 ATHENA-R1은 모든 케이스에서 도구를 호출했는데, 이는 도구의 가용성이 아니라 도구 출력을 바탕으로 추론하는 능력이 핵심 제한 요소임을 보여줍니다.

전문가 평가

• 희귀 질환 맥락: 28개 희귀 질환 단체의 전문가 23명이 참여한 블라인드 평가에서, ATHENA-R1은 8가지 모든 평가 기준에서 참조 모델들보다 선호되었습니다. 가장 큰 격차는 **인지적 추적 가능성(95.5% 선호)**과 **근거의 유용성(94.5% 선호)**에서 나타났습니다.
• 임상 사례: 현직 의사들은 심혈관 관리 및 감염병과 같이 상충하는 생리학적 제약(예: CABG 수술 후 신장 질환 환자, 항응고제 복용 중인 감염 환자)을 고려해야 하는 복잡한 입원 환자 케이스에서 ATHENA-R1을 긍정적으로 평가했습니다.

인구 규모 검증

저자들은 ATHENA-R1이 생성한 이상 사례 가설을 540만 명의 환자로부터 추출한 종단적 EHR 데이터를 통해 테스트했습니다.

• 가설 생성: 시스템은 기존 약물 감시 증거가 제한적인 예측을 우선순위에 두었습니다.
• 통계적 검증: 예측된 연관성은 고위험 하위 집단에서 **1.48~1.84의 조정 오즈비(OR)**에 도달했습니다.
  • 예시: 고혈압과 통풍을 앓으며 $\beta$-차단제를 복용하는 환자에게서 급성 신부전(OR 1.84) 및 고칼륨혈증(OR 1.78) 위험 증가.
  • 예시: 당뇨병과 허혈성 심장 질환을 앓으며 DPP-4 억제제를 복용하는 환자에게서 간세포암(OR 1.48) 위험 증가.
• 대조군: 음성 대조군은 위험 상승을 보이지 않았으며, 양성 대조군은 확립된 임상 효과를 성공적으로 회복함으로써, 이 파이프라인이 유의미한 신호와 노이즈를 구분하는 능력을 검증했습니다.

의의 및 주장

본 논문은 치료 추론이 결론을 내리기 전에 어떤 증거를 찾아야 하는지 알아야 한다는 점에서 AI가 어려워했던 영역이지만, 이를 반복적인 증거 수집의 학습 가능한 과정으로 재정의할 수 있다고 주장합니다.

주요 주장은 다음과 같습니다:

1. 추론과 지식의 분리: 추론과 지식 저장을 분리함으로써, ATHENA-R1은 정적인 모델 가중치가 아닌 검색 가능한 소스에 근거하여 결론이 형성되는 투명한 기록을 유지하면서도 새로운 증거를 통합할 수 있습니다.
2. 스케일링보다 효과적인 RL: 8B 파라미터 모델 기반인 ATHENA-R1이 훨씬 더 큰 모델인 DeepSeek-R1(671B 파라미터)보다 일관되게 높은 성능을 보이는 것은, 반복적 증거 수집이 필요한 과업에서는 단순한 규모 확장보다 타겟팅된 도구 사용 훈련과 강화 학습이 더 효과적일 수 있음을 시사합니다.
3. 의료 분야 AI의 재정의: 이 결과는 강화 학습이 단순한 사실 검색을 넘어, 복잡하고 맥락을 고려한 임상 추론을 뒷받침하는 증거 탐색 과정을 학습할 수 있음을 입증합니다.

저자들은 도구 라이브러리의 커버리지 의존성, 명시적인 불확실성 정량화의 부재, 인간이 작성한 것이 아닌 생성된 궤적에 의존한다는 점 등을 한계로 언급했습니다. 이들은 ATHENA-R1을 진단 현장에서 사용하는 참조 도구가 아니라, AI가 치료적 의사결정에 필요한 반복적 과정을 학습할 수 있음을 보여주는 연구용 시스템으로 규정합니다.