PaperlayerView exact revision on arXiv

Source-linked AI summary

Importance-Aware Feature Sparsification for Wireless Split Learning

Bumjun Kim, Yoon Huh, Wan Choi

arXiv:2609.39194v1cs.LGeess.SP

TL;DR

Wireless split learning reduces on-device computation, but repeated transmission of intermediate features remains a bottleneck, and existing selection methods can add client processing without accounting for task relevance. ICS uses server-computed Grad-CAM scores and class-balanced importance vectors reused by clients for sparsification. Experiments report consistent gains over baselines, with a more pronounced advantage under heterogeneous, label-skewed data.

  • Problem

    Wireless split learning repeatedly transmits high-dimensional intermediate features, while existing client-side selection is generally task-agnostic and can be biased under non-i.i.d. data.

    Source: §I · pp. 1–2

  • Method

    ICS uses server-computed Grad-CAM scores aggregated into a class-balanced, label-agnostic importance vector that clients reuse to sparsify features.

    Source: §I · p. 2

  • Results

    ICS consistently outperforms baseline sparsification methods, with the performance gap becoming more pronounced as client data become more heterogeneous.

    Source: §VI · pp. 10–11, §VII · p. 14

  • Takeaways & Limitations

    The reported advantages are especially evident under severe label skew, where class-balanced importance aggregation reduces the influence of locally dominant labels.

    Source: §VI · p. 11

Abstract

from arXiv · show

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.

I. INTRODUCTION

Split learning reduces clients’ computation but leaves repeated transmission of intermediate features as a communication bottleneck, and existing sparsification can add client work while overlooking task relevance and label skew. The introduction motivates ICS, which uses server-computed, class-balanced feature importance to address these concerns and extends to parallel split learning.

1) Distributed Learning Frameworks:

  • FL requires full on-device training, while split learning moves heavier model layers to the server but still transmits intermediate features at every iteration.Sequential client participation also causes waiting time, and parallel participation does not eliminate per-client feature transmissions.Source: §I · pp. 1–2

2) Parallel Split Learning Methods:

  • Prior PSL work addresses waiting time, resource allocation, splitting, or transmission control, while parallel uploads require greater aggregate uplink resources.The introduction also describes client clustering with sequential inter-cluster aggregation as an approach to reducing latency.Source: §I · p. 1, 2) Parallel Split Learning Methods · p. 3

I. INTRODUCTION

  • ICS uses server-computed Grad-CAM importance and a class-balanced, label-agnostic vector that clients reuse to select task-relevant features without extra client computation.The introduction also presents a convergence analysis of the sparsification error and the interaction between mini-batch size and sparsification ratio under a fixed communication budget.Source: §I · p. 2

3) Communication-Efficient Split Learning Methods:

  • Existing communication-efficient SL methods use clustering, feature statistics, random selection, or magnitude-based selection, generally computed on clients without explicit task-specific importance.The introduction identifies client preprocessing burden and possible loss of task-relevant information under non-i.i.d. data as limitations; ICS instead computes Grad-CAM-based importance at the server.Source: 3) Communication-Efficient Split… · p. 3

B. Contributions

The paper presents ICS as a server-driven, Grad-CAM-based feature sparsification method, with class-balanced importance aggregation and extensions to theoretical analysis and parallel split learning. The system model describes client-server training through intermediate-feature transmission, server computation, returned gradients, and forwarding of updated parameters and importance.

B. Contributions

  • ICS uses server-computed Grad-CAM importance vectors from previous iterations as surrogates for client-side feature compression, avoiding additional client forward–backward passes.Class-balanced aggregation produces a label-agnostic vector that mitigates class bias under non-i.i.d. data and supports sparsification when ground-truth labels are unavailable.Source: B. Contributions · p. 3
  • The analysis compares communication and computational overhead and derives a non-asymptotic convergence bound for sparsification error and the batch-size–sparsification-ratio trade-off under a fixed uplink budget.The overhead analysis states that ICS reduces transmitted intermediate-feature overhead without sending mask indices and shifts importance-estimation cost from clients to the server.Source: B. Contributions · p. 3

