TRACE: Transformer Residual Activation Codebook Encoding for Communication-efficient Distributed Edge Inference

Muhammad Azlan Qazi; Alexandros Iosifidis; Qi Zhang

cs.LGOct 15, 2026

TRACE: Transformer Residual Activation Codebook Encoding for Communication-efficient Distributed Edge Inference

Authors: Muhammad Azlan Qazi, Alexandros Iosifidis, Qi Zhang

Organizations: DIGIT and the Department of Electrical and Computer Engineering, Aarhus University, Denmark · the Data Science Research center, Tampere University, Finland

Abstract

Distributed Transformer inference across co-located edge devices is fundamentally communication-bound, as exchanging activations between sequence-parallel partitions dominates inference latency under commodity edge bandwidths. Existing compression schemes reduce this cost but retain a structural bottleneck: the received data must still undergo key/value projection before attention can proceed, serializing computation behind communication. This paper presents Transformer Residual Activation Codebook Encoding (Trace), a framework that addresses these limitations. Trace places a Vector Quantizer (VQ) at the post-attention residual of every Transformer layer and precomputes codebook-indexed key/value representations. During inference, exchanged activations are encoded into compact codebook indices, allowing receivers to directly recover keys and values through table lookup instead of executing the downstream computations required to derive them. By eliminating the received data downstream computation, Trace enables communication--computation overlap and reduces a key latency bottleneck in distributed Transformer inference. A hierarchical clustering further structures the learned codebook into a multi-resolution hierarchy, enabling runtime adaptation of the communication bitrate without retraining. The framework supports two operating variants: one prioritizing accuracy and the other balancing accuracy and efficiency through buffer-compatible quantization. Experiments on ViT-Base (ImageNet-1K and CIFAR-100) and ViViT-B (UCF-101) show that Trace achieves up to 1.89% higher accuracy and up to 1.27x end-to-end inference speedup compared with the state-of-the-art VQ-based distributed inference framework. Compared with uncompressed sequence-parallel inference, Trace reduces communication overhead by up to 22,624x.

Figures & tables

Fig. 1: Trace system overview. ( Top ) Offline pipeline: finetune Transformer with VQ modules, build a multi-level hierarchy tree via hierarchical k-medoids clustering. ( Bottom-left ) Pre-compute per-layer $\mathbf{K}^{\mathsf{vqcb}}$ and $\mathbf{V}^{\mathsf{vqcb}}$ buffers with frozen weights. ( Bottom-center ) Per-session adaptive inference: select resolution level $m$ from network conditions, exchange $r$ -bit indices between devices overlapped with local FFN computation. ( Bottom-right ) Adaptive encoding with 8-bit metadata header; receiver decodes indices and gathers precomputed K/V buffers — no MatMul on cross-device tokens.

Fig. 2: Comparison of distributed Transformer inference frameworks across two sequence-parallel devices. (a) Trace -b (ours) : transmits only a codebook index; the receiver fetches precomputed $\mathbf{K}$ / $\mathbf{V}$ entries from a buffer and concatenates them with local keys and values, eliminating receiver-side projection and enabling full compute/communication overlap. (b) Astra [ 20 ] : quantizes pre-attention embeddings via VQ; receiver-side projection remains a hard dependency. (c) Prism [ 22 ] : transmits Segment Means (token-group averages); receiver still applies $\mathbf{K}$ / $\mathbf{V}$ projections before attention. (d) Voltage [ 11 ] : transmits the full hidden-state partition; no compression.

Symbol	Description	Symbol	Description
System and dimensions
$P$	Number of co-located edge devices	$L$ , $\ell$	Number of Transformer layers; layer index
$p$ , $p^{\prime}$	Local device index; peer device index ( $p^{\prime}{\neq}p$ )	$N$ , $N_{p}$	Total / per-device sequence length ( $N_{p}{=}N/P$ )
$n$	Token index within a device’s local partition	$D$	Model (hidden-state) dimension
$H$ , $d$	Number of attention heads; per-head dimension ( $d{=}D/H$ )	$d_{\mathrm{ff}}$	Feed-forward (MLP) hidden dimension
Transformer activations

TABLE I: Summary of notation used throughout this paper.

Fig. 3: Compute–communication timing across distributed inference methods for Transformer blocks across devices

p_{0}

p_{1}

p_{2}

. (a) Trace -b (ours): index encoding and transmission run in parallel with the local FFN, making communication and computation overlap and saving communication latency. (b) Astra : compact VQ codewords’ indices are transmitted, but communication and computation are still non-overlapping. (c) Prism : Segment Means compression reduces communication volume but execution remains sequential. (d) Voltage : computation and communication are fully sequential; all devices idle during full-feature AllGather.

