Communication-Efficient Multi-Device Inference Acceleration for Transformer Models

## Abstract

Abstract Transformer models power many AI applications but suffer from high inference latency, limiting their use in real-time settings. Multi-device inference can reduce latency by parallelizing computation. Yet, existing methods require high inter-device bandwidth, making them impractical for bandwidth-constrained environments. We propose Astra , a communication-efficient framework that accelerates Transformer inference through a novel integration of sequence parallelism and a Mixed-Precision Attention mechanism designed to minimize inter-device communication. Astra compresses non-local token embeddings via vector quantization and preserves task accuracy through two optimizations, Noise-Augmented Quantization and Distributed Class Tokens. Experiments on ViT and GPT2 across vision and NLP tasks show that Astra achieves up to 2.64 × \times × speedups over single-device inference and up to 15.25 × \times × speedups over state-of-the-art multi-device inferences, while operating under bandwidths as low as 10 Mbps. Astra is open-sourced at https://github.com/xl1990/Astra .

## 1 Introduction

Transformer models dosovitskiy2020vit; devlin2019bert; radford2019gpt have become central to modern AI applications in both vision and language domains. As models grow in scale to improve accuracy, their inference time becomes prohibitive, particularly for real-time applications or latency-sensitive user experiences. To address this bottleneck, a rich body of work has focused on optimizing single-device inference through techniques like quantization liu2021quantize, pruning kwon2022prune, and knowledge distillation lin2022distill. However, the scope for improvement on a single device remains fundamentally limited by hardware capacity. Multi-device inference complements the above optimization techniques by parallelizing model execution across multiple computing devices.
In theory, it can reduce latency nearly linearly with the number of devices. This setup is especially attractive in practical settings where resource-constrained edge devices or distributed consumer-grade hardware can collaborate to process sporadic inference requests.

