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

# 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 HP-NTU Digital Manufacturing Corporate Lab (I1801E0028).
Source: [https://arxiv.org/html/2607.08013](https://arxiv.org/html/2607.08013)
## Collate: Collaborative Neural Network Learning for Latency\-Critical Edge Systems††thanks:© 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\. 627\-634, 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 HP\-NTU Digital Manufacturing Corporate Lab \(I1801E0028\)\.

Shuo Huai34, Di Liu4, Hao Kong34, Xiangzhong Luo3, Weichen Liu3, Ravi Subramaniam2, Christian Makaya2 and Qian Lin2

###### 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 proposeCollate, 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\.

## IIntroduction

Deep Neural Networks \(DNNs\) have brought significant breakthroughs in many different applications, such as image recognition and natural language processing\[[1](https://arxiv.org/html/2607.08013#bib.bib1)\]\. 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](https://arxiv.org/html/2607.08013#bib.bib1)\]\. Meanwhile, DNN models require a huge amount of training data to improve accuracy\[[2](https://arxiv.org/html/2607.08013#bib.bib2)\], 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](https://arxiv.org/html/2607.08013#bib.bib3)\]\. 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](https://arxiv.org/html/2607.08013#bib.bib4)\]\. With diverse 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](https://arxiv.org/html/2607.08013#S1.F1), when deploying the model into various edge systems, the latency differs from13\.613\.6ms to236\.0236\.0ms\. When this application has a latency constraint \(e\.g\.,30\.030\.0ms\), 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\.

![Refer to caption](https://arxiv.org/html/2607.08013v1/x1.png)Figure 1:Latency of the federated\-trained AlexNet on various edge systems\.Take a practical scenario from WeBank as an example\[[5](https://arxiv.org/html/2607.08013#bib.bib5)\]\. 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](https://arxiv.org/html/2607.08013#bib.bib6),[7](https://arxiv.org/html/2607.08013#bib.bib7),[8](https://arxiv.org/html/2607.08013#bib.bib8),[9](https://arxiv.org/html/2607.08013#bib.bib9),[10](https://arxiv.org/html/2607.08013#bib.bib10)\]\. 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](https://arxiv.org/html/2607.08013#bib.bib9)\]or an extra dataset for prototype learning\[[10](https://arxiv.org/html/2607.08013#bib.bib10)\]\. Moreover, the absence of the same global model lead to an accuracy drop of up to10%10\\%\[[11](https://arxiv.org/html/2607.08013#bib.bib11)\]\. 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](https://arxiv.org/html/2607.08013#S1.F1)\) to improve their accuracy\[[12](https://arxiv.org/html/2607.08013#bib.bib12)\], 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 ofCollateto discover the optimal DNN architecture for each edge system\.
- •We demonstrate the effectiveness ofCollatewith extensive experiments\. Compared to the state\-of\-the\-art methods and under the same latency constraints, our extended models can improve the accuracy by1\.96%1\.96\\%on average, and the accuracy of shrunk models outperforms others by3\.09%3\.09\\%on average\.

![Refer to caption](https://arxiv.org/html/2607.08013v1/x2.png)Figure 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\)\.![Refer to caption](https://arxiv.org/html/2607.08013v1/x3.png)Figure 3:The process ofCollate\. The top depicts the latency guarantee component, while the bottom shows the heterogeneous FL component\.
## IIBackground & Related Work

In this section, we introduce the preliminaries of FL, heterogeneous FL and latency prediction, including the advantages and disadvantages of some related work\.

### II\-AFederated Learning

In the traditional FL algorithm \(FedAvg\[[13](https://arxiv.org/html/2607.08013#bib.bib13)\]\), 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 everyl​elelocal training epochs\. Then the server executes model aggregation\. Considering this FL process includesnnsystems, the loss function for systemiiisfif\_\{i\}\. Letwwrepresent the model parameter, and then the training optimization goal becomes:

m​i​nwf​\(w\):=1n​∑i=1nfi​\(w\)\\mathop\{min\}\\limits\_\{w\}f\(w\):=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}f\_\{i\}\(w\)\(1\)And the training process can be formulated as:

wit\+\(e\+1\)=wit\+e−ηit\+e​∇fi​\(wit\+e\)w\_\{i\}^\{t\+\(e\+1\)\}=w\_\{i\}^\{t\+e\}\-\\eta\_\{i\}^\{t\+e\}\\nabla f\_\{i\}\(w\_\{i\}^\{t\+e\}\)\(2\)wt\+1=1n​∑i=1nwit\+l​ew^\{t\+1\}=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}w\_\{i\}^\{t\+le\}\(3\)wherewit\+ew\_\{i\}^\{t\+e\}denotes the local model parameter on systemiiafterttcommunication rounds andeelocal training steps\.η\\etais the learning rate, and∇\\nablarepresents the derivative\. From Eq\. \([3](https://arxiv.org/html/2607.08013#S2.E3)\), we can see that the traditional FL\[[13](https://arxiv.org/html/2607.08013#bib.bib13)\]is based on the fact that all local model architectures of different systems are the same for aggregation\.

### II\-BHeterogeneous Federated Learning

Two major challenges involved in the traditional FL process are statistical heterogeneity and hardware heterogeneity among different systems\[[14](https://arxiv.org/html/2607.08013#bib.bib14)\]\. Statistical heterogeneity means the data on different systems are not independent and identically distributed \(Non\-IID\)\[[15](https://arxiv.org/html/2607.08013#bib.bib15)\], 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\. Wanget al\.\[[16](https://arxiv.org/html/2607.08013#bib.bib16)\]proposed to fine\-tune some or all parameters of a trained global model using each client’s local data\. Jianget al\.\[[17](https://arxiv.org/html/2607.08013#bib.bib17)\]proposed to combine meta\-learning with FL for customizing models to recognize each local data pattern\. Khodaket al\.\[[18](https://arxiv.org/html/2607.08013#bib.bib18)\]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\. Caldaset al\.\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\]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\. Liet al\.\[[20](https://arxiv.org/html/2607.08013#bib.bib20)\]proposed Hermes to find small sub\-networks for some devices by applying the structured pruning\. Jeonget al\.\[[21](https://arxiv.org/html/2607.08013#bib.bib21)\]proposed to use knowledge distillation to reduce model size in some low\-end devices\. Xuet al\.\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\]proposed Helios to accelerate the devices with weak computational capacities by dynamically compressing the global model into an expected volume during the training stage\. Diaoet al\.\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]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](https://arxiv.org/html/2607.08013#bib.bib22)\]\. 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](https://arxiv.org/html/2607.08013#S4)\.

### II\-CLatency Prediction

Optimizing models based on latency can better explore hardware characteristics, providing additional advantages in the trade\-off between accuracy and latency\[[23](https://arxiv.org/html/2607.08013#bib.bib23)\]\. 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](https://arxiv.org/html/2607.08013#bib.bib24)\], especially for off\-device training\[[22](https://arxiv.org/html/2607.08013#bib.bib22)\]\. And as shown in Fig\.[2](https://arxiv.org/html/2607.08013#S1.F2), 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](https://arxiv.org/html/2607.08013#bib.bib25)\], look\-up table \(LUT\)\-based\[[24](https://arxiv.org/html/2607.08013#bib.bib24),[23](https://arxiv.org/html/2607.08013#bib.bib23),[26](https://arxiv.org/html/2607.08013#bib.bib26)\]or neural network\-based\[[27](https://arxiv.org/html/2607.08013#bib.bib27),[28](https://arxiv.org/html/2607.08013#bib.bib28)\]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, LUT\-based 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](https://arxiv.org/html/2607.08013#bib.bib28)\]to build a hardware\-customized latency predictor based on the Backpropagation \(BP\)\[[29](https://arxiv.org/html/2607.08013#bib.bib29)\]neural network\.

## IIIMethodology

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](https://arxiv.org/html/2607.08013#S1.F3), 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](https://arxiv.org/html/2607.08013#S1.F3)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\.

### III\-ALatency Predictor & Model Extension

Latency Predictor:Based on\[[28](https://arxiv.org/html/2607.08013#bib.bib28)\], we build a hardware\-customized 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](https://arxiv.org/html/2607.08013#S3.E4)\), wherexxis the vector of inputs,WLW^\{L\}is the weight matrix,LLis the index of layers, andσL\\sigma^\{L\}is the activation function at layerLL\. 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​\(W2​σ1​\(W1​σ0​\(W0​x\)\)\)g\(x\)=\\sigma^\{2\}\(W^\{2\}\\sigma^\{1\}\(W^\{1\}\\sigma^\{0\}\(W^\{0\}x\)\)\)\(4\)
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%\.

![Refer to caption](https://arxiv.org/html/2607.08013v1/x4.png)Figure 4:The model Extension scheme ofCollate\.Model Extension:In this work, we directly optimize models under the inference latency constraint, soCollatenot only reduces latency for low\-end systems but also extends models for powerful systems to improve accuracy\. As illustrated in Fig\.[4](https://arxiv.org/html/2607.08013#S3.F4), when the latency of the original model is lower than the latency constraint, we can uniformly extend the model\[[30](https://arxiv.org/html/2607.08013#bib.bib30)\]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](https://arxiv.org/html/2607.08013#S4)\), but it introduces heterogeneity among different local models\.

### III\-BLatency Constraint Guarantee – Local Training

As shown in Fig[3](https://arxiv.org/html/2607.08013#S1.F3), 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](https://arxiv.org/html/2607.08013#alg1)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δ\\deltain 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 Zeroizing process, we first sort all contribution factorsδ\\deltaaccording 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\\delta\_\{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\.

Input:model, training settings, current\_round:

cc, start\_zero round:

ss, local\_epoch:

l​ele, latency constraint:

ll, prototype array:

p​tpt\.

Output:updated local model, output prototype

1:Load private training data;

2:if

c<sc<sthen

3:

t​r​a​i​n​\_​w​i​t​h​\_​p​r​o​t​o​\(p​t\)train\\\_with\\\_proto\(pt\)for

l​eleepochs;

4:else

5:for

e←1e\\leftarrow 1to

l​eledo

6:

t​r​a​i​n​\_​w​i​t​h​\_​p​r​o​t​o​\(p​t\)train\\\_with\\\_proto\(pt\);

7:if

e%​2\!=0e\\%2\\ \!=0or

eeis the last epochthen

8:

C​o​n​\_​r​a​n​k←Con\\\_\{rank\}\\leftarrowSort kernels by

\|δ\|\|\\delta\|;

9:

S​h​r​i​n​k​\_​r​a​t​i​o←Shrink\\\_\{ratio\}\\leftarrowPredictor\(

l,C​o​n​\_​r​a​n​kl,Con\\\_\{rank\}\);

10:Find threshold of

δ\\deltaby

C​o​n​\_​r​a​n​kCon\\\_\{rank\},

S​h​r​i​n​k​\_​r​a​t​i​oShrink\\\_\{ratio\};

11:Zeroize

δi\\delta\_\{i\}if

δi\\delta\_\{i\}<<threshold;

12:endif

13:endfor

14:endif

Algorithm 1Latency Guarantee – LocalTrainFollowing each Zeroizing process, we propose a Recovering process \(Line 7\) to allow those zeroized parameters \(i\.e\.,δi\\delta\_\{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](https://arxiv.org/html/2607.08013#S3.E5)\) indicates the calculation formula in a convolution layer withδ\\deltaand activation function, whereLLis the index of layers,ALA^\{L\}is the output of layerLL,WLW^\{L\}is the weight between layerL−1L\-1andLL, andσL\\sigma^\{L\}is the activation function at layerLL\.⊛\\circledastdenotes the convolution operation\.

AL=σL​\(δL⋅\(AL−1⊛WL\)\)A^\{L\}=\\sigma^\{L\}\(\\delta^\{L\}\\cdot\(A^\{L\-1\}\\circledast W^\{L\}\)\)\(5\)In the Recovering processing, we mainly focus on the updating ofδL\\delta^\{L\}that are zeroized during the Zeroizing process\. Eq\. \([6](https://arxiv.org/html/2607.08013#S3.E6)\) shows the gradients ofδL\\delta^\{L\}to the final loss function\.

∂l​o​s​s∂δL=∂l​o​s​s∂AL⋅σL′⋅\(AL−1⊛WL\)\\frac\{\\partial loss\}\{\\partial\\delta^\{L\}\}=\\frac\{\\partial loss\}\{\\partial A^\{L\}\}\\cdot\{\\sigma^\{L\}\}^\{\{\}^\{\\prime\}\}\\cdot\(A^\{L\-1\}\\circledast W^\{L\}\)\(6\)
Since we have zeroized some values inδL\\delta^\{L\}and the most often used activation function isReLU, the corresponding output ofσL\\sigma^\{L\}is likewise zero underReLU\. AlthoughReLUis not differentiable at zero, it is widely accepted that its derivative is also zero\. Thus, we can easily derive Eq\. \([7](https://arxiv.org/html/2607.08013#S3.E7)\)\.

δL=0→σL′=0→∂l​o​s​s∂δL=0\\delta^\{L\}=0\\rightarrow\{\\sigma^\{L\}\}^\{\{\}^\{\\prime\}\}=0\\rightarrow\\frac\{\\partial loss\}\{\\partial\\delta^\{L\}\}=0\(7\)However, most training algorithms employ momentum in the optimizer to speed up convergence, as the momentum accumulates the gradients of the past steps to determine the direction to go in, rather than using only the gradient of the current step\[[31](https://arxiv.org/html/2607.08013#bib.bib31)\]\. Eq\. \([8](https://arxiv.org/html/2607.08013#S3.E8)\)\-\([9](https://arxiv.org/html/2607.08013#S3.E9)\) illustrates the updating rules of weight with momentum, in whichη\\etais the learning rate,zkz\_\{k\}is the updated value of the last step, andmmis the accumulation coefficient\. Thus,wkw\_\{k\}are easily updated towk\+1w\_\{k\+1\}by combining current and past gradients\.

zk\+1=m⋅zk\+∂l​o​s​s∂wkz\_\{k\+1\}=m\\cdot z\_\{k\}\+\\frac\{\\partial loss\}\{\\partial w\_\{k\}\}\(8\)wk\+1=wk−η⋅zk\+1w\_\{k\+1\}=w\_\{k\}\-\\eta\\cdot z\_\{k\+1\}\(9\)
![Refer to caption](https://arxiv.org/html/2607.08013v1/x5.png)Figure 5:The contribution values of a local model\.↑\\uparrowpoints to the communication round that starts zeroizing\.\}

packages each local training\.Therefore, even though the current gradients ofδL\\delta^\{L\}are zero, they can recover from zero by following the last updating directions\. This process is fully automated and efficient, as no additional steps are required\. These kernels zeroized in the previous zero epoch may become significant and we will not zeroize them in the future Zeroizing process\. Thus, the architecture of the local model is changed, as illustrated in Fig\.[5](https://arxiv.org/html/2607.08013#S3.F5)\. Moreover, the training process tries to minimize loss and improve accuracy, allowing the model architecture to improve for better accuracy\. Our local learning gradually cultivates an optimal architecture for each system through the training process to maximize accuracy under the latency constraint\.

Training Overhead:Compared to the local training of traditional FL, we add mask layers and a zeroizing process\. As they are both related to the kernels in convolution layers, we mainly focus on the convolution\. The time complexity of convolution layers for an input isO​\(∑l=1LMl2​Kl2​Cl−1​Cl\)O\(\\sum\_\{l=1\}^\{L\}M\_\{l\}^\{2\}K\_\{l\}^\{2\}C\_\{l\-1\}C\_\{l\}\), whereLLis the number of convolution layers,ClC\_\{l\}is the number of the kernels in layerll, that is, the number of output channels of this layer,MlM\_\{l\}is the output size of layerllandKlK\_\{l\}represents the kernel size of layerll\. Mask layers have a time complexity ofO​\(∑l=1LMl2​Cl\)O\(\\sum\_\{l=1\}^\{L\}M\_\{l\}^\{2\}C\_\{l\}\), which is much lower than convolution operations\. The zeroizing operation includes a sorting step with a time complexity ofO​\(n​l​o​g​n\),n=∑l=1LClO\(nlogn\),n=\\sum\_\{l=1\}^\{L\}C\_\{l\}\. Generally,l​o​g​nlognis much smaller thanMl2​Kl2​Cl−1M\_\{l\}^\{2\}K\_\{l\}^\{2\}C\_\{l\-1\}of each layerll, so its complexity is lower than the convolution layers\. The latency predictor only includes hundreds of multiplication and addition operations, which is much smaller than convolution operations\. Thus, the overhead of our local training can be neglected\.

### III\-CHeterogeneous Aggregation – Global Training

FL maintains a global model from all locally trained models, combining all patterns learned from each private dataset\. The parameters of local models are represented by \{Wl1,Wl2,…,WlnW\_\{l\}^\{1\},W\_\{l\}^\{2\},\.\.\.,W\_\{l\}^\{n\}\}, wherellmeans different layers, andnnis the total number of local models\. The global model receives all local model parameters and aggregates them into a global modelWg​lW\_\{gl\}\. As shown in Fig\.[3](https://arxiv.org/html/2607.08013#S1.F3), after several local training epochs, an aggregation process is conducted to combine all local parameters\. Each local training \- aggregation is a communication round, and it iterates multiple rounds\. In the traditional identical FL algorithm, the aggregation is formulated asWg​lt\+1=1n​∑i=1n\(Wli\)t\+l​eW\_\{gl\}^\{t\+1\}=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}\(W\_\{l\}^\{i\}\)^\{t\+le\}at the end of iterationtt\. At iterationt\+1t\+1,Wg​lt\+1W\_\{gl\}^\{t\+1\}is transmitted to each local model and sets the local model parameters as\(Wli\)t\+1=Wg​lt\+1\(W\_\{l\}^\{i\}\)^\{t\+1\}=W\_\{gl\}^\{t\+1\}\.

![Refer to caption](https://arxiv.org/html/2607.08013v1/x6.png)Figure 6:Heterogeneous federated learning aggregation scheme ofCollate\.In this work, our primary motivation is to learn models for each system to meet its latency constraint while obtaining as much accuracy as possible\. The details of our global training are shown in Algorithm[2](https://arxiv.org/html/2607.08013#alg2)\. We exploit model extension \(Line 4\) and shrinking \(Line 10\) to adjust each local model to suit its corresponding system, and these local models have a similar model architecture\. Parameters of the smaller model are a subset of parameters of the larger model\. Inspired by HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\], we first use an intuitive way to aggregate these heterogeneous local models \(Line 13\-16\), as shown in the rectangle in Fig\.[6](https://arxiv.org/html/2607.08013#S3.F6)\. The model architecture mainly includes two parts, where the heterogeneous layers are shrunk or extended to change the latency, and the identical layer is for proto\-correction\. The identical layer is a fully\-connected layer, and this layer has the same shape among all local models\. It can be aggregated following the traditional FL aggregation scheme as all local models share the same architecture\. We focus more on the aggregation of parametersWHW\_\{H\}in heterogeneous layers\. For notational convenience, we drop the iteration indextt\. Also, we aggregate parameters layer by layer, and we take one heterogeneous layer as an example\.WHiW\_\{H\}^\{i\}is the parameters of a heterogeneous layer of local modelii, andWHi∖WHjW\_\{H\}^\{i\}\\setminus W\_\{H\}^\{j\}means the parameters inWHiW\_\{H\}^\{i\}but not inWHjW\_\{H\}^\{j\},WHi∩WHjW\_\{H\}^\{i\}\\cap W\_\{H\}^\{j\}represents the parameters in bothWHiW\_\{H\}^\{i\}andWHjW\_\{H\}^\{j\}, whileWHi∪WHjW\_\{H\}^\{i\}\\cup W\_\{H\}^\{j\}represents the parameters inWHiW\_\{H\}^\{i\}orWHjW\_\{H\}^\{j\}\.nnis the total number of local models, andnan\_\{a\}is the number of local models that come before local modelaain the complexity rank \(Line 6\)\. Thus, the aggregation process can be represented in Eq\. \([10](https://arxiv.org/html/2607.08013#S3.E10)\)\-\([11](https://arxiv.org/html/2607.08013#S3.E11)\)\.

Wg​Ha,a\+1\\displaystyle W\_\{gH\}^\{a,a\+1\}=1na\+1​∑i=1n\(WHi∖WHa\+1∩WHa\),a<n\\displaystyle=\\frac\{1\}\{n\_\{a\+1\}\}\\sum\_\{i=1\}^\{n\}\(W\_\{H\}^\{i\}\\setminus W\_\{H\}^\{a\+1\}\\cap W\_\{H\}^\{a\}\),a<n\(10\)Wg​Hn,n\+1\\displaystyle W\_\{gH\}^\{n,n\+1\}=1n​∑i=1n\(WHi∩WHn\)\\displaystyle=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}\(W\_\{H\}^\{i\}\\cap W\_\{H\}^\{n\}\)Wg​H=Wg​H1,2∪Wg​H2,3∪⋯∪Wg​Hn,n\+1W\_\{gH\}=W\_\{gH\}^\{1,2\}\\cup W\_\{gH\}^\{2,3\}\\cup\\cdot\\cdot\\cdot\\cup W\_\{gH\}^\{n,n\+1\}\(11\)
Input:model, training settings:

t​sts, local\_epoch:

l​ele, start\_zero round:

ss, latency constraint:

ll, client number:

nn\.

Output:local models and their parameters

1:Initialize the original model and the prototype array

p​tpt;

2:Duplicate the model for each client;

3:for

m←1m\\leftarrow 1to

nndo

4:Extend

m​o​d​e​lmmodel\_\{m\}if its latency on

c​l​i​e​n​tmclient\_\{m\}less than

ll;

5:endfor

6:Sort models by their complexity in descending order;

7:for

c←1c\\leftarrow 1to

R​o​u​n​dm​a​xRound\_\{max\}do

8:for

m←1m\\leftarrow 1to

nnin paralleldo

9:Send

m​o​d​e​lmmodel\_\{m\}to the corresponding

c​l​i​e​n​tmclient\_\{m\};

10:

m​o​d​e​lm′,p​tm←model\_\{m\}^\{\{\}^\{\\prime\}\},pt\_\{m\}\\leftarrowLocalTrain\(

m​o​d​e​lm,t​s,c,s,l​e,l,p​tmodel\_\{m\},ts,c,s,le,l,pt\);

11:endfor

12:if

c≠R​o​u​n​dm​a​xc\\neq Round\_\{max\}then

13:Derive each part of

Wg​HW\_\{gH\}by Eq\. \([10](https://arxiv.org/html/2607.08013#S3.E10)\) and all

m​o​d​e​lm′model\_\{m\}^\{\{\}^\{\\prime\}\};

14:for

m←1m\\leftarrow 1to

nndo

15:

m​o​d​e​lm←∪mnWg​Hm,m\+1model\_\{m\}\\leftarrow\\cup\_\{m\}^\{n\}W\_\{gH\}^\{m,m\+1\};

16:endfor

17:Merge all

p​tmpt\_\{m\}and update

p​tptby Eq\. \([12](https://arxiv.org/html/2607.08013#S3.E12)\);

18:endif

19:endfor

Algorithm 2FL Aggregation – GlobalTrainHowever, this method may lead to lower accuracy in a larger model architecture\. Take Fig\.[6](https://arxiv.org/html/2607.08013#S3.F6)as an example, the parameters inWg​H1,2W\_\{gH\}^\{1,2\}are only trained on private dataset11\. These parameters impact the convergence of local mode11severely, reducing its accuracy, particularly on the Non\-IID data\. Thus, we also proposed to use the prototype scheme\[[32](https://arxiv.org/html/2607.08013#bib.bib32)\]to correct these parameters and enable these parameters to learn the pattern in other datasets \(Line 10, 17\)\. As shown in the circle of Fig\.[6](https://arxiv.org/html/2607.08013#S3.F6),OmiO\_\{m\}^\{i\}is the output of the identical fully\-connected layer near the final layer, and it can be used to represent a class by a high\-dimension vector\. This vector is called the prototype of this class\. As shown in Eq\. \([12](https://arxiv.org/html/2607.08013#S3.E12)\), we first aggregate all prototypes by class from all local models at training iterationtt\. Then, in the next training iterationt\+1t\+1, we add a penalty item between the current prototype of a class and the prototype of this class from previous iterations into the loss function, as illustrated in Eq\. \([13](https://arxiv.org/html/2607.08013#S3.E13)\), whereL1L\_\{1\}represents the loss function of the model outputs to labels,L2L\_\{2\}means the loss function of the prototype andβ\\betais a pre\-defined coefficient\. This is the proto\-training used in Algorithm[1](https://arxiv.org/html/2607.08013#alg1)\. In this way, larger models can perceive patterns of all private datasets to improve accuracy\.

t:proto\[lj\]t=1n∑i=1nOmi\[label==lj\],lj∈labelt:proto\[l\_\{j\}\]^\{t\}=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}O\_\{m\}^\{i\}\[label==l\_\{j\}\],l\_\{j\}\\in label\(12\)t\+1:L​o​s​s=L1​\(l​a​b​e​l,o​u​t\)\+β​L2​\(Om,p​r​o​t​ot\)t\+1:Loss=L\_\{1\}\(label,out\)\+\\beta L\_\{2\}\(O\_\{m\},proto^\{t\}\)\(13\)
![Refer to caption](https://arxiv.org/html/2607.08013v1/x7.png)
![Refer to caption](https://arxiv.org/html/2607.08013v1/x8.png)

Figure 7:The test accuracy of LeNet on MNIST under different complexity levels \(From left to right are1\.51\.5,1\.251\.25,1\.01\.0,0\.50\.5,0\.250\.25\)\. The top five show results on IID data and the bottom five are on Non\-IID data\. • points to the communication round that starts the zeroizing\-recovering process ofCollate\.Training Overhead:The global training adds the proto\-corrected scheme\. Derived from Eq\. \([12](https://arxiv.org/html/2607.08013#S3.E12)\)\-\([13](https://arxiv.org/html/2607.08013#S3.E13)\), the time complexity of this scheme isO​\(I​V\)O\(IV\), whereIIis the total number of training data andVVis the dimension ofOmiO\_\{m\}^\{i\}\. This complexity is not comparable to a traditional aggregation scheme, which is only related to the number of model parameters and clients\. But when considering the local training, its complexity is much lower thanO​\(I​∑l=1LMl2​Kl2​Cl−1​Cl\)O\(I\\sum\_\{l=1\}^\{L\}M\_\{l\}^\{2\}K\_\{l\}^\{2\}C\_\{l\-1\}C\_\{l\}\)\. The model extension introduces overhead to the learning process, but it is intended to improve accuracy for powerful systems\. And the model does not belong to the training algorithm, so the model extension is not included in the overhead of the training algorithm\.

## IVExperiments

In this section, we first present details of our implementation and the experimental environment, including hardware platforms, benchmark models and datasets\. Then we discuss the performance ofCollate, compared to state\-of\-the\-art methods: Expanding the Reach of Federated Learning \(ERFL\)\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\], Helios\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\]and HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]\. Their brief is in Section[II](https://arxiv.org/html/2607.08013#S2)\.

### IV\-AExperimental Setup

Hardware Platforms:We collaboratively train models for five distinct devices to evaluateCollate\. Each device and its inferred model constitute an edge system\. These devices are HP ProBook 440 G6, NVIDIA Jetson TX2, NVIDIA Jetson Nano, Raspberry Pi 4B and Samsung Galaxy Note10\. HP ProBook is equipped with an Intel Core i5\-8265U CPU with 8GB of memory\. TX2 is equipped with an NVIDIA Pascal GPU and 8GB of memory\. Nano has an NVIDIA Maxwell GPU and 4GB of memory\. Pi owns an ARM Cortex\-A72 embedded CPU and 4GB of memory, and Note10 is equipped with a Qualcomm Snapdragon 855 processor and 8GB of memory\. Also,Collatedoes not limit the number of clients\.

Datasets and Models:The datasets used for the benchmark are MNIST, CIFAR\-10, CIFAR\-100 and Human Activity Recognition \(HAR\)\[[33](https://arxiv.org/html/2607.08013#bib.bib33)\], which are widely used in evaluating Federated Learning\[[7](https://arxiv.org/html/2607.08013#bib.bib7),[8](https://arxiv.org/html/2607.08013#bib.bib8),[6](https://arxiv.org/html/2607.08013#bib.bib6),[9](https://arxiv.org/html/2607.08013#bib.bib9),[10](https://arxiv.org/html/2607.08013#bib.bib10),[19](https://arxiv.org/html/2607.08013#bib.bib19),[20](https://arxiv.org/html/2607.08013#bib.bib20),[21](https://arxiv.org/html/2607.08013#bib.bib21)\]\. The first three are image classification datasets, which are popular computer vision applications\. And HAR is an attractive feature for smartphones using data collected from different types of onboard sensors\. We split the training data of each dataset into five parts, and each device owns a part and keeps it private\. The trained model is evaluated on all test data of the corresponding dataset\. For the generation of IID data, we randomly split the training data, and each class in a dataset exists on all five devices\. And for the generation of Non\-IID data, we force each class of a dataset to be divided into at most two parts, thus, a class only exists on at most two devices\. The used models are LeNet, AlexNet, ResNet\-18 and IMU\[[34](https://arxiv.org/html/2607.08013#bib.bib34)\], and they are trained on MNIST, CIFAR\-10, CIFAR\-100 and HAR, respectively\.

Implementation Details:The proposed method is implemented using the Pytorch library\[[35](https://arxiv.org/html/2607.08013#bib.bib35)\]\. We set the local training epoch to66, and the communication rounds are2020,3030,4040and100100for LeNet, IMU, AlexNet and ResNet\-18, respectively\. And for all datasets, the optimizer is SGD with a momentum of0\.90\.9and no weight decay\. The learning rates are set to0\.050\.05,0\.0030\.003,0\.0010\.001and0\.010\.01for LeNet, IMU, AlexNet and ResNet\-18, respectively, which are widely used for these models\. To exhibit the superiority ofCollateunder the latency constraint, we re\-implement ERFL\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\], Helios\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\]and HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]according to their papers or open\-source codes\. We use the latency predictor for each method and modify them to guarantee the latency constraint of each local model\. And we use the same training settings for all methods\.

### IV\-BAccuracy Evaluation

Our method can also learn models with pre\-defined complexities, reflected by changing theS​h​r​i​n​k​\_​r​a​t​i​oShrink\\\_ratioto a constant in local training \(Line 9of Algorithm[1](https://arxiv.org/html/2607.08013#alg1)\)\. Due to HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]using pre\-defined complexities to learn models, their method can easily include the model extension\. We follow their work to set five discrete complexity levels of1\.51\.5,1\.251\.25,1\.01\.0,0\.50\.5,0\.250\.25to evaluateCollateand HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]\. However, ERFL\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\]and Helios\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\]can only shrink the model size, thus, we use1\.01\.0,0\.50\.5,0\.250\.25complexity levels for these two works\.

TABLE I:The test accuracy on HAR under different complexity levels\.Accuracy \(%\) of Different ComplexitiesDataDISTMethodC\-1\.50xC\-1\.25xC\-1\.00xC\-0\.50xC\-0\.25xOriginalN\.A\.N\.A\.94\.53N\.A\.N\.A\.ERFL\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\]N\.A\.N\.A\.93\.8789\.6481\.53Helios\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\]N\.A\.N\.A\.94\.3391\.2888\.72HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]93\.5794\.6094\.8193\.6390\.17IIDCollate96\.4096\.3394\.4094\.2691\.68OriginalN\.A\.N\.A\.76\.33N\.A\.N\.A\.ERFL\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\]N\.A\.N\.A\.74\.2171\.2666\.05Helios\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\]N\.A\.N\.A\.75\.1671\.6367\.87HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]75\.0375\.7075\.6773\.6869\.32Non\-IIDCollate77\.5478\.1576\.3875\.5771\.53Results on IID data:The top five in Fig\.[7](https://arxiv.org/html/2607.08013#S3.F7)show thatCollateachieves the highest accuracy compared to HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\], Helios\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\], and ERFL\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\]under different complexity levels for the MNIST dataset\. The original method is the traditional identical FL\[[13](https://arxiv.org/html/2607.08013#bib.bib13)\], and it obtains only one global model with the same complexity \(1\.01\.0\)\. We put its same curve in these five sub\-figures for easy comparison\. Compared with HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]on model extension,Collatewith the proto\-corrected scheme can achieve higher accuracy than the original method, while the accuracy of HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]is lower than the original method\. Compared with Helios\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\]and ERFL\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\],Collatecan achieve the highest accuracy in the shrunk models through the well\-designed Zeroizing\-Recovering process\. Table[I](https://arxiv.org/html/2607.08013#S4.T1)\(IID part\) presents that on the HAR dataset,Collatecan also perform well in most cases\. All highest accuracy is marked inboldface\. The accuracy of our extended models can outperform the original one by about2%2\\%and our shrunk models achieve the lowest accuracy loss\.

Results on Non\-IID data:The bottom five in Fig\.[7](https://arxiv.org/html/2607.08013#S3.F7)demonstrate thatCollatestill performs well on the Non\-IID data of the MNIST dataset\. Different from IID data, for each method and each dataset, we use the last10%10\\%of the total communication rounds to fine\-tune each local model in a traditional FL way\[[13](https://arxiv.org/html/2607.08013#bib.bib13)\]\. This is because final models are generated after local training, and aggregation updates zeroized weights into non\-zero in all methods \(except HeretoFL, but for fairness, we also fine\-tune HeteroFL\)\. Updating only by the local dataset impacts its accuracy on all test sets, especially under Non\-IID data\. Fig\.[7](https://arxiv.org/html/2607.08013#S3.F7)also shows that some of our models perform well without fine\-tuning by the proto\-corrected scheme\. Table[I](https://arxiv.org/html/2607.08013#S4.T1)\(Non\-IID part\) shows that on the HAR dataset,Collatecan still achieve the highest accuracy\.

TABLE II:The test accuracy on Cifar\-10 under the latency constraint\.
### IV\-CLatency\-Accuracy Evaluation

To demonstrate the effectiveness of our method for multiple latency\-critical edge systems, we evaluate the accuracy ofCollateand state\-of\-the\-art heterogeneous FL methods under a pre\-defined latency constraint\. As Helios used Nvidia Jetson Nano as the non\-straggler\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\], we also pick the original model inference latency on Jetson Nano as the constraint, which means we adjust all local models to achieve a similar latency as that of the original corresponding model on Jetson Nano\.

Table[II](https://arxiv.org/html/2607.08013#S4.T2)presents that compared with ERFL\[[19](https://arxiv.org/html/2607.08013#bib.bib19)\], Helios\[[7](https://arxiv.org/html/2607.08013#bib.bib7)\]and HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\],Collatecan achieve the best accuracy for almost all the five platforms on IID or Non\-IID data of the CIFAR\-10 dataset\. As HeteroFL\[[8](https://arxiv.org/html/2607.08013#bib.bib8)\]and our method can extend models to larger than the original ones, while others can only shrink models, we useEE\-flag to represent model extension\. Additionally, each local model in Table[II](https://arxiv.org/html/2607.08013#S4.T2)can meet the latency constraint \(3838ms\) and achieve a latency in the range of36\.8±0\.7​m​s36\.8\\pm 0\.7ms\. For methods without extension, the latency of their models is lower than the latency constraint on powerful devices\. Compared to the original model, for the IID data, we extend the model on HP ProBook from about23\.023\.0ms and improve the accuracy by1\.02%1\.02\\%\. Also, the latency of the original model on Note10 is around82\.682\.6ms, and we reduce its latency by about54\.0%54\.0\\%with only a1\.00%1\.00\\%accuracy drop\. For the Non\-IID data, since it is more difficult to learn the pattern of the dataset, the benefit ofCollateis less than that of IID data, but our method still outperforms other methods\.

Table[III](https://arxiv.org/html/2607.08013#S5.T3)shows the evaluation results ofCollateon the CIFAR\-100 dataset\. Each local model in Table[III](https://arxiv.org/html/2607.08013#S5.T3)is under a latency constraint of about7070ms and they achieve a latency in the range of67\.8±1\.8​m​s67\.8\\pm 1\.8ms\. Among these models,Collateobtains the best accuracy under all conditions\. Compared to the original model, we extend it on Jetson TX2 from about27\.427\.4ms and obtain a2\.02%2\.02\\%accuracy improvement\. Also, the latency of the original model on Pi is around152\.6152\.6ms, and we reduce its latency by about54\.1%54\.1\\%with only a1\.71%1\.71\\%accuracy drop\. Specifically, from Table[II](https://arxiv.org/html/2607.08013#S4.T2)\-[III](https://arxiv.org/html/2607.08013#S5.T3), we can conclude that, compared to other methods\[[7](https://arxiv.org/html/2607.08013#bib.bib7),[8](https://arxiv.org/html/2607.08013#bib.bib8),[19](https://arxiv.org/html/2607.08013#bib.bib19)\],Collateimproves the accuracy by1\.96%1\.96\\%on average for extended models, and our shrunk models can obtain a3\.09%3\.09\\%accuracy improvement on average\.

## VConclusion

In this work, we have presented a collaborative neural network learning framework for multiple latency\-critical edge systems to directly satisfy the latency constraint of DNNs on their devices\. The collaboration augmented the training data for each system to improve accuracy while also protecting their data privacy\. Our framework can satisfy the latency constraints of each system with one training process\. It contained two components: a dynamic zeroizing\-recovering algorithm and a proto\-corrected aggregation scheme\. The dynamic algorithm can extract optimal model architectures to achieve less accuracy drop under a large latency reduction\. Our aggregation scheme improved the accuracy of extended models\. We also included the latency predictor and model extension to guide and optimize this framework\. Experiments verified that our framework has obtained remarkable improvements\.

TABLE III:The test accuracy on Cifar\-100 under the latency constraint\.
## References

- \[1\]S\. Pouyanfar, S\. Sadiq, Y\. Yan, H\. Tian, Y\. Tao, M\. P\. Reyes, M\.\-L\. Shyu, S\.\-C\. Chen, and S\. S\. Iyengar, “A survey on deep learning: Algorithms, techniques, and applications,”*ACM Computing Surveys*, 2018\.
- \[2\]X\. Zhu, C\. Vondrick, C\. C\. Fowlkes, and D\. Ramanan, “Do we need more training data?”*International Journal of Computer Vision*, 2016\.
- \[3\]Q\. Yang, Y\. Liu, T\. Chen, and Y\. Tong, “Federated machine learning: Concept and applications,”*ACM Transactions on Intelligent Systems and Technology \(TIST\)*, vol\. 10, no\. 2, pp\. 1–19, 2019\.
- \[4\]S\. Wang, G\. Ananthanarayanan, Y\. Zeng, N\. Goel, A\. Pathania, and T\. Mitra, “High\-throughput cnn inference on embedded arm big\. little multicore processors,”*IEEE Transactions on Computer\-Aided Design of Integrated Circuits and Systems*, vol\. 39, no\. 10, pp\. 2254–2267, 2019\.
- \[5\]Y\. Cheng, Y\. Liu, T\. Chen, and Q\. Yang, “Federated learning for privacy\-preserving ai,”*Communications of the ACM*, 2020\.
- \[6\]J\. Zhang, S\. Guo, X\. Ma, H\. Wang, W\. Xu, and F\. Wu, “Parameterized knowledge transfer for personalized federated learning,”*Advances in Neural Information Processing Systems*, vol\. 34, 2021\.
- \[7\]Z\. Xu, F\. Yu, J\. Xiong, and X\. Chen, “Helios: heterogeneity\-aware federated learning with dynamically balanced collaboration,” in*2021 58th ACM/IEEE Design Automation Conference \(DAC\)*\. IEEE, 2021\.
- \[8\]E\. Diao, J\. Ding, and V\. Tarokh, “Heterofl: Computation and communication efficient federated learning for heterogeneous clients,”*arXiv preprint arXiv:2010\.01264*, 2020\.
- \[9\]D\. Li and J\. Wang, “Fedmd: Heterogenous federated learning via model distillation,”*arXiv preprint arXiv:1910\.03581*, 2019\.
- \[10\]Y\. Tan, G\. Long, L\. Liu, T\. Zhou, Q\. Lu, J\. Jiang, and C\. Zhang, “Fedproto: Federated prototype learning across heterogeneous clients,” in*AAAI Conference on Artificial Intelligence*, vol\. 1, 2022\.
- \[11\]B\. Yuan, C\. R\. Wolfe, C\. Dun, Y\. Tang, A\. Kyrillidis, and C\. M\. Jermaine, “Distributed learning of deep neural networks using independent subnet training,”*arXiv preprint arXiv:1910\.02120*, 2019\.
- \[12\]M\. Tan and Q\. Le, “Efficientnet: Rethinking model scaling for convolutional neural networks,” in*International conference on machine learning*\. PMLR, 2019, pp\. 6105–6114\.
- \[13\]B\. McMahan, E\. Moore, D\. Ramage, S\. Hampson, and B\. A\. y Arcas, “Communication\-efficient learning of deep networks from decentralized data,” in*Artificial intelligence and statistics*\. PMLR, 2017\.
- \[14\]T\. Li, A\. K\. Sahu, A\. Talwalkar, and V\. Smith, “Federated learning: Challenges, methods, and future directions,”*IEEE Signal Processing Magazine*, vol\. 37, no\. 3, pp\. 50–60, 2020\.
- \[15\]Y\. Zhao, M\. Li, L\. Lai, N\. Suda, D\. Civin, and V\. Chandra, “Federated learning with non\-iid data,”*arXiv preprint arXiv:1806\.00582*, 2018\.
- \[16\]K\. Wang, R\. Mathews, C\. Kiddon, H\. Eichner, F\. Beaufays, and D\. Ramage, “Federated evaluation of on\-device personalization,”*arXiv preprint arXiv:1910\.10252*, 2019\.
- \[17\]Y\. Jiang, J\. Konečnỳ, K\. Rush, and S\. Kannan, “Improving federated learning personalization via model agnostic meta learning,”*arXiv preprint arXiv:1909\.12488*, 2019\.
- \[18\]M\. Khodak, M\.\-F\. F\. Balcan, and A\. S\. Talwalkar, “Adaptive gradient\-based meta\-learning methods,”*Advances in Neural Information Processing Systems*, vol\. 32, 2019\.
- \[19\]S\. Caldas, J\. Konečny, H\. B\. McMahan, and A\. Talwalkar, “Expanding the reach of federated learning by reducing client resource requirements,”*arXiv preprint arXiv:1812\.07210*, 2018\.
- \[20\]A\. Li, J\. Sun, P\. Li, Y\. Pu, H\. Li, and Y\. Chen, “Hermes: an efficient federated learning framework for heterogeneous mobile clients,” in*Proceedings of the 27th Annual International Conference on Mobile Computing and Networking*, 2021, pp\. 420–437\.
- \[21\]E\. Jeong*et al\.*, “Communication\-efficient on\-device machine learning: Federated distillation and augmentation under non\-iid private data,”*arXiv preprint arXiv:1811\.11479*, 2018\.
- \[22\]D\. S\. Wibowo and S\. J\. Fong, “Mobile blockchain\-empowered federated learning: Current situation and further prospect,” in*2021 Third International Conference on Blockchain Computing and Applications \(BCCA\)*\. IEEE, 2021, pp\. 19–25\.
- \[23\]T\.\-J\. Yang*et al\.*, “Netadapt: Platform\-aware neural network adaptation for mobile applications,” in*ECCV*, 2018\.
- \[24\]X\. Dai*et al\.*, “Chamnet: Towards efficient network design through platform\-aware model adaptation,” in*CVPR*, 2019\.
- \[25\]A\. Gholami*et al\.*, “Squeezenext: Hardware\-aware neural network design,” in*CVPRW*, 2018\.
- \[26\]B\. Wu*et al\.*, “Fbnet: Hardware\-aware efficient convnet design via differentiable neural architecture search,” in*CVPR*, 2019\.
- \[27\]D\. Justus*et al\.*, “Predicting the computational cost of deep learning models,” in*2018 IEEE international conference on big data \(Big Data\)*\. IEEE, 2018, pp\. 3873–3882\.
- \[28\]S\. Huai*et al\.*, “Zerobn: Learning compact neural networks for latency\-critical edge systems,” in*2021 58th ACM/IEEE Design Automation Conference \(DAC\)*\. IEEE, 2021, pp\. 151–156\.
- \[29\]R\. Hecht\-Nielsen, “Theory of the backpropagation neural network,” in*Neural Networks*, 1992\.
- \[30\]A\. Gordon, E\. Eban, O\. Nachum, B\. Chen, H\. Wu, T\.\-J\. Yang, and E\. Choi, “Morphnet: Fast & simple resource\-constrained structure learning of deep networks,” in*Proceedings of the IEEE conference on computer vision and pattern recognition*, 2018, pp\. 1586–1595\.
- \[31\]G\. Goh, “Why momentum really works,”*Distill*, 2017\.
- \[32\]W\. Xu, Y\. Xian, J\. Wang, B\. Schiele, and Z\. Akata, “Attribute prototype network for zero\-shot learning,”*Advances in Neural Information Processing Systems*, vol\. 33, pp\. 21 969–21 980, 2020\.
- \[33\]D\. Anguita, A\. Ghio, L\. Oneto, X\. Parra Perez, and J\. L\. Reyes Ortiz, “A public domain dataset for human activity recognition using smartphones,” in*Proceedings of the 21th international European symposium on artificial neural networks, computational intelligence and machine learning*, 2013, pp\. 437–442\.
- \[34\]F\. Moya Rueda*et al\.*, “Convolutional neural networks for human activity recognition using body\-worn sensors,” in*Informatics*, vol\. 5, no\. 2\. Multidisciplinary Digital Publishing Institute, 2018, p\. 26\.
- \[35\]A\. Paszke*et al\.*, “Pytorch: An imperative style, high\-performance deep learning library,” in*NIPS*, 2019\.