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

Resumo Técnico: PRISM-VFL

Problema
O Aprendizado Federado Vertical (VFL) permite o treinamento colaborativo de modelos entre instituições que detêm características complementares de pacientes sem compartilhar dados brutos. No entanto, os frameworks de VFL existentes enfrentam duas limitações críticas em ambientes clínicos: (1) eles carecem de suporte para a otimização simultânea de tarefas de predição heterogêneas (Aprendizado Multitarefa, ou MTL) sob garantias formais de privacidade, e (2) frequentemente assumem uma única parte ativa ou estruturas de tarefas uniformes. Na saúde real, as instituições podem deter conjuntos de características disjuntos (por exemplo, sinais vitais versus resultados laboratoriais) para os mesmos pacientes enquanto buscam objetivos clínicos distintos (por exemplo, predição de mortalidade versus fenotipagem). Além disso, integrar a Privacidade Diferencial (DP) ao MTL Federado introduz trocas complexas; a aplicação uniforme de ruído pode degradar a utilidade para tarefas de alta prioridade, enquanto a interação entre a sensibilidade específica da tarefa e as representações compartilhadas permanece subexplorada.

Metodologia
O estudo introduz o PRISM-VFL (Privacy-Resilient Institutional Split Multi-task VFL), um framework projetado para predição clínica multitarefa heterogênea.

• Arquitetura: O sistema segue um paradigma de aprendizado dividido (split learning). Cada site participante (simulado como três hospitais) processa suas características de séries temporais locais usando um codificador de Long Short-Term Memory (LSTM) de duas camadas, produzindo um embedding de 64 dimensões. Esses embeddings são transmitidos para um servidor central.
• Processamento no Lado do Servidor: O servidor concatena os embeddings de todos os sites em um vetor de 192 dimensões. Este vetor é processado por um módulo de Mistura de Especialistas com Múltiplas Portas (MMoE), composto por quatro subredes de especialistas compartilhadas e uma porta softmax aprendida para cada tarefa. Esta arquitetura permite que as tarefas roteiem através de diferentes especializações de especialistas, mitigando a transferência negativa.
• Mecanismo de Privacidade: O framework aplica o Gradiente Descendente Estocástico com Diferencial de Privacidade (DP-SGD) especificamente na camada de corte (cut layer). Os gradientes de embedding por amostra ($\partial L/\partial \langle k$) retornados do servidor para os clientes são clipados e perturbados com ruído Gaussiano antes de atualizar os codificadores locais. O estudo avalia tanto a alocação de ruído uniforme quanto a alocação de ruído estratificada por tarefa, onde os multiplicadores de ruído são ajustados com base na urgência clínica (por exemplo, menor ruído para Mortalidade Intra-hospitalar para preservar a utilidade).
• Conjunto de Dados e Simulação: O framework foi avaliado em uma versão verticalmente particionada do conjunto de dados MIMIC-III. As 14 variáveis clínicas foram distribuídas entre três sites simulados:
  • Site A: Sinais vitais (7 variáveis) com a tarefa de Mortalidade Intra-hospitalar (IHM).
  • Site B: Resultados laboratoriais (4 variáveis) com a tarefa de estratificação de risco de Descompensação.
  • Site C: Medições estáticas no momento da admissão (3 variáveis) com a tarefa de Fenotipagem (25 grupos de doenças).
  • O alinhamento de pacientes foi alcançado via simulação de Interseção de Conjuntos Privados (PSI), retendo apenas pacientes com estadias na UTI $\ge$ 48 horas.
• Avaliação: O estudo utilizou um conjunto abrangente de métricas, incluindo AUC-ROC, escore de Brier e um novo escore de Utilidade Multitarefa Ajustado por Privacidade-Comunicação (PCMU). Também introduziu um limite de inferência de rótulo multitarefa derivado teoricamente, $g(\sigma, \rho) = \Phi(C\sqrt{1+\rho}/\sigma)$, e o validou contra ataques de espaço de embedding empíricos (inferência de rótulo e inferência de pertencimento).

Principais Contribuições

