Source-linked AI summary
Learning Task-Oriented Communication for Edge Inference: An Information Bottleneck Approach
Jiawei Shao, Yuyi Mao, Jun Zhang
TL;DR
Limited bandwidth and device computation make low-latency edge inference difficult when communication is designed for data reconstruction. The paper jointly designs task-oriented feature communication using IB and VIB, adds sparsity and variable-length encoding, and reports better rate-distortion tradeoffs with lower latency under dynamic channels.
Problem
Edge inference requires informative, compact feature transmission because device computation and bandwidth are limited, while data-oriented reconstruction can preserve task-irrelevant information.
Method
The paper uses an IB rate-distortion objective, a tractable VIB approximation with a sparsity-inducing prior for VFE, and dynamic neural networks for variable-length VL-VFE.
Results
VFE and VL-VFE outperform traditional communication design and learning-based JSCC baselines across static and dynamic image-classification channels, with VL-VFE achieving higher accuracy and lower latency than DeepJSCC.
Takeaways & Limitations
Task-oriented communication can discard task-irrelevant information while adapting transmitted feature dimensions to channel conditions for edge inference.
Abstract
from arXiv · showhide
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.
I. INTRODUCTION
Edge inference must balance limited device resources and transmission latency, motivating communication schemes optimized for downstream tasks rather than data reconstruction. The paper develops IB-based, learning-driven feature communication with adaptive transmission for changing channels.
- Edge-only inference incurs transmission latency from uploading large data, while device-only inference is constrained by limited on-device computation.
- Device-edge co-inference reduces overhead by extracting compact features on devices, but pretrained-network splitting typically leaves compression and transmission to separate communication modules.
- Data-oriented communication can waste bandwidth by preserving task-irrelevant information for reconstructing raw inputs instead of directly supporting inference.
- The paper formalizes task-oriented communication with an information bottleneck rate-distortion tradeoff between task-relevant information and representation compactness.
- Variational feature encoding uses a tractable VIB approximation and sparsity-inducing prior to prune redundant feature dimensions, while VL-VFE adapts active dimensions to channel conditions.
- Experiments in static and dynamic channels show VFE and VL-VFE outperform traditional communication design and learning-based JSCC baselines on image classification tasks.
D. Notations
The system models device-edge co-inference as noisy feature transmission between device and server DNNs. Its IB objective treats nonzero feature dimensions as communication overhead and balances inference uncertainty against preserved input information.
- The device and server cooperate on tasks such as image classification, with X and Y denoting input data and target variables, and Z, Ẑ, and Ŷ denoting encoded, received, and inferred variables.
- The device jointly trains a feature extractor and JSCC encoder, transmitting task-relevant features over a noisy channel to a server network.
- The channel model adds Gaussian noise to transmitted features and constrains each feature dimension's power, while the framework can extend to other channels with an estimable transfer function.
- The communication overhead is measured by the number of nonzero JSCC-encoder dimensions, creating a tradeoff between inference accuracy and transmission latency.
- The IB objective minimizes inference uncertainty through -I(Ẑ,Y) while penalizing preserved input information with βI(Ẑ,X), thereby expressing an edge-inference rate-distortion tradeoff.
C. Main Challenges
Applying the information bottleneck to practical feature encoding faces challenges in mutual-information estimation, sparsity control, and adaptation to dynamic channels. The VIB framework provides a tractable approximation, while noise-corrupted features and variational distributions support optimization.
- High-dimensional mutual-information estimation is challenging because empirical probability estimation may require samples increasing exponentially with dimension.
- Reducing input-feature redundancy does not directly ensure sparse feature activations, so nuisance information must be aggregated into expendable dimensions.
- Dynamic channels require feature lengths to adapt, but changing neuron activation is nontrivial because neural-network structures are fixed after initialization.
- The encoded feature is modeled as a noise-corrupted representation generated by the on-device network and channel model.
- VIB approximates intractable IB distributions with variational distributions and yields an upper-bound objective that can be optimized using reparameterization, Monte Carlo sampling, and stochastic gradient descent.
- Minimizing the VIB objective is intended to prune task-irrelevant dimensions, making it suitable for task-oriented communication.
B. Redundancy Reduction and Feature Sparsification
The method addresses redundancy reduction and feature sparsification by pairing the VIB KL term with a sparsity-inducing variational prior. The log-uniform prior encourages task-irrelevant information to concentrate in dimensions that can be pruned.
- Minimizing the KL term reduces redundancy but does not guarantee sparse activations when relevant information remains distributed across dimensions.
- A log-uniform variational prior aggregates task-irrelevant information into expendable dimensions and induces sparsity in the encoded feature.
- The mean-field approximation factorizes the prior across feature dimensions to alleviate computational complexity.
- Because the KL divergence lacks a closed form, the method uses an approximation to evaluate the sparsity-promoting term.
- The approximate KL term is minimized as α_i approaches infinity, driving z_i toward zero so dimensions identically equal to zero can be pruned.
C. Variational Pruning on Dimension Importance
VFE identifies removable feature dimensions through a learned importance vector rather than checking every network parameter. Thresholding the importance scale prunes dimensions whose outputs carry no valid information beyond channel noise.
- The dimension importance vector γ provides an efficient way to determine which output neurons can be pruned.
- Each augmented neuron weight row is represented by a normalized direction multiplied by its importance scale γ_i.
- When γ_i is near zero, z_i approaches zero and its received distribution degrades to channel noise without valid information.
- VFE prunes dimension i when γ_i is below threshold γ_0, using the training procedure's final threshold test.
- The encoder uses a Tanh activation, constraining the peak transmitted power P to 1.
IV. VARIABLE-LENGTH VARIATIONAL FEATURE ENCODING
VL-VFE adapts the transmitted feature length to channel conditions through a dynamic neural network whose dimension importance depends on channel noise variance. Consecutive activation avoids transmitting activation indexes while supporting channel-dependent feature lengths.
- VL-VFE adjusts the output feature length under different channel conditions using a dynamic neural network.
- More dimensions are activated during adverse channel conditions, while selective activation balances inference performance against communication overhead.
- The encoder receives channel noise variance σ^2 as an input, which represents channel quality because the Tanh output amplitude is constrained to 1.
- The importance value γ_i(σ^2) varies with channel noise variance, allowing thresholding to select active dimensions for each channel condition.
- The training algorithm samples channel noise for each datapoint, deactivates dimensions below γ_0, and updates encoder and server parameters through backpropagation.
- VL-VFE activates dimensions consecutively from the first dimension, avoiding extra communication for transmitting activation indexes.
C. Training Procedure for the Dynamic Neural Network
VL-VFE trains a dynamic neural network by averaging losses across sampled channel conditions and selectively activating feature dimensions.
- Each training iteration samples σ^2 from the possible PSNR range to represent different channel conditions.
- VL-VFE temporarily deactivates dimensions whose γ_i(σ^2) values fall below γ0 instead of permanently pruning them during unstable training.
V. PERFORMANCE EVALUATION
The evaluation uses image-classification benchmarks and compares the proposed schemes with learning-based communication baselines under a task-oriented objective.
- Experiments evaluate task-oriented communication schemes on MNIST and CIFAR-10 image-classification tasks.The study also reports additional Tiny ImageNet experiments in Appendix D.
- Baselines: The baselines are DeepJSCC and learning-based Quantization for device-edge co-inference.
- Baselines: DeepJSCC maps input data directly to channel symbols, with communication cost proportional to the feature encoder's output dimension.
- Baselines: Learning-based Quantization uses low-precision feature representations and requires an additional channel-coding step for error correction.
- The experiments use a scalar Gaussian channel capacity upper bound to characterize the communication setting.
3) Metrics:
Performance is measured through a classification-accuracy versus communication-latency tradeoff under bandwidth-limited edge-inference settings and fixed neural-network architectures.
- Metrics: Classification accuracy denotes inference performance, while communication latency serves as the rate indicator.
- Neural Network Architecture: All methods use the same network architecture for fair comparisons despite device computation and memory constraints.
- Neural Network Architecture: MNIST uses one fully connected layer on a sub-0.5 MB microcontroller and an MLP server network.
- Neural Network Architecture: CIFAR-10 uses convolutional layers on a single-board computer and an n-dimensional encoded feature produced by a final fully connected layer.
- Results for Static Channel Conditions: Static-channel curves vary β to adjust feature length, with results reported for MNIST at PSNR values of 10 dB and 20 dB and likewise for CIFAR-10.
- Results for Static Channel Conditions: Under t≤3.25 ms, accuracy is evaluated across PSNR values, while theoretical image-transmission latency ranges from 8 to 20 ms for MNIST and 70 to 180 ms for CIFAR-10.
C. Results for Dynamic Channel Conditions
Under dynamic channel conditions, VL-VFE adapts feature length to channel noise and achieves lower latency and higher accuracy than DeepJSCC, while the prior choice affects sparsity and accuracy.
- Dynamic Channel Conditions: VL-VFE is evaluated as PSNR varies from 10 to 25 dB, with channel noise variance σ^2 varying from 3 × 10^-3 to 0.1.
- Dynamic Channel Conditions: VL-VFE achieves higher accuracy and lower latency than DeepJSCC across the dynamic-channel evaluation.
- Dynamic Channel Conditions: VL-VFE adaptively changes the activated feature dimension according to instantaneous channel noise, reducing latency in the high-PSNR regime.
- Robustness: The study also tests robustness to inaccurate estimates of the channel noise variance σ^2.
- Ablation Study: A Gaussian variational prior can shrink informative dimensions and degrade inference accuracy, whereas sparsity induction is evaluated with a log-uniform prior.
- Conclusion: The paper frames IB as a candidate for task-oriented communication while identifying multiple-device systems, theory, and robustness as future directions.
APPENDIX A DERIVATION OF THE VARIATIONAL UPPER BOUND
The VIB objective is a variational upper bound on the IB objective, while the experiments assess VL-VFE under inaccurate channel-noise variance and dynamic conditions.
- Variational upper bound: The VIB objective provides a tractable variational upper bound on the IB objective because KL-divergence is nonnegative and the entropy of Y is constant.This connects the computationally tractable VIB formulation to the original IB optimization.
- Noise-variance mapping: VL-VFE constrains its noise-variance mapping to be nonnegative and increasing, using nonnegative MLP weights to guarantee these properties.The construction sets each MLP weight matrix to the element-wise absolute value of its actual parameters.
- Dynamic-channel results: The proposed method achieves lower communication latency than DeepJSCC in all three tested dynamic-channel cases.The experiments compare error rate and feature-transmission latency across channel PSNR values.
- Inaccurate channel knowledge: With 8 pilot symbols, inaccurate noise-variance estimation degrades performance, but the proposed method still achieves a much better rate-distortion tradeoff than DeepJSCC.The estimator becomes more accurate as the number of pilot symbols increases.
- Inaccurate channel knowledge: With no transmitter knowledge of the noise variance, VL-VFE still shows comparable performance to DeepJSCC.This case sets the feature-encoding PSNR to 10 dB rather than estimating the variance from pilots.
APPENDIX D ADDITIONAL EXPERIMENTS ON TINY IMANGENET
Additional Tiny ImageNet experiments evaluate VFE and VL-VFE with a ResNet18 backbone. VL-VFE achieves higher accuracy and lower latency than DeepJSCC in dynamic channel conditions.
- Experimental setup: Tiny ImageNet contains 200 classes, 100,000 training images, and 10,000 validation images, with all images sized 64 × 64.The selected ResNet18 backbone achieves around 50.5% top-1 accuracy on this task.
- Static-channel results: VFE outperforms the baselines by achieving a better rate-distortion tradeoff.This evaluates the proposed variational feature encoding method on the Tiny ImageNet classification task.
- Dynamic-channel results: VL-VFE achieves higher accuracy and lower latency than DeepJSCC in dynamic channel conditions.The dynamic-channel training uses β = 5 × 10^-4 while PSNR changes from 10 dB to 25 dB.