🤖 machine learning

Conformal Prediction for Molecular Properties under Label Shift

이 논문은 모델 재학습 없이도 AI 기반 신약 개발의 신뢰성과 규제 준수성을 향상시키기 위해, 주변 레이블 확률 비율로 점수를 가중치 부여함으로써 분자 특성에 대한 통계적으로 엄격한 예측 구간을 생성하는 레이블 시프트(label shift)에 특화된 컨포멀 예측 프레임워크를 소개한다.

원저자: Hyeonsu Lee, Juyeon Kim, Erkhembayar Jadamba, Seungjin Choi, Hyunjin Shin

게시일 2026-08-19

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Hyeonsu Lee, Juyeon Kim, Erkhembayar Jadamba, Seungjin Choi, Hyunjin Shin

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

화학 화합물을 생명을 구하는 의약품으로 바꾸는 여정은 불확실성의 마라톤이다. 이 과정은 10년 이상 걸릴 수 있고 수십억 달러의 비용이 들 수 있으며, 컴퓨터 모델에서는 완벽해 보였던 분자가 살아있는 생체 내에서 예측 불가능하게 행동하기 때문에 종종 실패로 끝나기도 한다. 이를 가속화하기 위해 과학자들은 새로운 분자가 어떻게 작용할지(예를 들어 물에 녹을 수 있는지 또는 세포에 독성을 가질지 등)를 예측하기 위해 인공지능에 점점 더 많이 의존하고 있다. 그러나 이러한 디지털 예측은 그 예측이 학습한 데이터가 얼마나 우수한가에 달려 있다. 모델이 학습했던 사례들과 크게 다른 새로운 유형의 분자를 마주할 때, 모델의 자신감은 오히려 독이 되어 단 하나의 정밀한 숫자를 확신 있게 제시하지만 실제로는 틀린 답을 내놓을 수 있다. 이러한 도구들을 실제 환경에서 안전하게 사용하기 위해서, 연구자들은 단순히 정답뿐만 아니라, 특히 게임의 규칙이 변할 때 그 정답이 얼마나 신뢰할 수 있는지를 측정할 방법이 필요하다.

최근 모감 생물 의학 연구소(Mogam Institute for Biomedical Research)와 인티코드(Intellicode)의 연구진은 기술적 용어로 '레이블 시프트(label shift)'라고 불리는, 규칙이 변하는 이 특정한 문제를 다루었다. 이는 모델이 예측하려는 속성의 분포가 훈련 단계와 실제 적용 단계 사이에서 변할 때 발생한다. 예를 들어, 모델이 다양한 일반 분자들을 학습했는데, 이후 매우 희귀하면서도 효능이 극도로 높은 화합물을 찾아내라는 요청을 받는 경우를 들 수 있다. 이 시나리오에서는 불확실성을 측정하는 표준적인 방식이 무너지며, 종종 예측 구간이 실제 값을 포착하기에는 너무 좁게 설정되는 문제가 발생한다. 연구팀은 값비싼 기존 모델을 다시 학습시킬 필요 없이 AI 예측의 신뢰 구간을 조정하는 새로운 프레임워크를 개발했다. 단일 점 추정치 대신 가능한 값들의 범위를 구축하는 '컨포멀 예측(conformal prediction)'이라는 통계적 기법을 사용하여, 데이터 분포가 변하더라도 신뢰성을 유지하는 시스템을 만들어냈다.

연구진은 알려진 용해도 수준을 가진 약 1만 개의 화합물로 구성된 대규모 데이터셋을 사용하여 이들의 방법을 테스트했는데, 용해도는 약물 설계의 핵심 요소다. 그들은 수억 개의 화학 구조를 학습한 정교한 언어 모델을 사용하여 초기 예측을 수행했다. 레이블 시프트라는 현실 세계의 과제를 시뮬레이션하기 위해, 연구진은 테스트 데이터의 용해도 분포를 인위적으로 변경하여 훈련 세트와 다르게 만들었다. 이처럼 변화된 데이터에 조정되지 않은 표준 방식을 적용했을 때, 시스템은 예측된 범위 안에 실제 값을 포착해야 하는 횟수만큼 포착하지 못했으며, 이는 결과가 위험할 정도로 과도한 자신감을 갖게 함을 보여주었다. 테스트 데이터가 훈련 데이터와 유사할 것이라고 가정하는 전통적인 방식은 이러한 변화를 따라잡을 수 없었다.

