Mehta, Shao develop federated learning framework for heterogeneous data

MechSE doctoral student Manan Mehta and Associate Professor Chenhui Shao recently published a new study in IEEE Transactions on Industrial Informatics. Their study, “A Greedy Agglomerative Framework for Clustered Federated Learning,” is the second major study they’ve published in the area of federated learning (FL) and manufacturing, following one published in the Journal of Manufacturing Systems in 2022.

Federated learning has received widespread attention for supporting the training of deep learning models across multiple IoT devices while preserving privacy of data. However, recent studies have shown that the quality of global FL models, which is significant for industrial big data in applications like healthcare, smart manufacturing, autonomous driving, and robotics, deteriorates in the presence of non-IID (i.e., non-independent and identically distributed) data. This type of data can cause the model to learn incorrect patterns, which decreases its accuracy and leads to poorer prediction performance.

For example, training a model for managing the quality control of 3D-printed parts using computer vision requires accounting for variations among defects, and their pixel intensities, across different printers and printing technologies. Without sufficient data diversity, models would have a limited view of the spectrum of printers and their defects. However, with significant data diversity, model performance could encompass all printers.

“We observed that when [non-IID data] factories build an FL model, their fault classification accuracy is poor due to the statistically heterogeneous data at each factory,” Mehta said.

The researchers present a novel clustered FL framework called Federated Learning via Agglomerative Client Clustering (FLACC), which greedily agglomerates similar clients and groups of clients based on gradient updates while learning a global FL model. The framework keeps clients with dissimilar data in separate clusters while clustering clients with similar data, which allows those clients to benefit from one another.

“Our framework identifies similar clients automatically,” Mehta said. “It is robust in that only a fraction of clients can participate in a training round; thus, not all clients need to be online at all times.”

The researchers next intend to explore how privacy-preserving mechanisms like differential privacy or homomorphic encryption impact overall model performance. “We are also focusing on making the framework end-to-end by designing a system-level architecture linking hardware, networking, and software,” Mehta said. The pair’s framework is applicable in both IID and non-IID scenarios.

“Overall, our framework is designed keeping realistic industrial scenarios in mind, [which allows it to] overcome some crucial obstacles for the widespread industrial adoption of FL,” Mehta said.

“The broad utility of FLACC is demonstrated not only in manufacturing settings, but also in other domains such as computer vision, thus contributing to the algorithmic development in the [field of] FL,” Shao said. “FLACC will be an essential part of the collaborative learning architecture consisting of hardware, software, and cyber-physical infrastructure that is being developed in our ongoing research.”