PRISM-VFL: A Differentially Private Vertical Federated Framework for Heterogeneous Multi-Task Clinical Prediction

기술 요약: PRISM-VFL

문제 정의
수직적 연합 학습(Vertical Federated Learning, VFL)은 상호 보완적인 환자 피처를 보유한 기관들이 원시 데이터를 공유하지 않고도 협력적인 모델 학습을 가능하게 한다. 그러나 기존의 VFL 프레임워크는 임상 환경에서 두 가지 결정적인 한계에 직면해 있다: (1) 공식적인 프라이버시 보장 하에 이질적인 예측 작업(Multi-Task Learning, MTL)의 동시 최적화를 지원하는 기능이 부족하며, (2) 단일 활성 참여자 또는 균일한 작업 구조를 가정하는 경우가 많다. 실제 의료 현장에서는 동일한 환자에 대해 서로 다른 피처 세트(예: 생체 징후 vs 검사 결과)를 보유하면서도 서로 다른 임상 목표(예: 병원 내 사망률 예측 vs 표현형 분류)를 추구할 수 있다. 또한, 차분 프라이버시(Differential Privacy, DP)를 연합 MTL에 통합하는 것은 복잡한 트레이드오프를 야기한다. 균일한 노이즈 적용은 우선순위가 높은 작업의 유용성을 저하시킬 수 있으며, 작업별 민감도와 공유된 표현(representation) 사이의 상호작용은 아직 충분히 탐구되지 않았다.

방법론
본 연구는 이질적인 다중 작업 임상 예측을 위해 설계된 PRISM-VFL(Privacy-Resilient Institutional Split Multi-task VFL) 프레임워크를 소개한다.

• 아키텍처: 시스템은 분할 학습(split learning) 패러다임을 따른다. 각 참여 사이트(세 곳의 병원으로 시뮬레이션됨)는 2계층 LSTM 인코더를 사용하여 로컬 시계열 피처를 처리하고, 64차원의 임베딩을 생성한다. 이 임베딩들은 중앙 서버로 전송된다.
• 서버 측 처리: 서버는 모든 사이트의 임베딩을 하나의 192차원 벡터로 결합한다. 이 벡터는 4개의 공유 전문가 서브 네트워크와 각 작업에 대한 학습된 소프트맥스 게이트(softmax gate)로 구성된 Multi-gate Mixture-of-Experts (MMoE) 모듈에 의해 처리된다. 이 아키텍처를 통해 각 작업은 서로 다른 전문가 전문화 경로를 거치게 되어, 부정적 전이(negative transfer)를 완화한다.
• 프라이버시 메커니즘: 프레임워크는 컷 레이어(cut layer)에 특화된 **차분 프라이비시 확률적 경사 하강법(DP-SGD)**을 적용한다. 서버에서 클라이언트로 반환되는 per-sample 임베딩 그래디언트($\partial L/\partial \langle k$)는 클리핑(clipping)된 후 가우시안 노이즈가 추가되어 로컬 인코더를 업데이트한다. 본 연구는 균일한 노이즈 할당과, 임상적 긴급도(예: 유용성 보존을 위해 병원 내 사망률에 대해 더 낮은 노이즈 적용)에 따라 노이즈 승수를 조정하는 **작업 계층적 노이즈 할당(task-stratified noise allocation)**을 비교 평가한다.
• 데이터셋 및 시뮬레이션: 프레임워크는 MIMIC-III 데이터셋의 수직 분할 버전을 사용하여 평가되었다. 14개의 임상 변수는 다음과 같이 세 곳의 시뮬레이션 사이트에 분배되었다:
  • 사이트 A: 생체 징후 (7개 변수), 병원 내 사망률(IHM) 작업 수행.
  • 사이트 B: 실험실 검사 결과 (4개 변수), 탈보상(Decompensation) 위험 계층화 작업 수행.
  • 사이트 C: 입원 시점 정적 측정값 (3개 변수), 표현형(Phenotyping, 25개 질병 그룹) 작업 수행.
  • 환자 정렬은 개인 정보 보호 집합 교집합(Private Set Intersection, PSI) 시뮬레이션을 통해 달성되었으며, 중환자실 체류 시간이 48시간 이상인 환자만을 유지하였다.
• 평가: 본 연구는 AUC-ROC, Brier score, 그리고 새로운 프라이버시-통신 조정 다중 작업 유용성(PCMU) 점수를 포함한 포괄적인 지표 세트를 활용하였다. 또한, 다중 작업 레이블 추론에 대한 이론적으로 도출된 경계값 $g(\sigma, \rho) = \Phi(C\sqrt{1+\rho}/\sigma)$를 도입하였으며, 이를 임베딩 공간 공격(레이블 추론 및 멤버십 추론)에 대한 경험적 검증을 통해 확인하였다.

주요 기여

