Variance Reduction for Stochastic Switching in Functionally Constrained Federated Optimization

Antesh Upadhyay; Sang Bin Moon; Zhankun Luo; Abolfazl Hashemi

cs.LGJan 23, 2026

Variance Reduction for Stochastic Switching in Functionally Constrained Federated Optimization

Authors: Antesh Upadhyay, Sang Bin Moon, Zhankun Luo, Abolfazl Hashemi

Organizations: School of Electrical and Computer Engineering, Purdue University, West Lafayette, IN 47907, USA

Abstract

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 (‘ $\bullet$ ’) represent ours (direct, momentum, and STORM), while crosses (‘ $\times$ ’) and triangles (‘ $\blacktriangle$ ’) represent penalty-based ( $\mu)$ and primal-dual ( $\lambda_{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(w_{k})$ (left) and constraint $G(w_{k})$ (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(w_{k})$ and constraint $G(w_{k})$ over communication rounds, respectively. From top to bottom, the rows vary the constraint-evaluation batch size $B_{\mathrm{val}}$ , client participation rate $m/n$ , and compression rate $s/d$ . The red dashed line marks the constraint threshold $\theta=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 ( $\theta$ )	0.05
Client participation ( $m_{\mathrm{up}}=m_{\mathrm{val}}$ )	6	Total number of clients ( $n$ )	20
Compression rate ( $s/d$ )	1.0	Classifier dimension ( $d$ )	30
Gradient update batch ( $B_{\mathrm{g}}$ )	8	Constraint value batch ( $B_{\mathrm{val}}$ )	1
Step size (centralized/direct, $\eta$ )	1.0	Step size (momentum/STORM, $\eta$ )	1.5
Momentum parameter ( $\alpha$ )	0.35	Number of runs	3

Appendix

Table 1: Detailed setting of NP classification experiment

Hyperparameter	Value	Hyperparameter	Value
Dirichlet parameter	2.0	Tolerance ( $\theta$ )	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 ( $B_{\mathrm{g}}=B_{\mathrm{val}}$ )	128	Client participation ( $m_{\mathrm{up}}=m_{\mathrm{val}}$ )	5
Learning rate ( $\eta$ )	0.1	Momentum parameter ( $\alpha$ )	0.2
Penalty parameter ( $\mu$ )	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 (‘ $\bullet$ ’) represent ours (direct, momentum and STORM), while crosses (‘ $\times$ ’) and triangles (‘ $\blacktriangle$ ’) 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 ( $\theta$ )	30
Episode length	200	Discount factor ( $\gamma$ )	1.0
Neural network dimension ( $d$ )	17,410	Total number of clients ( $n$ )	10
Batch size per client ( $B_{\mathrm{g}}=B_{\mathrm{val}}$ )	8	Client participation ( $m_{\mathrm{up}}=m_{\mathrm{val}}$ )	4
Number of runs	3	Momentum paramter ( $\alpha$ )	0.2

Appendix

Table 3: Detailed setting of CMDP experiment on Cartpole

Federated Learning

Federated learning trains shared models across data holders without requiring them to pool their raw data.

See all 943 papers

First-Order Softmax Weighted Switching Gradient Method for Distributed Stochastic Minimax Optimization with Stochastic Constraints
Mar 6, 2026Zhankun Luo, Antesh Upadhyay, Sang Bin Moon +1Stochastic Optimization Federated Learning
School of Electrical and Computer Engineering, Purdue University
Decision-Focused Federated Learning Under Heterogeneous Objectives and Constraints
Apr 21, 2026Konstantinos Ziliaskopoulos, Alexander VinelPredict-Then-Optimize Learning to Optimize
Auburn University
A Locally Penalized Cross-Estimate Federated Method with Guarantees for Constrained Personalized Learning
Mar 20, 2026Mohammadjavad Ebrahimi, Daniel Burbano, Farzad YousefianFederated Learning Heterogeneous Federated Learning
Department of Industrial and Systems Engineering, Rutgers University, USA. · Department of Electrical and Computer Engineering, Rutgers University, USA.

Mar 6, 2026J/K move · Enter open · S save

First-Order Softmax Weighted Switching Gradient Method for Distributed Stochastic Minimax Optimization with Stochastic Constraints

Zhankun Luo, Antesh Upadhyay, Sang Bin Moon, Abolfazl Hashemi

School of Electrical and Computer Engineering, Purdue University

Stochastic Optimization Federated Learning

Open paper AskPDF arXiv

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

\tilde{\mathcal{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

\mathcal{O}(\log\frac{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.

Abstract

Federated Learning

Explore similar work

First-Order Softmax Weighted Switching Gradient Method for Distributed Stochastic Minimax Optimization with Stochastic Constraints

Decision-Focused Federated Learning Under Heterogeneous Objectives and Constraints

A Locally Penalized Cross-Estimate Federated Method with Guarantees for Constrained Personalized Learning