[Paper Review] Analysis and Optimal Edge Assignment For Hierarchical Federated Learning on Non-IID Data

Naram Mhaisen, Awad, Alaa+3 more|arXiv (Cornell University)|Dec 10, 2020

Privacy-Preserving Technologies in Data26 references4 citations

TL;DR

This paper proposes an optimal user-edge assignment strategy in hierarchical federated learning (HFL) to reduce model performance degradation caused by non-IID data. By minimizing the statistical divergence between edge-level data distributions and the global distribution, the method improves convergence and accuracy through a heuristic optimization that balances class distributions across edge nodes.

ABSTRACT

Distributed learning algorithms aim to leverage distributed and diverse data stored at users' devices to learn a global phenomena by performing training amongst participating devices and periodically aggregating their local models' parameters into a global model. Federated learning is a promising paradigm that allows for extending local training among the participant devices before aggregating the parameters, offering better communication efficiency. However, in the cases where the participants' data are strongly skewed (i.e., non-IID), the local models can overfit local data, leading to low performing global model. In this paper, we first show that a major cause of the performance drop is the weighted distance between the distribution over classes on users' devices and the global distribution. Then, to face this challenge, we leverage the edge computing paradigm to design a hierarchical learning system that performs Federated Gradient Descent on the user-edge layer and Federated Averaging on the edge-cloud layer. In this hierarchical architecture, we formalize and optimize this user-edge assignment problem such that edge-level data distributions turn to be similar (i.e., close to IID), which enhances the Federated Averaging performance. Our experiments on multiple real-world datasets show that the proposed optimized assignment is tractable and leads to faster convergence of models towards a better accuracy value.

Motivation & Objective

• Address the performance degradation in federated learning due to non-IID (non-independent and identically distributed) data distributions across user devices.
• Formalize the user-edge assignment problem in hierarchical federated learning as an optimization task to minimize edge-level data distribution divergence from the global distribution.
• Leverage edge computing infrastructure to enable reconfigurable, data-balancing assignments that enhance FedAvg performance.
• Develop tractable solutions—both exact (branch and bound) and heuristic—for the NP-hard assignment problem under real-world constraints.
• Demonstrate that balancing edge-level data distributions leads to faster convergence and higher model accuracy in practice.

Proposed method

• Formulate the user-edge assignment problem as minimizing the weighted L2 distance between each edge’s data distribution and the global data distribution.
• Propose a branch-and-bound-based solution for a simplified linear version of the problem with equal-sized edge assignments.
• Introduce a greedy heuristic algorithm that iteratively assigns users to edges to equalize class distributions across edge nodes.
• Use the FedAvg algorithm at both the user-edge and edge-cloud layers, with model aggregation weighted by local dataset size.
• Integrate network topology constraints by restricting user assignments to only those edge nodes within communication range.
• Validate the optimization using real-world datasets under controlled simulation settings with multiple edge nodes per user.

Experimental results

Research questions

• RQ1 How does the statistical divergence between local data distributions and the global distribution affect model convergence in hierarchical federated learning?
• RQ2 Can user-edge assignment be optimized to minimize edge-level data distribution divergence and thereby improve FedAvg performance?
• RQ3 What are the trade-offs between minimizing user-edge distribution distance and minimizing edge-cloud distribution distance in a dual-objective optimization framework?
• RQ4 How effective is a heuristic-based user-edge assignment strategy in achieving balanced data distributions across edge nodes?
• RQ5 To what extent does optimized user-edge assignment improve convergence speed and final model accuracy in non-IID settings?

Key findings

• The upper bound of deviation from the centralized model is proportional to the weighted distance between each node’s data distribution and the global distribution, validating the need to minimize this divergence.
• The proposed heuristic user-edge assignment method achieves faster convergence and higher final accuracy compared to random or unoptimized assignments across multiple real-world datasets.
• On datasets where users are within range of only two edge nodes, the optimized assignment leads to significant performance improvements, demonstrating practical feasibility.
• The heuristic approach yields sub-optimal solutions that are close to the optimal solution in the simplified case, especially when class distributions are balanced.
• The problem is NP-hard, but the heuristic provides a tractable and scalable solution suitable for real-time deployment in mobile and edge learning systems.
• Future work should consider dual-objective optimization balancing user-edge and edge-cloud distribution distances, particularly when user-edge communication delay becomes significant.