Federated Learning with Ad-hoc Adapter Insertions: The Case of Soft-Embeddings for Training Classifier-as-Retriever (2509.16508v1)

Abstract: When existing retrieval-augmented generation (RAG) solutions are intended to be used for new knowledge domains, it is necessary to update their encoders, which are taken to be pretrained LLMs. However, fully finetuning these large models is compute- and memory-intensive, and even infeasible when deployed on resource-constrained edge devices. We propose a novel encoder architecture in this work that addresses this limitation by using a frozen small LLM (SLM), which satisfies the memory constraints of edge devices, and inserting a small adapter network before the transformer blocks of the SLM. The trainable adapter takes the token embeddings of the new corpus and learns to produce enhanced soft embeddings for it, while requiring significantly less compute power to update than full fine-tuning. We further propose a novel retrieval mechanism by attaching a classifier head to the SLM encoder, which is trained to learn a similarity mapping of the input embeddings to their corresponding documents. Finally, to enable the online fine-tuning of both (i) the encoder soft embeddings and (ii) the classifier-as-retriever on edge devices, we adopt federated learning (FL) and differential privacy (DP) to achieve an efficient, privacy-preserving, and product-grade training solution. We conduct a theoretical analysis of our methodology, establishing convergence guarantees under mild assumptions on gradient variance when deployed for general smooth nonconvex loss functions. Through extensive numerical experiments, we demonstrate (i) the efficacy of obtaining soft embeddings to enhance the encoder, (ii) training a classifier to improve the retriever, and (iii) the role of FL in achieving speedup.

• The paper introduces an adapter-based soft embedding approach that replaces full model fine-tuning for efficient domain adaptation in federated settings.
• It proposes a classifier-as-retriever mechanism that significantly improves retrieval accuracy over traditional MIPS, achieving up to 99.95% accuracy on benchmark datasets.
• By integrating differential privacy with federated learning, the method enables scalable, privacy-preserving training on resource-constrained edge devices with minimal accuracy trade-offs.

Federated Learning with Ad-hoc Adapter Insertions: Soft-Embeddings for Training Classifier-as-Retriever

Introduction and Motivation

This work addresses the limitations of current retrieval-augmented generation (RAG) systems, particularly in the context of resource-constrained edge devices and privacy-preserving distributed training. The authors propose a novel architecture that leverages a frozen small LLM (SLM) with lightweight, trainable adapters and a classifier head, enabling efficient adaptation to new domains without the computational and memory overhead of full model fine-tuning. The approach is further integrated with federated learning (FL) and differential privacy (DP) to support distributed, privacy-preserving training across multiple clients.

The key contributions are: (1) the insertion of small adapter networks for parameter-efficient fine-tuning (PEFT) to obtain soft embeddings, (2) the introduction of a classifier-as-retriever (CaR) mechanism as a trainable alternative to maximum inner-product search (MIPS), and (3) the integration of FL and DP for scalable, privacy-preserving training.

Figure 1: Overview of the CaR approach, showing adapter insertion for PEFT and the classifier head for retrieval, enabling two training modes: classifier-only and joint adapter-classifier fine-tuning.

Methodology

Adapter Insertion and Soft Embeddings

The architecture replaces full LLM fine-tuning with a frozen SLM encoder, into which a small, trainable adapter (a square transformation matrix) is inserted between the token embedding and the first transformer block. This adapter is responsible for learning corpus-specific soft embeddings, steering the representation space toward the target domain with minimal parameter updates. The adapter is significantly smaller than the full model, reducing both memory and communication costs in FL settings.

Classifier-as-Retriever (CaR)

Instead of relying on heuristic similarity functions such as MIPS, the retriever is augmented with a classifier head. This classifier is trained to map input queries to their corresponding documents, effectively learning a data-driven similarity function. During inference, the top-$K$ classifier outputs are used for retrieval, which are then passed to the generator for response synthesis. This approach enables more flexible and accurate retrieval, especially in domain adaptation scenarios.

