Personalized Federated Learning with Bidirectional Communication Compression via One-Bit Random Sketching

AI summary

Overview

Research area: Federated learning, specifically the intersection of personalized federated learning (PFL) and communication-efficient federated learning (CEFL), with supporting theory from randomized numerical linear algebra and non-convex optimization.

Technical level: Advanced. The paper combines a bilevel optimization formulation, subgradient-based convergence analysis, and structured random sketching (Subsampled Randomized Hadamard Transform), so it assumes familiarity with federated optimization and convergence proofs.

Scope: The paper proposes, analyzes, and empirically evaluates pFed1BS, a framework in which clients upload one-bit random sketches of their local models and the server broadcasts a one-bit global consensus vector, targeting extreme bidirectional communication compression under non-i.i.d. client data.

What This Paper Is About

Standard federated learning transmits full-precision, high-dimensional model parameters in both directions, which is infeasible on bandwidth-limited devices, and it degrades when client data is non-i.i.d. Existing one-bit compression methods (OBDA, OBCSAA, zSignFed) achieve extreme compression but are designed to train a single global model, so they do not address data heterogeneity, while PFL methods (pFedMe, Ditto, FedRep, DisPFL) personalize models but still exchange full-precision vectors. The paper's goal is to design a single framework that simultaneously achieves extreme bidirectional one-bit communication and native per-client personalization.

Key Contributions

1. A joint bilevel formulation of personalization with one-bit bidirectional communication. The authors define a client-side objective (local empirical risk + a sign-based regularizer + an ℓ2 penalty) and a server-side objective for aggregating the compressed client signals, and they state this is the first such formulation.

2. An efficient matrix-free sketching implementation. By replacing dense Gaussian projection with a structured projection based on the Subsampled Randomized Hadamard Transform and the Fast Hadamard Transform, the client-side projection complexity drops from 𝒪(mn) to 𝒪(n log n) without performance degradation.

3. A convergence analysis of the alternating optimization scheme. The paper proves pFed1BS converges to a stationary neighborhood of a global potential function, with error terms that explicitly account for personalization, local stochastic updates, one-bit sketching, server aggregation, and partial client participation.

4. An empirical study across five benchmarks under non-i.i.d. partitioning. Experiments on MNIST, FMNIST, CIFAR-10, CIFAR-100, and SVHN compare pFed1BS against FedAvg, OBDA, OBCSAA, zSignFed, EDEN, and FedBAT in terms of top-1 accuracy and one-round communication cost.

Main Findings

• Bidirectional one-bit compression is achievable together with personalization. Table 1 positions pFed1BS as the only compared algorithm marked as having upload dimension reduction, upload one-bit quantization, download dimension reduction, download one-bit quantization, and personalization capability; FedAvg has none of these, OBDA and zSignFed have only upload one-bit quantization, and OBCSAA has upload dimension reduction plus upload one-bit quantization.

• The server aggregation step has an exact closed-form solution. Lemma 1 states that the minimizer of the server objective is a weighted majority vote, v* = sign(Σ_{k∈𝒮ᵗ} p_k z_k^{t+1}), so the aggregation is optimal given available information rather than heuristic.

• Communication cost is reduced by roughly 99.68–99.69% relative to FedAvg on most benchmarks. The reported one-round cost of pFed1BS is 0.10 MB on MNIST and FMNIST, 0.13 MB on CIFAR-10 and SVHN, and 7.30 MB on CIFAR-100, compared with 31.06 MB (MNIST, FMNIST), 42.85 MB (CIFAR-10, SVHN), and 2335.85 MB (CIFAR-100) for FedAvg.

• pFed1BS exceeds the one-bit CEFL baselines on every dataset while using far less bandwidth. On MNIST it reports 97.83 ± 0.02% at 0.10 MB versus OBDA 92.54 ± 0.32% at 0.97 MB, OBCSAA 92.20 ± 0.20% at 15.58 MB, and zSignFed 94.83 ± 0.07% at 16.01 MB; on CIFAR-100 it reports 52.88 ± 0.32% at 7.30 MB versus OBDA 42.47 ± 2.02% at 72.95 MB, OBCSAA 48.99 ± 0.54% at 1171.57 MB, and zSignFed 40.17 ± 2.32% at 1203.78 MB.

