기술 요약: MedLLM – 10억 미만 규모의 오픈 의료 언어 모델

문제 정의

의료 분야의 대규모 언어 모델(LLM)은 단일한 규모로 수렴해 왔습니다. 널리 사용되는 시스템들은 대부분 70억(7B) 개 이상의 파라미터를 가집니다. 이러한 추세는 10억 개 미만(sub-billion, 1B 미만)의 파라미터를 가진 모델의 영역을 미개척 상태로 남겨두었습니다. 현재의 의료 LLM은 비공개 소스이거나(예: GPT-4, Med-PaLM-2) 공개되어 있더라도 규모가 너무 커서(7B–540B) 접근성, 재현성 및 훈련 과정에 대한 감사가 제한적입니다.

모델 용량이 10배 감소할 때 어떤 의료적 역량이 생존하는지를 이해하는 데에는 결정적인 격차가 존재합니다. 일반 목적의 모델을 의료 분야에 적응시키는 것은 웹 텍스트와 의료 텍스트 사이의 분포 불일치(distributional mismatch)로 인해 까다로운데, 이는 롱테일 용어 처리, 구조화된 추론, 고정밀 요구사항을 다뤄야 하기 때문입니다. 기존 연구들은 바이오메디컬 코퍼스에 대한 지속적 사전 학습(continued pretraining)과 지시어 튜닝(instruction tuning)을 통해 이를 해결하려 했으나, 모든 노력이 7B 이상의 파라미터에서 수행되었습니다. 이러한 적응 레시피가 더 작은 모델에도 전이되는지, 아니면 감소된 용량으로 인해 특정 의료 역량이 다른 역량보다 먼저 붕괴되는지는 알려진 바 없습니다.

방법론

모델 아키텍처: MedLLM

MedLLM은 0.1B (100M) 파라미터를 가진 디코더 전용 트랜스포머입니다. 주요 아키텍처 선택 사항은 다음과 같습니다:

• 토큰화 및 임베딩: SentencePiece 어휘집($V=32,000$)과 임베딩 차원 $d=512$를 사용합니다. 높은 어휘 대비 너비 비율($V/d = 62.5$)을 고려하여, 저자들은 입력 임베딩을 출력 프로젝션에 결합(tie)했습니다. 이는 별도의 $V \times d$ 행렬을 제거하여 전체 파라미터의 14%를 절감하며 정규화 역할을 합니다.
• 트랜스포머 레이어: 20개의 레이어를 쌓았으며, 프리 노름(pre-norm) 대신 잔차 브랜치 내부(어텐션 후 및 FFN 후)에 RMSNorm을 배치했습니다. 이는 잔차 스트림(residual stream)을 통제된 규모로 유지하여, 깊고 좁은 모델의 안정성을 높여줍니다.
• 어텐션: 8개의 헤드를 사용하며, RoPE(Rotary Position Embeddings)와 QK-정규화를 사용하여 어텐션 로짓 크기를 제한함으로써 작은 모델에서 흔히 발생하는 훈련 불안정성을 방지합니다.
• 피드 포워드 네트워크(FFN): 확장 계수 4($d_{ff} = 2048$)를 갖춘 SwiGLU를 채택했습니다. 이는 모델 파라미터의 62.7%를 피드 포워드 하위 레이어에 할당하여, 너비보다는 깊이와 구성적 변환을 우선시합니다.

데이터 파이프라인: MedFineWeb

수동으로 큐레이션된 의료 소스(종종 추상적인 스타일의 글쓰기에 편향되거나 희소함)에 의존하지 않고 모델을 적응시키기 위해, 저자들은 **참조 가이드 선택(reference-guided selection)**을 통해 일반 웹 데이터에서 선별된 의료 평문 코퍼스인 MedFineWeb을 도입했습니다.

1. 후보 풀(Candidate Pool): 일반 웹 데이터(RedPajama-V2, C4, OpenWebText, Wikipedia)를 약 1,000단어 단위의 청크로 나눕니다.
2. 참조 세트(Reference Sets): 세 개의 의료 QA 데이터셋(MedMCQA, MedQA, PubMedQA)이 의미론적 참조 역할을 합니다. 이들의 프롬프트와 응답은 임베딩됩니다.
3. 선택: 각 후보 청크는 FAISS를 통한 코사인 유사도를 통해 참조 세트와의 의미론적 유사도로 점수가 매겨집니다. 각 참조 세트에 대해 상위 50%의 후보가 선택되어 병합되고 중복이 제거됩니다.
4. 결과: 의료 스타일과 주제를 포착하면서도 벤치마크 질문이나 답변을 직접 복사하지 않는 3B 토큰 규모의 코퍼스가 생성되었습니다.

훈련 파이싸인

훈련은 완전히 공개된 3단계 파이프라인을 따릅니다.

1. 일반 사전 학습(General Pretraining): 약 29.5B 토큰의 일반 웹 데이터($D_{gen}$)를 사용하여 인과적 언어 모델링을 수행하며, 커리큘럼 시퀀스 길이 스케줄(128 $\to$ 256 $\to$ 512 토큰)과 출력 로짓을 안정화하기 위한 z-loss 정규화자를 사용합니다.
2. 도메인 미세 조정(Domain Fine-Tuning, DFT): 분포를 일반 웹 텍스트에서 의료 텍스트로 전환하기 위해 MedFineWeb에서 25,000 스텝 동안 지속적 사전 학습을 진행합니다.
3. 선호도 정렬(Preference Alignment): 벤치마크당 2단계 프로세스를 거칩니다.
   • 지도 미세 조정(Supervised Fine-Tuning, SFT): 레이블 예측 목적 함수를 사용하여 벤치마크 훈련 세트에서 미세 조정합니다.
   • 직접 선호 최적화(Direct Preference Optimization, DPO): 별도의 보상 모델 없이, 정답을 오답보다 선호하도록 모델을 최적화합니다. KL 페널티 강도는 $\beta=0.1$로 스케일링되어 선호도 마진을 증폭시킵니다.

