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

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

• The paper introduces a novel framework that integrates architecture adaptation, latency-aware optimization, and heterogeneous aggregation to meet strict per-device latency constraints.
• It leverages a backpropagation-based latency predictor and dynamic zeroizing-recovering techniques to fine-tune model complexity on diverse edge systems.
• Experimental results demonstrate up to a 3.09% accuracy improvement while ensuring consistent latency budgets across varied datasets and platforms.

Collaborative Neural Network Learning for Latency-Critical Edge Systems

Introduction

This essay analyzes "Collate: Collaborative Neural Network Learning for Latency-Critical Edge Systems" (2607.08013), which addresses the limitations of existing Federated Learning (FL) paradigms in the context of highly heterogeneous, resource-constrained edge environments with strict end-to-end inference latency constraints. By leveraging a novel integration of architecture adaptation, latency-aware optimization, and communication-efficient aggregation, Collate advances the state-of-the-art for real-world deployment of collaborative deep neural networks (DNNs) on diverse edge hardware.

Motivation: Device Heterogeneity and Latency Constraints

Standard FL frameworks aggregate identical model architectures across clients to ensure privacy and robustness. However, the computational heterogeneity across edge systems renders a single architecture suboptimal or even infeasible for some clients. Empirical observations show a wide variance in inference latency when deploying the same federated-trained AlexNet across devices, ranging from sub-20 ms to over 200 ms. This discrepancy highlights the inability of traditional FL to simultaneously provide high accuracy and satisfy stringent latency requirements for all participants.

Figure 1: Latency of the federated-trained AlexNet on various edge systems shows massive variance, motivating architecture-level adaptation.

Moreover, the commonly assumed monotonic mapping between model complexity (FLOPs) and device latency is unreliable—nonlinear device-specific characteristics (e.g., memory bandwidth, specialized accelerators) obfuscate this relationship.

Figure 2: Relationship between latency and FLOPs for ResNet variants reveals non-uniform, device-specific latency profiles.

Collate Framework Overview

Collate orchestrates heterogeneous model learning while meeting latency constraints on each edge device, using a single global training process. The overall design tightly integrates a hardware-customized, neural-latency prediction module, a dynamic zeroizing-recovering strategy for local architecture adjustment, and a new proto-corrected aggregation protocol for merging heterogeneous local models.

Figure 3: Schematic of Collate’s process. The top branch implements latency guarantee per edge, while the bottom depicts the aggregation of heterogeneous models.

The Collate framework operates in several key stages:

1. Device-Specific Latency Prediction

A lightweight backpropagation (BP) neural latency predictor is first trained for each device on a pool of micro-benchmarked, single-layer structures. This module provides low-overhead, differentiable latency estimations crucial for real-time architectural adjustments during local training.

2. Model Extension

Contrary to prior works, Collate supports both model shrinking (for constrained clients) and uniform model extension for powerful systems, proportionally scaling each layer’s width while controlling for device-specific latency ceilings.

Figure 4: Collate’s extension scheme uniformly scales model width within latency constraints, enabling accuracy gains on high-end devices.

3. Latency-Aware Local Training: Zeroizing and Recovering

Each client iteratively prunes (zeroizes) bottom-ranked filters using a trainable mask layer, guided by the latency predictor to guarantee compliance with local latency constraints. Crucially, the recovering step allows previously pruned filters to re-activate if subsequent training signals reveal their importance—enabled intrinsically by optimizer momentum. This promotes dynamic, data-driven architecture exploration rather than static pruning.

Figure 5: Evolution of local model filter contributions, with zeroizing (pruning) and later recovering, coordinated with communication rounds.

4. Heterogeneous Model Aggregation

Local models, now architecturally distinct, are aggregated using a two-part protocol:

• Heterogeneous layers are merged by parameter set alignment (subsumed or shared), using a complexity-ranking criterion.
• A proto-corrected mechanism, employing class prototypes extracted from a commonly-shaped fully connected layer, penalizes deviation between local and global prototypes, thus regularizing large models and improving generalization, especially under data heterogeneity.

  Figure 6: Heterogeneous Federated Learning aggregation in Collate, illustrating parameter alignment and prototype-based correction among diverse clients.

Experimental Validation

Experiments span five heterogeneous edge platforms (e.g., HP ProBook, Jetson TX2/Nano, Pi, Note10), across widely adopted datasets (MNIST, CIFAR-10/100, HAR), and multiple architectures. Collate is benchmarked against ERFL, Helios, and HeteroFL—all state-of-the-art heterogeneous FL approaches.

Accuracy is tracked under both IID and realistic Non-IID data splits, using local complexity scaling and strict latency controls.

Figure 7: Test accuracy trajectories for LeNet/MNIST across complexity levels and data distributions (top: IID, bottom: Non-IID), demonstrating Collate’s convergence and accuracy advantages.

Key numerical findings:

• For extended models (e.g., when deploying on Jetson TX2 or ProBook), Collate improves accuracy by 1.96% on average compared to competitive methods under identical latency constraints.
• For shrunk models (suited to severely constrained devices like Raspberry Pi 4B), Collate yields a 3.09% mean accuracy improvement.
• All models satisfy predefined per-device latency budgets (e.g., $36.8 \pm 0.7$ ms on CIFAR-10, matching the reference Jetson Nano configuration), with no significant training overhead and robust generalization under Non-IID splits.

Implications and Future Directions

From a practical standpoint, Collate presents a reproducible way to maximize the utility of edge FL deployments, dynamically exploiting the full compute spectrum present in real-world device fleets. The architecture search (via zeroizing-recovering), powered by BP latency prediction, decouples inference efficiency from static FLOP-based pruning and leverages runtime empirical device characteristics. This scalable framework supports mixed-accuracy/latency targets in collaborative inference infrastructures—foundational for interactive, privacy-preserving AI in domains such as IoT, automotive, and personalized healthcare.

Theoretically, this work suggests new research lines for device-aware neural architecture search integrated with FL. Further, Collate’s abstracted aggregation via class prototypes could be further extended to more general forms of aggregation in cross-domain or lifelong federated training regimes.

Conclusion

Collate formalizes a comprehensive framework for collaborative heterogeneous neural network learning tailored for latency-critical edge environments. Through a synergy of neural latency prediction, adaptive local architecture evolution, and prototype-regularized aggregation, Collate outperforms existing approaches in latency-constrained accuracy, with minimal communication or computational penalty. The methodology presents direct solutions for practical federated intelligence on edge platforms while motivating theoretical advances in device-adaptive, communication-efficient collaborative learning.