📊 statistics

Variational Consensus Monte Carlo for Bayesian Mixture

This paper introduces a comprehensive Variational Consensus Monte Carlo pipeline for Bayesian mixture models in federated learning that extends existing methods to infer the number of clusters and all parameters without conjugacy, employs novel cluster-matching algorithms for cross-silo settings, and demonstrates superior accuracy in recovering small clusters compared to pooled data approaches using electronic health record data.

Original authors: Julie Fendler, Francesca L. Crowe, Tom Marshall, Sylvia Richardson, Paul D. W. Kirk

Published 2026-06-19

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Julie Fendler, Francesca L. Crowe, Tom Marshall, Sylvia Richardson, Paul D. W. Kirk

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive jigsaw puzzle, but the pieces are scattered across 30 different locked rooms. You cannot move the pieces out of the rooms, and you cannot show the pieces to anyone outside your own room. This is the challenge of Federated Learning: you have data split across many locations (like hospitals or clinics), but privacy laws or technical limits prevent you from pooling all the data into one giant computer.

This paper presents a clever new way to solve that puzzle without ever moving the pieces. Here is how they did it, explained simply.

The Problem: The "Locked Room" Puzzle

Usually, to find patterns in data (like grouping patients with similar diseases), you need to see everyone's data at once. But in healthcare, patient data is sensitive. You can't just email a spreadsheet of patient records from a hospital in London to a server in Birmingham.

The authors wanted to use a statistical tool called a Bayesian Mixture Model. Think of this as a machine that looks at a crowd of people and says, "Okay, these 50 people look like a 'Group A' (maybe they have heart issues), and these 20 look like 'Group B' (maybe they have diabetes)."

The problem is that if you run this machine in each locked room separately, you might get different results. Room 1 might find "Group A" and "Group B." Room 2 might find "Group A" and "Group C." How do you combine these local guesses into one big, accurate picture of the whole population without ever seeing the raw data?

The Solution: The "Consensus" Strategy

The authors use a method called Consensus Monte Carlo (CMC). Imagine a team of detectives, each working in their own locked room.

The Apply Step: Each detective runs their own investigation (a computer algorithm called MCMC) on their local data. They produce a list of "suspects" (clusters) and a description of what those suspects look like.
The Aggregate Step: The detectives send only their descriptions (not the raw data) to a central coordinator. The coordinator's job is to figure out: "Is the 'Group A' in Room 1 the same as the 'Group A' in Room 2?" and then blend these descriptions together to form the final truth.

The New Tricks: What This Paper Adds

Previous versions of this "Consensus" method had some big flaws. They assumed everyone knew exactly how many groups existed and that every group was present in every room. In the real world, that's rarely true. Some rooms might have a rare disease that no one else has.

This paper introduces four major improvements:

1. The "Over-Fitted" Safety Net
Instead of guessing the exact number of groups beforehand (e.g., "There are exactly 5 groups"), the authors tell the computer to guess too many groups (e.g., "Let's look for 20 groups").

The Analogy: Imagine you are sorting a pile of mixed nuts. Instead of guessing there are exactly 3 types, you set out 20 bowls. The computer will fill the bowls it needs and leave the empty ones empty. This allows the system to automatically figure out how many groups actually exist without human guessing.

2. The "Cluster Matching" Algorithms
This is the hardest part. If Room 1 finds a "Heart Group" and Room 2 finds a "Heart Group," how does the coordinator know they are the same?

The Old Way (Hungarian Algorithm): This tried to force a perfect 1-to-1 match. If Room 1 had 5 groups and Room 2 had 5 groups, it matched them all. But if Room 2 had a rare group that Room 1 didn't have, the whole system broke.
The New Ways: The authors invented two new matching strategies:
- Minimum Divergence: This tries to match groups that look statistically similar, minimizing the "distance" between their descriptions.
- Ball Matching: This is like rolling a ball. If a group in Room 1 is "close enough" (within a certain radius) to a group in Room 2, they get merged into the same cluster. This is particularly good at handling rare groups that only appear in one room.

3. Flexible Communication Rules
The paper offers different strategies for how the coordinator and the rooms talk to each other, depending on how strict the privacy rules are.

