Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sequential Compression Layers for Efficient Federated Learning in Foundational Models

Navyansh Mahla, Sunny Gupta, Amit Sethi

TL;DR

The authors address inefficiencies of LoRA in federated fine-tuning of foundational models by introducing a sequential compression layer that inserts a compact MLP between the first and second MLPs within the transformer FFN. By freezing the initial MLP and attention components and training only the compression layer and the subsequent projection, the method achieves a parameter-efficient update mechanism with a linear excess risk bound of $ abla$-flow through a controlled subspace, in contrast to LoRA’s quadratic bounds. Empirically, the approach delivers substantial performance gains on both text (MedQuAD, Dolly-15k with Gemma-2B/TinyLlama) and vision (Brain Tumor with SigLIP) tasks under highly non-IID Dirichlet partitions, outperforming FedSA-LoRA and FFA-LoRA while maintaining comparable parameter efficiency. The combination of theoretical guarantees and consistent cross-modal improvements suggests strong practical potential for privacy-preserving distributed learning of foundation models in resource-constrained FL deployments.

Abstract

Federated Learning (FL) has gained popularity for fine-tuning large language models (LLMs) across multiple nodes, each with its own private data. While LoRA has been widely adopted for parameter efficient federated fine-tuning, recent theoretical and empirical studies highlight its suboptimal performance in the federated learning context. In response, we propose a novel, simple, and more effective parameter-efficient fine-tuning method that does not rely on LoRA. Our approach introduces a small multi-layer perceptron (MLP) layer between two existing MLP layers the up proj (the FFN projection layer following the self-attention module) and down proj within the feed forward network of the transformer block. This solution addresses the bottlenecks associated with LoRA in federated fine tuning and outperforms recent LoRA-based approaches, demonstrating superior performance for both language models and vision encoders.

Sequential Compression Layers for Efficient Federated Learning in Foundational Models

TL;DR

The authors address inefficiencies of LoRA in federated fine-tuning of foundational models by introducing a sequential compression layer that inserts a compact MLP between the first and second MLPs within the transformer FFN. By freezing the initial MLP and attention components and training only the compression layer and the subsequent projection, the method achieves a parameter-efficient update mechanism with a linear excess risk bound of -flow through a controlled subspace, in contrast to LoRA’s quadratic bounds. Empirically, the approach delivers substantial performance gains on both text (MedQuAD, Dolly-15k with Gemma-2B/TinyLlama) and vision (Brain Tumor with SigLIP) tasks under highly non-IID Dirichlet partitions, outperforming FedSA-LoRA and FFA-LoRA while maintaining comparable parameter efficiency. The combination of theoretical guarantees and consistent cross-modal improvements suggests strong practical potential for privacy-preserving distributed learning of foundation models in resource-constrained FL deployments.

Abstract

Federated Learning (FL) has gained popularity for fine-tuning large language models (LLMs) across multiple nodes, each with its own private data. While LoRA has been widely adopted for parameter efficient federated fine-tuning, recent theoretical and empirical studies highlight its suboptimal performance in the federated learning context. In response, we propose a novel, simple, and more effective parameter-efficient fine-tuning method that does not rely on LoRA. Our approach introduces a small multi-layer perceptron (MLP) layer between two existing MLP layers the up proj (the FFN projection layer following the self-attention module) and down proj within the feed forward network of the transformer block. This solution addresses the bottlenecks associated with LoRA in federated fine tuning and outperforms recent LoRA-based approaches, demonstrating superior performance for both language models and vision encoders.

Paper Structure

This paper contains 14 sections, 12 equations, 4 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Federated Learning
  4. LoRA in Federated Learning
  5. Methodology
  6. Experiments
  7. Datasets
  8. Non-IID Data Preparation
  9. Models and Hyperparameters
  10. Experiment Results
  11. Conclusion
  12. Appendix
  13. Proofs
  14. Dataset Shard Class Distribution

Figures (4)

  • Figure 1: The block at the top represents the original architecture and the one at the bottom represents the architecture with compression layers added. MLP Layer 1 and Self Attention modules are frozen during training. The new compression layer is added between MLP Layer 1 and MLP Layer 2 such that $d\ll M,m$
  • Figure 2: From mahla2025exploringgradientsubspacesaddressing: Label distribution across shards for the MedQuAD dataset produced using Dirichlet Allocation with $\alpha=0.1$.
  • Figure 3: From mahla2025exploringgradientsubspacesaddressing: Label distribution across shards for the Dolly dataset produced using Dirichlet Allocation with $\alpha=0.1$.
  • Figure 4: From mahla2025exploringgradientsubspacesaddressing: Distribution of labels across shards for Brain Tumour Dataset produced using Dirichlet Allocation with $\alpha=0.1$