FedAD: adaptive federated disentanglement for domain generalization in unlabeled medical imaging

1 Introduction

The significant potential of deep learning in the field of medical imaging is profoundly obstructed by a trio of interrelated challenges that create a substantial barrier to the translation of research into practical clinical applications. Firstly, stringent regulations regarding data privacy result in unavoidable data partitions, which restrict models to smaller, less diverse datasets, potentially leading to algorithms that are biased and ineffective (1). Secondly, the considerable expense and time required for expert labeling result in a significant annotation bottleneck, which impedes the scalability of artificial intelligence solutions across the extensive array of diseases and imaging modalities. Lastly, when models are deployed in new clinical settings, they often struggle to generalize due to domain shift (2). This represents a critical issue concerning safety and reliability, as a model’s performance can gradually decline due to subtle variations in imaging equipment, patient demographics, or imaging protocols. Although powerful methodologies such as Federated Learning (FL), self-supervised contrastive learning (3), and Domain Generalization (DG) have surfaced as possible solutions, they frequently operate under conflicting assumptions. Recent meta-analyses confirm that robust generalization requires careful dataset harmonization and validation protocols that extend far beyond standard deployment paradigms (4).

A cohesive framework that effectively integrates these paradigms is noticeably lacking. Existing federated semi-supervised learning (SSL) approaches, such as MOON (5), concentrate on alleviating heterogeneity among current clients but do not explicitly optimize for generalization to completely new, unseen domains; they tend to prioritize intra-federation performance over extra-federation robustness. On the other hand, federated DG methods (6) generally assume that fully labeled datasets are available, which is often impractical and undermines the primary advantage of SSL. This delineates a clear research gap: the need for a framework capable of learning from unlabeled, decentralized data while simultaneously optimizing for robust generalization to previously unseen domains.

This study addresses this gap by introducing FedAD: Adaptive Federated Disentanglement. FedAD is a unified framework based on two fundamental, synergistic principles. Rather than relying on implicit assumptions of domain invariance, it utilizes FedSD, an innovative learning objective designed to compel the model to learn two orthogonal feature spaces: a domain-invariant semantic space and a domain-specific style space. This explicit separation allows the model to effectively quarantine scanner-specific noise and artifacts into the style space, yielding a cleaner and more robust semantic signal for any downstream clinical task. This disentanglement is the key to generalization, as the semantic space is trained to be invariant to the domain-specific ’noise’ captured by the style space. Second, it introduces an Adaptive Teacher-Student Alignment (ATSA) curriculum, which dynamically adjusts the strength of the generalization pressure based on the global model’s stability. This adaptive mechanism treats the global model as a student that gradually becomes a teacher, preventing the federated system from prematurely locking in errors from an initially naive consensus and allowing it to transition gracefully from a phase of client-level discovery to one of strong, collective learning.

The main contributions of this work are therefore as follows: The remainder of this paper is organized as follows: Section 2 provides an overview of the fundamental concepts and limitations associated with federated learning, self-supervised learning, and domain generalization from previous research. Section 3 identifies the research gap as revealed in the literature and elaborates on the why and what the design of FedAD is, to address this gap found. Section 4 offers a comprehensive formal approach to the FedAD framework, detailing the model architecture, the mathematical formulation of the FedSD objective, and the algorithm used for the ATSA controller. Section 5 outlines our experimental setup, encompassing the diverse, multi-modal datasets, evaluation protocols, and baseline models employed. Section 6 presents our ablation study to prove component synergy. Section 7 details our empirical findings, including an analysis of convergence dynamics, and the generalization performance of FedAD. Section 8 addresses the present challenges and limitations inherent in our framework. Section 9 concludes the paper by summarizing our key contributions and proposing promising avenues for future research.

• (1) We propose FedAD, a holistic framework that learns from decentralized, unlabeled data while explicitly optimizing for generalization to unseen domains.
• (2) We demonstrate how FedAD provides a novel, synergistic integration of privacy-preserving federated learning (FL), self-supervised learning (SSL), and domain generalization (DG) within a single, unified architecture.
• (3) We introduce FedSD (Federated Semantic Disentanglement), a novel, non-adversarial objective that applies a direct orthogonality constraint at the client level. This explicitly separates features into domain-invariant semantic vectors and domain-specific style vectors to produce a clean, generalizable representation.
• (4) We design an Adaptive Teacher-Student Alignment (ATSA) curriculum, a server-side mechanism that dynamically controls the generalization pressure based on global model stability. This adaptive schedule prevents premature convergence and ensures a stable transition from client-level discovery to collective learning.
• (5) We propose a realistic experimental setup incorporating both radiological and pathological datasets to demonstrate that FedAD significantly outperforms existing approaches in generalizing to unseen domains.

2 Literature review

To fully contextualize the proposed methodology’s novelty, a comprehensive review of the foundational concepts, key advancements, and remaining limitations within Federated Learning, Self-Supervised Learning, and Domain Generalization is needed.

2.1 Foundations of federated learning and client drift

Federated Learning is becoming the go-to approach for ensuring privacy in machine learning, especially when it comes to sensitive areas like healthcare. At its core, the FedAvg algorithm (1) sets up a smart protocol where local models are trained using private client data. Only the model parameters are sent to a central server for aggregation, which helps create a global model that learns from everyone’s insights without needing access to the original data. However, putting FL into practice does come with a major hurdle: the challenge of statistical heterogeneity. Unlike traditional distributed learning, the data in federated settings is typically non-IID (non-independent and identically distributed).

This heterogeneity can be split into two main categories: (1) System Heterogeneity, where clients may have different computing power, network speed, and availability; and (2) Statistical Heterogeneity, in which the local data distributions vary among different clients. This variation poses a significant challenge, often appearing as feature shifts, label shifts, and various intricate covariate shifts. Because of this non-IID characteristic, the local training goals can start to stray from the global objective. This issue, referred to as “client drift,” can lead to instability in training, slow down convergence, and in some cases, result in the global model completely diverging (7). Researchers have explored this divergence theoretically, with studies like (8) clearly showing the impact of non-IID data on FedAvg’s convergence guarantees, highlighting the need for improved aggregation and local training techniques.

2.2 Advanced techniques for mitigating heterogeneity

