FLoKD: Adaptive Knowledge Distillation for Federated Low-Rank LLM over Wireless Networks

I. INTRODUCTION

• Full-model exchange is impractical over constrained wireless links, while LoRA transmission still scales with model size and typically treats transformer blocks uniformly.Logit-based distillation also incurs communication costs tied to sequence length and vocabulary size, and captures only final prediction distributions.Source: §I · pp. 1–2
• FLoKD transmits intermediate LoRA activations instead of model parameters or output logits, using their richer internal representations for knowledge transfer.The activations are described as smaller than output logits, reducing communication cost while retaining transformer-block representations.Source: §I · p. 2
• Block importance scoring selects informative transformer blocks, while token-aware and loss-based strategies filter or prioritize public samples for distillation.The introduction presents token-aware selection as filtering distributionally misaligned samples and loss-based selection as prioritizing samples with high model loss.Source: §I · p. 2
• The paper frames block count and public-dataset size as coupled choices under a fixed communication budget, and characterizes their Pareto trade-off.The stated goal is to guide joint configuration of block and data selection for a given wireless bandwidth constraint.Source: §I · p. 2
• Experiments on WikiText-103, Penn Treebank, and DailyDialog report competitive perplexity under tight uplink budgets compared with parameter- and logit-based baselines.The introduction also notes ablations of block importance scoring and sample selection.Source: §I · p. 2

II. SYSTEM MODEL

• The system model has a server and distributed clients with a shared frozen backbone, client-specific LoRA parameters, and a shared public dataset; private client datasets are non-IID.Clients exchange LoRA activations computed on public data rather than model parameters.Source: §II · p. 3

A. Local Fine-Tuning and Upload

• Each client freezes the shared pretrained weight matrix and adds trainable low-rank matrices A_n and B_n, whose rank is much smaller than the layer dimensions.Local training updates these matrices on the client's private data.Source: A. Local Fine-Tuning and Upload · p. 3
• The local objective uses average per-token negative log-likelihood and may include a weighted regularizer to limit overfitting to limited local data.Clients perform several local iterations before interacting with the server.Source: A. Local Fine-Tuning and Upload · p. 4

B. Sample selection and Uplink Transmission

• For uplink transmission, clients compute LoRA activations on selected public samples and selected transformer blocks, then stack the activations into block-wise matrices and an uploaded tensor.The selections yield P public samples and K blocks per client; only those activations are sent.Source: B. Sample selection and Uplink… · p. 4

C. Global Aggregation and Knowledge Distillation

• The server aligns uploads by block index and aggregates each block only over clients that selected it, avoiding dilution from treating missing activations as zeros.The contributing-client set and aggregation weights are defined separately for each block.Source: C. Global Aggregation and Knowledge… · p. 4
• Cosine distance measures the directional discrepancy between global and aggregated client activations, focusing on representation direction despite scale differences.The paper also notes that cosine distance is bounded and that per-sample averaging gives selected samples equal weight rather than implicitly emphasizing larger-norm activations.Source: C. Global Aggregation and Knowledge… · pp. 4–5
• Server-side distillation updates global LoRA parameters while keeping the pretrained backbone frozen; unselected blocks remain unchanged for that round.After distillation, the server broadcasts the updated global LoRA matrices and selected sample indices to clients.Source: C. Global Aggregation and Knowledge… · p. 5

D. Local Update on Clients

• After server-side distillation, clients replace selected local LoRA parameters with the updated global parameters, keep the remaining blocks unchanged, and resume private-data fine-tuning.The passage states that direct downlink transmission is affordable in the considered wireless setting and requires no additional local optimization.Source: D. Local Update on Clients · p. 5

A. Problem Formulation

• The framework transmits intermediate LoRA activations on public data instead of parameters or logits, pooling token representations into one r-dimensional descriptor per sample and block.Mean pooling reduces the per-sample, per-block payload from T_i r to r, independent of sequence length T_i.Source: A. Problem Formulation · p. 5
• The uplink budget jointly constrains the number of selected blocks and samples, while their identities determine how much useful knowledge the transmitted activations carry.The proposed mechanisms select informative blocks and public samples, reducing both factors because the per-round payload scales with their product.Source: A. Problem Formulation · pp. 5–6

B. Block Importance Scoring

• Uploading all block activations can waste communication and let weakly informative signals dilute aggregation, degrading the distilled global model.The passage attributes this risk to blocks that respond weakly to the fine-tuning signal.Source: B. Block Importance Scoring · p. 6
• The importance score estimates first-order loss change from removing a block’s LoRA activation, using the gradient–activation inner product averaged over calibration samples.The criterion treats blocks whose removal would barely change the loss as less important for the current objective.Source: B. Block Importance Scoring · p. 6
• The score uses gradient–activation alignment because a large gradient and activation can contribute little when they are nearly orthogonal.The product of their norms would discard this alignment factor and over-rank such a block.Source: B. Block Importance Scoring · p. 6
• Selecting the Top-K blocks reduces the activation payload to K/L of full-block transmission while discarding the remaining blocks.The proposed score ranks blocks, and only the selected blocks’ activations are uploaded.Source: B. Block Importance Scoring · pp. 6–7

C. Public Data Selection

• Token-aware selection aligns public samples with the aggregate token distribution of clients to determine what types of data should be covered.It can help select samples that are distribution-representative for distillation.Source: C. Public Data Selection · p. 7
• Loss-based selection ranks public samples by their loss under the server model to identify samples more informative for transmission.The two criteria address different aspects of sample usefulness.Source: C. Public Data Selection · p. 7
• The criteria can be applied independently or combined, such as ranking by loss within a token-aligned candidate set.That combination retains samples that are both distribution-representative and informative for distillation.Source: C. Public Data Selection · p. 7

