HFed-MIL: Patch Gradient-Based Attention Distillation Federated Learning for Heterogeneous Multi-Site Ovarian Cancer Whole-Slide Image Analysis

Abstract

Ovarian cancer remains a significant global health concern, and its diagnosis heavily relies on whole-slide images (WSIs). Due to their gigapixel spatial resolution, WSIs must be split into patches and are usually modeled via multi-instance learning (MIL). Although previous studies have achieved remarkable performance comparable to that of humans, in clinical practice WSIs are distributed across multiple hospitals with strict privacy restrictions, necessitating secure, efficient, and effective federated MIL. Moreover, heterogeneous data distributions across hospitals lead to model heterogeneity, requiring a framework flexible to both data and model variations. This paper introduces HFed-MIL, a heterogeneous federated MIL framework that leverages gradient-based attention distillation to tackle these challenges. Specifically, we extend the intuition of Grad-CAM to the patch level and propose Patch-CAM, which computes gradient-based attention scores for each patch embedding, enabling structural knowledge distillation without explicit attention modules while minimizing privacy leakage. Beyond conventional logit distillation, we designed a dual-level objective that enforces both class-level and structural-level consistency, preventing the vanishing effect of naive averaging and enhancing the discriminative power and interpretability of the global model. Importantly, Patch-CAM scores provide a balanced solution between privacy, efficiency, and heterogeneity: they contain sufficient information for effective distillation (with minimal membership inference risk, MIA AUC ≈ 0.6) while significantly reducing communication cost (0.32 MB per round), making HFed-MIL practical for real-world federated pathology. Extensive experiments on multiple cancer subtypes and cross-domain datasets (Camelyon16, BreakHis) demonstrate that HFed-MIL achieves state-of-the-art performance with enhanced robustness under heterogeneity conditions. Moreover, the global attention visualizations yield sharper and clinically meaningful heatmaps, offering pathologists transparent insights into model decisions. By jointly balancing privacy, efficiency, and interpretability, HFed-MIL improves the practicality and trustworthiness of deep learning for ovarian cancer WSI analysis, thereby increasing its clinical significance.

1. Introduction

Ovarian cancer remains a significant global health concern affecting millions of women, and the timely and accurate diagnosis of ovarian cancer plays a crucial role in improving patient outcomes and reducing mortality rates. Ovarian cancer diagnosis heavily relies on whole-slide images (WSIs) [1,2,3,4]. Pathologists scan microscope pathological slides to obtain gigapixel WSIs, which require several gigabytes or more to maintain high-resolution features for accurate diagnosis. Even 500 gigapixel WSIs can be as large as the whole ImageNet. Pathologists meticulously analyze slides in clinical practice, which significantly burdens healthcare resources. However, with the surge in big data and artificial intelligence, especially in vast medical image datasets, deep learning algorithms have emerged as pivotal tools to automate pathological image analysis.

Nevertheless, the massive size and spatial consumption of WSIs pose challenges for all deep networks. Typically, existing approaches divide WSIs into patches and model them as multiple instance learning, where each patch is treated as an instance, and the entire WSI is considered a bag because patches within the same WSI are correlated. Common implementations encode the features of each patch using a pre-trained backbone network, then aggregate these patch features to obtain the WSI’s features, followed by classification using a classifier [5,6,7,8,9,10,11,12].

Although existing studies have achieved remarkable performance comparable to humans, they often focus on centralized learning scenarios and overlook that gigapixel WSIs are distributed across multiple hospitals with privacy restrictions in clinical practice. Federated learning [13,14,15,16] was specifically developed to address this issue. This method learns from decentralized data distributed across various sites while ensuring the protection of privacy. Model training and validation are conducted locally at each participating institution in this paradigm. Only essential information is transmitted, thus eliminating the demand for data owners to communicate their private datasets [17,18,19]. Moreover, data distribution variation in multiple hospitals, namely, data heterogeneity, poses additional challenges to federated learning, prompting the development of federated learning methods that transmit different types of information. The types of information transmitted in federated learning can be categorized into parameter transmission, gradient transmission, and proxy dataset prediction knowledge transmission.

Parameter transmission is the most prevalent category. FedAvg [13] employs an approach that averages the model parameters for global model aggregation. However, numerous studies have demonstrated that FedAvg’s performance deteriorates in non-IID settings [20,21,22]. Subsequent works, such as FedProx [23] and Xie et al. [24], have been introduced to address these challenges, incorporating global model parameters into the local optimization process as a regularization term to align local models with the global model.

Gradient transmission can also optimize performance declines under non-IID conditions. Initially, FedSGD was proposed as a baseline for FedAvg, referring to the implementation of previous methods [25]. This approach closely resembles centralized training, eliminating the heterogeneity caused by distributed training, but it leads to high communication costs. Therefore, subsequent methods have adopted strategies that approximate the global gradient, such as Scaffold [20], Fed-grab [26], and AOCC [27].

Proxy dataset prediction knowledge transmission is also named distillation-based federated learning [28]. Unlike the above methods, it does not solely focus on the heterogeneity of data distributions. Because data heterogeneity can ultimately lead to heterogeneity in the preferred models at different sites, methods of transmitting proxy dataset prediction knowledge have been developed. Proxy datasets can be either public or synthetic [29]. The transmission of public dataset prediction results is commonly modeled in the paradigm of distillation learning. Typically, local models are regarded as teacher models. In contrast, the global model is viewed as the student model, which completes knowledge aggregation by distilling knowledge from predictions on public data [30,31]. Since only information independent of the model structure is transmitted, past research has shown that this can adapt to model and data heterogeneity [9,32,33].

