Robust Conditional Conformal Prediction via Branched Normalizing Flow

기술적 요약: 분기형 정규화 흐름을 통한 강건한 조건부 컨포멀 예측

1. 문제 제기

컨포멀 예측 (CP) 은 보정 데이터와 테스트 데이터가 독립적이고 동일하게 분포되어 있다는 (i.i.d.) 가정 하에 한계 커버리지 보장을 제공하는 예측 집합을 제공합니다. 그러나 분포 이동 (테스트 분포 $Q_{XY}$가 보정 분포 $P_{XY}$와 다른 경우) 이 존재할 때, 기존 방법들은 일반적으로 한계 컨포멀 점수 분포만을 정렬합니다. 이는 한계 커버리지는 유지하지만 개별 테스트 입력에서의 조건부 커버리지 오차를 제어하지는 못합니다. 그 결과, 조건부 점수 분포가 불일치하는 특정 영역에서 CP 집합이 신뢰할 수 없게 되어, 특정 입력에 대해 커버리지 부족이나 과도한 오버커버가 발생할 수 있습니다. 핵심적인 과제는 추론 중 실제 테스트 레이블에 대한 지식이 필요 없이 결합 분포 이동 하에서 강건한 조건부 커버리지 보장 ($Pr(Y \in C(X) | X=x) \ge 1-\alpha$) 을 확보하는 것입니다.

2. 방법론

제안된 해결책은 분포 이동과 조건부 커버리지 저하를 연결하는 이론적 프레임워크와 이를 완화하기 위한 새로운 아키텍처를 도입합니다.

이론적 프레임워크:

• 조건부 커버리지 갭 (CCG): 저자들은 $Q_{XY}$ 하의 테스트 입력에 대한 조건부 커버리지와 $P_{XY}$ 하의 보장 간의 불일치를 측정하기 위해 CCG 를 정의합니다.
• 통합 커버리지 갭 (ICG): 테스트 특징 분포 $Q_X$에 대한 CCG 의 기댓값으로 정의되며, 전체적인 강건성 지표 역할을 합니다.
• 와세르슈타인 상한: 논문은 결합 보정 및 테스트 분포 ($P_{XY}$ 및 $Q_{XY}$) 간의 와세르슈타인 거리 ($W$) 를 사용하여 ICG 에 대한 상한을 유도합니다. 구체적으로, $ICG(P, Q) \le \sqrt{2\kappa L} (\eta \cdot W(P_{XY}, Q_{XY}) + 1/4)$입니다. 이 상한은 분포 간의 와세르슈타인 거리를 최소화하는 것이 이론적으로 강건한 조건부 커버리지를 촉진할 수 있음을 시사합니다.

알고리즘: 분기형 정규화 흐름 (BNF)
와세르슈타인 상한을 실현하기 위해 저자들은 테스트 분포 $Q_{XY}$를 보정 분포 $P_{XY}$와 일치하도록 변환하도록 설계된 가역 매핑 $f_\theta(x, y) = (f_{\theta_X}(x), f_{\theta_Y}(y)) 인 **분기형 정규화 흐름 (BNF)**을 제안합니다.

• 아키텍처: BNF 는 두 개의 독립적인 분기를 사용합니다. $X$-분기 ($f_{\theta_X}$) 는 테스트 입력을 정규화하고, $Y$-분기 ($f_{\theta_Y}$) 는 레이블을 변환합니다. 중요한 점은 분기들이 매개변수를 공유하지 않으므로, $X$의 변환이 관찰되지 않은 테스트 레이블 $Y$에 의존하지 않도록 보장한다는 것입니다.
• 추론 과정:
  1. 테스트 입력 $X_{n+1}$을 정규화된 입력 $\bar{X}_{n+1} = f_{\theta_X}(X_{n+1})$로 변환합니다.
  2. 보정 분포 $P_{XY}$를 사용하여 (예: 컨포멀라이즈드 양적 회귀를 통해) 조건부 예측 집합 $CC(\bar{X}_{n+1})$을 구성합니다.
  3. $f_{\theta_Y}^{-1}$을 사용하여 집합을 테스트 공간으로 역변환하여 최종 강건한 예측 집합 $C_{BNF}(X_{n+1})$을 얻습니다.