• Against full-precision FedAvg the accuracy picture is mixed. pFed1BS is higher on MNIST (97.83% vs 97.21%) but lower on FMNIST (84.15% vs 84.40%), CIFAR-10 (85.21% vs 87.78%), CIFAR-100 (52.88% vs 59.60%), and SVHN (95.07% vs 96.33%), while using a fraction of the communication.

• pFed1BS also exceeds the state-of-the-art PFL/communication-efficient baselines shown. It reports higher accuracy than EDEN on MNIST (97.83% vs 96.50%), CIFAR-10 (85.21% vs 84.91%), CIFAR-100 (52.88% vs 47.55%), and SVHN (95.07% vs 89.01%), and higher than FedBAT on the same four (96.42%, 81.20%, 46.89%, 88.89% respectively), while FMNIST is the exception (84.15% for pFed1BS vs 83.85% for EDEN and 83.70% for FedBAT, so pFed1BS is higher there as well; FedAvg remains highest on FMNIST).

• The convergence rate is 𝒪(1/(RT)) to a neighborhood. Theorem 1 bounds the time-averaged expected squared gradient norm; the neighborhood is governed by stochastic noise 𝒪(ηL_Fσ²), communication error 𝒪(Δ_max/(ηR)), and client sampling error 𝒪(λE_S/(ηR)), where Δ_max = 2λ(√m W C_Φ + m) and E_S = 2m√(K(K−S)/(S(K−1)) Σ p_k²).

• The regularization weight must scale inversely with model dimension. Remark 1 states that λ must satisfy λ = 𝒪(1/n), which simultaneously controls the smoothness constant L_F = L + λγC_Φ² + μ, the communication error term Δ_max, and the sampling error term E_S.

• Partial participation degrades the bound, and full participation recovers it. Remark 2 notes that the client sampling error E_S vanishes when S = K, in which case the bound reduces to the full-participation result.

• Sketching complexity is near-linear. The structured projection pads the model to the next power of two (n′ = 2^⌈log₂ n⌉), applies random sign flips D, the normalized Walsh-Hadamard transform H, and subsampling S′ = √(n′/m) S; both the forward projection and its adjoint are matrix-free at 𝒪(n log n), with spectral norm bound ‖Φ‖ ≤ √(n′/m), denoted C_Φ = 𝒪(√(n/m)).

Methodology in Plain English

