AE-PSL is a research toolchain for communication-efficient parallel split learning when fine-tuning pre-trained foundation models on distributed and resource-constrained clients. The main contribution is a learnable autoencoder-based compression module placed at the split layer, allowing clients and the server to exchange compressed intermediate activations and gradients instead of full hidden representations.
- Code and toolbox link: https://github.com/Nousphera/AE_PSL
- The repository contains the full training/evaluation pipeline, baseline implementations, autoencoder alignment code, and experiment scripts for reproducing the reported results.
- The tool can be used by cloning the repository and installing the Python dependencies listed in
requirements.txt. Example experiment scripts are provided in the repository.
AE-PSL addresses a practical bottleneck in split learning: even when only a small part of a pre-trained model is kept on the client, every training step still requires exchanging intermediate activations and their gradients with the server. This communication cost becomes especially important for transformer-based foundation models, where hidden representations can be large. Instead of using fixed heuristic compression such as quantization or sparsification, AE-PSL learns an autoencoder that compresses split-layer patch-token activations while preserving task-relevant information.
The tool is designed around a two-stage alignment procedure. First, the autoencoder is generally aligned on public data to match the feature manifold of the pre-trained model. Second, each client performs a lightweight client-specific alignment step on local data, after which the server uses an aggregated decoder while clients keep their adapted encoders. During distributed fine-tuning, the autoencoder modules are frozen and used to reduce the bidirectional communication of activations and gradients. In the reported experiments, AE-PSL preserves downstream accuracy at around 10x communication reduction and provides a stronger accuracy-communication trade-off than evaluated heuristic split-learning compression baselines.
Key features:
- Implements autoencoder-compressed parallel split learning for fine-tuning pre-trained vision transformers.
- Uses General Alignment and Client-Specific Alignment to make learnable compression compatible with off-the-shelf pre-trained models.
- Supports compressed distributed fine-tuning with LoRA, client/server model partitioning, and frozen aligned autoencoder modules.
- Includes baseline comparisons against heuristic split-learning compression methods such as Rand-Top-K, ADC, C3-SL, and no-compression settings.
- Provides experiment scripts for main results, baselines, autoencoder architecture studies, ablations, FLOP profiling, and no-compression reference runs.
The tool operates in the following stages:
- Select an experiment configuration, including dataset, pre-trained model, split layer, number of clients, compression method, AE latent dimension, random seed, and output directories.
- Load the selected vision dataset and construct distributed client partitions. The repository supports datasets such as CIFAR100, Food101, SUN397, FEMNIST, and ImageNet100 for alignment/pre-training workflows.
- Perform General Alignment by training the autoencoder on public data to reconstruct split-layer activations from the frozen pre-trained model.
- Optionally perform Client-Specific Alignment, where each client locally warms up the autoencoder on its own data and contributes decoder weights for server-side aggregation.
- Run autoencoder-compressed distributed fine-tuning. Clients send compressed activations to the server, the server reconstructs them for the server-side model, and compressed gradients are sent back for client-side backpropagation.
- Evaluate the resulting model under global or local evaluation settings, depending on the dataset and experiment type.
- Compare AE-PSL against compression baselines and analyze communication reduction, accuracy, and client-side computational overhead.
The main entrypoint is src/orchestrator.py. Typical executions use the provided scripts, such as run_basic_example.sh, run_main_results.sh, run_baselines.sh, run_ae_architecture.sh, run_ablation.sh, and run_no_compression.sh.
This tool is intended for:
- studying communication-efficient split learning for pre-trained model fine-tuning
- evaluating learnable compression of intermediate activations and gradients
- benchmarking AE-based compression against heuristic split-learning compression baselines
- analyzing the accuracy-communication-computation trade-off in distributed fine-tuning
- reproducing and extending AE-PSL experiments on vision classification benchmarks
Within the AIMS5.0 context, AE-PSL is especially relevant for edge AI and privacy-sensitive distributed learning scenarios where data cannot be centralized and clients may not have the compute capacity to fine-tune full foundation models locally. The toolbox provides a concrete framework for reducing communication overhead in collaborative fine-tuning while keeping the split-learning execution model intact.
- The current repository is a research prototype and experiment framework, not a production-ready split-learning platform.
- The evaluated implementation focuses on vision classification and ViT-style pre-trained models; other modalities are conceptually possible but not yet fully evaluated.
- The current approach keeps the
[CLS] token uncompressed in the main setup to preserve downstream accuracy, which limits how model-agnostic the compression path is.
- The workflow is primarily CLI- and script-driven; it does not provide an end-user graphical interface.
As this research is still in progress, a link will be provided once the paper is published.