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

Jul 16, 2026

Distributed stochastic minimax optimization under stochastic constraints is addressed with a first-order Softmax-Weighted Switching Gradient method designed for federated learning. With full client participation, the method attains oracle complexity O(ε−4) to meet a unified tolerance ε for both optimality gap and feasibility. The analysis extends to partial participation via a stochastic superiority assumption. Relaxing objective boundedness yields a tighter softmax hyperparameter lower bound, and a sharp O(log(1/δ)) high-probability convergence result. Experiments validate performance on Neyman–Pearson classification, fair classification, and federated safe reinforcement learning, using a stable primal-only switching alternative.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages59

Reading Time2 h 29 min

Abstract

Introduction

Federated learning objective and heterogeneity

Distributionally robust minimax formulation

Stochastic minimax with stochastic constraints

Key challenges: non-smoothness, worst-case constraints, and coupling

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 {luo333, aantesh, moon182, [abolfazl}@purdue.edu](abolfazl}@purdue.edu) arXiv :2603 .05774v2 [ cs .LG] 13 Jul 2026 Abstract 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 p˜articipation, 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 primaldual 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. 1 Introduction Federated Learning (FL) aims to solve distributed optimization problems of the form n mwnΘ Xpifi (w), (1) i=1 where Θ ⊆ Rd is a compact convex set, n is the number of clients, fi (w) := Eζ∼Di [fi (w,ζ)] denotes the local expected loss at client i and pi denote the probability weight of the client [McMahan et al., 2017, Kairouz and McMahan, 2021] . Under statistical heterogeneity, where local distributions {Di } are non-identical, this empirical risk minimization (ERM) objective inherently prioritizes average performance across clients [Li et al., 2020a, Mohri et al., 2019] . As a result, the learned model is biased toward dominant client distributions and may exhibit severely degraded performance on underrepresented or hard clients [Mohri et al., 2019, Hashimoto et al., 2018, Li et al., 2019] . To guarantee uniformly good performance across all devices, a powerful alternative is to frame the training process as a distributionally robust (or agnostic) optimization problem [Mohri et al., 2019, Deng et al., 2020, Duchi and Namkoong, 2021] . Instead of minimizing the average loss, the algorithm minimizes the maximum expected loss over a global set of adversarial weights n mwnΘ mλ X λifi (w), (2) i=1 where Λ := 􀀈λ ∈ R : P λi = 1 􀀉 represents the probability weight assigned to each local client i. Intuitively, the inner maximization concentrates probability mass on the worst-performing clients. This recovers the equivalent minimax formulation minw∈Θ maxi∈I fi (w) with I :={1, 2 ,..., n} which directly enforces robustness to client heterogeneity. Existing minimax formulations in federated settings typically optimize this worst-case loss in isolation. However, in many practical deployments, models must simultaneously satisfy strict client-wise operational requirements, such as fairness mandates, safety limits, resource budgets, or regulatory thresholds [Islamov et al., 2025b, Upadhyay et al., 2026] . Tracking n distinct stochastic constraints, i.e., gi (w) = Eζ∼Di [gi (w,ζ)] ≤ 0 , ∀i ∈ I, which encode client-specific operational constraints, is prohibitively expensive in federated environments, as it requires maintaining and synchro- nizing n distinct dual variables across a network with intermittent cl

FAQs

What optimization problem does the paper focus on?

It studies distributed stochastic minimax optimization subject to stochastic constraints in federated learning, aiming to control both worst-case performance and constraint violations across heterogeneous clients.

What is the proposed algorithmic approach?

The paper introduces a first-order Softmax-Weighted Switching Gradient method, tailored for federated learning, that uses a switching mechanism with a Softmax weighting to handle the worst-case structure.

What does the paper guarantee about complexity and convergence?

Under full client participation, it achieves oracle complexity O(ε−4) to satisfy a unified bound ε for both optimality gap and feasibility. It also provides a sharp high-probability convergence guarantee with rate O(log(1/δ)).

Jul 16, 2026

Upload by Lucas Martin