Each client holds its own model. Instead of sending that model to the server, the client multiplies it by a fixed random projection matrix and keeps only the signs of the resulting low-dimensional vector, so only one bit per projected coordinate is transmitted. The server collects these sign vectors from a randomly sampled subset of clients and forms a single global sign vector by taking a weighted majority vote (the paper proves this vote is the exact optimum of the server's objective). That one-bit consensus vector is broadcast back to all clients.

During local training, each client minimizes its normal training loss plus two extra terms: a penalty that punishes disagreement between the signs of its own projected model and the broadcast consensus vector, and a small penalty on the size of the model parameters to keep it from growing without bound. Because the sign penalty involves an ℓ1 norm that is not differentiable, the authors replace it with a smooth approximation based on log(cosh(·)), whose gradient turns out to be a clean expression involving tanh, which behaves like the sign function as the smoothing parameter grows. This leads to an alternating scheme: clients take R local stochastic gradient steps, then the server updates the consensus vector.

To avoid the cost of explicitly building a large random matrix, the projection is implemented as a sequence of fast operations: sign flips, a Walsh-Hadamard transform, and random subsampling. The backward pass (applying the transpose of the projection) uses the same components in reverse. The theoretical part establishes standard smoothness and variance assumptions, bounds the model norm under a learning-rate condition, and derives the convergence guarantee for the resulting bilevel algorithm, including the effects of partial client sampling.

Why This Matters

Impact on research. The paper argues there is a gap in the literature: no prior framework combined extreme bidirectional compression with native personalization, and Table 1 makes that gap explicit. It also shows that a structured random projection can make a signal-processing-style sketching scheme practical for large models, and it provides a convergence analysis for a bilevel scheme in which the upper-level variable is a discrete sign vector. The result that the regularization weight must scale as 𝒪(1/n) is a concrete constraint that future work on sign-based regularizers will have to confront.

Real-world applications:

• Massive Internet of Things deployments where devices operate under extremely constrained bandwidth.
• Vehicle-to-Everything (V2X) communications, where vehicles exchange model updates over limited links.
• Remote sensing networks, where distributed sensors with heterogeneous local conditions need tailored models.
• Any multi-organization deployment where client data distributions differ and full-precision model exchange is not feasible.

Industry relevance. The reported one-round costs (0.10 MB on MNIST/FMNIST, 0.13 MB on CIFAR-10 and SVHN, 7.30 MB on CIFAR-100) represent reductions of about 99.68–99.69% relative to FedAvg on four of the five benchmarks, which directly affects bandwidth budgeting and the number of clients a server can serve per round. Because compression is bidirectional, both the uplink and the downlink are cheap, which matters for asymmetric or downlink-constrained deployments. The random seed for the projection is shared once at initialization, and both directions are matrix-free, which simplifies implementation on constrained devices.

Future Directions

• The λ = 𝒪(1/n) requirement. Remark 1 ties the convergence error to a regularization weight that shrinks with model dimension. Whether this is compatible with strong personalization on very large models, and what practical schedule to use, is left open.
• Accuracy parity with full-precision FL on harder benchmarks. The largest accuracy gaps against FedAvg appear on CIFAR-100 (52.88% vs 59.60%) and SVHN (95.07% vs 96.33%), so closing the gap on more complex tasks and larger client populations is a natural next step.
• Interaction with other federated mechanisms. The paper does not report how one-bit sketching composes with secure aggregation, differential privacy, over-the-air computation, or heterogeneous client model architectures, and the sketch dimension m and client sampling size S are additional knobs whose optimal settings are not characterized.
• Beyond image classification. Evaluation is limited to MNIST, FMNIST, CIFAR-10, CIFAR-100, and SVHN with a two-layer MLP (MNIST, FMNIST) and VGG architectures (the rest); generalization to sequence models, federated language modeling, or federated regression is not reported.

Target Audience

Researchers and graduate students working on federated learning, distributed optimization, or communication-efficient training who want to understand how one-bit sketching can be combined with personalization and what convergence guarantees such a scheme admits. It is also relevant to practitioners building federated systems on bandwidth-constrained hardware, though they will need to tolerate the advanced theoretical presentation. Readers looking for a gentle introduction to FL should treat this as an advanced follow-up rather than a starting point.

Authors’ abstract

Federated Learning (FL) enables collaborative training across decentralized data, but faces key challenges of bidirectional communication overhead and client-side data heterogeneity. To address communication costs while embracing data heterogeneity, we propose pFed1BS, a novel personalized federated learning framework that achieves extreme communication compression through one-bit random sketching. In personalized FL, the goal shifts from training a single global model to creating tailored models for each client. In our framework, clients transmit highly compressed one-bit sketches, and the server aggregates and broadcasts a global one-bit consensus. To enable effective personalization, we introduce a sign-based regularizer that guides local models to align with the global consensus while preserving local data characteristics. To mitigate the computational burden of random sketching, we employ the Fast Hadamard Transform for efficient projection. Theoretical analysis guarantees that our algorithm converges to a stationary neighborhood of the global potential function. Numerical simulations demonstrate that pFed1BS substantially reduces communication costs while achieving competitive performance compared to advanced communication-efficient FL algorithms.