주요 기여

1. MedLLM: 최초의 오픈 0.1B 파라미터 의료 언어 모델로서, 기존 오픈 의료 LLM보다 10배 작은 재현 가능한 베이스라인을 제공합니다.
2. MedFineWeb: 의료 QA 예시를 통한 임베딩 유사도를 사용하여 일반 웹 데이터로부터 구축된 새로운 의료 사전 학습 코퍼스입니다. 이는 큐레이션된 의료 소스 없이도 구축 가능하며, 타겟 도메인의 소규모 QA 세트만 있으면 어떤 도메인에도 적응할 수 있습니다.
3. 태스크 구조적 해리(Task-Structured Dissociation): 압축 상황에서 의료 역량이 균일하게 저하되지 않는다는 것을 확인했습니다. 대신, 모델은 태스크 유형에 따라 분리됩니다. 즉, 문맥 기반 QA에서는 7B 모델과 경쟁 가능한 정확도를 달 하지만, 지식 회상(knowledge-recall) QA에서는 용량 제한을 받습니다.

결과

사전 학습 평가

태스크 특화 미세 조정 전에도(제로샷 top-token selection으로 평가), MedLLM은 7B 지시어 튜닝 모델(Zephyr-7B-$\beta$: 32.4%) 및 의료 적응 7B 모델(PMC-LLaMA-7B: 30.6%)을 포함한 모든 베이스라인 중 가장 높은 평균 정확도(35.1%)를 기록했습니다. 이 모델은 70배 더 작음에도 불구하고 회상 중심 벤치마크(MMLU 및 MedMCQA)에서 더 큰 모델들을 능가했습니다.

미세 조정 평가

SFT 및 DPO 이후의 결과는 태스크 유형에 따른 뚜렷한 분리를 보여주었습니다.

• 지식 회상 벤치마크 (MedMCQA, MMLU): MedLLM(DPO)은 모든 7B 및 7B 미만 베이스라인을 크게 앞질렀습니다. MedMCQA에서 34.9%에 도달했습니다(Zephyr-7B-$\beta$ 33.0%, PMC-LLaMA-7B 26.6% 대비). MMLU에서는 32.4%에 도달했습니다.
• 문맥 기반 QA (PubMedQA): MedLLM은 58.2%에 도달하여 지시어 튜닝된 7B 베이스라인(Zephyr-7B-$\beta$: 43.0%)을 넘어섰으나, 지속적 사전 학습 베이스라인(PMC-LLaMA-7B: 72.9%)에는 미치지 못했습니다.
• 복잡한 임상 사례 (MedQA): 모델은 어려움을 겪었으며, 4지 선다형 하한선인 25%에 근접한 28.1%에 그쳤습니다. 이는 지시어 튜닝된 7B 모델들에 비해 현저히 낮은 수치입니다.

훈련 역학 분석

• Perplexity vs. 정확도: 도메인 미세 조정 동안 의료 텍스트에 대한 검증 퍼플렉시티는 절반으로 감소(27.3 $\to$ 13.6)했으나, 벤치마크 정확도는 정체되었습니다. 이는 0.1B 규모에서 모델이 의료적 형태(언어 분포)는 빠르게 배우지만, 회상 벤치마크에 필요한 검색 가능한 사실을 습득하는 데는 실패함을 나타냅니다.
• DPO 효과: DPO는 정확도를 약간만 높였으나(예: MedMCQA에서 +1.0 pp), 정답과 오답 사이의 로그 확률 마진을 5.2배 증폭시켰습니다. 이는 이 규모에서의 DPO가 새로운 지식을 주입하기보다는 기존의 변별력을 정교화한다는 것을 시사합니다.

의의 및 주장

본 논문은 10억 미만 규모가 단순히 더 큰 모델의 열화된 버전이 아니라, 의료적 역량이 파라미터적 회상보다는 검색 및 제공된 문맥으로부터 가장 잘 얻어진다는 독특한 행동 양식을 보인다고 주장합니다.

• 용량 vs. 적응: 결과는 회상이 성공하는 곳(MedMCQA)에서는 결합 제약 요인이 적응 방법이 아닌 모델 용량임을 입증합니다. 회상이 실패하는 곳(MedQA)에서는 용량 한계가 절대적입니다.
• 태스크 특이적 해리: 이러한 성능의 "분리"는 7B 규모에서는 가려져 있다가(두 역량이 모두 존재함), 용량이 부족해질 때 비로소 드러납니다.
• 재현성: MedLLM과 MedFineWeb을 공개함으로써, 저자들은 거대한 파라미터 수 없이도 특정 의료 작업을 위해 모델을 훈련, 감사 및 배포할 수 있는 경로를 제공하며, 대규모 파라미터의 필요성에 도전합니다.

저자들은 10억 미만 규모에서 모델의 지식 폭은 본질적으로 제한되며, 성능은 내재된 파라미터적 지식보다는 제공된 문맥의 품질에 크게 의존한다고 결론짓습니다.