Federated Learning for Collaborative Inference Systems: The case of early exit networks

In today’s increasingly diverse computing landscape, end devices like sensors and smartphones are progressively equipped with AI models tailored to their local memory and computational constraints. Local inference reduces communication costs and latency; however, these smaller models typically underperform compared to more sophisticated models deployed on edge servers or in the cloud. Collaborative Inference Systems (CISs) address this performance trade-off by enabling smaller devices to offload part of their inference tasks to more capable devices. These systems often deploy hierarchical models that share numerous parameters, exemplified by deep neural networks that utilize strategies like early exits or ordered dropout. In such instances, Federated Learning (FL) may be employed to jointly train the models within a CIS. Yet, traditional training methods have overlooked the operational dynamics of CISs during inference, particularly the potential high heterogeneity in serving rates across the devices within a given CIS. To address this gap, we propose a novel FL approach that explicitly accounts for variations in serving rates within CISs. Our framework not only offers rigorous theoretical guarantees but also surpasses state-of-the-art training algorithms for CISs, especially in scenarios where end devices handle higher inference request rates and where data availability is uneven across devices. • We propose the first inference-aware Federated Learning framework for Collaborative Inference Systems (CISs), enabling smaller devices to offload part of their inference tasks to more capable devices. This framework incorporates real-time inference requests into the training process, coordinating the simultaneous training of heterogeneous models across devices with varying capacities. • We provide analytical guarantees for the error of our proposed training algorithm with respect to its key parameter: the aggregation weights. Our comprehensive error analysis covers generalization error, bias error, and optimization error. The resulting theoretical bound offers novel insights into the selection of aggregation weights, distinguishing our configuration rule from existing approaches and leading to improved performance. • Extensive experimental results across diverse inference request scenarios validate the theoretical insights and demonstrate the superiority of our approach over state-of-the-art methods for training CISs, particularly in scenarios involving high inference request rates and uneven data availability across devices. • Our rigorous theoretical results extend beyond the CIS training framework, applying broadly to approaches that jointly train models with shared parameters, such as ordered dropout, pruning, and other nested training methodologies.