CoFormer: Collaborating with Heterogeneous Edge Devices for Scalable Transformer Inference

TL;DR

CoFormer is a collaborative inference system that decomposes large transformer models into smaller parts for distributed execution on heterogeneous edge devices, significantly improving speed, reducing memory, and lowering energy consumption.

Contribution

It introduces a novel optimization and calibration framework for decomposing transformers, enabling efficient edge inference with minimal accuracy loss.

Findings

Achieves up to 3.1× inference speedup on large transformers.

Supports GPT2-XL inference on edge devices with 76.3% memory reduction.

Reduces energy consumption by approximately 40%.

Abstract

The impressive performance of transformer models has sparked the deployment of intelligent applications on resource-constrained edge devices. However, ensuring high-quality service for real-time edge systems is a significant challenge due to the considerable computational demands and resource requirements of these models. Existing strategies typically either offload transformer computations to other devices or directly deploy compressed models on individual edge devices. These strategies, however, result in either considerable communication overhead or suboptimal trade-offs between accuracy and efficiency. To tackle these challenges, we propose a collaborative inference system for general transformer models, termed CoFormer. The central idea behind CoFormer is to exploit the divisibility and integrability of transformer. An off-the-shelf large transformer can be decomposed into multiple smaller models for distributed inference, and their intermediate results are aggregated to generate the final output. We formulate an optimization problem to minimize both inference latency and accuracy degradation under heterogeneous hardware constraints. DeBo algorithm is proposed to first solve the optimization problem to derive the decomposition policy, and then progressively calibrate decomposed models to restore performance. We demonstrate the capability to support a wide range of transformer models on heterogeneous edge devices, achieving up to 3.1$\times$ inference speedup with large transformer models. Notably, CoFormer enables the efficient inference of GPT2-XL with 1.6 billion parameters on edge devices, reducing memory requirements by 76.3\%. CoFormer can also reduce energy consumption by approximately 40\% while maintaining satisfactory inference performance.

Tables

Table 1. TABLE I: Comparison of inference performance and latency for single-edge solution on ImageNet-1K.

Model	Device	Accuracy	Latency
EfficientFormer L7[26]	Jetson TX2	83.30%	145.84$\pm$4.27ms
	Jetson Nano		374.62$\pm$6.41ms
MobileViTv2_200[17]	Jetson TX2	81.17%	74.32$\pm$3.58ms
	Jetson Nano		180.83$\pm$4.99ms

Table 2. TABLE II: Comparison between our method and efficient transformer models with approximate FLOPs on ImageNet-1K.

Method	FLOPs $\downarrow$	Memory $\downarrow$	#Params. $\downarrow$	Accuracy $\uparrow$	Latency $\downarrow$	Speed-up $\uparrow$	Energy $\downarrow$	Reduction $\uparrow$
PoolFormer-M48	23.2 G	4.39 GB	56.0 M	82.50 %	176.87 $\pm$0.38 ms	1.00 $\times$	1156.93 $\pm$86.3 mJ	0.00 %
EfficientFormer L7	20.4 G	4.31 GB	82.1 M	83.30 %	128.84 $\pm$0.33 ms	1.37 $\times$	897.46 $\pm$23.1 mJ	22.43 %
T2T-ViTt-19	19.6 G	2.13 GB	39.2 M	81.90 %	119.70 $\pm$0.43 ms	1.48 $\times$	764.88 $\pm$46.5 mJ	33.89 %
CoFormer+Swin	20.3 G	2.09 GB	68.1 M	83.61 %	98.30 $\pm$0.97 ms	1.80 $\times$	705.12 $\pm$52.1 mJ	39.05 %
PoolFormer-M36	17.6 G	4.31 GB	56.0 M	82.10 %	134.32 $\pm$0.45 ms	1.00 $\times$	858.94 $\pm$63.7 mJ	0.00 %
T2T-ViT-19	17.0 G	2.12 GB	39.2 M	81.90 %	99.06 $\pm$0.57 ms	1.36 $\times$	609.22 $\pm$31.2 mJ	29.07 %
MobileViTv2-200	15.0 G	3.87 GB	18.5 M	81.17 %	74.32 $\pm$0.84 ms	1.81 $\times$	500.92 $\pm$26.9 mJ	41.68 %
CoFormer+DeiT	14.4 G	2.10 GB	37.8 M	82.26 %	54.89 $\pm$1.62 ms	2.45 $\times$	424.57 $\pm$31.8 mJ	50.57 %

Table 3. TABLE III: Ablation study results of main components on CIFAR-100. The ”✔” mark indicating we adopt the corresponding method. The second row represents the evaluation results of three decomposed sub-models.

Decompose	Aggregate	Accuracy (%)			Latency (ms)
✗	✗	91.3			123.5
✔	✗	52.2	68.7	76.5	34.2	25.9	16.4
✔	✔	90.3			51.8

Table 4. TABLE IV: Comparison of different aggregating methods. CoFormer achieves good balance between accuracy and latency.

Aggregating Method	Accuracy	Latency
DeiT-B	83.40 %	126.91 $\pm$0.46 ms
Average [30]	81.02 %	54.84 $\pm$0.53 ms
Majority Voting [30]	80.71 %	54.42 $\pm$0.24 ms
Attention [41]	82.53 %	61.44 $\pm$0.49 ms
SENet [42]	82.67 %	58.74 $\pm$0.38 ms
CoFormer	82.26 %	54.89 $\pm$0.62 ms

Table 5. TABLE V: Results of different device number on CIFAR-100. ”Num.=1” represents the inference of DeiT-B in Jetson TX2.

Num.	Accuracy	Latency	Energy
1	91.30 %	123.48 $\pm$0.43 ms	760.64 $\pm$57.3 mJ
2	90.36 %	85.64 $\pm$0.52 ms	618.94 $\pm$38.2 mJ
3	90.27 %	51.84 $\pm$0.81 ms	409.55 $\pm$42.7 mJ
4	90.03 %	45.48 $\pm$0.65 ms	367.46 $\pm$24.1 mJ

Table 6. TABLE VI: Notation.

Symbol	Description
$L$	The layer number of the large ViT.
$d$	The embedding dimension of the large ViT.
$h$	The number of attention heads in the large ViT.
$D$	The neuron dimension in the MLP block of the large ViT.
$N$	The number of heterogeneous edge devices.
$l_n$	The layer number of the sub-model $n$.
$d_n$	The embedding dimension of the sub-model $n$.
$h_n^k$	The number of attention heads in the $k$-th layer of the sub-model $n$.
$h_n^{1:l_n}$	The set of attention head numbers from the first layer to the $k$-th layer of the sub-model $n$.
$D_n^k$	The neuron dimension in the $k$-th MLP block of the sub-model $n$.
$D_n^{1:l_n}$	The neuron dimension from the first MLP block to the $k$-th block of the sub-model $n$.
$C_n$	The architecture specifications of sub-model $n$.
$𝒞$	The set of decomposition decisions.
$𝒩$	The set of sub-models.
$f(\cdot )$	The latency predictor.
${𝐗}_n$	The intermediate features of sub-model $n$.
$r_n$	The data transmission rate of the edge device $n$.
${𝐗}_{agg}$	The aggregated features from all sub-models.
$𝐖$	The weight of linear transformation in the Results Aggregation.
$g$	The available computing power for the central node $i$.
$t_n^1$	The latency of the Phase 1: Backbone Forward of sub-model $n$.
$t_n^2$	The latency of the Phase 2: Data Transmission of sub-model $n$.
$t_n^3$	The latency of the Phase 3: Results Aggregation of sub-model $n$.
$T$	The end-to-end inference latency.
${𝒟}_{train}$	The training dataset.
${𝒟}_{val}$	The validation dataset.
${ℒ}_n$	The validation loss of sub-model $n$
$ℒ$	The overall accuracy degradation for collaborative inference.
$\delta$	The balancing hyperparameter.
${\Omega }_n$	The maximum amount of computing resources for the edge device $n$.
${\Phi }_n$	The maximum amount of memory capacity for the edge device $n$.
$\omega \left(C_n\right)$	The floating point operations of the sub-model $n$.
$\varphi \left(C_n\right)$	The memory requirements of the sub-model $n$.
$\Omega$	The set of all sub-models’ floating point operations.
$\Gamma$	The set of all edge devices’ memory requirements.
$\hat{h}$	The average number of attention heads among all sub-models.
$\hat{D}$	The average neuron dimension of MLP among all sub-models.
$\Psi \left(𝒞\right)$	The objective function of latency and accuracy degradation minimization.
$ϵ$	Gaussian white noise.
$M$	The number of samples in the training dataset.
$w_i^n$	The weight of $i$-th sample in the training dataset when calibrating the $n$-th sub-model.
$W_n$	The set of samples’ weight when calibrating the $n$-th sub-model.
${ℒ}_{CE}$	The cross entropy loss.
${ℒ}_{Bo}^n$	The calibrating loss of the $n$-th sub-model.
$\Lambda (\cdot )$	The prediction results of ViT.

Table 7. TABLE VII: Comparison between resources on typical computing platforms and deployment demands for representative models.

