Stacking-Based Aggregation (FLoRA)

• Stacking-based aggregation is a method that precisely concatenates block-disjoint low-rank client updates to eliminate cross-term noise in federated learning.
• It forms stacked matrices from individual client updates, supporting heterogeneous adapter ranks without the need for zero-padding or rigid constraints.
• Empirical results show FLoRA improves fine-tuning and hyperparameter optimization efficiency while reducing communication overhead and scaling to many clients.

Stacking-based aggregation, as instantiated in the FLoRA method, refers to mathematically precise matrix stacking strategies for federated aggregation of heterogeneous low-rank model updates, particularly in LLM fine-tuning and federated hyperparameter optimization. This approach addresses key aggregation noise and rank heterogeneity issues in previous federated learning (FL) protocols, providing noise-free, scalable, and efficient update composition across diverse clients. The core principle is the elimination of cross-term aggregation error by concatenating and summing block-disjoint client updates, ensuring faithful and resource-appropriate federated model improvement. The stacking-based paradigm underpins two prominent works: FLoRA for federated LLM fine-tuning with arbitrary low-rank adapters (Wang et al., 2024) and FLoRA for single-shot federated hyperparameter optimization via surrogate regression stacking (Zhou et al., 2021).

1. Federated Fine-Tuning and the Aggregation Challenge

Federated fine-tuning of LLMs involves $K$ clients, each accessing a shared frozen model $W\in \mathbb{R}^{m\times n}$. Clients train local low-rank adapters $\Delta W_k= B_kA_k$ (with $A_k\in \mathbb{R}^{r_k\times n}$, $B_k\in \mathbb{R}^{m\times r_k}$, $r_k\ll \min(m,n)$), reflecting individual data and resource profiles. The server’s objective is to aggregate these $\Delta W_k$ into a unified global update $\Delta W$. Traditional approaches (notably FedAvg-LoRA/FedIT) average $A_k$ and $B_k$ independently and compute the product, leading to

$W\in \mathbb{R}^{m\times n}$0

which expands to include cross-terms $W\in \mathbb{R}^{m\times n}$1 ($W\in \mathbb{R}^{m\times n}$2), introducing "aggregation noise." This noise not only corrupts the desired weighted sum $W\in \mathbb{R}^{m\times n}$3 but enforces a rigid constraint that all $W\in \mathbb{R}^{m\times n}$4 be identical—a poor fit for heterogeneous client capability (Wang et al., 2024).

2. Stacking-Based Aggregation: Mathematical Principles

Stacking-based aggregation avoids cross-terms and supports arbitrary per-client ranks via direct blockwise concatenation. Given local adapters $W\in \mathbb{R}^{m\times n}$5, construct

• $W\in \mathbb{R}^{m\times n}$6 (vertical stack of $W\in \mathbb{R}^{m\times n}$7)
• $W\in \mathbb{R}^{m\times n}$8 (horizontal stack of $W\in \mathbb{R}^{m\times n}$9)

The global update is then

$\Delta W_k= B_kA_k$0

which, due to block-disjoint structure, reduces precisely to $\Delta W_k= B_kA_k$1. Weighting of client contributions is handled by scaling matrices before stacking, i.e., $\Delta W_k= B_kA_k$2 yields $\Delta W_k= B_kA_k$3. No zero-padding or block-diagonal encoding is necessary, and heterogeneous adapter ranks are seamlessly accommodated.

3. Aggregation Algorithm and Workflow

A single FLoRA round is characterized by the following protocol:

Server:

1. Broadcasts frozen global $\Delta W_k= B_kA_k$4 to clients.
2. Receives $\Delta W_k= B_kA_k$5 from each client $\Delta W_k= B_kA_k$6.
3. Forms $\Delta W_k= B_kA_k$7 (vertically concatenated, with client scaling $\Delta W_k= B_kA_k$8) and $\Delta W_k= B_kA_k$9 (concatenated horizontally).
4. Computes $A_k\in \mathbb{R}^{r_k\times n}$0.
5. Distributes $A_k\in \mathbb{R}^{r_k\times n}$1 to all clients for update integration.

Client $A_k\in \mathbb{R}^{r_k\times n}$2:

1. Initializes LoRA module $A_k\in \mathbb{R}^{r_k\times n}$3 with local $A_k\in \mathbb{R}^{r_k\times n}$4.
2. Fine-tunes locally for $A_k\in \mathbb{R}^{r_k\times n}$5 epochs (keeping $A_k\in \mathbb{R}^{r_k\times n}$6 frozen).
3. Sends $A_k\in \mathbb{R}^{r_k\times n}$7 to server, awaits $A_k\in \mathbb{R}^{r_k\times n}$8.
4. Updates local model by adding $A_k\in \mathbb{R}^{r_k\times n}$9 to $B_k\in \mathbb{R}^{m\times r_k}$0.

This workflow is preserved across rounds, supports arbitrary client configuration, and is communication- and computation-efficient (Wang et al., 2024).

4. Theoretical Properties and Correctness

