Adaptive Selection of LoRA Components in Privacy-Preserving Federated Learning

Abstract: Differentially private federated fine-tuning of large models with LoRA suffers from aggregation error caused by LoRA's multiplicative structure, which is further amplified by DP noise and degrades both stability and accuracy. Existing remedies apply a single update mode uniformly across all layers and all communication rounds (or alternate them on a fixed schedule), ignoring both the structural asymmetry between the two LoRA factors and the round-wise dynamics of training. We propose AS-LoRA, an adaptive framework defined by three axes (i) layer-wise freedom, in which each layer independently selects its active component, (ii) round-wise adaptivity, in which the selection updates over communication rounds, and (iii) a curvature-aware score derived from a second-order approximation of the loss. Theoretically, AS-LoRA eliminates the reconstruction-error floor of layer-tied schedules, accelerates convergence, implicitly biases solutions toward flatter minima, and incurs no additional privacy cost. Across GLUE, SQuAD, CIFAR-100, and Tiny-ImageNet under strict DP budgets and non-IID partitions, AS-LoRA improves over the federated LoRA baselines by up to $+7.5$ pp on GLUE and $+12.5$ pp on MNLI-mm for example, while matching or exceeding SVD-based aggregation methods at $33\text{--}180 \times$ lower aggregation cost and with negligible communication overhead. Code for the proposed method is available at https://anonymous.4open.science/r/as_lora-F75F/.

• The paper introduces AS-LoRA, a framework that adaptively selects LoRA components per layer to eliminate reconstruction errors in DP federated learning.
• It employs EMA-based round scheduling and curvature-aware random projection scoring to accelerate convergence and mitigate gradient outliers.
• Empirical results demonstrate significant accuracy improvements and reduced communication overhead across benchmarks like GLUE, SQuAD, and image datasets.

Adaptive Selection of LoRA Components in Privacy-Preserving Federated Learning

Motivation and Problem Statement

The communication and memory overhead of fine-tuning large-scale foundation models in Federated Learning (FL) is efficiently addressed by Parameter-Efficient Fine-Tuning (PEFT) methods such as Low-Rank Adaptation (LoRA). However, when combined with strict Differential Privacy (DP) constraints, the inherent multiplicative structure of LoRA induces severe aggregation errors that are further amplified by DP-SGD noise, leading to instability and accuracy degradation, especially in the presence of statistical heterogeneity across clients. Existing solutions in federated LoRA—such as FFA-LoRA, RoLoRA, and FedSVD—either select only one LoRA component for global updates across all layers or alternate them according to a fixed scheme, disregarding structural asymmetry and the temporal dynamics of training.

Adaptive Selection LoRA (AS-LoRA): Methodology

The paper introduces AS-LoRA, a novel federated LoRA framework guided by three main axes: layer-wise mode selection, round-wise adaptivity, and curvature-aware, data-driven scoring.

• Layer-wise Freedom: Each layer independently selects the active LoRA component ($A$ or $B$) for optimization at every round. This per-layer selection eliminates the irreducible reconstruction-error floor inherent in layer-tied schedules.
• Round-wise Adaptive Scheduling: Component selection is based on scores aggregated with an Exponential Moving Average (EMA) across communication rounds, adaptively updating as the relative importance of the two LoRA factors shifts over the course of training.
• Curvature-Aware, Random-Projection Scoring: At each selection point, the expected one-step loss improvement for each component is approximated using a second-order Taylor expansion: the score is the squared norm of the projected gradient minus a curvature penalty. Channel-wise gradient outliers, which are typical in LoRA layers, are mitigated by applying Gaussian random projections to gradient vectors prior to score computation, improving the reliability of component selection.

  Figure 1: Comparison of update patterns in FedLoRA, FFA-LoRA, RoLoRA, and the proposed Adaptive Selection LoRA (AS-LoRA).

  

  Figure 2: Schematic workflow of AS-LoRA's adaptive layer-wise mode selection and aggregation cycle.

  

  Figure 3: Details of the AS-LoRA score computation with random projections and temperature-scaled mode selection.

The approach operates at minimal communication overhead, with only a negligible increase in uplink and downlink traffic (per-layer scalar score and mode bit).

Theoretical Results

The proposed adaptive selection framework is shown to offer strictly improved reconstruction error, accelerated convergence, and a guarantee of implicit bias towards flatter minima:

