Technical Summary: PreLort – Prefix-Nested LoRA for Federated Fine-Tuning under Rank Heterogeneity

1. Problem Statement

Federated Learning (FL) enables collaborative training of Large Language Models (LLMs) without sharing raw data, but it faces significant challenges when adapting foundation models via parameter-efficient fine-tuning (PEFT) methods like Low-Rank Adaptation (LoRA). A primary obstacle is rank heterogeneity: in realistic cross-device settings, clients possess varying computational and memory constraints, leading them to initialize LoRA adapters with different ranks ($r$).

Existing solutions for heterogeneous federated LoRA primarily focus on the aggregation phase, often employing zero-padding to align adapters of different ranks before averaging. However, the authors identify a fundamental limitation in these approaches: representation misalignment. Even when aggregation is mathematically defined, combining independently trained models with different ranks leads to inconsistencies. Specifically, lower-rank clients do not benefit from the updates of higher-rank clients because the information is not distributed hierarchically across rank dimensions. Zero-padding further dilutes the contributions of high-rank clients by introducing zero-valued segments, resulting in suboptimal use of shared low-rank representations.

2. Methodology: PreLort

The authors propose PreLort (Prefix-Nested LoRA), a framework that addresses rank heterogeneity by structuring the training dynamics and aggregation rules jointly. The core philosophy is to organize adapter dimensions into a prefix hierarchy, ensuring that lower-rank dimensions encode task-relevant information shared across all clients, while higher-rank dimensions capture incremental capacity.

2.1 Nested Low-Rank Structure

PreLort imposes a prefix-structured organization on adapter dimensions. For a client with rank $r_k$, only the first $r_k$ dimensions are active. Crucially, during local training, the adapter is not optimized solely at its full rank. Instead, it is optimized under multiple truncated configurations (prefix sub-ranks).

• Mechanism: A client with rank $r_k$ is trained to perform well not only at $r_k$ but also at all lower rank levels present in the federation ($r \in \{r \mid r \le r_k\}$).
• Objective: This forces the model to concentrate task-relevant signals into the leading (low-rank) dimensions, ensuring that these dimensions remain informative and useful even when truncated.

2.2 Local Training Strategy

The local objective for a client $k$ is defined as the average of cross-entropy losses computed over all valid prefix sub-ranks $R_k$:
$$\mathcal{L}_k = \frac{1}{|R_k|} \sum_{r \in R_k} \mathcal{L}_{CE}(f(\cdot; \theta, r), \mathcal{D}_k)$$
In practice, each training step involves one full-rank forward/backward pass and additional passes for all sub-ranks. This strategy ensures that the prefix segments are explicitly optimized at the same granularity used during aggregation, encouraging a consistent hierarchy where lower dimensions capture the most critical information.

2.3 Segment-wise Aggregation

To avoid the dilution caused by zero-padding, PreLort introduces a segment-wise aggregation rule operating in the gradient (update) space.

• Process: Instead of averaging all clients for every dimension, the server aggregates gradients for each rank segment $[r_{l-1}, r_l)$ only over the subset of clients whose rank $r_k$ is sufficient to contribute to that segment (i.e., $r_k \ge r_l$).
• Formula: For the $l$-th segment, the aggregated gradient is:
  $$\nabla A_{global}^{r_{l-1}:r_l} = \frac{\sum_{k: r_k \ge r_l} n_k \nabla A_k^{r_{l-1}:r_l}}{\sum_{k: r_k \ge r_l} n_k}$$
• Benefit: This prevents signal dilution from zero-padded lower-rank clients and ensures that higher-rank dimensions are updated only by clients that actually possess the capacity to learn them, while lower dimensions benefit from the full population.

3. Key Contributions

The paper outlines three primary contributions:

1. Identification of Representation Misalignment: The authors pinpoint that the lack of structured alignment across rank dimensions is a key limitation in current heterogeneous federated LoRA methods, leading to inconsistent representations.
2. Nested Formulation and Unified Strategy: They propose a nested low-rank formulation that imposes a common prefix structure. This is coupled with a unified training and aggregation strategy that enables consistent aggregation across heterogeneous ranks without zero-padding.
3. Empirical Validation: The method is demonstrated to consistently improve model quality and stability over prior heterogeneous federated LoRA methods on instruction-tuning benchmarks.

4. Experimental Results

The authors evaluated PreLort on two base models (TinyLlama-1.1B and Qwen2.5-0.5B) across three datasets: Alpaca, Databricks-dolly-15, and 20 Newsgroups.

• Performance Metrics: PreLort consistently outperformed baselines (ZeroPad, HetLoRA, FLoRA) in Accuracy (MMLU, 20 Newsgroups) and ROUGE-L scores.
  • On TinyLlama-1.1B, PreLort achieved 33.86% accuracy on MMLU (Dolly) and 30.76% on Alpaca, surpassing the next best baselines.
  • On Qwen2.5-0.5B, it achieved 39.85% accuracy on MMLU (Dolly).
• Perplexity: PreLort achieved the lowest perplexity on Dolly and NewsGroup for both base models, demonstrating superior token-level likelihood.
• Rank Robustness: Analysis showed that PreLort maintains stable performance across different prefix truncation levels. In contrast, baselines like "Nested Aggregation Only" exhibited instability, particularly at higher ranks, confirming that aggregation alone is insufficient without the nested training objective.
• Per-Dimension Importance: Visualization of dimension importance revealed that PreLort concentrates task-relevant signal in leading dimensions (a structured distribution), whereas standard training yields a uniform distribution. This hierarchical prioritization facilitates stable updates across heterogeneous ranks.

5. Significance and Claims

The paper claims that PreLort represents a shift in addressing rank heterogeneity: rather than treating it as a post-hoc aggregation problem, it is solved by structuring the training process itself.

• Core Argument: The authors argue that addressing rank heterogeneity requires not only improved aggregation but also structured alignment of representations across ranks.
• Efficiency: The method allows low-rank clients to benefit from richer information contributed by higher-rank clients because the prefix dimensions are consistently learned and aggregated.
• Modesty on Limitations: The authors acknowledge that the nested training strategy introduces approximately 40% additional computational overhead due to multiple forward/backward passes per batch. They also note that the effectiveness depends on sufficient local optimization steps; with very limited local epochs, the model may not fully learn well-aligned low-rank prefix representations.

In conclusion, PreLort demonstrates that rank heterogeneity in federated LoRA is best handled by jointly structuring how adapters are trained (via nested prefix optimization) and how their updates are aggregated (via segment-wise rules), resulting in aligned, robust prefix representations that outperform existing heterogeneous aggregation methods.