1. Desenvolvimento de Framework: O PRISM-VFL é a primeira arquitetura de VFL que suporta simultaneamente conjuntos de tarefas heterogêneas através de sites verticalmente particionados com garantias explícitas de DP, utilizando MMoE para roteamento específico de tarefas.
2. Análise de Privacidade por Tarefa: O estudo fornece uma análise granular das trocas entre privacidade e utilidade, demonstrando que um único multiplicador de ruído uniforme falha em capturar a instabilidade específica da tarefa. Introduz um método para avaliar limiares de utilidade clínica ($\epsilon^*$) para tarefas individuais dentro de um cenário multitarefa.
3. Limites Teóricos e Empíricos: Um limite de referência teórico para inferência de rótulo multitarefa é derivado e validado contra um conjunto de ataques empíricos, mostrando que o roteamento MMoE mantém os riscos de inferência multitarefa dentro dos limites de tarefa única quando o acoplamento de gradiente é baixo.
4. Benchmark Reproduzível: Os autores estabelecem um benchmark de divisão vertical do MIMIC-III reproduzível que particiona tanto as características quanto os alvos de predição, facilitando pesquisas futuras em VFL multitarefa de múltiplos sites.

Resultados

• Desempenho: O PRISM alcançou uma média de AUC-ROC de 0,781 para IHM, 0,632 (com alta variância de semente cruzada) para Descompensação e 0,613 (macro-AUC) para Fenotipagem.
• Ganho Multitarefa: Ganhos significativos foram observados para Descompensação (+0,093 AUC) em comparação com baselines de tarefa única, atribuídos ao contexto fisiológico auxiliar fornecido por outros sites. IHM e Fenotipagem mostraram ganhos marginais (+0,003 e +0,002, respectivamente). Nenhuma transferência negativa sistemática foi detectada entre as configurações.
• Troca Privacidade-Utilidade:
  • IHM: Permaneceu clinicamente viável sob restrições de privacidade moderadas ($\epsilon \in (1, 2)$), mantendo um AUC acima do piso de utilidade de 0,75.
  • Descompensaçao e Fenotipagem: Falharam em atingir seus respectivos pisos de utilidade clínica (0,70 e 0,65) mesmo na ausência de restrições de DP ($\epsilon = \infty$), indicando que a arquitetura atual e a distribuição de dados são insuficientes para essas tarefas específicas nesta configuração.
  • Estabilidade: A estocasticidade da DP causou instabilidade assimétrica. A IHM mostrou pico de instabilidade em $\epsilon=1$ antes de se recuperar, enquanto a Descompensação e a Fenotipagem degradaram progressivamente.
• Eficácia da Privacidade: Sob privacidade forte ($\epsilon = 0,5$), ataques de inferência de rótulo foram suprimidos para um desempenho próximo ao aleatório (AUC 0,498). Ataques de inferência de pertencimento permaneceram ineficazes em todos os orçamentos, um resultado atribuído às propriedades arquiteturais da configuração de VFL, e não apenas à DP.
• Ablações: O estudo descobriu que o roteamento de porta aprendido no MMoE foi o principal motor dos benefícios multitarefa, não o aumento da capacidade. Curiosamente, dimensões de embedding menores (32 vs. 128) preservaram melhor a utilidade da IHM sob restrições de privacidade fortes.

Significância e Alegações
O artigo afirma que o PRISM-VFL demonstra a viabilidade de aprendizado multitarefa vertical federado, preservando a privacidade e sendo heterogêneo. No entanto, os autores enfatizam que a prontidão clínica é condicional à tarefa. Sob as condições testadas, a IHM é a única tarefa que atende aos limiares de utilidade clínica sob privacidade moderada, enquanto a Descompensação e a Fenotipagem requerem desenvolvimento arquitetônico adicional ou enriquecimento de dados antes da implementação.

O estudo posiciona o PRISM não apenas como um benchmark de desempenho, mas como um modelo metodológico para avaliar quais problemas de predição clínica são adequados para implantação multitarefa federada vertical. Ele destaca que estratégias de privacidade uniformes são insuficientes para tarefas heterogêneas e que sistemas futuros requerem mecanismos de calibração adaptáveis à tarefa. O trabalho estabelece uma linha de base para pesquisa em VFL de múltiplos sites, reconhecendo que os resultados em um conjunto de dados de instituição única (MIMIC-III) provavelmente representam uma estimativa otimista do desempenho no mundo real devido à falta de mudanças distributivas encontradas em implementações reais de múltiplos sites.