Source-linked AI summary
Robust Federated Learning: The Case of Affine Distribution Shifts
Amirhossein Reisizadeh, Farzan Farnia, Ramtin Pedarsani, Ali Jadbabaie
TL;DR
The paper addresses statistical heterogeneity and distribution shifts across federated users while data remains localized. It proposes FLRA with FedRobust optimization and theoretical guarantees, and reports improved robustness against affine shifts and adversarial attacks compared with standard approaches.
Problem
Federated learning must generalize across users whose localized training samples exhibit heterogeneous, device-dependent distribution shifts.
Method
FLRA models user data with affine shifts and solves the resulting distributed minimax problem using the FedRobust gradient descent ascent method.
Results
FedRobust achieves the best performance in most affine-shift evaluations, shields against PGD attacks, and outperforms FedAvg.
Takeaways & Limitations
Affine shifts provide a structured model of federated heterogeneity, and the proposed framework is empirically robust across affine and standard adversarial perturbations.
Takeaways & Limitations
The local affine variables are coupled through global-model averaging, creating a technical challenge for the distributed algorithm.
Abstract
from arXiv · showhide
Federated learning is a distributed paradigm that aims at training models using samples distributed across multiple users in a network while keeping the samples on users' devices with the aim of efficiency and protecting users privacy. In such settings, the training data is often statistically heterogeneous and manifests various distribution shifts across users, which degrades the performance of the learnt model. The primary goal of this paper is to develop a robust federated learning algorithm that achieves satisfactory performance against distribution shifts in users' samples. To achieve this goal, we first consider a structured affine distribution shift in users' data that captures the device-dependent data heterogeneity in federated settings. This perturbation model is applicable to various federated learning problems such as image classification where the images undergo device-dependent imperfections, e.g. different intensity, contrast, and brightness. To address affine distribution shifts across users, we propose a Federated Learning framework Robust to Affine distribution shifts (FLRA) that is provably robust against affine Wasserstein shifts to the distribution of observed samples. To solve the FLRA's distributed minimax problem, we propose a fast and efficient optimization method and provide convergence guarantees via a gradient Descent Ascent (GDA) method. We further prove generalization error bounds for the learnt classifier to show proper generalization from empirical distribution of samples to the true underlying distribution. We perform several numerical experiments to empirically support FLRA. We show that an affine distribution shift indeed suffices to significantly decrease the performance of the learnt classifier in a new test user, and our proposed algorithm achieves a significant gain in comparison to standard federated learning and adversarial training methods.
1 Introduction
Federated learning must handle device-dependent statistical heterogeneity while preserving localized data and managing communication constraints. This paper proposes FLRA and FedRobust to improve robustness to affine distribution shifts, with optimization, generalization, and empirical guarantees.
- Motivation: Federated learning preserves data locality but faces statistical heterogeneity, communication bottlenecks, and privacy-related algorithmic challenges.Non-i.i.d. device data can impair accuracy or cause standard methods such as FedAvg to diverge, especially when generalizing to unseen test users.
- Affine distribution shifts: FLRA models each device’s samples as shifted from a common distribution by a device-dependent affine transformation xi → Λixi + δi.The model captures shared distortions within a device while allowing different transformations across devices, including image changes in intensity, contrast, brightness, and blurring.
- Proposed framework: FLRA formulates robust learning as a minimax problem that minimizes cumulative loss under worst-case local affine transformations.The framework is designed to provide robustness beyond adversarial training models that use solely additive perturbations.
- Optimization: FedRobust uses gradient descent ascent to solve FLRA’s distributed minimax problem while avoiding full local adversarial optimization and excessive communication.The method is motivated by limited computation and communication resources at federated devices.
- Guarantees and evaluation: The paper characterizes FLRA’s generalization and robustness properties and reports empirical advantages over standard approaches.The analysis includes convergence guarantees, PAC-Bayes generalization bounds, and numerical demonstrations of efficiency and comparative performance.
2 Federated Learning Scenario
The federated learning scenario consists of devices holding local samples from heterogeneous distributions and collaborating through a central server. The paper models this heterogeneity with device-specific affine transformations and formulates robust training as a coupled minimax problem.
- Federated setting: A federated learning system contains n devices connected to a server, with each device holding m local training samples.The classical objective fits a global model to the combined nm samples through empirical risk minimization.
- Data model: Each device’s data distribution is modeled as an affine transformation of a common underlying distribution, capturing non-identically distributed local samples.The transformation acts on feature vectors and represents device-dependent shifts from shared data-generating structure.
- Data model: All samples on one device share the same affine transformation, while different devices may experience different transformations.This structure represents identical yet device-dependent distortions within each device and varying distortions across devices.
- Robust objective: The robust federated objective uses local affine transformations to induce high losses and an outer optimization to find the global model with the smallest cumulative loss.The minimax formulation therefore couples n inner local maximization problems with one outer global minimization problem.
3 The Proposed FedRobust Algorithm
FLRA replaces the original constrained maximization with a worst-case affine-shift loss and solves the resulting minimax problem using FedRobust. The algorithm alternates local updates with periodic server aggregation, reducing communication while introducing coupled affine variables that complicate convergence analysis.
- FLRA formulation: FLRA penalizes affine transformations by their squared distance from the identity and maximizes each node’s resulting loss before minimizing the averaged loss over the global model.This yields the paper’s minimax robust optimization formulation.
- FLRA formulation: FLRA is the paper’s Federated Learning framework Robust to Affine distribution shift.
- FedRobust updates: FedRobust solves FLRA using stochastic gradient descent ascent, with local ascent updates for affine parameters and descent updates for classifier parameters.The method is designed to be computationally and communication-wise efficient.
- FedRobust updates: After τ local iterations, devices upload local classifiers, the server averages them, and the resulting global model is sent back as the next initialization.The algorithm input includes initial local models and affine parameters, step sizes, τ, and the total iteration count T.
- Efficiency: Periodic communication every τ iterations reduces communication load by a factor of τ, while each node updates its perturbation parameters only once per iteration.
- Optimization challenge: The local affine variables remain on their devices but become coupled through the averaged global model, creating a technical challenge for FedRobust’s optimization guarantees.Updating one node’s affine variables affects other nodes in subsequent iterations through model averaging.
4 Theoretical Guarantees: Optimization, Generalization and Robustness
FLRA provides theoretical guarantees for optimization, generalization, and robustness under affine distribution shifts. The analysis covers convergence of FedRobust, margin-based generalization, and a Wasserstein robustness interpretation.
- FLRA’s guarantees cover FedRobust convergence, generalization to unseen samples, and robustness to Wasserstein shifts structured across nodes.
- Optimization guarantees: FedRobust finds saddle points for PL-PL and nonconvex-PL loss classes under stated assumptions and feasible step sizes.The nonconvex-PL result establishes convergence toward a stationary saddle point.
- Assumptions: The analysis assumes bounded non-iid degree, unbiased stochastic gradients, Lipschitz gradients, and PL-type conditions on the objective.Affine transformations are also constrained through bounds involving ε_1 and ε_2.
- Optimization guarantees: The PL-PL convergence bound decreases geometrically up to additive terms involving step sizes, heterogeneity, stochastic variance, and affine-shift parameters.The bound is expressed through the potential function P_t = a_t + βb_t, whose convergence toward zero indicates approach to a minimax point.
- Generalization guarantees: A PAC-Bayes margin bound controls the gap between empirical and underlying FLRA error for spectrally regularized deep neural networks.The bound depends on spectral norms and requires smooth, Lipschitz activations; the experiments also consider ReLU.
- Distributional robustness: Theorem 4 shows that FLRA’s minimax approach optimizes a lower bound on a distributionally robust objective based on Wasserstein cost.The Wasserstein cost measures the minimum expected transport cost between the original and perturbed distributions.
5 Numerical Results
The experiments evaluate FedRobust on CIFAR-10 and MNIST under affine distribution shifts and PGD perturbations, comparing it with FedAvg and adversarial-training baselines. FedRobust generally provides stronger affine-shift robustness while also achieving comparable adversarial robustness and faster training.
- Experimental setup: FedRobust was evaluated on MNIST and CIFAR-10 using AlexNet, Inception-Net, and mini-ResNet classifiers.The implementation used TensorFlow and tested robustness against both affine distribution shifts and adversarial perturbations.
- Experimental setup: The experiments simulated 10 federated nodes, each observing 5,000 training samples, with separate validation and test sets.CIFAR-10 inputs were normalized and mapped to [-1,1], with batch normalization and ADAM optimization.
- Experimental setup: Affine shifts were generated by applying node-specific random linear and translational transformations to training samples, while evaluation included no shift, adversarial affine shifts, and PGD perturbations.The affine transformations used Gaussian random matrices and vectors with σ = 0.01; adversarial evaluation used 100 projected-gradient steps.
- CIFAR-10 results: FedRobust outperformed baseline methods in most CIFAR-10 affine-shift experiments, improving over FedAvg by as much as 54%.It also exceeded distributed FGM and PGD training, suggesting that ordinary adversarial perturbations may not capture affine-shift complexity.
- MNIST results: On MNIST, FedRobust achieved the best performance in most affine-shift evaluations and comparable performance to PGD and FGM training under PGD attacks.The affine-shift evaluations varied either the maximum linear-shift norm or the maximum translation norm.
Appendix B Proof of Theorem 1
The proof of Theorem 1 combines contraction bounds for optimization error, local-model deviation, and auxiliary sequences. Under stated smoothness, PL, and step-size conditions, these bounds establish convergence of the FedRobust iterations.
- Proof strategy: The proof analyzes a contraction sequence combining the expected global objective gap with a weighted auxiliary error term.The sequence is formed using a quantity at = E[Φ(wt)] − Φ* and a weighted bt term.
- Supporting bounds: The proof controls local-model deviation through a bound on et and incorporates stochastic-gradient variance and client heterogeneity terms.The resulting bounds depend on step sizes, the number of local steps, variance, and the heterogeneity parameter.
Appendix C Proof of Theorem 2
The proof of Theorem 2 derives a finite-iteration convergence bound by combining descent and auxiliary-sequence inequalities. It then applies the required smoothness, PL, and step-size assumptions to conclude the theorem.
- Proof strategy: The proof begins by combining bounds on the change in the global objective with a contraction inequality for the auxiliary sequence bt.The argument tracks the expected gradient norm together with auxiliary error terms.
- Averaging argument: Summing the per-iteration inequalities over T iterations and dividing by T converts the recursive bounds into an average convergence statement.Subsequent lemmas replace the remaining auxiliary terms in this averaged inequality.
- Conclusion: The proof uses the step-size condition η1( L-hat + 40 L-tilde(τ − 1)^2 ) ≤ 1 to simplify the final bound.This condition removes an additional term before the theorem is concluded.
D.1 Proof of Lemma 1
The lemma proof establishes smoothness and bounds the progress and disagreement terms used in FedRobust’s convergence analysis. It combines Lipschitz-gradient assumptions, PL structure, stochastic-gradient variance control, and the algorithm’s descent/ascent updates.
- Smoothness bounds: The proof bounds the global function’s gradient using Lipschitz continuity of the local objectives and the corresponding block-gradient constants.It then uses the update rule to relate averaged-model changes to stochastic gradients.
- Stochastic estimates: Unbiased, variance-bounded stochastic gradients allow the analysis to take expectations and control the stochastic update error.The proof explicitly invokes the variance parameter σ²_w when bounding update differences.
- Auxiliary quantities: The disagreement term ht is bounded by decomposing local gradients around the averaged model and applying Lipschitz and PL inequalities.The resulting bound relates ht to bt, et, and the constants governing the maximization problem.
- Combined bound: The proof combines bounds on objective descent, local-gradient disagreement, and model deviation to obtain the contraction inequality used in Lemma 5.The final step substitutes the intermediate estimates into the preceding recursive inequalities.
D.6 Proof of Lemma 6
The proof bounds local-model deviations and gradient-related errors across communication rounds, then combines these bounds under step-size and communication assumptions to establish Lemma 6.
- The proof decomposes the error into terms controlled by stochastic-gradient variance, Lipschitz gradients, gradient diversity, and local-update duration.
- The resulting bounds are averaged across workers and iterations to complete the proof of Lemma 6.
- Across each communication interval, the error satisfies a recursive bound involving prior errors, gradient terms, and a residual coefficient C3.
- Assuming C1(τ −1) ≤1/2 permits rearrangement of the recursive inequalities into a finite aggregate bound.
- Proposition 2 bounds the diversity of local gradients using Assumptions 1 and 3.
Appendix E Proof of Theorem 3
The proof derives a margin-based generalization bound by controlling parameter perturbations, optimizing auxiliary variables, and applying covering arguments across all nodes.
- The analysis uses Lipschitz properties of the loss and its gradient to control derivatives with respect to the perturbation variables δ and Λ.
- Strong concavity of the maximization problem controls how the optimal auxiliary solutions δw and Λw change when weights are perturbed.
- A norm-bounded weight perturbation is propagated through the loss and network output to bound the resulting change in loss.
- The proof applies a Gaussian perturbation distribution and covering arguments to obtain a bound for feasible norm-bounded weights.
- A union bound extends the fixed-distribution result across the n nodes, yielding an average worst-case margin-based loss bound.
Appendix F Proof of Theorem 4
The proof bounds the optimal transport cost between the original data distribution and its affine transformation by constructing a coupling and using the covariance assumption.
- The transformed random vector is defined as U = ΛX + δ for evaluating the transport cost between PX and PU.
- The joint distribution of (X, ΛX + δ) provides a coupling between the original and transformed distributions.
- The transport-cost bound separates the affine transformation and translation effects using norm inequalities and trace identities.
- The covariance condition E[XX⊺] ≤λI controls the quadratic contribution of the transformed data.