Evidential Federated Learning for Skin Lesion Image Classification

TL;DR

FedEvPrompt offers a promising approach for federated learning, effectively addressing challenges such as data heterogeneity, imbalance, privacy preservation, and knowledge sharing, and is optimized within a round-based learning paradigm.

Abstract

We introduce FedEvPrompt, a federated learning approach that integrates principles of evidential deep learning, prompt tuning, and knowledge distillation for distributed skin lesion classification. FedEvPrompt leverages two sets of prompts: b-prompts (for low-level basic visual knowledge) and t-prompts (for task-specific knowledge) prepended to frozen pre-trained Vision Transformer (ViT) models trained in an evidential learning framework to maximize class evidences. Crucially, knowledge sharing across federation clients is achieved only through knowledge distillation on attention maps generated by the local ViT models, ensuring enhanced privacy preservation compared to traditional parameter or synthetic image sharing methodologies. FedEvPrompt is optimized within a round-based learning paradigm, where each round involves training local models followed by attention maps sharing with all federation clients. Experimental validation conducted in a real distributed setting, on the ISIC2019 dataset, demonstrates the superior performance of FedEvPrompt against baseline federated learning algorithms and knowledge distillation methods, without sharing model parameters. In conclusion, FedEvPrompt offers a promising approach for federated learning, effectively addressing challenges such as data heterogeneity, imbalance, privacy preservation, and knowledge sharing.

Figures (2)

• Figure 1: Overview of FedEvPrompt. During a round of Training (top), local data is used to optimize b-prompts, encoding general visual features, and t-prompts, encoding task-specific information, prepended to a frozen ViT encoder. Optimization is carried out by minimizing evidential loss ($\mathcal{L_{\epsilon}}$) and a knowledge distillation loss ($\mathcal{L_{KD}}$) between local attention maps and those of the federation available in the uncertainty-aware attention buffer. After a round of training at Inference (bottom), the client identifies, for each class $K$, its $M$ most informative attention rollout maps (sorted by lowest uncertainty) to contribute to the federated uncertainty-aware attention buffer.
• Figure 2: Distribution of the Fed-ISIC2019 dataset across clients.