FedAgain: A Trust-Based and Robust Federated Learning Strategy for an Automated Kidney Stone Identification in Ureteroscopy

1. Problem Statement

The paper addresses the critical challenge of deploying Artificial Intelligence (AI) for medical imaging in real-world, multi-institutional settings. Specifically, it focuses on automated kidney stone identification via ureteroscopy. The core problems identified are:

• Data Heterogeneity (Non-IID): Medical data across different hospitals varies significantly due to different imaging devices, acquisition protocols, patient demographics, and label distributions. Standard Federated Learning (FL) algorithms (like FedAvg) often struggle to converge or generalize under these conditions.
• Image Corruption and Noise: Endoscopic images are prone to specific artifacts such as motion blur, specular reflections, variable illumination, and occlusions. These "corruptions" degrade model performance and can lead to diagnostic errors.
• Client Reliability: In a federated setting, some clients may provide low-quality data (due to hardware issues) or malicious updates (Byzantine attacks). Standard aggregation methods treat all clients equally, allowing noisy or adversarial updates to poison the global model.
• Privacy Constraints: Centralizing sensitive patient data is often prohibited by regulations (e.g., HIPAA, GDPR), necessitating decentralized training methods that preserve privacy while ensuring model robustness.

2. Methodology: The FedAgain Framework

The authors propose FedAgain, a trust-based Federated Learning strategy designed to dynamically weight client contributions based on their reliability. Unlike standard FL which uses sample-size weighting, FedAgain introduces a dual-trust mechanism that evaluates clients on two dimensions: data quality and model consistency.

Core Components

1. Benchmark Error ($E_k^{(t)}$):
   • Before local training begins in round $t$, each client evaluates the current global model on its local validation dataset.
   • This metric quantifies how well the global model aligns with the client's data distribution. A high error suggests the client's data is either unrepresentative, corrupted, or the client is behaving anomalously.
2. Parameter Divergence ($D_k^{(t)}$):
   • After local training, the server calculates the Euclidean distance between the client's updated weights and the global model's weights.
   • This measures the deviation of the client's learning trajectory from the global optimization path. Large divergence may indicate overfitting to noisy local data or adversarial manipulation.
3. Trust Score Calculation:
   • The server computes a composite trust score $T_k^{(t)}$ for each client using the formula:
     $$T_k^{(t)} = \frac{1}{E_k^{(t)} \cdot D_k^{(t)} + \epsilon}$$
     (where $\epsilon$ is a small constant to prevent division by zero).
   • Logic: Clients with low benchmark error (good data alignment) and low parameter divergence (consistent updates) receive high trust scores. Conversely, clients with high error or divergence are down-weighted.
4. Trust-Weighted Aggregation:
   • The global model is updated using a weighted average where the weights are the normalized trust scores ($\hat{T}_k^{(t)}$) rather than just the number of samples.
   • This effectively filters out unreliable or corrupted updates during the aggregation phase, promoting stable convergence.

Experimental Setup

• Datasets: The framework was tested on five datasets:
  • Two public benchmarks: MNIST and CIFAR-10 (for generalizability).
  • Three kidney stone datasets: Michel Daudon (Dataset A), Jonathan El-Beze (Dataset B), and MyStone (Dataset C). These vary in acquisition devices (CCD vs. Endoscope), resolution, and environmental conditions (controlled vs. in-vivo).
• Corruption Simulation: The authors simulated realistic adversarial scenarios by corrupting a percentage ($\rho$) of clients (10% to 50%) with various distortions (Gaussian noise, blur, contrast changes, etc.) and varying severity levels (1–5).
• Baselines: FedAgain was compared against FedAvg (standard), FedProx (heterogeneity-aware), FedMedian (robust statistics), and Bulyan (Byzantine-robust).

3. Key Contributions

1. First Robust FL for In-Vivo Kidney Stone ID: This is the first federated learning approach specifically designed for the in-vivo recognition of kidney stone types that explicitly addresses model robustness against image corruption and client unreliability.
2. Dual-Signal Trust Mechanism: The paper introduces a novel aggregation rule that combines benchmark error (pre-training data quality check) and parameter divergence (post-training consistency check) to dynamically weight clients. This avoids the need for explicit outlier detection or attack-specific tuning.
3. Comprehensive Evaluation Framework: The authors provide a reproducible PyTorch codebase and a rigorous evaluation suite involving:
   • Seven stress-tests varying corruption rates (10%–50%) and severity.
   • Evaluation across non-IID distributions (Label Skew and Dirichlet).
   • Scalability tests with varying client counts (5 to 20).
4. Empirical Validation of Robustness: The study demonstrates that integrating anomaly detection signals directly into the federated loop (rather than as a passive filter) significantly improves reliability in heterogeneous and corrupted environments.

4. Experimental Results

The results demonstrate that FedAgain consistently outperforms standard and robust baselines, particularly under adverse conditions:

• Clean Data (IID): Under ideal conditions, FedAgain performs competitively with FedAvg and FedProx, showing that the trust mechanism does not degrade performance when data is clean.
• Corrupted Data (Non-IID & Adversarial):
  • Moderate to Heavy Corruption: As corruption levels increased (up to 40% of clients corrupted with high severity), FedAvg and FedProx suffered significant performance drops. FedAgain maintained stable accuracy (e.g., ~72–73% accuracy even with 40% corruption), while FedAvg dropped to ~68% and Bulyan to ~64%.
  • Stability: FedAgain exhibited lower variance (smaller error bars) across runs compared to FedAvg, indicating more predictable behavior in heterogeneous settings.
  • Precision: FedAgain showed a significant gain in precision (4–12% higher than FedAvg), reducing false positives, which is critical in medical diagnosis.
• Scalability: As the number of participating clients increased (from 5 to 20), FedAvg's performance degraded due to increased heterogeneity and noise. FedAgain scaled more gracefully, maintaining a performance margin of +3.8% to +8.3% over FedAvg in larger federations.
• Generalization: The method successfully generalized across different image modalities (grayscale MNIST, RGB CIFAR-10, and complex endoscopic images), proving its domain-agnostic nature.

5. Significance and Impact

• Clinical Viability: By ensuring robustness against image artifacts (blur, glare, noise) common in endoscopy, FedAgain moves AI-assisted diagnosis closer to real-world clinical deployment where data quality cannot be guaranteed.
• Privacy-Preserving Collaboration: It enables multiple hospitals to collaborate on training a kidney stone classifier without sharing sensitive patient data, overcoming regulatory barriers.
• Resilience to "Bad Actors": The framework provides a practical solution for handling both unintentional data degradation (e.g., old cameras, poor lighting) and intentional adversarial attacks without requiring complex cryptographic protocols or prior knowledge of the number of malicious clients.
• Future Direction: The work lays the foundation for trustworthy, scalable, and privacy-compliant AI in healthcare, suggesting future integration with differential privacy and secure aggregation for even stricter compliance.

In summary, FedAgain represents a significant step forward in making Federated Learning viable for high-stakes medical applications by actively managing data heterogeneity and corruption through a dynamic, trust-based aggregation strategy.