이를 해결하기 위해 연구진은 예측에 대한 '교정 렌즈' 역할을 하는 가중치 시스템을 도입했다. 그들은 먼저 가용한 데이터를 세 개의 뚜렷한 그룹으로 나누었다. 하나는 모델을 훈련시키기 위한 것, 하나는 데이터 분포가 어떻게 변했는지 추정하기 위한 것, 그리고 마지막 하나는 최종 예측 범위를 보정하기 위한 것이다. 두 번째 그룹을 사용하여, 연구진은 새로운 환경에서의 다양한 용해도 수준의 기대 빈도와 이전 환경에서의 빈도 사이의 비율을 계산했다. 그런 다음 이 비율을 모델이 보정 과정에서 범한 오류에 가중치로 적용했다. 이 과정은 시스템에 새로운 데이터에서 더 흔해지고 있는 유형의 오류에는 더 주목하고, 사라져가는 유형의 오류에는 덜 주목하도록 효과적으로 지시하는 역할을 했다. 이를 통해 연구진은 데이터가 어떻게 변하더라도 원하는 신뢰 수준에서 실제 값이 예측 범위 안에 들어오도록 하여 통계적 타당성을 유지하는 예측 구간을 구축할 수 있었다.

견고함을 보장하기 위해 1,000번 반복 수행된 시뮬레이션 결과, 명확한 개선이 나타났다. 표준 방식은 레이블 시프트 상황에서 실제 값을 자주 놓치는 구간을 생성한 반면, 새로운 가중치 적용 방식은 커버리지(coverage)를 목표 수준으로 일관되게 회복시켰다. 연구진은 필요한 가중치를 계산하는 데 '최대 가능도 추정(maximum likelihood estimation)'이라는 특정 통계 기법을 사용할 때 가장 효과적이라는 것을 발견했으나, 다른 방법들도 유망한 결과를 보였다. 흥미롭게도, 커버리지를 회복하는 데 가장 정확한 방법은 다른 방법들보다 약간 더 넓은 구간을 생성하는 경향이 있었는데, 이는 불확실성에 대한 더 정직한 평가를 반영하는 트레이드오프(trade-off)이다. 이 연구는 기존의 강력한 AI 모델을 막대한 비용이 드는 재학습 과정 없이도 새로운 변화하는 화학적 환경에 적응시키는 것이 가능하다는 것을 입증했다.

이 연구는 인공지능을 신약 개발이라는 고도의 전문성이 요구되는 환경에 통합하기 위한 실질적인 경로를 제시한다. 변화하는 조건에 적응할 수 있는 통계적으로 엄격한 신뢰 척도를 생성하는 방법을 제공함으로써, 이 프레임워크는 이론적인 모델 성능과 투명성에 대한 규제 요구 사항 사이의 간극을 메우는 데 도움을 준다. 이는 과학자들이 새로운 분자에 대한 예측을 볼 때, 단순히 숫자 하나를 보는 것이 아니라 위험 요소를 포함한 잠재력에 대한 현실적인 평가를 확인할 수 있도록 보장한다. 이러한 '맹목적인 자신감'에서 '측정된 신뢰성'으로의 전환은 AI를 환자들에게 새로운 약을 전달하기 위한 수십억 달러 규모의 과업에서 신뢰할 수 있는 파트너로 만드는 데 있어 중요한 단계이다.

기술 요약: 레이블 시프트 상황에서의 분자 특성에 대한 컨포멀 예측(Conformal Prediction)

문제 정의
신약 개발 분야에서 인공지능(AI)의 적용은 분포 변화(distribution shift)라는 결정적인 신뢰성 병목 현상에 직면해 있습니다. AI 모델은 분자의 특성(예: 용해도, 효능, 독성)을 예측하여 가속화하는 역할을 하지만, 테스트 데이터가 훈련 분포와 달라질 경우 성능이 저하되는 경우가 빈번합니다. 구체적으로, 본 연구는 레이블 시프트(label shift) 문제를 다룹니다. 이는 소스(훈련)와 타겟(테스트) 도메인 사이에서 조건부 분포 $P(x|y)$ 는 불변하지만, 타겟 속성의 주변 분포(marginal distribution)인 $P(y)$ 가 변하는 시나리오를 의미합니다. 이러한 현상은 연구 우선순위가 기존 훈련 데이터에서 과소 대표되었던 특정 속성 값(예: 극단적인 효능 또는 낮은 독성)을 가진 분자를 발견하는 방향으로 이동할 때 신약 개발 과정에서 흔히 발생합니다.

