Fed-Meta-Align: A Similarity-Aware Aggregation and Personalization Pipeline for Federated TinyML on Heterogeneous Data

TL;DR

Fed-Meta-Align is a four-phase federated learning framework that enhances fault classification accuracy on heterogeneous IoT devices by combining meta-initialization, similarity-aware aggregation, and personalization.

Contribution

It introduces a novel multi-stage pipeline with meta-initialization and similarity-based aggregation to improve federated learning in resource-constrained, non-IID IoT environments.

Findings

Achieves 91.27% average test accuracy on heterogeneous IoT datasets.

Outperforms FedAvg and FedProx by up to 3.87% and 3.37%.

Demonstrates robustness in real-time fault classification for TinyML.

Abstract

Real-time fault classification in resource-constrained Internet of Things (IoT) devices is critical for industrial safety, yet training robust models in such heterogeneous environments remains a significant challenge. Standard Federated Learning (FL) often fails in the presence of non-IID data, leading to model divergence. This paper introduces Fed-Meta-Align, a novel four-phase framework designed to overcome these limitations through a sophisticated initialization and training pipeline. Our process begins by training a foundational model on a general public dataset to establish a competent starting point. This model then undergoes a serial meta-initialization phase, where it sequentially trains on a subset of IOT Device data to learn a heterogeneity-aware initialization that is already situated in a favorable region of the loss landscape. This informed model is subsequently refined in a parallel FL phase, which utilizes a dual-criterion aggregation mechanism that weights for IOT devices updates based on both local performance and cosine similarity alignment. Finally, an on-device personalization phase adapts the converged global model into a specialized expert for each IOT Device. Comprehensive experiments demonstrate that Fed-Meta-Align achieves an average test accuracy of 91.27% across heterogeneous IOT devices, outperforming personalized FedAvg and FedProx by up to 3.87% and 3.37% on electrical and mechanical fault datasets, respectively. This multi-stage approach of sequenced initialization and adaptive aggregation provides a robust pathway for deploying high-performance intelligence on diverse TinyML networks.