Attention-aware Semantic Communications for Collaborative Inference

TL;DR

This paper tackles the high communication cost of edge-server inference for vision tasks by exploiting a dual ViT setup: a lightweight ViT (DeiT-Tiny) on the edge as a semantic encoder and a heavier ViT (DeiT-Base) on the server. It introduces two complementary strategies—attention-aware patch selection and entropy-aware image transmission—to minimize data transfer while maintaining server-level accuracy. Experiments on ImageNet show about a 68% reduction in communication with only a small loss in classification performance (80.84% vs 81.8% for the server model), and the analysis demonstrates meaningful server-side computational savings and latency improvements. The approach highlights the potential of semantic-encoder-like edge models and patch-level offloading to enable efficient, scalable edge intelligence with ViT-based architectures.

Abstract

We propose a communication-efficient collaborative inference framework in the domain of edge inference, focusing on the efficient use of vision transformer (ViT) models. The partitioning strategy of conventional collaborative inference fails to reduce communication cost because of the inherent architecture of ViTs maintaining consistent layer dimensions across the entire transformer encoder. Therefore, instead of employing the partitioning strategy, our framework utilizes a lightweight ViT model on the edge device, with the server deploying a complicated ViT model. To enhance communication efficiency and achieve the classification accuracy of the server model, we propose two strategies: 1) attention-aware patch selection and 2) entropy-aware image transmission. Attention-aware patch selection leverages the attention scores generated by the edge device's transformer encoder to identify and select the image patches critical for classification. This strategy enables the edge device to transmit only the essential patches to the server, significantly improving communication efficiency. Entropy-aware image transmission uses min-entropy as a metric to accurately determine whether to depend on the lightweight model on the edge device or to request the inference from the server model. In our framework, the lightweight ViT model on the edge device acts as a semantic encoder, efficiently identifying and selecting the crucial image information required for the classification task. Our experiments demonstrate that the proposed collaborative inference framework can reduce communication overhead by 68% with only a minimal loss in accuracy compared to the server model on the ImageNet dataset.

Figures (16)

• Figure 1: Overview of the ViT model Dosovitskiy2021image.
• Figure 2: The proposed collaborative inference framework utilizing pre-trained ViT models: A lightweight model (e.g., DeiT-Tiny) on the edge device and a complicated model (DeiT-Base) on the server. The edge device (client) evaluates the uncertainty of its inference results by computing the entropy. If this entropy exceeds a predetermined threshold, the edge device selects the most important patches of the image based on the attention scores and transmits these to the server.
• Figure 3: The visualization of attention-aware patch selection. The left is an image of the ImageNet dataset that the client model (DeiT-Tiny) inaccurately classifies as 'Hammerhead Shark'. The right shows the selected patches by attention-aware patch selection. These selected patches allow the server model (DeiT-Base) to correctly classify the image as 'White Shark'.
• Figure 4: The visualization of images and attention scores. In the left column, we have the original images of the ImageNet dataset. The middle column displays the attention score maps generated by DeiT-Tiny. The right column shows the patches selected by the attention-sum threshold selection. The images are labeled as 'Brambling', 'Irish Setter', and 'Snow Leopard', respectively.
• Figure 5: The comparison of image patches selected by DeiT-Tiny (middle column) and DeiT-Base (right column). The same number of patches are selected according to the mean attention scores. The left column displays the original images of the ImageNet dataset. The images are labeled as 'Miniature Schnauzer', 'Great Gray Owl', and 'Magpie', respectively.