기술 요약: 화학 언어 모델을 위한 불확실성 인지 강화 학습

문제 정의

강화 학습(RL)은 화학 언어 모델(CLM)이 화학적 공간을 탐색하고 특정 속성을 최적화할 수 있도록 하여, de novo 분자 설계의 강력한 패러다임으로 부상했습니다. 그러나 기존의 RL 프레임워크는 일반적으로 스코어링 함수를 결정론적 오라클(deterministic oracle)로 취급하며, 분자 특성 예측에 내재된 불확실성을 무시합니다. 이러한 스코어링 함수는 제한적이고 잠재적으로 편향된 데이터셋을 통해 훈련된 예측 모델에 의존하므로, 알레아토릭(aleatoric, 데이터 노이즈) 불확실성과 에피스테믹(epistemic, 모델 파라미터/커버리지) 불확실성을 모두 유발합니다.

RL 에이전트가 점 추정치(point estimates)만을 기준으로 최적화할 경우, 훈련 데이터로부터 멀리 떨어져 예측 불확실성이 높은 화학적 공간 영역을 탐색할 수 있습니다. 이는 훈련 데이터에 의해 충분히 뒷받침되지 않으면서도 높은 점수를 받는 분자들을 생성하게 만들어, 최적화 과정을 불안정하게 만들고 실제 값과는 거리가 먼 예측값을 산출할 위험이 있습니다. 불확실성 정량화는 QSAR(정량적 구조-활성 관계) 모델링에서는 잘 연구되어 있으나, RL 기반 분자 생성 맥락에서는 이에 대한 관심이 제한적이었습니다.

방법론

저자들은 REINVENT4 프레 프레임워크 내에서 예측 불확실성을 RL 루프에 통합하기 위한 두 가지 상호 보완적인 전략을 제안하고 평가합니다. 이 프레임워크는 분자를 샘플링하기 위해 순환 신경망(RNN) 프라이어를 사용하며, 샘플링된 분자들은 다중 파라미터 최적화(MPO) 스코어링 함수에 의해 평가됩니다.

1. 확률적 스코어링 함수

불확실성 전략을 도입하기에 앞서, 저자들은 표준 REINVENT 구현에서의 수학적 불일치를 해결합니다. 스코어링 함수가 전체 예측 분포(예: 평균-분산 추정법을 통한)를 제공할 때, 기댓값 자체에 비선형 변환을 적용하는 것($f(E[X])$)은 분포를 먼저 변환하는 것($E[f(X)]$)과 수학적으로 일치하지 않습니다. 저자들은 M-C 샘플링을 통해 분산을 올바르게 전파함으로써 MPO를 확률 변수로 취급하는 확률적 스코어링 접근 방식을 구현했습니다.

2. 불확실성 통합 전략

본 논문은 불확실성을 활용하는 두 가지 구별된 방법을 소개합니다.

• 스코어 변조 (Score Modulation, SM): 불확실성을 명시적인 추가 최적화 목표로 취급합니다. 예측 불확실성을 활성 및 물리화학적 특성과 함께 MPO 스코어링 함수에 직접 통합합니다. 이는 정책(policy)이 높은 점수(exploitation)와 신뢰성(low uncertainty) 사이에서 절충하도록 유도합니다.
• 손실 변조 (Loss Modulation, LM): 불확실성을 스코어에 더하는 대신, 정책 업데이트의 크기를 조절하는 데 사용합니다. 각 분자가 그래디언트 업데이트에 기여하는 정도를 불확실성에 반비례하여 가중치를 둡니다. 불확실성이 높은 분자는 정책에 미치는 영향력이 줄어들어, 신뢰할 수 없는 학습 신호를 약화시킵니다.

3. 실험 설정

제시된 전략들은 세 가지 설정에서 평가되었습니다:

1. 제어된 모델 시스템 (Model System, MS): 예측 오차가 훈련 데이터와의 거리(Applicability Domain 시뮬레이션)에 비례하는 가우시안 노이즈로 모델링된 합성 환경입니다. 참값(ground truth)을 알고 있으므로 정밀한 오차 측정이 가능합니다.
2. ChemProp 모델: 평균-분산 추정(MVE)으로 훈련된 ChemProp 모델을 사용하는 실제 데이터 시나리오입니다. "Predictor" 모델(소규모 서브셋으로 훈련됨)이 스코어링 함수 역할을 하며, "Oracle" 모델(전체 데이터셋으로 훈련됨)이 참값 근사치 역할을 합니다.
3. 컨포멀 예측 (Conformal Prediction, CP): Mondrian 유도 컨포멀 예측기로 래핑된 랜덤 포레스트 분류기를 사용하는 설정입니다. 불확실성은 외부의 참값 오라클 없이 예측 세트와 p-value를 통해 정량화됩니다.

