분자 그래프 기반 모델의 강건 파인튜닝 벤치마크와 새로운 방법 제안

초록

본 논문은 분자 그래프 기반 기초 모델(MGFM)의 파인튜닝에서 발생하는 과적합·라벨 희소성 문제를 해결하기 위해 8가지 파인튜닝 기법을 가중치 기반, 표현 기반, 부분 파인튜닝으로 분류하고, 다양한 회귀·분류 벤치마크와 라벨 수·분포 변화를 고려한 실험을 수행한다. 실험 결과를 토대로 기존 방법들의 장단점을 종합한 새로운 가중치 기반 기법 DWiSE‑FT를 제안하여, 회귀·분류 모두에서 기존 최고 성능을 능가하면서 사용 편의성도 유지한다.

상세 분석

본 연구는 분자 그래프 기초 모델(MGFM)이 직면한 두 가지 근본적인 제약, 즉 (1) 사전학습에 사용할 수 있는 데이터 규모가 수억 개 수준으로 제한적이며(이미지·텍스트 분야와 달리 수십억 수준에 비해 현저히 작음), (2) 실제 다운스트림 화학·생물학 과제에서 라벨이 수십에서 수백 개에 불과한 극심한 데이터 희소성을 가진다는 점에 주목한다. 이러한 환경에서는 전통적인 전면 파인튜닝(full‑FT)이 쉽게 과적합에 빠지거나 사전학습에서 얻은 일반화 능력을 손상시킨다.

논문은 먼저 8가지 파인튜닝 기법을 “가중치 기반(Weight‑based)”, “표현 기반(Representation‑based)”, “부분 파인튜닝(Partial‑FT)”이라는 세 메커니즘으로 체계화한다.

• 가중치 기반은 사전학습 가중치와 파인튜닝 후 가중치를 선형 혹은 정규화된 방식으로 결합한다. 대표적으로 WiSE‑FT는 사전학습 가중치와 파인튜닝 가중치를 α값으로 선형 보간하고, L2‑SP는 파인튜닝 과정에서 가중치가 사전학습 가중치와의 L2 거리보다 크게 변하지 않도록 정규화한다.
• 표현 기반은 파인튜닝 중에 모델의 내부 임베딩이 사전학습 임베딩과 크게 차이나지 않도록 제약한다. Feature‑Map은 임베딩 차이를 L2 손실로 penalize하고, BSS는 작은 특이값에 해당하는 스펙트럼 성분을 억제함으로써 전이 가능한 표현을 보존한다.
• 부분 파인튜닝은 전체 파라미터를 업데이트하지 않고, 헤드만 학습하거나 특정 레이어만 미세 조정한다. Linear Probing(LP), Sur‑FT, LP‑FT가 이에 해당한다.

다음으로, 논문은 세 종류의 사전학습 모델(자기지도 기반 GraphMAE, Mole‑BERT, 멀티모달 MoleculeSTM)과 두 종류의 지도학습 모델(Graphium‑Toy, Graphium‑Large, GraphGPS) 총 6개를 선정한다. 이들은 그래프‑CNN, 그래프‑Transformer, 텍스트‑멀티모달 등 구조적 다양성을 포괄한다.

다운스트림 평가는 8개의 분류 데이터셋(예: BBBP, Tox21, HIV 등)과 4개의 회귀 데이터셋(예: ESOL, Lipo, CEP, Malaria)을 사용한다. 각 데이터셋에 대해 랜덤, 스캐폴드, 사이즈 기반 3가지 분할을 적용해 ID와 OOD 상황을 모두 시뮬레이션한다. 또한 라벨 수를 50, 100, 500개의 few‑shot 설정과 전체 데이터를 활용하는 non‑few‑shot 설정으로 나누어, 라벨 희소성이 파인튜닝 성능에 미치는 영향을 정량화한다.

실험 결과는 다음과 같은 핵심 인사이트를 제공한다.

1. 사전학습 방식의 영향: few‑shot 상황에서는 지도학습(pre‑trained on task‑specific labels) 모델이 일반적으로 자기지도 모델보다 우수한 성능을 보인다. 반면 라벨이 충분히 확보된 non‑few‑shot 상황에서는 두 방식의 격차가 줄어들며, 특히 사전학습 과제가 다운스트림 과제와 유사할 경우 지도학습이 더 큰 이점을 갖는다.
2. 작업 유형에 따른 파인튜닝 메커니즘 차이: 회귀 과제는 정밀한 수치 예측이 요구되므로 과적합 위험이 낮고, 가중치 기반 기법이 특히 강력하게 작동한다. 반면 분류 과제는 라벨 불균형·희소성에 민감해 부분 파인튜닝이나 표현 기반 정규화가 도움이 된다.
3. 가중치 기반 vs 부분 파인튜닝 vs 표현 기반: 자기지도 사전학습 모델에서는 가중치 기반(특히 WiSE‑FT와 L2‑SP)의 조합이 가장 일관된 성능 향상을 제공한다. 지도학습 모델에서는 표현 기반(Feature‑Map)이 사전학습된 도메인 지식을 보존하면서도 다운스트림 적응을 가능하게 해 좋은 결과를 만든다. 부분 파인튜닝은 특히 few‑shot 회귀에서 과소적합을 초래해 성능이 저하되는 경향이 있다.

이러한 발견을 바탕으로 저자들은 DWiSE‑FT라는 새로운 가중치 기반 방법을 설계한다. DWiSE‑FT는 기존 WiSE‑FT의 단순 선형 보간에 L2‑SP가 강조하는 스펙트럼 정규화를 결합해, 회귀·분류 모두에서 강건한 전이 성능을 달성한다. 실험에서는 DWiSE‑FT가 기존 최고 성능 방법을 평균 2~4%p(percentage points) 정도 능가했으며, 특히 OOD(스캐폴드·사이즈) 상황과 라벨이 극히 제한된 few‑shot 설정에서 두드러진 개선을 보였다. 또한 DWiSE‑FT는 사전학습 가중치와 파인튜닝 가중치를 후처리 단계에서만 결합하므로 구현이 간단하고 하이퍼파라미터 튜닝 부담이 적다.

결론적으로, 본 논문은 MGFM 파인튜닝에 대한 체계적인 벤치마크와 메커니즘별 성능 특성을 제시함으로써, 연구자들이 사전학습 목표와 다운스트림 작업 유형에 맞는 최적의 파인튜닝 전략을 선택하도록 돕는다. 또한 제안된 DWiSE‑FT는 현재 가장 실용적이며 성능이 뛰어난 솔루션으로, 향후 다양한 화학·생물학 응용에 바로 적용될 수 있다.