표준 컨포멀 예측(conformal prediction)은 유한 샘플 유효성(finite-sample validity)을 갖는 통계적으로 엄격한 예측 구간을 생성하는 프레임워크이지만, 훈련 데이터와 테스트 데이터 사이의 교환 가능성(exchangeability) 가정을 전제로 합니다. 레이블 시프트 상황에서는 이 가정이 위배되어, 표준 방식은 공칭 신뢰 수준(nominal confidence level)보다 훨씬 낮은 커버리지 비율을 보이며 과도하게 확신에 찬(overconfident) 그러나 신뢰할 수 없는 예측을 생성하게 됩니다. 또한, 기존의 점 예측(point prediction)은 고위험 실험 설계에 제한적인 지침만을 제공하며, 기존의 분포 변화 해결책들은 비용이 많이 드는 모델 재학습을 요구하거나 연속적인 회귀 작업에서의 레이블 시프트보다는 주로 공변량 변화(covariate shift)에 초점을 맞추고 있습니다.

방법론
저자들은 모델을 재학습시키지 않고도 레이블 시프트가 발생하는 분자 특성 회귀 작업을 처리할 수 있도록 설계된 가중 컨포멀 예측(Weighted Conformal Prediction, WCP) 프레임워크를 제안합니다. 방법론은 다음 단계로 진행됩니다:

데이터 분할: 데이터 누수(data leakage)를 방지하기 위해 소스 데이터셋( $D_s$ )을 세 개의 서로 소인 하위 집합으로 분할합니다:
- 훈련 세트 ( $D_{train}$ ): 기본 예측 모델( $f$ )을 학습시키는 데 사용됩니다.
- 가중치 세트 ( $D_{weight}$ ): 타겟과 소스의 레이블 분포 비율을 나타내는 중요도 가중치(importance weights)를 추정하는 데 사용됩니다.
- 교정 세트 ( $D_{cal}$ ): 비정형성 점수(nonconformity scores)를 계산하는 데 사용됩니다.
- 타겟 데이터셋( $D_t$ )은 평가를 위해 변화가 없는 부분집합과 변화가 있는 부분집합으로 나뉩니다.
이산화 및 가중치 추정: 타겟 속성이 연속적이므로, $y$ 를 $K$ 개의 빈(bin)으로 이산화하여 의사 레이블(pseudo-labels)을 생성합니다. 중요도 가중치( $w$ )는 타겟 주변 레이벨 확률과 소스 주변 레이블 확률의 비율로 추정됩니다( $w(\tilde{y}) = q(\tilde{y})/p(\tilde{y})$ ). 저자들은 라벨링되지 않은 타겟 데이터로부터 타겟 분포 $q(\tilde{y})$ 를 추정하기 위해 세 가지 별도의 기술을 채택했습니다:
- 블랙박스 시프트 추정 (Black Box Shift Estimation, BBSE)
- 레이블 시프트 하의 정규화 학습 (Regularized Learning under Label Shifts, RLLS)
- 편향 수정 온도 스케일링(bias-corrected temperature scaling, BCTS)이 적용된 최대 가능도 추정 (Maximum Likelihood Estimation, MLE).
- MLE의 경우, 예측값을 모델 출력에 중심을 두고 표준 편차를 RMSE와 동일하게 하는 가우시안 분포로 모델링하고, 누적 분포 함수(CDF) 차이를 통해 빈 확률을 계산하는 확률적 접근 방식을 사용합니다.
가중 컨포멀 예측:
- 교정 세트에 대한 비정형성 점수( $s_i$ )는 일반적으로 절대 잔차( $|y_i - f(x_i)|$ )로 계산됩니다.
- 각 점수에는 실제 레이블의 빈(bin)을 기반으로 추정된 중요도 가중치( $\hat{w}_i$ )가 할당됩니다.
- 가중된 비율의 점수가 원하는 신뢰 수준( $1-\alpha$ )을 충족하도록 가중 분위수(weighted quantile)( $\hat{q}_w$ )를 계산합니다.
- 새로운 테스트 지점 $x_t$ 에 대한 최종 예측 구간은 $C(x_t) = [f(x_t) - \hat{q}_w, f(x_t) + \hat{q}_w]$ 로 구성됩니다.

