Collate: Collaborative Neural Network Learning for Latency-Critical Edge Systems

2026-07-08

AI Summary

This work addresses the challenge in real-time edge federated learning where device heterogeneity makes it difficult to simultaneously satisfy diverse system latency constraints and maintain high model accuracy. To this end, the authors propose Collate, a framework that jointly trains heterogeneous local models within a single training process to accommodate varying latency requirements while preserving performance. The key innovations include a dynamic zeroing-and-recovery mechanism that adaptively adjusts local model architectures and a prototype-corrected federated aggregation strategy that enables efficient fusion of heterogeneous models. Experimental results demonstrate that Collate improves the average accuracy of expanded and compressed models by 1.96% and 3.09%, respectively, with virtually no additional training overhead.

Abstract

Federated Learning (FL) empowers multiple clients to collaboratively learn a model, enlarging the training data of each client for high accuracy while protecting data privacy. However, when deploying FL in real-time edge systems, the heterogeneity of devices among systems has a severe impact on the performance of the inferred model. Existing optimizations on FL focus on improving the training efficiency but fail to speed up inference, especially when there is a latency constraint. In this work, we propose Collate, a novel training framework that collaboratively learns heterogeneous models to meet the latency constraints of multiple edge systems simultaneously. We design a dynamic zeroizing-recovering method to adjust each local model architecture for high accuracy under its latency constraint. A proto-corrected federated aggregation scheme is also introduced to aggregate all heterogeneous local models, satisfying the latency constraint of different systems with only one training process and maintaining high accuracy. Extensive experiments indicate that, compared to state-of-the-art methods and under a latency constraint, our extended models can improve the accuracy by 1.96% on average, and our shrunk models can also obtain a 3.09% accuracy improvement on average, with almost no extra training overhead. The related codes and data will be available at https://github.com/ntuliuteam/Collate

Problem

• Federated Learning
• Edge Systems
• Latency Constraint
• Model Heterogeneity
• Inference Latency

Innovation

• Federated Learning
• Latency Constraint
• Heterogeneous Models
• Dynamic Architecture Adjustment
• Proto-corrected Aggregation