Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

Resumo Técnico: Imputação Probabilística de Recursos e Agregação Federada Multimodal Consciente da Incerteza

1. O Problema

O aprendizado federado (FL) multimodal é crucial para aplicações de saúde, permitindo o treinamento colaborativo de modelos sem violar a privacidade dos dados (ex: HIPAA, GDPR). No entanto, uma barreira fundamental é a heterogeneidade de modalidades:

• Cenário Real: Centros médicos grandes possuem dados completos (imagens de raio-X + laudos textuais), enquanto clínicas menores ou hospitais rurais frequentemente possuem apenas um subconjunto (apenas imagens).
• Limitação das Abordagens Atuais: Métodos existentes utilizam redes de imputação de recursos determinísticas para sintetizar a modalidade ausente (ex: gerar texto a partir da imagem).
• Risco Crítico: Essas abordagens determinísticas produzem estimativas pontuais (valores únicos) sem medidas de confiabilidade. Em ambientes médicos de alta segurança, isso é perigoso, pois o modelo pode "alucinar" características com alta confiança mesmo quando a entrada é ambígua, levando a diagnósticos errôneos. Além disso, a agregação federada padrão (FedAvg) trata todas as atualizações de clientes como igualmente confiáveis, permitindo que dados imputados de baixa qualidade degradem o modelo global.

2. Metodologia Proposta

Os autores propõem o P-FIN (Probabilistic Feature Imputation Network) e uma estratégia de agregação global chamada Fed-UQ-Avg.

A. Rede de Imputação Probabilística (P-FIN)

• Abordagem: Em vez de prever um vetor fixo para a modalidade ausente, o P-FIN modela a distribuição condicional $p(z_T | z_I)$ como uma Gaussiana heterocedástica.
• Arquitetura:
  • Utiliza um codificador de imagem (ResNet-50) e um codificador de texto (BERT).
  • O núcleo da imputação é um Transformer (2 camadas, 4 cabeças) que recebe características da imagem e um token de consulta aprendível.
  • Saída Dupla: O modelo prevê tanto a média ($\mu$) quanto a variância ($\sigma^2$) para cada dimensão do recurso sintetizado. A variância representa a incerteza por dimensão.
• Treinamento: Utiliza a função de perda $\beta$-NLL (Negative Log-Likelihood). Isso é crucial para evitar que o modelo minimize a perda simplesmente prevendo uma incerteza infinita (explosão de variância), forçando-o a aprender representações significativas enquanto calibra a incerteza.

B. Fusão Local Consciente da Incerteza

• Portão de Incerteza (Uncertainty Gating): Antes da fusão, a variância prevista ($\sigma^2$) é usada para calcular um sinal de portão ($g = \text{sigmoid}(-\log \sigma^2)$).
  • Se a incerteza for alta, o portão se aproxima de zero, suprimindo as características imputadas não confiáveis.
  • Se a incerteza for baixa, o portão permanece próximo de um, permitindo a fusão.
• Fusão por Atenção: As características de imagem e texto (já filtradas pelo portão) são combinadas usando mecanismos de atenção cruzada bidirecional, permitindo que o modelo pondere dinamicamente a contribuição de cada modalidade.

C. Agregação Global: Fed-UQ-Avg

• Problema: O FedAvg padrão pondera clientes apenas pelo tamanho do conjunto de dados, ignorando a qualidade da imputação.
• Solução: O Fed-UQ-Avg introduz um esquema de ponderação de dois componentes:
  1. Peso de Dados: Proporcional ao tamanho do conjunto de dados local ($n_k$).
  2. Peso de Confiança: Baseado na incerteza média de imputação do cliente ($\bar{\sigma}^2_k$). Clientes com baixa incerteza (alta confiança) recebem pesos maiores.
• Fórmula: O peso final combina ambos, com um parâmetro de equilíbrio $\alpha$ (definido como 0.6), priorizando a confiabilidade da imputação.

3. Principais Contribuições

1. P-FIN: Primeira rede de imputação probabilística projetada especificamente para cenários federados, fornecendo estimativas de incerteza calibradas junto com os recursos sintetizados.
2. Mecanismo de Dupla Nível de Incerteza:
   • Local: Supressão de ruído na fusão de características via portão de incerteza.
   • Global: Agregação ponderada por qualidade via Fed-UQ-Avg, protegendo o modelo global de clientes com imputação pobre.
3. Validação em Cenários Críticos: Demonstração de que a quantificação de incerteza é essencial para segurança em aplicações médicas, alinhando-se com diretrizes regulatórias (ex: FDA) que exigem transparência sobre a confiança do modelo.

4. Resultados Experimentais

Os experimentos foram realizados em cenários federados simulados utilizando três conjuntos de dados de raios-X torácicos: CheXpert, NIH Open-I e PadChest.

• Configuração: 10 clientes com diferentes proporções de dados unimodais (apenas imagem) vs. multimodais (imagem + texto), variando de 80% a 40% de clientes unimodais.
• Desempenho:
  • O método proposto (P-FIN + Fed-UQ-Avg) superou consistentemente todas as linhas de base determinísticas (incluindo Zero-filling, Uniform-filling e FIN + FedAvg).
  • Melhoria Significativa: Na configuração mais desafiadora (80% de clientes sem texto), houve um ganho absoluto de 5,36% na AUC (Área Sob a Curva) em comparação com a melhor abordagem determinística.
  • Contribuição da Agregação: A adição do Fed-UQ-Avg ao P-FIN (sem ele) trouxe ganhos adicionais de ~2,4% a 2,8% em todas as configurações, provando que a ponderação baseada em incerteza é vital.
• Calibração: A análise mostrou que as incertezas previstas estão bem calibradas (ECE = 0,0422) e correlacionadas positivamente com o erro de reconstrução real (alta incerteza = baixo desempenho de imputação).

5. Significado e Impacto

Este trabalho estabelece um novo paradigma para o aprendizado federado multimodal em saúde: "aprender a desconfiar".

• Segurança Clínica: Ao quantificar a incerteza, o sistema permite que radiologistas confiem nas previsões quando a incerteza é baixa e investiguem manualmente quando é alta, mitigando riscos de diagnósticos errôneos por alucinações de IA.
• Robustez: A abordagem torna os sistemas de IA mais resilientes à disparidade de recursos entre instituições, permitindo que hospitais com dados limitados participem de colaborações globais sem degradar o modelo coletivo.
• Futuro: O trabalho abre caminho para a integração de incerteza epistêmica e síntese bidirecional de modalidades em cenários de distribuição fora do padrão (out-of-distribution).

Em resumo, o P-FIN transforma a imputação de dados de um processo cego e arriscado em uma ferramenta transparente e segura, essencial para a adoção clínica de IA federada.