A significant body of research aims to mitigate client drift. One major line of work modifies the local client objective. FedProx (9) introduced a proximal term to the local loss, regularizing local updates to stay closer to the global model, which helps account for system heterogeneity and, to a lesser extent, statistical drift. A more sophisticated solution, SCAFFOLD (10), introduced control variates at both client and server levels to correct for local update directions based on an estimate of the local gradient drift. More recently, methods like FedDyn have introduced dynamic regularizers that are updated at the server to align local and global stationary points.

Other methods tackle heterogeneity by modifying the model architecture or aggregation process itself. FedBN (11), for example, addresses feature shift by simply keeping batch normalization layers local to each client, preventing the aggregation of domain-specific statistics. Fed-CP (12) learns both a global consensus model and client-specific parameters to better personalize models for each client’s unique data distribution.

Another line of research involves more intelligent server-side aggregation. For example, FedMA (13) matches and averages layers of deep neural networks based on structural similarity, while FedCA (14) introduces a cross-attention mechanism to adaptively aggregate model features for medical image segmentation. While powerful, these methods are typically (1) designed for supervised learning and (2) primarily aim to improve performance on the existing client distributions. They do not explicitly solve the challenge of learning from unlabeled, heterogeneous clients to optimize for generalization to unseen domains.

2.3 The imperative for self-supervision in medical imaging

The prohibitive cost and logistical challenges of acquiring expert-annotated data have made SSL an indispensable tool, especially in medicine. SSL provides a path to leverage vast archives of unlabeled imaging data to learn rich feature representations that can then be used for downstream tasks with very few labels. This paradigm shifts the learning burden from manual annotation to algorithmic discovery. While generative SSL approaches (e.g., masked autoencoders) are powerful, contrastive methods have become dominant for learning separable, transferable representations that are highly effective for downstream classification and segmentation, making them critical for unlocking the full potential of large-scale, unlabeled medical datasets.

2.4 Mechanics of contrastive representation learning

Among various SSL paradigms, Contrastive Learning has emerged as the state-of-the-art. Foundational frameworks like SimCLR (3) operate on an instance discrimination task. For any image xi (the “anchor”), a stochastic augmentation pipeline creates two distorted but semantically identical views (the “positive pair”). All other images in the batch are treated as “negatives.” The model maps these views into an embedding space, and the contrastive loss function, NT-Xent, pulls the representations of the positive pair closer together while pushing them far apart from all negative representations. A key finding was the importance of the projection head, which maps the learned features to a space where the loss is applied, but which is discarded after training to reveal more robust features. This concept is extended to FL by MOON (5), which explicitly addresses client drift within the SSL framework by implementing a model-level contrastive loss to align local and global model representations. Building on this, the combination of FL and SSL has been further investigated in various research. For example, (15) explicitly addresses domain shifts in medical imaging using federated SSL. In order to increase performance on unsupervised medical segmentation tasks, more sophisticated frameworks such as P-FCL (16) build customized federated contrastive learning, which again divides the model into global and local components.

2.5 The challenge of domain shift in clinical deployment

A model’s ability to perform reliably on data from a source different from its training data is known as Domain Generalization (DG). It is a critical requirement for deploying safe AI in the real world (17). In a clinical context, domain shift arises from differences in scanner manufacturers (e.g., Siemens vs. GE), imaging protocols, staining reagents (e.g., batch effects), and patient populations. A model trained at one site may experience a silent and catastrophic drop in performance when deployed at another. The difficulty of this problem has been highlighted by benchmarks like DomainBed (2), which show that many sophisticated algorithms fail to consistently outperform simple empirical risk minimization (ERM).

2.6 Disentanglement as a strategy for domain generalization

One of the most powerful strategies for DG is to learn disentangled representations. The core idea is to train a feature extractor that explicitly separates the causal, semantic factors of variation (e.g., tumor presence) from the non-causal, domain-specific factors (e.g., scanner noise). The foundational work here is Domain-Adversarial Training of Neural Networks (DANN) (18). DANN introduced an adversarial objective where a feature extractor is trained to “fool” a domain classifier via a gradient reversal layer. The equilibrium of this game is a feature representation that is predictive for the main task but contains no discernible information about its source domain. Other approaches include Variational Autoencoders (VAEs) that penalize the KL-divergence of the latent space to encourage factorization, or information-theoretic methods that directly minimize the mutual information between the feature vector and the domain label.

2.7 Federated and non-adversarial DG approaches

Applying DG in a federated context is uniquely challenging because data from different domains cannot be pooled. This prevents the direct use of centralized DG methods like DANN (18), as a global domain classifier cannot see data from all domains simultaneously. While federated DANN variants exist, they often require complex sharing of gradient statistics or domain labels, which can be unstable. Several works have begun to design frameworks to address this. FedDG (19) pioneered this by introducing a domain generalization approach specifically for federated medical image segmentation. Recent methods have explicitly adopted non-adversarial disentanglement. For example, (20) uses a VAE-based federated approach to disentangle features for skin lesion segmentation. Similarly, FedDis (21) proposes a disentangled federated learning framework that separates the model into a global, domain-invariant component and a local, domain-specific component for classification tasks. Given the inherent instability of adversarial training (e.g., non-convergence, mode collapse, sensitivity to hyperparameters), which is exacerbated in a distributed, non-IID setting, a direct, non-adversarial approach is highly desirable for FL. This motivates the exploration of alternative, more stable disentanglement strategies.

3 Bridging the research gap

A significant gap between current solutions has been identified by the examination of the literature. Although there are techniques for FL, SSL, and DG separately, they work under different presumptions. Current federated SSL techniques, such as MOON (5), focus on reducing drift among current clients rather than explicitly optimizing for generalization to completely unseen domains. Conversely, most federated DG methods (e.g., (6)) presume the availability of fully labeled data, which is often infeasible in medical contexts and defeats the purpose of SSL. Finally, classical DG methods like DANN (18) are centralized, and their adversarial nature makes them unstable and difficult to adapt to a federated setting.

This paper addresses this critical gap. We propose a framework that, for the first time, synergistically unifies all three paradigms, with a comparative analysis of these paradigms detailed in Table 1. It is (1) Federated by design for privacy, (2) Self-Supervised to learn from unlabeled data, and (3) a Domain Generalization method that explicitly optimizes for unseen domains. This is achieved through a novel, non-adversarial disentanglement objective stabilized by an adaptive curriculum.

