Back to Research

Open threads

Attention-Aware Semantic Communications for Collaborative Inference

A practical edge-server inference scheme where a small ViT decides when to call the server and which patches are worth sending.

ENstudied
semantic-communicationscollaborative-inferenceedge-aivision-transformer

This paper starts from a practical observation: split inference is less natural for ViTs than for CNNs. In CNN pipelines, an intermediate feature map can be much smaller than the raw input, so sending that feature to the server may save bandwidth. In ViTs, the dimensions across layers are more uniform, so simply cutting the model in the middle does not automatically reduce communication.

The paper avoids that trap with two models. A small ViT runs on the edge device and a stronger ViT remains on the server. The edge model decides whether a sample is uncertain enough to justify server inference; if it is, the same model also decides which patches are worth uploading. The attention matrix is therefore not just an internal model artifact:

A=softmax⁡(qk⊤dh)A = \operatorname{softmax}\left(\frac{qk^{\top}}{\sqrt{d_h}}\right)

For upload selection, class-token attention becomes the patch-importance signal:

a=softmax⁡(qclskp⊤dh)a = \operatorname{softmax}\left(\frac{q_{\mathrm{cls}}k_p^{\top}}{\sqrt{d_h}}\right)

The second control signal is uncertainty. The paper compares Shannon entropy,

gs(x)=−∑y∈Ypθ(y∣x)log⁡2pθ(y∣x)g_s(x) = -\sum_{y\in\mathcal{Y}} p_{\theta}(y|x)\log_2 p_{\theta}(y|x)

with min-entropy,

gm(x)=−log⁡max⁡y∈Ypθ(y∣x)g_m(x) = -\log\max_{y\in\mathcal{Y}}p_{\theta}(y|x)

and uses the uncertainty score as a gate:

g(x)≥η⟹send selected patches to the server.g(x)\ge\eta \Longrightarrow \text{send selected patches to the server.}

The design is attractive because it reuses signals the edge model already computes. Easy samples stay local. Ambiguous samples ask the server for help. If the edge model is weak as a final classifier, it can still be useful as a semantic filter that identifies which patches matter.

The main risk is calibration. If the small model is confidently wrong, the system will skip the server at exactly the wrong moment. Patch selection is also task-dependent: classification may tolerate missing background patches, while detection, segmentation, video tracking, or multimodal reasoning may not. I would combine this paper with IAQ so that selected patches are not all sent at the same quality. High-attention patches should get more bits, marginal patches fewer bits, and both the attention threshold and entropy threshold should adapt to battery level, network load, latency target, and the cost of an incorrect decision.