KMM-CP: 변수 이동(Covariate Shift) 환경에서의 실용적인 공변 예측(Conformal Prediction) - 선택적 커널 평균 매칭(Kernel Mean Matching) 기반

KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching

불확실성 정량화는 과학 연구 및 의료와 같은 고위험 영역에서 머신러닝 모델을 배포하는 데 필수적입니다. 공변 예측(CP)은 교환 가능성(exchangeability)이라는 가정 하에서 유한 샘플 보장(finite-sample coverage guarantee)을 제공하지만, 실제로는 분포 이동(distribution shift)으로 인해 이 가정이 종종 위반됩니다. 변수 이동(covariate shift)의 경우, 유효성을 회복하려면 중요도 가중치(importance weighting)가 필요하지만, 정확한 밀도 비율 추정(density-ratio estimation)은 훈련 및 테스트 분포 간의 제한적인 지지 영역 중첩(support overlap)을 보이는 경우 불안정해질 수 있습니다. 본 연구에서는 변수 이동을 보정하기 위한 커널 평균 매칭(KMM) 기반의 공변 예측 프레임워크인 KMM-CP를 제안합니다. KMM이 명시적인 가중치 제약 조건 하에서 RKHS 모멘트 불일치(RKHS moment discrepancy)를 최소화함으로써 공변 예측 오차를 지배하는 편향-분산(bias-variance) 구성 요소를 직접 제어하며, 온화한 조건 하에서 점근적 보장(asymptotic coverage guarantee)을 제공함을 보여줍니다. 또한, 신뢰할 수 있는 지지 영역 중첩을 식별하고 공변 보정을 이 부분 집합으로 제한하는 선택적 확장(selective extension)을 도입하여, 중첩이 낮은 경우의 안정성을 더욱 향상시킵니다. 현실적인 분포 이동이 있는 분자 특성 예측 벤치마크 실험에서 KMM-CP는 기존 접근 방식과 비교하여 50% 이상의 커버리지 격차(coverage gap)를 줄이는 것으로 나타났습니다. 코드: https://github.com/siddharthal/KMM-CP

Uncertainty quantification is essential for deploying machine learning models in high-stakes domains such as scientific discovery and healthcare. Conformal Prediction (CP) provides finite-sample coverage guarantees under exchangeability, an assumption often violated in practice due to distribution shift. Under covariate shift, restoring validity requires importance weighting, yet accurate density-ratio estimation becomes unstable when training and test distributions exhibit limited support overlap. We propose KMM-CP, a conformal prediction framework based on Kernel Mean Matching (KMM) for covariate-shift correction. We show that KMM directly controls the bias-variance components governing conformal coverage error by minimizing RKHS moment discrepancy under explicit weight constraints, and establish asymptotic coverage guarantees under mild conditions. We then introduce a selective extension that identifies regions of reliable support overlap and restricts conformal correction to this subset, further improving stability in low-overlap regimes. Experiments on molecular property prediction benchmarks with realistic distribution shifts show that KMM-CP reduces coverage gap by over 50% compared to existing approaches. The code is available at https://github.com/siddharthal/KMM-CP.