Table 1. Comparative analysis of representative paradigms, highlighting the novelty of the proposed FedAD framework.

4 Proposed method: the FedAD framework

The proposed FedAD framework is designed to learn domain-invariant semantic representations from unlabeled medical images distributed across multiple clients, without sharing raw data. The methodology is built to explicitly handle domain shift and ensure stable convergence in a federated setting by integrating two core mechanisms: client-side FedSD and a server-side adaptive controller, ATSA.

4.1 Formal problem definition

We consider a federated network of K clients, where each client k possesses a local, unlabeled dataset Dk={xik}i=1nk. The data distribution Dk for each client is distinct, representing a unique domain (e.g., a specific hospital or scanner). The objective is to learn a single, global, domain-generalizable feature extractor fg(⋅;θg) by training local extractors fk(⋅;θk) on each client and aggregating them at a central server. The goal is that the learned representations hg=fg(x) are effective for a downstream task (e.g., classification) on data from a completely unseen target domain Dt, which was not involved in the federated training process. The challenge lies in ensuring that the aggregated θg converges to a state that represents domain-invariant semantic features, rather than an ’average’ of domain-specific biases. A high-level overview of this federated setup is presented in Figure 1, and a simplified schematic of the learning loop is shown in Figure 2.

Clinical Rationality of Heterogeneous Client Federation: The federated clients constructed in this study are sourced from radiological imaging, histopathology, dermatoscopy, and endoscopic datasets. Although they exhibit significant variations in imaging modalities, anatomical regions, pathological types, and imaging mechanisms, they can serve as highly effective co-learning clients under a unified collaborative medical framework. The rationale is that medical imaging, regardless of modality, relies heavily on shared foundational visual primitives—such as edge gradients, textures, contour formations, and cellular pattern distributions. By federating these disparately heterogeneous clients under a unified SSL objective, the framework actively prevents the global model from memorizing modality-specific shortcuts. Instead, it forces the extraction of universally applicable representational kernels, whilst the FedSD mechanism explicitly quarantines the highly variant imaging mechanisms into the domain-specific subspace.

Figure 2 illustrates the core, privacy-preserving cycle of our framework. The process begins with the Central Server, which runs the ATSA controller, sending the current global model weights (θg) and the dynamically computed alignment weight (λt) to a selection of clients (Step 1). Each client then performs its local training (Step 2) using our novel FedSD objective, which is detailed in Section 4.3. Crucially, after local training, clients only return their updated extractor weights (θk) to the server (Step 3). The local disentanglement head (gk) and any client-specific data remain private, ensuring no raw data is ever shared. The server then aggregates these weights to form the next global model (θgt+1).

4.2 Model architecture and notation consistency

The model architecture, as depicted in Figure 1, decouples the feature extractor from a local disentangling head. To ensure strict notation consistency, let θg(t) denote the global extractor weights at communication round t. At the beginning of each communication round, clients explicitly initialize their local extractor weights by copying the global weights: θk←θg(t). During local optimization, θk is actively updated via gradient descent, while θg(t) is maintained locally as a strictly frozen teacher network, hereafter denoted as θ~g. Upon completion of local training epochs, only the updated parameters θk are transmitted back to the server for aggregation.

4.2.1 Global feature extractor

The global model, fg(⋅;θg), is a standard feature extractor whose parameters θg are the only ones aggregated by the server. We use a ResNet-18 backbone, pre-trained on ImageNet, with its final fully-connected classification layer removed. This model’s role is to learn a general-purpose feature representation h=fg(x)∈RF, where F=512 for ResNet-18.

4.2.2 Local disentangling head

Each client k maintains a local projection head gk(⋅;ϕk), which is not shared with the server. This head takes the features hk=fk(x) produced by the client’s local extractor and projects them into two distinct subspaces: This head gk is a simple two-layer MLP (Linear → ReLU → Linear) that maps hk∈RF to Rds+dd, which is then split. Crucially, the projection dimensions ds=128 and dd=128 are fixed and strictly invariant across all datasets and clients. This enforces reproducible representational capacity, preventing domain-specific models from altering the structural geometry of the semantic space. By keeping gk local, we prevent the server from aggregating client-specific, non-generalizable information learned in the final layers, a common issue in FedAvg (1) where the entire model is averaged.

• A semantic projection zs∈Rds
• A domain projection zd∈Rdd

4.3 Client-side training: the FedSD objective

At the start of round t, client k receives the global extractor weights θgt and sets its local extractor θk←θgt. It also receives the adaptive weight λt from the ATSA controller. The client then trains both its local extractor fk and its local head gk for E epochs using a composite loss LkFedAD.

A detailed schematic of this local update data flow is presented in Figure 3. The process begins with a single “Image xi”, which is passed through an “Augment T()” block to create two distinct views, xi(a) and xi(b). These views are then processed in parallel. The “Local Extractor fk” (the student) produces “Local Features hk”, while the frozen “Global Extractor fg” (the teacher, denoted θ~g) produces “Global Features hg”. The local features hk are passed to the “Local Head gk”, which disentangles them into “Semantic Projections zs” and “Domain Projections zd”. Finally, these three sets of vectors (zs, zd, and zg) are used to compute the three loss components.

For a batch of N images, generating 2N augmented views, the client performs the required forward passes and calculates the total loss for the batch as the sum of three components (Equation 1):

LkFedAD=LS-CL+LS-Align+LOrtho(1)

4.3.1 Semantic contrastive loss (LS-CL)

This loss enforces instance discrimination strictly on the semantic projections. To ensure mathematical tightness, the NT-Xent loss is indexed rigorously over the 2N augmented samples within the batch. Let p(i) denote the index of the corresponding positive augmented view for sample i (Equation 2):where sim(⋅,⋅) is cosine similarity and τ is the temperature. This loss pulls semantic projections of the same image together, while pushing all others apart, forcing fk and gk to encode semantic content into zs. As shown in Figure 3, this loss solely relies on the “Semantic Projections zs”.