Method	Payload (bits)	Depends on $D$ ?
Voltage [ 11 ]	$32N_{p}D$	Yes
Prism [ 22 ]	$32N_{p}D/s$	Yes
Astra [ 20 ] †	$N_{p}\lceil\log_{2}K\rceil$	No
Trace -b (ours)	$N_{p}r+8$	No
Trace -nb (ours)	$N_{p}r+8$	No

TABLE II: Per-device per-layer communication cost (bits transmitted by each device to all peers).

N_{p}{=}N/P

;

D

: model dimension;

s

: PRISM segment size;

K

: codebook size;

r{=}\lceil\log_{2}K_{m}\rceil

: resolution bit-width.

	ImageNet		CIFAR-100
Method	$P{=}2$	$P{=}3$	$P{=}2$	$P{=}3$
Voltage [ 11 ]	80.32	80.32	92.30	92.30
Astra ( $K{=}1024$ ) [ 20 ]	78.45	77.14	89.34	88.12
Prism [ 22 ]	79.66	79.47	90.10	89.69
Trace -nb ( $K{=}1024$ )	80.23	80.11	91.23	90.58
Trace -b ( $K{=}1024$ )	79.90	79.70	90.55	89.90

TABLE III: Top-1 accuracy (%) on ImageNet and CIFAR-100 (ViT-Base/16). Single-device baseline:

80.32\%

(ImageNet),

92.30\%

(CIFAR-100).

Panel of Figure 7 — Fig. 4: Top-1 accuracy on ImageNet-1K and CIFAR-100 using ViT-Base/16 with $P=2$ and $P=3$ devices. The single-device baselines are $80.32\%$ on ImageNet-1K and $92.30\%$ on CIFAR-100. At $K=1024$ , Trace achieves the highest accuracy among the evaluated compression methods in all four settings.

Method	$P$	Compute	Comm.	Effective †
Voltage [ 11 ]	2	92.5	2770.31	2862.8
Prism ( $s{=}10$ ) [ 22 ]	2	95.2	140.63	235.9
Astra ( $K{=}1024$ ) [ 20 ]	2	115.0	1.81	116.8
Trace -nb ( $K{=}1024$ )	2	119.0	1.34	120.3
Trace -b ( $K{=}1024$ )	2	106.0	1.14	106.0
Trace -b ( $K{=}32$ )	2	92.0	0.58	92.0

TABLE IV: Per-device latency breakdown (ms) at

10

Mbps on ViT-Base/16 (2-core CPU).

Panel of Figure 9 — Fig. 5: Effective per-device latency (ms) at 10 Mbps on ViT-Base/16. Dark bars show compute latency; light extensions indicate communication latency when executed sequentially. Trace -b at $K\leq 32$ reaches the single-device compute floor (92 ms at $P=2$ , 63 ms at $P=3$ ).

Fig. 6: Communication latency (ms) versus link bandwidth (Mbps) on ViT-Base/16 (log scale). Trace (

K=2

) achieves

22{,}624\times

lower communication latency than Voltage and

15\times

lower than Astra at

P=2

. The Trace curves indicate that its communication advantage is maintained consistently across the evaluated 1–20 Mbps bandwidth range.

Fig. 7: Compression ratio for ViT-Base/16 at

P=2

and

P=3

. Trace (

K=1024

) achieves

2{,}426\times

P=2

, a

1.59\times

improvement over Astra (

1{,}528\times

) at the same codebook size. Coarser levels yield near-proportional gains, reaching

22{,}624\times

K=2

Panel of Figure 12 — Fig. 8: Compression ratio for ViViT-B/16 on UCF-101 across $P\in\{2,4,10\}$ devices. The compression ratio at each resolution level varies by less than $0.5\%$ across all device counts, confirming that Trace ’s compression mechanism is invariant to the number of participating devices.

Fig. 9: Top-1 accuracy on UCF-101 (ViViT-B/16,

K{=}1024

) as the number of devices scales from

2

10

. Trace degrades by only

1.53

percentage points from

P{=}2

(

90.25\%

) to

P{=}10

(

88.72\%

) relative to the single-device baseline (

90.77\%

, dashed line).

Devices	Without buffers	With buffers	Gain (%)
2	3441	2329	32.3
3	2281	1596	30.0
4	1766	1320	25.2
6	1213	952	21.5
8	917	719	21.6
10	742	627	15.5

TABLE V: Effect of precomputed K/V buffers on per-device compute latency (ms) for ViViT (

K{=}1024

). “Without” reconstructs cross-device K/V via the full projection chain; “With” uses buffer lookup.

