Search papers, labs, and topics across Lattice.
The paper introduces FLoRG, a federated learning framework for fine-tuning LLMs that uses a single low-rank matrix and aggregates its Gram matrix to avoid aggregation errors and reduce communication overhead. To address decomposition drift during server-side factor recovery, FLoRG incorporates a Procrustes alignment approach to ensure consistent updates across fine-tuning rounds. Theoretical analysis demonstrates a tighter convergence bound with Procrustes alignment, and experiments show that FLoRG outperforms existing methods in downstream task accuracy while significantly reducing communication costs.
By aggregating Gram matrices instead of low-rank factors, FLoRG sidesteps decomposition drift in federated LoRA, achieving up to 2041x communication savings and improved accuracy.
Parameter-efficient fine-tuning techniques such as low-rank adaptation (LoRA) enable large language models (LLMs) to adapt to downstream tasks efficiently. Federated learning (FL) further facilitates this process by enabling collaborative fine-tuning across distributed clients without sharing private data. However, the use of two separate low-rank matrices in LoRA for federated fine-tuning introduces two types of challenges. The first challenge arises from the error induced by separately aggregating those two low-rank matrices. The second challenge occurs even when the product of two low-rank matrices is aggregated. The server needs to recover factors via matrix decomposition, which is non-unique and can introduce decomposition drift. To tackle the aforementioned challenges, we propose FLoRG, a federated fine-tuning framework which employs a single low-rank matrix for fine-tuning and aggregates its Gram matrix (i.e., the matrix of inner products of its column vectors), eliminating the aggregation error while also reducing the communication overhead. FLoRG minimizes the decomposition drift by introducing a Procrustes alignment approach which aligns the decomposed matrix between consecutive fine-tuning rounds for consistent updates. We theoretically analyze the convergence of FLoRG and prove that adopting the Procrustes alignment results in a tighter convergence bound. Experimental results across multiple LLM fine-tuning benchmarks demonstrate that FLoRG outperforms five state-of-the-art baseline schemes in the downstream task accuracy and can reduce the communication overhead by up to 2041$\times$.