LS-CL=−12N∑i=12Nlog⁡exp⁡(sim(zs(i),zs(p(i)))/τ)∑j=12N⊮[j≠i]exp⁡(sim(zs(i),zs(j))/τ)(2)

4.3.2 Semantic alignment & anti-alignment loss (LS-Align)

Our novel, diagram-informed loss (see Figure 3) uses the frozen global extractor as a “teacher” (θ~g) that holds the global consensus knowledge. To ensure mathematical consistency for cosine similarity computations, the global feature vector hg∈R512 is mapped to a lower-dimensional subspace zg∈R128 via an auxiliary linear projection layer pg(⋅), ensuring its dimensions exactly match those of the local projections. The local model is trained as a “student” that must learn to separate its knowledge. The loss forces the local semantic projection zs to align with the global projection zg, while simultaneously forcing the local domain projection zd to anti-align with zg. We formulate this using cosine similarity, averaged over both views v∈{a,b} (Equation 3):This dual-pressure objective is a core innovation. The alignment term acts as a form of federated knowledge distillation, pulling the local semantic vector towards the global consensus. Simultaneously, the anti-alignment term explicitly enforces the local model to expel global consensus information from its domain vector.

LS−Align=λt2∑v∈{a,b}[(1−cos⁡(zs(v),zg(v)))+(1+cos⁡(zd(v),zg(v)))](3)

Justification for Anti-Alignment: Unlike margin-based or information-theoretic penalties which are heavily dependent on batch statistics and computationally expensive to estimate reliably in highly non-IID federated settings, cosine anti-alignment provides a direct, bounded gradient. It explicitly forces the domain vector to orthogonalize against the global consensus without the oscillatory instability inherent to adversarial projections. This dual pressure encourages the local model to distill the global model’s general knowledge into zs, while explicitly identifying and isolating local, non-consensus (i.e., domain-specific) information into the zd vector.

4.3.3 Orthogonality loss (LOrtho)

To complete the disentanglement, we penalize the squared cosine similarity between the two local projections, forcing them to be linearly independent:This simple, non-adversarial loss ensures that zs and zd capture mutually exclusive information. It serves as a powerful regularizer, ensuring the two projection spaces do not collapse or become redundant. Without this, the model could ’cheat’ by encoding semantic information in zd (or vice-versa) to minimize the other losses.

LOrtho=γ2∑v∈{a,b}cos2⁡(zs(v),zd(v))(4)

The conceptual goal of this loss is illustrated in Figure 4. The total feature space hk is forced to factorize into two orthogonal vectors: a “Semantic Vector (zs)” representing “The What” (i.e., the underlying domain-invariant pathology) and a “Domain Vector (zd)” representing “The Where” (i.e., domain-specific scanner artifacts or noise). The LOrtho loss (Equation 4) is the mechanism that enforces this linear independence, effectively quarantining domain-specific information into zd and leaving zs as a clean, generalizable representation.

4.4 Server-side controller: the ATSA mechanism

The server orchestrates the federated process and adaptively controls the learning curriculum by dynamically tuning the alignment weight λt. The logic of this server-side controller is visualized in the flowchart in Figure 5.

As shown in Figure 5, at the start of a new round t+1, the server first receives the local weights (θkt+1) from the participating clients. It then performs two key steps. First, it aggregates these weights to create the new global model (θgt+1). Second, it immediately calculates the model’s “Instability” (Δt+1) based on the magnitude of this update (as defined in Equation 5). Third, this instability score is fed into the curriculum scheduler (Equation 6) to “Compute the New Weight” (λt+1). This new alignment weight is not used immediately; it is stored and sent to clients in the next round (t+2), thus completing the adaptive feedback loop.

4.4.1 Stability measurement

After aggregating the client extractors fk(⋅;θkt+1) to produce the new global model fg(⋅;θgt+1), the server measures parameter instability Δ(t+1):Here, vec(⋅) is the vectorization operator, and ε (e.g., 10−8) prevents division by zero. This L2-norm of the weight delta is a direct measure of the magnitude of the consensus update. Empirical validation across our experiments demonstrates that L2-norm parameter drifts serve as a highly reliable, label-free proxy for training maturity, consistently correlating with loss plateaus in non-IID settings. This metric keeps the framework’s self-supervised and federated integrity intact by avoiding the impractical assumption of a central validation set. Significant client disagreement and quick changes in the global state, which are typical of the early exploration phase, are indicated by high values of this parameter drift.

Δ(t+1)=‖vec(θg(t+1))−vec(θg(t))‖2‖vec(θg(t))‖2+ε(5)

4.4.2 Curriculum scheduling of alignment

The instability score Δ(t+1) is used to compute the alignment weight λt+1 for the next round via a smoothed, exponential decay schedule.This update has two parts: Hyperparameters such as α, k, λmin, and λmax were systematically selected based on grid search during the 5-fold cross-validation phase to ensure optimal schedule timing.

λt+1=αλt+(1−α)[λmin+(λmax−λmin)⋅e−k⋅Δ(t+1)](6)

• Adaptive Target: The term in brackets calculates a target λ. k controls the steepness of the transition. When instability Δ is high (early training), the decay term e−kΔ approaches 0, and λ correctly maps to the lower bound λmin. Conversely, when Δ approaches 0 (stable convergence), the decay term e0 approaches 1, allowing the formula to correctly map to the strict upper bound λmax.
• EMA Smoothing: The outer expression is an exponential moving average (EMA) with smoothing factor α. This prevents drastic oscillations in λt due to any single noisy round, ensuring a smooth curriculum.

4.5 Algorithm summary

The complete FedAD training procedure is summarized in Algorithm 1. The process is orchestrated by the server, which begins each communication round t by calculating the adaptive alignment weight λt+1 based on the model’s instability from the previous round (Equations 5 & 6). It then sends the current global extractor weights θgt and the new weight λt+1 to a subset of clients. Each client, in the ClientUpdate procedure, copies the global model to act as a frozen “teacher” (θ~g) and trains its local model for E epochs. During local training, the client computes the three-component FedAD loss: the semantic contrastive loss (LS-CL), the orthogonality loss (LOrtho), and the adaptive alignment loss (LS-Align). After its local updates, each client returns only its updated extractor weights (θkt+1) to the server, which aggregates them to form the next global model θg(t+1).