Capabilities of Typical Computing Platforms						Overheads for Deploying Representative Models
Devices	Memory	Computation Performance	TDP	Cost		Models	Required Memory	FLOPs	#Params.	Storage
Commercial GPUs						CNNs
V100 SXM3	32 GB	16.4 TFLOPS	300 W	$ 4.4/hr		ResNet50	2.2 GB	8.3 G	26 M	97.7 MB
A100 SXM4	80 GB	19.5 TFLOPS	400 W	$ 12.0/hr		ResNet101	2.4 GB	15.7 G	45 M	170.0 MB
H100 SXM5	80 GB	66.9 TFLOPS	700 W	$ 20.7/hr		MobileNet V2	2.1 GB	0.6 G	4 M	13.5 MB
Edge Devices						Transformers
Raspberry Pi-4B [46]	8 GB	13.5 GFLOPS	7.3 W	$ 99		Swin-L [4]	3.3 GB	103.9 G	197 M	751.0 MB
NVIDIA Jetson Nano [47]	4 GB	235.8 GFLOPS	10 W	$ 60		ViT-L/16 [23]	5.3 GB	123.1 G	304 M	1.1 GB
NVIDIA Jetson TX2 [47]	8 GB	665.6 GFLOPS	15 W	$ 249		Flan-T5 Large [5]	4.2 GB	1780 G	751 M	2.9 GB
NVIDIA Jetson Orin Nano [47]	4 GB	640.0 GFLOPS	10 W	$ 199		GPT2-XL [6]	7.8 GB	3340 G	1560 M	5.8 GB

Equations

$${\min }_{\mathcal{C}}\mathcal{L}(\mathcal{C})+\delta [T_c(\mathcal{C})+T_t(\mathcal{C})],$$

$$\text{(C0)}r(C_n)\le R_n,\forall n\in [1,N],C_n\in \mathcal{C},$$

$$X_{agg}=\text{Pool}(\mathbf{W}\cdot \text{Concat}\left({\mathbf{X}}_1,{\mathbf{X}}_2,...,{\mathbf{X}}_N\right)+\mathbf{b}),$$

$$T={\max }_{n\in \mathcal{N}}\left(t_n^1+t_n^2\right)+t_n^3,$$

$$t_n^1=f\left(C_n\right)=f\left(l_n,d_n,h_n^{1:l_n},D_n^{1:l_n}\right),\forall n\in \mathcal{N},$$

$$t_n^2=|{\mathbf{X}}_n|/r_n,\forall n\in \mathcal{N},$$

$$t_n^3=2M{d}_i{d}_{agg}/g,i\in \mathcal{N}.$$

$${\mathcal{L}}_{val}=\frac{1}{N}{\sum }_{n=1}^N{\sum }_{x\in {\mathcal{D}}_{val}}{\mathcal{L}}_n(x,C_n),C_n\in \mathcal{C},$$

$${\min }_{\mathcal{C}}{\mathcal{L}}_{val}(\mathcal{C})+\delta T(\mathcal{C})$$

$$\text{(C1)}{l}_n\le L,\forall n\in \mathcal{N},$$

$$\text{(C2)}{\sum }_{n=1}^N{d}_n\le d,$$

$$\text{(C3)}{\sum }_{n=1}^N{h}_n^k\le h,\forall k\in [1,l_n],$$

$$\text{(C4)}{\sum }_{n=1}^N{D}_n^k\le D,\forall k\in [1,l_n],$$

$$\text{(C5)}\omega (C_n)\le {\Omega }_n,\forall n\in \mathcal{N},C_n\in \mathcal{C},$$

$$\text{(C6)}\varphi (C_n)\le {\Phi }_n,\forall n\in \mathcal{N},C_n\in \mathcal{C},$$

$$\Psi (\mathcal{C})\sim \text{GP}(\mu (\cdot ),K(\cdot ,\cdot )),$$

$$K({\mathcal{C}}_i,{\mathcal{C}}_j)=\frac{{\Theta }^{\nu }{K}_{\nu }\left(\Theta \right)}{\Gamma (\nu )2^{\nu -1}},\Theta =\sqrt{2\nu ||{\mathcal{C}}_i-{\mathcal{C}}_j|{|}_2}/l.$$

$$\Psi ({\mathcal{C}}_{1:r}^{\prime })|{\mathcal{C}}_{1:r}^{\prime },{\hat{y}}_{1:r}\sim \mathbf{N}({\hat{y}}_{1:r},\mathbf{K}+{\sigma }^2\mathbf{I}).$$

$$\Psi (\mathcal{C})|{\mathcal{C}}_{1:r}^{\prime },{\hat{y}}_{1:r}$$

$$\mu (\mathcal{C})$$

$${\sigma }^2(\mathcal{C})$$

$${\mathcal{C}}_{i+1}^{\prime }$$

$$u_i(\mathcal{C})$$

$$w_i^n=w_i^{n-1}\cdot \text{exp}\left[(1/M-1){\mathcal{L}}_{Bo}^{n-1}\right],i=1,...,M,$$

$${\mathcal{L}}_{Bo}^n=\frac{W_n}{2}\left[{\mathcal{L}}_{CE}(\Lambda (Y_{s_n}),y)+{\mathcal{L}}_{CE}(\Lambda (Y_{s_n}),y_t)\right],$$

$$\Psi (\mathcal{C})\sim \text{GP}(\mu (\cdot ),K(\cdot ,\cdot )),$$

$$K({\mathcal{C}}_i,{\mathcal{C}}_j)=\frac{{\Theta }^{\nu }{K}_{\nu }\left(\Theta \right)}{\Gamma (\nu )2^{\nu -1}},\Theta =\sqrt{2\nu ||{\mathcal{C}}_i-{\mathcal{C}}_j|{|}_2}/l.$$

$$\Psi ({\mathcal{C}}_{1:r}^{\prime })|{\mathcal{C}}_{1:r}^{\prime },{\hat{y}}_{1:r}\sim \mathbf{N}({\hat{y}}_{1:r},\mathbf{K}+{\sigma }^2\mathbf{I}).$$

$$\Psi (\mathcal{C})|{\mathcal{C}}_{1:r}^{\prime },{\hat{y}}_{1:r}$$

$$\mu (\mathcal{C})$$

$${\sigma }^2(\mathcal{C})$$

$${\mathcal{C}}_{i+1}^{\prime }$$

$$u_i(\mathcal{C})$$

Full text

I Introduction

With the rapid development of 5G wireless networks and the artificial intelligence of things (AIoT), deploying deep neural networks (DNNs) on edge devices to enable intelligent applications has become prevalent, such as virtual/augmented reality [1], autonomous vehicles [2], object tracking [3]. The transformer models have emerged as the fundamental backbones in most intelligent applications for their impressive performance across diverse benchmarks [4, 5]. However, the substantial computational demands and stringent latency requirements of transformer models pose significant challenges to deploy them on edge devices. For example, the GPT2-XL [6], a representative transformer architecture, requires 7.8GB of memory and 3340 Giga floating point operations (GFLOPs) for inference, surpassing the capabilities of the typical NVIDIA Jetson Nano, which offers only 4GB of memory and a processing capability of 235.8 Giga floating point operations per second (GFLOPS). To facilitate efficient inference, models are segmented and deployed on edge devices or cloud servers to conduct inference together, called collaborative inference [7]. This method helps mitigate the resource constraints and computational burdens on a single device, enabling more feasible on-device inference.

Existing collaborative inference methods [8, 9, 10, 11] utilize a pipe-edge solution, depicted in Figure 2 (a), which offloads tasks to cloud servers equipped with high-performance GPUs or distributes segments to various devices. However, this sequential execution results in high end-to-end latency. The end-to-end latency, comprising both the processing time and the communication delay among edge devices, can extend into seconds. A device can only commence inference after receiving outputs from preceding devices, causing significant delays. This delay is unacceptable for real-time interactive applications in edge systems, which typically tolerate only tens to hundreds of milliseconds of latency [12]. To minimize these delays and fully utilize computational resources across all edge devices, several researchers [13, 14, 15] have proposed a distri-edge solution, illustrated in Figure 2 (b). This strategy, inspired by model parallel, involves decomposing models across layers and deployed segments on multiple devices for distributed inference. The decomposition scheme destroys the original structure and hinders the capability to extract information, necessitating frequent inter-device communication during inference. Another line of works [16, 17], the single-edge solution, shown in Figure 2 (c), utilize model compression [18, 19, 20] or neural architecture search (NAS) [21, 22] techniques to obtain a compact and small model deployed on a single device. While this method maximizes resource utilization on a single device, striking a balance between performance and efficiency remains challenging.

Motivated by the divisibility and integrability of transformer models, we propose a collaborative inference system for general transformer models, termed CoFormer. It belongs to the aggregate-edge solution depicted in Figure 2 (d), where an off-the-shelf large transformer model is automatically decomposed into several smaller models. These decomposed models can be readily deployed on edge devices for concurrent inference, and the intermediate results are transmitted only once and aggregated to produce the final output. The scheme significantly lessens the resource demands on any single device and reduces transmission costs in collaborative inference. We formulate an optimization problem to concurrently minimize inference latency and performance degradation, while accommodating the heterogeneous hardware resources. To address this problem, we propose the DeBo algorithm, which leverages Bayesian optimization to derive an efficient decomposition strategy. Furthermore, DeBo incorporates a progressive knowledge distillation process to transfer knowledge from the original large transformer to the sub-models, thereby calibrating potential accuracy sacrifice caused by model decomposition.