C. Notations

  • Table III is identified as the reference for the paper’s key notations and their descriptions.The supplied system-model passages describe mini-batch client updates, feature transmission to the server, returned activation gradients, and forwarding of updated parameters and importance for later use.Source: Table III · p. 4, §II · pp. 4–5

B. Contributions

  • For PSL, the server aggregates Grad-CAM-based scores across participating clients to construct a shared class-balanced importance vector.The paper also reports validation across diverse tasks and datasets, with consistent accuracy gains and faster convergence over baselines, especially under non-i.i.d. data.Source: B. Contributions · p. 4

III. IMPORTANCE-AWARE CLASS- BALANCED FEATURE SPARSIFICATION

The section outlines conventional Grad-CAM as a post-hoc attribution method for identifying image regions relevant to a class prediction. It derives channel importance from target-class gradients, combines those weights with feature maps, and applies ReLU to retain positive contributions.

A. Conventional Grad-CAM for Model Interpretability

  • Conventional Grad-CAM is applied post-hoc to a fixed, pre-trained network and does not account for ongoing learning iterations.It backpropagates a target-class score into an intermediate convolutional layer to assess feature contributions.Source: A. Conventional Grad-CAM for Model… · p. 5
  • Gradients of the target-class score with respect to feature maps indicate their sensitivity, and aggregated gradients yield an importance weight for each channel.The channel weights are used in a weighted combination of feature maps to form a class-discriminative localization map.Source: A. Conventional Grad-CAM for Model… · p. 5
  • Applying ReLU produces a saliency map that highlights image regions positively contributing to the prediction for the target class.The paper later uses these maps to validate the proposed framework.Source: A. Conventional Grad-CAM for Model… · p. 5

B. Importance-Aware Class-Balanced Sparsification

ICS uses server-computed Grad-CAM importance scores to select intermediate feature channels, then balances those scores across classes to reduce bias from non-i.i.d. client data. Clients reuse the prior round’s importance vector, while the server updates it from true-class gradients.

B. Importance-Aware Class-Balanced Sparsification

  • ICS ranks channels using Grad-CAM importance derived from the true-class logit, targeting features relevant to correct prediction rather than magnitude alone.Magnitude can prioritize features unrelated to the true label.Source: B. Importance-Aware Class-Balanced… · pp. 5–6
  • Clients apply the previous round’s importance vector to retain the top-N feature channels, without computing current-round Grad-CAM before transmission.The server can compute importance only after receiving the client’s intermediate features.Source: B. Importance-Aware Class-Balanced… · p. 6
  • The server forms per-class importance vectors by averaging sample scores for each label, then maintains an exponential moving average for each client and class.The EMA uses a momentum coefficient β and is initialized to zero.Source: B. Importance-Aware Class-Balanced… · p. 7
  • A uniform average of class-specific EMA memories creates a class-balanced, label-agnostic importance vector for reuse in the next training step.Equal weighting is applied to classes that have contributed at least once.Source: B. Importance-Aware Class-Balanced… · p. 7
  • ICS extends beyond CNNs to architectures with a sparsifiable feature dimension, including transformers by treating embedding dimensions as channels and pooling over tokens.The importance scores use gradients of the true-class logit with respect to intermediate features, aggregated over non-feature axes.Source: B. Importance-Aware Class-Balanced… · p. 7

A. Communication Overhead Analysis

ICS reduces uplink traffic by transmitting selected feature channels without extra mask indices, and shifts importance estimation from clients to the server. Its client-side selection cost is O(C), with server-side importance estimation hidden behind the model update.