주요 결과

모델 시스템 결과

제어된 설정에서 표준 "Noisy Component" 베이스라인은 스코어링 불안정성으로 인해 낮은 불확실성 영역으로 이동하는 일종의 암묵적 자기 규제(self-regularization) 현상을 보였으나, 이는 느리고 초기 조건에 의존적이었습니다.

• SM 및 LM 성능: 두 명시적 전략 모두 베이스라인에 비해 저불확실성 영역으로의 이동을 가속화했습니다.
• Hit 품질: LM 전략이 가장 우수한 성능을 보였으며, 베이스라인과 비교하여 총 히트 대비 실제 히트(true hits)의 비율을 유의미하게 높였습니다.
• 결합 접근 방식 (SM&LM): 이 방식은 저불확실성 영역으로 가장 빠르게 이동했지만, 누적된 히트 스캐폴드(hit-scaffolds)의 총수가 크게 감소했습니다. 이는 제어된 설정에서 다양성과 탐색을 저해하는 과잉 교정(over-correction)이 발생했음을 시사합니다.

실제 데이터 결과 (ChemProp)

ChemProp 시나리오에서 베이스라인(Predictor 전용)은 높은 활성도를 보였으나 Oracle 점수와 상당한 괴리를 보였습니다. 이는 모델이 Predictor가 활성도를 과대평가하는 고불확실성 영역을 악용하고 있음을 나타냅니다.

• 신뢰성: 불확실성 가이드 전략(특히 LM)은 Predictor와 Oracle 점수 사이의 간극을 더 작고 안정적으로 유지했습니다.
• 실제 히트 (True Hits): LM 전략은 가장 많은 실제 히트를 생성했으며, 실제 히트 대 총 히트 비율을 (0.5에서 0.75로) 거의 두 배 가까이 높이며 가장 큰 개선을 보였습니다.
• 불확실성: 모든 가이드 전략은 생성된 분자들의 예측 불확실성을 성공적으로 감소시켰습니다.

실제 데이터 결과 (Conformal Prediction)

참값 오라클이 없는 상태에서 "실제 히트"는 컨포멀 예측기에 의해 고유한 클래스 레이블(active)이 할당된 화합물로 정의되었습니다.

• 모호성 감소: 베이스라인 최적화는 모호한 예측(active와 inactive 클래스 모두에 할당됨)을 다수 생성했습니다. 불확실성 가이드 전략은 이러한 모호성을 유의미하게 감소시켰습니다.
• 히트 확신도: 총 히트 수는 모든 전략에서 유사했으나, 불확실성 가이드 접근 방식은 더 높은 비율의 확신 있는 고유 클래스 예측을 생성했습니다.

의의 및 주장

본 논문은 예측 불확실성을 RL 기반 분자 생성에 통합하는 것이 다양한 실험 설정에서 일관되고 의미 있는 개선을 가져온다고 주장합니다. 주요 의의는 불확실성이 단순한 진단 도구가 아니라 최적화 과정의 능동적인 구성 요소가 될 수 있음을 입증한 데 있습니다.

• 안정성 및 신뢰성: 불확실성 인지 전략은 스코어링 함수가 신뢰할 수 있는 영역을 선호하게 함으로써 탐색을 안정화하고, 허위 예측을 악용할 위험을 줄입니다.
• 손실 변조의 우월성: 저자들은 손실 변조 (Loss Modulation, LM) 전략이 가장 견고한 접근 방식임을 확인했습니다. 스코어와 불확실성 추정치를 혼합하여 보상 랜드스케이프를 변경하는 스코어 변조와 달리, LM은 원래의 스코어 목적 함수의 무결성을 유지하면서 신뢰할 수 없는 샘플이 정책 학습에 미치는 영향을 감쇄시킵니다.
• 실질적 영향: 이 접근 방식은 분자 점수를 훼손하지 않으면서 히트 발견의 신뢰성을 높이며, RL이 분자 최적화를 가이드하고 예측 불확실성을 추정할 수 있는 모든 생성 프레임워크에 폭넓게 적용될 수 있습니다.

저자들은 결론적으로, 훈련 데이터와 거리가 먼 화학적 공간을 탐색할 때 분자의 품질과 신뢰성을 높이기 위해 불확실성을 고려하는 것이 필수적임을 강조합니다.