To verify the effectiveness of our system, we conducted comprehensive experiments across three widely-adopted tasks, including image classification, object detection, and language understanding, using both vision transformer backbones (ViT [23], DeiT [18], and Swin [4]) and language transformer backbones (BERT [24], GPT [6], and Flan-T5 [5]) on heterogeneous edge devices. Our system supports the deployment of large transformer models on resource-constrained edge devices by significantly reducing computational burdens, thereby achieving a satisfactory balance between efficiency and performance. Notably, CoFormer enables the deployment of GPT2-XL, a huge transformer, on edge devices with only 4 GB memory capacity, reducing memory requirements by 76.3%. Furthermore, as shown in Figure 1, CoFormer accelerate inference speed by 1.7$\times$ to 3.1$\times$ and reduces energy consumption by 36.3% to 63.8% with less than 2% performance sacrifice compared to the large transformer.

Our main contributions are summarized as follows:

• •

We propose a collaborative inference system for general transformers, termed CoFormer, automatically decomposing an off-the-shelf large transformer into multiple small models deployed on heterogeneous edge devices.

• •

We formulate an optimization problem to minimize both inference latency and accuracy degradation under the heterogeneous hardware constraints to achieve the trade-off between efficiency and performance.

• •

We develop the DeBo algorithm to first solve the optimization problem to derive the decomposition policy, and then progressively calibrate decomposed models for performance enhancement.

• •

We perform extensive experiments on a real-world collaborative edge computing system consisting of heterogeneous edge devices using three representative transformer backbones across three widely-used tasks, which accelerates inference speed by 2.5$\times$ and saves energy consumption by 53.2% on average.

II Rethinking Resource Utilization in Collaborative Inference

In this section, we first define the problem encountered in collaborative inference systems and rethink resource utilization in existing collaborative inference schemes. Then, we are inspired by two observations to remedy the problems introduced by resource utilization.

II-A Problem Definition

Collaborative inference [7, 9] involves the cooperation of multiple edge devices or cloud servers to execute model inference tasks efficiently. In this approach, a pre-trained model is strategically decomposed, and components are distributed across various edge devices or servers. Upon receiving an inference request, the initial segments of the model are activated to process the data. The resultant intermediate data is then forwarded to subsequent devices or cloud servers for further processing. This process continues until the final result is produced. The primary objective of collaborative inference is to minimize inference latency and accuracy degradation within the resource constraints of edge devices. Given a transformer model, the number of devices is $N$. The objective function is

[TABLE]

where $\mathcal{L}$, $T_{c}$ and $T_{t}$ represents accuracy degradation, computation latency and communication latency, respectively. (C0) is resource constraints of edge devices, *e.g., *memory capacity, computational capability. $r(C_{n})$ and $R_{n}$ denote the resource requirements of the sub-model and resource constraints of corresponding devices, respectively. The decomposition policy is denoted as $\mathcal{C}=\{C_{1},C_{2},...,C_{N}\}$. Note that $\mathcal{C}=\varnothing$ represents utilizing a single edge device to execute inference.

II-B Resource Utilization in Collaborative Inference

Given the constrained hardware resources on edge devices, collaborative inference systems should aim for the trade-off between accuracy and efficiency, and maintain the quality of service (QoS) in edge computing applications. However, existing solutions fail to meet all goals stated in §II-A, resulting in the following problems:

Significant idle latency. Several methods [8, 9, 10, 11] partition models by layer into multiple sequential execution blocks, termed pipe-edge solution. These sub-models are tailored to the resource limitations of edge devices, easing the computational load on individual devices. However, the manner of sequential execution neglects the inherently parallel nature of neural networks, resulting in computational redundancies due to device idleness. We partition DeiT-B at the third and sixth layers into three sub-models and profile the inference latency on three edge devices, as shown in Figure 3. It reveals that idle time constitutes over 70% of the total inference process of DeiT-B on ImageNet-1K. The significant idle latency contributes to the overall increase in end-to-end inference latency.

Substantial transmission overheads. To enhance resource utilization and minimize device idleness, an alternative distri-edge solution partitions models across layers for distributed inference [13, 14, 15, 25]. This setup requires sub-models to frequently communicate between devices for feature interaction, necessitating a stable network with adequate bandwidth and resulting in significant transmission costs. We profile representative collaborative inference methods [13, 14, 15] on widely-used edge devices, Jetson Nano and Jetson TX2, with a network bandwidth of 2 Mb/s and measure their transmission costs as shown in Figure 4. This breakdown illustrates the transmission delay constitutes more than 40% of the total latency, and this percentage significantly increases on the Jetson TX2. Thus, substantial network transmission delay emerges as the principal bottleneck of inference latency.

Poor performance-efficiency trade-off. Another line of methods such as model compression [17, 19, 20] or NAS techniques [21, 22] strive to optimize models to fully harness a single device’s hardware resources, termed single-edge solution. We compare inference performance and latency of two typical methods, EfficientFormer L7 [26] and MobileViTv2_200 [17] on the Jetson Nano and Jetson TX2 platforms using ImageNet-1K, as shown in Table I. While EfficientFormer L7 boosts classification accuracy by 2.13%, it also increases inference latency by 107.2% on the Jetson Nano, a trend consistent on the Jetson TX2. These results highlight the challenging balance between performance and efficiency for single-edge solution.

II-C Insight: Divisibility and Integrability of Transformer

We conduct preliminary experiments to investigate the transformer models and find the divisibility and integrability.

Divisibility of transformer. Due to the significant computational demands of transformer models, which exceeds the capacities of individual edge devices, we propose to collaborate with multiple edge devices by leveraging the divisibility of transformer. Recent studies demonstrate that distinct attention heads within transformer models are capable of encoding diverse semantic information [27, 28], and certain neurons are critical for capturing features of particular samples [29]. Building on these insights, we delve into the role of different attention heads as depicted in Figure 5 (a). Figure 5 (a) reveals only a subset of attention heads are crucial for specific tokens, allowing for the feasible decomposition of heads to extract varied semantic information. Figure 5 (b) illustrates the performance degradation resulting from this decomposition. A minority of heads are not essential for certain tasks, enabling their segregation with minimal performance impact. However, increasing the decomposition ratio from 0.4 to 0.5 results in a 14.1% reduction in performance, highlighting that decomposing certain target heads can lead to significant performance degradation. Consequently, we intend to decompose a large transformer model into multiple smaller models to enable collaborative inference.

Integrability of transformer. To mitigate the performance degradation resulting from decomposing large transformer models, an intuitive method is to integrate multiple decomposed models. Model ensemble is a strategy that combines several individual models to produce a final output [30], including model averaging [31], majority voting [31], etc. This technique not only boosts the overall performance but also reduces the variance across models [30]. We assessed the ensemble efficacy of various transformer models through weighted averaging on the ImageNet-1K dataset, as illustrated in Figure 6 (a) and (b). Aggregating multiple homogeneous transformer models of varying sizes yielded a 2.94% increase in accuracy compared to DeiT-S. Moreover, integrating two heterogeneous transformer models led to a 9.89% accuracy improvement over DeiT-Ti. However, the ensemble of multiple transformers leads to a more than 200% increase in inference latency. While model integration significantly boosts performance, the overall inference speed is limited by the slowest model in the ensemble.

Inspired by these two observations, we further propose CoFormer, a collaborative inference system by decomposing large transformer models and deploying decomposed small models on multiple heterogeneous edge devices for parallel inference. The intermediate results from these devices are aggregated to gain the final output. However, realizing CoFormer suffers from a set of challenges:

• •

**Challenge 1: **How to decompose the large transformer model to make the most use of hardware resources among heterogeneous edge devices?

• •

**Challenge 2: **How to reduce transmission overheads during aggregating intermediate results?

• •

**Challenge 3: **How to alleviate the synchronization delay of parallel execution?

III CoFormer Design

In this section, we first present an overview of CoFormer, and construct the system model to formulate inference latency and accuracy degradation minimization problem. Then, we illustrate the DeBo algorithm to solve the problem.

III-A Design Overview

In terms of resource-harvesting characteristics of transformer inference and existing problems of collaborative inference methods, we propose CoFormer, a collaborative inference system for general transformers, to achieve efficient and reliable inference on heterogeneous edge devices. CoFormer requires multiple heterogeneous edge devices, a central node, a target large transformer, a model pool consisting of different sizes of transformer with the same architecture as the target transformer, as well as utilizes a decomposer to partition the transformer, an evaluator to evaluate decomposed models and a booster to calibrate these models. Figure 7 depicts an overview of CoFormer which features three primary stages:

**(i) Preprocessing Stage: ** (Step ➀) Preprocessing stage is an offline procedure that runs once before deployment. All edge devices execute transformers in the model pool using calibration data as input to record the runtime profile. The model configurations (*e.g., *the number of model parameters, FLOPs) and hardware resources (*e.g., *computational capabilities, memory budgets) are utilized to construct evaluator for the decomposition stage.

**(ii) Decomposition Stage: ** (Step ➁) According to the hardware constraints of various edge devices, the large transformer is first decomposed into multiple distinct sub-models of varying sizes (*e.g., *tiny, small, and medium transformers) by decomposer. (Step ➂) These sub-models are evaluated on the evaluator to profile inference latency and accuracy. Evaluation results are used to update the decomposer. The decomposer generates new decomposition policies until finding efficient decomposed models. (Step ➃) These models are progressively calibrated in the booster to restore performance.