Scenario A: If you can share a tiny summary (like a count of how many people have which symptoms), the coordinator can do the math easily.
Scenario B: If you can't share even a summary, the rooms can send "directions" (gradients) to the coordinator, who then figures out the best way to combine them without ever seeing the data.

4. Handling the "Small Clusters"
One of the most surprising findings is that this method is actually better at finding rare groups than just dumping all the data into one big computer.

The Analogy: Imagine looking for a specific rare bird. If you look at a massive forest all at once, the rare bird might get lost in the noise. But if you split the forest into small patches, and the rare bird happens to be in one specific patch, the local detective in that patch sees it clearly. When the coordinator combines the reports, that rare bird is identified with high confidence, whereas the "big computer" might have missed it.

The Real-World Test: Elderly Health Records

The authors tested this on real data from the UK: health records of nearly 300,000 elderly people (over 80 years old). They wanted to find patterns of "multi-morbidity" (people having multiple diseases at once).

The Result: The system split the data into 30 "rooms" (simulating different hospitals). It successfully identified 27 distinct groups of patients.
The Discovery:
- One giant group (48% of people) didn't have a specific pattern; they were just the "average" elderly patient.
- Other groups had clear themes: one group was heavily characterized by stroke and HIV; another by dementia and heart issues; a tiny group (only 31 people) was characterized by a specific mix of pancreatitis, arthritis, and erectile dysfunction.
- Crucially, the system found these small, specific groups even though they were hidden in the massive dataset.

The Bottom Line

This paper provides a "pipeline" (a step-by-step recipe) for solving complex data puzzles when the pieces are locked in different rooms. It proves that you don't need to break the locks (share private data) to get a clear picture. In fact, by keeping the data separate and using their new "matching" and "over-fitting" tricks, they can sometimes find rare patterns better than if they had just combined everything into one giant pile.

They compared their method to other existing tools and found that while some other tools are faster, their method is more accurate at finding the true structure of the data, especially when the data is messy or the groups are small.

Technical Summary: Variational Consensus Monte Carlo for Bayesian Mixture Models

Problem Statement
The paper addresses the challenge of performing Bayesian inference for mixture models in a Federated Learning (FL) setting, specifically the "cross-silo" scenario where data is naturally partitioned across multiple nodes (e.g., medical centers) and cannot be pooled due to privacy or regulatory constraints. While Consensus Monte Carlo (CMC) offers a framework for distributed Bayesian inference by aggregating local posterior samples, existing methods face significant limitations when applied to unsupervised clustering. Specifically, the Variational Consensus Monte Carlo (VCMC) approach proposed by Rabinovich, Angelino, and Jordan (2015) assumes that the number of clusters and key mixture parameters are known and that every cluster is present in every data shard. These assumptions rarely hold in practice, particularly in health data applications where cluster prevalence varies across sites and the true number of latent subgroups is unknown. Furthermore, standard CMC aggregation (e.g., weighted means) does not guarantee an optimal approximation of the global posterior, and existing federated clustering methods often rely on mean-field approximations requiring conjugate priors or fail to handle non-conjugate settings without accessing the full dataset.

Methodology
The authors propose a comprehensive pipeline extending VCMC to over-fitted Bayesian mixture models of categorical distributions. The framework operates in two main phases:

The Apply Step (Local Inference):
- An MCMC algorithm (specifically a Gibbs sampler) is run independently within each data shard.
- To handle the unknown number of clusters, the model is "over-fitted" by setting the number of mixture components ( $C'$ ) larger than the anticipated number of true clusters ( $C$ ).
- A fractionated prior is used, where the global prior is split as $p(\theta)^{1/S}$ across shards. This ensures that the product of local posteriors approximates the global posterior.
- The authors implement a parallelized C++17 Gibbs sampler to accelerate local inference.
The Aggregate Step (Global Inference):
- Cluster Matching: Since mixture models are invariant to label permutations, a critical step is matching clusters across shards. The paper introduces and evaluates three algorithms:
  - Hungarian Matching: Enforces a one-to-one correspondence but fails if not all clusters appear in every shard.
  - Minimum Divergence Matching: Minimizes the Kullback-Leibler (KL) divergence between local marginal posteriors.
  - Ball Matching: A novel algorithm that groups clusters that are "close" (within a radius $r$ ) in parameter space, based on convergence rates of over-fitted mixtures. This method allows for merging clusters within a shard if they are indistinguishable, mitigating local over-estimation of cluster counts.
- Variational Aggregation: The aggregation is framed as a variational inference problem to minimize the KL divergence between the approximated and true global posterior. The aggregation function is restricted to a weighted sum of local parameters ( $\sum \upsilon_s \theta_s$ ).
- Optimization: The optimal aggregation weights are found using projected stochastic gradient descent. The paper details four strategies for this optimization to accommodate different FL constraints:
  1. Sharing raw data to the global node (least restrictive).
  2. Federated gradient descent (sharing gradients only).
  3. Limited communication (parallel runs with few iterations and multiple starting points).
  4. Sharing a contingency matrix (summary statistics) to the global node.

Key Contributions
The paper makes four primary methodological contributions:

Extension to Over-fitted Models: It extends VCMC to infer the number of clusters and all model parameters simultaneously without requiring conjugacy, relaxing the assumption that the number of clusters is known.
Novel Cluster-Matching Algorithms: It introduces matching strategies (specifically Minimum Divergence and Ball matching) suitable for cross-silo settings where clusters may be absent in some shards, a scenario where standard Hungarian matching fails.
Flexible Inference Strategies: It proposes multiple aggregation strategies tailored to different federated constraints (e.g., limited communication bandwidth or strict privacy requirements).
Practical Guidelines: It provides empirical guidelines for selecting matching algorithms and inference strategies based on data characteristics (e.g., cluster separability and shard heterogeneity).

Results
The framework was evaluated through a comprehensive simulation study and an application to real-world electronic health record (EHR) data.

Simulation Study:
- Matching Performance: In heterogeneous settings where not all clusters appear in every shard, the Hungarian algorithm failed to recover the correct matching (0% success in non-nested scenarios). The Ball matching and Minimum Divergence matching performed significantly better. The Ball matching was particularly effective at mitigating the over-estimation of the number of clusters caused by local MCMC noise.
- Clustering Quality: VCMC achieved Adjusted Rand Index (ARI) scores comparable to or slightly better than state-of-the-art alternatives (FedMerDel and Federated K-modes) and the "gold standard" full-data MCMC. Notably, VCMC outperformed full-data MCMC in recovering small clusters when the data distribution across shards was heterogeneous.
- Parameter Estimation: VCMC provided low relative bias for cluster weights. In scenarios with poorly separable clusters and heterogeneous shard distributions, VCMC estimated the weights of small clusters more accurately than full-data MCMC, likely because these small clusters were large enough within specific shards to be identified with lower uncertainty.
- Comparison to Alternatives: While FedMerDel was faster, it relied on mean-field approximations that underestimated posterior variances. Federated K-modes showed high bias and poor performance in estimating cluster structures for categorical data.
Real-World Application (THIN Cohort):
- The method was applied to 289,821 geriatric patients from the UK THIN database to identify multi-morbidity patterns.
- The algorithm identified 27 clusters. It successfully characterized clinically meaningful groups, such as clusters dominated by stroke, dementia, or specific cancer types.
- Crucially, it identified small but distinct clusters (e.g., one with only 31 individuals characterized by pancreatitis, rheumatoid arthritis, and erectile dysfunction) that aligned with documented medical correlations, demonstrating the method's ability to detect rare subgroups that might be diluted in a pooled analysis.

Significance and Claims
The authors claim that their framework provides a robust solution for unsupervised clustering in privacy-sensitive environments. The significance lies in the ability to:

Infer the number of clusters and model parameters without data pooling.
Recover small, rare clusters with greater accuracy than standard pooled MCMC when local data shards contain sufficient signal for those specific subgroups.
Operate under strict privacy constraints (cross-silo FL) without requiring conjugate priors or access to the full dataset.

The paper modestly notes that while the framework is theoretically extendable to other mixture distributions, the current validation is specific to categorical distributions. It also highlights that the choice of cluster matching algorithm depends on the specific goal: Ball matching is preferred for estimating the number of clusters, while Minimum Divergence matching may be better if avoiding the merging of distinct clusters is the priority. The authors conclude that while VCMC can be computationally intensive compared to variational methods like FedMerDel, it offers superior accuracy in parameter estimation and cluster recovery in complex, heterogeneous federated settings.