Variational Consensus Monte Carlo for Bayesian Mixture

Resumo Técnico: Variational Consensus Monte Carlo para Modelos de Mistura Bayesianos

Definição do Problema
O artigo aborda o desafio de realizar inferência bayesiana para modelos de mistura em um cenário de Aprendizado Federado (FL), especificamente o cenário "cross-silo", onde os dados são naturalmente particionados entre múltiplos nós (ex: centros médicos) e não podem ser agrupados devido a restrições de privacidade ou regulatórias. Embora o Consensus Monte Carlo (CMC) ofereça uma estrutura para inferência bayesiana distribuída através da agregação de amostras posteriores locais, os métodos existentes enfrentam limitações significativas quando aplicados ao agrupamento (clustering) não supervisionado. Especificamente, a abordagem de Variational Consensus Monte Carlo (VCMC) proposta por Rabinovich, Angelino e Jordan (2015) assume que o número de clusters e os principais parâmetros da mistura são conhecidos e que cada cluster está presente em todos os fragmentos (shards) de dados. Essas suposições raramente se sustentam na prática, particularmente em aplicações de dados de saúde, onde a prevalência de clusters varia entre os locais e o verdadeiro número de subgrupos latentes é desconhecido. Além disso, a agregação padrão do CMC (ex: médias ponderadas) não garante uma aproximação ótima da posterior global, e os métodos de clustering federados existentes frequentemente dependem de aproximações de campo médio (mean-field) que exigem priors conjugados ou falham em lidar com configurações não conjugadas sem acessar o conjunto de dados completo.

Metodologia
Os autores propõem um pipeline abrangente que estende o VCMC para modelos de mistura superajustados (over-fitted) de distribuições categóricas. A estrutura opera em duas fases principais:

1. A Etapa de Aplicação (Inferência Local):

   • Um algoritmo MCMC (especificamente um amostrador de Gibbs) é executado independentemente dentro de cada fragmento de dados.
   • Para lidar com o número desconhecido de clusters, o modelo é "superajustado" definindo o número de componentes da mistura ($C'$) maior do que o número antecipado de clusters reais ($C$).
   • Um prior fracionado é utilizado, onde o prior global é dividido como $p(\theta)^{1/S}$ entre os fragmentos. Isso garante que o produto das posteriores locais aproxime a posterior global.
   • Os autores implementam um amostrador de Gibbs paralelizado em C++17 para acelerar a inferência local.

2. A Etapa de Agregação (Inferência Global):

   • Correspondência de Clusters (Cluster Matching): Como os modelos de mistura são invariantes a permutações de rótulos, uma etapa crítica é a correspondência de clusters entre os fragmentos. O artigo introduz e avalia três algoritmos:
     • Hungarian Matching: Impõe uma correspondência um-para-um, mas falha se nem todos os clusters aparecerem em cada fragmento.
     • Minimum Divergence Matching: Minimiza a divergência de Kullback-Leibler (KL) entre as marginais posteriores locais.
     • Ball Matching: Um algoritmo inovador que agrupa clusters que estão "próximos" (dentro de um raio $r$) no espaço de parâmetros, baseado nas taxas de convergência de misturas superajustadas. Este método permite fundir clusters dentro de um fragmento se eles forem indistinguíveis, mitigando a superestimação local da contagem de clusters.
   • Agregação Variacional: A agregação é formulada como um problema de inferência variacional para minimizar a divergência KL entre a aproximada e a verdadeira posterior global. A função de agregação é restrita a uma soma ponderada dos parâmetros locais ($\sum \upsilon_s \theta_s$).
   • Otimização: Os pesos de agregação ideais são encontrados usando gradiente descendente estocástico projetado. O artigo detalha quatro estratégias para essa otimização para acomodar diferentes restrições de FL:
     1. Compartilhamento dos dados brutos com o nó global (menos restritivo).
     2. Gradiente descendente federado (compartilhando apenas os gradientes).
     3. Comunicação limitada (execuções paralelas com poucas iterações e múltiplos pontos de partida).
     4. Compartilhamento de uma matriz de contingência (estatísticas de resumo) com o nó global.

Principais Contribuições
O artigo apresenta quatro contribuições metodológicas primárias:

1. Extensão para Modelos Superajustados: Estende o VCMC para inferir o número de clusters e todos os parâmetros do modelo simultaneamente sem exigir conjugação, relaxando a suposição de que o número de clusters é conhecido.
2. Novos Algoritmos de Correspondência de Clusters: Introduz estratégias de correspondência (especificamente Minimum Divergence e Ball matching) adequadas para configurações cross-silo onde clusters podem estar ausentes em alguns fragmentos, um cenário onde o Hungarian matching padrão falha.
3. Estratégias de Inferência Flexíveis: Propõe múltiplas estratégias de agregação adaptadas a diferentes restrições federadas (ex: largura de banda de comunicação limitada ou requisitos rigorosos de privacidade).
4. Diretrizes Práticas: Fornece diretrizes empíricas para a seleção de algoritmos de correspondência e estratégias de inferência com base nas características dos dados (ex: separabilidade de clusters e heterogeneidade dos fragmentos).

Resultados
O framework foi avaliado através de um estudo de simulação abrangente e uma aplicação a dados reais de registros eletrônicos de saúde (EHR).

• Estudo de Simulação:

  • Desempenho de Correspondência: Em cenários heterogêneos onde nem todos os clusters aparecem em cada fragmento, o algoritmo de Hungarian falhou em recuperar a correspondência correta (0% de sucesso em cenários não aninhados). O Ball matching e o Minimum Divergence matching tiveram um desempenho significativamente superior. O Ball matching foi particularmente eficaz em mitigar a superestimação do número de clusters causada pelo ruído local do MCMC.
  • Qualidade de Clustering: O VCMC alcançou escores de Índice Rand Ajustado (ARI) comparáveis ou ligeiramente superiores aos de alternativas de estado da arte (FedMerDel e Federated K-modes) e ao "padrão ouro" do MCMC de dados completos. Notavelmente, o VCMC superou o MCMC de dados completos na recuperação de pequenos clusters quando a distribuição dos dados entre os fragmentos era heterogênea.
  • Estimativa de Parâmetros: O VCMC forneceu baixo viés relativo para os pesos dos clusters. Em cenários com clusters mal separáveis e distribuições de fragmentos heterogêneas, o VCMC estimou os pesos de pequenos clusters com mais precisão do que o MCMC de dados completos, provavelmente porque esses pequenos clusters eram grandes o suficiente dentro de fragmentos específicos para serem identificados com menor incerteza.
  • Comparação com Alternativas: Embora o FedMerDel fosse mais rápido, ele dependia de aproximações de campo médio que subestimavam as variâncias posteriores. O Federated K-modes mostrou alto viés e baixo desempenho na estimativa de estruturas de clusters para dados categóricos.

• Aplicação no Mundo Real (Coorte THIN):

  • O método foi aplicado a 289.821 pacientes geriátricos do banco de dados britânico THIN para identificar padrões de multimorbidade.
  • O algoritmo identificou 27 clusters. Ele caracterizou com sucesso grupos clinicamente significativos, como clusters dominados por acidente vascular cerebral (AVC), demência ou tipos específicos de câncer.
  • Crucialmente, identificou pequenos mas distintos clusters (ex: um com apenas 31 indivíduos caracterizado por pancreatite, artrite reumatoide e disfunção erétil) que se alinhavam com correlações médicas documentadas, demonstrando a capacidade do método de detectar subgrupos raros.

Significância e Alegações
Os autores afirmam que seu framework fornece uma solução robusta para o agrupamento não supervisionado em ambientes sensíveis à privacidade. A significância reside na capacidade de:

• Inferir o número de clusters e os parâmetros do modelo sem a união de dados (data pooling).
• Recuperar pequenos clusters raros com maior precisão do que o MCMC padrão de dados agrupados, quando os fragmentos de dados locais contêm sinal suficiente para esses subgrupos específicos.
• Operar sob restrições rigorosas de privacidade (FL cross-silo) sem exigir priors conjugados ou acesso ao conjunto de dados completo.

O artigo observa modestamente que, embora o framework seja teoricamente extensível a outras distribuições de mistura, a validação atual é específica para distribuições categóricas. Também destaca que a escolha do algoritmo de correspondência de clusters depende do objetivo específico: o Ball matching é preferível para estimar o número de clusters, enquanto o Minimum Divergence matching pode ser melhor se a prioridade for evitar a fusão de clusters distintos. Os autores concluem que, embora o VCMC possa ser computacionalmente intensivo em comparação com métodos variacionais como o FedMerDel, ele oferece precisão superior na estimativa de parâmetros e na recuperação de clusters em configurações federadas complexas e heterogêneas.