PRISM: Distributed Inference for Foundation Models at Edge

The gist The PRISM method proposes a communication-efficient and compute-aware strategy for distributed Transformer inference on edge devices by leveraging Segment Means representation to approximate intermediate output features, restructuring self-attention to eliminate redundant computations, and designing a partition-aware causal masking scheme.

How it works

1. PRISM introduces the use of Segment Means representations for compressing input and intermediate feature data, which reduces inter-device communication by up to 89.9% on the ViT model with P = 2 and CR = 9.9, enabling efficient Transformer inference in bandwidth-constrained edge environments with only a minor impact on accuracy, which can be mitigated through model finetuning The Segment Means computation procedure is outlined in Algorithm 2.

2. The method eliminates redundant Key and Value matrix computation across devices by limiting it to local partition data and their corresponding received Segment Means vectors.

3. To preserve the model output shape and accuracy despite the impact of compression on the Key and Value matrices, PRISM restructures the self-attention equation by introducing a scaling-aware softmax computation, which reduces overall computation by up to 50% and 68% for 2 and 3 devices, respectively.

4. For autoregressive models like GPT-2, a partition-aware causal mask generation strategy based on the self-attention permutation invariance is proposed to ensure accurate distributed inference.

Segment Means Representation

The Segment Means representation is used to compress intermediate features and reduce interdevice communication during distributed inference. The procedure involves dividing each partitioned input sequence into non-overlapping segments of approximately equal size, and then computing the column-wise mean for each segment to yield the compressed representation Zp. This yields a compact contextual information with minimal overhead. The expanded representation becomes Yp, which is constructed by concatenating local partition data Xp with the Segment Means received from other edge devices.

Self-Attention Optimization

To enable distributed inference across edge devices, PRISM adopts position-wise layer partitioning and extends it by addressing communication challenges through an optimized mechanism. The self-attention equation is restructured using the property of permutation invariance, allowing each device to operate on a reordered view of the input data without affecting the final result. Each device constructs an augmented matrix Xˆ p by concatenating its local partition data Xp with the Segment Means received from other edge devices. This selective augmentation reduces redundant computation and enables scalable, parallel self-attention across devices. The attention equation is then simplified by leveraging the repetitive structure of Segment Means to eliminate redundant operations through scaling factors in the softmax expression.

Partition-Aware Causal Masking

For autoregressive models such as GPT-2, a causal mask is required to prevent attention to future tokens. To avoid incorrect masking, each device constructs its causal mask using its global position within the original sequence of length N by communicating the global partition index p of each partition. The defined causal mask Mp ensures that only future tokens are masked appropriately on each device, preserving causal structure in a distributed and potentially out-of-order processing environment.

Experimental Evaluation

The PRISM method was evaluated on ViT, BERT, and GPT-2 models across diverse datasets including CIFAR-10, CIFAR100, ImageNet-1K, GLUE, and CBT. Results demonstrate substantial reductions in communication overhead (up to 99.2% for BERT at compression rate CR = 128) and per-device computation (51.24% for BERT at the same setting), with only minor accuracy degradation. PRISM consistently outperforms the Voltage [20] method across all bandwidths, even at 200 Mbps, reducing latency by 43.3% for P = 2 with CR = 9.9. For GPT-2, PRISM achieves a 66.7% reduction in per-device computation and 90% reduction in communication overhead with P = 3 and CR = 10. These results highlight the broad applicability of PRISM across different Transformer architectures and tasks.

The research concludes that PRISM significantly reduces communication overhead with minor or no impact on model accuracy, demonstrating its strong potential for real-time, edge-level, and multi-device deployment scenarios. This approach sets the foundation for future improvements to communication efficient, distributed Transformer inference.

--- Page 1 ---

JOURNAL OF LATEX CLASS FILES, VOL. XX, NO. X, XXX XXXX 1

