💻 computer science

MuCALD-SplitFed: Causal-Latent Diffusion for Privacy-Preserving Multi-Task Split-Federated Medical Image Segmentation

This paper introduces MuCALD-SplitFed, a novel multi-task Split Federated Learning framework that integrates causal representation learning and latent diffusion to enhance medical image segmentation performance, ensure convergence, and mitigate privacy vulnerabilities in decentralized clinical workflows.

Original authors: Chamani Shiranthika, Hadi Hadizadeh, Parvaneh Saeedi

Published 2026-05-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Chamani Shiranthika, Hadi Hadizadeh, Parvaneh Saeedi

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 a group of hospitals, each holding a unique set of medical images (like photos of embryos, skin lesions, or lungs). They all want to build a super-smart AI to help doctors spot diseases, but they can't share their actual patient photos because of strict privacy laws.

This is where Federated Learning comes in. It's like a group of chefs trying to perfect a secret recipe without ever leaving their own kitchens. They send only their "notes" (math updates) to a central hub, which mixes them together to improve the master recipe, then sends the improved version back.

However, the paper points out a problem with the current version of this system, called SplitFed:

The "One-Size-Fits-All" Problem: Current systems assume every hospital is trying to solve the exact same puzzle. But in reality, one hospital might be looking at skin cancer while another is looking at lung issues. Trying to force them to learn the same thing creates confusion and instability, like trying to teach a chef to bake a cake while simultaneously teaching them to fix a car engine.
The "Glass Wall" Problem: Even though they aren't sharing photos, the "notes" they send (the middle layer of the AI) can sometimes be reverse-engineered. A hacker could look at the notes and reconstruct the original patient photo, breaking privacy.

The Solution: MuCALD-SplitFed

The authors propose a new system called MuCALD-SplitFed. Think of it as upgrading the chefs' communication system with three special tools:

1. The "Causal Detective" (Causal Representation)

Instead of just memorizing that "red spots usually mean skin cancer" (which might be a coincidence or a trick of the lighting), this system learns the actual cause-and-effect relationships.

Analogy: Imagine a detective who doesn't just say, "The suspect was seen near the bank, so they must be the robber." Instead, they ask, "Did the suspect cause the alarm to go off?"
In the paper: The system extracts specific, meaningful features (like the shape of a bubble in an embryo image) and builds a map of how these features truly cause the medical condition. This helps the AI ignore "noise" and focus on what actually matters, even when the hospitals have very different types of data.

2. The "Confetti Machine" (Latent Diffusion)

To stop hackers from reconstructing the patient photos from the notes, the system adds a layer of "confetti" or noise before sending the data.

Analogy: Imagine sending a postcard with your secret message, but before you mail it, you throw a handful of glitter and confetti over it. The recipient (the server) has a special "eraser" (a diffusion decoder) that knows exactly how to remove the confetti and read the message. But a hacker trying to guess the message from the glittery mess sees only static.
In the paper: They use a "diffusion" process to scramble the data sent between the hospital and the server. The server can clean it up to learn, but the scrambled version is useless for reconstructing the original private image.

3. The "Universal Translator" (Domain-Adversarial Alignment)

Since every hospital has different equipment and patient demographics, the data looks different (like different dialects of the same language).

Analogy: Imagine a translator who is so good at their job that they strip away the accent. They ensure the AI learns the meaning of the words, not the specific accent of the speaker. This prevents the AI from accidentally learning "This hospital always uses blue lights, so blue means cancer."
In the paper: This part of the system forces the AI to ignore which hospital the data came from, ensuring the model works well for everyone, regardless of their specific setup.

The Results: What Happened?

The researchers tested this new system against the old "SplitFed" method and other popular AI learning methods using five different types of medical images (embryos, skin, fetal heads, lungs, and polyps).

Stability: The old system was like a shaky tower; it often failed to learn anything useful when mixing different tasks. MuCALD-SplitFed was a solid foundation, converging smoothly and consistently.
Accuracy: The new system significantly outperformed the others. For example, in one test, the old system barely recognized the target (0.045 score), while the new system was highly accurate (0.556 score).
Privacy: When they tried to "reconstruct" the original images from the data sent over the network, the new system made it nearly impossible. The "confetti" (diffusion) worked so well that the reconstructed images were just blurry noise, whereas the old system's data could be somewhat reconstructed.

In Summary

MuCALD-SplitFed is a new way for hospitals to collaborate on AI without sharing patient data. It fixes the confusion of mixing different medical tasks by focusing on true causes rather than coincidences, and it locks down privacy by scrambling the data so it can't be reverse-engineered. The result is a system that is smarter, more stable, and much safer for patient privacy.

Technical Summary: MuCALD-SplitFed

Problem Statement

Medical image segmentation is critical for clinical workflows, yet practical AI systems face significant hurdles in decentralized settings. Hospitals and clinics operate with heterogeneous data (varying modalities, protocols, and demographics) and rarely share raw data due to privacy constraints. While Federated Learning (FL) and Split Federated Learning (SplitFed) enable collaborative training without data sharing, standard implementations assume all clients solve the same task. This "non-IID" assumption is unrealistic in real-world clinical settings where institutions often handle different, domain-specific tasks.

Existing Multi-Task FL and SplitFed approaches attempt to address this by allowing clients to handle different objectives. However, these methods frequently suffer from:

Convergence Instability: Gradients from heterogeneous tasks conflict during aggregation, leading to oscillatory updates.
Feature Entanglement: Shared latent representations at split points mix task-specific features, causing cross-task interference.
Privacy Vulnerabilities: The transmission of latent activations at split points is susceptible to reconstruction attacks, membership inference, and property inference, risks that are amplified by diverse task semantics.