Our analysis, however, reveals a critical bottleneck that has been overlooked: in bandwidth-constrained settings (<=100 Mbps), inter-device communication dominates the existing multi-device inference time, accounting for 58.55-93.47% of total latency (see Figure [3](https://arxiv.org/html/2505.19342v1#S4.F3)).
This communication overhead negates the theoretical speedups promised by existing multi-device approaches.
Existing methods, such as Tensor Parallelism (TP) in Megatron-LM shoeybi2019megatron, Sequence Parallelism (SP) in Voltage hu2024voltage, and Block Parallelism (BP) in DeTransformer du2024detransformer, rely on high-volume communication bandwith to achieve modest speedups (see Figure [1](https://arxiv.org/html/2505.19342v1#S1.F1)).
This requirement far exceeds what is commonly available in bandwidth-constrained environments. For example, indoor Wi-Fi networks (e.g., Wi-Fi 4/5/6) typically deliver practical throughput in the range of 50–300 Mbps, depending on interference, distance from the router, and device capabilities. As a result, current approaches are ill-suited for wireless or edge deployments, where lower and more variable bandwidth is the norm wifi.

Figure: Figure 1: Latency speedup for existing multi-device inference methods and our proposed method Astra with different groups $G$ under different bandwidths with 4 devices and 1024 input tokens. BP: Block Parallelism, SP: Sequence Parallelism, TP: Tensor Parallelism, AG: Allgather. A smaller $N_{b}$ means fewer communication for BP.
Refer to caption: x1.png

In this work, we propose Astra, a communication-efficient framework for accelerating Transformer inference across multiple devices. Astra rethinks how attention computation is distributed by building on sequence parallelism, which assigns different input tokens to different devices. The key innovation in Astra is a Mixed-Precision Attention mechanism that dramatically reduces communication cost: local attention is computed using full-precision embeddings, while remote tokens are encoded using low-bit vector quantization before transmission. These compressed embeddings are decoded on the receiving end and used for approximate attention computation.

To preserve accuracy, Astra incorporates two key designs.
First, a Noise-Augmented Vector Quantization strategy injects Gaussian noise to vector-quantized token embeddings during training, enhancing the diversity of the quantized feature space and improving generalization to unseen inputs.
Second, a Distributed Class Token scheme assigns each device its own local class token, which attends to uncompressed local tokens with full precision and vector-quantized tokens from other devices.
The outputs from these distributed class tokens are then aggregated as a more comprehensive embedding before the final prediction.
Through these innovations, Astra makes multi-device inference viable in low-bandwidth scenarios while maintaining high accuracy.

Our main contributions are:

- •
Astra— a new multi-device Transformer inference framework that significantly reduces communication overhead via a novel Mixed-Precision Attention mechanism. We design two techniques, Noise-Augmented Quantization and Distributed Class Tokens, to maintain high accuracy under aggressive communication compression.
- •
*Comprehensive Evaluation* — We evaluate Astra across different Transformer models, including ViT dosovitskiy2020vit and GPT2 radford2019gpt, on both vision tasks (e.g., CIFAR-100 krizhevsky2009lcifar and ImageNet-1K deng2009imagenet) and NLP tasks (e.g., Wikipedia and Wikitext-103 merity2016wikitext). Astra reduces the bandwidth requirement for multi-device speedup from 500 Mbps to as low as 10 Mbps, achieving up to $2.64\times$ end-to-end latency speedup even under constrained bandwidth where prior methods fail (Figure [1](https://arxiv.org/html/2505.19342v1#S1.F1)). It maintains task accuracy within $3.58\%$ of the original model under extreme communication compression.
- •
*Design Analysis* — Through detailed ablation studies, we demonstrate the effectiveness of each key design component.
The Mixed-Precision Attention Computation Mechanism achieves a latency speedup of up to $171.82\times$ compared to the original attention under constraint bandwidth, the Noise-Augmented Vector Quantization improves accuracy by up to $0.86\%$ compared to naive vector quantization, and the Distributed Class Tokens design further enhances accuracy by up to $7.13\%$ compared to using a single class token.

## 2 Background and Related Work

This section introduces existing techniques for multi-device inference acceleration and then the vector quantization (VQ) technique that is necessary to understanding Astra. A broader survey of related work is provided in Appendix [A](https://arxiv.org/html/2505.19342v1#A1).

Parallelization in Multi-Device Inference.
To accelerate inference latency, many recent approaches explore distributing Transformer computations across multiple devices. Techniques originally developed for training, such as data parallelism, pipeline parallelism, and model parallelism, have been repurposed for inference.
Data parallelism li2014communication and pipeline parallelism huang2019gpipe improve throughput by processing different input samples or partitioning layers across devices. However, these methods are primarily suited to batch processing and offer limited benefits for latency-sensitive, single-request inference.
Model parallelism, in contrast, splits the computation of individual model layers across devices, enabling true per-request acceleration. Notable examples include:
Megatron-LM shoeybi2019megatron, which partitions model weights (tensor parallelism), Voltage hu2024voltage, which partitions token sequences (sequence parallelism), and DeTransformer du2024detransformer, which decomposes layers into smaller blocks for distributed execution.
These approaches reduce some communication costs by tailoring the parallelization granularity, but they still require substantial bandwidth, such as 500 Mbps, to realize modest speedups. While this bandwidth may be achievable in high-performance server clusters, it exceeds the practical limits of many real-world environments, such as indoor Wi-Fi, where sustained throughput often falls lower than 300 Mbps depending on interference wifi.
Our work diverges from prior efforts by focusing explicitly on reducing the communication cost, which we identify as a primary bottleneck in multi-device Transformer inference.

Vector Quantization (VQ).
Astra leverages vector quantization to compress token embeddings before inter-device transmission, significantly reducing communication overhead. VQ maps continuous vectors to discrete indices in a fixed codebook, allowing compact representation of information with a small number of bits vq.

Vanilla VQ partitions the feature space into $K$ clusters, each represented by a centroid.
Let $\mathbf{e}\in\mathbb{R}^{K\times D}$ be the codebook of centroids. For an input vector $\mathbf{z}\in\mathbb{R}^{D}$, VQ selects the closest centroid by $i=\arg\min_{i}\|\mathbf{e}_{i}-\mathbf{z}\|_{2}^{2}.$
and transmits only the index $i$, requiring $\log_{2}K$ bits per token. The codebook $\mathbf{e}$ can be learned via $K$-means clustering or updated online using exponential moving averages vqvae. This compact representation greatly reduces the communication cost between devices.

Grouped VQ GVQ extends this idea by dividing the input vector into $G$ equal-length sub-vectors and quantizing each independently using separate codebooks. It increases expressiveness in the compressed representation and improves task accuracy, at the cost of higher communication overhead, $G\cdot\log_{2}K$ bits per token.
In experiments, Astra with Grouped VQ outperforms Vanilla VQ counterpart in accuracy, offering a tunable trade-off between bandwidth usage and model performance.

## 3 The Astra Framework

We present Astra, a communication-efficient multi-device inference framework for Transformer models. Astra is designed to minimize inter-device communication while preserving accuracy, enabling fast inference even in bandwidth-constrained environments. The framework achieves this through three key innovations: (1) Mixed-Precision Attention, (2) Noise-Augmented Vector Quantization, and (3) Distributed Class Tokens. We begin with an overview of the inference workflow, then describe each core design in detail.

### 3.1 Overview of Astra

Figure: Figure 2: Overview of Astra with two devices. We introduce three key innovations: (1) Mixed-Precision Attention, (2) Noise-Augmented Vector Quantization, and (3) Distributed Class Tokens to achieve communication-efficient multi-device inference. Astra can be applied to transformers for both deterministic and generative tasks.
Refer to caption: x2.png

Figure [2](https://arxiv.org/html/2505.19342v1#S3.F2) illustrates the inference procedure of the Astra framework.
Given an input sequence consisting of a class token (optional) and $T$ content tokens, Astra first partitions the content tokens evenly across $N$ devices, assigning $T/N$ tokens to each device.
To support classification and similar global tasks, the class token CLS is replicated to each device (*Distributed Class Token* in §[3.3](https://arxiv.org/html/2505.19342v1#S3.SS3)).
Each device thus holds a disjoint subset of the input sequence, along with its own class token copy, and maintains a full copy of the Transformer model.

Within each Transformer block, the inference procedure proceeds in parallel across devices.
Each device first applies our *Noise-Augmented Vector Quantization* (see §[3.3](https://arxiv.org/html/2505.19342v1#S3.SS3)) to its local tokens, and transmits the corresponding low-bit indices to other devices.
Each device now has access to the full input sequence, full-precision for local tokens, and vector quantized versions for non-local tokens.
It then performs *Mixed-Precision Attention Computation* (see §[3.2](https://arxiv.org/html/2505.19342v1#S3.SS2)), computing attention maps over this hybrid set of representations.
Since the feed-forward network (i.e., MLP) is position-wise independent, it is executed locally on each device without additional communication.

After all Transformer blocks, all class tokens, replicated across devices, are aggregated via average pooling to form a single unified representation.
For classification tasks, this pooled class token is passed to a downstream prediction head to generate the final output.
For generative tasks such as next-token prediction, there is no class token. Instead, the input sequence is evenly partitioned across devices for parallel encoding, which accelerates the initial digestion time when inference.
Once the token encoding is complete, autoregressive decoding proceeds sequentially on a single device that holds the final (i.e., most recent) token in the input sequence.

### 3.2 Extreme Communication Compression

Transformer attention requires global context aggregation, which poses a communication bottleneck when the input tokens are partitioned across devices.
In detail, the self-attention layers must use information from all tokens in the sequence, including those stored on other devices, to compute attention for each local token.
As a result, each Transformer block must perform an all-to-all exchange of embeddings across devices, leading to significant communication overhead.
For example, if the input tokens are evenly partitioned across $N$ devices, with each holding $T/N$ tokens, then transmitting full-precision embeddings requires sending $T/N\times D\times r$ bits per device per block, where $D$ is the hidden dimension and $r$ is the precision (e.g., float32). Such overhead becomes prohibitive under realistic bandwidth constraints.

Mixed-Precision Attention.
To address this challenge, we propose to leverage both full-precision and compressed token representations during attention computation.
Specifically, each device computes full-precision attention among its local tokens, and approximates attention interactions with non-local tokens (i.e., those stored on other devices) using vector-quantized embeddings.
Only the low-bit indices of the vector-quantized embeddings are transmitted between devices.

Formally, for each local query $\mathbf{q}$, we compute attention over a mixed set of key and value pairs, full-precision representations from local tokens and vector-quantized representations from non-local tokens.
Therefore, the attention map is computed as:

$$ $\text{Attention}(\mathbf{Q},\mathbf{K},\hat{\mathbf{K}},\mathbf{V},\hat{ \mathbf{V}})=\text{softmax}\left(\frac{\mathbf{Q}[\mathbf{K}\mid\hat{\mathbf{K }}]^{\top}}{\sqrt{d_{k}}}\odot\mathbf{M}\right)[\mathbf{V}\mid\hat{\mathbf{V}}],$ (1) $$

where $\hat{\mathbf{K}}$ and $\hat{\mathbf{V}}$ are derived from the vector-quantized embeddings $\hat{\mathbf{X}}$ via linear projections. The operator $\mid$ denotes row-wise concatenation.
During the training stage, the attention mask $\mathbf{M}$ ensures that full-precision interactions are only applied between local tokens, while interactions with non-local tokens use their compressed counterparts.

Vector-Quantized Non-Local Tokens.
The non-local compressed embeddings $\hat{\mathbf{X}}$ in Mixed-Precision Attention are produced by a vector quantization (VQ) module.
Prior to transmission, each token embedding $\mathbf{X}$ is quantized by a codebook to an index $i$ via nearest-neighbor lookup.
Since the codebook is shared across devices for each Transformer block, the receiving device can reconstruct $\hat{\mathbf{X}}$ using only the transmitted index $i$.
This reduces the per-token communication cost from $rD$ bits to $\log_{2}K$ bits, where $K$ is the codebook size, resulting in a compression ratio of $2457.6\times$ when $r=32$, $D=768$, and $K=1024$.

The VQ module is jointly trained with the Transformer model.
Specifically, the codebook is initialized by running $K$-means clustering over intermediate token embeddings from the pretrained model, and is further updated via exponential moving average during model fine-tuning.
Following VQVAE vqvae, we further introduce a commitment loss to encourage token embeddings to stay close to their assigned centroids. The overall training objective is:

$$ $\mathcal{L}=\mathcal{L}_{t}+\beta\|\mathbf{X}-\text{sg}(\hat{\mathbf{X}})\|_{2 }^{2},$ (2) $$

where $\mathcal{L}_{t}$ is the task loss, $\text{sg}(\cdot)$ denotes the stop-gradient operation, and $\beta$ controls the strength of the commitment term.
This design encourages the model to align token embeddings with their corresponding quantized representations, which improves downstream task performance even under aggressive compression.
Appendix [E](https://arxiv.org/html/2505.19342v1#A5) further empirically demonstrates that the commitment term and a well-configured loss weight is necessary for maintaining accuracy.

### 3.3 Task Accuracy Preservation

Noise-Augmented Vector Quantization.
Quantizing embeddings introduces discretization error, which can degrade model generalization.
To mitigate this, we propose a novel regularization strategy called Noise-Augmented Vector Quantization (NAVQ), which adds Gaussian noise to quantized embeddings during training.
This technique is inspired by the Vicinal Risk Minimization (VRM) principle chapelle2000vicinal, which improves generalization by exposing the model to perturbations around each data point in input space.
NAVQ extends this philosophy into the latent quantized embedding space.
Instead of directly using the deterministic quantized embedding $\hat{\mathbf{X}}$ during training, we compute a noise-augmented version,
$\tilde{\mathbf{X}}=\hat{\mathbf{X}}+\lambda\xi$,
where $\lambda\in(0,1]$ controls the noise magnitude and $\xi\sim\mathcal{N}(\boldsymbol{\mu},\boldsymbol{\Sigma})$ is Gaussian noise sampled from the distribution of quantization residuals $\boldsymbol{\varepsilon}:=\mathbf{X}-\hat{\mathbf{X}}$. This residual captures the error introduced by quantization, and the noise distribution is fit with empirical mean $\boldsymbol{\mu}$ and covariance $\boldsymbol{\Sigma}$ over training data.

By injecting noise into quantized embeddings during training, NAVQ restores a degree of continuity to the otherwise discrete latent space, encouraging the model to generalize across small perturbations and reducing sensitivity to codebook boundaries. At inference time, the noise is omitted and the model operates deterministically using $\hat{\mathbf{X}}$. We theoretically justify this approach in Appendix [B](https://arxiv.org/html/2505.19342v1#A2), where we prove the following:

###### Theorem 1 (Noise-Augmented Embeddings Improve Distributional Fidelity) .

Let $\hat{\mathbf{X}}$ denote the quantized embedding of $\mathbf{X}$, and let $\tilde{\mathbf{X}}=\hat{\mathbf{X}}+\lambda\xi$ with $\xi\sim\mathcal{N}(\boldsymbol{\mu},\boldsymbol{\Sigma})$ sampled from the quantization residuals. Then the 2-Wasserstein distance between the original embedding distribution $P_{\mathbf{X}}$ and the perturbed distribution $P_{\tilde{\mathbf{X}}}$ satisfies:

$$ $W_{2}^{2}(P_{\mathbf{X}},P_{\tilde{\mathbf{X}}})<W_{2}^{2}(P_{\mathbf{X}},P_{ \hat{\mathbf{X}}}),$ (3) $$

i.e., the noise-augmented distribution is statistically closer to the true distribution than the raw quantized embedding.

Empirically, NAVQ reduces overfitting and improves generalization. As shown in our ablation study (§[4.4](https://arxiv.org/html/2505.19342v1#S4.SS4)), setting $\lambda=1.0$ improves validation accuracy by 0.86% compared to training without noise, demonstrating the effectiveness of this regularization under extreme compression.

Distributed Class Tokens.
In Transformer-based classification models, such as ViT dosovitskiy2020vit, a special class token is prepended to the input sequence and used to aggregate information from all other tokens through attention.
However, in the context of our Mixed-Precision Attention mechanism, attention between tokens on different devices is performed using vector-quantized embeddings.
If the class token is assigned to a single device, it will attend to full-precision local tokens but only see vector-quantized representations from other devices.
This asymmetric access to information introduces a bias in the class token’s representation, potentially limiting its ability to effectively summarize the entire input sequence.

To address this issue, we introduce the Distributed Class Token mechanism.
Instead of assigning the class token to a single device, we replicate it across all devices, creating one local copy per device.
Each replica computes attention with full-precision local tokens and quantized non-local tokens.
At the end of the model, all class token replicas are aggregated (e.g., via mean pooling) into a single vector, which is then passed to the final prediction head. This approach not only restores symmetry in access to information but also reduces estimation error in the attention output, improving robustness to quantization artifacts. We formally justify this mechanism in Appendix [C](https://arxiv.org/html/2505.19342v1#A3) and prove the following:

###### Theorem 2 (Variance Reduction via Distributed Class Tokens) .

Let $\mathbf{h}$ denote the class token embedding of a full-precision global attention computation. Let $\tilde{\mathbf{h}}_{\mathrm{single}}$ be the output of a single-device class token using Mixed-Precision Attention, and let $\tilde{\mathbf{h}}_{\mathrm{dist}}$ be the average of $N$ distributed class token outputs. Then:

$$ $\mathbb{E}\bigl{[}\lVert\tilde{\mathbf{h}}_{\mathrm{dist}}-\mathbf{h}\rVert_{2 }^{2}\bigr{]}=\frac{1}{N}\mathbb{E}\bigl{[}\lVert\tilde{\mathbf{h}}_{\mathrm{ single}}-\mathbf{h}\rVert_{2}^{2}\bigr{]},$ (4) $$

i.e., distributed class tokens reduce the expected attention output error by a factor of $1/N$.

Empirically, our ablation study (§[4.4](https://arxiv.org/html/2505.19342v1#S4.SS4)) confirms that Distributed Class Tokens consistently outperform the single-token variant across all evaluated settings, yielding accuracy improvements between 0.37% and 7.13% depending on the compression level and commitment loss weight.

###### Theorem 1 (Noise-Augmented Embeddings Improve Distributional Fidelity) .

###### Theorem 2 (Variance Reduction via Distributed Class Tokens) .

## 4 Empirical Evaluation

This section evaluates the effectiveness of Astra by answering the following question: (1) How well does Astra maintain model accuracy under aggressive token compression compared to the original model? (2) Can Astra deliver substantial inference speedups compared to existing multi-device inference methods with minimal bandwidth usage? (3) How effective are the optimizations in Astra?
We answer these questions through extensive experiments across Transformer models (ViT and GPT2), application domains (vision and NLP tasks), and deployment conditions (varying bandwidth, device count, and compression settings).

### 4.1 Experimental Setup

Environment.
Astra is implemented in PyTorch 2.5 and trained on a single L40S GPU with 40GB memory. For deployment, we simulate distributed inference on personal laptops provisioned with an NVIDIA 1660Ti GPU. We emulate a range of network conditions by enforcing bandwidth caps, enabling realistic evaluation under constrained environments. Unless stated otherwise, experiments use 4 devices.

Transformer Models.
We evaluate Astra across both encoder and decoder Transformer architectures:
For encoder architecture, we focus on vision tasks with Vision Transformer (ViT-Base) dosovitskiy2020vit.
For decoder architecture, we conduct experiments on NLP tasks with GPT2-Small (GPT2-S) and GPT2-Medium (GPT2-M) radford2019gpt.

Datasets and Metrics.
We evaluate Astra on both vision and language tasks.
For vision, we use CIFAR-100 krizhevsky2009lcifar and ImageNet-1K deng2009imagenet, reporting top-1 classification accuracy.
For language modeling, we perform next-word prediction using two datasets, English Wikipedia and Wikitext-103 merity2016wikitext, and report perplexity (PPL; lower is better).
The evaluation includes three settings: training and evaluating on Wikipedia, training and evaluating on Wikitext-103, and a zero-shot evaluation where the model is trained on Wikipedia but directly evaluated on the Wikitext-103 validation set without further fine-tuning.
The last zero-shot setting follows the evaluation used in the original GPT2 radford2019gpt and serves to assess the model’s generalization to unseen domains.

Baselines.
We compare Astra with both single-device and three multi-device inference approaches.

- •
Original Model: The baseline model runs entirely on a single device using float32 precision.
- •
Tensor Parallelism (TP): Represented by Megatron-LM shoeybi2019megatron, which partitions weight matrices across devices and requires two allreduce operations per Transformer layer.
- •
Sequence Parallelism (SP): Introduced by Voltage hu2024voltage, which partitions the input sequence and performs one allgather operation per layer.
- •
Block Parallelism (BP): Proposed by DeTransformer du2024detransformer, which replaces Transformer blocks with multiple parallel sublayers for distributed execution.

For BP, we evaluate two most efficient design variants proposed in du2024detransformer:
(i) BP+Allgather (BP+AG) minimizes communication by performing more local computation, and
(ii) BP+SequenceParallel (BP+SP) reduces local computation at the cost of moderate communication overhead.
Both variants include a hyperparameter $N_{b}$ that controls the number of original Transformer blocks retained.
A smaller $N_{b}$ leads to fewer communications and thus lower latency.
We compare against BP configurations with $N_{b}=1$ and $N_{b}=4$.

Astra Settings.
For the Noise-Augmented Vector Quantization in Astra, the codebook size is 1024, representing each transmitted token with 10 bits (i.e., $\mathrm{log}_{2}1024$).
The noise magnitude $\lambda$ is 1.0 in the main results and we test other settings including $\lambda\in\{0.0,0.1,0.3\}$ in §[4.4](https://arxiv.org/html/2505.19342v1#S4.SS4).
We further evaluate the use of Grouped VQ, as introduced in §[2](https://arxiv.org/html/2505.19342v1#S2), which splits each input vector into $G$ groups and applies vector quantization independently to each group using separate codebooks.
We experiment with group sizes of 16 and 32, in addition to Vanilla VQ of a single group. Increasing the number of groups leads to a higher bits per token and thus reduces the overall compression ratio proportionally.

We load the pre-trained weights for all the transformer models from the HuggingFace official model zoo.
Then Astra is fine-tuned for additional iterations on each dataset using the Adam optimizer kingma2014adam.
Specifically, for vision tasks, Astra is fine-tuned on CIFAR-100 and ImageNet-1K for 32 and 4 epochs, respectively.
For NLP tasks, Astra is fine-tuned on 1 million samples from English Wikipedia and the complete Wikitext-103 dataset for 1 epoch.
During fine-tuning, we test with different commitment loss weights $\beta\in\{0.0001,0.0002,0.0005\}$ in §[4.4](https://arxiv.org/html/2505.19342v1#S4.SS4) and report the best accuracy performance in §[4.2](https://arxiv.org/html/2505.19342v1#S4.SS2).

### 4.2 Results on Accuracy and Communication Costs

We evaluate the accuracy performance of Astra across three Transformer architectures, ViT-Base, GPT2-S, and GPT2-M, on vision and NLP benchmarks.
Notice that we only report the baseline accuracy for the original model.
Existing multi-device baselines, including Megatron-LM shoeybi2019megatron and Voltage hu2024voltage mentioned in §[4.1](https://arxiv.org/html/2505.19342v1#S4.SS1), do not incur any accuracy loss since they merely reorganize computation without altering the model’s numerical outputs.
Therefore, their results are equivalent to the original model and are omitted here for clarity.
Alongside accuracy, we also report the associated communication overhead, measured as the total amount of data exchanged per token during a single forward pass (i.e., Total Bits per Token).
Results are summarized in Tables [1](https://arxiv.org/html/2505.19342v1#S4.T1) and [3](https://arxiv.org/html/2505.19342v1#S4.T3).

ViT-Base.
Astra maintains high accuracy on ViT-Base for image classification (CIFAR-100 and ImageNet-1K), with less than 3.58% degradation even under 2457.6$\times$ compression, as shown in Table [1](https://arxiv.org/html/2505.19342v1#S4.T1).
With 32 groups, Astra achieves 91.64% on CIFAR-100 and 80.05% on ImageNet-1K, closely matching the original performance of 92.53% and 80.32%.
To further assess scalability, we fix the compression configuration (32 groups) and evaluate Astra on CIFAR-100 using varying numbers of devices.
As shown in Table [2](https://arxiv.org/html/2505.19342v1#S4.T2), Astra preserves model accuracy within 1.39% of the original model across different device counts.

Figure: Table 1: Task accuracy and communication overhead on CIFAR-100 and ImageNet-1K with ViT-Base.

**Table 3: Task performance (i.e., perplexity) and communication overhead on Wikipedia and Wikitext-103 with GPT2.**
| Model | #Groups | Datasize<br>per Token | Compression<br>Ratio | Wikipedia | Wikitext-103 |  |
| --- | --- | --- | --- | --- | --- | --- |
| Fine-Tuned | Zero-Shot |  |  |  |  |  |
| GPT2 - S | - | 294912 | - | 15.79 | 18.96 | 58.91 |
| Astra | 1 | 120 | 2457.6 | 21.46 | 25.98 | 120.7 |
| 16 | 1920 | 153.6 | 18.84 | 23.22 | 94.8 |  |
| 32 | 3840 | 76.8 | 17.39 | 20.95 | 76.24 |  |
| GPT2 - M | - | 786432 | - | 12.16 | 14.8 | 43.22 |
| Astra | 1 | 240 | 3276.8 | 17.86 | 21.97 | 96.99 |
| 16 | 3840 | 204.8 | 14.43 | 18.03 | 75.29 |  |
| 32 | 7680 | 102.4 | 13.83 | 16.84 | 62.29 |  |

GPT2.
Table [3](https://arxiv.org/html/2505.19342v1#S4.T3) summarizes the perplexity (PPL) results of Astra on the next-token prediction task, i.e., Wikipedia and Wikitext-103, using GPT2-S and GPT2-M.
Overall, Astra achieves competitive performance under aggressive communication compression.
Note that perplexity is an exponential function of the language modeling loss, i.e., $\text{PPL}=\exp(\mathcal{L})$, and thus small differences in loss can result in amplified changes in PPL.
Specifically, on GPT2-M, the PPL on Wikitext-103 increases from 14.8 (loss = 2.70) to 16.84 (loss = 2.82), reflecting only a 4.4% increase in loss despite a 102.4$\times$ compression ratio.
Similarly, on Wikipedia, the loss increases marginally from 2.5 to 2.63 (PPL from 12.16 to 13.83), confirming that much of the accuracy is preserved under significant transmission savings.

We also evaluate Astra in the zero-shot setting by directly testing the model trained on Wikipedia on the Wikitext-103 validation set.
Here, we observe a larger performance drop compared to the original model. For example, GPT2-M’s zero-shot PPL rises from 43.22 to 62.29 with Astra at 32 groups.
This discrepancy is more significant than in the fine-tuned setting and suggests a limitation of Astra in zero-shot generalization.
This is due to the discretization introduced by vector quantization, which reduces the diversity of token representations and hinders generalization to out-of-distribution data.

### 4.3 Results on Inference Latency

We evaluate the inference latency of multi-device Transformer inference under various system configurations.
Specifically, we benchmark a 12-layer Transformer encoder with 768 hidden dimensions, which is equivalent to the feature extraction stage of ViT-Base and the input token digestion stage of GPT2-S.
We compare Astra with existing multi-device inference baselines and evaluate their latency across three key dimensions, varying bandwidth, device count, and input token length, in Figure [1](https://arxiv.org/html/2505.19342v1#S1.F1), [4](https://arxiv.org/html/2505.19342v1#S4.F4), and [5](https://arxiv.org/html/2505.19342v1#S4.F5). All results are averaged over 20 runs to ensure stability.

Figure: Figure 3: Latency breakdown of local computation and inter-device communication time. The red dashed line represents the single-device latency.
Refer to caption: x3.png

Varying Bandwidth.
Figure [1](https://arxiv.org/html/2505.19342v1#S1.F1) in §[1](https://arxiv.org/html/2505.19342v1#S1) presents the latency speedup of various multi-device inference methods over single-device inference, evaluated across inter-device bandwidths ranging from 10 Mbps to 500 Mbps.
We fix the number of devices to 4 and the input token length to 1024. Additional configurations with different device counts and sequence lengths are provided in Appendix [D](https://arxiv.org/html/2505.19342v1#A4).
Across the bandwidth spectrum, Astra consistently outperforms all baselines and maintains substantial speedup even under extremely limited bandwidth.
For instance, Astra achieves a speedup of $1.27-2.74\times$ at 20 Mbps, while all other baselines perform even worse than single-device inference.
Even at 10 Mbps, our method with 16 and 1 quantization groups still delivers $1.26-2.65\times$ speedup, demonstrating strong scalability to bandwidth bottlenecks.

**Table 4: Astra’s speedup over baseline methods on 4 devices with 1024 Tokens.**
| Bandwidth<br>(Mbps) | TP | SP | BP+AG, Nb=1 | BP+SP, Nb=1 |
| --- | --- | --- | --- | --- |
| 10 | 342.74 | 171.82 | 15.25 | 29.37 |
| 20 | 177.89 | 89.41 | 8.41 | 15.66 |
| 50 | 73.14 | 37.05 | 4.07 | 6.95 |
| 100 | 37.19 | 19.08 | 2.58 | 3.96 |
| 200 | 19.02 | 9.99 | 1.83 | 2.45 |
| 500 | 8.05 | 4.51 | 1.37 | 1.53 |

We also visualize the absolute latency breakdown in Figure [3](https://arxiv.org/html/2505.19342v1#S4.F3), decomposing each method’s latency into computation and communication components.
Specifically, we depict the latency breakdown for the two fastest baselines, BP+AG and BP+SP when $N_{b}=1$, as well as Astra with different groups.
We can see that the communication time dominates in total latency for BP+AG and BP+SP, accounting for as much as $58.55-93.47\%$ of total runtime at low-bandwidth settings below 100 Mbps.
In contrast, Astra effectively mitigates this communication bottleneck via aggressive compression, thereby significantly reducing total latency.

We further summarize the relative speedup of Astra over each baseline across different bandwidth in Table [4](https://arxiv.org/html/2505.19342v1#S4.T4).
Astra’s advantage becomes more significant under stricter bandwidth.
Specifically, among all baselines, Sequence Parallelism (SP) shares the same parallelization granularity as Astra but performs full-precision attention across all tokens.
Thus, the speedup of Astra over SP directly reflects the benefit of our proposed Mixed-Precision Attention mechanism, which contributes up to $171.82\times$ latency reduction under low-bandwidth settings.

Varying Device Counts.
Figure [4](https://arxiv.org/html/2505.19342v1#S4.F4) shows the latency speedup comparison as the number of devices increases from 2 to 8.
We fix the input token length to 1024 and illustrate two representative bandwidth settings, 20 Mbps and 200 Mbps (more results see Appendix [D](https://arxiv.org/html/2505.19342v1#A4)).
For both bandwidth, Astra consistently achieves higher speedup than all baselines.
As the number of devices increases, more computation can be parallelized, leading to greater latency reduction.
For example, under 20 Mbps, Astra with 1 group improves from $1.72\times$ speedup with 2 devices to $3.69\times$ with 8 devices.

Figure: (a) 20 Mbps
Refer to caption: x4.png

Varying Input Token Length.
Figure [5](https://arxiv.org/html/2505.19342v1#S4.F5) presents the latency speedup comparison as the input token length increases from 256 to 4096.
Similarly, we fix the number of devices to 4 and evaluate under 20 Mbps and 200 Mbps (more bandwidth see Appendix [D](https://arxiv.org/html/2505.19342v1#A4)).
Across all sequence lengths, Astra consistently outperforms existing methods and our superiority becomes more significant at longer input lengths.
In real-world applications, longer input token lengths tend to form a more substantial barrier to achieving low-latency inference.
At 512 tokens and 20 Mbps bandwidth, for instance, Astra achieves a latency speedup of $1.98\times$ compared to the fastest baseline BP-AG of $0.25\times$, highlighting the practical value of Astra in real deployment scenarios.

### 4.4 Ablation Study

**Table 5: The impact of noise magnitude $\lambda$ on classification accuracy. Gap = Train - Val.**
| $\lambda$ | Train | Val | Gap |
| --- | --- | --- | --- |
| 0.0 | 99.98 | 89.91 | 10.07 |
| 0.1 | 99.97 | 90.02 | 9.95 |
| 0.3 | 99.98 | 90.13 | 9.85 |
| 1.0 | 99.98 | 90.77 | 9.21 |

Varying Noise Magnitude $\mathbf{\lambda}$.
The noise magnitude $\lambda$ in Noise-Augmented Vector Quantization controls the scale of noise added to the quantized embeddings during training.
Table [5](https://arxiv.org/html/2505.19342v1#S4.T5) shows the effect of varying $\lambda\in\{0.0,0.1,0.3,1.0\}$ on both training and validation accuracy. All other hyperparameters are fixed (e.g., 16 groups, commitment loss weight $\beta=0.0005$).
As $\lambda$ increases, the gap between training and validation accuracy consistently decreases, indicating that injecting noise improves generalizability by preventing the model from overfitting to discrete embedding patterns.
Notably, when $\lambda=1.0$, the validation accuracy improves by 0.86% compared to $\lambda=0$, demonstrating the effectiveness of our proposed strategy over naive vector quantization.

Distributed Class Token VS Single Class Token.
Table [6](https://arxiv.org/html/2505.19342v1#S4.T6) reports the classification accuracy of Astra using either a single class token or distributed class tokens across devices.
The distributed strategy consistently outperforms the single-token baseline, with accuracy gains ranging from 0.37% to 7.13% across different group configurations and commitment loss weights.
This demonstrates that allowing class tokens to attend to all full-precision context tokens in a distributed manner significantly enhances their ability to aggregate global information.

**Table 6: Distributed Class Token VS Single Class Token under different group configurations and commitment loss weights $\beta$.**
| $\beta$ | #Groups=1 | #Groups=16 | #Groups=32 |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Single | Distributed | $\Delta$ Acc. | Single | Distributed | $\Delta$ Acc. | Single | Distributed | $\Delta$ Acc. |  |
| 0.0001 | 82.39 | 88.95 | 6.56 | 89.11 | 90.37 | 1.26 | 90.39 | 91.60 | 1.21 |
| 0.0002 | 81.48 | 88.60 | 7.12 | 89.02 | 90.38 | 1.36 | 90.84 | 91.21 | 0.37 |
| 0.0005 | 81.82 | 88.95 | 7.13 | 88.93 | 90.77 | 1.84 | 90.79 | 91.64 | 0.85 |

## 5 Conclusion

We present Astra, a communication-efficient framework for accelerating Transformer inference in multi-device settings.
By integrating sequence parallelism with a novel Mixed-Precision Attention mechanism, Astra significantly reduces inter-device communication while preserving accuracy through Noise-Augmented Vector Quantization and Distributed Class Tokens.
Extensive experiments across both vision and NLP tasks demonstrate that Astra delivers substantial end-to-end latency improvements over existing baselines, achieving up to 2.64$\times$ speedup over single-device inference and up to 15.25$\times$ over state-of-the-art multi-device methods, under constrained bandwidth conditions as low as 10 Mbps.
Our results highlight the potential of Astra for practical deployment of large Transformer models in real-world, bandwidth-limited environments.

## References

- (1)
Understanding wi-fi 4/5/6/6e/7 (802.11 n/ac/ax/be).
- (2)
Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Olatunji Ruwase, Shaden Smith, Minjia Zhang, Jeff Rasley, et al.
Deepspeed-inference: enabling efficient inference of transformer models at unprecedented scale.
In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–15. IEEE, 2022.
- (3)
Olivier Chapelle, Jason Weston, Léon Bottou, and Vladimir Vapnik.
Vicinal risk minimization.
Advances in neural information processing systems, 13, 2000.
- (4)
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.
Imagenet: A large-scale hierarchical image database.
In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
- (5)
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
Bert: Pre-training of deep bidirectional transformers for language understanding.
In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pages 4171–4186, 2019.
- (6)
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.
An image is worth 16x16 words: Transformers for image recognition at scale.
arXiv preprint arXiv:2010.11929, 2020.
- (7)
Jiangsu Du, Yuanxin Wei, Shengyuan Ye, Jiazhi Jiang, Xu Chen, Dan Huang, and Yutong Lu.
Co-designing transformer architectures for distributed inference with low communication.
IEEE Transactions on Parallel and Distributed Systems, 2024.
- (8)
Allen Gersho and Robert M Gray.
Vector quantization and signal compression, volume 159.
Springer Science & Business Media, 2012.
- (9)
Chenghao Hu and Baochun Li.
Distributed inference with deep learning models across heterogeneous edge devices.
In IEEE INFOCOM 2022-IEEE Conference on Computer Communications, pages 330–339. IEEE, 2022.
- (10)
Chenghao Hu and Baochun Li.
When the edge meets transformers: Distributed inference with transformer models.
In 2024 IEEE 44th International Conference on Distributed Computing Systems (ICDCS), pages 82–92. IEEE, 2024.
- (11)
Yang Hu, Connor Imes, Xuanang Zhao, Souvik Kundu, Peter A Beerel, Stephen P Crago, and John Paul Walters.
Pipeedge: Pipeline parallelism for large-scale model inference on heterogeneous edge devices.
In 2022 25th Euromicro Conference on Digital System Design (DSD), pages 298–307. IEEE, 2022.
- (12)
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al.
Gpipe: Efficient training of giant neural networks using pipeline parallelism.
Advances in neural information processing systems, 32, 2019.
- (13)
Diederik P Kingma and Jimmy Ba.
Adam: A method for stochastic optimization.
arXiv preprint arXiv:1412.6980, 2014.
- (14)
Alex Krizhevsky, Geoffrey Hinton, et al.
Learning multiple layers of features from tiny images.
2009.
- (15)
Woosuk Kwon, Sehoon Kim, Michael W Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gholami.
A fast post-training pruning framework for transformers.
Advances in Neural Information Processing Systems, 35:24101–24116, 2022.
- (16)
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut.
Albert: A lite bert for self-supervised learning of language representations.
International Conference on Learning Representations, 2020.
- (17)
Mu Li, David G Andersen, Alexander Smola, and Kai Yu.
Communication efficient distributed machine learning with the parameter server.
Advances in neural information processing systems, 27, 2014.
- (18)
Sihao Lin, Hongwei Xie, Bing Wang, Kaicheng Yu, Xiaojun Chang, Xiaodan Liang, and Gang Wang.
Knowledge distillation via the target-aware transformer.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10915–10924, 2022.
- (19)
Zhenhua Liu, Yunhe Wang, Kai Han, Wei Zhang, Siwei Ma, and Wen Gao.
Post-training quantization for vision transformer.
Advances in Neural Information Processing Systems, 34:28092–28103, 2021.
- (20)
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher.
Pointer sentinel mixture models, 2016.
- (21)
Paul Michel, Omer Levy, and Graham Neubig.
Are sixteen heads really better than one?
Advances in neural information processing systems, 32, 2019.
- (22)
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.
Language models are unsupervised multitask learners.
OpenAI blog, 1(8):9, 2019.
- (23)
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro.
Megatron-lm: Training multi-billion parameter language models using model parallelism.
arXiv preprint arXiv:1909.08053, 2019.
- (24)
Aaron Van Den Oord, Oriol Vinyals, et al.
Neural discrete representation learning.
Advances in neural information processing systems, 30, 2017.
- (25)
Dongchao Yang, Songxiang Liu, Rongjie Huang, Jinchuan Tian, Chao Weng, and Yuexian Zou.
Hifi-codec: Group-residual vector quantization for high fidelity audio codec.
2023.
- (26)
Shengyuan Ye, Jiangsu Du, Liekang Zeng, Wenzhong Ou, Xiaowen Chu, Yutong Lu, and Xu Chen.
Galaxy: A resource-efficient collaborative edge ai system for in-situ transformer inference.
In IEEE INFOCOM 2024-IEEE Conference on Computer Communications, pages 1001–1010. IEEE, 2024.
- (27)
Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat.
Q8bert: Quantized 8bit bert.
In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing-NeurIPS Edition (EMC2-NIPS), pages 36–39. IEEE, 2019.
- (28)
Liekang Zeng, Xu Chen, Zhi Zhou, Lei Yang, and Junshan Zhang.
Coedge: Cooperative dnn inference with adaptive workload partitioning over heterogeneous edge devices.
IEEE/ACM Transactions on Networking, 29(2):595–608, 2020.
- (29)
Shuai Zhang, Sheng Zhang, Zhuzhong Qian, Jie Wu, Yibo Jin, and Sanglu Lu.
Deepslicing: Collaborative and adaptive cnn inference with low latency.
IEEE Transactions on Parallel and Distributed Systems, 32(9):2175–2187, 2021.
- (30)
Zhuoran Zhao, Kamyar Mirzazad Barijough, and Andreas Gerstlauer.
Deepthings: Distributed adaptive deep learning inference on resource-constrained iot edge clusters.
IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems, 37(11):2348–2359, 2018.

## Appendix A Additional Related Work

Deploying Transformers on Edge Devices.
Substantial efforts have been devoted to enabling Transformer models on edge devices through model compression and architecture simplification.
For example, Michel et al. [21] proposed pruning attention heads to reduce computational cost, while Q8BERT [27] quantizes BERT weights from 32-bit to 8-bit to accommodate memory-limited environments.
Other approaches, such as parameter factorization in ALBERT [16] and knowledge distillation methods [18], aim to construct lightweight variants of Transformer architectures suitable for resource-constrained hardware.
These techniques focus on discovering a compact model that maintains acceptable task performance under tight latency or memory budgets.

In contrast, Astra targets distributed inference while preserving the original model architecture.
The compressed transformer models from the above techniques can also leverage Astra ’s distributed inference system for further acceleration, as long as they retain the core transformer architecture.
This makes Astra an orthogonal solution that offers further performance improvements without requiring re-design or re-training for new hardware targets.

Distributed Inference Systems.
Distributed inference has emerged as a practical strategy to accelerate computation.
Early works such as DeepThings [30] exploited the partial receptive fields of convolutional neural networks (CNNs) to parallelize inference by splitting intermediate feature maps across multiple devices.
The follow-up works, including CoEdge [28], DeepSlicing [29], and EdgeFlow [9], further incorporated network heterogeneity and device resource profiling to optimize system throughput.
However, these methods are designed specifically for CNN-based models and are not applicable to the self-attention structure in Transformers.

Recent work has begun to explore multi-device inference for Transformers by adapting techniques from distributed training.
PipeEdge [11] utilizes pipeline parallelism to improve throughput, but its efficiency relies on large batch sizes and does not benefit per-request latency.
Other systems, such as DeepSpeed [2] and Megatron-LM [23] apply tensor parallelism by splitting weight matrices across devices, which leads to frequent and expensive inter-device communication.
To reduce this cost, Voltage [10] introduces sequence parallelism by distributing input tokens across devices and minimizing the number of cross-device interactions per Transformer block.
Galaxy [26] further combines tensor and sequence parallelism, while DeTransformer [7] even modifies the Transformer block structure to enable more efficient distribution.

Despite their contributions, these methods still rely on high-bandwidth connections to achieve meaningful speedups.
In contrast, Astra significantly reduces the required bandwidth to only 10 Mbps, while still achieving $2.64\times$ end-to-end latency speedup, making it far more practical for real-world deployments in constrained edge environments.

## Appendix B Proof for Noise-Augmented Vector Quantization

###### Theorem 1 (Noise-Augmented Embeddings Improve Distributional Fidelity) .

Let $\hat{\mathbf{X}}$ denote the quantized embedding of $\mathbf{X}$, and let $\tilde{\mathbf{X}}=\hat{\mathbf{X}}+\lambda\xi$ with $\xi\sim\mathcal{N}(\boldsymbol{\mu},\boldsymbol{\Sigma})$ sampled from the quantization residuals. Then the 2-Wasserstein distance between the original embedding distribution $P_{\mathbf{X}}$ and the perturbed distribution $P_{\tilde{\mathbf{X}}}$ satisfies:

$$ $W_{2}^{2}(P_{\mathbf{X}},P_{\tilde{\mathbf{X}}})<W_{2}^{2}(P_{\mathbf{X}},P_{ \hat{\mathbf{X}}}),$ (5) $$

i.e., the noise-augmented distribution is statistically closer to the true distribution than the raw quantized embedding.

###### Proof of Theorem 1 .

Let

$$ $\displaystyle\mathbf{m}_{X}=\mathbb{E}[\mathbf{X}],\quad\mathbf{m}_{\hat{X}}= \mathbb{E}[\hat{\mathbf{X}}],\quad\mathbf{m}_{\tilde{X}}=\mathbb{E}[\tilde{ \mathbf{X}}],$ (6) $$

and let
$C_{X},\;C_{\hat{X}},\;C_{\tilde{X}}$
denote the corresponding covariance matrices.
By definition of the quantization error
$\boldsymbol{\varepsilon}=\mathbf{X}-\hat{\mathbf{X}}$
we have
$\mathbf{m}_{X}=\mathbf{m}_{\hat{X}}+\mu$
and
$C_{X}=C_{\hat{X}}+\Sigma.$
As the injected noise
$\xi\sim\mathcal{N}(\mu,\Sigma),$,
we know that

$$ $\displaystyle\mathbf{m}_{\tilde{X}}=\mathbf{m}_{\hat{X}}+\lambda\mu,\quad C_{ \tilde{X}}=C_{\hat{X}}+\lambda^{2}\Sigma.$ (7) $$

The Wasserstein distance between two Gaussian distributions can be computed by

$$ $\displaystyle W_{2}^{2}(P_{1},P_{2})$ $\displaystyle=\|m_{1}-m_{2}\|_{2}^{2}+\mathop{\mathrm{trace}}\left(C_{1}+C_{2} -2\left(C_{2}^{\frac{1}{2}}C_{1}C_{2}^{\frac{1}{2}}\right)^{\frac{1}{2}}\right)$ (8) $\displaystyle=\|m_{1}-m_{2}\|_{2}^{2}+d_{B}^{2}(C1,C2),$ $$

where $m$ and $C$ are mean and covariance of the distributions, $d_{B}$ is the Bures metric.

To prove $W_{2}^{2}(P_{\mathbf{X}},P_{\tilde{\mathbf{X}}})<W_{2}^{2}(P_{\mathbf{X}},P_{
\hat{\mathbf{X}}})$, we will first show the mean term of $\tilde{\mathbf{X}}$ is smaller, then the Bures term is smaller.

Step 1 (mean term is smaller).
Using the mean part of the Gaussian $W_{2}$ formula, we have

$$ $\displaystyle\|\mathbf{m}_{X}-\mathbf{m}_{\hat{X}}\|_{2}^{2}-\|\mathbf{m}_{X}- \mathbf{m}_{\tilde{X}}\|_{2}^{2}=\|\mu\|_{2}^{2}-(1-\lambda)^{2}\|\mu\|_{2}^{2 }=(2\lambda-\lambda^{2})\|\mu\|_{2}^{2}>0,$ (9) $$

because $0<\lambda\leq 1$. Then we prove that

$$ $\displaystyle\|\mathbf{m}_{X}-\mathbf{m}_{\tilde{X}}\|_{2}^{2}<\|\mathbf{m}_{X }-\mathbf{m}_{\hat{X}}\|_{2}^{2}$ (10) $$

Step 2 (Bures term is smaller).

For analytical clarity, we assume that the quantization
errors $\varepsilon$ are independent and identically distributed across
dimensions, i.e. $\varepsilon_{k}\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\mathcal{N}(0,
\sigma^{2})$.
Under this assumption, the global error covariance becomes
$\Sigma=\sigma^{2}I$, which commutes with $C_{\hat{X}}$, namely $C_{\hat{X}}\Sigma=\Sigma C_{\hat{X}}$.
Consequently, $C_{\hat{X}}$ and $\Sigma$ can be diagonalized by the same orthonormal eigenbasis $U$.
Then we have

$$ $\displaystyle U^{\!\top}C_{\hat{X}}U=\operatorname{diag}(\sigma_{\hat{X},i}^{2 }),\quad U^{\!\top}\Sigma U=\operatorname{diag}(\sigma_{i}^{2}),\quad\sigma_{ \hat{X},i},\sigma_{i}\geq 0,$ (11) $$

then we obtain

$$ $\displaystyle U^{\!\top}C_{X}U=\operatorname{diag}(\sigma_{X,i}^{2}),\;\sigma_ {X,i}^{2}=\sigma_{\hat{X},i}^{2}+\sigma_{i}^{2},\qquad U^{\!\top}C_{\tilde{X}} U=\operatorname{diag}(\sigma_{\tilde{X},i}^{2}),\;\sigma_{\tilde{X},i}^{2}= \sigma_{\hat{X},i}^{2}+\lambda^{2}\sigma_{i}^{2}.$ (12) $$

For diagonal matrices the Bures term in the $W_{2}$ expression reduces to a sum of squared differences of *standard deviations*:

$$ $\displaystyle d_{B}^{2}(C_{A},C_{B})=\sum_{i}\bigl{(}\sigma_{A,i}-\sigma_{B,i} \bigr{)}^{2}.$ (13) $$

Then we have

$$ $\displaystyle d_{B}^{2}(C_{X},C_{\hat{X}})-d_{B}^{2}(C_{X},C_{\tilde{X}})$ $\displaystyle=\sum_{i}\Bigl{(}(\sigma_{X,i}-\sigma_{\hat{X},i})^{2}-(\sigma_{X ,i}-\sigma_{\tilde{X},i})^{2}\Bigr{)}$ (14) $\displaystyle=\sum_{i}(\sigma_{\tilde{X},i}-\sigma_{\hat{X},i})\Bigl{(}2\sigma _{X,i}-(\sigma_{\tilde{X},i}+\sigma_{\hat{X},i})\Bigr{)}.$ $$

Because
$\sigma_{\hat{X},i}\leq\sigma_{\tilde{X},i}\leq\sigma_{X,i}$,
we have $\sigma_{\tilde{X},i}-\sigma_{\hat{X},i}>0$ and
$2\sigma_{X,i}-(\sigma_{\tilde{X},i}+\sigma_{\hat{X},i})>0$,
therefore Equation [14](https://arxiv.org/html/2505.19342v1#A2.E14) is positive. Thus, we prove that

$$ $\displaystyle d_{B}^{2}(C_{X},C_{\tilde{X}})<d_{B}^{2}(C_{X},C_{\hat{X}}).$ (15) $$

Summary. Since both the mean part and the Bures part are strictly smaller for
$(\mathbf{m}_{\tilde{X}},C_{\tilde{X}})$ than for $(\mathbf{m}_{\hat{X}},C_{\hat{X}})$, hence we have completed the proof for

$$ $\displaystyle W_{2}^{2}(P_{\mathbf{X}},P_{\tilde{\mathbf{X}}})<W_{2}^{2}(P_{ \mathbf{X}},P_{\hat{\mathbf{X}}}),$ (16) $$

which completes the proof. ∎

###### Theorem 1 (Noise-Augmented Embeddings Improve Distributional Fidelity) .

###### Proof of Theorem 1 .

## Appendix C Proof for Distributed Class Tokens

###### Theorem 2 (Variance Reduction via Distributed Class Tokens) .

Let $\mathbf{h}$ denote the class token embedding of a full-precision global attention computation. Let $\tilde{\mathbf{h}}_{\mathrm{single}}$ be the output of a single-device class token using Mixed-Precision Attention, and let $\tilde{\mathbf{h}}_{\mathrm{dist}}$ be the average of $N$ distributed class token outputs. Then:

$$ $\mathbb{E}\bigl{[}\lVert\tilde{\mathbf{h}}_{\mathrm{dist}}-\mathbf{h}\rVert_{2 }^{2}\bigr{]}=\frac{1}{N}\mathbb{E}\bigl{[}\lVert\tilde{\mathbf{h}}_{\mathrm{ single}}-\mathbf{h}\rVert_{2}^{2}\bigr{]},$ (17) $$

i.e., distributed class tokens reduce the expected attention output error by a factor of $1/N$.

###### Proof of Theorem 2 .

Setup.
Tokens are evenly partitioned: $\mathbf{X}=\bigcup_{i=1}^{N}\mathbf{X}^{(i)}$, $|\mathbf{X}^{(i)}|=T/N$.
Each device stores local keys $\mathbf{k}_{j}$ and values $\mathbf{v}_{j}$ in full precision for $j\in\mathbf{X}^{(i)}$, and transmits quantized versions $\tilde{\mathbf{k}}_{j}=\mathbf{k}_{j}+\delta\mathbf{k}_{j}$ and $\tilde{\mathbf{v}}_{j}=\mathbf{v}_{j}+\delta\mathbf{v}_{j}$ to other devices, where $\delta\mathbf{k}_{j}$ and $\delta\mathbf{v}_{j}$ are the error introduced by quantization.
For every non-local token, we assume

$$ $\mathbb{E}[\delta\mathbf{k}_{j}]=\mathbf{0},\quad\mathbb{E}[\delta\mathbf{v}_{ j}]=\mathbf{0},\quad\operatorname{Cov}(\delta\mathbf{k}_{j})=\sigma_{k}^{2} \mathbf{I},\;\operatorname{Cov}(\delta\mathbf{v}_{j})=\sigma_{v}^{2}\mathbf{I},$ (18) $$

and errors are mutually independent.

Full-Precision Attention.
For a query $\mathbf{q}$, the attention logits are $a_{j}=\mathbf{q}^{\top}\mathbf{k}_{j}/\sqrt{d}$, attention weights $\alpha_{j}=\operatorname{softmax}(a_{j})$, and the output $\mathbf{h}=\sum_{j=1}^{T}\alpha_{j}\,\mathbf{v}_{j}$.

Mixed-Precision Attention via First-Order Taylor Expansion.
For a non-local token, the logit is perturbed by the key noise,

$$ $\tilde{a}_{j}=\frac{\mathbf{q}^{\top}(\mathbf{k}_{j}+\delta\mathbf{k}_{j})}{ \sqrt{d}}=a_{j}+\underbrace{\frac{\mathbf{q}^{\top}\delta\mathbf{k}_{j}}{\sqrt {d}}}_{=:e^{(k)}_{j}}.$ (19) $$

Because $e^{(k)}_{j}$ is small, we could linearise the softmax function by first-order Taylor expansion.
Specifically, the softmax funcion is $\alpha_{j}=\exp(a_{j})/\!\sum_{j}\exp(a_{j})$, and its Jacobian is $\partial\alpha_{j}/\partial a_{k}=\alpha_{j}(\delta_{jk}-\alpha_{k})$, where $\delta_{jk}$ is the Kronecker delta.
Therefore, we have the perturbed attention weights,

$$ $\tilde{\alpha}_{j}\approx\alpha_{j}+\sum_{k}\alpha_{j}(\delta_{jk}-\alpha_{k}) \,e^{(k)}_{k}=\alpha_{j}\;+\;\alpha_{j}\!\Bigl{(}e^{(k)}_{j}-\sum_{k}\alpha_{k }e^{(k)}_{k}\Bigr{)}=:\alpha_{j}+\delta\alpha_{j}.$ (20) $$

The terms $\delta\alpha_{j}$ remain zero-mean and mutually independent.

Then the mixed-precision output is

$$ $\tilde{\mathbf{h}}=\sum_{j}\tilde{\alpha}_{j}\,(\mathbf{v}_{j}+\delta\mathbf{v }_{j}),$ (21) $$

where $\delta\mathbf{v}_{j}=\mathbf{0}$ if $j$ is the local token index.

Attention Output Error.
Subtracting $\mathbf{h}$ and discarding higher-order noise products, the first-order output error is

$$ $\boldsymbol{\delta}:=\tilde{\mathbf{h}}-\mathbf{h}=\sum_{\text{non-local}\ j} \bigl{(}\alpha_{j}\,\delta\mathbf{v}_{j}+\delta\alpha_{j}\,\mathbf{v}_{j}\bigr {)}\;=\;\underbrace{\sum_{j}\alpha_{j}\,\delta\mathbf{v}_{j}}_{\text{V-error}: \ \mathbf{e}^{(v)}}+\underbrace{\sum_{j}\alpha_{j}\!\bigl{(}e^{(k)}_{j}-\!\sum _{k}\alpha_{k}\,e^{(k)}_{k}\bigr{)}\mathbf{v}_{j}}_{\text{K-propagated error}: \ \mathbf{e}^{(k)}}.$ (22) $$

Both error components are zero-mean random vectors and each coordinate has variance bounded by $C_{1}\sigma_{v}^{2}+C_{2}\sigma_{k}^{2}$, where $C_{1}$ and $C_{2}$ are deterministic constants.

Specifically, for the first value-error component $\mathbf{e}^{(v)}=\sum_{j\in\mathcal{N}}\alpha_{j}\,\delta\mathbf{v}_{j},$ where $\mathcal{N}$ is the set of $m=\tfrac{N-1}{N}T$ non-local tokens per device and each $\delta\mathbf{v}_{j}$ is an independent, zero-mean random vector with $\operatorname{Cov}(\delta\mathbf{v}_{j})=\sigma_{v}^{2}\,\mathbf{I}$.
Then for an arbitrary coordinate $c\in\{1,\dots,d\}$, since the noises are independent, we have

$$ $\operatorname{Var}\!\bigl{(}[\mathbf{e}^{(v)}]_{c}\bigr{)}=\operatorname{Var} \Bigl{(}\,\sum_{j\in\mathcal{N}}\alpha_{j}\,\bigl{[}\delta\mathbf{v}_{j}\bigr{ ]}_{c}\Bigr{)}=\sum_{j\in\mathcal{N}}\alpha_{j}^{2}\,\operatorname{Var}\bigl{( }[\delta\mathbf{v}_{j}]_{c}\bigr{)}=\sigma_{v}^{2}\sum_{j\in\mathcal{N}}\alpha _{j}^{2}.$ (23) $$

Since the attention weights $0\leq\alpha_{j}\leq 1$ and there are exactly $m$ non-local tokens, we have

$$ $\operatorname{Var}\!\bigl{(}[\mathbf{e}^{(v)}]_{c}\bigr{)}\;\leq\;\sigma_{v}^{ 2}\,m\,\max_{j\in\mathcal{N}}\alpha_{j}^{2}\;=\;C_{1}\,\sigma_{v}^{2},\qquad \text{with }\ C_{1}:=m\,\max_{j\in\mathcal{N}}\alpha_{j}^{2}.$ (24) $$

The constant $C_{1}$ is deterministic since it depends only on the current softmax weights.
And every coordinate of the value-error term is bounded in variance by $C_{1}\sigma_{v}^{2}$.

For the second key-propagated term, recall the first-order perturbation of each softmax weight

$$ $\delta\alpha_{j}\;=\;\alpha_{j}\Bigl{(}e^{(k)}_{j}-\sum_{k}\alpha_{k}\,e^{(k)} _{k}\Bigr{)},\quad e^{(k)}_{j}:=\tfrac{\mathbf{q}^{\!\top}\delta\mathbf{k}_{j} }{\sqrt{d}},$ (25) $$

where the key-noise scalars $e^{(k)}_{k}$ are i.i.d., zero-mean with variance $\sigma_{k}^{2}$.
Then for one output coordinate $c$, we need the variance of
$\bigl{[}\mathbf{e}^{(k)}\bigr{]}_{c}=\sum_{j\in\mathcal{N}}\delta\alpha_{j}\,v
_{j,c}$.

First, since $\delta\alpha_{j}$ is a linear combination of independent noises,

$$ $\operatorname{Var}[\delta\alpha_{j}]=\alpha_{j}^{2}\,\sigma_{k}^{2}\Bigl{(}1+ \textstyle\sum_{k}\alpha_{k}^{2}\Bigr{)}\;\leq\;2\,\alpha_{j}^{2}\sigma_{k}^{2},$ (26) $$

because $\sum_{k}\alpha_{k}^{2}\leq 1$.

Then since each addend $\delta\alpha_{j}\,v_{j,c}$ is zero-mean, we have

$$ $\operatorname{Var}\bigl{(}[\mathbf{e}^{(k)}]_{c}\bigr{)}=\sum_{j\in\mathcal{N} }v_{j,c}^{\,2}\,\operatorname{Var}[\delta\alpha_{j}]\;\leq\;2\,\sigma_{k}^{2} \bigl{(}\max_{j\in\mathcal{N}}v_{j,c}^{\,2}\bigr{)}\sum_{j\in\mathcal{N}} \alpha_{j}^{2}.$ (27) $$

With $\sum_{j\in\mathcal{N}}\alpha_{j}^{2}\leq m\max_{j}\alpha_{j}^{2}$,

$$ $\operatorname{Var}\bigl{(}[\mathbf{e}^{(k)}]_{c}\bigr{)}\;\leq\;2\,\sigma_{k}^ {2}\,m\,\bigl{(}\max_{j\in\mathcal{N}}v_{j,c}^{\,2}\bigr{)}\bigl{(}\max_{j\in \mathcal{N}}\alpha_{j}^{2}\bigr{)}\;=\;C_{2}\,\sigma_{k}^{2},$ (28) $$

where

$$ $C_{2}:=2\,m\,\max_{j\in\mathcal{N}}\bigl{(}\alpha_{j}^{2}v_{j,c}^{\,2}\bigr{)}.$ (29) $$

The constant $C_{2}$ is deterministic since it depends only on the current softmax weights and the fixed value vectors, not on the random noise. And every coordinate of the key-propagated error is bounded in variance by $C_{2}\sigma_{k}^{2}$.

Single Class Token Attention Output Error.
Its attention output error vector has $m=\tfrac{N-1}{N}T$ independent coordinates (i.e., the number of non-local tokens), each with per-coordinate variance $\sigma^{2}:=\operatorname{Var}\!\bigl{(}[\mathbf{e}^{(v)}]_{c}\bigr{)}+
\operatorname{Var}\!\bigl{(}[\mathbf{e}^{(k)}]_{c}\bigr{)}$. Hence

$$ $\mathbb{E}\!\bigl{[}\lVert\boldsymbol{\delta}_{\mathrm{single}}\rVert_{2}^{2} \bigr{]}=m\,d\,\sigma^{2}.$ (30) $$

Distributed Class Tokens Attention Output Error.
Let $\boldsymbol{\delta}^{(i)}$ be the error vector for the $i$-th device. Vectors $\boldsymbol{\delta}^{(i)}$ are independent and identically distributed. The averaged class token output is $\bar{\boldsymbol{\delta}}=\tfrac{1}{N}\sum_{i=1}^{N}\boldsymbol{\delta}^{(i)}$, then we have

$$ $\mathbb{E}\!\bigl{[}\lVert\bar{\boldsymbol{\delta}}\rVert_{2}^{2}\bigr{]}= \frac{1}{N^{2}}\sum_{i=1}^{N}\mathbb{E}\!\bigl{[}\lVert\boldsymbol{\delta}^{(i )}\rVert_{2}^{2}\bigr{]}=\frac{1}{N}\,m\,d\,\sigma^{2}.$ (31) $$

Attention Error Comparison.
According to Equations [30](https://arxiv.org/html/2505.19342v1#A3.E30) and [31](https://arxiv.org/html/2505.19342v1#A3.E31), using distributed class tokens yields a factor of $1/N$ in the expected attention output error compared to the single class token, indicating a more accurate attention computation under the mixed-precision attention mechanism.

∎

###### Theorem 2 (Variance Reduction via Distributed Class Tokens) .

###### Proof of Theorem 2 .

## Appendix D Results on Inference Latency Cont.

In the main paper, we evaluate the inference latency of Astra under three key factors that impact multi-device performance:
(1) Inter-device bandwidth, which affects the cost of communication across devices;
(2) Number of devices, which determines the degree of parallelism; and
(3) Input token length, which scales both computation and communication demands.

While the main text focuses on a representative configuration for each factor, this section presents additional results across a broader range of settings to further validate our findings.
We include more granular evaluations varying each of the three dimensions to provide a comprehensive view of Astra’s latency behavior under diverse deployment scenarios.

Varying Bandwidth.
To assess the scalability of Astra under different communication constraints, we evaluate its end-to-end latency speedup across a wide range of inter-device bandwidths from 10 to 500 Mbps.
Figure [6](https://arxiv.org/html/2505.19342v1#A4.F6) presents the results across varying numbers of devices i.e., 2, 4, 6, 8) with a fixed input token length of 1024, while Figure [7](https://arxiv.org/html/2505.19342v1#A4.F7) shows the results under varying input lengths (i.e., 256, 512, 1024, 2048, 4096) with 4 devices.

Across all configurations, Astra consistently outperforms existing multi-device baselines, with its advantage becoming increasingly significant as bandwidth decreases.
For instance, in the 4-device and 1024-token setting (Figure [6](https://arxiv.org/html/2505.19342v1#A4.F6)(b)), Astra achieves a speedup of $2.64\times$ at 10 Mbps, while the strongest baseline (BP-SP, $N_{b}=1$) only reaches $0.17\times$.
This trend also holds for different token lengths. As shown in Figure [7](https://arxiv.org/html/2505.19342v1#A4.F7), the relative speedup of Astra grows with increasing sequence length.
These results confirm that Astra maintains high efficiency even under constrained bandwidth, thanks to its aggressive communication reduction strategy.
In addition, grouped quantization (e.g., $G=16,32$) offers a tunable balance between compression and accuracy, enabling consistent latency benefits across a wide range of system settings.

Figure: (a) 2 Devices
Refer to caption: x8.png

Figure: (a) 256 Input Tokens
Refer to caption: x11.png

Varying Device Counts.
We further examine how the number of participating devices affects end-to-end latency speedup under varying bandwidth conditions.
Figure [8](https://arxiv.org/html/2505.19342v1#A4.F8) presents results for device counts ranging from 2 to 8 across bandwidths from 10 to 500 Mbps, with the input length fixed at 1024 tokens.

Across all bandwidth settings, Astra exhibits steadily increasing speedup as more devices are involved, demonstrating its ability to effectively utilize parallel computation.
Thanks to its communication-efficient design, where non-local token embeddings are transmitted in compact vector-quantized form, Astra consistently delivers strong gains even under constrained bandwidth, highlighting its scalability across a range of deployment scenarios.

Figure: (a) 10 Mpbs
Refer to caption: x15.png

Varying Input Token Length.
We evaluate how input token length affects latency speedup under different bandwidth conditions.
Figure [9](https://arxiv.org/html/2505.19342v1#A4.F9) shows results for input lengths from 256 to 4096 tokens across bandwidths from 10 to 500 Mbps using 4 devices.
Across all bandwidth settings, Astra consistently achieves higher speedup than baseline methods and our improvements enhance as the input length increases.
This trend highlights the communication bottleneck in existing methods, which becomes more significant with longer sequences, while Astra effectively mitigates this overhead through aggressive compression.

Figure: (a) 10 Mpbs
Refer to caption: x21.png

## Appendix E More Ablation Study

Varying the Commitment Loss Weights $\mathbf{\beta}$.
Recall from §[3.2](https://arxiv.org/html/2505.19342v1#S3.SS2) that we include a commitment loss term to encourage the original token embeddings to remain close to their assigned codebook entries following VQVAE [24].
While the original VQVAE typically sets the commitment loss weight to 0.25, we adopt much smaller values in our setting, as we apply vector quantization separately at each Transformer block.
Table [6](https://arxiv.org/html/2505.19342v1#S4.T6) in the main paper reports the results of Astra under different commitment weights $\beta\in\{0.0001,0.0002,0.0005\}$.
Here we further compare with two control variants: one without commitment loss (i.e., $\beta=0$), and one using an excessively large weight (i.e., $\beta=0.25$).
As shown in Table [7](https://arxiv.org/html/2505.19342v1#A5.T7), either omitting or misconfiguring the commitment term slightly degrades accuracy, with performance drops ranging from 0.1% to 1.67%, confirming the importance of tuning $\beta$ appropriately.

Table: Table 7: The impact of commitment loss weight $\beta$.

## Appendix F Impact Statement and Limitation

Potential Societal Impact.
Our work aims to make large Transformer models more deployable in real-world environments by enabling efficient multi-device inference under limited bandwidth. This can expand the accessibility of powerful AI models to edge and consumer-grade devices, potentially benefiting applications in healthcare, accessibility, and low-connectivity regions.
However, multi-device deployment may introduce new robustness and maintenance challenges. Unlike single-device inference, distributed inference requires reliable synchronization and communication among devices. Failures in individual devices or unstable connections can lead to degraded performance or unpredictable outputs. These issues can make such systems harder to debug, monitor, and guarantee correctness, especially in safety-critical applications.

Limitation and Future Work.
While Astra achieves strong performance across vision and language tasks, we observe a degradation in zero-shot generalization in the GPT experiments (see §[4.2](https://arxiv.org/html/2505.19342v1#S4.SS2)). We hypothesize this is due to the limited expressiveness of the discrete embedding space introduced by vector quantization. Future work may explore hybrid compression strategies that retain generalization ability while still reducing communication costs.
Additionally, our grouped vector quantization design requires maintaining a separate codebook for each group, which increases the overall storage footprint and may limit deployment flexibility across heterogeneous environments. As a future direction, we aim to investigate codebook sharing mechanisms or dynamically composable codebooks to reduce storage costs and enable bandwidth-aware adaptation without retraining.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 4
questions: - It would be helpful to further motivate the wireless and edge deployment motivation for this framework. Why are we doing inference requests on wifi?
soundness: 3
strengths: - The paper is well-presented and easy to follow
- Communication overhead is significant in large Transformer model distributed settings
- The use of codebooks is interesting
summary: The paper introduces a way to improve communication efficiency for multi device Transformer inference using sequence parallelism and mixed-precision. They show significant speedups on ViT and GPT2 scale models.
weaknesses: - The models used for inference are small and it is not clear to me that these hold at scale.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please address the implicitly listed questions in the weakness.

Can you provide a decision tree or heuristic for practitioners to select:
1. Number of groups G based on task type (vision vs. NLP) and target accuracy?
2. Commitment loss weight ε based on model architecture and dataset?
3. When to use distributed vs. single class tokens?
soundness: 3
strengths: 1. Addresses a real bottleneck: The paper identifies and tackles a genuine problem, that communication dominates latency (58.6-93.5%) in bandwidth-constrained multi-device inference.
2. Novel compression approach: The Mixed-Precision Attention mechanism is creative, using full-precision for local tokens and VQ for remote tokens.
3. Good evaluation: Extensive experiments across multiple architectures (ViT, GPT-2), tasks (classification, language modeling), and conditions (bandwidth, device count, heterogeneity) demonstrate broad applicability.
4. Practical compatibility: The framework integrates with existing quantization methods (8-bit, 4-bit), showing additional speedups of 1.35 to 2.73x when combined.
summary: This paper introduces ASTRA, a framework for accelerating transformer inference across multiple devices in bandwidth-constrained environments. The key part is a Mixed-Precision Attention mechanism that computes local attention with full precision while using vector-quantized embeddings for non-local tokens, reducing communication overhead. To preserve accuracy under aggressive compression, the authors propose Noise-Augmented Vector Quantization (NAVQ) and Distributed Class Tokens. Experiments on ViT and GPT-2 models across vision and NLP tasks demonstrate speedups of up to 2.64x over single-device inference and 15.25x over existing multi-device methods at bandwidths as low as 10 Mbps, while maintaining accuracy within 3.58% of the original models.
weaknesses: 1. Limited architectural types: The evaluation focuses only on ViT and GPT-2, which are relatively small and dated models. Modern applications use much larger models (e.g., LLaMA variants). The scalability claims are weakened without evidence on contemporary, production-scale models.

2. Severe zero-shot degradation: Table 3 shows large performance drops in zero-shot settings (e.g., GPT-2M perplexity increases from 43.22 to 62.29, a 44% degradation). This is a critical limitation for practical deployment where generalization is essential. 

3. Baselines potentially unfair: The comparison with BP, TP, and SP assumes these methods use full float32 precision. However, these methods could also be combined with standard compression techniques (gradient compression, activation compression). A more fair comparison would evaluate "BP+8bit quantization" vs "ASTRA+8bit quantization" to isolate ASTRA's contribution. Additionally, recent methods like FlexGen or DistServe are not compared, making it unclear how ASTRA compares to the current sota.

4. Communication model oversimplified: The paper assumes fixed bandwidth and doesn't account for real-world network variability, packet loss, or latency jitter. The latency model appears to assume perfect overlap of computation and communication, which is rarely achievable. Dynamic bandwidth fluctuations common in WiFi environments (cited as the target deployment) could significantly impact the practical speedups. The authors should evaluate under realistic network conditions with variable bandwidth and packet loss, or at minimum discuss how ASTRA degrades under non-ideal conditions.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: see weakness
soundness: 3
strengths: This paper observes a real bottleneck in multi-device Transformer inference for low-bandwidth or edge environments, which is increasingly relevant for real-time AI applications.
summary: This paper introduces ASTRA, a communication-efficient framework for multi-device Transformer inference under bandwidth-constrained settings. Existing multi-device methods suffer from high inter-device communication overhead, which dominates latency when bandwidth is limited. ASTRA addresses this by combining sequence parallelism with a Mixed-Precision Attention mechanism: local attention is computed at full precision, while remote tokens are compressed via low-bit vector quantization.
weaknesses: 1. ASTRA integrates known techniques (sequence parallelism + token quantization + noise augmentation), so the main contribution is in practical integration and bandwidth optimization, not a fundamentally new inference algorithm.
2. Lacks formal characterization of attention approximation error due to vector quantization and noise injection. Consider adding error bounds or theoretical analysis of how quantization and noise affect attention computation and model accuracy.
3. Experiments assume stable bandwidth; network variability or heterogeneous device scenarios are not tested.
4. Latency and energy claims are based on simulation; real-world multi-device deployment is not evaluated. Incorporate hardware-level profiling (multi-GPU, FPGA, or edge devices) to substantiate speedup and efficiency claims.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: see above.
soundness: 2
strengths: The paper is pretty interesting wtih 2 proposed optimizations to speedup inference across multiple device. I am not a systems expert, but the math behind the optimizations is correct and makes sense.
summary: This paper addresses the significant challenge of high inference latency in Transformer models, particularly in multi-device scenarios where inter-device communication becomes a dominant bottleneck in low-bandwidth environments. The authors propose ASTRA, a communication-efficient inference framework that builds on sequence parallelism but introduces a novel Mixed-Precision Attention mechanism to drastically reduce communication overhead. This mechanism computes attention using full-precision embeddings for local tokens while using low-bit vector-quantized (VQ) representations for non-local tokens transmitted between devices. To preserve model accuracy under such aggressive compression, ASTRA introduces two key optimizations: Noise-Augmented Quantization (NAVQ), a training-time regularization strategy that injects noise into quantized embeddings to improve generalization , and Distributed Class Tokens, which replicates the class token to each device and aggregates the outputs to mitigate information bias. Experiments on ViT and GPT-2 models show that ASTRA achieves substantial end-to-end speedups -- up to 2.64x over single-device inference and 15.25x over other multi-device methods -- in bandwidth-constrained settings (as low as 10 Mbps), while incurring only minor accuracy degradation.
weaknesses: The biggest weakness and drawback from the paper is the lack of large-scale experiments. Espeically the model choices in this paper is so small that none of these optimizations matter. Given the prevelance of open models of much larger sizes even ViTs and Qwen series models, it is empirically needed to make sure the proposed optimizations are quality neutral while providing the benefits. Without these results it hard to make a case for accepting the paper. I hope the authors can scale up the benchmarking further.


Decision: Reject