Algorithm 1.The FedAD Framework. 1: Server executes: 2: Initialize global extractor weights θg0 3: Initialize λ0←λmin 4: for each communication round t=0,1,…,T−1 do 5: Select a subset of M clients St 6: if t>0 then 7: Compute instability Δ(t) using Eq. 5 8: Compute adaptive weight λt+1 using Eq. 6 9: else 10: λt+1←λmin 11: foreachclientk k∈St in parallel do 12: Send θgt and λt+1 to client k 13: θkt+1←ClientUpdate(k,θgt,λt+1) 14: Aggregate θg(t+1)←∑k∈StnkNθkt+1 ▹ Only extractor 15: 16: procedure ClientUpdate(k,θg,λ) 17: Initialize local extractor θk←θg 18: Initialize local head ϕk (if t=0) 19: Freeze global extractor θ~g←θg 20: for each local epoch e=1,…,E do 21: for batch of data {xi}⊂Dk do 22: Generate augmented views {xi(a),xi(b)} 23: Compute LS-CL, LS-Align(θk,ϕk,θ~g,λ), LOrtho 24: LkFedAD←LS-CL+LS-Align+LOrtho 25: Update θk,ϕk using ∇θk,ϕkLkFedAD 26: return θk ▹ Return only extractor weights

5 Implementation and experimentation

This section details the training protocols, datasets, and evaluation methodologies used to validate the performance of the FedAD framework.

5.1 Federated datasets and client configuration

To create a diverse and heterogeneous federated network, we select six distinct, publicly available datasets. This diverse, multi-modal, and multi-scale collection is designed to create a highly challenging, non-IID environment that mirrors real-world federated collaborations, preventing the model from overfitting to a single modality. All labels are strictly withheld during the federated training phase. (See Table 2).

Table 2. Federated client datasets used for SSL training.

5.2 Data preprocessing and augmentation

All images are pre-processed by resizing to 224×224 pixels using bicubic interpolation and normalizing using ImageNet statistics.

The stochastic augmentation pipeline T(⋅) used to generate views (a and b) is identical to SimCLR. It includes a random resized crop to 224×224 (with a scale from 0.08 to 1.0), followed by a random horizontal flip (with p=0.5). We then apply color jitter with a 0.4 strength to brightness, contrast, saturation, and hue (with p=0.8), and randomly convert images to grayscale (with p=0.2). Finally, a Gaussian blur using a 23-pixel kernel and a sigma between 0.1 and 2.0 is applied (with p=0.5).

This specific pipeline is strong yet ’content-preserving,’ meaning it alters style (color, orientation) without destroying the underlying anatomical or pathological features, which is critical for SSL in medicine.

5.3 Unseen target domains for generalization testing

To address the necessity of demonstrating generalization across multiple distinct medical modalities, we evaluate FedAD on two unseen target domains completely excluded from the federated training process.

5.3.1 Radiology target

We use the PadChest dataset (28) to evaluate generalization within the respiratory modality. A balanced test set of 10,000 images is used for multi-class disease classification to assess robustness against novel scanner types.

5.3.2 Histopathology target

We use the Camelyon17 dataset (29). This dataset is a standard benchmark for domain generalization in histopathology, as its data comes from 5 hospitals with distinct staining and acquisition protocols. We use data from two hospitals as our labeled training split (approx. 40,000 patches) and data from a third, unseen hospital as the test set (approx. 12,000 patches) for binary tumor classification.

5.4 Training protocol and hyperparameters

Federated training is run for 40 communication rounds. In each round, the server samples M=4 out of K=6 clients using uniform random sampling. Each selected client performs E=3 local epochs of training using the Adam optimizer with a learning rate of 5×10−4, weight decay of 10−5, and betas of (0.9,0.999). The local batch size is 64.

5.5 Evaluation protocol: linear probing and K-fold CV

After federated training, the final global extractor fg is frozen. We evaluate its representation quality using linear probing on the unseen target domain’s labeled training split. All quantitative metrics reported in our tables (e.g., values following ±) strictly represent the Standard Deviation calculated across 5 independent experimental runs utilizing different random seeds to ensure reproducibility. Statistical significance for reported improvements is validated via paired t-tests (p<0.05).

• (1) Feature Extraction: We pass the entire target training split (e.g., 40,000 Camelyon17 patches) through the frozen fg to get a set of 512-dim feature vectors and their corresponding labels.
• (2) K-Fold CV for Tuning: We perform 5-fold cross-validation on these features. In each fold, a simple linear classifier (Logistic Regression with a small L2 penalty of 10−4) is trained on 80% of the features and validated on 20%.
• (3) Classifier Training: The classifier is trained using the L-BFGS optimizer for a maximum of 100 iterations. We chose L-BFGS as it is a standard, deterministic optimizer for convex problems like logistic regression, ensuring reproducible evaluation results.
• (4) Hyperparameter Selection: The mean F1-Score across these 5 folds is used to select the best hyperparameter for γ, as shown in Table 3 and Figure 6.
• (5) Final Test: The linear classifier is re-trained on 100% of the target training features. This single trained classifier is then evaluated once on the held-out target test set to produce the final metrics.

Table 3. Key hyperparameters for FedAD loss components and ATSA.

Figure 6 provides a clear visual summary of the 5-fold cross-validation results presented in Table 3. This bar chart plots the mean F1-score against the orthogonality weight (γ). It is evident that the model’s performance is sensitive to this hyperparameter, with γ=0.0 (no disentanglement) yielding the lowest score. The performance peaks at γ=2.0, which achieves the highest mean F1-score of 0.910. The error bars, representing the standard deviation across the 5 folds, are tightest for γ=2.0, indicating a stable and reliable performance gain. This result strongly justifies our selection of γ=2.0 as the optimal value.

5.6 Baseline models for comparison

All baselines use the same ResNet-18 backbone and are trained for a comparable number of optimization steps to ensure a fair comparison.