Parameter and gradient transmissions could indirectly expose sensitive personal information if leaked or attacked. Merely anonymizing data does not provide sufficient protection [34] as parts of the training data can be reconstructed through gradients [35,36] or model parameters [37,38]. Pathology images are often considered capable of indirect identity recognition, posing a threat to patient privacy and usually subject to legal restrictions. Generally, communicating less data is more conducive to privacy protection, and transmitting knowledge from public datasets is more secure. Moreover, federated learning must also consider the cost of communication, necessitating that the fewer data shared during communication, the better. Combined with the demand for relative flexibility to model heterogeneity, transmitting only the knowledge of predictions on public datasets is the most appropriate choice compared with model parameters or gradient transmission. On the other hand, for the effectiveness of the training model, we need to share as much information as possible to convey knowledge that is as rich as possible to mitigate the performance decline brought by heterogeneity, especially when models process vasts numbers of slices in WSIs. Merely using non-resultant prediction knowledge, namely, pseudo-labels, is insufficient, and we also demand secure, lightweight, structured prediction knowledge as distillation knowledge.

For additional distillation information selection, directly aggregating intermediate parameters or their gradients for distillation is not possible due to model heterogeneity, and sharing raw input data is also prohibited due to privacy concerns. Therefore, we selected patch feature gradient-based attention in the proxy dataset during training, which commonly exists in any MIL framework. We drew inspiration from Grad-CAM [39], measuring the importance of each patch based on the average gradient and then distilling its structural knowledge. This approach was chosen partly because fully exposing gradient information could increase the risk of model leakage and partly because lesion tissue patterns generally differ with healthy tissues, leading to variations in gradient-based patch attention scores. This represents a form of knowledge that generalizes well, even across heterogeneous ovarian cancer WSI data distributions. We name the proposed mechanism Patch-CAM.

To address these issues, this study presents HFed-MIL, which uses federated learning based on gradient attention measured distillation. To be more precise, we start by inducing and formulating a widely used variant of MIL. Next, we distill the structural knowledge of each patch based on its average gradient and use it to measure its importance. This allows us to be relatively flexible when modeling heterogeneity and achieve a balance between security and effectiveness in HFed-MIL. This method enables sufficient information distillation without compromising privacy issues since patch gradient significance scores contain more information than labels but less than features, and they are also more dimensionally adjustable than features. Our results show that the suggested model outperforms the state-of-the-art. HFed-MIL raises the clinical relevance of deep learning by making it more applicable to ovarian cancer and more trustworthy in the medical field.

The key contributions of this work are summarized as follows:

• Patch-CAM-based structural distillation. We extend the intuition of Grad-CAM to the patch level in WSIs and propose Patch-CAM, which computes gradient-based attention scores for each patch embedding. This design enables knowledge distillation in heterogeneous MIL frameworks without requiring explicit attention modules, while avoiding the privacy risks associated with directly sharing gradients or features.
• Dual-level distillation objective (class-level and structural-level). Beyond conventional logit distillation, we introduce structural distillation that enforces consistency of patch-level attention distributions across clients. This prevents the vanishing effect of naive averaging and strengthens the discriminative power and interpretability of the global model.
• Balanced trade-off between privacy, efficiency, and heterogeneity. Patch-CAM scores lie between logits and raw features: they provide sufficient information for effective distillation while leaking minimal private data (MIA AUC ≈ 0.6, close to random guessing). In addition, HFed-MIL reduces communication cost to 0.32 MB per round, which is orders of magnitude smaller than parameter or gradient sharing, making it practical for real-world federated pathology.
• Enhanced interpretability and generalization in heterogeneous federated WSI analysis. Extensive experiments on multiple cancer subtypes and cross-domain datasets (Camelyon16, BreakHis) demonstrate that HFed-MIL yields more robust performance across heterogeneous models. Moreover, the global attention visualizations show sharper and clinically meaningful heatmaps, offering pathologists transparent insights into model decisions.

The rest of the study is outlined as follows: Section 2 discusses literature work followed by Section 3, which discusses the dataset and the proposed work’s methodology together. In Section 4 experiments are presented. Lastly, Section 5 concludes the current study and discusses limitations with possible future directions.

2. Related Work

In this section, we briefly discuss current advancements in ovarian cancer, followed by a discussion on attention mechanisms, and finally, we present challenges in multi-site data distribution.

2.1. Advancements in Ovarian Cancer

Ovarian cancer remains a significant health concern, with its late-stage diagnosis contributing to high mortality due to complex subtype classification [1,40,41]. Traditional diagnostic methods, including histopathological examination of whole-slide images (WSIs), have been enhanced by deep learning techniques, allowing for improved accuracy and efficiency in cancer detection. Earlier works demonstrate the potential of CNNs in histopathological image classification—subsequent studies, such as those by [2,3,42,43,44,45], have focused on the application of deep learning frameworks to the analysis of ovarian cancer, achieving remarkable predictive performances.

However, despite these advancements, several critical gaps remain unaddressed. Notably, the heterogeneity of the data across different clinical sites poses challenges for the generalizability of machine learning models [46,47]. A few studies have highlighted that variability in image quality and annotation standards across institutions can lead to biased model performance [48]. Furthermore, privacy concerns limit the sharing of sensitive patient data, necessitating innovative approaches that can leverage decentralized data for training while adhering to ethical guidelines [49]. HFed-MIL aims to bridge these gaps by utilizing federal learning to facilitate collaborative model training without compromising patient privacy.

2.2. Attention Mechanisms in Deep Learning