• 증강된 BNF: 단순한 역 매핑에 필요한 단변수 단조 변환으로는 불충분한 복잡한 분포를 처리하기 위해 저자들은 증강된 BNF를 도입합니다. 이 변형은 $Y$-분기 변환 ($f_{\theta_Y}^{aug}(y; \epsilon)$) 에 확률적 노이즈 성분 $\epsilon \sim \mathcal{N}(0,1)$을 추가합니다. 이는 표현력을 높이지만, 변환이 더 이상 엄격하게 단조롭지 않으므로 예측 집합을 구성하기 위해 추론 중 근 찾기 문제 (또는 그리드 검색) 를 해결해야 합니다.
• 다중 소스 도메인 일반화 (MSDG): 테스트 분포가 $K$개의 소스 분포의 혼합인 시나리오에서, 이 방법은 균일한 혼합 가정을 전제로 보정 세트와 변환된 소스 세트 간의 평균 와세르슈타인 거리를 최소화합니다.

3. 주요 기여

1. 강건성의 정량화: 조건부 커버리지 갭 (CCG) 과 통합 커버리지 갭 (ICG) 의 정의는 분포 이동 하에서 조건부 커버리지 강건성을 평가하기 위한 이론적 지표를 제공합니다.
2. 이론적 상한: 결합 분포 간의 와세르슈타인 거리에 대한 ICG 의 상한 유도로서, 수송 최소화와 조건부 커버리지 보존 간의 이론적 연결을 확립합니다.
3. 분기형 정규화 흐름 (BNF): 테스트 레이블을 요구하지 않고 테스트 입력을 보정과 유사한 분포로 변환할 수 있는 이분기 가역 아키텍처를 도입하여, 역변환 시 조건부 보장을 유지합니다.
4. 증강된 BNF: 복잡한 분포에 대한 적합 능력을 향상시키기 위해 확률적 증강을 사용한 확장이며, 단조성 손실에도 불구하고 조건부 보장을 회복하는 방법을 결합합니다.

4. 실험 결과

이 방법은 다중 소스 도메인 일반화 (MSDG) 설정 하에서 합성 이동과 실제 세계의 도전 과제 (판매 예측, 교통 예측, 의료 의사결정, 전염병 모델링) 를 다루는 9 개의 데이터셋에서 평가되었습니다.

• 지표: 성능은 한계 커버리지, 최악 슬라이스 커버리지 (WSC), 그리고 제안된 최악 슬라이스 커버리지 갭 (WSCG) 을 사용하여 측정되었으며, 이는 커버리지 부족과 과잉 커버리지 모두를 패널티로 부과합니다.
• 결과:
  • 강건성: 증강된 BNF 와 컨포멀라이즈드 양적 회귀 (CQR) 의 조합은 모든 데이터셋과 신뢰 수준에서 일관되게 가장 낮은 WSCG 를 달성하여, Split CP, Importance-Weighted CP, Worst-Case CP, Wasserstein-Regularized CP 를 포함한 베이스라인을 능가했습니다.
  • 커버리지: 이 방법은 한계 커버리지와 WSC 를 목표 신뢰 수준 ($1-\alpha = 0.9$) 에 가깝게 유지하면서 조건부 슬라이스에서 목표와의 편차를 크게 줄였습니다.
  • 효율성: 표준 증강된 BNF 는 소스 불확실성을 고려하여 더 큰 예측 집합을 생성했지만, "증강-조건부" 변형 (추론 중 알려진 소스에 조건을 부여함) 은 커버리지를 유지하면서 집합 크기를 크게 줄였습니다.
  • 일반화: 이 방법은 다양한 샘플 크기, 소스 도메인 수 ($K$), 그리고 다른 신뢰 수준에 걸쳐 강건성을 입증했습니다.

5. 중요성과 주장

이 논문은 기존 한계 중심 방법들이 largely 다루지 않았던 분포 이동 하의 컨포멀 예측 조건부 커버리지 저하를 해결하는 데 그 주요 의미가 있다고 주장합니다. 결합 분포 간의 와세르슈타인 거리를 통해 조건부 무효성을 상한으로 묶고 이를 가역 분기 흐름을 통해 구현함으로써, 이 연구는 데이터 분포가 이동하더라도 개별 테스트 입력에 대한 불확실성 정량이 신뢰할 수 있도록 보장하는 이론적으로 근거 있고 경험적으로 효과적인 메커니즘을 제공합니다. 저자들은 이를 분포 이동이 불가피하고 조건부 신뢰성이 최우선인 핵심 영역 (예: 의료, 교통, 전염병 모델링) 에서 더 안전한 AI 응용을 향한 한 걸음으로 위치시킵니다.

논문은 증강된 BNF 의 근 찾기에 따른 계산 비용, 예측 집합의 확률적 성격, 그리고 다변량 입력을 요구하는 아키텍처 제약과 같은 한계를 인정합니다. 또한 보정 데이터가 흐름 훈련에 재사용된다는 점 (잠재적인 튜닝 편향) 을 지적하지만, 제거 연구 (ablation studies) 는 실제적으로 이 영향이 미미함을 시사한다고 덧붙입니다.