• Centralized Supervised: ResNet-18 trained on pooled source data with labels. Violates privacy, serves as a theoretical upper bound.
• Centralized SSL (SimCLR) (3): SimCLR trained on pooled source data (no labels).
• FedAvg-SSL (1): A standard federated SimCLR (both backbone and projection head are aggregated).
• FedProx-SSL (9): We add a proximal term (FedProx (9)) to FedAvg-SSL to show the effect of a basic drift-mitigation technique.
• MOON (5): Re-implemented MOON. For a fair comparison, we apply its model-contrastive loss between the feature extractor outputs (hk and hg), using their recommended static alignment weight μ=1.0.
• FedDG (19): A federated domain generalization approach adapted for the classification task.
• FedDis (21): A federated disentangled representation approach that separates global domain-invariant components from local domain-specific ones.

5.7 Computational environment

All experiments were developed and simulated using a custom federated learning framework built on PyTorch 2.1 and Python 3.10. Initial development and small-scale tests were run locally. Due to computational demands, the main federated training experiments were executed on a cloud compute instance, utilizing a single NVIDIA A100 (40GB) GPU. The final linear probing evaluations were performed locally.

6 Ablation study

To precisely quantify the individual and synergistic contributions of our proposed mechanisms, we conduct a rigorous ablation study. The FedAD framework is built on two core innovations: We systematically deconstruct the full model to isolate and analyze the impact of each component. We evaluate four distinct configurations on the unseen histopathology target test set: The performance of these configurations is summarized in Table 4.

• (1) FedSD: The explicit, non-adversarial disentanglement objective, primarily enforced by the LOrtho loss.
• (2) ATSA: The adaptive teacher-student curriculum, which dynamically schedules the LS-Align loss weight (λt).

• FedAvg-SSL (1): The baseline federated self-supervised model. This configuration has no explicit disentanglement (γ=0) and no alignment curriculum (λ=0).
• FedAD w/o Orthogonality: This model uses the full ATSA curriculum but disables the explicit disentanglement by setting the orthogonality weight to zero (γ=0). This tests if alignment alone is sufficient.
• FedAD w/o ATSA: This model uses the full FedSD disentanglement architecture (γ=2.0) but replaces the adaptive ATSA controller with a static, manually-tuned alignment weight (λ=0.5). This tests if the architecture alone is sufficient.
• FedAD (Full Model): Our complete proposed framework, with both FedSD (γ=2.0) and ATSA (λt) active.

Table 4. Ablation study performance on unseen histopathology target domain (test set).

Bold values indicate the highest performance achieved across the compared methods.

6.1 Analysis of ablation results

The results in Table 4 are unequivocal. They reveal that both components are crucial for high performance, and their combined effect is synergistic, not merely additive. The full FedAD model significantly outperforms all ablated variants.

6.1.1 The criticality of explicit disentanglement (FedSD)

The most significant performance drop occurs when the orthogonality constraint is removed. The “FedAD w/o Orthogonality” model, despite having the adaptive ATSA curriculum, witnesses its F1-score reduce from 0.910 to 0.853. This result is critical, as it validates our core hypothesis: explicit disentanglement is essential for generalization.

This configuration tests if alignment alone is sufficient, and the answer is a clear “no.” Without the LOrtho pressure, the model is incentivized to “cheat.” It leaks non-generalizable, domain-specific features (e.g., scanner artifacts, staining variations) into the semantic vector (zs). It does this because these spurious correlations make the local instance discrimination (contrastive) task easier to solve. However, this creates a “contaminated” semantic vector that is hopelessly overfit to the source domains, causing it to fail dramatically when generalizing to a new, unseen domain.

This finding is further corroborated by our hyperparameter tuning in Table 3 and Figure 6. Both show that γ=0.0 (the “No Disentangle” case) yields the lowest cross-validation performance, and performance peaks at our selected value of γ=2.0.

6.1.2 The necessity of an adaptive curriculum (ATSA)

Removing the adaptive curriculum also causes a significant performance drop. The “FedAD w/o ATSA” model, which uses a static λ=0.5, achieves an F1-score of only 0.871. This demonstrates that the training dynamic is just as important as the model architecture.

A static, fixed weight is a poor compromise that fails at both ends of the training process due to the following concerns: The ATSA mechanism perfectly solves this dilemma. As shown in our results, it intelligently keeps λt low during the high-instability phase (Figure 11) to encourage discovery, then smoothly increases it as the model stabilizes (Figure 10) to enforce consensus.

• Premature Convergence: In the initial, unstable training phase (clearly shown by the high instability in Figure 9), a static weight of 0.5 is too restrictive. It forces premature convergence, compelling clients to align with a naive, under-trained global model that has not yet learned a meaningful consensus. This disrupts the crucial early-stage local exploration and can lock in optimization errors.
• Insufficient Consensus: Conversely, in the later, stable phase, a weight of 0.5 is too lenient and broad. It fails to enforce a strong, unified consensus, leaving the global model as a loose average of specialized local models rather than a truly robust and generalizable representation.

6.2 Synergistic effect

In summary, our ablation study demonstrates that FedAD’s components are highly synergistic rather than being a mere addition. FedSD offers the ideal framework for distinguishing between domain noise and semantic information. To properly study inside that framework, ATSA offers the appropriate training curriculum. Sub-optimal generalization results from breaking the system by removing either component. This meticulously crafted dual-component system directly contributes to the state-of-the-art 0.910 F1-score of the entire model.

7 Results and discussion

This section presents a detailed analysis of the experimental results, now fully aligned with the provided figures.

7.1 FedAD convergence vs. baselines

A primary indicator of a federated method’s effectiveness is its convergence speed and stability. In Figure 7, we plot the training loss (on a log scale) for FedAD against two standard baselines, FedAvg (1) and FedProx (both adapted for SSL).

The results are clear: FedAvg (1) exhibits the slowest convergence and highest volatility, which is characteristic of unconstrained client drift in a heterogeneous (non-IID) setting. FedProx, by adding a proximal term, improves stability and achieves a lower loss than FedAvg (1), but still struggles. FedAD (Ours) shows markedly superior performance. Its loss descends more rapidly and smoothly, achieving a final loss value nearly 2x lower than FedAvg (1). The superior convergence of FedAD is likely due to the FedSD objective. By explicitly forcing domain-specific information into a separate vector (zd) that is not aggregated, FedAD reduces the amount of conflicting, domain-specific ’noise’ in the gradients that are sent to the server. This leads to a cleaner, more stable, and faster-converging global model.

