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

May 25, 2026

Profiling-Driven Adaptive Distributed Transformer Inference on Embedded Edge Deployment

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.

Muhammad Azlan Qazi, Alexandros Iosifidis, Qi Zhang

Aarhus University · Aarhus, Denmark · Tampere University +1

Efficient Transformer Inference Transformer Inference Adaptive Inference Edge Inference

Ask PDF

May 15, 2026

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

Transformer models are rapidly becoming a cornerstone of modern Internet of Things (IoT) applications, yet their computational and memory demands far exceed the capabilities of a single typical ultra-low-power IoT device. We present CATS, a framework for distributed transformer inference on ultra-low-power wireless devices, enabling multiple devices to collaboratively execute models far larger than what a single device can sustain. At its core, CATS is a communication-aware distributed transformer inference scheme co-designed across transformer partitioning, wireless communication and training. It employs SomeGather, a new pruned communication primitive that selectively broadcasts activation columns to reduce communication bandwidth and RAM usage without sacrificing model accuracy. Building on SomeGather, we design a partitioning method that exploits this primitive for efficient model parallelism. To cope with unreliable wireless communication, CATS employs message-dropout during training, which mimics packet losses and yields models that are robust to message loss during inference. In real-world experiments, we show that CATS brings distributed transformer inference to ultra-low-power wireless devices for the first time, with deployments on up to 16 devices that collaboratively execute transformer models up to 14 times larger than what a single device can run.

Alexander Gräfe, Ding Huo, Vincent de Bakker +3

1RWTH Aachen University · 2TU Darmstadt

Efficient Transformer Inference Transformer Inference Edge Computing On-Device Inference

Ask PDF

Mar 20, 2026

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

The key-value (KV) cache is widely treated as essential state in transformer inference, and a large body of work engineers policies to compress, evict, or approximate its entries. We prove that this state is entirely redundant: keys and values at every layer are deterministic projections of the residual stream, and recomputing them from a single residual vector per token incurs exactly zero reconstruction error, not approximately, but bit-identically. We verify this across six models from four architecture families (135M to 4B parameters). Cross-task residual patching at every layer produces D_KL = 0 between patched and original output distributions, confirming that the residual stream satisfies a Markov property and is the sole information-carrying state. Removing the cache entirely and recomputing from scratch yields token-identical output under greedy decoding on all models tested. We build on this result with KV-Direct, a bounded-memory inference scheme that checkpoints residual vectors (5 KB per token on Gemma 3-4B) instead of full KV pairs (136 KB), recomputing keys and values on demand. Over 20 conversation turns, KV-Direct holds peak memory at 42 MB while the standard cache grows past 103 MB. Against five eviction baselines (H2O, StreamingLLM, SnapKV, TOVA, window-only), KV-Direct maintains 100% token match at every cache budget; all baselines degrade to 5-28%. A per-operation latency analysis shows recomputation runs up to 5x faster than reading cached tensors at moderate batch sizes. Code is available at https://github.com/Kaleemullahqasim/KV-Direct.

Kaleem Ullah Qasim, Jiashu Zhang, Muhammad Kafeel Shaheen +2

School of Computing and Artificial Intelligence, Southwest Jiaotong University, Chengdu, 611756, China

Transformer Inference Memory-Efficient Inference KV-Cache Compression

Ask PDF

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