Attention mechanisms have emerged as a powerful tool in deep learning, particularly for enhancing model interpretability and performance in complex tasks such as image segmentation [50], breast cancer diagnosis [51], image captioning [52], generative modeling [52], remote sensing [53], foreign object detection [54], and facial micro-expression recognition [55]. These methods showcase their versatility and effectiveness in handling diverse challenges in computer vision. Further advancements in attention in WSI-based works have yielded significant results, demonstrating applicability across multiple domains, including ovarian cancer. Vaswani et al. [56] initially proposed the Transformer architecture, fundamentally altering how attention is integrated into neural networks. The dynamics of attention allow the model to focus on relevant parts of an image, enhancing its ability to capture intricate patterns.

Notably, Vision Transformer (ViT) [57] and its hierarchical variant, Swin Transformer [58], have demonstrated strong capability in modeling global and multi-scale context. Transformer-based attention techniques have been adapted for WSI analysis, offering superior modeling of long-range dependencies [7,59]. Furthermore, cross-attention networks have been developed to leverage inter-slide relationships, particularly beneficial for rare ovarian cancer subtypes [60,61]. Gupta et al. [62] integrated gradient-based visualization, enabling efficient training of very large images and potentially advancing pathologists’ trust in model predictions. Beyond canonical Transformer designs, hybrid architectures combining CNN and Transformer modules have also shown strong potential. For instance, Huo et al. [53] proposed FCIHMRT, which employs a ViT branch to capture global dependencies and a Res2Net branch for multi-scale local feature representation, merging their outputs to enhance remote sensing scene classification.

Multiple instance learning (MIL) is central to WSI analysis, allowing models to learn from bags of image patches with only slide-level labels. Incorporating attention within MIL enhances focus on diagnostically relevant regions, improving interpretability and performance [5,6]. Recent frameworks integrate probabilistic attention [63], multi-head attention [8], and cascaded attention [64] to refine feature extraction and localization. However, most of these approaches neglect the federated setting, particularly the heterogeneous nature of multi-institutional WSI data, where variations in staining protocols, scanner types, and population distributions pose significant challenges. In such scenarios, striking a balance between effective knowledge sharing and strict privacy preservation remains a significant challenge. Moreover, gradients derived from attention weights can inadvertently expose sensitive patient data [35,36]. To address these issues, HFed-MIL introduces a privacy-preserving patch gradient attention distillation framework tailored for heterogeneous federated MIL, ensuring that the benefits of attention mechanisms are retained without compromising robustness across sites or patient confidentiality.

2.3. Federated Learning with WSIs

Whole-slide images (WSIs) have revolutionized digital pathology by enabling high-resolution visualization of tissue samples. However, due to their inherent complexity, direct centralized learning on WSIs is computationally expensive and raises substantial data privacy concerns, particularly in medical applications, where patient confidentiality is paramount. In response, federated learning (FL) has gained attention in computational pathology as a privacy-preserving approach to train models on distributed data points without data sharing [16,65].

While existing works have shown significant results within the intersection of federated learning with WSIs, this intersection introduces several unique challenges, which are summarized as follows:

• Meaningful Gigapixel Analysis: WSIs require patch-wise modeling and MIL-based aggregation, which complicates communication and synchronization in federated settings.
• Data and Model Heterogeneity: Differences in data distributions, annotation quality, and model architecture across hospitals cause a decrease in performance in conventional FL pipelines, further discussed in the subsequent Section 2.4.
• Privacy vs. Performance trade-off: Existing methods often compromise diagnostic performance for stronger privacy or vice versa.

To address these issues, we propose HFed-MIL, a novel gradient attention-based knowledge distillation framework for federated MIL in ovarian cancer WSI analysis. Instead of relying solely on logits or weight updates, HFed-MIL extracts and transfers gradient-derived patch attention scores, which provide richer semantic context than predictions but are less privacy-sensitive than raw features. This design balances model interpretability, data privacy, and robustness to non-IID distributions—making HFed-MIL particularly suited for real-world multi-institutional pathology workflows.

2.4. Challenges in Multi-Site Data Distribution

The deployment of machine learning models in multi-site healthcare settings presents unique challenges, particularly concerning data distribution. As discussed above, federated learning has emerged as a promising solution, allowing models to be trained across decentralized data sources while maintaining privacy. However, the heterogeneous nature of medical data—characterized by variations in imaging protocols, demographic distributions, and annotation practices—further complicates the training process. Several studies have outlined these difficulties, emphasizing the challenges of achieving significant results and model convergence under such conditions.

Moreover, existing federated learning frameworks often struggle with knowledge distribution, where the goal is to effectively convey insights from diverse data sources. Techniques such as federated averaging have been widely adopted; however, they may not adequately address the complexities of multi-instance learning scenarios inherent in whole-slide image (WSI) analysis. Our method seeks to overcome these limitations by introducing a standardized approach that not only accommodates the variability of multi-site data but also enhances the efficiency of knowledge transfer across federated networks, facilitating broader insights.

In conclusion, while significant advancements have been made in ovarian cancer diagnosis, the integration of patch gradient-based attention mechanisms with federated learning remains an underexplored area. This innovative approach holistically addresses the challenges of privacy, data heterogeneity, and efficient learning. Our method represents a substantial enhancement in this domain, promising to improve the accuracy and applicability of cancer diagnostics across diverse clinical environments.

3. Dataset and Methodology

3.1. Dataset