Case	Quantized-copy treatment	$P{=}2$	$P{=}3$
1 A.i	$\mathbf{X}_{p}^{(0)}$ residual at block 1; $\tilde{\mathbf{Z}}_{p}^{(\ell-1)}$ thereafter ( Trace -nb)	91.23	90.58
2 B	Post-residual VQ; $\mathrm{VQ}_{0}$ K/V at block 1 ( Trace -b)	90.51	89.90
2 A	Post-residual VQ; full $\mathbf{X}$ K/V at block 1	90.52	88.61
1 A.ii	No residual at block 1; $\tilde{\mathbf{Z}}_{p}^{(\ell-1)}$ thereafter	90.78	89.03
1 A.iii	No residual in any block	90.51	88.59
1 B	$\mathrm{VQ}_{0}$ as the block-1 residual and K/V	90.90	89.56

TABLE VI: VQ placement ablation on CIFAR-100 (Accuracy %) using ViT-Base/16. Trace uses post-residual VQ with

\mathrm{VQ}_{0}

for the first-block K/V.

		Resolution level $m$
$P$	Metric	0	1	2	3	4
2	Euclidean	84.91	89.18	89.89	90.52	90.55
	Cosine	83.30	88.40	89.53	90.32	90.55
	Frequency	85.56	89.29	89.73	89.72	90.55
3	Euclidean	83.33	85.81	87.70	88.54	89.90
	Cosine	82.59	85.30	86.28	87.60	88.90

TABLE VII: Hierarchy construction metric ablation: top-1 accuracy (%) on CIFAR-100 at each resolution level

m

(level

4

= full codebook,

K{=}1024

Update strategy	$P{=}2$	$P{=}3$
Gradient-based	90.47	88.44
EMA ( Trace default)	90.51	89.00
EMA + dim. reduction ( $d_{q}{=}256$ )	90.79	89.30

TABLE VIII: Codebook update strategy ablation: top-1 accuracy (%) on CIFAR-100 (ViT-Base/16). Single-device baseline:

92.30\%

Efficient Transformer Inference

Efficient Transformer Inference studies methods and systems that reduce the latency, memory, and energy needed to run Transformer models.

See all 121 papers

Profiling-Driven Adaptive Distributed Transformer Inference on Embedded Edge Deployment
May 25, 2026Muhammad Azlan Qazi, Alexandros Iosifidis, Qi ZhangEfficient Transformer Inference Transformer Inference
Aarhus University · Aarhus, Denmark · Tampere University +1
Going Beyond the Edge: Distributed Inference of Transformer Models on Ultra-Low-Power Wireless Devices
May 15, 2026Alexander Gräfe, Ding Huo, Vincent de Bakker +3Efficient Transformer Inference Transformer Inference
1RWTH Aachen University · 2TU Darmstadt
The Residual Stream Is All You Need: On the Redundancy of the KV Cache in Transformer Inference
Mar 20, 2026Kaleem Ullah Qasim, Jiashu Zhang, Muhammad Kafeel Shaheen +2Transformer Inference Memory-Efficient Inference
School of Computing and Artificial Intelligence, Southwest Jiaotong University, Chengdu, 611756, China

May 25, 2026J/K move · Enter open · S save

Profiling-Driven Adaptive Distributed Transformer Inference on Embedded Edge Deployment

Muhammad Azlan Qazi, Alexandros Iosifidis, Qi Zhang

Aarhus University · Aarhus, Denmark · Tampere University+1

Efficient Transformer Inference Transformer Inference Adaptive Inference Edge Inference

Open paper AskPDF arXiv

Distributing Transformer inference across embedded edge devices can alleviate individual memory and compute constraints, yet practical benefits on real hardware remain unclear: prior work relies largely on simulations that overlook hardware-specific communication overheads. We present a hardware prototype study on NVIDIA Jetson Orin Nano devices connected over WiFi. Our key finding is that the dominant bottleneck is not just network bandwidth but also the CPU-GPU staging during communication. Because Jetson's integrated GPU architecture lacks the PCIe/NVLink pathway that NCCL requires, all inter-device data communication should be routed through GLOO and staged in CPU memory; an overhead that scales with communication data volume and makes full-tensor exchange slower than single-device inference across the batch sizes for medium sized models such as ViT. We therefore evaluate Prism by combining Segment Means compression with lightweight offline profiling to adaptively select between local and distributed execution at runtime. Experiments show that this strategy reduces latency by 65%-77% and energy consumption by 34%-52% relative to full-tensor exchange in static distributed execution setup, demonstrating that profiling-driven adaptation is essential for practical distributed Transformer inference on embedded hardware.

Abstract

Efficient Transformer Inference

Explore similar work

Profiling-Driven Adaptive Distributed Transformer Inference on Embedded Edge Deployment

Going Beyond the Edge: Distributed Inference of Transformer Models on Ultra-Low-Power Wireless Devices

The Residual Stream Is All You Need: On the Redundancy of the KV Cache in Transformer Inference