The stacking method’s correctness follows from linearity and the mutual orthogonality of block partitions in $B_k\in \mathbb{R}^{m\times r_k}$1 and $B_k\in \mathbb{R}^{m\times r_k}$2. Specifically:

• Each $B_k\in \mathbb{R}^{m\times r_k}$3 occupies distinct row ranges, $B_k\in \mathbb{R}^{m\times r_k}$4 distinct column ranges; off-diagonal products vanish.
• Weighted stacking ($B_k\in \mathbb{R}^{m\times r_k}$5 scaling) produces exactly the intended $B_k\in \mathbb{R}^{m\times r_k}$6 aggregation.
• No quadratic terms ($B_k\in \mathbb{R}^{m\times r_k}$7) or aggregation noise from cross-terms appears.
• No information from any client is lost, and each is embedded in a unique submatrix.

The block-matrix view formalizes that $B_k\in \mathbb{R}^{m\times r_k}$8 sums only the correct local updates, bypassing constraints and inaccuracies inherent in previous federated LoRA aggregation schemes.

5. Empirical Evaluation and Results

Experiments on MMLU (QA), MT-bench (chat), and standard LLM backbones (TinyLlama-1.1B, Llama-7B) demonstrate that stacking-based FLoRA outperforms baseline FedIT (FedAvg-LoRA) in both homogeneous and heterogeneous rank configurations:

• On MMLU-Dolly with TinyLlama-1.1B: FedIT achieves $B_k\in \mathbb{R}^{m\times r_k}$9, FLoRA reaches $r_k\ll \min(m,n)$0.
• On TinyLlama (MT-bench): FedIT $r_k\ll \min(m,n)$1, FLoRA $r_k\ll \min(m,n)$2.
• Llama-7B shows consistent improvements of $r_k\ll \min(m,n)$3–$r_k\ll \min(m,n)$4 points.
• Heterogeneous ranks ([64,32,16,8,4,...]): FedIT with zero-padding degrades (MMLU-Alpaca $r_k\ll \min(m,n)$5), while FLoRA maintains high performance (high-20’s to low-30’s on MMLU, $r_k\ll \min(m,n)$6–$r_k\ll \min(m,n)$7 on MT-bench).
• FLoRA+AdaLoRA demonstrates further reduction of total rank budget (from $r_k\ll \min(m,n)$8) with negligible accuracy loss.
• Scaling factor $r_k\ll \min(m,n)$9 has no universal optimum; optimality is dataset- and model-dependent ($\Delta W_k$0–$\Delta W_k$1 explored) (Wang et al., 2024).

Global models consistently outperform any constituent local model across all ablation studies, and in some tasks stacking-based aggregation even slightly outperforms centralized LoRA, plausibly due to decreased overfitting from better-regularized aggregation.

6. Communication, Computation, and Scalability Considerations

FLoRA's stacking-based setup imposes only marginal overhead:

• Each round transmits $\Delta W_k$2 elements, a fraction of $\Delta W_k$3 for full model transfer.
• Over three rounds, FLoRA sends $\Delta W_k$4–$\Delta W_k$5 fewer bytes than full fine-tuning, only $\Delta W_k$6–$\Delta W_k$7 more than FedIT.
• The stacking operation is $\Delta W_k$8 and negligible in the context of LLM computation.
• FLoRA scales to $\Delta W_k$9 clients and arbitrary $\Delta W$0 values without modification, and is compatible with secure aggregation, encryption, and differential privacy, as only adapters are transmitted (Wang et al., 2024).

7. Stacking-Based Aggregation in Federated Hyperparameter Optimization

A parallel application of stacking-based aggregation appears in "FLoRA: Single-shot Hyper-parameter Optimization for Federated Learning" (Zhou et al., 2021). Here, the stacking construct is used for surrogate loss surface aggregation in federated HPO:

• Each client $\Delta W$1 locally fits a regressor $\Delta W$2 (e.g., random forest, GP) to observed $\Delta W$3 pairs.
• The aggregator combines these via four possible strategies, one of which—APLM ("average of per-client models")—is a stacking-style ensemble: $\Delta W$4.
• The aggregated surrogate guides a global hyperparameter choice in a single communication round, minimizing overhead and achieving low regret and robust performance as $\Delta W$5 grows.
• Empirical results on gradient-boosted trees and neural networks validate stacking’s effectiveness and communication efficiency in federated HPO (Zhou et al., 2021).

Summary

Stacking-based aggregation, as developed in FLoRA, constitutes a mathematically rigorous and resource-aware solution to federated aggregation of heterogeneous low-rank adaptations and local surrogate models. By precisely partitioning and summing blockwise contributions, stacking eliminates aggregation noise, enables flexible per-client participation, and achieves superior communication and computational efficiency. Its principles are central both to federated LLM fine-tuning with LoRA adapters (Wang et al., 2024) and to efficient single-shot federated HPO via ensemble surrogates (Zhou et al., 2021), marking a significant advancement in scalable and heterogeneous federated learning.