The study compiled a dataset from the UBC-Ovarian Cancer Challenge, (https://www.kaggle.com/competitions/UBC-OCEAN/overview, (accessed on 14 May 2025)), which included histopathological images categorized into five ovarian cancer subtypes: Clear Cell Carcinoma (CC), Mucinous Carcinoma (MC), Low-Grade Serous Carcinoma (LGSC), High-Grade Serous Carcinoma (HGSC), and Endometrioid Carcinoma (EC). To handle the large WSIs, we performed downscaled patch sampling to reduce the computational load. We chose 50 representative tiles per WSI, each sized $512\times 512$ pixels. The selection focused on diagnostically significant regions, eliminating extraneous background. The preprocessing stage involved reducing noise and normalizing the data to improve the model’s accuracy without sacrificing efficiency. Figure 1 shows sample WSIs from UBC-Ovarian Cancer, while Figure 2 depicts the distribution of tissue types, and Figure 3 emphasizes the need to effectively preprocess WSI into smaller patches as it shows WSI gigapixel information.

Furthermore, this dataset poses unique challenges in the context of federated learning compared to other WSI cancer datasets. The limited availability of ovarian cancer cases results in a smaller dataset, complicating model training. Additionally, class imbalance among different subtypes of ovarian cancer further hinders the performance of FL models. Addressing these specific challenges is crucial for developing practical machine learning approaches in ovarian cancer diagnosis and treatment. The variability in imaging techniques across institutions can lead to inconsistencies in slide quality, while the complexity of histopathological features necessitates the use of advanced feature extraction methods.

3.2. Local WSI MIL Modeling

Let W represent a whole-slide image (WSI), which is decomposed into a collection of patches $\{P_1,P_2,\dots ,P_n\}$. Each patch $P_i$ is processed through a backbone neural network f (e.g., ResNet [66], ViT [67]) that serves as an encoder. This network maps each patch to a feature vector in a high-dimensional space:

$${\mathbf{x}}_i=f(P_i;\theta )$$

(1)

where $\theta$ denotes the parameters of the backbone network, $x_i$ is the feature vector with any shape or dimension with the total size of d for the i-th patch.

The feature vectors $\{{\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_n\}$ in the W are then aggregated using an aggregation operator g, which could be self-attention, max pooling, or so on, to form a single feature vector $\mathbf{v}$ that represents the entire WSI [68]:

$$\mathbf{v}=g({\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_n;\varphi )$$

(2)

Here, $\varphi$ represents the parameters of the aggregation mechanism, and $\mathbf{v}\in {\mathbb{R}}^d$ is the aggregated feature vector for the WSI. And the g can be anonymous to the server.

Finally, a classifier h, parameterized by $\psi$, takes the aggregated feature vector $\mathbf{v}$ and produces a prediction y for the WSI:

$$z=h(\mathbf{v};\psi )$$

(3)

where z is the predicted logits for the WSI, which have not been normalized with the softmax function.

$$y=\mathrm{softmax}\left(z\right)$$

(4)

where $z_c$ is the c-th dimension of z, and predicted probability y is the logits z normalized by the softmax function.

Given the prediction y for the WSI, we define the target label as t, which is a categorical label representing the true class of the WSI. The cross-entropy loss function L is used to quantify the error between the predicted probabilities $\widehat{y}$ from the classifier and the true distribution t. The cross-entropy loss for a single instance is given by

$$L_{cls}=L(y,t)=-\sum _{c=1}^C\mathbb{1}\{c=t\}log\left(y_c\right)$$

(5)

where C is the number of classes, $\mathbb{1}\{y=t_c\}$ is an indicator function that equals 1 if the predicted class y is equal to the true class $t_c$ and 0 otherwise, and ${\widehat{y}}_c$ is the predicted probability of the class c as output by the softmax layer of the classifier.

3.3. Label Distillation Fed-MIL

In a federated learning setting, consider several local models, each hosted on a client, denoted by ${\mathcal{M}}^{\left(i\right)}$ where $i=1,2,\dots ,n$. These models, probability differing in architecture, generate logits ${\mathbf{z}}^{\left(i\right)}$ for given inputs. The global model, hosted on a server, is denoted by $\mathcal{G}$.

The distillation process involves treating the local models as teacher models. The predicted logits from each local model are averaged to form pseudo-labels:

$$\overline{\mathbf{z}}=\sum _{i=1}^n{\displaystyle \frac{|D^{\left(i\right)}|}{\left|D\right|}}{\mathbf{z}}^{\left(i\right)}$$

(6)

$$\left|D\right|=\sum _{i=1}^n\left|D^{\left(i\right)}\right|$$

(7)

Here, $D^{\left(i\right)}$ denotes the local dataset on client i, and $|D^{\left(i\right)}|$ is the number of WSIs (samples) in that dataset. The term $\overline{\mathbf{z}}$ represents the aggregated logits, which serve as pseudo-labels for supervising the global model at the class prediction level.

The global model $\mathcal{G}$ learns from these pseudo-labels through distillation. The loss function used is the Kullback–Leibler (KL) divergence, which measures the difference between the predictive distribution of the global model and the distilled pseudo-labels:

$$L_{label}=\mathrm{KL}(\mathrm{softmax}({\mathbf{z}}^{\mathcal{G}}/\tau )\left|\right|\mathrm{softmax}(\overline{\mathbf{z}}/\tau ))$$

(8)

Here, ${\mathbf{z}}^{\mathcal{G}}$ represents the logits output by the global model, and $\mathrm{softmax}(·)$ denotes the temperature-scaled softmax with parameter $\tau$.

The objective is to minimize this loss, aligning the global model’s predictions with the pseudo-labels derived from the local models, thus leveraging the collective knowledge of all participating models while maintaining data privacy.

3.4. Patch-CAM-Based Distillation

For distillation information selection, directly aggregating intermediate parameters or their gradients is infeasible in federated settings due to model heterogeneity, and sharing raw input data is prohibited due to privacy concerns. Whole-slide images (WSIs) are extremely large and must be divided into hundreds of patches, most of which do not contain lesions. This makes attention mechanisms particularly important as they enable the model to highlight the small fraction of diagnostically relevant patches. However, not all WSI models include explicit attention modules. To ensure generality across heterogeneous clients, we draw inspiration from Grad-CAM [39] and extend it to the patch level, leading to our Patch-CAM design.

Specifically, we compute gradient-based importance scores for each patch feature in the proxy dataset during training. The importance score is derived from the gradients of the maximum predicted logit ${\widehat{z}}_c$ with respect to each feature component and then averaged across channels to obtain a stable estimate of patch relevance. This adaptation follows the spirit of Grad-CAM but is redefined at the patch level, allowing for universal application to any MIL framework. Moreover, compared with exposing full gradients, sharing only Patch-CAM scores is less privacy-sensitive since they represent coarse structural knowledge rather than fine-grained feature information. Importantly, lesion tissue typically induces different gradient responses compared with normal tissue, which makes these scores both discriminative and generalizable across heterogeneous data distributions. In addition, the Patch-CAM scores provide interpretable heatmaps that help clinicians understand which regions influenced the model’s decision, further enhancing transparency.

Formally, let $A_{ij}$ denote the j-th feature component of the i-th patch $x_i$, where each patch has a feature dimension d, and the total number of patches in the slide is $\left|W\right|$. The patch importance weight $\alpha$ is obtained by averaging the gradient of the maximum predicted logit ${\widehat{z}}_c$ with respect to all feature components. Following the Grad-CAM logic, we define the Patch-CAM score as

$$\mathrm{score}\left(P\right)=\sum _i^{\left|W\right|}\sum _j^d{\displaystyle \frac{\partial {\widehat{z}}_c}{\partial A_{ij}}}{A}_{ij},$$

(9)

where the multiplication between patch features and their corresponding gradient-derived weights highlights the relative importance of each patch in the global decision.

Every WSI can be regarded as the collection of patch features x. We compute a normalized importance score for each patch and treat it as additional structural distillation knowledge:

$$\mathbf{S}\left(W\right)=\mathrm{softmax}(x_1,x_2,\dots ,x_n)$$

(10)

The structural distillation process aggregates patch-level attention distributions from clients:

$$\overline{\mathbf{S}}=\sum _{i=1}^n{\displaystyle \frac{|D^{\left(i\right)}|}{\left|D\right|}}{\mathbf{S}}^{\left(i\right)}$$

(11)

$$L_{struct}=\mathrm{KL}(\mathbf{S}\left(W^G\right)\left|\right|\overline{\mathbf{S}})$$

(12)

Here, $\mathbf{S}\left(W\right)$ denotes the attention-based importance distribution over patches within a WSI, and ${\mathbf{S}}^{\left(i\right)}$ is the corresponding distribution computed from client i. Similar to the label distillation, the weighting factor $|D^{\left(i\right)}|$ ensures that clients with larger datasets contribute more. Unlike $L_{\mathrm{label}}$, which aligns class-level predictive distributions, $L_{\mathrm{struct}}$ aligns the patch-level attention distributions, thereby transferring structural knowledge across heterogeneous clients.

The final loss is

$$L_{global}=L_{label}+L_{struct}$$

(13)

3.5. Model Architecture Diagram

The architecture diagram as shown in Figure 4 depicts the proposed HFed-MIL, a specialized federated learning framework designed for ovarian cancer diagnosis using whole-slide images (WSIs) from whole-slide images (WSIs) across multiple medical sites. The procedure begins by converting whole-slide images (WSIs) into smaller sections called patches at each location. Data augmentation techniques are then employed to improve the resilience of the model. Each site calculates attention scores based on gradients for its patches in a way that protects privacy. These scores are then used to inform the global model’s process of distilling knowledge. This allows the central model to get valuable information without directly accessing sensitive data, maintaining privacy. The approach achieves a compromise between sharing information and maintaining anonymity. This allows the federated system to effectively manage various data distributions and model preferences while delivering state-of-the-art medical picture analysis performance. The corresponding algorithm is shown in Algorithm 1.

Algorithm 1: Algorithm for Proposed HFed-MIL

4. Experiments and Discussion

In this section, we present the results of our experiments, focusing on the impact of model heterogeneity on the performance of different federated learning methods for ovarian cancer analysis.

4.1. Baseline and Parameter Setting

We compare HFed-MIL with other state-of-the-art FL approaches, including FedAvg [21], FedDF [28], FedMD [30], and FedRAD [32]. FedAvg aggregates local models by averaging their parameters. FedDF distills the ensemble knowledge of all client models into a single global model. FedMD resembles FedDF but maintains no global model; each client independently trains its local model using aggregated soft predictions from the server. FedRAD improves robustness by using the median aggregation of client predictions to filter out low-quality knowledge. Finally, HFed-MIL refines the global model by integrating both class-level and structural knowledge from client models.

The following hyperparameters are used throughout all experiments without special mention:

• Optimizer: Adam with learning rate $5\times {10}^{-4}$;
• Number of communication rounds: 20;
• Distillation temperature: $\tau =1$;
• Patch sampling: 100 patches per WSI, each of size $512\times 512$;
• Number of clients: 4;
• Batch size: 1 WSI bag per iteration (due to memory constraints with large patch sets);
• GPU: 4× Titan RTX.

This setting serves as the default configuration across all compared methods.

4.2. Experiment Type—A: Comparative Analysis

In this subsection, we compare the performance of different federated learning methods in the context of ovarian cancer analysis. The following methods were evaluated: Centralized, FedAvg, FedDF, FedMD, FedRAD, and our proposed method, HFed-MIL, as shown in Table 1 for four clients. For the heterogeneous model setting, we apply ResNet, ViT, EffcientViT [69], and ConVit [70], and for the homogeneous model setting, we apply four ResNet models. Following the non-IID setting of [71], we apply Dirichlet distribution to split the local datasets parameterized by $\alpha$. The homogeneous data setting $\alpha$ is 1, and the heterogeneous data setting $\alpha$ is 10; further relevant experimental parameters are discussed earlier in Section 4.1.

Starting with the centralized approach, we achieved an accuracy of 76.7% with a loss of 0.80. This serves as an upper bound for comparison with the federated learning methods, which denote performance without any heterogeneity. The experimental results indicate significant variability in performance across different federated learning methods when dealing with heterogeneous data and models. Specifically, FedAvg [21] exhibited a substantial decline in performance with heterogeneous data, recording accuracy of only 37.2%, and it does not support heterogeneous model scenarios, as shown by the absence of results in these cases. This highlights its limitations in environments with diverse data distributions.

Table 1. Comparison of federated learning methods on different models and data types. The relevant loss and accuracy curves of the heterogeneous model and data are presented in Figure 5 and Figure 6, respectively. Abbreviations: Homo = Homogeneous, Heter = Heterogeneous.

Method	Model	Data	Accuracy	Loss	Sensitivity	Specificity
Centralized	-	-	75.7%	0.70	76.2%	75.0%
FedAvg [21]	Homo	Homo	74.0%	0.74	73.5%	74.2%
		Heter	47.2%	1.54	46.0%	48.0%
	Heter	Homo	-	-	-	-
		Heter	-	-	-	-
FedDF [28]	Homo	Homo	56.5%	1.07	57.0%	56.0%
		Heter	42.0%	1.33	41.5%	42.2%
	Heter	Homo	61.5%	1.08	62.0%	61.0%
		Heter	50.1%	1.32	50.5%	49.8%
FedMD [30]	Homo	Homo	54.7%	1.14	55.0%	54.5%
		Heter	47.1%	1.37	47.5%	46.8%
	Heter	Homo	55.7%	1.25	56.0%	55.2%
		Heter	49.3%	1.31	49.0%	49.5%
FedRAD [32]	Homo	Homo	48.3%	1.31	48.5%	48.0%
		Heter	43.5%	1.45	44.0%	43.0%
	Heter	Homo	54.0%	1.23	54.5%	53.8%
		Heter	50.0%	1.20	50.5%	49.8%
HFed-MIL	Homo	Homo	66.1%	0.91	66.5%	65.8%
		Heter	57.3%	1.03	57.0%	57.5%
	Heter	Homo	57.2%	0.88	57.5%	57.0%
		Heter	53.9%	0.95	54.0%	53.8%

Table 1. Comparison of federated learning methods on different models and data types. The relevant loss and accuracy curves of the heterogeneous model and data are presented in Figure 5 and Figure 6, respectively. Abbreviations: Homo = Homogeneous, Heter = Heterogeneous.

Method	Model	Data	Accuracy	Loss	Sensitivity	Specificity
Centralized	-	-	75.7%	0.70	76.2%	75.0%
FedAvg [21]	Homo	Homo	74.0%	0.74	73.5%	74.2%
		Heter	47.2%	1.54	46.0%	48.0%
	Heter	Homo	-	-	-	-
		Heter	-	-	-	-
FedDF [28]	Homo	Homo	56.5%	1.07	57.0%	56.0%
		Heter	42.0%	1.33	41.5%	42.2%
	Heter	Homo	61.5%	1.08	62.0%	61.0%
		Heter	50.1%	1.32	50.5%	49.8%
FedMD [30]	Homo	Homo	54.7%	1.14	55.0%	54.5%
		Heter	47.1%	1.37	47.5%	46.8%
	Heter	Homo	55.7%	1.25	56.0%	55.2%
		Heter	49.3%	1.31	49.0%	49.5%
FedRAD [32]	Homo	Homo	48.3%	1.31	48.5%	48.0%
		Heter	43.5%	1.45	44.0%	43.0%
	Heter	Homo	54.0%	1.23	54.5%	53.8%
		Heter	50.0%	1.20	50.5%	49.8%
HFed-MIL	Homo	Homo	66.1%	0.91	66.5%	65.8%
		Heter	57.3%	1.03	57.0%	57.5%
	Heter	Homo	57.2%	0.88	57.5%	57.0%
		Heter	53.9%	0.95	54.0%	53.8%

Figure 5. Comparative plots of accuracy curves for different federated aggregation methods. Each colored curve corresponds to one client.

Figure 5. Comparative plots of accuracy curves for different federated aggregation methods. Each colored curve corresponds to one client.

Figure 6. Comparative plots of training loss curves for different federated aggregation methods. Each colored curve corresponds to one client.

Figure 6. Comparative plots of training loss curves for different federated aggregation methods. Each colored curve corresponds to one client.

In contrast, methods employing knowledge distillation, such as FedDF [28] and our proposed HFed-MIL, demonstrated more robust performance regarding data and model heterogeneity. Notably, HFed-MIL outperformed other methods in heterogeneous settings, achieving an accuracy of 63.9% on heterogeneous data within heterogeneous models. This superior performance can be attributed to effectively utilizing diverse knowledge inherent in heterogeneous models, which is crucial for distillation-based approaches. The distillation process benefits significantly from these models’ rich and varied information, enhancing learning efficacy and model robustness across diverse federated networks.

Finally, our proposed method, HFed-MIL, outperformed all other federated learning methods regarding accuracy and loss. When a homogeneous model with homogeneous data was used, HFed-MIL achieved an accuracy of 66.1% with a loss of 0.91, surpassing the performance of other methods. However, when dealing with heterogeneous data, HFed-MIL, like other methods, faced challenges and experienced a drop in performance, achieving an accuracy of 63.9% with a loss of 0.95. Nevertheless, HFed-MIL demonstrated superior performance when employing a heterogeneous model with homogeneous data, achieving an accuracy of 67.2% with a loss of 0.88.

In the context of model heterogeneity, we emphasize that model diversity contributes to improved performance in our distillation-based models. In contrast, the performance of FedAvg, which uses “transmit” parameters, suffers when models are heterogeneous. Notably, incorporating diverse knowledge becomes crucial for models employing distillation and transfer learning with proxy datasets, particularly when dealing with heterogeneous data. This diversity enhances the overall performance of our proposed approach.

4.3. Experiment Type—B: Global Model Attention Visualization

To further investigate the interpretability of the proposed method, we visualize the attention maps produced by the global model. The motivation is to examine whether the gradient-based attention design can highlight discriminative regions in whole-slide images (WSIs) and provide consistent structural cues across heterogeneous clients. By comparing the visualization results with those from baseline methods, we aim to demonstrate that HFed-MIL not only achieves competitive accuracy but also yields more coherent and semantically meaningful attention patterns, thereby enhancing the transparency of the federated aggregation process.

As illustrated in Figure 7, the top row corresponds to the proposed HFed-MIL, the middle row shows the results of FedDF, and the bottom row presents FedRAD. HFed-MIL produces sharper and more distinctive attention patterns that emphasize discriminative lesion regions, thereby mitigating the vanishing effect often caused by averaging attention across heterogeneous clients. In contrast, FedDF yields smoother but less informative responses due to the averaging of knowledge distillation. At the same time, FedRAD exhibits an even more uniform distribution of attention, which may further dilute critical diagnostic cues. These findings indicate that HFed-MIL not only improves predictive performance but also enhances interpretability by providing pathologists with more precise and clinically meaningful visual evidence of the model’s decision rationale.

4.4. Experiment Type—C: Ablation Study of Heterogeneity

As shown in Figure 8, We tested the three methods previously mentioned—FedDF, FedRAD, and HFed-MIL—with the same settings as the heterogeneous models discussed earlier, except for $\alpha$ variation, to validate the model performance with various degrees of heterogeneity. The experiments indicate that all methods experience a decline in performance as heterogeneity increases. However, HFed-MIL consistently achieves the best results regardless of the level of heterogeneity.

4.5. Experiment Type—D: Ablation Study of Client Number

Figure 9 presents the ablation study investigating how the number of clients affects model performance under heterogeneous conditions. As the number of clients increases with the same settings as the heterogeneous models discussed earlier except the client number, the data becomes more fragmented, and the heterogeneity among clients is exacerbated, leading to a general decline in accuracy for all methods. Despite this trend, HFed-MIL consistently outperforms both FedDF and FedRAD across all settings, demonstrating superior robustness and generalization capability in the presence of increasing client numbers.

4.6. Experiment Type—E: Generalization Test with Proxy Dataset from Another Domain

To further evaluate cross-domain generalization, we constructed proxy datasets by mixing a small portion of in-domain data with external sources. Specifically, in addition to $10\%$ of the original training data, we consider two external domains: Camelyon16 and BreakHis. Three proxy compositions are tested with the same total size N = 100: (i) Camelyon-only, where the remaining samples are from Camelyon16; (ii) BreakHis-only, where the remaining samples are from BreakHis; and (iii) both, where the remaining samples are equally split between Camelyon16 and BreakHis. In all cases, label spaces are aligned with the target task by mapping external samples to the closest categories. We keep the backbone (ResNet-18), client number (4), and all optimization hyperparameters identical to previous experiments for a fair comparison. Distillation-based methods use the proxy only at the server side, while FedAvg does not rely on a proxy. The experiement result is shown in

Table 2. Cross-domain proxy generalization (10% in-domain + external). Metrics are percentages (%).

Proxy Composition	Method	Accuracy	Sensitivity	Specificity
Camelyon16	FedDF	31.0	30.2	31.8
	FedMD	30.2	29.5	31.0
	FedRAD	31.8	32.5	31.1
	HFed-MIL	34.1	34.0	33.5
BreakHis	FedDF	25.0	26.0	24.3
	FedMD	29.0	28.2	29.8
	FedRAD	26.3	27.0	25.5
	HFed-MIL	32.0	32.8	31.2
Both (Camelyon16 + BreakHis)	FedDF	28.2	27.6	28.9
	FedMD	29.5	29.0	30.1
	FedRAD	29.0	29.1	28.8
	HFed-MIL	35.0	35.6	34.3

.

This setup simulates realistic scenarios where publicly available proxy data may derive from different sources or acquisition conditions than the clients’ local data. We expect Camelyon-only to benefit tumor-region localization due to structural similarity, while BreakHis-only emphasizes appearance shift at the patch level. The mixed proxy stabilizes performance under distribution shift conditions, and HFed-MIL is expected to retain stronger robustness by leveraging structural attention in addition to class-level logits.

4.7. Experiment Type—F: Privacy and Security Analysis

To substantiate our claim that gradient-based attention scores are less privacy-sensitive than raw features, we conducted a dedicated privacy analysis within the federated setting. The objective is to examine whether an adversary can exploit the information shared in HFed-MIL to infer private data from local clients. We compared four types of shared information: (1) raw gradients, (2) class logits, (3) gradient-based attention scores, and (4) our proposed HFed-MIL, which combines logits with attention.

Among various privacy attacks, we focus on the Membership Inference Attack (MIA) [72], where an adversary attempts to determine whether a given sample belongs to the training set. MIA is the most broadly applicable since it can exploit any type of shared signal. Other attack families such as gradient inversion [35,36] or GAN-based reconstruction are not applicable in our setting: gradient inversion requires direct access to raw gradients, which HFed-MIL never exposes, while GAN-based conditional generation could at best produce vague class-level prototypes rather than patient-specific WSI patches. Therefore, we only evaluated MIA, which remains the most realistic threat with our federated design.

To evaluate MIA, we followed the standard shadow-model setup. For each type of shared information, we trained a binary attack classifier using shadow clients with disjoint datasets, where positive examples correspond to training samples and negative examples to held-out samples. The attack model was then applied to target clients, and MIA performance was reported in terms of AUC. We do not report results for reconstruction-style attacks since such methods require raw gradients or embeddings that are not shared in our design.

As shown in

Table 3. Membership inference performance for different shared information types. Higher AUC indicates greater privacy leakage. An AUC of 0.5 corresponds to random guessing, meaning no privacy leakage. Reconstruction-style attacks are infeasible without direct access to the gradient.

Shared Information	MIA AUC
Gradients	0.85
Patch Embedding Feature	0.80
Class Logits	0.55
Gradient-Based Attention	0.60
HFed-MIL (Logits + Attention)	0.62

, exposing raw gradients leads to severe privacy leakage with MIA AUC∼0.85. Sharing logits alone reduces leakage to nearly random guessing (AUC∼0.55), while attention scores exhibit slightly higher vulnerability (AUC∼0.60). HFed-MIL, which combines logits and attention, yields AUC∼0.62. Although this is marginally higher than either signal alone, it remains close to the random baseline of 0.5 and substantially lower than the gradient case. This indicates that HFed-MIL achieves a practical balance between utility and privacy preservation: it improves predictive performance and interpretability without incurring significant privacy risks.

4.8. Experiment Type—G: Computation and Communication Cost Analysis

To ensure applicability across heterogeneous local model architectures, HFed-MIL includes gradient-based attention scores. While this design provides a unified way of distilling structural knowledge across clients, it may introduce additional computational overhead. Therefore, we conducted an experiment to quantify the computation and communication cost, where all methods were evaluated with the same setting using four clients with ResNet-18 as the backbone model. For distillation-based methods, a proxy dataset of size 100 was used, while for FedAvg, each client was assigned 100 local samples.

The results in

Table 4. Comparison of computation and communication cost across different federated aggregation methods. Reported values are averaged over 20 communication rounds.

Method	Training Time/Round (s)	Communication Cost	GPU Memory (MB)
FedAvg	20.4	0.37 GB	3100
FedDF	28.4	7.8 KB	3100
FedMD	29.1	15.6 KB	3250
FedRAD	28.9	7.8 KB	3180
HFed-MIL	32.5	0.32 MB	3450

show that HFed-MIL incurs only a modest additional overhead, with about $15\%$ higher training time and around $10\%$ higher GPU memory usage per round compared with other distillation-based methods (FedDF, FedMD, and FedRAD). This gap mainly derives from the computation of gradient-based attention scores, which requires intermediate gradients to be retained. Nevertheless, since the global model can also be equipped with an explicit attention head to directly align with the distilled structural targets, the overhead can be bounded without doubling the training cost.

Regarding communication, all distillation-based methods drastically reduce the cost compared with FedAvg since they exchange only predictions rather than full model parameters. Among them, HFed-MIL requires larger communication than FedDF/FedMD/FedRAD because in addition to class-level logits, it uploads patch-level attention scores. Concretely, for each client with N proxy samples (each containing P patches and C classes), FedDF/FedRAD only transmit $N\times C$ logits per round, whereas HFed-MIL transmits $(N\times C)+(N\times P)$ values (i.e., one patch attention vector per sample). Aggregated across K clients, the server receives $K\times \left[\right(N\times C)+(N\times P\left)\right]$ values each round. This explains the ∼0.32 MB per client per round cost of HFed-MIL, which still remains orders of magnitude smaller than FedAvg (∼0.37 GB).

Overall, the results demonstrate that the proposed gradient-based attention design is computationally feasible in practice, while providing additional robustness and interpretability.

5. Conclusions

Ovarian cancer remains a significant global health concern, and its diagnosis heavily relies on whole-slide images (WSIs). Due to WSI gigapixel spatial consumption and medical privacy restrictions, the Fed-MIL framework is highly appropriate. This study summarizes and categorizes the current federated learning into three types based on the information transmitted: parameters, gradients, and proxy dataset prediction knowledge. However, due to model heterogeneity, we can only adopt the transmission of proxy dataset prediction knowledge, namely, the distillation-based method. Furthermore, the performance of distillation-based federated learning highly depends on the effectiveness of the distilled information. For the substantial spatial consumption of WSIs, in addition to general unstructured result information, we also require attention to information that can be measured within heterogeneous models to ensure performance. This paper introduces HFed-MIL, which applies gradient attention measured-distillation-based federated learning to tackle these challenges. Specifically, we first induce and formulate a commonly used form of MIL. Then, to enable relative flexibility to model heterogeneity and reach the balance between security and effectiveness in HFed-MIL, we not only distill the information of prediction results but also measure the importance of each patch based on the average gradient and then distill its structural knowledge. Since the patch gradient importance scores contain more information than labels but less than features and are more dimensionally flexible than features, this method allows for sufficient information distillation without compromising privacy risks. Our experiments demonstrate that the proposed model achieves state-of-the-art performance. HFed-MIL improves the practicality of deep learning in ovarian cancer and further enhances the trustworthiness of medical deep learning, thereby increasing its clinical significance.

Although HFed-MIL achieves notable performance improvements on ovarian cancer WSIs under data and model heterogeneity conditions, this study’s limitation is that the shared proxy dataset must have uniformly sampled patch positions and sizes, which reduces the framework’s flexibility. We plan to improve on this aspect in future work. Furthermore, comprehensive evaluation on an external multi-institutional dataset is necessary to fully validate HFed-MIL’s generalization ability and robustness in practical clinical scenarios.