A. Communication Overhead Analysis

  • Client-side selection methods require extra uplink index information, with overhead depending on whether selection is over channels, feature positions, or clustered subvectors.The stated overheads are O(BRCUV + BRC log C) for channel selection, O(BRCUV + RCUV log(CUV)) for SplitFC, and O(Qdsub + Bnsub log Q) for FedLite.Source: A. Communication Overhead Analysis · p. 7
  • O(BRCUV) is the per-client uplink overhead for ICS, with no additional mask-index transmission at the same sparsification ratio.The shared server-generated importance vector lets the server infer the deterministic top-N selection.Source: A. Communication Overhead Analysis · p. 7

B. Computational Complexity Analysis

  • O(C) is ICS’s additional client-side selection cost, while the server’s importance estimation can run in parallel with its update and has zero effective additional runtime.The analysis also measures per-batch client- and server-side method-specific runtime, excluding common forward and backward computation.Source: B. Computational Complexity Analysis · p. 8

C. Convergence Analysis

The analysis derives a non-asymptotic convergence bound under stated assumptions and identifies a trade-off between stochastic gradient variance and sparsification error. Under a fixed communication budget, increasing batch size requires a lower sparsification ratio, so larger batches are not always better.

C. Convergence Analysis

  • Theorem 1 gives a non-asymptotic convergence rate under the stated assumptions and sparsification ratio.The bound includes a sparsification-induced error term.Source: C. Convergence Analysis · p. 8
  • Under a fixed communication budget, increasing batch size reduces stochastic gradient variance but forces a lower sparsification ratio and raises sparsification error, so larger batches are not always better.The bound also contains a data-heterogeneity term that does not vanish with T.Source: C. Convergence Analysis · p. 8
  • The convergence analysis is limited to mini-batch SGD; adaptive optimizers such as Adam require a different treatment of their moment-based dynamics.The sparsification mechanism remains compatible with adaptive optimizers, and the stated variance–sparsification trade-off also arises with Adam.Source: C. Convergence Analysis · p. 8

V. EXTENSION TO PARALLEL SPLIT LEARNING

The PSL extension replaces sequential client-wise importance forwarding with a shared iteration-wise vector computed from participating clients’ mini-batches. The server broadcasts that vector for use in sparsifying features in the next round.

V. EXTENSION TO PARALLEL SPLIT LEARNING

  • In PSL, the server aggregates Grad-CAM-based importance scores across participating clients into a shared iteration-wise vector, which it broadcasts for next-round sparsification.Unlike sequential SL, the vector is not passed along in client order.Source: §V · p. 9
  • Each client uses the shared vector from the previous round to sparsify its intermediate features and sends the sparsified features and labels to the server.The server processes the participating clients’ features and updates its server-side model; client-side updates occur in parallel.Source: §V · p. 9
  • The server computes class-specific importance by averaging samples of each class across participating clients, then averages class memories over observed classes to form the shared vector.The server maintains per-class EMA memories for temporal stability.Source: §V · p. 9

VI. NUMERICAL RESULTS

Across classification and segmentation experiments, ICS outperforms the compared sparsification methods, with its advantage more pronounced under severe non-i.i.d. data. Tests also examine class-balanced aggregation, the batch-size/sparsification trade-off, and extensions to parallel and transformer-based split learning.

Comparison of Test Accuracy under i.i.d. Datasets

  • ICS consistently outperforms all baselines on CIFAR-100 test accuracy under i.i.d. data.It also improves test accuracy faster in early training.Source: Figure 2 · p. 10, §VI · p. 10

Comparison of Test Accuracy under Non-i.i.d. Datasets

  • Under non-i.i.d. partitions, ICS achieves higher final accuracy and faster convergence, with a more pronounced advantage as label skew increases.The reported gap is most evident at Dirichlet α = 0.05; all methods generally degrade as α decreases.Source: §VI · pp. 10–11, Figure 3 · p. 11

Effects of Class-Balanced Importance Aggregation

  • ICS outperforms ICS w/o CB and trains more stably at α = 0.1 and α = 0.05.The comparison attributes the difference to class-balanced aggregation mitigating mini-batch label imbalance.Source: Figure 6 · p. 11, §VI · p. 11