7.2 Analysis of FedAD’s loss components

To understand our model’s internal learning process, Figure 8 plots the evolution of the four components of the LkFedAD composite loss, averaged across all participating clients per round.

The plot shows a stable and successful multi-objective optimization: This stable, parallel convergence confirms that the different loss pressures are not conflicting. The smooth, monotonic decrease of all loss components, especially the LOrtho loss, strongly suggests that our multi-objective optimization is stable and finds a good equilibrium, unlike the oscillatory behavior often seen in adversarial methods. It is noteworthy that the LS-CL (red) and LS-Align (green) losses descend in parallel. This suggests the local instance discrimination task (learning semantics) and the global alignment task (matching consensus) are not in conflict, but are mutually beneficial.

• Total Loss (blue) starts high (approx. 2.4) and steadily converges to a low value (approx. 0.4), driven by its components.
• Contrastive Loss (red) is the primary driver of feature learning, starting at approx. 1.4 and decreasing to approx. 0.2. This indicates the model is successfully learning to discriminate between instances.
• Alignment Loss (green) starts at approx. 0.6 and smoothly decreases, showing the local models are progressively agreeing with the global model’s representations.
• Orthogonality Loss (purple) starts lowest (approx. 0.4) and decreases to near-zero (approx. 0.05), demonstrating that the framework successfully forces the semantic (zs) and domain (zd) projections into orthogonal subspaces.

7.3 Analysis of ATSA training dynamics and stability

The effectiveness of the ATSA mechanism is validated by interlinking the three plots in Figures 9, 10, and 11. These plots collectively tell the story of the adaptive curriculum.

• (1) Phase 1: Unstable Exploration (Rounds 1-27): As shown in Figure 9, the Model Instability (orange line) starts very high (approx. 0.195), indicating rapid, large changes in the global model weights as clients learn foundational features. Figure 10 confirms the system is in the “Unstable Phase” (red dots). In response, the ATSA controller (Figure 11) keeps the Alignment Weight (λt) low (from 0.1 to 0.7). This “low-pressure” curriculum is crucial: it allows local clients to explore their own data distributions (discover) without being prematurely and aggressively forced to consensus by a naive, under-trained global model.
• (2) Phase 2: Stable Enforcement (Rounds 28-40): At round 28, two things happen simultaneously. The instability score (Figure 9) has decayed to approx. 0.05. This triggers the “Stability Achieved” point, and the system state transitions to the “Stable Phase” (Figure 10, green dots). In direct response, the ATSA controller (Figure 11) begins to “enforce” the consensus. The alignment weight λt now sharply increases, plateauing at a high value. This “high-pressure” phase forces the local models to strongly align their semantic features with the now-mature and stable global model. This step is critical for the model’s final generalization capability, as it ’polishes’ the global representation, forcing all clients to agree on a single, robust semantic feature space that is free of their individual domain biases.

This dynamic, two-phase process confirms that ATSA is functioning as intended from the design, balancing local discovery with global consensus.

7.4 Per-client training dynamics

To analyze the impact of our federated framework on the heterogeneous clients, we tracked the local accuracy of each client’s model throughout the 40 communication rounds. Figure 12 plots these individual learning trajectories, while Table 5 provides a numerical log of their progress at 5-round intervals.

Table 5. Numerical log of simulated local client accuracy at 5-round intervals.

This data corresponds to the plotted trajectories in Figure 12.

These results clearly demonstrate two key findings. First, all clients improve over time, as seen by the general upward trend of all lines in Figure 12. This confirms that participating in the federated network is beneficial for all members, even those with vastly different data, and that knowledge is being successfully aggregated and redistributed.

Second, our framework’s non-independent and identically distributed (non-IID) character is confirmed by the trajectories’ significant degree of variation. Each client learns at a different pace and attains a varied level of final accuracy. Given that their histopathology data closely resembles the final evaluation job, it makes sense that Client 3 (PatchCam) and Client 4 (Breast) start out with great performance and finish with the highest accuracies (0.894 and 0.875, respectively). Client 5 (HAM10000), on the other hand, has more erratic performance, which is consistent with its smaller and more diverse dermatoscopic dataset. This finding highlights the FedAD framework’s resilience in handling substantial statistical variation, allowing every client to participate and profit without jeopardizing the global model’s stability.

It is strictly important to clarify that the data presented in Figure 12 and Table 5 represent a simulated, illustrative analysis of local learning trajectories. These values are not logs obtained from actual, physical client-side training environments. Therefore, they serve solely to illustrate the theoretical convergence dynamics and heterogeneous learning behaviors across different data distributions and are not intended to be used as direct empirical evidence supporting the main generalization and performance conclusions of the FedAD framework.

7.5 Primary generalization performance

Table 6 presents the final performance on the held-out histopathology test set. FedAD (Ours) achieves an F1-score of 0.910, significantly outperforming all other federated methods. Notably, it surpasses the state-of-the-art baselines including FedDG and FedDis. This underscores the benefit of our explicit disentanglement (FedSD) and adaptive curriculum (ATSA) over static, implicit alignment, as proven in the preceding Ablation Study (Section 6). FedAD’s performance even approaches the theoretical (and privacy-violating) “Centralized Supervised” upper bound, despite using no labels.

Table 6. Main performance comparison on the unseen histopathology target domain (Camelyon17 test set).

Bold values indicate the highest performance achieved across the compared methods.

All values following ± in our main results tables represent the standard deviation across 5 independent experimental runs to verify robustness against random weight initialization variations.

To further validate that the core objective of generalization across unseen domains has been met comprehensively, we present experimental results from the independent PadChest dataset in Table 7, which represents a radiological shift entirely unseen by the global model during the federated pre-training phase.

Table 7. Performance comparison on unseen radiological domain (PadChest test set).

Values are Mean ± SD.

Bold values indicate the highest performance achieved across the compared methods.

This primary result is visualized in the bar chart in Figure 13. This chart clearly illustrates the performance gap between the methods. Our proposed method, FedAD (Ours), achieves a substantially higher F1-score than all other federated and centralized self-supervised baselines. The tight error bar on the FedAD result indicates a stable and consistent performance.

