Federated Co-tuning Framework for Large and Small Language Models

1. Problem Statement

The paper addresses three critical challenges in adapting Large Language Models (LLMs) to domain-specific tasks within a federated setting:

• Data Privacy: Downstream clients cannot share sensitive, domain-specific labeled data with LLM owners for fine-tuning due to privacy regulations and security concerns.
• Resource Constraints: Fine-tuning massive LLMs requires prohibitive computational and storage resources, making them inaccessible for small and medium-sized enterprises (SMEs). Consequently, these clients are often restricted to training smaller models (SLMs) locally.
• Lack of Mutual Enhancement: Existing methods typically treat LLMs and SLMs in isolation. There is a gap in frameworks that enable mutual knowledge transfer, where server-side LLMs guide client SLMs, and client SLMs simultaneously enrich server LLMs with domain-specific insights.

2. Methodology: FedCoLLM

The authors propose FedCoLLM, a parameter-efficient federated co-tuning framework that creates a symbiotic relationship between a server-side LLM and multiple client-side SLMs.

Core Architecture

• Server Side: Hosts a large LLM ($f_\psi$) and an auxiliary distillation dataset ($D_a$).
• Client Side: $K$ clients, each possessing a Small Language Model ($g_\phi$) and private local data ($D_k$).
• The Bridge: A shared SLM architecture acts as the communication bridge. Instead of training the entire model, clients and the server utilize Parameter-Efficient Fine-Tuning (PEFT), specifically LoRA (Low-Rank Adaptation) adapters.

Training Workflow

The framework operates in iterative communication rounds involving two distinct phases:

1. Federated Fine-Tuning (Client Side):

   • The server broadcasts a global LoRA adapter ($\theta$) to all clients.
   • Each client $k$ fine-tunes its local adapter ($\theta_k$) using its private data $D_k$ while keeping the base SLM parameters frozen.
   • Clients upload only the updated adapter parameters to the server.
   • The server aggregates these updates using SecureAggregation (e.g., SecureAvg) to create a new global adapter, preserving data privacy.

2. Mutual Knowledge Distillation (Server Side):

   • The server performs Mutual Knowledge Distillation (MKD) between the global SLM (with the aggregated adapter $\theta$) and the server's LLM (with its own adapter $\omega$).
   • This process utilizes the auxiliary dataset $D_a$.
   • Loss Functions:
     • Supervised Fine-Tuning (SFT): Both models are trained on $D_a$ using Cross-Entropy loss ($\ell_{CE}$).
     • Knowledge Distillation: The models learn from each other's probability distributions using Kullback-Leibler (KL) divergence ($\ell_{KL}$). The LLM guides the SLM, and the SLM (trained on diverse domain data) provides feedback to the LLM.
   • The updated global adapter $\theta$ is then redistributed to clients for the next round.

Optimization Objective

The framework minimizes two objectives simultaneously:

1. Client Objective: Minimize task loss on local data using the global adapter.
2. Server Objective: Minimize a combined loss of SFT and Mutual KD on the auxiliary dataset to enhance both the LLM and the global SLM.

3. Key Contributions

• Novel Framework: FedCoLLM is the first framework to enable bidirectional knowledge transfer between server LLMs and client SLMs in a federated setting.
• Parameter Efficiency: By leveraging LoRA adapters, the framework drastically reduces the number of parameters requiring updates, making it feasible for resource-constrained clients.
• Privacy Preservation: It adheres to strict privacy standards by never transmitting raw data. It combines Federated Learning (FL) with SecureAggregation and PEFT to minimize information leakage.
• Resource Optimization: It eliminates the need for clients to host or fine-tune full LLMs, offloading heavy computation to the server while still allowing clients to benefit from domain-specific adaptation.

4. Experimental Results

The authors evaluated FedCoLLM using various models (GPT-2, OPT, LLaMa2) across four NLP question-answering datasets (CommonsenseQA, OpenBookQA, ARC-C, ARC-E).

• Server-Side LLM Performance:
  • FedCoLLM significantly outperformed Zero-Shot baselines (e.g., +47% improvement on GPT-2-Large).
  • It achieved performance comparable to Centralized Fine-Tuning (where the server has access to all client data directly), reaching 97–99% of the centralized performance across different models.
• Client-Side SLM Performance:
  • Client SLMs trained with FedCoLLM outperformed Standalone (local-only) and FedAvg baselines.
  • Improvements ranged from 2% to 6% over Standalone and 2% to 4% over FedAvg, demonstrating the value of the server's LLM guidance.
• Communication Efficiency:
  • FedCoLLM reduced communication costs by ~99.7% compared to full-model fine-tuning.
  • For example, on LLaMa2, only 0.23% of the total parameters were transmitted per round.

5. Significance

FedCoLLM represents a significant step forward in the democratization of LLMs. It solves the "privacy vs. performance" trade-off by allowing enterprises to leverage powerful server-side LLMs without exposing their proprietary data. Simultaneously, it allows server-side models to evolve by absorbing diverse, real-world domain knowledge from clients. This framework is particularly impactful for industries with strict data governance (e.g., finance, healthcare) where SMEs need high-performance AI but lack the infrastructure to train or host large models independently.

The code is open-sourced within the FATE (Federated AI Technology Enabler) project, facilitating reproducibility and adoption.