Wireless split learning improves communication with importance-aware feature selection
Importance-Aware Feature Sparsification for Wireless Split Learning
Machine Learning
Summary
Wireless split learning helps devices do less work by sending some tasks to a server, but sending the right data back and forth can be slow and costly. The authors introduce a new way to pick the most important parts of the data using signals from the server, which helps keep the learning accurate and balanced even when data is very different across devices. This method saves time on the device side and improves overall learning speed and quality. Their approach also works with different types of machine learning models and setups.
What this means in practice
- •For mobile app developers: Reduce communication costs and speed up training for apps using split learning on devices with limited resources and diverse user data.
- •For edge computing engineers: Improve split learning deployments by applying importance-aware feature selection to balance accuracy and communication load across heterogeneous edge devices.
Authors
Bumjun Kim, Yoon Huh, Wan Choi
Abstract
Wireless split learning (SL) reduces on-device computation by offloading upper layers to a server, yet transmitting high-dimensional intermediate features at each iteration remains a major communication bottleneck. Existing methods select features at the client side using task-agnostic criteria such as magnitude, statistics, or clustering, which increases client-side processing and often degrades accuracy under non-independent and identically distributed (non-i.i.d.) client data. We propose importance-aware class-balanced sparsification (ICS), a lightweight approach in which the server ranks feature channels using Grad-CAM-based scores obtained from the true-class logit during backpropagation. The per-class scores are aggregated into a class-balanced, label-agnostic importance vector that mitigates head-class bias under label skew, and each client reuses this vector in the next round to retain the top-$N$ feature channels, incurring no additional client-side forward or backward passes. We further derive a non-asymptotic convergence bound that isolates the sparsification-induced error and characterizes how the sparsification ratio and mini-batch size jointly affect convergence under a fixed communication budget, and we analyze the communication and computational overhead of ICS against representative baselines. Beyond sequential CNN-based SL, we extend ICS to parallel split learning and to transformer-based models. Experiments show that ICS consistently outperforms the baselines, with larger gains under severe non-i.i.d. partitions.