Federated learning (FL) is increasingly used in safety-critical and fairness-sensitive applications, where learning objectives must be optimized subject to functional constraints. This setting is challenging due to costly projection onto the resulting feasible set, stochastic constraint information, partial client participation, and limited communication at edge devices. We propose a federated stochastic switching framework based on primal-only updates, avoiding dual-variable tuning for functionally constrained federated optimization. For convex objectives, the framework further supports compression with error feedback to control compression-induced bias. We establish high-probability convergence and feasibility guarantees under stochastic gradients and partial client participation, and extend the analysis to weakly convex objectives through proximal-stationarity guarantees. Our bounds isolate the effects of stochastic gradients, constraint-evaluation noise, and partial participation, and show that direct switching can retain a nonvanishing constraint-estimation error even with exact local constraint evaluations. To overcome this limitation, we develop variance-reducing constraint estimators that yield vanishing feasibility and stationarity guarantees with fixed per-round client participation and constraint-evaluation batch sizes. Experiments on Neyman-Pearson classification, fairness-constrained learning on the ADULT dataset, and constrained deep reinforcement learning demonstrate the effectiveness of the proposed framework.
Figures & tables
Figure 1: NP classification. Comparison of our method against penalty-based FedAvg and the primal-dual method. Each point represents a result from a separate experimental run with different random seeds. Circles (‘ ∙ ’) represent ours (direct, momentum, and STORM), while crosses (‘ × ’) and triangles (‘ ▲ ’) represent penalty-based ( μ) and primal-dual ( λ0 ) FedAvg, respectively.
v_{k}^{i}=\begin{cases}\widehat{h}_{k}^{i},&\mathtt{comp}=\mathrm{off},\\ \mathcal{Q}_{i}\bigl(r_{k}^{(i)}+\widehat{h}_{k}^{i}\bigr),&\mathtt{comp}=\mathrm{on},\end{cases}\qquad\text{and, if \mathtt{comp}=\mathrm{on}, set }r_{k+1}^{(i)}\leftarrow r_{k}^{(i)}+\widehat{h}_{k}^{i}-v_{k}^{i}.
Algorithm 1 Federated Stochastic Switching with Bi-Directional Compression
Figure 2: NP classification. Evolution of the objective F(wk) (left) and constraint G(wk) (right) over communication rounds. We compare direct, momentum, and STORM switching variants and the centralized setting. The red dashed line marks the constraint threshold.
Figure 3: NP classification. Sensitivity of (a) direct, (b) momentum, and (c) STORM switching. Within each panel, the left and right columns show the objective F(wk) and constraint G(wk) over communication rounds, respectively. From top to bottom, the rows vary the constraint-evaluation batch size Bval , client participation rate m/n , and compression rate s/d . The red dashed line marks the constraint threshold θ=0.05 .
Figure 4: Fair classification. We compare our methods (blue/orange/green solid lines) against the centralized setting (black dashed line), federated penalty-based baseline (purple solid line), and IPP-ConEx (brown dashed line). Our momentum and STORM switching approach achieves the best trade-off: they converge to the lowest loss while satisfying the fairness constraint, matching the results of the full-information centralized setting.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
Value
Hyperparameter
Value
Communication rounds ( K )
2500
Tolerance ( θ )
0.05
Client participation ( mup=mval )
6
Total number of clients ( n )
20
Compression rate ( s/d )
1.0
Classifier dimension ( d )
30
Gradient update batch ( Bg )
8
Constraint value batch ( Bval )
1
Step size (centralized/direct, η )
1.0
Step size (momentum/STORM, η )
1.5
Momentum parameter ( α )
0.35
Number of runs
3
Appendix
Table 1: Detailed setting of NP classification experiment
Hyperparameter
Value
Hyperparameter
Value
Dirichlet parameter
2.0
Tolerance ( θ )
0.05
Communication rounds ( K )
500
Number of runs
3
Neural network dimension ( d )
6,529
Total number of clients ( n )
10
Batch size per client ( Bg=Bval )
128
Client participation ( mup=mval )
5
Learning rate ( η )
0.1
Momentum parameter ( α )
0.2
Penalty parameter ( μ )
10
Dual step (IPP-ConEx)
200
Appendix
Table 2: Detailed setting of fair classification experiment on Adult dataset
Figure 5: Fair classification. Comparison of our method against penalty-based FedAvg and primal-dual method. Each point represents a result from a separate experimental run with different random seeds. Circles (‘ ∙ ’) represent ours (direct, momentum and STORM), while crosses (‘ × ’) and triangles (‘ ▲ ’) represent penalty-based and primal-dual FedAvg, respectively.
Figure 6: Constrained MDP. Evolution of episodic reward (left) and cost (right) on the Cartpole experiment. We compare direct and momentum switching (blue and orange solid lines) against the centralized baseline (black dashed line). The red dashed line marks the safety budget.
Hyperparameter
Value
Hyperparameter
Value
Communication rounds ( K )
500
Tolerance ( θ )
30
Episode length
200
Discount factor ( γ )
1.0
Neural network dimension ( d )
17,410
Total number of clients ( n )
10
Batch size per client ( Bg=Bval )
8
Client participation ( mup=mval )
4
Number of runs
3
Momentum paramter ( α )
0.2
Appendix
Table 3: Detailed setting of CMDP experiment on Cartpole
This paper addresses the distributed stochastic minimax optimization problem subject to stochastic constraints. We propose a novel first-order Softmax-Weighted Switching Gradient method tailored for federated learning. Under full client participation, our algorithm achieves the standard O~(ε−4) oracle complexity to satisfy a unified bound ε for both the optimality gap and feasibility tolerance. We extend our theoretical analysis to the practical partial participation regime by quantifying client sampling noise through a stochastic superiority assumption. Furthermore, by relaxing standard boundedness assumptions on the objective functions, we establish a strictly tighter lower bound for the softmax hyperparameter. We provide a unified error decomposition and establish a sharp O(logδ1) high-probability convergence guarantee. Ultimately, our framework demonstrates that a single-loop primal-only switching mechanism provides a stable alternative for optimizing worst-case client performance, effectively bypassing the hyperparameter sensitivity and convergence oscillations often encountered in traditional primal-dual or penalty-based approaches. We verify the efficacy of our algorithm via experiment on the Neyman-Pearson (NP) classification, fair classification, and federated safe reinforcement learning tasks.
Zhankun Luo, Antesh Upadhyay, Sang Bin Moon +1
School of Electrical and Computer Engineering, Purdue University
We consider Decision-Focused Federated Learning (DFFL), a predict-then-optimize setting in which multiple clients collaboratively train predictive models for downstream linear optimization problems without exchanging raw data. Besides the data heterogeneity typical of standard federated learning, clients may also have different objective functions and feasible regions. Building on the SPO+ surrogate loss, we derive heterogeneity bounds that separate objective shift, measured through cost-vector distances, from feasible-set shift, measured through support-function and shape-distance terms. We show that, for general compact feasible sets, small objective perturbations can still induce nonvanishing decision-focused loss discrepancies, while strongly convex feasible regions yield sharper stability-based bounds. We then lift these pointwise bounds to a local-versus-federated excess-risk comparison, showing that federation is beneficial when the statistical advantage of pooling exceeds a client-specific heterogeneity penalty. Computational experiments on polyhedral and strongly convex problems confirm that federation is substantially more robust under strongly convex feasible regions. Finally, we evaluate a simple validation-based interpolation between local and federated DFFL models. This interpolation mitigates the theoretical tradeoff and reduces aggregate regret and worst-client harm in both synthetic experiments and a PJM energy-pricing case study.
Federated learning (FL) is a communication-efficient framework for solving distributed optimization problems. Standard FL methods aggregate local updates into a single global model that is shared by all agents, effectively imposing consensus. Under heterogeneous local constraints, however, a common feasible model may be overly restrictive or may not exist. Existing constrained FL methods largely retain this shared-model structure and therefore do not directly address personalization under heterogeneous agent-specific feasible sets. We study a constrained personalized FL problem that assigns a distinct feasible model to each agent while coupling the models through a collaborative objective. We propose Locally Penalized Cross-Estimate Federated Averaging (PCE-FedAvg), where each agent maintains a multi-block vector containing its own model and estimates of the other agents' models while applying the feasibility penalty only locally to its own block. The server aggregates corresponding blocks separately, thereby preserving personalization without requiring agents to explicitly disclose their local constraint sets. For any ε>0, we establish finite-time upper and lower suboptimality bounds and agent-wise squared infeasibility bounds, yielding communication complexities of O(ε−2) and O(ε−1), respectively. Experiments on MNIST and CIFAR-10 support our theoretical results.
Mohammadjavad Ebrahimi, Daniel Burbano, Farzad Yousefian
Department of Industrial and Systems Engineering, Rutgers University, USA. · Department of Electrical and Computer Engineering, Rutgers University, USA.