Collate: Collaborative Neural Network Learning for Latency-Critical Edge Systems Shuo Huai‡§ , Di Liu§ , Hao Kong‡§ , Xiangzhong Luo‡ , Weichen Liu‡ , Ravi Subramaniam† , Christian Makaya† and Qian Lin†
‡ School of Computer Science and Engineering, Nanyang Technological University, Singapore § HP-NTU Digital Manufacturing Corporate Lab, Nanyang Technological University, Singapore † HP Inc., Palo Alto, California, USA
arXiv:2607.08013v1 [cs.LG] 9 Jul 2026
Email: {shuo001, liu.di, kong.hao, xiangzho001, liu}@ntu.edu.sg, {ravi.subramaniam, christian.makaya, qian.lin}@hp.com
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. Index Terms—edge devices, edge intelligence, neural network learning, federated learning, inference efficiency
I. I NTRODUCTION Deep Neural Networks (DNNs) have brought significant breakthroughs in many different applications, such as image recognition and natural language processing [1]. With the emphasis on data privacy and concerns over transmission stability, current DNN applications are increasingly deployed on edge devices, such as autonomous vehicles, healthcare devices, etc [1]. Meanwhile, DNN models require a huge amount of training data to improve accuracy [2], but data in most industries are protected by privacy laws and thus are required to be in the form of isolated islands. Federated Learning (FL) is designed to coordinate multiple clients to train a DNN model collaboratively without sharing their original local data, and it is capable of preserving the data privacy and achieving better accuracy than each individual client training with only its local data [3]. This addresses the issue of insufficient individual training data. However, in addition to accuracy, latency is also an important metric for edge intelligent systems [4]. With diverse © 2022 IEEE. Personal use of this material is permitted. Permission from IEEE must be obtained for all other uses, in any current or future media, including reprinting/republishing this material for advertising or promotional purposes, creating new collective works, for resale or redistribution to servers or lists, or reuse of any copyrighted component of this work in other works. This is the author’s accepted version of the article published in the 2022 IEEE 40th International Conference on Computer Design (ICCD), pp. 627634, 2022, DOI: 10.1109/ICCD56317.2022.00097. This study is supported under the RIE2020 Industry Alignment Fund – Industry Collaboration Projects (IAF-ICP) Funding Initiative, as well as cash and in-kind contribution from the industry partner, HP Inc., through the HPNTU Digital Manufacturing Corporate Lab (I1801E0028).
Federated Training on Private Data
Inference on Multiple Heterogeneous Devices
Private Data 1
HP ProBook 440 G6
Private Data 2
Model Architecture
Private Data 3 Private Data 4 Private Data 5
NVIDIA Jetson TX2
Federated-Trained Identical Model
NVIDIA Jetson Nano
Raspberry Pi 4B
23.0 ms
13.6 ms
37.5 ms
236.0 ms
Samsung Galaxy Note10 82.6 ms
Fig. 1. Latency of the federated-trained AlexNet on various edge systems.
edge devices that demonstrate different computational capabilities emerging, the identical model trained by FL is not efficient for all participating edge systems and cannot achieve a good balance between latency and accuracy. As demonstrated in Fig. 1, when deploying the model into various edge systems, the latency differs from 13.6 ms to 236.0 ms. When this application has a latency constraint (e.g., 30.0 ms), some systems can infer more complex models for higher accuracy, whereas others cannot even deploy the model due to the limited memory and computational resources. Take a practical scenario from WeBank as an example [5]. They need an edge FL framework in which surveillance video data collected and stored in the edge cloud of each surveillance company are not required to be uploaded to a central cloud for centralized model training. After each local training iteration, only the model parameters from each surveillance company are sent to the FL server for aggregation. The final federated-trained model is distributed to the participating surveillance companies for object detection. However, different companies can use distinct edge devices and this application features real-time, so it is necessary to train heterogeneous models for different companies to meet their latency constraints. Some efforts have been made to use different DNN architectures to fit various clients in the training stage, known as Heterogeneous FL [6]–[10]. These methods mainly fall into two categories: one is to fine-tune models for different clients from an identical global model, and the other is to directly learn from heterogeneous models without the same global model. However, these existing approaches are designed for accelerating the training stage and cannot optimize the inference latency directly. When the first one is employed in a
latency-critical system, it can only train a specific network for one system at a time until it trains for all systems using multiple FL processes, which imposes enormous training overhead. Although the second one can directly provide heterogeneous model architectures for different systems, it requires some public datasets for transfer learning [9] or an extra dataset for prototype learning [10]. Moreover, the absence of the same global model lead to an accuracy drop of up to 10% [11]. To guarantee the generality and accuracy of the FL scheme, our method should be based on the first one. Meanwhile, when optimizing models under the latency constraint, we should not only reduce latency for low-end systems but also extend models for powerful systems (e.g., Jetson TX2 in Fig. 1) to improve their accuracy [12], in contrast to existing heterogeneous FL that only reduces the training cost by shrinking models. Thus, we integrate the model extension into our learning framework to better utilize each client. To our best knowledge, this is the first paper to optimize FL to simultaneously meet the latency constraints of all participating systems while obtaining high accuracy. Specifically, our main contributions are summarized as: • We propose a novel model learning framework, Collate, that cultivates optimal DNN architectures collaboratively for multiple edge systems to obtain higher accuracy and satisfy their latency constraints with only one training process. • We present a proto-corrected aggregation scheme in the global training process to effectively aggregate all heterogeneous models from each edge system for higher accuracy. • We design a latency-aware local training scheme by a dynamic zeroizing-recovering training process. It extends the exploration space of Collate to discover the optimal DNN architecture for each edge system. • We demonstrate the effectiveness of Collate with extensive experiments. Compared to the state-of-the-art methods and under the same latency constraints, our extended models can improve the accuracy by 1.96% on average, and the accuracy of shrunk models outperforms others by 3.09% on average. II. BACKGROUND & R ELATED W ORK In this section, we introduce the preliminaries of FL, heterogeneous FL and latency prediction, including the advantages and disadvantages of some related work. A. Federated Learning In the traditional FL algorithm (FedAvg [13]), multiple edge systems collaborate to train the same model on their respective local training data to solve issues of data isolated islands and protect privacy. During this training step, different systems transmit model parameters to the server after every le local training epochs. Then the server executes model aggregation. Considering this FL process includes n systems, the loss function for system i is fi . Let w represent the model parameter, and then the training optimization goal becomes: n 1X min f (w) := fi (w) (1) w n i=1 And the training process can be formulated as: t+(e+1)
wi
= wit+e − ηit+e ∇fi (wit+e )
wt+1 =
n X
1 wt+le n i=1 i
(2) (3)
where wit+e denotes the local model parameter on system i after t communication rounds and e local training steps. η is the learning
rate, and ∇ represents the derivative. From Eq. (3), we can see that the traditional FL [13] is based on the fact that all local model architectures of different systems are the same for aggregation.
B. Heterogeneous Federated Learning Two major challenges involved in the traditional FL process are statistical heterogeneity and hardware heterogeneity among different systems [14]. Statistical heterogeneity means the data on different systems are not independent and identically distributed (Non-IID) [15], and different systems have different information to learn, while hardware heterogeneity refers to various devices participating in the FL process having different computational abilities. Different heterogeneous FL methods are proposed for these two challenges. Statistical Heterogeneity: Most existing works on statistical heterogeneity aim to achieve higher accuracy on each local data. They use different or identical model architectures with different parameters for each participating system by customizing the global model. Wang et al. [16] proposed to fine-tune some or all parameters of a trained global model using each client’s local data. Jiang et al. [17] proposed to combine meta-learning with FL for customizing models to recognize each local data pattern. Khodak et al. [18] proposed to improve the accuracy of statistically heterogeneous FL via online convex optimization theory and meta-learning. These methods focus on improving the accuracy of local models on the local dataset of each system. They are applied to the scenario where most of the inference data of each system match the pattern of its corresponding training set. This paper aims to learn different local models for all participating edge systems to meet their latency requirements. Thus, when considering statistical heterogeneity, these aforementioned methods can be used as the back-end of our method to achieve higher accuracy. Hardware Heterogeneity: Current research on hardware heterogeneity concentrates on simplifying the model architecture to accelerate training and reduce the model size on devices with low computational ability. Caldas et al. [19] proposed to randomly select small subsets of the global model with the expected model volume to provide a reduction in both client-to-server communication and local computation during the training. Li et al. [20] proposed Hermes to find small sub-networks for some devices by applying the structured pruning. Jeong et al. [21] proposed to use knowledge distillation to reduce model size in some low-end devices. Xu et al. [7] proposed Helios to accelerate the devices with weak computational capacities by dynamically compressing the global model into an expected volume during the training stage. Diao et al. [8] proposed HeteroFL to assign models with different computational complexities to each device in advance according to the computational ability of each device. Also, they proposed a method to aggregate heterogeneous local models to produce a global model during the training stage. These methods are proposed to reduce the training time or the communication cost when applying on-device training. However, FL is not just on-device training [22]. And typically, the usage (inference) of a model lasts longer than its generation (training). Thus, inference optimization is more critical for the FL framework. Although some of these methods can be modified to accelerate inference on low-end devices, these methods designed for the global model lead to a large accuracy loss for each local model. Meanwhile, these methods did not consider improving accuracy for powerful devices by extending models, and their aggregation methods introduce an accuracy drop to large models, as illustrated in Section IV.
7.5 5.0 0
243 380 200 400 FLOPs (M)
24 22 20 180
217 457 200 400 FLOPs (M)
70 60 50 40 30 20 0
203 483 200 400 FLOPs (M)
150
70
125
60
Latency (ms)
10.0
Latency (ms)
Latency (ms)
Latency (ms)
12.5
26
Latency (ms)
28
15.0
100 75 50 25 0
258 405 200 400 FLOPs (M)
50 40 30 0
255 302 200 400 FLOPs (M)
Fig. 2. Relationship between latency and FLOPs of ResNet variants on devices (From left to right are HP Probook, Jetson TX2, Jetson Nano, Pi and Note10).
Device n FL Process
Local Training
Aggregation
Local Training
Aggregation
Zerorizing
n
Latency Predictor
Latency Predictor
Training Epoch
Latency Predictor
Device 2
Recovering
Latency Predictor
Heterogeneous Federated Learning Aggregation
Device 1 (Extension)
Zerorizing Latency Predictor
Initializing
Heterogeneous Federated Learning Aggregation
0
Latency Predictor
Latency Guarantee
Local Training
Finalizing
Fig. 3. The process of Collate. The top depicts the latency guarantee component, while the bottom shows the heterogeneous FL component.
C. Latency Prediction
Optimizing models based on latency can better explore hardware characteristics, providing additional advantages in the trade-off between accuracy and latency [23]. To obtain high accuracy under latency constraints, we dynamically adjust the local model architectures during their local training, and the new model architecture needs to satisfy the latency constraint. However, measuring latency on a device will interrupt the training process and usually takes minutes [24], especially for off-device training [22]. And as shown in Fig. 2, the relationship between latency and floating-point operations (FLOPs) of a model is weak, making it unlikely to calculate the latency with an arithmetic function with respect to the FLOPs. A latency predictor is necessary for latency-critical systems. As the overhead of on-device measurement is huge, previous works for latency-critical systems use hardware simulator-based [25], lookup table (LUT)-based [23], [24], [26] or neural network-based [27], [28] predictors to estimate the latency of DNN on a certain device. Most commercial devices are black boxes to users, so it is difficult to emulate hardware by analyzing its resources and scheduling algorithms for building hardware simulators. On the other hand, LUTbased predictors calculate the latency of DNN by summing up the recorded latency of each layer, resulting in that it only suiting models made up of pre-defined layer structures. During our latency-critical learning process, the number of input channels and kernels, and other size-related parameters of each layer will be updated continuously to find the best model architecture under the latency constraint. It is impractical for LUT to preserve the whole design space. Thus, in this work, we follow [28] to build a hardware-customized latency predictor based on the Backpropagation (BP) [29] neural network.
III. M ETHODOLOGY In this section, we demonstrate our model learning framework, which collaboratively trains heterogeneous DNNs for different edge devices to meet their respective latency constraints simultaneously, while optimizing the model architecture for high accuracy. As shown in Fig. 3, this framework mainly includes two parts: the latency guarantee for each participant and the heterogeneous federated learning process. We also employ model extension for powerful clients, and we use the latency predictor to optimize this flow. Fig. 3 also illustrates the change of local models on each device during our training process, where we use kernels in a layer as an example to show this training process, and each oblique rectangle represents a kernel in this layer. The dashed line means this kernel is shielded in the current epoch. In the following, we first briefly introduce the hardware-customized latency predictor and the model extension. Then we present the latency guarantee, which is integrated into the local training of the FL process with almost no extra training overhead. Finally, we introduce our heterogeneous FL, mainly focusing on the heterogeneous aggregation algorithm. A. Latency Predictor & Model Extension Latency Predictor: Based on [28], we build a hardwarecustomized latency predictor using a three-layer BP network. The BP network is a lightweight neural network, and its forward propagation is shown in Eq. (4), where x is the vector of inputs, W L is the weight matrix, L is the index of layers, and σ L is the activation function at layer L. Thus, hundreds of multiplication and addition operations are enough to obtain the latency with the BP network. Also, the updating scheme of the BP network is simple, so its training is not time-consuming. g(x) = σ 2 (W 2 σ 1 (W 1 σ 0 (W 0 x)))
(4)
Latency > Constraint? Latency Predictor
NO
Latency Predictor
Latency Predictor
YES
Original Model
Model Duplication
YES
Model Exceeding Latency Constraint
Fig. 4. The model Extension scheme of Collate.
We first build a pool of models, which comprises thousands of single-layer models with various operations like Convolution and Pooling, configured with different configurations, to train the latency predictor. The latency of single-layer models is only a few milliseconds, so their measurement does not take much time. After training, by feeding the configuration of a layer, the latency predictor can predict its latency. For a real model, adding the latency of all layers in the latency-weighted longest path can get the whole model latency. The latency predictor is trained in advance for each device and is used in model learning. Assessing the latency predictor of Jetson TX2 with real models shows its variation is about 6.12%. Model Extension: In this work, we directly optimize models under the inference latency constraint, so Collate not only reduces latency for low-end systems but also extends models for powerful systems to improve accuracy. As illustrated in Fig. 4, when the latency of the original model is lower than the latency constraint, we can uniformly extend the model [30] on this device to obtain a larger model. The uniform extension means all layers are extended with a constant width multiplier, which enables the new model to exceed the system latency constraint. Then our following learning method can ensure the latency constraint is satisfied again. The model extension improves the accuracy for powerful devices (see Section IV), but it introduces heterogeneity among different local models. B. Latency Constraint Guarantee – Local Training As shown in Fig 3, the latency guarantee flow is integrated into the local training of the FL process. This flow can be divided into three phases: Initializing, Zeroizing, and Recovering. Algorithm 1 demonstrates the details of local training with the latency constraint guarantee. We first train the model for some communication rounds, called Initializing process, to let model parameters have informative values to reflect their contributions (Line 2-4) instead of randomly initialized numbers. The proto-training will be introduced in the next section. We design a trainable mask layer following each convolution layer, and each value δ in the mask layer is multiplied into the output of each kernel in its previous convolution layer. Thus, the contribution metric we define is that kernels with larger absolute mask values provide more impact on the final accuracy. And zeroizing the mask value means removing the corresponding kernel from the previous convolution layer. Note that mask values can be fused into the weights of kernels before inference to eliminate their overhead. In the
Algorithm 1: Latency Guarantee – LocalTrain Input: model, training settings, current round: c, start zero round: s, local epoch: le, latency constraint: l, prototype array: pt. Output: updated local model, output prototype 1: Load private training data; 2: if c < s then 3: train with proto(pt) for le epochs; 4: else 5: for e ← 1 to le do 6: train with proto(pt); 7: if e%2 ! = 0 or e is the last epoch then 8: Con rank ← Sort kernels by |δ|; 9: Shrink ratio ← Predictor(l, Con rank); 10: Find threshold of δ by Con rank, Shrink ratio; 11: Zeroize δi if δi < threshold; 12: end if 13: end for 14: end if
Zeroizing process, we first sort all contribution factors δ according to their absolute values to get the global contribution rank of kernels in the entire local model (Line 8). During the whole training process, these contribution factors are jointly optimized with the network weights, hence the training scheme can automatically identify the contribution of each kernel. Next, we calculate the contribution factor threshold (Line 10) by the latency predictor (Line 9) and the contribution rank of the entire model. Finally, we zeroize those δi corresponding to the unimportant parts with these thresholds (Line 11) to temporarily remove these weights for the current epoch. After this Zeroizing process, we extract a compact model that meets the latency constraint. Following each Zeroizing process, we propose a Recovering process (Line 7) to allow those zeroized parameters (i.e., δi ) to recover instead of being permanently removed. If any weights that are eliminated in the previous Zeroizing process but are potentially essential, the Recovering process can help them escape from zero and play a crucial part in the subsequent training process, the aggregation process, and even the inference processes. The Recovering process gives our method a chance to learn more efficient architecture with higher accuracy. Eq. (5) indicates the calculation formula in a convolution layer with δ and activation function, where L is the index of layers, AL is the output of layer L, W L is the weight between layer L − 1 and L, and σ L is the activation function at layer L. ⊛ denotes the convolution operation. AL = σ L (δ L · (AL−1 ⊛ W L ))
(5)
In the Recovering processing, we mainly focus on the updating of δ L that are zeroized during the Zeroizing process. Eq. (6) shows the gradients of δ L to the final loss function. ∂loss ∂loss L ′ = · σ · (AL−1 ⊛ W L ) ∂δ L ∂AL
(6)
Since we have zeroized some values in δ L and the most often used activation function is ReLU, the corresponding output of σ L is likewise zero under ReLU. Although ReLU is not differentiable at zero, it is widely accepted that its derivative is also zero. Thus, we can easily derive Eq. (7). ′
δL = 0 → σL = 0 →
∂loss =0 ∂δ L
(7)
...