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

Technische Samenvatting: PRISM-VFL

Probleemstelling
Vertical Federated Learning (VFL) maakt collaboratieve modeltraining mogelijk over instellingen die complementaire patiëntkenmerken bezitten zonder ruwe data te delen. Bestaande VFL-frameworks kampen echter met twee kritieke beperkingen in klinische settings: (1) ze missen ondersteuning voor de gelijktijdige optimalisatie van heterogene voorspellingsopdrachten (Multi-Task Learning, of MTL) onder formele privacygaranties, en (2) ze gaan vaak uit van één actieve partij of uniforme taakstructuren. In de echte gezondheidszorg kunnen instellingen disjuncte feature-sets bezitten (bijv. vitale functies versus laboratoriumresultaten) voor dezelfde patiënten, terwijl zij verschillende klinische doelen nastreven (bijv. mortaliteitvoorspelling versus fenotypering). Bovendien introduceert de integratie van Differential Privacy (DP) in Federated MTL complexe afruilmechanismen; uniforme ruistoepassing kan de bruikbaarheid voor hoog-prioritaire taken verminderen, terwijl de interactie tussen taakspecifieke sensitiviteit en gedeelde representaties nog onvoldoende is onderzocht.

Methodologie
De studie introduceert PRISM-VFL (Privacy-Resilient Institutional Split Multi-task VFL), een framework ontworig voor heterogene multi-task klinische voorspelling.

• Architectuur: Het systeem volgt een split learning-paradigma. Elke deelnemende locatie (gesimuleerd als drie ziekenhuizen) verwerkt zijn lokale tijdreekskenmerken met behulp van een twee-laagse Long Short-Term Memory (LSTM) encoder, wat een 64-dimensionale embedding produceert. Deze embeddings worden naar een centrale server verzonden.
• Server-zijde Verwerking: De server concateneert de embeddings van alle locaties tot een 192-dimensionale vector. Deze vector wordt verwerkt door een Multi-gate Mixture-of-Experts (MMoE) module, bestaande uit vier gedeelde expert sub-netwerken en een geleerde softmax gate voor elke taak. Deze architectuur stelt taken in staat om via verschillende expert-specialisaties te routeren, wat negatieve transfer mitigeert.
• Privacy Mechanisme: Het framework past Differentially Private Stochastic Gradient Descent (DP-SGD) specifiek toe op de 'cut layer'. Per-sample embedding gradiënten ($\partial L/\partial \langle k$) die van de server naar de clients worden geretourneerd, worden geclipt en geperturbeerd met Gaussische ruis voordat de lokale encoders worden bijgewerkt. De studie evalueert zowel uniforme ruistoewijzing als taak-gestratificeerde ruistoewijzing, waarbij ruismultipliers worden aangepast op basis van klinische urgentie (bijv. minder ruis voor In-Hospital Mortality om de bruikbaarheid te behouden).
• Dataset en Simulatie: Het framework werd geëvalueerd op een verticaal gepartitioneerde versie van de MIMIC-III dataset. De 14 klinische variabelen werden verdeeld over drie gesimuleerde locaties:
  • Locatie A: Vitale functies (7 variabelen) met de taak In-Hospital Mortality (IHM).
  • Locatie B: Laboratoriumresultaten (4 variabelen) met de taak Decompensatie risico-stratificatie.
  • Locatie C: Admissie-tijd statische metingen (3 variabelen) met de taak Fenotypering (25 ziektegroepen).
  • Patiënt-alignement werd bereikt via Private Set Intersection (PSI) simulatie, waarbij alleen patiënten met een IC-verblijf $\ge$ 48 uur werden behouden.
• Evaluatie: De studie maakte gebruik van een uitgebreid pakket aan metrieken, inclusocief AUC-ROC, Brier score, en een nieuwe Privacy-Communication-adjusted Multi-task Utility (PCMU) score. Er werd ook een theoretisch afgeleide multi-task label-inferentie grens geïntroduceerd, $g(\sigma, \rho) = \Phi(C\sqrt{1+\rho}/\sigma)$, en deze gevalideerd tegen empirische embedding-space aanvallen (label-inferentie en membership-inferentie).

Belangrijkste Bijdragen