**(iii) Inference Stage: ** (Step ➄) These calibrated models are deployed the corresponding edge devices. When receiving an inference task, the edge devices conduct inference with sub-models in parallel. Intermediate results are transmitted to central device and aggregated to produce the inference output.

III-A1 Inference Workflow

The bottom of Figure 7 illustrates how edge devices conduct collaborative inference tasks. After sub-models are calibrated and deployed on corresponding edge devices, inter-device communication is enabled through either wireless (*e.g., *Wi-Fi, Zigbee, Bluetooth) or wired networks. One of edge devices is chosen as the central node. Upon receiving an inference task, the collected data (*e.g., *images, videos, text) are initially transmitted across edge devices as input. The whole inference workflow includes three phases:

• •

Phase 1 (Backbone Forward): All edge devices concurrently execute the sub-models backbone for inference.

• •

Phase 2 (Data Transmission): Intermediate results from other devices are transmitted to the central node.

• •

Phase 3 (Results Aggregation): The central node aggregates these intermediate results to produce the inference output by utilizing aggregation module.

Results aggregation. We present an efficient ensemble method to aggregate intermediate results from edge devices, significantly reducing transmission overhead. The down-sampled features from the final layer of the transformer are leveraged within the aggregation module, which substantially decreases data volume and communication latency. Let the $i$-th device serve as the central node. In the aggregation module, intermediate features $\{\mathbf{X}_{1},\mathbf{X}_{2},...,\mathbf{X}_{N}\}$ from all devices are first concatenated at the central node. Then, an MLP block facilitates the interaction and fusion of these features through a linear transformation $\mathbf{W}\in\mathbb{R}^{d_{agg}\times d_{i}}$, where $d_{agg}$ represents the cumulative dimension of intermediate features and $d_{i}$ denotes the dimension of features at the central node. Consequently, the aggregation module is formulated as:

[TABLE]

where $\text{Pool}(\cdot)$ signifies an average pooling operation and $\mathbf{b}$ is the bias vector. The aggregated features are finally passed to the task-specific heads111The task-specific heads refer to specialized output layers in the neural network that are only responsible for making predictions for different downstream tasks. The heads are directly inherited from the large transformer model, exclude from the decomposition process. to generate predictions. Our aggregation module guarantees a satisfactory inference performance with slight transmission overheads.

III-B Evaluator: System Modeling

To minimize both inference latency and accuracy degradation, we pose an optimization problem focusing on these metrics to achieve efficient collaborative inference.

III-B1 Inference Latency Model

A large transformer is decomposed into multiple sub-models for deployment across various edge devices. These devices collaboratively execute the sub-models to perform inference tasks. Given a system composed of $N$ edge devices, the large transformer is decomposed into $N$ sub-models under hardware constraints of edge devices. Each device executes its corresponding sub-model backbone for inference. We represent the set of sub-models as $\mathcal{N}=\{1,2,...,N\}$. The architecture configurations of sub-model $n$ are denoted as $C_{n}=\{l_{n},d_{n},h_{n}^{1:l_{n}},D_{n}^{1:l_{n}}\}$, where $l_{n}$ and $d_{n}$ are the number of layers and embedding dimension, respectively. $h_{n}^{1:l_{n}}$ and $D_{n}^{1:l_{n}}$ are two vectors consisting of the number of attention heads and the MLP dimension from the first to the $l_{n}$-th layer, respectively. The set of decomposition decisions is denoted as $\mathcal{C}=\{C_{1},C_{2},...,C_{N}\}$, which indicates the large transformer is partitioned into $N$ sub-models with these specific specifications. When multiple edge devices execute inference tasks concurrently, their completion times can vary, requiring a delay in aggregation until the slowest device has completed its data transmission. We can formulate the overall latency of collaborative inference as

[TABLE]

where $t_{n}^{1}$, $t_{n}^{2}$ and $t_{n}^{3}$ represents the latency of sub-model $n$ at different phases.

Phase 1. The inference latency $t_{n}^{1}$ for the backbone of sub-model $n$ is formulated as:

[TABLE]

where function $f(\cdot)$ is a latency predictor developed from actual measurement data profiled on edge devices. This predictor can estimate the latency of sub-models based on their configurations, as detailed in §A-A of supplementary material.

Phase 2. The latency $t_{n}^{2}$ incurred by edge device $n$ when transmitting the intermediate feature $\mathbf{X}_{n}$ to the central node is calculated as follows:

[TABLE]

where $|\mathbf{X}_{n}|$ signifies the size of $\mathbf{X}_{n}$. The data transmission rate between device $n$ and the central node is $r_{n}$.

Phase 3. The aggregated features are initially concatenated into $\mathbf{X}_{con}\in\mathbb{R}^{S\times d_{agg}}$. Then the features undergo a linear transformation $\mathbf{W}\in\mathbb{R}^{d_{agg}\times d_{i}}$ and an average pool layer to extract the target features. The available computational power for the central node $i$, denoted by $g$ (in FLOPs/ms), determines the latency $t_{n}^{3}$ for results aggregation:

[TABLE]

III-B2 Accuracy Degradation Model

The inference accuracy of collaborative inference primarily depends on the average accuracy of these sub-models[30]. The decomposition strategy $\mathcal{C}$, determines the configurations of sub-models, which in turn affects their performance. Typically, generalization error is employed to evaluate the performance degradation of neural network models. Ideally, each model should be trained from scratch and then assessed for accuracy on a test dataset. However, training all generated sub-models exhaustively is impractical due to significant computational overheads and time constraints. Therefore, we adopt the average validation loss of all sub-models on the validation dataset $\mathcal{D}_{val}$, as a surrogate for generalization error to estimate accuracy degradation. The total accuracy degradation is defined as follows:

[TABLE]

where $\sum_{x\in\mathcal{D}_{val}}\mathcal{L}_{n}$ is the validation loss of sub-model $n$.

III-B3 Inference Latency and Accuracy Degradation Minimization Problem

Based on the above modeling, we formulate the overall optimization problem aim to minimize both inference latency and accuracy degradation in collaborative inference system. In our scenario, each edge device receives input data at the beginning. Our goal is to find a decomposition policy $\mathcal{C}$ that optimizes both the end-to-end inference latency and accuracy degradation in collaborative inference. Specifically, given $N$ heterogeneous devices, a large transformer with configuration $\{L,d,h,D\}$ is decomposed into $N$ sub-models, each deployed on one of these devices. These sub-models must meet the hardware constraints of the respective devices. Thus, the minimization problem can be formulated as follows:

[TABLE]

where $\delta>0$ is a balancing hyperparameter. Constraint (C1), (C2), (C3) and (C4) ensure that the number of layers, the embedding dimension, the number of attention heads and the MLP dimension of each sub-model are smaller than those of the original transformer. Constraint (C5) and (C6) limit the maximum computing resources and memory capacity for device $n$ to $\Omega_{n}$ and $\Phi_{n}$, respectively.

Remark. The inference latency model and accuracy degradation model are employed as the evaluator to assess decomposition policies. In order to solve the formulated problem, applying standard optimization methods is impractical due to two reasons: (i) it is a high-order optimization of $\mathcal{C}$ since we need to find all configurations of $N$ sub-models, which requires significant computational and memory costs; (ii) inference latency and accuracy degradation are intractable to express in a closed-form and are not differentiable. To tackle these issues, we treat the objective function in equation (8) as a black-box function and optimize it using Bayesian optimization based on the observed decomposition dynamics.

III-C DeBo Algorithm

Algorithm 1 presents our DeBo algorithm to optimize inference latency and accuracy degradation, which utilizes decomposer to generates appropriate decomposition policies and decompose the large transformer into multiple sub-models. Then, booster progressively calibrates these sub-models.

III-C1 Decomposer: Bayesian Decomposition

We utilize the decomposer to generate decomposition policies for the large transformer as shown in Lines 1 to 11 of Algorithm 1. These policies are evaluated on the evaluator. Evaluation results are utilized to update decomposer until finding the optimal policy. Decomposer decomposes the target transformer into multiple sub-models according to the final policy.

The process of decomposing transformer models is illustrated in the §A-B of supplementary material due to space limitation. Then we introduce how to gain the optimized decomposition policy.

As analyzed in §III-B3, we treat the objective function in the equation (8) as a black-box-function $\varPsi(\mathcal{C})=\mathcal{L}(\mathcal{C})+\delta T(\mathcal{C})$ and employ BO technique to optimize. It applies exploration and exploitation to the objective function by sequentially and actively querying the function values of some input instances. In this work, we utilize Gaussian process (GP) as a prior model to fit the black-box objective function $\varPsi(\mathcal{C})$ in a closed form. It updates the distribution of $\varPsi(\mathcal{C})$ to gain posterior prediction by using its likelihood on newly evaluated $(\mathcal{C},y=\varPsi(\mathcal{C})+\epsilon)$ pairs, where $y$ is a noisy observation of $\varPsi(\mathcal{C})$ and is the sum of inference latency and accuracy degradation in our case. Then it determines the next decomposition policy $\mathcal{C}$ by minimizing an acquisition function, which is computed from the updated posterior. The acquisition function performs a trade-off between exploration and exploitation in evaluating the candidates of decomposition policy. The above process is repeated until achieving a precise posterior predictive distribution of $\varPsi(\mathcal{C})$.