주요 기여

레이블 시프트 적응: 본 논문은 공변량 변화에 비해 상대적으로 미개척 영역이었던, 레이블 시프트 상황에서의 연속적 회귀 작업을 위한 컨포멀 예측의 구체적인 적응 방식을 도입했습니다.
모델 불가지론적(Model-Agnostic) 프레임워크: 이 방법은 다양한 기본 모델(여기서는 대규모 언어 모델(LLM)로 입증됨)과 호환되며, 비용이 많이 드는 모델 재학습을 요구하지 않습니다. 오직 교정 단계를 재가중하는 것에만 의존합니다.
추정 기술의 통합: 본 프레임워크는 $q(\tilde{y})$ 를 추정하기 위해 BBSE, RLLS, MLE를 체계적으로 평가하고 통합하며, 이러한 기술들이 예측 모델의 출력만을 사용하여 어떻게 분포 변화를 교정할 수 있는지 보여줍니다.
통계적 엄밀성: 이 접근 방식은 속성 분포가 변하더라도 주변 커버리지 보장( $P(y_{test} \in C(x_{test})) \geq 1-\alpha$ )을 복구하여, 분포 변화 시에도 통계적으로 유효한 불확실성 정량화를 제공합니다.

실험 결과
프레임워크는 약 2억 개의 SMILES 문자열로 사전 학습되고 용해도(logS) 예측을 위해 미세 조정된 화학적 LLM(BART 아키텍처 기반)을 사용하여 AqSolDB 데이터셋(9,982개 화합물)에 대해 평가되었습니다.

베이스라인 실패: 표준 분할 컨포멀 예측(CP)은 레이블 시프트가 발생한 테스트 데이터에서 커버리지가 크게 감소하여, 평균 커버리지가 목표치보다 현저히 낮아지는 현상을 보였습니다.
성능 회복: 제안된 가중 컨포멀 예측(WCP) 방법들은 커버리지를 성공적으로 회복했습니다. WCP의 커버리지 분포는 표준 CP에 비해 오른쪽으로 이동하였으며, 이는 예측 구간이 실제 레이블을 더 자주 포함하고 있음을 나타냅니다.
추정기 비교: 가중치 추정 기술 중 MLE가 가장 높은 커버리지를 달s성했으며, RLLS와 BBSE가 그 뒤를 이었습니다. MLE는 이산화에 사용된 빈(bin)의 개수에 관계없이 견고함을 보여주었습니다.
트레이드오프(Trade-offs): MLE는 가장 우수한 커버리지를 제공했지만, BBSE나 RLLS에 비해 약간 더 넓은 예측 구간을 생성하는 결과를 낳았습니다.

의의 및 주장
저자들은 본 연구를 신약 개발 분야에서 AI의 규제 준수 및 실질적인 도입을 위한 중요한 단계로 규정합니다. 단일 값 정확도를 넘어 실행 가능하고 통계적으로 엄격한 신뢰 측도를 제공함으로써, 이 방법은 AI 시스템에 대해 "적절한 신뢰 구간"과 "불확실성 추정치"를 요구하는 FDA의 2025년 가이드라인을 충족합니다.

본 논문은 이 프레임워크가 수십억 달러 규모의 개발 파이프라인에서 AI 기반 예측의 신뢰성을 높인다고 주장합니다. 이는 속성 분포가 변하는 상황에서도 강력한 불확실성 정량화를 가능하게 하여, 새로운 화학적 공간을 탐색할 때 발생하는 고질적인 장애물을 직접적으로 해결합니다. 궁극적으로, 이 작업은 예측 모델링을 투명성과 신뢰성에 대한 요구에 부합하도록 함으로써, 고위험 과학적 의사결정 과정에서 모델이 신뢰할 수 있는 상태를 유지하도록 하는 일반화 가능하고 모델 불가지론적인 솔루션을 제공합니다.

유사한 논문