Figure 2: Comparison between conventional MIPS-based retrieval and the proposed Classifier-as-Retriever methodology.

Federated Learning and Differential Privacy

The FL setup distributes training across $m$ clients, each holding a local dataset. Only the adapter and classifier head parameters are updated and communicated, while the SLM encoder and generator remain frozen. The aggregation is performed via FedAvg, and both fixed and adaptive DP mechanisms are applied to the parameter updates. DP is enforced by gradient clipping and Gaussian noise addition, ensuring client-level privacy before aggregation.

Theoretical analysis establishes convergence guarantees for the proposed FL+DP training under general smooth non-convex loss functions, with explicit bounds on the stationarity gap as a function of learning rate, clipping threshold, and noise variance.

Experimental Results

Comparison with MIPS and Baseline Architectures

Empirical results on the SMS spam, AG News, and arXiv metadata datasets demonstrate the superiority of the proposed approach. On the SMS spam dataset, MIPS with a frozen LLM achieves only 12.36% retrieval accuracy. Inserting a trainable adapter boosts this to 96.79%, while training only the classifier head achieves 97.02%. Joint training of both components yields 99.95% accuracy, indicating a strong synergistic effect.

Figure 3: Training loss for joint soft embedding adapter and classifier head optimization.

Federated and Distributed Training

FL experiments show that distributed training achieves comparable accuracy to centralized baselines, with significant speedup proportional to the number of clients. For example, on the AG News dataset with Llama-3.2-3B, FL with three clients achieves a 2.4x speedup over centralized training, with only a minor reduction in accuracy (94.18% centralized vs. 93.39% FL). The integration of DP results in a small additional drop in accuracy, but preserves privacy guarantees.

Figure 4: Performance of Llama-3.2-3B on AG News with pre-classifier and mean pooling, highlighting the impact of model size and pooling strategy.

Scaling to Larger Models and Datasets

The methodology scales to larger models (e.g., Llama-3.2-3B) and datasets (e.g., arXiv metadata with 173 classes). Larger models provide higher-quality embeddings, and mean pooling over token outputs further improves classifier performance. The approach supports distributed inference, where clients extract hidden states from a remote, frozen LLM and train lightweight classifier heads locally.

Figure 5: Results on arXiv metadata with Llama-3.2-3B, demonstrating improved performance with larger models and mean pooling.

Theoretical Analysis

The paper provides a rigorous convergence analysis for the FL+DP training procedure. Under standard smoothness and bounded variance assumptions, the method achieves an $\mathcal{O}(1/T)$ convergence rate to a stationarity neighborhood, with explicit dependence on the learning rate, clipping threshold, and noise scale. The analysis covers both fixed and adaptive DP regimes, and recovers standard FL results in the absence of clipping and noise.

Practical Implications and Future Directions

The proposed architecture enables efficient, scalable, and privacy-preserving RAG systems suitable for deployment on edge devices. By decoupling the heavy transformer blocks (which can be hosted in distributed clusters) from the lightweight, trainable components (adapters and classifier heads), the method supports both resource-constrained and distributed environments. The classifier-as-retriever paradigm offers a flexible alternative to MIPS, particularly advantageous in domain adaptation and large-scale retrieval scenarios.

The authors discuss potential extensions, including distributed load offloading of transformer blocks and the use of mixture-of-experts (MoE) architectures to handle large-scale classification tasks. These directions are promising for further reducing memory and compute requirements while maintaining retrieval accuracy.

Conclusion

This work introduces a parameter- and communication-efficient RAG retriever architecture, combining adapter-based soft embeddings and a classifier-as-retriever head, integrated with federated learning and differential privacy. The approach achieves strong empirical performance, significant training speedup in distributed settings, and theoretical convergence guarantees. The methodology is well-suited for scalable, privacy-preserving deployment of RAG systems on edge devices and in federated environments, and provides a foundation for future research on efficient, adaptive retrieval mechanisms in LLMs.