Prior model. GP is utilized as the prior to model the objective function. A GP prior is specified by its mean function $\mu(\cdot)=\mathbb{E}[\cdot]$ and covariance function $K(\cdot,\cdot)$. The prior of objective function is defined as follows:

[TABLE]

We set $\mu(\cdot)=0$ and $K(\cdot,\cdot)$ to be the Matern kernel with smoothness factor $\nu=1.5$ and length scale $l=1$. $K_{\nu}(\cdot)$ and $\Gamma(\cdot)$ are a modified Bessel and gamma function, respectively.

Posterior prediction. We sample $r$ different decomposition policy $\mathcal{C^{\prime}}$ satisfying the constraints $\varOmega=\{\Omega_{1},...,\Omega_{N}\}$ and $\varPhi=\{\Phi_{1},...,\Phi_{N}\}$ of all devices, and obtain $r$ sets of decomposed sub-models. For notation simplicity, vectors composed of $\{\mathcal{C}^{\prime}\}_{i=1}^{r}$ and $\{\varPsi(\mathcal{C}^{\prime})\}_{i=1}^{r}$ are denoted as $\mathcal{C}^{\prime}_{1:r}$ and $\varPsi(\mathcal{C}^{\prime}_{1:r})$, respectively. We evaluate the objective function denoted by $\hat{y}_{i}$ as a noisy observation of $\varPsi(\mathcal{C}^{\prime})$, i.e., $\hat{y}_{i}=\varPsi(\mathcal{C}^{\prime})+\epsilon$ where Gaussian white noise $\epsilon\thicksim\mathbf{N}(0,\sigma^{2})$. Given the noisy observations $\hat{y}_{1:r}$, the GP posterior can be updated as follows:

[TABLE]

Given a new decomposition policy $\mathcal{C}$ and $\mathbf{k}_{i}=K(\mathcal{C},\mathcal{C}^{\prime}_{i})$, the GP posterior can be used to predict the distribution of $\varPsi(\mathcal{C})$:

[TABLE]

Acquisition function. Given the posterior predictive distribution of $\varPsi(\mathcal{C})$, BO finds the next decomposition policy $\mathcal{C}^{\prime}_{i+1}$ to evaluate based on an acquisition function $u_{i}(\mathcal{C})$ defined by the posterior mean $\mu_{i}(\mathcal{C})$ and standard deviation $\sigma_{i}(\mathcal{C})$. In our algorithm, expected improvement (EI) is utilized as acquisition function. For problem (P1), we can determine the next decomposition policy $\mathcal{C}^{\prime}_{i+1}$ by minimizing EI:

[TABLE]

where $Z$ and $H$ represent the probability density function and cumulative distribution function of the standard normal distribution, respectively. $\varPsi^{*}$ is the best policy in $\mathcal{C}^{\prime}_{1:r}$.

Analysis of search space. Consider a large transformer model with $L$ layers, $h$ attention heads, embedding dimension $d$, and MLP hidden dimension $D$. The configuration of a sub-model deployed on device $n$ is defined as $C_{n}={l_{n},d_{n},h_{n}^{1:l_{n}},D_{n}^{1:l_{n}}}$, where $l_{n}$ and $d_{n}$ denote the number of layers and embedding dimension, respectively. $h_{n}^{1:l_{n}}$ and $D_{n}^{1:l_{n}}$ are two vectors consisting of the number of attention heads and the MLP dimension from the first to the $l_{n}$-th layer, respectively. This configuration must satisfy the hardware constraints of device $n$, including memory $\phi(C_{n})\leq\Phi_{n}$ and computation $\omega(C_{n})\leq\Omega_{n}$. Accordingly, the search space can be characterized as $l_{n}^{3}d_{n}h_{n}^{1:l_{n}}D_{n}^{1:l_{n}}$, constrained by both memory and computational budgets. The decomposition is an offline procedure that runs once before deployment.

III-C2 Booster: Progressively Boosting Distillation

To mitigate the performance degradation resulting from decomposition, we introduce the booster, employing a progressively boosting distillation method to improve performance as shown in Lines 12 to 15 of Algorithm 1.

The booster consists of a large transformer, $N$ decomposed sub-models and a training dataset. We denote the training dataset as $\mathcal{D}_{train}=\{(x_{i},y_{i}),i\in[1,M]\}$, where $M$ and $R$ are the number of samples and the total number of sample classes, respectively. For the $n$-th iteration, the weights of all training samples are $W_{n}=\{w_{i}^{n}|i=1,...,M\}$. The progressively boosting distillation calibrates each sub-model sequentially, comprising two primary stages per iteration: weight update and model distillation.

Weight update. In each iteration, the weights of training samples are adjusted based on the prediction errors of the last calibrated sub-model. Note that in the first iteration, sample weights are uniformly initialized as $W_{1}=\{w_{i}^{1}|w_{i}=1/M,i=1,...,M\}$. For the $n$-th iteration, the sample weights $W_{n}$ are updated as follows:

[TABLE]

where $\mathcal{L}_{Bo}^{n-1}$ is the distillation loss to calibrate the $(n-1)$-th sub-model computed by aligning it with the large transformer.

Model distillation. The updated weights are then used in the distillation process. Following the approach in DeiT [18], knowledge from the large transformer is transferred to the current sub-model. For the $n$-th iteration, the objective function of distillation is defined as

[TABLE]

where $\varLambda(\cdot)$ is the softmax function and $\mathcal{L}_{CE}$ is the cross entropy loss. The prediction logits of the large transformer and $n$-th sub-model are denoted as $Y_{t}$ and $Y_{s_{n}}$, respectively. $y_{t}$ is the hard decision of the large transformer. This process is repeated until all sub-models are sequentially calibrated.

Discussion. Traditional ensemble learning methods typically need to train all models independently from scratch and then aggregate their predictions to generate the final output. In contrast, CoFormer decomposes a pre-trained large transformer into multiple sub-models. These sub-models are calibrated by employing a progressive distillation process to transfer knowledge from the original large transformer model. These calibrated sub-models perform inference in parallel on different edge devices. Their intermediate results from these devices are aggregated by the central device using an efficient aggregation module to produce the final output. Our method significantly reduces the inference latency and communication overhead compared to traditional ensemble methods, while maintaining high accuracy.

IV Evaluation

IV-A Experimental Setup

Prototype. We implement our system on a real-world collaborative edge computing testbed that consits of an NVIDIA Jetson Nano with 4 GB memory, an NVIDIA Jetson TX2 with 8 GB memory, an NVIDIA Jetson Orin Nano with 4 GB memory and a switch. These devices represent the edge devices with different hardware resources and computing capabilities. The detailed configurations of these devices are shown in the lower left part of Table VII of supplementary material. Figure 8 shows the implemented hardware platform for CoFormer. These devices are connected via a gigabyte switch TP-LINK TL-SG1008D. We use the traffic control tool tc222Tc-show, https://www.linux.com/tutorials/tc-show-manipulate-traffic-control-settings/ to control bandwidth, which can help simulate the diverse network conditions within realistic edge environments under different network conditions to evaluate the performance.

Transformer backbones, tasks and datasets. To evaluate the generality of our system on vision and language tasks, we selected three typical intelligent applications.

• •