1) Token-Aware Data Selection:

• Non-IID client token distributions motivate selecting public data using collective client token-level statistics rather than relying on a fixed, uniformly sampled pool.The passage states that a mismatch may yield activations with little client-specific signal and an unrepresentative distillation target.Source: 1) Token-Aware Data Selection · p. 7
• Clients upload only the indices of their Top-M most frequent tokens, which the server aggregates into a token-importance profile.This avoids transmitting the full token-frequency vector.Source: 1) Token-Aware Data Selection · p. 7
• The server retains Top-P public samples by cosine similarity between each sample’s normalized token-frequency vector and the aggregated profile, using only a small index-only client upload.The selected samples better match the clients’ shared token characteristics.Source: 1) Token-Aware Data Selection · p. 7

2) Loss-Based Data Selection:

• Higher-loss public samples may expose regions the current model captures poorly, while confidently predicted examples offer less corrective information,.The server ranks samples using only public data and the global model, without additional client-side computation or communication.Source: 2) Loss-Based Data Selection · p. 8
• The server computes each sample’s length-normalised predictive loss, making scores comparable across sequences of different lengths.A large score indicates that the global model explains the sample poorly.Source: 2) Loss-Based Data Selection · p. 8
• At each communication round, the server retains the Top-P samples by descending loss and updates the selected subset as the global model evolves.Source: 2) Loss-Based Data Selection · p. 8

D. Communication Analysis

• Mean pooling collapses the sequence axis, making the compact activation upload independent of the public-data sequence length.Per-round uplink cost depends on the selected sample count P, block count K, and activation dimension r.Source: D. Communication Analysis · p. 8
• The two selection mechanisms reduce distinct payload dimensions, yielding a joint reduction factor of (P/|Dp|)·(K/L); either remains effective alone.Token-aware selection reduces samples from |Dp| to P, while block scoring limits uploads to K of L blocks; token indices add negligible one-off overhead.Source: D. Communication Analysis · p. 8
• The relative communication saving grows with model depth because the fixed activation budget covers a smaller fraction K/L of deeper models’ blocks.The retained blocks are the most informative ones, according to the passage.Source: D. Communication Analysis · p. 8

E. Pareto-Optimal Block-Data Allocation under a Bandwidth Budget

• The loss model combines omitted block signal, a per-block redundancy penalty, and diminishing returns from adding distribution-relevant samples.It is written as L(K, P) = L0 + c1K^-β + νK + c2P^-γ, with the block and sample gains modeled as diminishing.Source: E. Pareto-Optimal Block-Data… · pp. 8–9
• Because additional samples reduce modeled loss, the optimum uses the full activation budget at PK = Q, coupling the numbers of transmitted blocks and samples.The budget is measured in pooled-vector units; each vector costs ρr bits.Source: E. Pareto-Optimal Block-Data… · pp. 8–9
• Even without a bandwidth constraint, retaining every block is not optimal: blocks are kept only while marginal gain exceeds noise, up to a budget-independent bound K0.At finite budgets, K⋆ is below K0; as the budget loosens, K⋆ approaches K0 without attaining it. Model parameters are fitted by least squares on a small calibration grid of (K, P) pairs.Source: E. Pareto-Optimal Block-Data… · p. 9
• Substituting P = Q/K yields a one-variable objective whose stationarity condition balances truncation loss, low-importance-block noise, and displaced sample coverage.The continuous relaxation has a unique positive global minimizer K⋆.Source: E. Pareto-Optimal Block-Data… · p. 9
• For the discrete implementation, each integer K from 1 to L is evaluated and the block-sample pair with minimum loss is selected.Source: E. Pareto-Optimal Block-Data… · p. 9
• Sweeping the budget traces the Pareto frontier: the optimal split minimizes loss at each cost, and K⋆ increases toward K0 as the budget grows.Cost is C(K, P) = ρrPK; the passage states that the frontier is identical whether loss or its monotone transform, perplexity, is used.Source: E. Pareto-Optimal Block-Data… · p. 9

• Experiments use GPT-2 on WikiText-103, PTB, and DailyDialog with 30 non-IID clients and a shared 2,000-sample public dataset.Source: §IV · p. 9
• FLoKD achieves substantially lower PPL than all baselines across bandwidth settings and reduces communication cost by more than 50% versus FedPET at all three target PPL levels.Source: Figure 2 · p. 10, §IV · p. 11, Figure 4 · p. 11
• Under unlimited bandwidth, FLoKD settles near 25 PPL, versus FedPET around 27 and FedMKT and FedSA around 38.Source: Figure 3 · p. 11, §IV · p. 11
• Top-K block selection achieves lower server PPL than Random-K and Bottom-K across selection sizes, with the best performance at a moderate number of blocks.Source: Figure 5 · p. 11, §IV · p. 12
• 38.97 PPL for magnitude-based selection at 40% and 36.56 at 50% contrast with token-aware's 32.57 at 40% and loss-based's 30.55 at 90%.Source: Figure 6 · p. 12, §IV · p. 12
• At 0.50 MB, K=3 reaches 24.21 PPL, while K=12 uses 1.50 MB for 27.54 PPL; under tighter budgets, K=1 is optimal.Source: Figure 7 · p. 12, §IV · p. 12

V. CONCLUSION

• Communication overhead falls by 50–65% compared with baselines, while FLoKD achieves fast convergence and better PPL across language-generation tasks and multiple datasets.The conclusion reports these results for the proposed framework overall.Source: §V · p. 13
• FLoKD transmits intermediate LoRA activations as its distillation signal instead of full model parameters or output logits.Source: §V · p. 13
• A Transformer block-importance scoring framework selects informative block activations, and token-aware selection chooses the shared public dataset.Source: §V · p. 13