생물의학 지식 정렬을 위한 균형 파인튜닝

초록

본 논문은 의료·생물학 분야의 희소하고 파편화된 텍스트 데이터를 효과적으로 학습하기 위해 “균형 파인튜닝(Balanced Fine‑Tuning, BFT)”이라는 새로운 사후 학습 방식을 제안한다. BFT는 토큰 수준에서 예측 확률에 기반한 손실 스케일링과 샘플 수준에서 최소 그룹 신뢰도(minimum group confidence)를 활용한 가중 학습을 결합해 과적합을 방지하고 어려운 샘플을 강조한다. 실험 결과, BFT는 기존의 지도학습(SFT)보다 의료 질문 응답·생물학적 과정 추론에서 현저히 높은 성능을 보이며, GeneAgent을 능가한다. 또한 BFT가 만든 텍스트 임베딩은 유전자 상호작용 예측·단일세포 교란 반응 예측 등 다양한 다운스트림 작업에 바로 활용 가능함을 입증한다.

상세 분석

본 연구는 대규모 언어 모델(LLM)을 생물의학 분야에 정렬시키는 과정에서 두 가지 근본적인 한계를 지적한다. 첫째, 생물학적 메커니즘은 실험 데이터가 제한적이고 텍스트로는 파편화된 형태로 존재한다는 점이다. 기존의 지도학습(Supervised Fine‑Tuning, SFT)은 표면적인 명령어 패턴에 과도하게 적합(overfit)되어, 이러한 희소한 지식을 내재화하지 못한다. 둘째, 강화학습(Reinforcement Learning, RL)은 보상 설계가 현실적인 실험 검증을 필요로 하므로, 생명과학에서는 실시간 피드백을 제공하기 어렵다. 이러한 문제점을 해결하기 위해 제안된 BFT는 두 단계의 가중치 메커니즘을 도입한다.

1️⃣ 토큰 수준 가중치: 각 토큰의 손실을 해당 토큰이 모델이 예측한 확률에 비례하여 스케일링한다. 구체적으로, 손실 L_i를 1/p_i (p_i는 토큰 i에 대한 예측 확률) 로 곱함으로써, 모델이 자신 있게 예측한 토큰은 손실이 감소하고, 불확실한 토큰은 손실이 증폭된다. 이는 그래디언트가 특정 토큰에 집중되는 현상을 완화하고, 전체 학습 과정의 안정성을 높인다.

2️⃣ 샘플 수준 가중치: 학습 데이터는 “그룹”으로 묶이며, 각 그룹 내 최소 신뢰도(minimum group confidence)를 계산한다. 그룹은 동일한 생물학적 주제(예: 특정 신호 경로, 약물 반응) 혹은 난이도 기반으로 정의된다. 최소 신뢰도가 낮은 그룹은 전체 샘플에 대해 높은 가중치를 부여받아, 어려운 샘플이 학습에 충분히 반영되도록 한다. 이는 SFT에서 흔히 발생하는 ‘쉬운 샘플에만 과도하게 학습’되는 현상을 방지한다.

BFT는 별도의 보상 함수 없이도 복잡한 추론 능력을 향상시킨다. 실험에서는 MedQA, PubMedQA 등 의료 질문 응답 데이터와, Gene Ontology 기반 생물학적 과정 추론 과제를 사용하였다. BFT 기반 모델은 SFT 대비 평균 7~12%의 정확도 상승을 보였으며, 특히 GeneAgent이 강세를 보이던 ‘생물학적 과정 추론’에서는 4% 이상의 절대 성능 차이로 우위를 점했다.

또한, BFT가 생성한 텍스트 임베딩은 별도 파인튜닝 없이도 유전자‑유전자 상호작용 예측, 단일세포 교란 반응 예측 등에 적용되었다. 이때 임베딩 기반 선형 모델은 기존 특화된 그래프 기반 방법과 경쟁하거나 이를 능가하는 성능을 기록했다.

종합적으로, BFT는 (1) 토큰‑샘플 이중 가중치를 통한 과적합 억제, (2) 희소하고 복잡한 생물의학 지식의 효율적 내재화, (3) 별도 보상 설계 없이도 강화학습 수준의 성능 향상을 가능하게 하는 혁신적 사후 학습 프레임워크라 할 수 있다.