1. 프레임워크 개발: PRISM-VFL은 명시적인 DP 보장과 함께 MMoE를 활용하여 작업별 라우팅을 지원하는, 수직적으로 분할된 사이트 간의 이질적인 작업 세트를 동시에 지원하는 최초의 VFL 아키텍처이다.
2. 작업별 프라이버시 분석: 본 연구는 프라이버시-유용성 트레이드오프에 대한 세밀한 분석을 제공하며, 단일 균일 노이즈 승수가 작업별 불안정성을 포착하는 데 실패함을 입증한다. 또한 다중 작업 환경 내 개별 작업에 대한 임상적 유용성 임계값($\epsilon^*$)을 평가하는 방법을 제시한다.
3. 이론적 및 경험적 경계: 다중 작업 레이블 추론에 대한 이론적 참조 경계가 도출되었으며, 이를 경험적 공격 스위트에 대해 검증하였다. 이는 그래디언트 결합(gradient coupling)이 낮을 때 MMoE 라우팅이 다중 작업 추론 위험을 단일 작업 수준으로 유지함을 보여준다.
4. 재현 가능한 벤치마크: 저자들은 피처와 예측 타겟을 모두 분할한 재현 가능한 MIM-III 수직 분할 벤치마크를 구축하여, 향후 다중 사이트 다중 작업 VFL 연구를 용이하게 한다.

결과

• 성능: PRISM은 IHM에 대해 0.781의 평균 AUC-ROC, 탈보상에 대해 0.632(높은 교차 시드 분산 동반), 표현형에 대해 0.613(macro-AUC)을 달성했다.
• 다중 작업 이득: 탈보상 작업에서 단일 작업 베이스라인 대비 상당한 이득(+0.093 AUC)이 관찰되었는데, 이는 다른 사이트로부터 제공된 부가적인 생리학적 맥락 덕분이다. IHM과 표현형은 각각 미미한 이득(+0.003 및 +0.002)을 보였다. 모든 구성에서 체계적인 부정적 전이는 감지되지 않았다.
• 프라이버시-유용성 트레이드오프:
  • IHM: 중간 정도의 프라이버시 제약 조건($\epsilon \in (1, 2)$) 하에서도 임상적으로 실행 가능한 수준을 유지하며, 0.75 유용성 하한선 이상의 AUC를 유지했다.
  • 탈보상 및 표현형: DP 제약이 없는 경우($\epsilon = \infty$)에도 각각의 임상 유용성 하한선(0.70 및 0.65)에 도달하지 못했다. 이는 현재의 아키텍처와 데이터 분포가 이 특정 작업들에에는 불충분함을 나타낸다.
  • 안정성: DP의 확률적 특성은 비대칭적 불안정성을 유발했다. IHM은 $\epsilon=1$에서 정점의 불안정성을 보인 후 회복된 반면, 탈보상과 표현형은 점진적으로 저하되었다.
• 프라이버시 효능: 강력한 프라이버시($\epsilon = 0.5$) 하에서 레이블 추론 공격은 무작위 성능에 가까운 수준(AUC 0.498)으로 억제되었다. 멤버십 추론 공격은 모든 예산 범위에서 효과가 없었으며, 이는 DP 자체보다는 VFL 구성의 아키텍처적 특성에 기인한 결과로 판단된다.
• 절제 연구(Ablations): 연구 결과, MMoE의 학습된 게이팅(learned gating)이 용량 증가가 아닌 다중 작업 이득의 주요 동력임을 발견했다. 흥//미롭게도, 더 낮은 임베딩 차원(32 vs 128)이 강력한 프라이버시 제약 하에서 IHM 유용성을 더 잘 보존했다.

의의 및 주장
본 논문은 PRISM-VFL이 프라이버시가 보장되는 이질적 수직 연합 다중 작업 학습의 실행 가능성을 입증한다고 주장한다. 그러나 저자들은 **임상적 준비도가 작업 조건부(task-conditional)**임을 강조한다. 테스트된 조건 하에서 IHM만이 중간 정도의 프라이버시 하에서 임상 유용성 임계치를 충족하는 유일한 작업이며, 탈보상과 표현형은 배포 전 추가적인 아키텍처 개발이나 데이터 보강이 필요하다.

본 연구는 PRISM을 단순한 성능 벤치마크가 아니라, 어떤 임상 예측 문제가 연합 다중 작업 배포에 적합한지를 평가하기 위한 방법론적 템플릿으로 위치시킨다. 연구는 균일한 프라이버시 전략이 이질적인 작업에는 불충 sufficient하지 않으며, 향후 시스템에는 작업 적응형 캘리브레이션 메커니즘이 필요함을 강조한다. 이 작업은 실제 다중 사이트 배치에서 발견되는 분포 변화(distributional shifts)가 결여된 단일 기관 데이터셋(MIMIC-III)을 사용했기에, 결과가 실제 환경에서의 성능에 대해 낙관적인 추정치일 수 있음을 인정하면서도 다중 사이트 VFL 연구를 위한 베이스라인을 확립한다.