[Paper Review] Visual Transformer Meets CutMix for Improved Accuracy, Communication Efficiency, and Data Privacy in Split Learning

Sihun Baek, Jihong Park+4 more|arXiv (Cornell University)|Jul 1, 2022

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes CutMixSL, a novel split learning framework for Vision Transformers (ViT) that enhances accuracy, communication efficiency, and data privacy by introducing CutSmashed data—randomly masked and compressed hidden representations. By leveraging CutMix-style data augmentation during split learning, CutMixSL reduces upload payload by 20–50%, cuts privacy leakage by ~8×, and improves top-1 accuracy by up to 18.5% compared to baselines.

ABSTRACT

This article seeks for a distributed learning solution for the visual transformer (ViT) architectures. Compared to convolutional neural network (CNN) architectures, ViTs often have larger model sizes, and are computationally expensive, making federated learning (FL) ill-suited. Split learning (SL) can detour this problem by splitting a model and communicating the hidden representations at the split-layer, also known as smashed data. Notwithstanding, the smashed data of ViT are as large as and as similar as the input data, negating the communication efficiency of SL while violating data privacy. To resolve these issues, we propose a new form of CutSmashed data by randomly punching and compressing the original smashed data. Leveraging this, we develop a novel SL framework for ViT, coined CutMixSL, communicating CutSmashed data. CutMixSL not only reduces communication costs and privacy leakage, but also inherently involves the CutMix data augmentation, improving accuracy and scalability. Simulations corroborate that CutMixSL outperforms baselines such as parallelized SL and SplitFed that integrates FL with SL.

Motivation & Objective

• To address the high communication cost and privacy leakage in split learning for Vision Transformers (ViT), which suffer from large, unstructured smashed data similar to input images.
• To overcome the limitations of existing split learning and federated learning in handling large ViT models due to model size and communication overhead.
• To improve scalability and accuracy in distributed ViT training by integrating data augmentation into the split learning pipeline.
• To develop a privacy-preserving mechanism that reduces reconstruction-based privacy leakage without sharing raw data.

Proposed method

• Clients generate CutSmashed data by applying mutually exclusive random masks to the hidden representations (smashed data) at the split-layer, using a shared pseudo-random sequence generator.
• The server aggregates the masked CutSmashed data from multiple clients to form CutMix data, which is used for forward and backward propagation.
• A shared pseudo-random sequence ensures mask mutual exclusivity across clients and server without exposing raw data.
• The method integrates Patch CutMix data augmentation into the split learning process, enhancing model generalization and inductive bias.
• Clients upload only non-zero masked patches, reducing communication payload and limiting information leakage.
• The framework supports scalability by avoiding update imbalance issues seen in parallel split learning.

Experimental results

Research questions

• RQ1 Can CutSmashed data reduce communication overhead in ViT-based split learning while preserving model accuracy?
• RQ2 How does CutMixSL compare to parallel split learning and SplitFed in terms of privacy leakage and communication efficiency?
• RQ3 To what extent does integrating CutMix-style augmentation into split learning improve ViT performance?
• RQ4 Does the use of masked, shuffled data in CutSmashed representations reduce the risk of reconstruction attacks?
• RQ5 Can CutMixSL scale effectively with increasing numbers of clients without performance degradation?

Key findings

• CutMixSL reduces communication payload by 20–50% compared to baseline parallel split learning, depending on the masking ratio and group size.
• The framework reduces privacy leakage by approximately 8×, as measured by reconstruction mean-squared error (MSE), compared to standard smashed data.
• Top-1 accuracy improves by up to 18.5% compared to baseline split learning, with CutMixSFL achieving 80.97% accuracy at 10 clients.
• CutMixSL achieves superior scalability, with a 12.9% performance gain at 10 clients, outperforming both parallel SL and SplitFed.
• Shuffled CutMix masks amplify privacy protection, increasing reconstruction loss by 3× compared to standard CutMix.
• The method effectively mitigates the update imbalance problem in split learning, enabling consistent performance gains with more clients.