# Research into Robust Federated Learning Methods Driven by Heterogeneity Awareness
## 4 Experiment
### Implementation Details  
This section provides a detailed overview of the experimental setup, including the federated learning framework, model architecture, optimization settings, and client-server interaction. All experiments in this work are conducted using the Flower federated learning framework, which supports a flexible and extensible environment for federated learning experiments. 1. Federated Learning Framework: - Flower Framework: We implement the proposed heterogeneity-aware optimization method using the Flower framework, which allows for the easy integration of federated learning components, including client-server communication, model aggregation, and distributed training. - Client-Side Setup: Each client is assigned a local dataset, which is non-IID to simulate realistic conditions in federated learning. Clients independently perform local model training on their datasets and periodically upload their model updates and heterogeneity metrics to the server. Figure 4: Data distribution plot, with alpha = 0.5 above and alpha = 0.1 below. Model Architecture: A Convolutional Neural Network (CNN) with a simple and lightweight architecture is employed across all the clients for a fair comparison of methods. Specifically, the model consists of three convolutional layers, followed by a fully connected layer to output the class predictions. This architecture ensures that the focus remains on the aggregation strategy and heterogeneity-aware optimization rather than model complexity. The layer details are as follows: - Convolutional Layers: The model uses three convolutional layers with ReLU activations and max-pooling after each convolution. - Fully Connected Layer: A fully connected layer is used to output the final class prediction, followed by a softmax activation to generate probabilities for classification tasks. This choice of architecture is designed to be lightweight, ensuring the experiments are computationally feasible, while maintaining sufficient complexity to test the core ideas of the proposed method. 3. Federated Learning Procedure: - Federated Rounds: The system operates over a series of federated rounds. In each round, all clients perform local training on their datasets, calculate their heterogeneity metrics (based on feature statistics and manifold structures), and upload both the model parameters and heterogeneity scores to the server. - Server-Side Aggregation: The server aggregates the model parameters from all the clients using the heterogeneity-aware weighted aggregation strategy described in Section 3.4. This aggregation method adjusts the contribution of each client's update based on its local heterogeneity, reducing the impact of clients with highly imbalanced data distributions. - Global Model Update: After receiving updates from all the clients, the server combines the model parameters using the weighted strategy, updating the global model. The updated global model is then sent back to the clients for the next round of local training. 4. Client Participation: - All Clients Participate in Every Round: In each federated round, all clients participate in the model update process. This ensures that the global model benefits from diverse data distributions and that the model learns from all the clients' datasets over time. - Communication Efficiency: The communication overhead is minimized by uploading only the model parameters and heterogeneity metrics, rather than the entire dataset. This ensures the privacy of client data while enabling effective federated learning. 5. Baselines: For comparison, FedAvg is used as the baseline aggregation method. FedAvg is a widely used approach for federated learning that aggregates local model updates by averaging, without accounting for client heterogeneity. The performance of the proposed method is compared against FedAvg to evaluate the improvements achieved by the heterogeneity-aware optimization framework.