PRISM: Distributed Inference for Foundation Models at Edge Muhammad Azlan Qazi Alexandros Iosifidis and Qi Zhang Abstract—Foundation models (FMs) have achieved remarkable success across a wide range of applications, from image classification to natural langurage processing, but pose significant challenges for deployment at edge. This has sparked growing interest in developing practical and efficient strategies for bringing foundation models to edge environments. In this work, we propose PRISM, a communication-efficient and computeaware strategy for distributed Transformer inference on edge devices. Our method leverages a Segment Means representation to approximate intermediate output features, drastically reducing inter-device communication. Additionally, we restructure the selfattention mechanism to eliminate redundant computations caused by per-device Key/Value calculation in position-wise partitioning and design a partition-aware causal masking scheme tailored for autoregressive models. We evaluate PRISM on ViT, BERT, and GPT-2 across diverse datasets, namely CIFAR-10, CIFAR100, ImageNet-1k, GLUE, and CBT. Our results demonstrate substantial reductions in communication overhead (up to 99.2% for BERT at compression rate CR = 128) and per-device computation (51.24% for BERT at the same setting), with only minor accuracy degradation. This method offers a scalable and practical solution for deploying foundation models in distributed resource-constrained environments Index Terms—Distributed Inference, Edge Inference, Foundation model, Transformer, Causal mask, Large Language Models, and AIoT. I INTRODUCTION T HE introduction of the Transformer architecture has revolutionized the field of artificial intelligence leading to groundbreaking advancements across natural language processing NLP computer vision and multimodal learning Foundation models FMs such as GPT [1] [2] LLaMA [3] [4] BERT [5] and ViT [6] DeepSeek’s R1 [7], and Veo-3 have significantly transformed how people interact with AI becoming integral to everyday applications from document summarization and dialogue systems to image and video generation However the extraordinary capabilities of these models come at the cost of enormous model size These FMs typically contain hundreds of billions of parameters which are essential for capturing complex patterns and long-range dependencies in the input data As a result they require high-performance accelerators e.g. GPUs with large memory capacity which are typically hosted on cloud servers This cloud-based deployment introduces several limitations suffering from large and time-varying communication latency imposing high communication network traffic and ongoing maintenance costs More importantly it raises critical privacy concerns particularly Muhammad Azlan Qazi and Qi Zhang are with the Department of Electrical and Computer Engineering Aarhus University Denmark (e-mail: maq@ece.au.dk; qz@ece.au.dk) Alexandros Iosifidis is with the Faculty of Information Technology and Communication Sciences Tampere University Finland (e-mail: alexandros.iosifidis@tuni.fi) Jobs Terminal Device Edge Devices Input Embedding Positional Embedding Partitioning Segment Mean (SM) Inter- Device Com. Terminal Device Self Attn. FFN Self Attn. FFN Self Attn. FFN Output SM Fig 1 Overview of the proposed model architecture Input tokens are embedded and partitioned on the terminal device with each partition summarized using Segment Means and distributed to edge devices Each edge device performs local self-attention and feedforward network computation computes its own Segment Means and exchanges it with other devices This process is repeated across Transformer blocks Finally output partitions are aggregated at the terminal device to produce the final result Fig 1 Overview of the proposed model architecture Input tokens are embedded and partitioned on the terminal device, with each partition summarized using Segment Means and distributed to edge devices, with each edge device performs local self-attention and feedforward network computation, computes its own Segment Means, and exchanges it with other devices This process is repeated across Transformer blocks. Finally output partitions are aggregated at the terminal device to produce the final result Fig 1 Overview of the proposed model architecture Input tokens are embedded and partitioned on the terminal device, with each partition summarized using Segment Means and distributed to edge devices, with each edge device performs local self-attention and feedforward network computation, computes its own Segment Means, and exchanges it with other devices This process is repeated across Transformer blocks. Finally output partitions are aggregated at the terminal device to produce the final result Fig 1 Overview of the proposed model architecture Input tokens are embedded and partitioned on the terminal device, with each partition summarized using Segment Means and distributed to edge devices, with each edge device performs local self-attention and feedforward network computation, computes its own Segment Means, and exchanges it with other devices This process is repeated across Transformer blocks. Finally output partitions are aggregated at the terminal device to produce the final result Fig 1 Overview of the proposed model architecture Input tokens are embedded and partitioned on the terminal device, with each partition summarized using Segment Means and distributed to edge devices, with each edge device performs local self-attention and feedforward network computation, computes its own Segment Means, and exchanges it with other devices This process is repeated across Transformer blocks.

1. Improved distributed inference for Foundation Models by leveraging Segment Means representation to approximate intermediate output features, which drastically reduc[es] inter-device communication. This allows for efficient Transformer inference on edge devices by reducing communication overhead up to 99.2% (for BERT at CR = 128).

2. Elimination of redundant computation across devices by restructuring the self-attention mechanism to eliminate redundant computations caused by per-device Key/Value calculation in position-wise partitioning. This optimization reduces overall computation by up to 50% and 68% for 2 and 3 devices, respectively.

3. Implementation of a partition-aware causal masking strategy tailored for autoregressive models, which ensures each device generates the appropriate causal mask for its partition, preserving communication efficiency while supporting correct generation in distributed settings.