Revisiting Federated Fine-Tuning: A Single Communication Round is Enough for Foundation Models

TL;DR

The paper tackles the high communication costs of federated fine-tuning for foundation models and demonstrates that a single aggregation round (one-shot federated fine-tuning) can achieve performance comparable to, or better than, multi-round federated learning for large models. A theoretical bound shows the one-shot gap scales with the Lipschitz smoothness $L$, update magnitude $\tau$, total horizon $T k$, and client count $m$ as $\|\varepsilon\| \le (L\tau T k m) \|w^{(0,0)}\|$, with large FMs exhibiting smaller $L$ and $\tau$ and requiring fewer fine-tuning epochs $T k$. Empirically, experiments on models from 1B to 13B parameters across MMLU, Wizard, ARC, and text-to-image tasks reveal that one-shot FL can match multi-round FL while dramatically reducing communication (roughly by a factor of $1/T$), enabling asynchronous aggregation, and reducing certain client-side privacy risks. The results suggest a practical pathway to efficient, scalable federated adaptation of FMs, with broad implications for deployment under bandwidth and privacy constraints.

Abstract

The recent advancement of foundation models (FMs) has increased the demand for fine-tuning these models on large-scale cross-domain datasets. To address this, federated fine-tuning has emerged, allowing FMs to be fine-tuned on distributed datasets across multiple devices while ensuring data privacy. However, the substantial parameter size and the multi-round communication in federated learning algorithms result in prohibitively high communication costs, challenging the practicality of federated fine-tuning. In this paper, we identify and analyze, both theoretically and empirically, that the traditional multi-round aggregation algorithms may not be necessary for federated fine-tuning large FMs. Our experiments reveal that a single round of aggregation (i.e., one-shot federated fine-tuning) yields a global model performance comparable to that achieved through multiple rounds of aggregation. Through rigorous mathematical and empirical analyses, we demonstrate that large FMs, due to their extensive parameter sizes and pre-training on general tasks, achieve significantly lower training loss in one-shot federated fine-tuning compared to smaller models. Our extensive experiments show that one-shot federated fine-tuning significantly reduces communication costs. It also has the potential to enable asynchronous aggregation, enhances privacy, and maintains performance consistency with multi-round federated fine-tuning on both text generation and text-to-image generation tasks. Our findings provide insights to revolutionize federated fine-tuning in practice, enhancing efficiency, reducing costs, and expanding accessibility for FMs.

Figures (8)

• Figure 1: The distinct performances of one-shot federated learning between small models and large FMs. The horizontal axis represents multi-round FL accuracy, while the vertical axis represents one-shot FL accuracy. The ResNet-18 and LSTM are trained and tested on CIFAR-10 and Shakespeare respectively. Other models are fine-tuned on Wizard dataset and tested on ARC Easy. The closer the points are to the dashed line means the accuracy of one-shot and multi-round FL are closer in the corresponding model.
• Figure 2: Experiment results on $L$, $\tau$, and $\|\bm{w}^{(0,0)} \|$ in different models. We use the CIFAR-10 dataset to compute the gradient on ResNet18 he2016deep. We use the WizardLM dataset xu2023wizardlm to compute the gradient on the language models. Models to the left of the red dashed line are small models, while those to the right are foundation models (FMs). The figures indicate that large FMs have significantly smaller $L$ and $\tau$ values compared to small models. Additionally, $\|\bm{w}^{(0,0)} \|$ does not increase proportionally with the model size. In conclusion, without considering other unrelated influencing factors, the value of $\Gamma \|\bm{w}^{(0,0)}\|$ decreases as the model size increases.
• Figure 3: Aggregation in FL with two clients. (a): one round aggregation with $Tk$ local epochs of both the FM and the small model (SM). Although the FM and the SM may differ in practice, we depict both using the same setup as in (a) to make the comparison easier to understand. (b): FL of FM under $T$ round aggregation with $k$ local epochs, where the landscape is smooth and fine-tuning makes small changes, so the new local directions remain close to the one-shot direction. (c): FL of SM under $T$ round aggregation with $k$ local epochs, where the training direction changes markedly after each global update. The final global model deviates significantly from the one-shot global model. This figure visualizes the intuition behind our analysis that one round works well for FMs.
• Figure 4: The estimated $\log \| \varepsilon \|$ in different models calculated by $\log \| \varepsilon \| = \log (L\tau Tk \|\bm{w}^{(0,0)} \|)$.
• Figure 5: "A photo of a dog in a bucket" generated by LoRA fine-tuned stable diffusion models. The left column shows images generated by the one-shot FL models, along with their CLIP scores; the right column shows the results of multi-round FL.