Going Beyond the Edge: Distributed Inference of Transformer Models on Ultra-Low-Power Wireless Devices
AuthorsAlexander Gräfe, Ding Huo, Johannes Berger, Marco Zimmerling, Sebastian Trimpe
Resources
This paper shows how tiny wireless IoT devices can team up to run transformer models far too large for any one device by co-designing communication, partitioning, and loss-tolerant training.
Key results
CATS was demonstrated on a BLE mesh of up to 16 nRF52840 MCUs collaborating in inference.
The system can execute a ViT far larger than what a single MCU can run.
SomeGather prunes activation columns to reduce inter-device communication volume while preserving accuracy.
SomeGather lowers the activation memory each device must store compared with unpruned DTI.
Under 10% message loss, MD reduces the relative test-error increase across datasets.
With 90% pruning, CATS accelerates the residual block on the hardware testbed.
What the paper found
The paper introduces CATS, the first end-to-end framework for distributed inference of Transformer models on ultra-low-power wireless MCUs, demonstrated on up to 16 nRF52840 devices over a BLE mesh. Its core novelty is SomeGather, a pruned AllGather primitive that broadcasts only selected activation columns, so cross-device communication is expressed as column-pruned all-to-all exchanges; this simultaneously reduces bandwidth, per-device activation RAM, and flash footprint. CATS partitions multi-head attention along attention heads and feature columns so each device computes local layer normalization, Q/K/V projections, and head-wise attention with minimal inter-device traffic. To handle packet loss, CATS adds message dropout (MD), which trains the model under realistic wireless loss patterns across all layers rather than a single split point. On four time-series forecasting benchmarks, ETT-h2, ICD, London-smart-meters, and Traffic, SomeGather preserves prediction error while normal pruning degrades accuracy as communication savings increase. Quantitatively, SomeGather cuts communication volume by up to 90% and per-device activation RAM by 67.5% compared with unpruned DTI, and MD reduces the test-error increase under 10% message loss from as much as 200% to 23.6% across datasets. In hardware experiments, CATS executes a 14-million-parameter Vision Transformer, about 14× larger than what a single MCU can run; with 90% pruning, latency improves by up to 4.68× on the residual block and 4.37× on the attention block, while the 16-device system remains communication-efficient on a two-hop wireless mesh.
Original abstract
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.
Read the original paper