To further break down the high-performing F1-score of our FedAD model, we present the confusion matrix in Figure 14 and the Receiver Operating Characteristic (ROC) curve in Figure 15.

The confusion matrix in Figure 14 provides a detailed look at the model’s predictions on the unseen test set. The model correctly identified 5,430 tumor patches (True Positives) and 5,528 normal patches (True Negatives). The errors are well-balanced, with 570 false negatives (tumors missed) and 472 false positives (normal patches misclassified as tumors). These raw numbers directly lead to the high precision (0.920) and recall (0.905) reported in Table 6, demonstrating a strong and balanced classifier.

The ROC curve in Figure 15 reinforces this conclusion. It plots the True Positive Rate against the False Positive Rate across all possible classification thresholds. The curve for our FedAD model (blue line) bows sharply towards the top-left corner, indicating high sensitivity and specificity. The Area Under the Curve (AUC) is an excellent 0.921, which is significantly better than the 0.500 of a random guess (dashed grey line). This high AUC score confirms that the features learned by FedAD are highly discriminative for the downstream task, allowing a simple linear classifier to effectively separate the two classes.

7.6 Computational and communication overhead

A critical requirement for deploying federated systems is understanding scalability constraints. While FedAD achieves state-of-the-art generalization, it incurs approximately a 1.45× increase in wall-clock training time per local epoch compared to standard FedAvg-SSL. This is directly attributable to the dual forward passes required for the actively training student network and the explicitly frozen teacher network (θ~g). Consequently, client-side memory usage increases by roughly 12% to accommodate the retention of θ~g alongside the standard optimizer state. However, the communication payload transmitted across the network remains perfectly identical to standard FedAvg; only the local feature extractor weights θk are synced across the bandwidth constraint, completely omitting the transmission of the projection head ϕk.

8 Challenges and limitations

While our proposed FedAD framework demonstrates a significant step towards robust, privacy-preserving, and self-supervised learning, several challenges remain which have to be addressed, opening up promising avenues for future research.

• Assumption of Linear Disentanglement: While the explicit LOrtho loss utilizes cosine orthogonality, this strictly implies linear independence but does not guarantee true statistical independence or causal separation of domains. Future iterations will explicitly evaluate the disentanglement efficacy using diagnostics such as Mutual Information Neural Estimation (MINE), domain predictability from zs, or reconstruction leakage evaluation from zd.
• Evaluation via Linear Probing: Furthermore, the reliance on frozen-feature linear probing via logistic regression, while standard for SSL evaluation, may not entirely reflect clinical deployment performance where end-to-end task-relevant fine-tuning is widely utilized in downstream hospitals.
• Informative Missingness Bias: The framework currently does not explicitly model unlabeled data distributions. As highlighted by (30), informative missingness in semi-supervised learning can introduce substantial bias if not carefully mitigated.
• Architectural Baselines: Exploring hybrid architectural paradigms like ViT-GCN may yield superior baseline anatomical structure modeling compared to traditional CNN backbones, suggesting our current approach may be architecturally under-explored relative to recent diagnostic advances (31).
• Hyperparameter Sensitivity: FedAD introduces new hyperparameters related to the disentanglement and adaptive curriculum (γ, λmin, λmax, k). As shown in our tuning, these parameters (especially γ) are critical to performance and may require careful, domain-specific tuning.
• Scalability to Non-Image Data: FedAD’s core contrastive engine depends on visual augmentations, in addition to its strength in visual tasks. It would require a completely different augmentation approach and possibly alternative disentanglement objectives to extend this paradigm to additional data modalities, such as 1D time-series (EEG) or 3D volumetric data (fMRI).

9 Conclusion and future work

We have addressed various significant issues that potentially impede the advancement of clinical artificial intelligence (AI) in this study, including limited labeling, data privacy, and domain generalization. We are thrilled to present FedAD, a new federated framework designed to learn robust, domain-invariant representations from decentralized and unlabeled data sources. The two components of the framework—FedSD, which effectively distinguishes semantic content from domain noise, and an Adaptive Training Schedule Algorithm (ATSA), which cleverly directs the training process to produce a stable and broad consensus—are what make it so effective.

Our primary contribution is a comprehensive approach that integrates explicit, non-adversarial disentanglement, an adaptable training schedule, and self-supervised learning. In order to urge local semantic projections to align with the overall consensus while simultaneously encouraging local domain projections to behave in opposition, we have developed a new alignment loss that employs the global model as a teacher. Our results demonstrate how effective this method has been in removing client-specific noise while extracting generalizable knowledge.

When applied to unseen domains, FedAD greatly outperforms current federated approaches such as FedAvg and MOON, given the experiment results, which were validated using extensive k-fold cross-validation. Our ablation tests validate the critical roles that the ATSA curriculum and FedSD disentanglement play in attaining this improved performance. Additionally, a visual study of the training dynamics has demonstrated the stability and intelligent two-phase learning process of the framework.

The FedAD architecture represents a major advancement in the development of useful and reliable medical AI that is resilient, label-efficient, and privacy-preserving. It makes it possible to transition from isolated data silos to a collaborative learning environment, opening the door for technologies that are equitable, efficient, and safe. In the future, semi-supervised extensions where labeled clients can anchor the semantic space and Personalized FedAD (pFedAD) with local heads for few-shot adaptation will be investigated. Furthermore, exploring the creation of more sophisticated non-linear disentanglement techniques and extending the FedAD framework to real multi-modal learning, incorporating knowledge from other sources such as clinical reports and radiology images.

Generative AI statement

The author(s) declared that generative AI was not used in the creation of this manuscript.

Any alternative text (alt text) provided alongside figures in this article has been generated by Frontiers with the support of artificial intelligence and reasonable efforts have been made to ensure accuracy, including review by the authors wherever possible. If you identify any issues, please contact us.

Publisher's note

All claims expressed in this article are solely those of the authors and do not necessarily represent those of their affiliated organizations, or those of the publisher, the editors and the reviewers. Any product that may be evaluated in this article, or claim that may be made by its manufacturer, is not guaranteed or endorsed by the publisher.

Footnotes

Edited by: Tuo Zhang, Northwestern Polytechnical University, China

Reviewed by: Jinran Wu, The University of Queensland, Australia

Fubao Zhu, Zhengzhou University of Light Industry, China