[논문 리뷰] UltraMedical: Building Specialized Generalists in Biomedicine

Kaiyan Zhang, Sihang Zeng외 12명|arXiv (Cornell University)|2024. 06. 06.

Clinical Reasoning and Diagnostic Skills인용 수 6

한 줄 요약

본 논문은 대규모 생의학 지시 데이터셋 UltraMedical와 Llama-3 모델을 생의학 일반화 전문가로 미세조정하기 위한 오픈 소스 워크플로를 제시하며, 생의학 보상 모델과 반복 선호 학습의 도움을 받는다. 경쟁력 있는 의료 벤치마크 성능을 보고하고 데이터셋과 모델을 공개한다.

ABSTRACT

Large Language Models (LLMs) have demonstrated remarkable capabilities across various domains and are moving towards more specialized areas. Recent advanced proprietary models such as GPT-4 and Gemini have achieved significant advancements in biomedicine, which have also raised privacy and security challenges. The construction of specialized generalists hinges largely on high-quality datasets, enhanced by techniques like supervised fine-tuning and reinforcement learning from human or AI feedback, and direct preference optimization. However, these leading technologies (e.g., preference learning) are still significantly limited in the open source community due to the scarcity of specialized data. In this paper, we present the UltraMedical collections, which consist of high-quality manual and synthetic datasets in the biomedicine domain, featuring preference annotations across multiple advanced LLMs. By utilizing these datasets, we fine-tune a suite of specialized medical models based on Llama-3 series, demonstrating breathtaking capabilities across various medical benchmarks. Moreover, we develop powerful reward models skilled in biomedical and general reward benchmark, enhancing further online preference learning within the biomedical LLM community. Datasets and models are available at https://github.com/TsinghuaC3I/UltraMedical

연구 동기 및 목표

• 데이터 중심의 미세조정과 선호 학습을 통해 오픈 소스 GPT-4 수준 생의학 일반 목표 대형 모델을 구축한다.
• 수작업 프롬프트와 합성 프롬프트를 결합한 고품질의 다양하고 복잡한 생의학 지시를 생성한다.
• 온라인 선호 학습과 반복적 모델 개선을 가능하게 하는 보상 모델을 개발한다.
• 표준 의학 벤치마크 및 일반 도메인 과제에서 UltraMedical LMs를 평가하여 전문화와 일반화의 트레이드오프를 특징화한다.
• 생의학 생성 AI 커뮤니티의 협업을 촉진하기 위해 데이터셋과 모델을 공개한다.

제안 방법

• 수작업 및 합성 의학 프롬프트를 합쳐 약 410k 지시로 구성된 UltraMedical를 구축한다.
• 일부(~100k)의 프롬프트에 모델이 완성한 응답과 GPT-4에서 파생된 선호도에 주석화하여 순위를 매긴다.
• UltraMedical(및 UltraMix)을 사용하여 Llama-3-8B/70B 모델을 감독 학습(SFT)으로 미세조정한다.
• 프롬프트당 8개 완성 구성을 사용하고 선택/거부를 이진화한 쌍으로 DPO, KTO 등의 선호 학습 기법을 적용한다.
• UltraMedical 선호도 및 기타 Ultra-시리즈 데이터를 기반으로 생의학 보상 모델(RM)을 학습하여 온라인 선호 학습 및 BoN 샘플링을 가능하게 한다.
• 점진적으로 모델 성능을 향상시키기 위해 반복적인 온라인 선호 학습 및 Best-of-N 샘플링을 수행한다.

Figure 1 : The UltraMedical Datasets, Models and Performance on MedQA.

실험 결과

연구 질문

• RQ1 데이터셋 설계와 선호 학습을 통해 오픈 소스 생의학 LLM이 의학 분야에서 GPT-4 수준의 능력을 달성할 수 있는가?
• RQ2 의학 데이터와 일반 도메인 데이터를 혼합하는 것이 SFT 및 선호 최적화에 어떤 영향을 미쳐 전문화된 일반 전문가에게 어떤 효과가 있는가?
• RQ3 Llama-3 기반 모델의 의료 영역 능력을 끌어올리는 데 보상 모델과 온라인 선호 학습이 얼마나 효과적인가?
• RQ4 전문화된 생의학 일반가를 개발할 때 의료 작업 성능과 일반 도메인 능력 사이에 어떤 트레이드오프가 나타나는가?
• RQ5 공개적으로 공개된 UltraMedical 데이터세트가 경쟁력 있는 오픈 생의학 LLM을 얼마나 촉진할 수 있는가?

주요 결과

• UltraMedical 8B/70B 모델은 의료 벤치마크에서 진보된 성능을 달성하여 독점 모델과의 격차를 좁힌다.
• 70B UltraMedical 모델은 MedQA-USMLE에서 86.5를 달성하여 일부 설정에서 더 큰 독점 모델과 비슷하다.
• SFT 및 xPO에서 의학 데이터와 일반 도메인 데이터를 혼합하면 의학 작업 성능이 향상되지만 일부 일반 도메인 트레이드오프가 발생한다.
• UltraMedical 선호도 기반으로 학습된 보상 모델은 온라인/반복적 선호 학습 및 BoN 샘플링을 지원하여 의학 및 일반 보상 벤치마크에서 RM 기반의 경쟁력 있는 결과를 낳는다.
• 일부 경우에서 보상 모델이 재랭킹을 위한 self-consistency 앙상블보다 우수할 수 있지만, 모델 크기에 따라 효과가 다르다.
• UltraMedical 데이터셋과 모델의 공개는 생의학 생성 AI 커뮤니티의 진전을 촉진한다.

Figure 2 : The Construction Pipeline for the UltraMedical Dataset.