• Zero Additional Privacy Cost: The data-adaptive selection is a post-processing step of DP-SGD outputs and, by standard invariance, incurs no extra privacy cost.
• Strict Dominance: AS-LoRA eliminates the reconstruction error floor of fixed-component and layer-tied schedules. For any fixed selection scheme, there exists an inherent lower bound on achievable population loss, which AS-LoRA can drive to zero with sufficient rounds given well-conditioned input distributions and adequate sampling diversity.
• Accelerated Convergence: The score-based selection can be interpreted as a block-coordinate descent with an adaptive Gauss–Southwell-type rule, achieving strictly faster convergence rates than uniform or static block selection (modulo stochastic variance and curvature penalty scaling).
• Flatness Bias: By prioritizing components with lower curvature penalties, AS-LoRA implicitly regularizes towards flatter loss landscapes, as observed in sharpness-aware minimization (SAM).

Empirical Results

AS-LoRA is evaluated on a spectrum of benchmarks, including GLUE, SQuAD, CIFAR-100, and Tiny-ImageNet, with non-IID partitions and strict differential privacy ($\epsilon \in \{1, 3\}$). Robust performance gains are observed across tasks:

• On GLUE, AS-LoRA achieves up to +7.5 percentage points improvement over other federated LoRA methods under $\epsilon = 3$, and +12.5 pp over FFA-LoRA on MNLI-mm at $\epsilon=1$.
• On SQuAD v1.1 and v2.0, AS-LoRA outperforms FFA-LoRA and RoLoRA by a clear margin in both EM and F1 metrics.
• On CIFAR-100 and Tiny-ImageNet, consistent gains ($+2.34$pp and $+1.19$pp, respectively) are reported.

Compared to server-side SVD-based aggregation (FedSVD), AS-LoRA achieves matching or superior accuracy with aggregation cost reduced by two orders of magnitude and maintains negligible communication overhead per round.

Robustness, Overhead, and Ablation

• Gradient Outlier Mitigation: Random projections of local gradients are essential when batch sizes are small; their contribution diminishes with larger batches due to inherent variance reduction.
• Computation Scheduling: Curvature-based scoring is expensive but can be practically scheduled (using one-sided finite difference or late-phase periodic computation) to reduce overhead below 10% with negligible performance loss.
• Mode Selection Granularity: Strictly per-client mode selection is suboptimal due to exacerbation of aggregation errors; a global, layer-wise strategy is optimal in practice.
• Task- and Layer-Specific Dynamics: Adaptive selection uncovers that layer and task preferences for LoRA components are non-uniform and persistently structured.

  Figure 4: Loss landscape visualization of FFA-LoRA, RoLoRA, and AS-LoRA on MNLI under DP-SGD ($\epsilon=3$); AS-LoRA finds flatter minima.

  

  Figure 5: Effects of random projection on LoRA-A score ratio across batch sizes, highlighting improved robustness of score estimates at small batches.

  

  Figure 6: Analysis of data heterogeneity (Dirichlet $\alpha$) and LoRA rank $r$ on MNLI; AS-LoRA is robust to high non-IID and low-rank regimes.

Analysis of Mode Selection

Extensive analysis indicates that the preferred active component alternates in a temporally and layer-structured manner, varying significantly across language and vision tasks. For certain NLP tasks, highly correlated selection patterns emerge in intermediate layers, implying potential for future structured selection models exploiting cross-layer dependencies.

Implications and Future Directions

AS-LoRA supplies a robust, privacy-compatible, communication-efficient, and computationally tractable solution for federated fine-tuning of large models under DP constraints. By eliminating irreducible reconstruction errors, accelerating convergence rates, and regularizing for solution flatness, it fundamentally raises the attainable utility-efficiency frontier in practical privacy-preserving FL contexts.

Future work may extend the independent layer-wise selection to joint or structured mode scheduling, potentially capturing global or cross-layer dependencies supported by empirical task-specific selection patterns. Additionally, further reductions in curvature estimation overhead or tighter integration with adaptive optimizer design remain promising avenues.

Conclusion

AS-LoRA advances the state of parameter-efficient, privacy-preserving federated adaptation for foundation models, defining a new optimality class among LoRA scheduling protocols. Its theoretical and empirical strengths demonstrate that dynamic, data-driven layer-wise mode selection is critical for stabilizing and maximizing accuracy in differentially private FL with non-IID data, with immediate applicability to both language and vision domains (2605.05769).