Current solutions often rely on correlation-based representations, which fail to distinguish between task-relevant structures and spurious correlations, limiting robustness under distribution shifts.

Methodology: MuCALD-SplitFed

The authors propose MuCALD-SplitFed, a multi-task SplitFed framework designed for privacy-preserving medical image segmentation. The architecture extends the standard SplitFed model (partitioning into Front-End [FE], Server-Side [SS], and Back-End [BE] components) by integrating three novel stages:

1. Proxy-Label Construction and Causal Graph Discovery

To facilitate weakly supervised causal learning without requiring explicit causal labels, the system extracts domain-specific morphological and intensity features from each dataset to serve as "proxy labels."

Datasets: Blastocyst (embryo), FHPsAOPMSB (fetal head), HAM10K (skin lesions), MosMed (lung CT), and Kvasir-SEG (polyps).
Process: Using Notears-MLP, the system learns dataset-specific causal graphs from these proxy features. The three strongest causal relations per dataset are selected to guide the model (e.g., "Lesion size $\to$ Asymmetry").

2. Causal Representation and Diffusion Module (CRDM)

This module operates on the client side (FE) to generate robust, privacy-preserving latent representations before transmission to the server.

Exogenous Encoder: Maps FE feature maps ( $z_{fe}$ ) to low-dimensional exogenous variables ( $u$ ).
Neural SCM: Learns a causal factorization to produce causally structured latents ( $z_{causal}$ ).
Forward Diffusion: Applies a forward diffusion step to the encoder features ( $z_{fe}$ ) to obfuscate transmitted activations, creating noisy latents ( $z_{fe\_noisy}$ ).
Causal Diffusion Decoder (LDDM-1): Denoises the noisy latents conditioned on the causal latents ( $z_{causal}$ ) to generate $z_{denoised}$ .
Benefit: This process exposes less reconstructive information than standard SplitFed activations, mitigating reconstruction and membership inference attacks.

3. Domain-Adversarial Causal Alignment (DACA)

Located on the server side (SS), this block enforces domain invariance to suppress task-specific biases.

Mechanism: It employs a Gradient Reversal Layer (GRL) and a lightweight domain discriminator.
Training Objective: The segmentation network attempts to "fool" the discriminator, producing domain-agnostic representations, while the discriminator penalizes domain-specific leakage. This improves cross-task generalization and privacy.

Optimization

The model is trained by jointly minimizing a composite loss function (Equation 1) that includes:

Segmentation loss ( $L_{seg}$ )
Proxy-alignment losses ( $L_{proxy}$ )
Diffusion regularization losses ( $L_{diff}$ )
KL divergence terms for exogenous and causal latents ( $L_{KL}$ )
Adversarial domain losses ( $L_{adv}$ )

After each communication round, the SS models, LDDMs, encoders, and Neural-SCMs are aggregated, while FE and BE models remain local to preserve personalization.

Key Contributions

Causal Integration in SplitFed: The first integration of causal modeling within the SplitFed latent space specifically for multi-task segmentation in heterogeneous clinical settings.
Privacy-Enhanced Architecture: A novel combination of causal representations, latent diffusion, and adversarial denoising to reduce information leakage at split points.
Feature Disentanglement: A mechanism to separate task-relevant causal factors from domain-specific information, enhancing both segmentation performance and privacy.
Comprehensive Evaluation: Extensive testing on five heterogeneous medical datasets, comparing against Baseline SplitFed, state-of-the-art (SoTA) personalized FL, and multi-task FL methods.

Experimental Results

The framework was evaluated across five clients using diverse datasets (Blastocyst, HAM10K, FHPsAOPMSB, MosMed, Kvasir-SEG) over 24 communication rounds.

Convergence Stability: Baseline SplitFed failed to converge consistently, exhibiting unstable Intersection over Union (IoU) scores (e.g., final average IoU N/B of 0.001 for UNet). In contrast, MuCALD-SplitFed converged stably, achieving a final average IoU N/B of 0.487 for the UNet architecture.
Segmentation Performance: MuCALD-SplitFed significantly outperformed Baseline SplitFed and SoTA methods (FedPer, FedRep, MOCHA, etc.).
- Dice Score: MuCALD-SplitFed achieved 0.816 (average across clients), compared to 0.394 for Baseline SplitFed.
- IoU N/B: MuCALD-SplitFed achieved 0.579, compared to 0.045 for Baseline SplitFed.
Privacy Metrics: Reconstruction quality metrics (MSE, PSNR, SSIM) indicated that feature maps in MuCALD-SplitFed were far more obfuscated. For instance, the average PSNR for reconstruction attacks dropped significantly (e.g., ~35.7 for MS vs ~7.0 for BS), rendering reconstruction and membership inference attacks effectively infeasible.
Ablation Studies: Removing the diffusion component caused the largest performance degradation, highlighting its critical role in stability and robustness. Both CRDM and DACA provided complementary benefits.

Significance and Claims

The paper positions MuCALD-SplitFed as a necessary evolution for decentralized medical AI, addressing the misalignment between standard FL assumptions and real-world clinical workflows. The authors claim that by moving beyond correlation-based representations to causal–latent diffusion, the framework successfully:

Resolves convergence instability caused by gradient conflicts in multi-task settings.
Mitigates privacy risks inherent in split-point communications.
Improves generalization across heterogeneous domains without requiring raw data sharing.

The work is presented as an "initial step" toward comprehensive privacy-preserving multi-task SplitFed learning. The authors modestly note that future work will involve evaluating stronger adversarial models, assessing scalability to larger client populations, and extending the framework to broader causal settings and adaptive privacy mechanisms.