1. Framework Ontwikkeling: PRISM-VFL is het eerste VFL-architectuur die simultaan heterogene taaksets ondersteunt over verticaal gepartitioneerde locaties met expliciete DP-garanties, gebruikmakend van MMoE voor taakspecifieke routing.
2. Per-Taak Privacy Analyse: De studie biedt een granulair onderzoek naar de privacy-utility afruil, waarbij wordt aangetoond dat een enkele uniforme ruismultiplier niet in staat is om taakspecifieke instabiliteit te vangen. Het introduceert een methode voor het evalueren van klinische bruikbaarheidsdrempels ($\epsilon^*$) voor individuele taken binnen een multi-task setting.
3. Theoretische en Empirische Grenzen: Een theoretische referentiegrens voor multi-task label-inferentie is afgeleid en gevalideerd tegen een empirische aanvalssuite, waarbij wordt aangetoond dat MMoE-routing de multi-task inferentie-risico's binnen de grenzen van single-task houdt wanneer de gradiënt-koppeling laag is.
4. Reproduceerbare Benchmark: De auteurs vestigen een reproduceerbare MIMIC-III verticale split benchmark die zowel features als predictiedoelen partitioneert, wat toekomstig onderzoek in multi-site multi-task VFL faciliteert.

Resultaten

• Prestaties: PRISM bereikte een gemiddelde AUC-ROC van 0,781 voor IHM, 0,632 (met hoge cross-seed variantie) voor Decompensatie, en 0,613 (macro-AUC) voor Fenotypering.
• Multi-Task Winst: Significante winsten werden waargenomen voor Decompensatie (+0,093 AUC) vergeleken met single-task baselines, toegeschreven aan de hulp van de fysiologische context van andere locaties. IHM en Fenotypering vertoonden marginale winsten (+0,003 en +0,002, respectievelijk). Er werd geen systematische negatieve transfer gedetecteerd over de configuraties heen.
• Privacy-Utility Afruil:
  • IHM: Bleef klinisch levensvatbaar onder matige privacy-restricties ($\epsilon \in (1, 2)$), waarbij een AUC boven de 0,75 bruikbaarheidsvloer werd behouden.
  • Decompensatie en Fenotypering: Bereikten hun respectievelijke klinische bruikbaarheidsvloeren (0,70 en 0,65) niet, zelfs in de afwezigheid van DP-restricties ($\epsilon = \infty$), wat aangeeft dat de huidige architectuur en datadistributie onvoldoende zijn voor deze specifieke taken in deze configuratie.
  • Stabiliteit: DP-stochastiek veroorzaakte asymmetrische instabiliteit. IHM vertoonde piekinstabiliteit bij $\epsilon=1$ alvorens te herstellen, terwijl Decompensatie en Fenotypering progressief verslechterden.
• Privacy Effectiviteit: Onder sterke privacy ($\epsilon = 0,5$) werd label-inferentie aanvallen onderdrukt tot bijna random prestaties (AUC 0,498). Membership-inferentie aanvallen bleven ineffectief over alle budgetten, een resultaat dat wordt toegeschreven aan de architecturale eigenschappen van de VFL-configuratie in plaats van enkel aan DP.
• Ablaties: De studie vond dat de geleerde gating in MMoE de primaire drijfveer was voor de multi-task voordelen, en niet een verhoogde capaciteit. Interessant genoeg behielden lagere embedding-dimensies (32 vs. 128) de IHM-bruikbaarheid beter onder sterke privacy-restricties.

Betekenis en Claims
Het artikel claimt dat PRISM-VFL de haalbaarheid demonstreert van privacy-beschermende, heterogene verticale federated multi-task learning. De auteurs benadrukken echter dat klinische gereedheid taak-conditioneel is. Onder de geteste condities is IHM de enige taak die aan klinische bruikbaarheidsdrempels voldoet onder matige privacy, terwijl Decompensatie en Fenotypering verdere architecturale ontwikkeling of dataverrijking vereisen voordat ze ingezet kunnen worden.

De studie positioneert PRISM niet louter als een prestatie-benchmark, maar als een methodologisch sjabloon voor het evalueren van welke klinische voorspellingsproblemen geschikt zijn voor federated multi-task deployment. Het benadrukt dat uniforme privacy-strategieën onvoldoende zijn voor heterogene taken en dat toekomstige systemen behoefte hebben aan taak-adaptieve kalibratiemechanismen. Het werk vestigt een baseline voor multi-site VFL-onderzoek, waarbij erkend wordt dat resultaten op een enkele instellingsdataset (MIMIC-III) waarschijnlijk een optimistische schatting van de real-world prestaties vertegenwoordigen vanwege het gebrek aan distributieve verschuivingen die men in echte multi-site deployments vindt.