Learning Task-Oriented Communication for Edge Inference: An Information Bottleneck Approach (2102.04170v3)

Abstract: This paper investigates task-oriented communication for edge inference, where a low-end edge device transmits the extracted feature vector of a local data sample to a powerful edge server for processing. It is critical to encode the data into an informative and compact representation for low-latency inference given the limited bandwidth. We propose a learning-based communication scheme that jointly optimizes feature extraction, source coding, and channel coding in a task-oriented manner, i.e., targeting the downstream inference task rather than data reconstruction. Specifically, we leverage an information bottleneck (IB) framework to formalize a rate-distortion tradeoff between the informativeness of the encoded feature and the inference performance. As the IB optimization is computationally prohibitive for the high-dimensional data, we adopt a variational approximation, namely the variational information bottleneck (VIB), to build a tractable upper bound. To reduce the communication overhead, we leverage a sparsity-inducing distribution as the variational prior for the VIB framework to sparsify the encoded feature vector. Furthermore, considering dynamic channel conditions in practical communication systems, we propose a variable-length feature encoding scheme based on dynamic neural networks to adaptively adjust the activated dimensions of the encoded feature to different channel conditions. Extensive experiments evidence that the proposed task-oriented communication system achieves a better rate-distortion tradeoff than baseline methods and significantly reduces the feature transmission latency in dynamic channel conditions.

Learning Task-Oriented Communication for Edge Inference

The paper "Learning Task-Oriented Communication for Edge Inference: An Information Bottleneck Approach" addresses the efficient transmission of feature vectors from edge devices to edge servers in wireless communication networks, focusing on the challenges posed by bandwidth limitations and dynamic channel conditions. The authors propose a novel learning-driven communication framework using the Information Bottleneck (IB) principle to prioritize task relevance over data reconstruction in transmission schemes.

Key Contributions

1. Task-Oriented Communication Model: Traditional communication systems optimize for data reconstruction between transmitters and receivers, typically employing joint source-channel coding (JSCC) techniques. The proposed model diverges from this by shifting the focus towards preserving information pertinent to downstream inference tasks. The authors achieve this through an IB framework that precisely quantifies the tradeoff between the informativeness of encoded features and communication overhead.
2. Variational Information Bottleneck (VIB) Framework: Direct computation of mutual information under the IB framework is intractable for high-dimensional neural networks. To address this, the paper introduces VIB, leveraging variational approximations to create a tractable upper bound on the IB objective. This variational approach enables end-to-end learning of the feature encoding process, adapting both the feature extraction and transmission to optimize inference tasks.
3. Sparsity-Inducing Encoding: To reduce communication overhead, a sparsity-inducing distribution is adopted as the variational prior. This distribution facilitates the pruning of redundant dimensions in the feature vector—a critical aspect for efficient bandwidth usage. The method, Variational Feature Encoding (VFE), juxtaposes task-informed compression against raw data preservation.
4. Adaptive Feature Encoding for Dynamic Channel Conditions: Recognizing the variability inherent in wireless channels, the paper introduces Variable-Length Variational Feature Encoding (VL-VFE). VL-VFE dynamically adjusts the length of transmitted features via a model incorporating dynamic neural networks, thus ensuring robust inference performance across fluctuating channel states without reconstructing the transmitted data.

Experimental Validation

Through extensive experiments on image classification tasks (e.g., MNIST, CIFAR-10), the authors demonstrate the superiority of VFE and VL-VFE over classical data-reconstruction-based methods like Deep JSCC in terms of achieving a more favorable rate-distortion tradeoff. Notably, VFE achieves significant reductions in redundant feature transmission without sacrificing classification accuracy. Similarly, VL-VFE adapts efficiently to dynamic PSNR conditions, maintaining low latency and high accuracy.

Implications and Future Directions

The proposed framework's success underscores the viability of task-oriented communication systems for edge computing applications, where communication overhead posed by data reconstruction becomes a bottleneck. This work lays the groundwork for further research into adaptive systems compatible with complex environments, such as integration with multiple devices or evolving IoT landscapes.

Future research may delve into optimizing the computational processes further, enhancing robustness against even more adverse conditions, or expanding the model's applicability to diverse AI-driven applications beyond image classification. Additionally, theoretical investigations could yield insights into the underpinning rate-distortion metrics for predictive inference tasks, facilitating the development of universally adaptable encoding protocols.

By leveraging an interdisciplinary approach combining machine learning with information theory principles, this paper opens new avenues for efficient communication system design in next-generation wireless networks.