Effects of Mini-batch Size and Sparsification Ratio

  • Under a fixed communication budget, larger mini-batches improve early accuracy growth, while higher sparsification ratios narrow the per-epoch accuracy gap later.The study varies B ∈ {128, 32, 16} and R ∈ {0.1, 0.4, 0.8}; the reported trade-off reflects sample variety versus per-sample feature information.Source: §VI · pp. 11–12, Figure 7 · p. 12

Evaluation under PSL

  • In PSL with Dirichlet α = 0.5 and mini-batches of 32, ICS achieves the best performance among the compared methods.The authors suggest its shared importance vector helps maintain consistency across parallel client updates.Source: Figure 8 · p. 12, §VI · p. 12

Generalization to Transformer-based SL Models

  • On CIFAR-100 with transformer-based SL, ICS trains stably while RS fails to train reliably under the same communication budget.Sparsification is applied along the token embedding dimension.Source: Figure 9 · p. 12, §VI · p. 12

Generalization to a Different Dataset

  • On CIFAR-10 classification, ICS consistently outperforms the alternative sparsification schemes.The experiment uses a mini-batch size of 256 and runs for 50 iterations.Source: Figure 10 · p. 13, §VI · p. 13

Generalization to a Different Task

  • On Oxford-IIIT Pet segmentation, ICS achieves the highest mIoU and pixel accuracy, while TS and RTS degrade relative to other methods.The discussion links magnitude-based selection's weaker performance to background-dominant channels displacing object-dominant features.Source: Figure 11 · p. 13, Figure 12 · p. 13, §VI · p. 13

VII. CONCLUSION

The conclusion presents ICS as server-generated, class-balanced feature sparsification for wireless split learning and reports consistent gains over existing communication-efficient baselines across tasks, data heterogeneity levels, and architectures.

VII. CONCLUSION

  • ICS uses true-class-logit Grad-CAM scores aggregated into a class-balanced, label-agnostic vector that clients reuse next round to transmit discriminative channels without extra client-side computation.The conclusion states that the class-balanced design mitigates bias under non-i.i.d. data.Source: §VII · p. 14
  • ICS extends to parallel split learning through a shared server-generated importance vector and to transformer models along the embedding dimension.Source: §VII · p. 14
  • Across classification and semantic segmentation, heterogeneity levels, and architectures, ICS consistently outperforms existing communication-efficient split-learning baselines while keeping client-side overhead minimal.Source: §VII · p. 14

APPENDIX A PROOF OF THEOREM 1

The proof develops bounds for model updates with sparsified stochastic gradients, accounting for client/server update error and averaging the resulting inequality over iterations.

APPENDIX A PROOF OF THEOREM 1

  • The proof represents the model update using concatenated stochastic gradients with sparsification from the client and server.Source: Appendix A · p. 14
  • The derivation upper-bounds terms (a) and (b) in (26), using S-smoothness, Assumption 3 for an equality, Assumption 4 for an inequality, and Lemma 1 for term (b).Source: Appendix A · pp. 14–15
  • The proof averages the resulting inequality over iterations from 0 to T − 1 using ηt = 1 KE T.Source: Appendix A · p. 16

APPENDIX B PROOF OF LEMMA 1

The lemma proof bounds the update error caused by sparsification by relating features before and after sparsification through the mean value theorem and stated spectral-norm bounds.

APPENDIX B PROOF OF LEMMA 1

  • The server-side sparsification error is expressed using the difference between intermediate features before and after sparsification.The proof then invokes the mean value theorem.Source: Appendix B · p. 16
  • The client-side update error is likewise expressed in terms of sparsified and unsparsified features and logits.The stated inequalities use the mean value theorem and Assumptions 1 and 5.Source: Appendix B · p. 16
  • Combining the server- and client-side bounds yields an error bound with Λ2 = Λ2 0(Λ2 1 + Λ2 2Λ2 3).Source: Appendix B · p. 16

Related papers

Loading 2609.39194v1…