Generic-category image recognition. The vision task aims to recognize the generic category of an image. We select three representative ViT backbones (ViT[23], DeiT[18], Swin[4]) and two widely-used classification datasets (CIFAR-100333CIFAR dataset, https://www.cs.toronto.edu/~kriz/cifar.html and ImageNet-1K444ImageNet dataset, https://image-net.org/).

• •

Real-time object detection. This type of vision tasks aims to detect objects in real-time with fast processing time while maintaining a required level of accuracy. We choose Mask R-CNN [32] with ViT, DeiT and Swin backbone as the detection model in the most widely-used dataset, MS COCO 2017555MS COCO dataset, https://cocodataset.org/, to evaluate.

• •

General language understanding. The language task aims to understand natural language. We select three language transformer backbones (BERT [24], GPT2 [6] and Flan-T5 [5]) and a representative General Language Understanding Evaluation (GLUE) benchmark666GLUE benchmark, https://gluebenchmark.com/ to evaluate our method. The GLUE benchmark consists of nine language understanding tasks and cover a diverse range of dataset sizes and degrees of difficulty.

Baselines. We compare CoFormer with state-of-the-art edge deployment methods for transformer models by implementing three types of baselines. (i) Comparison with large transformer models: We compare three vision transformer models (ViT-L/16, DeiT-B and Swin-L) and three language transformer models (BERT-Large, GPT2-XL and Flan-T5-Large) as the decomposed models to demonstrate the effectiveness of our methods. (ii) Comparison with efficient transformer models: We compare our methods with several representative lightweight transformer models (*i.e., *single-edge solution illustrated in Figure 2 (b)) to show the superiority, including T2T-ViT[33], EfficientFormer [26], PoolFormer [34] and MobileViTv2 [17]. (iii) Comparison with collaborative inference methods: DeViT [35] is the most relevant baseline to our work, designed for vision transformer inference on multiple homogeneous edge devices. Galaxy [15] employs tensor parallel across attention and MLP modules for collaborative inference on edge devices. DeTransformer [36] combine block parallel and tensor parallel techniques to balance inference accuracy and communication cost. EdgeShard [37] splits transformer into a pipeline architecture to collaborate with edge devices.

Evaluation metrics. We employ Top-1 classification accuracy, mean average precision (mAP) and evaluation scores on GLUE benchmark as the overall performance metric of image classification, object detection and language understanding tasks, respectively. The end-to-end inference latency and energy consumption are utilized as the efficiency metrics. (i) Measurements of inference latency: We measure the end-to-end execution time in the inference process with 50 runs without break using PyTorch Profiler tool777PyTorch profiler, https://pytorch.org/docs/stable/profiler.html. We record the average execution time as the end-to-end inference latency. (ii) Measurements of energy consumption: We measure the total energy consumption of all edge devices. We use a Monsoon High Voltage Power Monitor888Monsoon Solutions Inc., High voltage power monitor, https://www.msoon.com/high-voltage-power-monitor connected to edge devices to obtain accurate power readings over running a forward pass of each model for 50 times. For each device, we derive the average energy consumption by first subtracting background power consumption (*i.e., *power readings when not running any model) and calculate the integral over the individual power samples as energy consumption. The average energy consumption is recorded as the final statistics of a model. The measurement follows [38] for fair comparison.

Implementation. CoFormer is fully implemented with $~2,000$ LOC in Python in total atop PyTorch. Although we use PyTorch for auto-differentiation and computation graph execution, CoFormer is extensible and can work well with other lightweight ML frameworks such as TF-Lite and MNN. The employed transformer models are trained following DeBo algorithm and deployed on all devices in advance. The communication module is implemented based on gPRC999gRPC, https://grpc.io/. All the compared approaches are run with timm101010timm, https://github.com/rwightman/pytorch-image-models (vision transformers models) or transformers111111transformers, https://github.com/huggingface/transformers (language transformer models), for fair comparison.

IV-B End-to-End Performance

Comparison with large transformer models. We compare the performance of our method with large transformers on three representative tasks including image classification, object detection and language understanding. The results of inference performance, inference latency, average energy consumption and memory usage are shown in Figure 9.

(1) Image recognition: The original large transformer models used for the image recognition task include ViT-L/16, DeiT-B, and Swin-L. Since ViT-L/16 and Swin-L are impractical for deployment on Jetson Nano and Jetson Orin Nano due to their high memory demands and computational costs, we report the evaluation results of all large transformer models on Jetson TX2 for a fair comparison. The evaluation results on ImageNet-1K and CIFAR-100 are presented in Figure 9 (a-d) and Figure 15 of supplementary material, respectively. Our method demonstrates the effectiveness of decomposing large transformer models for various backbones into smaller, more efficient models, enabling collaborative inference with constrained hardware resource requirements while maintaining satisfactory accuracy. Specifically, our decomposed model, based on DeiT-B, accelerates inference speed by $2.35\times$, reduce energy consumption by 45.3% and save memory usage by 35.6%, with only 1.14% accuracy sacrifice on ImageNet-1K.

(2) Object detection: We also evaluate our system on the object detection task to validate its generality and performance. We utilize CoFormer to gain decomposed ViT, DeiT and Swin as backbones in Mask-RCNN pipeline and evaluate over MS COCO 2017 dataset. These backbones are first initialized with ImageNet-1K pretrained weights and then fine-tuned on MS COCO 2017. Following the training protocol in [4, 39], we use AdamW optimizer with an initial learning rate of 1e-4 and a weight decay as 0.05. Note that we follow [40] to construct hierarchical feature maps for DeiT and ViT because they only produce a single resolution of feature maps, which is unavailable for the detection task. The results are shown in Figure 9 (e-h). Similar to the classification task, our method can achieve fast inference speed and low energy consumption, with minimal performance degradation.

(3) Language understanding: We employ BERT-Large, GPT2-XL and Flan-T5-Large as transformer backbones to evaluate on the language understanding task. Evaluation results on the GLUE benchmark, depicted in Figure 9 (i-l), demonstrates that our approach enables the effective deployment of huge transformer models on resource-constrained edge deviecs. Notably, our method support the GPT2-XL model, which comprises 1.6 billion model parameters, to perform efficient inference on resource-constrained edge devices, achieving a reduction in memory requirements by up to 76.3%. Furthermore, for BERT-Large, our approach can improve execution speed by $1.87\times$, reduce energy consumption by 40.2%, and decrease memory usage by 20.1%.

Comparison with efficient transformer models. Table II presents the evaluation results between our method and efficient transformer models with similar FLOPs to ensure a fair comparison on ImageNet-1K These models are all evaluated on the Jetson TX2 to maintain consistency. We can observe that our method can achieve the lowest end-to-end inference latency and energy consumption with comparable accuracy to other efficient models under varying hardware constraints. For instance, using the DeiT backbone, our method accelerates inference speed by 26.1%, saves energy consumption by 15.2% and improves accuracy by 1.09% with the lowest memory requirements and computational costs compared to state-of-the-art efficient transformer model, MobileViTv2.

Comparison with collaborative inference methods. We implement several representative collaborative inference methods as baselines to validate the superiority of our proposed approach. Figure 10 presents the classification accuracy, average inference latency, average energy consumption and peak memory requirements on ImageNet-1K, using the DeiT backbone. The results indicate CoFormer can achieve the fastest inference speed and lowest energy consumption while maintaining competitive accuracy. For instance, DeViT [35], although conceptually similar, fails to consider the heterogeneity of edge devices, leading to significant synchronization delays. While Galaxy [15] improves classification accuracy by 0.97% over CoFormer but incurs an 82% increase in inference latency due to intensive communication in its tensor-parallel design. DeTransformer [36] reduces communication costs and achieves a 1.73$\times$ speedup over Galaxy, but its inference latency remains 64.6% higher than CoFormer and comes with increased memory consumption. EdgeShard[37] exhibits lower inference speed than CoFormer due to its sequential execution strategy. Overall, these results highlight CoFormer offers an effective trade-off between accuracy and efficiency.

IV-C Ablation Study

Ablation of main components. We design experiments to verify the effectiveness of the proposed transformer decomposition and results aggregation. Table III presents the inference accuracy and latency utilizing different components for DeiT-B. These results show that these components of our system are all significant. Specifically, we can observe that transformer decomposition accelerates inference speed, leading to considerable performance degradation. By utilizing results aggregation, the inference efficiency and accuracy can achieve good trade-off.

Transformer decomposition. In order to evaluate the effectiveness of the proposed transformer decomposition, we compare random decomposition and uniform decomposition with the same computational constraints. Uniform decomposition represents we decompose the large transformer into the multiple sub-models with the same structure. The results of accuracy and latency in the decomposition process are shown in Figure 11. We can observe that our method performs the best with the highest accuracy and lowest inference latency. Although the uniform decomposition convergences faster than others, its inference latency is higher due to the considerable synchronization delay. The random decomposition is unstable and uncontrollable, which cannot be utilized to decompose.

Results aggregation. We compare our methods with other representative aggregation methods, and they all utilize the same decomposed small transformers. The accuracy and inference latency on ImageNet-1K are shown in Table IV. Traditional ensemble methods, such as average and majority voting [30], provide fast aggregation with limited performance improvement. In contrast, attention [41] and SENet [42] methods can focus on target features during aggregation. Our method achieves the lowest inference latency with only 1.14% accuracy sacrifice, demonstrating its ability to effectively balance latency and accuracy.

**Impact of network bandwidths. ** We evaluate the end-to-end inference latency using the DeiT backbone on CIFAR-100 across different network bandwidths, *i.e., *100 Mbps, 500 Mbps, and 1 Gbps, as illustrated in Figure 12. The inference latency of DeiT-B is measured on the Jetson TX2 platform. The results show that our approach consistently reduces inference latency across all tested bandwidths. Notably, higher bandwidths alleviate communication overhead, further amplifying the performance gains. For instance, CoFormer achieves $2.98\times$ speedup at 100 Mbps, and the improvement increases to $3.62\times$ at 1 Gbps compared to DeiT-B. Our method outperforms other collaborative inference baselines and achieves the lowest inference latency across different bandwidths. Specifically, CoFormer accelerates inference speed by $5.65\times$ speedup at 100 Mbps and $1.76\times$ at 1 Gbps compared to Galaxy. These findings highlight the robustness and effectiveness of our method across diverse network environments.

Impact of resource constraints. Figure 13 illustrates the performance of our method for DeiT-B under several resource constraints. We impose computational limits of 30%, 40%, and 50% for edge deployments. At a computational cost of 30%, our method improves inference speed by 3.05$\times$ and reduces energy consumption by 52.3%, with only a slight accuracy sacrifice of 1.56%. Given that accuracy losses up to 2% are acceptable for substantial latency improvements in most real-world applications, it demonstrates the flexibility and efficiency of our method under varied resource constraints.

Impact of device quantity. We examine the impact of altering the quantity of devices on inference performance, as depicted in Table V. The reported energy consumption is the total energy consumed across all devices. To ensure a fair comparison, the total FLOPs across devices are kept consistent. As the number of devices increases, a slight reduction in accuracy is observed, not surpassing a 2% decrease, demonstrating the robustness of our method to changes in device quantity. CoFormer achieves lower inference latency and total energy consumption using multiple edge devices compared to single-device inference for large transformer models. These benefits stem from parallel execution, which shortens overall inference time and reduces energy consumption, even with more devices involved. However, the marginal gains in latency and energy diminish as the device count grows, due to increasing communication and idle overheads.

V Related Work

Transformer deployment. Deploying transformers on resource-constrained edge devices is challenging due to substantial computational requirements. Recently several works designed lightweight architectures for intelligent applications on such devices. Maaz et al. [16] introduced a convolutional structure into transformers to leverage both convolutional and attention mechanisms, developing an attention mechanism with linear complexity [17]. On the other hand, various methods have employed model compression techniques, including network pruning [19, 20], knowledge distillation [43, 18], low-bit quantization [20] and NAS [21, 22] to gain compact transformer models. For instance, Hao et al. [43] projected intermediate features into an aligned latent space to help heterogeneous student models imitate teacher models.

Collaborative inference. This technique involves dividing a DNN into multiple segments for deployment across different devices to facilitate inference. Recent research in collaborative inference falls into two categories: *(i) Layer-wise splitting. * A DNN is partitioned by layer into multiple groups and distributed across edge nodes, sometimes involving direct transmission of input data to the cloud for inference. Most work focuses on finding an optimal partitioning point and scheduling strategy for balancing performance and efficiency [8, 9, 10, 11]. Yang et al. [8] proposed a deep reinforcement learning algorithm for optimizing task offloading strategies. (ii) Across-layer splitting. It involves dividing the DNN weights into multiple unconnected groups for parallel execution on various edge devices [35, 44, 25, 15, 36, 45]. These methods require an additional training procedure where publicly available pre-trained models cannot be used after splitting, compared to layer-wise splitting. Tuli et al. [44] utilized Multi-Armed-Bandits for determining the optimal division and distribution of DNNs across specific devices.

Existing strategies for deploying transformer models primarily focus on compressing models or designing lightweight architectures tailored for individual edge devices. In this work, we explore an alternative direction by leveraging the inherent divisibility and integrability of transformer architectures to enable collaborative inference across multiple edge devices. Specifically, we propose to decompose large transformer models into segments that are distributed for parallel execution, thereby reducing the per-device computational burden. Among existing collaborative inference techniques, the most closely related approaches are Galaxy [15], DeTransformer [36, 45] and DeViT [35]. Galaxy adopts a tensor-parallel strategy by partitioning the weights of attention and MLP modules across devices. DeTransformer employs block-parallel execution to alleviate communication overhead. However, both methods require multiple rounds of inter-device communication during inference, which introduces significant latency. In contrast, our approach requires only a single round of communication for final result aggregation, substantially reducing communication costs. DeViT decomposes vision transformer for distributed inference on multiple homogeneous edge devices. However, to the best of our knowledge, CoFormer is the first work to automatically decompose off-the-shelf large transformer models for collaborative inference across heterogeneous edge devices to achieve a good balance between efficiency and accuracy.

VI Conclusion

In this paper, we propose a collaborative inference system for general transformers to facilitate efficient deployment across heterogeneous edge devices. An off-the-shelf large transformer can be automatically decomposed into multiple smaller models tailored to the specific hardware constraints. To balance inference performance and efficiency, we formulate an optimization problem to minimize both inference latency and accuracy degradation. We also develop the DeBo algorithm to address this optimization problem, leveraging Bayesian optimization to derive an optimal decomposition policy and progressively calibrate sub-models. Extensive experiments demonstrates our system can accelerate inference speed by 1.7$\times$ to 3.1$\times$ and reduces energy consumption by 36.3% to 63.8% with satisfactory performance.

Appendix A Appendix

In this supplementary material, we introduce how to construct the latency predictor. Besides, we show the process of decomposing the large transformer into multiple sub-models and utilize Bayesian optimization to gain the efficient decomposition policy. The detailed experimental setup and additional experimental results are presented. Table VI is a summary of the used notation.

A-A Latency Predictor Training

In order to find an optimal policy to minimize collaborative inference latency on heterogeneous edge devices, we need to evaluate the hardware latency of decomposed sub-models. There are usually two ways: (1) online measurement: during the optimization process, directly measure the latency on the target device; (2) offline prediction: utilize a predictor to provide the latency without on-device measurement. For the online measurement, a single decomposed sub-model requires hundreds of inferences to get an accurate latency, which lasts for minutes and slows the optimization. We apply the offline prediction here because it is fast and accurate. The architecture of sub-models is encoded into a feature vector, and predict its latency instantly with a multi-layer perceptron (MLP). The process of training the latency predictor is as follows.

Latency dataset collection. Several researchers indicated linear transformation, patch embedding and attention operations are most time-consuming in one transformer inference [26]. The computational cost of these operations mainly depends on the number of transformer layer $l_{n}$, the embedding dimension $d_{n}$, the number of attention heads $h_{n}$ and the neuron dimension $D_{n}$ of MLP. Hence, we evaluate the actual inference latency of transformer models with different architectures on the specific edge device and collect the dataset of (sub-model architecture, measured latency) pairs for all edge devices.

Latency predictor fitness. The collected latency dataset is utilized to train the latency predictor. The input of latency predictor is a feature vector of sub-model architecture with four elements: the number of transformer layer $l$, the embedding dimension $d$, the average number of attention heads $\bar{h}=\sum_{i=1}^{N}h_{i}$ and the average neuron dimension of MLP $\bar{D}=\sum_{i=1}^{N}D_{i}$. We utilize a three-layer MLP with 600 hidden dimension and ReLU activation as the predictor $f(l,d,\bar{h},\bar{D})$ to provide latency because the three-layer MLP is more efficient and accurate. Trained with thousands of real latency data points, the predictor yields high prediction accuracy. Note that the predicted latency is only used in the optimization process. Compared with deducing a closed-form latency model for each hardware, the latency predictor method is more general and faster.

A-B Decomposer: Bayesian Decomposition

We utilize the decomposer to generate decomposition policies for the large transformer as shown in Lines 1 to 11 of Algorithm 1. These policies are evaluated on the evaluator. Evaluation results are utilized to update decomposer until finding the optimal policy. Decomposer decomposes the target transformer into multiple sub-models according to the final policy. We first present the process of decomposing the target transformer model, and then illustrate how to gain efficient decomposition policies.

A-B1 Transformer Decomposition

Given $N$ heterogeneous edge devices, the large ViT is decomposed into $N$ distinct sub-ViTs. As depicted in Figure 14, the decomposition procedure for sub-ViT $n$ encompasses:

• •

Block decomposition: The number of layers in the sub-ViT is decreased to $l_{n}$.

• •

MLP decomposition: Some unimportant neurons in the MLP block are eliminated. The revised neuron dimension for the $j$-th MLP block is $D_{n}^{j}$, where $j\in\{1,2,...,l_{n}\}$.

• •

Head decomposition: Some redundant attention heads are removed, resulting in a reduced number of heads in the $j$-th layer to $h_{n}^{j}$, where $j\in\{1,2,...,l_{n}\}$.

• •

Embedding decomposition: The dimension of embedding is diminished to $d_{n}$.

A-B2 Optimize Decomposition Policy

As analyzed in §III-B3, we treat the objective function in the equation (8) as a black-box-function $\varPsi(\mathcal{C})=\mathcal{L}(\mathcal{C})+\delta T(\mathcal{C})$ to optimize. BO [48] is one of the state-of-the-art techniques for black-box optimization. It applies exploration and exploitation to the objective function by sequentially and actively querying the function values of some input instances. Specifically, BO uses Gaussian process [49] as a prior model to fit the black-box objective function $\varPsi(\mathcal{C})$ in a closed form. It updates the distribution of $\varPsi(\mathcal{C})$ to gain posterior prediction by using its likelihood on newly evaluated $(\mathcal{C},y=\varPsi(\mathcal{C})+\epsilon)$ pairs, where $y$ is a noisy observation of $\varPsi(\mathcal{C})$ and is the sum of inference latency and accuracy degradation in our case. Then it determines the next decomposition policy $\mathcal{C}$ by minimizing an acquisition function, which is computed from the updated posterior. The acquisition function performs a trade-off between exploration and exploitation in evaluating the candidates of decomposition policy. BO repeats the above process until achieving a precise posterior predictive distribution of $\varPsi(\mathcal{C})$.

Prior model. We utilize the Gaussian process (GP) as the prior to model the objective function. A GP prior is specified by its mean function $\mu(\cdot)=\mathbb{E}[\cdot]$ and covariance function $K(\cdot,\cdot)$. The prior of objective function is defined as follows:

[TABLE]

We set $\mu(\cdot)=0$ and $K(\cdot,\cdot)$ to be the Matern kernel [50] with smoothness factor $\nu=1.5$ and length scale $l=1$. $K_{\nu}(\cdot)$ and $\Gamma(\cdot)$ are a modified Bessel and gamma function, respectively. The covariance function is defined as follows:

[TABLE]

Posterior prediction. We sample $r$ different decomposition policy $\mathcal{C^{\prime}}$ satisfying the constraints $\varOmega=\{\Omega_{1},...,\Omega_{N}\}$ and $\varPhi=\{\Phi_{1},...,\Phi_{N}\}$ of all devices, and obtain $r$ sets of decomposed sub-models. For notation simplicity, vectors composed of $\{\mathcal{C}^{\prime}\}_{i=1}^{r}$ and $\{\varPsi(\mathcal{C}^{\prime})\}_{i=1}^{r}$ are denoted as $\mathcal{C}^{\prime}_{1:r}$ and $\varPsi(\mathcal{C}^{\prime}_{1:r})$, respectively. The GP prior indicates a Gaussian distribution over function values, i.e., $\varPsi(\mathcal{C}^{\prime}_{1:r})|\mathcal{C}^{\prime}_{1:r}\thicksim\mathbf{N}(0,\mathbf{K})$, where $\mathbf{K}_{i,j}=K(\mathcal{C}^{\prime}_{i},\mathcal{C}^{\prime}_{j}),\forall i,j\in[1,r]$. We evaluate the objective function denoted by $\hat{y}_{i}$ as a noisy observation of $\varPsi(\mathcal{C}^{\prime})$, i.e., $\hat{y}_{i}=\varPsi(\mathcal{C}^{\prime})+\epsilon$ where Gaussian white noise $\epsilon\thicksim\mathbf{N}(0,\sigma^{2})$. Given the noisy observations $\hat{y}_{1:r}$, the GP posterior of the black-box function can be updated as follows:

[TABLE]

Given a new decomposition policy $\mathcal{C}$ and $\mathbf{k}_{i}=K(\mathcal{C},\mathcal{C}^{\prime}_{i})$, the GP posterior can be used to predict the distribution of $\varPsi(\mathcal{C})$:

[TABLE]

Acquisition function. Given the posterior predictive distribution of $\varPsi(\mathcal{C})$, BO finds the next decomposition policy $\mathcal{C}^{\prime}_{i+1}$ to evaluate based on an acquisition function $u_{i}(\mathcal{C})$ defined by the posterior mean $\mu_{i}(\mathcal{C})$ and standard deviation $\sigma_{i}(\mathcal{C})$. In our algorithm, expected improvement (EI) [51] is utilized as acquisition function. For problem (P1), we can determine the next decomposition policy $\mathcal{C}^{\prime}_{i+1}$ by minimizing EI:

[TABLE]

where $Z$ and $H$ represent the probability density function and cumulative distribution function of the standard normal distribution, respectively. $\varPsi^{*}$ is the best policy in $\mathcal{C}^{\prime}_{1:r}$.

A-C Detailed Experimental Setup

Prototype. We implement our system on a real-world collaborative edge computing testbed that consits of an NVIDIA Jetson Nano with 4 GB memory, an NVIDIA Jetson TX2 with 8 GB memory, an NVIDIA Jetson Orin Nano with 4 GB memory and a switch. These devices represent the edge devices with different hardware resources and computing capabilities. The detailed configurations of these devices are shown in the lower left part of Table VII. Figure 8 shows the implemented hardware platform for CoFormer. These devices are connected via a gigabyte switch TP-LINK TL-SG1008D. For bandwidth control, we use the traffic control tool tc[52], which is able to limit the bandwidth under the setting value. The maximum bandwidth between devices is fixed at 2 Mb/s.

Transformer backbones, tasks and datasets. To evaluate the generality of our system on vision and language tasks, we selected three most typical applications in mobile vision and language systems.

• •

Generic-category image recognition. The vision task aims to recognize the generic category of an image. We select three representative ViT backbones (ViT[23], DeiT[18], Swin[4]) and two widely-used classification datasets (CIFAR-100 [53] and ImageNet-1K [54]). CIFAR-100 contains 50K training images and 10K testing images, uniformly categorized into 100 classes. ImageNet-1K spans 1000 object classes and contains 1,281,167 training images, 50K validation images and 100K test images. They represent a small and large scale image classification tasks, respectively.

• •

Real-time object detection. This type of vision tasks aims to detect objects in real-time with fast processing time while maintaining a required level of accuracy. We choose Mask R-CNN [32] with ViT, DeiT and Swin backbone as the detection model in the most widely-used dataset, MS COCO 2017 [55], to evaluate. MS COCO 2017 contains 0.33 million images, 1.5 million objects, and 80 categories.

• •

General language understanding. The language task aims to understand natural language. We select three language transformer backbones (BERT [24], GPT2 [6] and Flan-T5 [5]) and a representative General Language Understanding Evaluation (GLUE) [56] benchmark to evaluate our method. The GLUE benchmark consists of nine language understanding tasks and cover a diverse range of dataset sizes and degrees of difficulty.

Baselines. We compare CoFormer with state-of-the-art edge deployment methods for transformer models by implementing three types of baselines.

• •

Comparison with large transformer models: Our method can decompose large transformer models for collaborative inference on multiple edge devices. We select three vision transformer models (ViT-L/16, DeiT-B and Swin-L) and three language transformer models (BERT-Large, GPT2-XL and Flan-T5-Large) to demonstrate the effectiveness of our methods.

• •

**Comparison with efficient transformer models:**We compare our methods with several representative lightweight transformer models (*i.e., *single-edge solution illustrated in Figure 9 (b)) to show the superiority: T2T-ViT[33] incorporates a layer-wise transformation to structuralize the image to tokens. EfficientFormer [26] introduces a dimension-consistent pure transformer as a design paradigm to reduce latency. PoolFormer [34] replaces the original attention module in transformers with the simple spatial pooling operator. MobileViTv2 [17] proposes a separable self-attention method with linear complexity and incorporates CNN structures.

• •

Comparison with collaborative inference methods: DeViT [35] is the most relevant baseline of our work, designed for vision transformer inference on multiple homogeneous edge devices. We also compare the ensemble of heterogeneous transformer models including EfficientFormer L1, PoolFormer-S36 and MobileViTv2-100, denoted as Hete-Model-Ens. WDPruning [57] is a weight pruning methods for efficient transformer inference. We adapt WDPruning in different pruning ratios to gain multiple compressed models for collaborative inference.

Evaluation metrics. We employ the Top-1 classification accuracy, mean average precision (mAP) and evaluation scores on GLUE benchmark as the overall performance metric of image classification, object detection and language understanding tasks, respectively. The end-to-end latency, energy consumption and average power are utilized as the efficiency metrics. (i) Measurements of latency: Before measurement, we switch the NVIDIA Jetson devices to the max power mode. For example, the max power of NVIDIA Jetson Nano is 10 Watt, and turn off the dynamic voltage and frequency scaling to ensure a steady measurement environment. For each transformer model, we run it for once as a warm-up and then record the averaged execution time with 50 runs without break using PyTorch Profiler tool [58]. The aim of warm-up running is to alleviate the impact of weight loading and PyTorch initiation since we have omitted these overheads. We use full precision and batch 1 to execute all operations, which can help robustly execute most models and baselines. (ii) Measurements of energy and power: We use a Monsoon High Voltage Power Monitor [59] connected to the edge device to obtain accurate power readings over the course of running a forward pass of each test model 50 times. The device settings are the same as that in latency measurements. From the power signal reported by the monitor, we derive the average per-inference power and energy consumption by first subtracting background power consumption (i.e., power readings when not running any model) and then identifying 50 continuous regions of significantly higher power draw. Each region like that is considered a single inference and we calculate its total energy as the integral over the individual power samples. Analogously we also calculate average power consumption by averaging over the same set of samples. After energy and power are calculated for each inference, the final statistics of a model are obtained by again averaging over the 50 identified runs. The measurement method follows [38] for fair comparison.

Implementation. CoFormer is fully implemented with $~2,000$ LOC in Python in total atop PyTorch [60]. Although we use PyTorch for auto-differentiation and computation graph execution, CoFormer is extensible and can work well with other lightweight ML frameworks such as TF-Lite[61] and MNN[62]. The employed transformer models are trained following DeBo algorithm and deployed on all devices in advance. The communication module is implemented based on gPRC [63]. All the compared approaches are run with timm [64] (vision transformers models) or transformers [65] (language transformer models), for fair comparison.

A-D Additional Experimental Results

Performance comparison on CIFAR-100. The performance comparison of image classification on the small dataset, CIFAR-100 is presented on Figure 15. The original large transformer models used for the image recognition task include ViT-L/16, DeiT-B, and Swin-L. Our method can decompose large transformer models into smaller models and achieve efficient collaborative inference while maintaining satisfactory classification accuracy. Specifically, our decomposed model, based on Swin-L, can accelerate inference speed by 3.11$\times$, reduce energy consumption by 64.0% and save memory usage by 51.9%, on CIFAR-100.

Latency predictor and performance proxy. We compare predicted latency and real inference latency to validate the performance of the latency predictor. The results of NVIDIA Jetson TX2 on CIFAR-100 dataset are shown in Figure 16 (top). We can observe that predicted latency is very close to real latency with an average prediction error (RMSE) of 8.1 ms. As a result, the predictor can yield highly accurate results of inference latency, trained with thousands of real latency data points. On the other hand, we utilize the average validation loss of decompose sub-models as the proxy of accuracy degradation. We conduct experiments to show the effectiveness of the proxy in Figure 16 (bottom). The average validation loss of sub-models without training is relative to the error rate of sub-models trained from scratch. The large the validation loss of sub-models, the lower the error rate. Therefore, the average validation loss is a good proxy of accuracy degradation for sub-models trained from scratch.