AutoEncoder-Compressed Parallel Split Learning for Pre-trained Model Fine-Tuning

2026-07-20

AI Summary

This work addresses the challenges of insufficient computational capacity and high communication overhead in distributed fine-tuning of large foundation models on resource-constrained edge devices. The authors propose the AE-PSL framework, which seamlessly integrates a lightweight, learnable autoencoder into the partitioned layers of off-the-shelf pre-trained models without requiring joint training. This design enables efficient compression of intermediate activations and gradients while bridging the gap between client data distributions and the pre-trained feature manifold through a two-stage feature alignment mechanism. The approach significantly reduces communication costs and maintains or even improves fine-tuning performance, outperforming task-agnostic compression strategies such as sparsification and quantization.

Abstract

Distributed Fine-Tuning (DFT) of large-scale Foundation Models (FMs) on resource-constrained edge devices is limited by local compute constraints and communication overhead. Parallel Split Learning (PSL) reduces client-side computation by keeping few model layers on each client and offloading the remaining computation to the server; however, clients must exchange intermediate activations and gradients with the server at every training step. Existing SL communication-compression methods mainly rely on task-agnostic heuristics, such as sparsification and quantization. While learnable SL compressors can better adapt to intermediate representations, they require co-training with the target model. Therefore, directly inserting them into off-the-shelf FMs introduces feature-distribution misalignment and degrades DFT performance. To address this, we propose AE-PSL, a communication-efficient PSL framework that compresses intermediate activations and gradients using a lightweight AutoEncoder (AE) placed at the split layer. To ensure compatibility of AE compression with pre-trained FMs, AE-PSL introduces a novel two-stage alignment mechanism, which adapts the AE to the pre-trained model's feature manifold and client-specific feature distributions before DFT.

Problem

• Distributed Fine-Tuning
• Foundation Models
• Split Learning
• Communication Compression
• Feature Distribution Alignment

Innovation

• AutoEncoder
• Parallel Split Learning
• Communication Compression
• Feature Alignment
• Distributed Fine-Tuning