FedTreeLoRA: Reconciling Statistical and Functional Heterogeneity in Federated LoRA Fine-Tuning
Abstract
Federated Learning (FL) with Low-Rank Adaptation (LoRA) has become a standard for privacy-preserving LLM fine-tuning. However, existing personalized methods predominantly operated under a restrictive Flat-Model Assumption: they addressed client-side statistical heterogeneity but treated the model as a monolithic block, ignoring the functional heterogeneity across LLM layers. We argue that these two statistical (horizontal) and functional (vertical) dimensions, are orthogonal in source yet coupled in interaction, implying that the optimal depth of parameter sharing is functionally dependent on client similarity. To address this, we propose FedTreeLoRA, a framework employing tree-structured aggregation for fine-grained, layer-wise alignment. By dynamically constructing an aggregation hierarchy, FedTreeLoRA allows clients to share broad consensus on shallow ’trunks‘ while progressively specializing on deep ‘branches'. Experiments on NLU and NLG benchmarks demonstrate that FedTreeLoRA significantly outperforms state-of-the-art methods by effectively reconciling generalization and personalization.
Lay Summary
Large language models are increasingly being adapted on private and distributed data, such as information stored on personal devices or across different organizations. Federated learning makes this possible by allowing many users to collaboratively improve a model without sharing their raw data. However, existing methods often struggle when users have very different types of data, because they treat the entire model as if every part should be shared in the same way. In our work, we show that different layers inside a language model play different roles. Earlier layers tend to learn general language patterns that can be safely shared across users, while deeper layers capture more personalized or task-specific knowledge. Based on this observation, we develop FedTreeLoRA, a new federated fine-tuning method that organizes users into a tree-like structure and automatically decides how much information should be shared at each layer of the model. Our approach allows similar users to collaborate more closely while still preserving personalized behavior when needed. Experiments on multiple language understanding and text generation tasks show that FedTreeLoRA consistently improves performance over existing federated learning methods, while adding very little extra computational cost. This work offers a more flexible and efficient way to personalize large AI models in privacy-sensitive environments.