Federated Conditional Stochastic Optimization
Authors: Xidong Wu
Jianhui Sun
Zhengmian Hu
Junyi Li
Aidong Zhang
Heng Huang
Venue: NeurIPS 2023 poster (2023)
Keywords: Federated Learning
Conditional Stochastic Optimization
Nonconvex Optimization

Abstract: Conditional stochastic optimization has found applications in a wide range of machine learning tasks, such as invariant learning, AUPRC maximization, and meta-learning. As the demand for training models with large-scale distributed data grows in these applications, there is an increasing need for communication-efficient distributed optimization algorithms, such as federated learning algorithms. This paper considers the nonconvex conditional stochastic optimization in federated learning and proposes the first federated conditional stochastic optimization algorithm (FCSG) with a conditional stochastic gradient estimator and a momentum-based algorithm (\emph{i.e.}, FCSG-M). To match the lower bound complexity in the single-machine setting, we design an accelerated algorithm (Acc-FCSG-M) via the variance reduction to achieve the best sample and communication complexity. Compared with the existing optimization analysis for Meta-Learning in FL, federated conditional stochastic optimization considers the sample of tasks. Extensive experimental results on various tasks validate the efficiency of these algorithms.

Federated Conditional Stochastic Optimization
Xidong Wut Department of ECE University of Pittsburgh Pittsburgh, PA 15213 sidong_wu@outlonn com
Jianhui Sunt Computer Science University of Virginia Charlottesville, VA 22903 jsšguvvirginia. edu
Zhengmian Hu Computer Science University of Maryland
College Park, MD 20742 muzhengmian@gne com
Junyi Li Department of ECE University of Pittsburgh Pittsburgh, PA 15213 junyili. ai @gmail com
Aidong Zhang Computer Science University of Virginia Charlottesville, VA 22903
Heng Huang"
Computer Science
University of Maryland
College Park, MD 20742
aidong@virginia.co
henghuanghh@gmaill
Abstract
Conditional stochastic optimization has found applications in wide range of machine learning tasks, such as invariant learning, AUPRC maximization, and meta-learning. As the demand for training models with large-scale distributed data grows in these applications, there is an increasing need for communication-effici distributed optimization algorithms, such as federated learning algorithms. This paper considers the nonconvex conditional stochastic optimization in federated learning and proposes the first federated conditional stochastic optimization algo- rithm (FCSG) with conditional stochastic gradient estimator and momentum- based algorithm (i.e., FCSG-M). To match the lower bound complexity in the single-machine setting, design accelerated algorithm (Acc-FCSG-M) via the variance reduction to achieve the best sample and communication complexity. Compared with the existing optimization analysis for Meta- Meta-Learning in FL, feder- ated conditional stochastic optimization considers the sample of tasks. Extensive experimental results various tasks validate the efficiency of these algorithms.
Introduction
The conditional stochastic optimization arises throughout wide range of machine learning tasks, such as the policy evaluation in reinforcement learning [5], invariant learning [16), instrumental variable regression in causal inference [30,, Model-Agnostic Meta- Meta-Learning (MAML) [10], AUPRC maximization 28 and so on. Recently many efficient conditional stochastic optimization algorithms have been developed [16 28
B4 32 to solve the corresponding machine learning problems and applications. However, all existing conditional stochastic optimization algorithms were only designed for centralized learning (i.e., model and data both deployed at single machine) or finite- sum optimization, without considering the large-scale online distributed scenario. Many federated proposed since FL is communication- efficient training paradigm for large-scale machine learning training preserving data privacy. In federated learning, clients update the model locally, and the global server aggregates the model
learning algorithms (26] 22
29
23
44
39
25
45
I 'qqual contribution
TThis work was partially supported by NSF CNS 2213700 and CCF 2217071 at UVA. "This work was partially supported by NSF IIS 1838627, 1837956, 1956002, 221 1492, CNS 2213701, CCF 2217003, DBI 2225775 at Pitt and UMD.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

parameters periodically, Although federated learning has been actively applied to numerous real- world applications in the past years, the federated conditional stochastic optimization problem is still underexplored. To bridge this gap, in this paper we study the following federated conditional stochastic optimization problem:
min F(x)(()) :- N Σ
Eg f (By Si E");
(I)
#EX
where X' C Id is a closed convex set, Be- (t) Rdd + R is the outer-layer function on the n-th device with the randomness ç", and Ie (..5)) Rd Rdd is the inner-layer function on the n-th device with respect to the conditional distribution of 17" g We assume (-) and Si (-.() are continuously differentiable. The objective subsumes two stochastic functions in ( where the inner functions rely on the randomnesses of both inner and outer layers, and & and 1) are not independent, which makes the federated conditional stochastic optimization more challenging compared with the standard federated learning optimization problems.
Federated conditional stochastic optimization contains the standard federated learning optimization as a special situation when the inner-layer function 9.j.. (I.(") ) = x. In addition, federated stochastic compositional optimization is similar federated conditional stochastic optimization given that both problems contain two-layer nested expectations. However, they are fundamentally different. In federated stochastic compositional optimization, the inner randomness 1) and the outer randomness & are independent and data samples of the inner layer are available directly from 1) (instead of a conditional distribution as in Problem (.). Therefore, when randomnesses 7) and & are independent and Pijn (z, -) 9h2- (I), ( is converted into federated stochastic compositional optimization [II].
Recently, to solve the conditional stochastic optimization problem efficiently, [16] studied the sample complexity of the sample average approximation for conditional stochastic optimization. Afterward,
[17] proposed the algorithm called biased stochastic gradient descent (BSGD) and an accelerated algorithm called biased SpiderBoost (BSpiderBoost). The convergence guarantees of BSGD and BSpiderBoost under different assumptions are established. More recently, 128
34
37
32
[4]
reformulated the AUC maximization into a finite-sum version of conditional stochastic optimization and introduced algorithms solve it. In increasing amount of distributed computing settings, efficient federated learning algorithms are absolutely necessary but still lacking. An important example of conditional stochastic optimization is MAML. In meta-learning, we attempt to train models that can efficiently adapt to unseen tasks via learning with metadata from similar tasks [TO. When the tasks are distributed at different clients, a federated version of MAML would be beneficial to leverage information from all workers B A lot of existing efforts 19 1 have been made to convert FL MAML into federated compositional optimization. Nonetheless, they ignore the sample of tasks in MAML, and federated conditional stochastic optimization problems have never been studied. Thus, there exists a natural question: Can we design federated algorithms for conditional stochastic optimization while maintaining the fast convergence rate to solve Problem (
In this paper, we give an affirmative answer to the above question. We propose a suite of approaches to solve Problem (I) and establish their corresponding convergence guarantee. To our best knowledge, this is the first work that thoroughly studies federated conditional stochastic optimization problems and provides completed theoretical analysis. Our proposed algorithm matches the lower- sample complexity in single-mas machine setting and obtains convincing results in empirical experiments. Our main contributions are four-fold:
1) we propose the federated conditional stochastic gradient (FCSG) algorithm to solve Problem (T.. We establish the theoretical convergence analysis for FCSG. In the general nonconvex setting, we prove that FCSG has a sample complexity of O(e and communication complexity of O to reach e-stationary point, and achieves appealing linear speedup w.zt the number of clients.
2) To further improve the empirical performance of our algorithm, we introduce a momentum- based FCSG algorithm, called FCSG-M since the momentum-basee estimator could reduce noise from samples with history information. FCSG-M algorithm obtains the same theoreti- cal guarantees FCSG.
3) To reach the lower bound of sample complexity of the single-machine counterpart 17 we propose an accelerated version of FCSG-M (Acc-FCSG-M) based on the momentum- based variance reduction technique. We prove that Acc-FCSG-M has sample complexity
2

---

of "), and communication complexity of O) which matches the best sample complexity attained by single-machine algorithm BSpiderBoost with variance reduction.
4) Experimental results on the robust logistic regression, MAML and AUPRC maximization tasks validate the effectivvness of our proposed algorithms.
2 Related Work
2.1 Conditional Stochastic Optimization
[6 studied the generalization error bound and sample complexity of the sample average approxima- tion (SAA) for conditional stochastic optimization. Subsequently, [17] proposed a class of efficient stochastic cradient-based methods for general conditional stochastic optimization to reach either global optimal point in the convex setting or stationary point in the nonconvex setting, respectively. In the nonconvex setting, BSGD has the sample complexity of O(E 6) and a variance reduction algorithm (BSpiderBoost) has the sample complexity of O(e [8 utilized the Monte Carlo method to achieve better results compared with the vanilla stochastic gradient method, Recently, (28 converted AUPRC maximization optimization into the finite-sum version of the conditional stochastic optimization and propose adaptive and non-adaptive stochastic algorithms to solve Similarly, recent work (34] used moving average techniques to improve the convergence rate of AUPRC maximization optimization and provide theoretical analysis for the adaptive algorithm. Furthermore, (2] focused on finite-sum coupled compositional stochastic optimization, which limits the outer-layer function to the finite-sum structure. The algorithms proposed 332 improved oracle complexity with the parallel speed-up. More recently, [14] use federated learning to solve AUC maximization. However, algorithms proposed in [28] 34 32 [4] for AUC maximization have a significant limitation because they maintain an inner state for each data point. As result, its convergence rate depends on the number of data points and cannot be extended to other tasks and large-scale model training. It is also not applicable to online learning due the dependence on each local data point. (38) consider the decentralised online AUPRC maximization but the theoretical analysis cannot be applied into the federated learning.
2.2 Stochastic Compositional Optimization
Recently, a related optimization problem, stochastic compositional optimization, has attracted widely attention 36 and solve the following objective:
min F(x) == (En,((E)))))))))
x
To address this problem, 336 developed SCGD, which utilizes the moving average technique to
estimate the inner-layer function value. (35) further developed an accelerated SCGD method with the were presented
extrapolation-smoothie scheme. Subsequently, series of algorithms (20 43 15 41 to improve the complexities using the acceleration or variance reduction techniques.
More recently. [19 and [TI] studied the stochastic compositional problem in federated learning. 19 transformed the distributionally robust federated learning problem (i.e., a minimax optimization problem) into a simple compositional optimization problem by using KL divergence regularization and proposed the first federated learning compositional algorithm and analysis. model personalization problem in federated learning as model-agnostic meta- problem. In personalized federated learning, each client's task assignment fixed and there no task sampling on each client in the training procedure. The sampling of the inner layer and outer layer are independent. Therefore, personalized federated learning is formulated as the stochastic compositional optimization 119. 33 solves personalized federated learning utilizing SCGD, in contrast to SGD in 19), to reduce the convergence complexities. However, the algorithm in (33] has a drawback in that keeping an inner state for each task is necessary, which is prohibitively expensive in large-scale settings. More recently, [ proposed a momentum- like method for nonconvex problems with better complexities to solve the stochastic compositional problem in the federated learning setting. Although [I]] claims their algorithm can be used in the MAML problem, it does not consider the two fundamental characteristics in MAML, i.e., task sampling and the dependency of inner data distribution on the sampled task.
K formulated the
Overall, problems T and 2 differ two aspects: i) in Problem ( the inner randomness 1) and the outer randomness are independent, while in Problem (1). 17 is conditionally dependent on the E: and ii) in Problem ( the inner function depends on both & and 1). Therefore, Problem (2) can be regarded as special case of (I). Thus, the conditional stochastic optimization I is more general.
3

---

Table 1: Complexity summary of proposed federated conditional stochastic optimization algorithms reach an e-stationary point. Sample complexity is defined as the number of calls to the First-order Oracle (IFO) by clients to reach an c-stationary point. Communication complexity denotes the total number of aack-and-forth communication rounds between each client and the central server required to reach an E-stationary point.
Algorithm
Sample
Communication
FCSG
0
0 (
FCSG-M
0 (
(
Lower Bound [6]
0 ( (5-
Acc-FCSG-M
0 c
0 (E
3 Preliminary
For solving the problem I.. we first consider the local objective (²)(x) and its gradient. We have F"(x)) (x) Een f (E 9yt ))
VF" Eç" [En g. (z; (")))) f (e (x, [x)]]
Since there are two layers of stochastic functions, the standard stochastic gradient estimator is not an unbiased estimation for the full gradient. Instead of constructing an unbiased stochastic gradient estimator, 17 considered biased estimator of V (x)) using one sample of & and samples of
V Fm (x; E", B.)
Σ
Big (z, (x,,")))
Σ EB,
Pap (xr,
1/ BB,
where B, = ! And F (x: Br) is the gradient of an empirical objective such that
Fm (x; B,) f m 3T EB,
Σ
Ing (z)
3.1 Assumptions
Assumption 3.1. (Smoothness) Vn E [N], the function (.) is Sy-Lipschitz smooth, and the function Si- (.) is LypLhshc smooth, ie., for sample 8d and m samples (1]² yy=1 from the conditional distribution P(r7" {)), Vo,, 22 dom f*(-), and Vyr.32 E dom g" (c), there exist S₁ > 0 and Sg > 0 such that
V 1 (x1) V f (x2)ll S S5]| -2ll
E| S E") Viyj (v2,<">)l S S,llvi -2ll
Assumption 3.1 is widely used assumption optimization analysis. Many single-maching stochastic algorithms use this assumption, such as BSGD (17), SPIDER [9], STORM [4) ADSGD M. and D2SG [12]. distributed learning, the convergence analysis of distributed learning algorithms. such as DSAL [2), and many federated learning algorithms such as MIME (21), Fed-GLOMO (6. STEM 23 and FAFED (39] also depend on
Assumption 3.2. (Bounded gradient) Vn E [NV, the function fr((-) is L Ly-Lisschitz continuous, and the function g" (-) is L, -iiischctt continuous, i.e., Vr dom fm(+). and Vy dom (-), the second moments of functions are bounded as below:
E 1 (x) || < L3
E Vg ( <² < D3
Assumption 3.2 is typical assumption in the multi-layer problem optimization to constrain the upper bound of the gradient of each layer, as in [36]
28
Assumption 3.3. (Bounded variance) [17 Vn E [M]. and X:: Eygig= 1ggn (x, 5")) Ey Io (a, en) o²
sup E
where 03 < +00. Assumption 3.3 indicates that the random vector By has bounded variance.

---

Algorithm FCSG and FCSG-M Algorithm
Input: Parameters: momentum weight p. learning rate a, the number of local updates q, inner batch size m and outer batch size b, as well as the initial outer batch size B
2: Initialize: Tü = Zo N Eb= z. Draw B samples of (Si,1 Stt
".l.,,.., (t,n)
and draw
m samples B
{
m j=1
from P(r)" | (6,) for each S6 E ( ((,B):
1 + V F (rr; S0, B.,)).
3: for l
................................ T do
4: for n = 1, 2,
do then
5:
mod (t. Server Update:
6:
7:
" the N Ei= uz
8:
24 Jt. = N En 1 - x2u
9:
else
10:
27 x/-1 - auu
11:
end if
12:
Draw b samples of {Ei, ... 51,61
13:
Draw samples Bin { ]==1 from P(r (.,) for each Ei "",,,,"","0"
14:
t = - V V (x ; Ei.3 Bi,))
15:
"f (1 B)u Σ b
V F ("f:: BY,)
16: end for
17:
end for
18:
Output: chosen uniformly random from (
4 Proposed Algorithms
In the section, we propose class federated first-order methods to solve the Problem We first design federated conditional stochastic gradient (FCSG) algorithm with biased gradient estimator and the momentum-
algorithm FCSG-M. To further accelerate our algorithm and achieve the lower bound of sample complexity of the single-machine algorithm, design the Acc-FCSG-M with variance reduction technique. Table summarizes the complex details of each algorithm.
4.1 Federated Conditional Stochastic Gradient (FCSG)
First, design a federated conditional stochastic gradient (FCSG) algorithm with the biased gradient estimator. We leverage mini-batch of conditional samples to construct the gradient estimator with controllable bias as (6). At each iteration, clients update their local models 2; with local data, which can be found in Line 9-14 of Algorithm Once every local iterations, the server collects local models and returns the averaged models each client, as Line 5-8 of Algorithm [ Here, the number of local update steps y is greater than such that the number of communication rounds is reduced to T/y. The details the method are summarized Algorithm [ Then we study the convergence properties of our new algorithm FCSG. Detailed proofs are provided in the supplementary materials.
Theorem 4.1. Suppose Assumptions
3.1
3.2 and
3.3 hold, if S 6ySF'
1
FCSG has the following
convergence result
([F(To) F(TT)) aT
2035300
205r13L23 N
|VF
42(4 1) 11)
m
Corollary 4.2. We choose
N T and q (T/N3/1/4 have
63p
2 125p(F( - F(a.)]
2T2
1.
1913
(/2)
m
6(NT)1/2
9(NT)
5

---

Remark 4.3. We choose B = b = 0(1) 2 1, and m = 0(= according to Corollary 4.2 to let T E FF = we get T -0(N-'=-)) O(N-1=-1) indicates the linear speedup of our algorithm. Given = (T/N5)3/1, the communication complexity is (NT)3/+= 0
Then the sample complexity is mT = O(N-'e
4.2 Federated Conditional Stochastic Gradient with Momentum (FCSG-M)
Next, we propose momentum-based based local updates algorithm (FCSG-M) for federated conditional stochastic optimization problems. Momentum is a popular technique widely used in practice for training deep neural networks. The motivation behind it in local updates is use the historic information (i.e., averaging of stochastic gradients) to reduce the effect of stochastic gradient noise. The details of our method are shown in Algorithm [
Initially, each device utilizes the standard stochastic gradient descent method to update the model in Line of Algorithm|] [ Afterward, compared with FCSG, each step, each client uses momentum- based gradient estimators u, to update the local model, which can found in Lines 9-15 of Algorithm [ The coefficient 3 for the update of u, should satisfy 0 < 8 < 1. In every q iterations, the clients communicate (Ie, ue} to the server, which computes the (Te, ūe ). and returns the averaged model and gradient estimator to each client, Lines 5-8 of Algorithm [ Then we study the convergence properties of our new algorithm FCSG-M. The details of the proof are provided in the supplementary materials.
parameter,
Theorem 4.4. Suppose Assumptions the following convergence result
3.1
3.2 and 3.3 hold, a S 6g5p and B = 5S pn. FCSG-M has
1
F F(oo) F(TT)
E VF()))11 2
T
aT
96S
4L227003
812,12 BBT
8,13,12 N
p
a
(1
+ N
+ +..,31,91
5
m
Corollary 4.5. We choose a
7. g = 1/4, we have
G5F
[F - F(2.)
VF
7
(NT))/2
11217 4125700
48L3,2
20L3
3(NT) +
m
5(
3(NT)1/2
Remark 4.6. We choose b = 0(1), (1),, and = 0(= According to Corollary 4.5 to make + II- [/ S e², we get ONN-----)) Given q = (T/N3)1/4, the commu- nication complexity is I = (NT)3/4 (s- The sample complexity is m = O(N-18-)) which indicates FCSG-M also has the linear speedup with respect to the number of clients.
4.3 Acc-FCSG-M
In the single-machine setting, 17 presents that under the general nonconvex conditional stochastic optimization objective, the lower bound of sample complexity is 0(s-5). means that the sample complexity achieved by FCSG and FCSG-M could be improved for nonconvex smooth conditional stochastic optimization objectives. To match the above lower bound of sample complexity, we propose accelerated version of the FCSG-M (Acc-FCSG-M) based on the momentum-batt based variance reduction technique. The details of the method are shown in Algorithm 2
Similar to the FCSG-M, in the beginning, each client initializes the model parameters and utilizes the stochastic gradient descent method calculate the gradient estimator. Subsequently, in every g iterations, all clients perform communication with the central server, and the model parameters and gradient estimators are averaged. The key difference is that at Line 14 in Acc-FCSG-M, we use the momentum- based variance reduction gradient estimator w/+1 to track the gradient and update the model. where B E (0, 1). We establish the theoretical convergence guarantee of our new algorithm Acc-FCSG-M. All proofs are provided in the supplementary materials.

---

Algorithm 2 Acc-FCSG-M Algorithm
1: Input: T, momentum weight B, learning rate a, the number of local updates g, inner batch size m and outer batch size b, as well as the initial outer batch size B
EN
2: Initialize: =0 = N k== z0- Draw B samples of ,..,..,, and draw m sam-
ples B = {
from P(n" (T) for each 5 E t ; then of
j-1
Di=1
B v
(x)) 50. BG,)
for [N].
3:
for trrt = 1,2
do
4:
for = 1,2, ..... do
5:
if mod (t.g) 0 then
6:
Server Update:
7:
"i Ur N ol
8:
xt It E (1 onl?)
9:
else
10:
Fin r/-1 au!
11:
end
if
12:
Draw b samples of {5t,1) St,bJ
13:
Draw m samples B = { j= from P(7" (2,) for each S1, E (5i,1 ... ,i,,)
14:
"h+1 = + V F (17: 57, Bi,) ) - (u, E V
(2f (xf
1i,,,, B(,))
15:
end for
16:
end for
17:
Output: x chosen uniformly random from { 1==
Theorem 4.7. Suppose Assumptions 3.1
3.2 and
3.3
and 3 5Spa. Acc-FCSG-M hold, S 6q5F
has the following
[FF 2(F( F(
3L312
13
3L35703
63L;
a²
Ta
BBNT
6.57
m
Nb
Corollary 4.8. We choose q= (T/N2) 1/3
Therefore, a 12g8F And let B T//3 Therefore, have
128
N2/3
since we
30S2 bN
have 8 ca2
5/'/3 2472738
[F < 24S (F(t)) - F(i.))
(NTY2/3
12
3251.
312
5T
+
+ 24h² 24h (TN)2/3
5(NT)2/3
4b2 (NTT2/3
m
Remark 4.9, We choose b 0(1) 2 1) and m 0 According Corollary 4.8 make
+ EI- [VF(2)] < E², we get ((N-g-3) and 2 O(E 2). The sample
complexity is O(N-1e-5). The communication complexity is 0( O(N -yy-g -3) indicates the linear speedup of our algorithm.
5
Experiments
The experiments are run on CPU machines with AMD EPYC 7513 32-Core Processors as well as NVIDIA RTX A6000. The code is available x and Federated Online AUPRC maximization task follow 38
5.1 Invariant Logistic Regression
Invariant learning has an important role in robust classifier training 22 In this section, We compare the performance of our algorithms, FCSG and FCSG-M, on the distributed invariant logistic regression
7

---

Tost Accuracy
Teet Accuracy
PCS0 (c, (1,=1) =1
FCSG 5
FCS0 (r, =1
-FCGG-M 2-
FFCGGM l,
am
FCSG-M 6:=-
PCSG (r, -1.9
FC5G (1 1.5
PCSG (r, 1..9
FCB0-M -1.8
-
(-,=
FCSG -2)
FC5G (---
FCSG (v,
FCS0-M byg 22
-FC80-M -
-FC80-M 6.
Communication Round
Communication Round
Communication Round
(a) m
(b) m
(c) m
Figure 1: Test accuracy vs the number communication rounds for different inner mini-batch (m 1, 10, 100) under different noise ratios (02/01 1.5.
evaluate the benefit from momentum and the effect of inner batch size, and the problem was formulated by [16]
N Σ
min
N (a. g(x)
n=l
where (x:) log(1 expl (xappbBB. [67")) [()"))=)))
g(x) = =
y
y
1=1
where (x) the logistic loss function and g(x) is a non-convex regularization. We follow the experimental protocols [17 and set the dimension of the model 10 over 16 clients. We construct the dataset <" = (a, b) and v as follow: We sample a ~ N(0, af Ta). set b = {+1} according to the sign of a x², then sample 7/ ~ N (a, oā. We choose σ1 1, and consider the 02/01 from {1, 1,5, 2}. At each local iteration, we use a fixed mini-batch size m from {1, 10, 100). The outer batch size is set 1. We test the model with 50000 outer samples report the test accuracy. We carefully tune hyperparameters for both methods. À 0.001 and a = 10. We run a grid search for the learning rate and choose the learning rate in the set {0.01, 0.005, 0.001), in FCSG-M chosen from the set {0.001, 0.01, 1.110.111 0.1, 1,11,00 0.5, 0.9). The local update step is set as 50. The experiments are presented in Figure 1
Figure 1 shows that when the noise ratio 02/01 increases, larger inner samples m are needed, as suggested by the theory because large batch size could reduce sample noise. In addition, when m = 100, FCSG and FCSG-M have similar performance. However, when batch size small, compared with FCSG, FCSG-M has more stable performance, because FCSG-M can use historic information to reduce the effect of stochastic gradient noise.
5.2 Federated Model-Agnostic Meta-nearning
Next, we evaluate our proposed algorithms with the few shot image classification task over the dataset with baselines: Local-SCGD and Local-SCGDM [1] MOML [33]
suitable for this task since MOML requires maintaining an inner state for each task which is not permitted due to the large number tasks in the Omniglot dataset. Local-SCGD is the federated version of SCGD (6). This
task can be effectively solved via Mgnet-Agnostic Meta-Learning [10]
Meta-learning aims train model various learning tasks, such that the model can easily adapt
to new task using few training samples, Model-agnostic meta-learning (MAML) (10) is a popular meta-learning method to learn good initialization with gradient- based update, which can be formulated as the following conditional stochastic optimization problem:
min Ei-PPaa
Dian L, (E Dimren (ar Xz,
where Punk denotes the learning tasks distribution, D Dimpport and Dí query support (training) dataset and query (testing) dataset of the learning task i, respectively. Le(., D.) is the loss function on dataset D, of task And the is fixed meta step size. Assume { = (V, a) and y = b, the MAMI. problem an example of conditional stochastic optimization where the support (training) samples in the inner

---

BOB S0OR
ToW BOOD
FCRG
PORD-M
Tes Aeterroo
PORD
PORO
Thabbor of forations
Number llerations
Humtbrr Reeatione
Numbbrr of Ranationn
Figure 2: Convergence results of the 5-way-1-saa --shot case over Omniglot Dataset.
9060
local doioo
Trans
FORG
FOSG
Top Aacooon
COGG AcS--000-1
Namber of lerooaoe
Nambbr Rocotiono
Number Nerations
Number lltrrtioos
Figure 3: Convergence results of the 5-way-5-sher case over Omniglot Dataset.
layer for the meta-gradient update are drawn from the conditional distribution of P(n { based on the sampled task in the outer layer.
Given there are a large number pre-train tasks MAML, federated learning a good training strategy to improve efficiency because can evenly distribute tasks over various clients and the global server coordinates clients to learn a good initial model collaboratively like MAML. Therefore, in Federated MAML, it is assumed that each device has part of the tasks. The optimization problem is defined as follows:
min ZERd
1
F"
(z)
EST (g"
N
N
(5)
where (x) = Dy appen
[x AVET (f""")) = Ei-pa, an an Di, qary
C. (y:a") a").
In this part, we apply our methods to few-shot image classification on the Omniglot 24 10. The Omniglot dataset contains 1623 different handwritten characters from 50 different alphabets and each of the 1623 characters consists of 20 instances drawn by different persons. We divide the characters to train/validation/tesl with 1028/172/423 by Torchmeta (7) and tasks are evenly partitioned into disjoint sets and we distribute tasks randomly among 16 clients. We conduct the fast learning of N-way- K-shot classification following the experimental protocol in [31]. The N-way-K-she classification denotes we sample N unseen classes, randomly provide the model with K different instances of each class for training, and evaluate the model's ability classify with new instances from the same N classes. We sample 15 data points for validation. We use 4-layer convolutional neural model where each layer has 3 3 convolutions and 64 filters, followed by a ReLU nonlinearity and batch normalization [10] The images from Omniglot are downsampled to 28 28. For all methods, the model trained using a single gradient step with learning rate of 0.4. The model was evaluated using 3 gradient steps [10]. Then we use grid search and carefully tune other syper-pprameters for each method. We choose the learning rate from the set (0.1, 0.05, 0.01) and r) as [1]. We select the inner state momentum coefficient for Local-SCGD and Local-SCGDM from {0.1, 0.5, 0.9] and outside momentum coefficient for Local-SCGDM, FCSG-M and Acc-FCSG-M from {0.1, 0.5, 0.9).
Figures 2 and 3 show experimental results in the 5-way-1-shb and 5-way-5-shot cases, respectively, Results show that our algorithms outperform baselines by a large margin. The main reason for Local- SCGD and Local-SCGDM to have bad performance is that converting the MAML optimization to the stochastic compositional optimization unreasonable. It ignores the effect of task sampling on the training and inner training data distribution changes based on the sampled tasks in the outer layer. The use of the momentum-lite inner state to deal with the MAML will slow down the convergence and we have to tune the extra momentum coefficient for the inner state. In addition, the momentum-lic inner state also introduces extra communication costs because the server needs to average the inner state in Local-SCGDA, In addition, comparing the results in Figures 2 and can when

---

Table 2: Final averaged AP scores on the testing data.
Datasets
FedAvg
CODA+
FCSG
FCSG-M
Acc-FCSG-M
MNIST
0.9357
0.9733
0.9868
0.9878
0.9879
CIFAR-I0
0.5059
0.6039
0.7130
0.7157
0.7184
the K increases in the few-shot learning, the training performance is improved, which matches the theoretical analysis that large inner batch-size m benefits the model training.
5.3 Federated Online AUPRC maximization
AUROC maximization in FL has been studied 13
40 and AUPRC maximization is also
used to solve the imbalanced classification. Existing AUPRC algorithms maintain inner state for each data point. (38] consider the online AUPRC in the decentralized learning. For the large-scale distributed data over multiple clients, algorithms for online AUPRC maximization in FL is necessary.
Following 288 and 38 the surrogate function of average precision (AP) for online AUPRC maxi-
mization is:
BENDI (u =
AP = Eg WD+ EęND C (x)
where (a;; =t, 2) (max h(r; = + +h(1;;) 2), 0})² and h(:; is the prediction score function of input z with model x. Federated Online AUPRC maximization could be reformulated as:
min F(x) min N Σ
B f(E D. (x:
(6)
where Sn (=n; Vn)) D, and 5it = (zt: Mit) D are samples drawn from the whole datasets and positive datasets, respectively. is a two-level problem and the inner objective depends on both É and Since federated online AUPRC is special example of federated conditional stochastic optimization, our algorithms could be directly applied to it.
We choose MNIST dataset and CIFAR-10 datasets. As AUROC maximization in federated settings has been demonstrated in existing works 13 42 we use CODA+ in 42 as a baseline. Another baseline is the FedAvg with cross-entropy loss. Since AUPRC is used for binary classification, the first half the classes in the MNIST and CIFARIO 10 datasets are designated to be the negative class, and the rest half of the classes are considered to be the positive class. Then, we remove 80% of the positive examples in the training set to make it imbalanced, while keeping the test set unchanged. The results in Table 2 show that our algorithms could be used to solve the online AUPRC maximization in FL and it largely improves the model's 's performance.
Conclusion
In this paper, we studied federated conditional stochastic optimization under the general nonconvex setting. To the best of our knowledge, this is the first paper proposing algorithms for the federated conditional stochastic optimization problem. We first used the biased stochastic first-order gradient to design an algorithm called FCSG, which we proved have a sample complexity of O(e ((c-)), and communication complexity of (6) to reach an e-stationary point. FCSG enjoys an appealing linear speedup with respect to number of clients. To improve the empirical performances of FCSG, we also proposed a novel algorithm (i.e., FCSG-M), which achieves the same theoretical guarantees as FCSG. To fill the gap from lower-bound complexity, we introduced an accelerated version of FCSG-M, called Acc-FCSG-M, using variance reduction technique, which is optimal for the nonconvex smooth federated conditional stochastic optimization problems as matches the best possible complexity result achieved by the single-machine method. It has sample complexity of Occ=)). ", and communication complexity of O(c-2) And the communication complexity achieves the best communication complexity of the nonconvex optimization in federated learning. Experimental results on the machine learning tasks validate the effectiveness of our algorithms.
Limitation Conditional stochastic optimization has much broader applications and more com- prehensive evaluation of our proposed algorithms on other use cases would be a promising future direction. In addition, FCSG-M and accelcrated FCSG-M require higher communication overhead.
10

---

References
[1] Runxue Bao, Bin Gu, and Heng Huang. accelerated doubly stochastic gradient method with faster explicit model identification. In Proceedings of the 31xt ACM International Conference on Information Knowledge Management, pages 57-66, 2022.
[2] Runxue Bao, Xidong Wu, Wenhan Xian, and Heng Huang. Doubly sparse asynchronous learning for stochastic composite optimization. In Thirty-Fire First International Joint Conference on Artificial Intelligence (IJCAI), pages 1916-1922, 2022.
[3] Fei Chen, Mi Luo. Zhenhua Dong, Zhenguo Li, and Xiuqiang He. Federated meta- meta-learning with fast convergence and efficient communication. arXiv preprint arXiv:1802.07876. 2018.
[4] Ashok Cutkosky and Francesco Orabona. Momentum- based variance reduction in non-convex sgd. Advances in neural information processing systems, 32, 2019.
[5] Bo Dai, Niao He, Yunpeng Pan, Byron Boots, and Le Song. Learning from conditional distributions via dual embeddings. Artificial Intelligence and Statistics, pages 1458 1467. PMLR, 2017.
[6] Rudrajit Das, Anish Acharya, Abolfazl Hashemi, Sujay Sanghavi, Inderjit Dhillon, and Ufuk Topcu. Faster non-convex federated learning via global and local momentum. In Uncertainty in Artificial Intelli- gence, pages 496-00. PMLR, 2022.
[7] Tristan Deleu, Tobias Würfl, Mandana Samiei, Joseph Paul Cohen, and Yoshua Bengio. Torchmeta: A Meta-Learning library PyTorch, 2019. nattps://arxdvv.rr//bbbbs/ 06576 Available at: https://ithub..com http://hnwb.ommmiiiiimdddeeeeeee
[8] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach. arXiv preprint (tt:::00220094455 2020.
[9] Cong Fang, Chris Junchi Li, Zhouchen Lin, and Tong Zhang. Spider: Near-optimal non-convex opti- mization stochastic path-integrated differential estimator. Advances in Neural Information Processing Systems, 31, 2018
[10] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1126-1135 1135. PMLR, 2017.
[11] Hongchang Gao, Junyi Li, and Heng Huang. On the convergence of local stochastic compositional gradient descent with momentum. In International Conference Machine Learning, pages 7017-7035. PMLR, 2022.
[12] Bin Gu, Runxue Bao, Chenkang Zhang, and Heng Huang. New scalable and efficient online pairwise learning algorithm. IEEE Transactions on Neural Networks and Learning Systems, 2023.
[13] Zhishuai Guo, Mingrui Liu, Zhuoning Yuan, Li Shen, Wei Liu, and Tianbao Yang. Communication- distributed stochastic maximization with deep neural networks. In International conference on muchine learning, pages 3864-3874. PMLR, 2020.
[14] Zhishuai Guo, Rong Jin, Jiebo Luo, and Tianbao Yang. Fedxl: Provable federated learning for deep x-risk optimization. 2022.
[15] Wenqing Hu, Chris Junchi Li. Xiangru Lian, Ji Liu, and Huizhuo Yuan. Efficient smooth non-convex stochastic compositional optimization via stochastic recursive gradient descent. Advances in Neural Information Processing Systems. 32, 2019.
[16] Yifan Hu, Xin Chen, and Niao He. Sample complexity of sample average approximation for conditional stochastic optimization. SIAM Journal on Optimization, 30(3):1033-2333, 2020.
[17] Yifan Hu, Siqi Zhang, Xin Chen, and Niao He. Biased stochastic first-order methods for conditional stochastic optimization and applications in meta learning. Advances in Neural Information Processing Systems, 33:2759-2770, 2020.
[18] Yifan Hu, Xin Chen, and Niao He. On the bias-variance-cost tradeoff stochastic optimization, Advances in Neural Information Processing Systems, 34:22119 19-22131, 2021.
[19] Feihu Huang, Junyi Li. and Heng Huang. Compositional federated learning: Applications in distributionally robust averaging and meta learning. arXiv preprint 1iii:2106.11266 2021.

---

[20] Zhouyuan Huo, Bin Gu, Ji Liu, and Heng Huang. Accelerated method for stochastic composition optimization with nonsmooth regularization. In Shirty-Second AAAI Conference Artificial Intelligence. 2018.
[21] Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale. Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint arXiv:2008. 03608.00006...0000 2020.
[22] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In Inter- national conference on machine learning, pages 1132-5143. PMLR, 2020.
[23] Prashant Khanduri, Pranay Sharma. Haibo Yang. Mingyi Hong, Jia Liu, Ketan Rajawat, and Pramod Varshney. Stem: A stochastic two-sided momentum algorithm achieving near-optimal sample and com- munication complexities for federated learning. Advances in Neural Information Processing Systems, 34: 6050-6061, 2021.
[24] Brenden Lake, Ruslan Salakhutdinov, Jason Gross, and Joshua Tenenbaum. One shot learning of simple visual concepts. Proceedings of the annual meeting of the cognitive science society, volume 33, 2011.
[25] Tian Li. Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429-4500 2020.
[26] Brendan McMahan, Eider Moore. Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efffcceee learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282 PMLR, 2017.
[27] Youssef Mroueh, Stephen Voinea, and Tomaso A Poggio. Learning with group invariant features: A kernel perspective. Advances in neural information processing systems, 28, 2015.
[28] Qi Qi, Youzhi Luo, Zhao Xu, Shuiwang Ji, and Tianbao Yang. Stochastic optimization of areas under precision recall curves with provable convergence. Advances in Neural Information Processing Systems, 34:1752-1766 2021.
[29] Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295, 2020.
[30] Rahul Singh, Maneesh Sahani, and Arthur Gretton. Kernel instrumental variable regression. Advances in Neural Information Processing Systems, 32, 2019.
[31] Oriol Vinyals. Charles Blundell, Timothy Lillicrap. Daan Wierstra, et al. Matching networks for one shot learning. Advances in neural information processing systems, 29, 2016.
[32] Bokun Wang and Tianbao Yang. Finite-sum coupled compositional stochastic optimization: Theory and applications. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland. USA, volume 162 of Proceedings of Machine Learning Research, pages 13292-233313 PMLR. 2022.
[33] Bokun Wang, Zhuoning Yuan, Yiming Ying. and Tianbao Yang. Memory- based optimization methods for manes-ag meta-leaming. arXiv preprint arXiv:2106.04911 2021.
[34] Guanghui Wang, Ming Yang, Lijun Zhang, and Tianbao Yang. Momentum accelerates the convergence of stochastic aupre maximization. In International Conference Artificial Intelligence and Statistics, pages 3753 3771. PMLR, 2022.
[35] Mengdi Wang, Ji Liu, and Ethan Fang. Accelerating stochastic composition optimization. Advunces in Neural Information Processing Systems, 29, 2016
[36] Mengdi Wang, Ethan X Fang, and Han Liu. Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. Mathematical Programming. 161(():419-49) 2017.
[37] Xidong Wu. Feihu Huang, and Heng Huang. Fast stochastic recursive momentum methods for imbalanced data mining. In 2022 IEEE International Conference on Data Mining (ICDM), pages 578-587. IEEE, 2022.
[38] Xidong Wu, Zhengmian Hu, Jian Pei, and Heng Huang. Serverless federated aupre optimization for multi-pan party collaborative imbalanced data mining. In Proceedings of the 29th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pages (648-2659) 2023.
12

---

[39] Xidong Wu, Feihu Huang, Zhengmian Hu, and Heng Huang. Faster adaptive federated learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 37, pages 10379-10387, 2023.
[40] Xidong Wu, Jianhui Sun, Zhengmian Hu, Aidong Zhang, and Heng Huang. Solving class of non- convex minimax optimization in federated learning. Advances in Neural Information Processing Systems (Neu
[41] Huizhuo Yuan and Wenqing Hu. Stochastic recursive momentum method for non-convex compositional optimization. arXiv preprint arXiv:2006. 01688, 2020.
[42] Zhuoning Yuan, Zhishuai Guo, Yi Xu, Yiming Ying, and Tianbao Yang. Federated deep auc maximization hetergeneous data with constant communication complexity. International Conference on Machine Learning, pages 22299. PMLR, 2021.
[43] Junyu Zhang and Xiao. composite randomized incremental gradient method. International Conference on Machine Learning, pages 7464-7462. PMLR, 2019.
[44] Hanhan Zhou, Tian Lan, Guru Prasadh Venkataramani, Wenbo Ding. Federated learning with online adaptive heterogeneous local models. In Workshop on Federated Learning: Recent Advances and New Challenges (in Conjunction with NeurIPS 2022), 2022.
[45] Hanhan Zhou, Tian Lan, Guru Prasadh Venkataramani, and Wenbo Ding. Every parameter matters: Ensuring the convergence of federated learning with dynamic heterogeneous models reduction. In Advances in Neural Information Processing Systems, volume 36, 2023.
13


Reviews:
[{"review_id": "Zrqy9sgEcy", "reviewer": "Reviewer_L363", "summary": "Authors propose methods for problem of stochastic optimisation of composition of functions' expected values, which can be addressed with conditional gradient concept. Federated algorithm, which uses biased gradient estimation, and its modifications with momentum and acceleration are proposed. Theoretical guarantees on iteration and oracle complexities match lower bounds or correspond to best-known results.", "questions": "It may be better if convergence curves of algorithms were provided with shadows showing deviation of trajectory for different runs with different realisations of randomness.", "limitations": "Everything is okay.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The algorithms proposed are efficient theoretically and their efficiency is confirmed by numerical experiments. The experiments show practically interesting problems statements.", "weaknesses": "Empirical study could be more extensive nd helpful for practitioners who may want to use your method in their tasks if the comparison contained also different approaches designed specifically for MAML.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are deeply thankful for your comments. I will add deviation of trajectory for different runs with different realisations of randomness in our final version. "}, {"review_id": "fykYOPfcVA", "reviewer": "Reviewer_VagV", "summary": "This research delves into the realm of conditional stochastic optimization within the framework of federated learning. It begins by introducing the Federated Conditional Stochastic Gradient (FCSG) algorithm, which utilizes a conditional stochastic gradient estimator. Subsequently, the algorithm is refined by incorporating a momentum term into the update rules, thus successfully attenuating the influence of gradient estimate noise. Moreover, a variance reduction technique is validated to enhance sample complexity. The efficacy of the proposed algorithms is assessed through the evaluation of their performance on both real-world and synthetic datasets.", "questions": "See weaknesses.", "limitations": "certain limitations are discussed.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The objective function mentioned in equation (1) appears to be a generic representation capable of encompassing a broad spectrum of stochastic optimization problems, including stochastic compositional stochastic optimization. Moreover, it lends itself well to implementations based on federated learning.\n\n- The accelerated algorithm exhibits a sample complexity that aligns with the lower bound established for the corresponding single-machine counterpart.\n\n- Simulated tests on the robust logistic regression and MAML tasks show the superior performance of the proposed approaches.", "weaknesses": "- The proof of convergence to the stationary points, as presented in the supplementary material, can sometimes be challenging to follow due to the lack of clarity regarding the variables with respect to which the gradients are computed. For example, in equation (6), it is unclear which variable or function each gradient corresponds to. Additionally, there appears to be a mismatch in the dimensionalities of the terms, as the multiplication involves two $d \\times 1$ vectors (where $d$ represents the size of the vector $x$).\n\n- The submission states that the MAML problem is an example of conditional stochastic optimization. It is better to discuss the related works about federated meta-learning to highlight the novelty of this paper.\n\n- Compared with the FCSG algorithm, the FCSG-M and accelerated FCSG-M approaches require both the local model parameters and gradients to be transmitted to the central server for the aggregation, which may lead to higher communication overhead.\n\n- For the experimental results on invariant logistic regression, the performance of the proposed accelerated FCSG-M method should be involved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are deeply thankful for your appreciation and valuable comments.\n\n> Q1. \"in equation (6), it is unclear which variable or function each gradient corresponds to ...\" \n\nThanks. I will fix these typos in the final version. Since our objective function is a CSO and its gradient follow the chain rule in the derivative of the composition of two differentiable functions.  \n\nGiven that Inner-layer function  $\\mathbb{E}\\_{\\eta^n \\mid \\xi^n} g\\_{\\eta^n}^{n}(\\cdot, \\xi^n) : \\mathbb{R}^{d} \\rightarrow \\mathbb{R}^{d^{\\prime}}$ and the outer-layer function $\\mathbb{E}\\_{\\xi^n} f\\_{\\xi^n}^{n}(\\cdot): \\mathbb{R}^{d^{\\prime}} \\rightarrow \\mathbb{R}$, the dimension of $\\nabla g$ is $d^{\\prime} \\times d$ and  dimension of $\\nabla_g f$ is $d^{\\prime} \\times 1$. we rewrite (6) as \n\n$ \\nabla F^n(x) = \\nabla \\mathbb{E}\\_{\\xi^n}\\left[f^n\\_{\\xi^n}(g^n(x, \\xi^n))\\right] = \\mathbb{E}\\_{\\xi^n}\\left[\\nabla\\left(f^n\\_{\\xi^n}(g^n(x, \\xi^n))\\right)\\right]$   = $\\mathbb{E}\\_{\\xi^n}  [ [\\nabla g^n(x, \\xi^n) ]^{\\top} \\cdot \\nabla_g f^n_{\\xi^n}(g^n(x, \\xi^n)) ]$\n\n>Q2. \"It is better to discuss the related works about federated meta-learning to highlight the novelty of this paper.\"\n\nThanks for your suggestion. We talked about some existing works [8]. Due to the page limitation,  we will add more related works in the final version. \n\n> Q3. \" ...FCSG-M and accelerated FCSG-M approaches require both the local model parameters and gradients to be transmitted to .., which may lead to higher communication overhead\"\n\nFCSG-M and accelerated FCSG-M need to send gradient estimators. But it is very common in existing FL algorithms [2][3][4]. I will add this in the limitation part.\n\n> Q4.   \"For the experimental results on invariant logistic regression, the performance of the proposed accelerated FCSG-M method should be involved.\"\n\nThe target of the first experiment is to evaluate the benefit of momentum and the effect of inner batch size. The FCSG-M method is relatively simple to explore. Since we cannot upload the external link to present images, I will add the results with the accelerated FCSG-M method in the final version.\n\n[1] On the convergence of local stochastic compositional gradient descent with momentum. In International Conference on Machine Learning, pages 7017–7035. PMLR,327 2022\n\n[2] Khanduri, Prashant, et al. \"Stem: A stochastic two-sided momentum algorithm achieving near-optimal sample and communication complexities for federated learning.\" Advances in Neural Information Processing Systems 34 (2021): 6050-6061. \n\n[3] Karimireddy, Sai Praneeth, et al. \"Scaffold: Stochastic controlled averaging for federated learning.\" International Conference on Machine Learning. PMLR, 2020.\n\n[4] Sharma, P., Panda, R., Joshi, G. and Varshney, P., 2022, June. Federated minimax optimization: Improved convergence analyses and algorithms. In International Conference on Machine Learning. PMLR."}, {"review_id": "0mX48dbG1q", "reviewer": "Reviewer_w7D4", "summary": "This paper proposed algorithms for federated conditional stochastic optimization problem. Noteably, the sample complexity is of $O(\\epsilon^{-5})$ and the communication complexity is of $O(\\epsilon^{-2})$ to obtain $\\epsilon$-stationary point. This complexity matches by the single-machine algorithm BSpiderBoost with variance reduction that has been studied in Hu et. al. ", "questions": "N/A", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The study of federated conditional stochastic optimization problem is new in literature. The proposed algorithm has sharp sample complexity rate compared with the algorithms in literature.", "weaknesses": "- In related work part, the reference [13] actually introduced MLMC method that has sample complexity $O(\\epsilon^{-4})$, which is the optimal sample complexity. The authors omitted this fact and actually, I am wondering if it could be possible to utilize MLMC mehod to solve the federated conditional stochastic optimization problem? Then the sample complexity rate of $O(\\epsilon^{-5})$ could be potentially reduced.\n- In Fig. (2c) and (2d), the test loss seems does not become stable at last few iterations. The authors are suggested to increase the number of iterations to let readers see the whole trend.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and valuable comments. \n\nQ1. \"... the reference [13] actually introduced MLMC method that has sample complexity ...\"\n\n[13] utilized the Monte Carlo method to achieve better results compared with the vanilla stochastic gradient method. In our paper, we focus on the stochastic gradient method because it is very popular in ML training. Thanks for your suggestion and we will consider it in our future work.\n"}, {"review_id": "Zy0LvFCdmX", "reviewer": "Reviewer_tZ8Z", "summary": "In this paper, the authors delve into the complexities of nonconvex federated conditional stochastic optimization problems. A suite of algorithms, ranging from the simplest form to accelerated versions that meet the ideal lower bounds, is presented. The study also investigates the sample and communication complexities associated with these algorithms, they also provided numerical experiments to rationalize the outperformance of proposed algorithms.", "questions": "See above", "limitations": "Not applicable.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper combines both conditional stochastic optimization and federated learning.\n2. The proposed accelerated algorithm matches the lower bound.\n3. The writing and flow of the paper are good in general.", "weaknesses": "1. A bunch of assumptions are imposed here for the theoretical analysis, like Lipschitz continuity, which remedies a bunch of common issues like heterogeneity and client drift. Also it seems like the experiments you consider may not fit with your settings. I believe it is better that authors can further rationalize these assumptions.\n2. As far as I can see, your paper borrows a lot of settings in the previous conditional stochastic optimization literature, which suggests that the strong assumptions you impose may be unavoidable to some extent. But with that, my second issue comes that your work seems to be just a simple extension of existing works in the centralized setting, the algorithm is obvious in my opinion; and the analysis, with the success in the centralized literature, is something that as expected I think. So the current result looks not that exciting to me in terms of significance. Could you please elaborate more on the nontrivial part in the analysis. Thank you.\n\nSome typos:\n1. Several big-O notations are missing, e.g., Line 170, 172, 191, 193, 197, 213, 214\n2. Assumption 3.2, I think the rightmost term should be $L_f^2L_g^2$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for your time and comments. \n\n1. \"A bunch of assumptions are imposed here for the theoretical analysis, like Lipschitz continuity ... \"\n\nAs you mention in Q2, we follow the assumptions in the previous work and do not add any extra assumptions. For Assumption Lipschitz continuous objective (gradient bounded), it is widely used in optimization analysis in related two-level optimization. Many typical centralized two-level optimizations use this assumption [2][3][4] and FL optimization [1] [5].\n\n2. \"... your work seems to be just a simple extension of existing works in the centralized setting .. \"\n\nWe respectfully disagree with the statement. While FCSG is the federated counterpart of the BSGD technique introduced in [1], we take it a step further by integrating momentum technology, leading to our method named FCSG-M. We also conduct a comprehensive theoretical analysis to underscore its good performance in practical scenarios. In contrast to the centralized variance reduction approach (BSpiderBoost) outlined in [1], which mandates a substantial batch size, our proposed variance reduction method (Acc-FCSG-M) operates effectively without such a requirement.\n\nFurthermore, providing theoretical justification for these approaches is a nontrivial part. On the one hand, we not only need to consider the complicated structure of the CSO problem but also need to handle the local update and global update separately.  On the other hand, in the case of Acc-FCSG-M, a direct combination of variance reduction mechanisms like STORM [6] with FL conditional stochastic optimization is not feasible. Unlike the specially calibrated step sizes in STORM [7], our approach involves a more lenient step size condition, necessitating a distinct theoretical proof. The advantage lies in our ability to incorporate variance reduction technology. Concurrently, our results omit logarithmic terms, which are present in the outcomes of STORM-like methodologies [6] [7] (noted in a footnote on page 2 [7]).\n\n\n\n\n[1] Gao, H., Li, J., and Huang, H. On the convergence of local stochastic compositional gradient descent with momentum. In International Conference on Machine Learning, PMLR, 2022\n\n[2] Qi, Q., Luo, Y., Xu, Z., Ji, S., and Yang, T. Stochastic optimization of areas under precision-recall curves with provable convergence. Advances in Neural Information Processing Systems, 2021. https://arxiv.org/pdf/2104.08736.pdf\n\n[3] Wang, M., Fang, E. X., and Liu, H. Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. Mathematical Programming, 161(1):419–449, 2017.\n\n[4] Hu, Y., Zhang, S., Chen, X., and He, N. Biased stochastic first-order methods for conditional stochastic optimization and applications in meta learning. Advances in Neural Information Processing Systems, 33:2759–2770, 2020b.\n\n[5] Huang, F., Li, J., and Huang, H. Compositional federated learning: Applications in distributionally robust averaging and meta learning. arXiv preprint arXiv:2106.11264, 2021\n\n[6] Cutkosky, Ashok, and Francesco Orabona. \"Momentum-based variance reduction in non-convex sgd.\" Advances in neural information processing systems 32 (2019).\n\n[7] Khanduri, Prashant, et al. \"Stem: A stochastic two-sided momentum algorithm achieving near-optimal sample and communication complexities for federated learning.\" Advances in Neural Information Processing Systems 34 (2021): 6050-6061. https://proceedings.neurips.cc/paper/2021/file/3016a447172f3045b65f5fc83e04b554-Paper.pdf\n"}, {"review_id": "AuMKrkwFWP", "reviewer": "Reviewer_jBXN", "summary": "This paper explores different algorithms for optimization problems in the FL setting with 2 layers of stochasticity, where the inner stochasticity is dependent on the outer layer. An interesting application for these algorithms is the MAML problem, where we aim to find a global initialization over clients that carry data from different tasks. Inspired by recent work in stochastic conditional optimization, the authors propose a simple algorithm that uses biased gradient estimates and also propose momentum and variance reduction extensions.\nFor all 3 algorithms, both a convergence guarantee, communication complexity and sample complexity is given.\nFinally empirical evidence is provided on invariant logistic regression and the MAML problem on the Omniglot dataset. They compare against Stochastic Compositional optimization algorithms on the MAML problem.", "questions": "- Does the algorithm allow for partial participation?\n- How does this algorithm perform in the one-task-per-client setting, as [1] and [2] are focussing on?\n- In section 5.2, do you only consider $q=1$, concluding from line 274?\n- Are the experimental results concluded from multiple runs from different initializations?", "limitations": "- As stated before, it is not clear whether the algorithm also applies to partial participation settings. If not, then this is a limitation in the cross-device FL setting", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is understandable for someone who is not familiar with stochastic conditional optimization problems\n- Conditional stochastic optimization in FL is a novel problem formulation that is relevant, especially when viewed from the perspective of MAML, as non-iid-ness of client distributions is still one of the most important challenges to tackle in FL.\n- The convergence analysis is extensive", "weaknesses": "- The proposed method is only evaluated on problems where it is assumes multiple tasks can be sampled at clients instead of one task being assumed for each client. While tackling this problem setting can be a valid contribution, it is questionable to me how realistic this assumed setting is. The authors motivate this assumption around line 258 by saying that 'we can evenly distribute the tasks along the worker nodes', but this is not how the FL setting works in practice. If the proposed algorithm also shows strong performance on the ‘one task per client’ problem, which is a common way of tackling the non-iid-ness issue in FL, this would make for stronger contribution, but experiments on this problem setting are missing.\n- The proposed algorithms seem to assume that all clients participate in each communication round, which is impractical and unscalable in the cross-device FL setting where a high number of clients is typical and the presence of stragglers is probable.\n- The paper contains some grammar mistakes\n- Adding some naive baselines on the MAML problem such as e.g [1][2] and FedAvg would give a better insight in the contribution of the proposed method.\n- No results are given of the eventual ‘converged’ accuracy/loss of the algorithms and the baselines, merely convergence results are shown. Therefore the claim ‘our algorithms outperform baselines by a large margin’ is a little premature in my opinion, especially looking at the convergence results in figure 3.\n- Some interesting questions are not answered by the empirical results, such as:\n    - How do the amount of local updates influence performance?\n    - How do dataset sizes of clients influence performance?\n\n[1] Fallah, A., Mokhtari, A., & Ozdaglar, A. (2020). Personalized federated learning: A meta-learning approach. arXiv preprint arXiv:2002.07948.\n\n[2] Jiang, Y., Konečný, J., Rush, K., & Kannan, S. (2019). Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv:1909.12488.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for your time and valuable review.\n\n> Q1. \"The proposed method is only evaluated on problems where it is assumes multiple tasks can be sampled at clients instead of one task being assumed for each client. \" \n\nOur paper focus on conditional stochastic optimization and FL MAML is an example while Personalized federated learning is different.  In the related work, we mention some work that applies the personalization problem in federated learning in section 2.2. “Stochastic Composition Optimization” and it could be regarded as federated stochastic composition optimization. \n\nWe would like to reiterate one key difference between Stochastic Composition Optimization (SCO) and Conditional Stochastic Optimization (CSO): in the SCO the sampling from the inner randomness and the outer randomness are $\\mathbf{independent}$, while in Problem CSO, they are not. Task sampling is the main challenge in MAML and that is why we need to study Federated Conditional Stochastic Optimization. Some works [2][3] have studied stochastic composition optimization and propose related theoretical analysis. It is not in the scope of our work.\n\n\n> Q2. \"Does the algorithm allow for partial participation?\"\n\nBy adding the clients sampling, our algorithms can consider partial participation. The theoretical analysis of  partial participation will be consider in future work.\n\n\n> Q2. \"How does this algorithm perform in the one-task-per-client setting, as [1] and [2] are focussing on?\"\n\nSeen as Q1. \n\n> Q3.  \"In section 5.2, do you only consider $q = 1$, concluding from line 274?\n\nNo. It is completely different. q is the local training steps in federated learning. This one step means an update step in the inner layer in the MAML task where we follow the setting in [4]. \n\n> Q4. \"Are the experimental results concluded from multiple runs from different initialization?\" \n\nI run each test 3 times and report the average results. We will update the image with the std shadow in the final version. Thanks for your mentions. \n\n\n\n\n[1] Fallah, A., Mokhtari, A., & Ozdaglar, A. (2020). Personalized federated learning: A meta-learning approach. arXiv preprint arXiv:2002.07948.\n\n[2] Gao, H., Li, J., and Huang, H. On the convergence of local stochastic compositional gradient descent with momentum. In International Conference on Machine Learning, PMLR, 2022\n\n[3] Huang, F., Li, J., and Huang, H. Compositional federated learning: Applications in distributionally robust averaging and meta learning. arXiv preprint arXiv:2106.11264, 2021\n\n\n[4] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1126–1135. PMLR, 2017. https://arxiv.org/pdf/1703.03400.pdf"}, {"review_id": "pXlVuctd0S", "reviewer": "Reviewer_pkoz", "summary": "This work proposes the first provably efficient federated conditional stochastic optimization solutions with a biased gradient estimator, including 3 variants: FCSG, FCSG-M (with momentum), and Acc-FCSG-M (with a variance reduction technique). The sample complexity and communication complexity are analyzed for all 3 algorithms. Notably, Acc-FCSG-M achieves the lower bound of sample complexity of the single-machine counterpart. Additionally, the paper demonstrates the empirical performance of these algorithms on robust logistic regression and MAML tasks.", "questions": "I don't have any additional suggestions besides the ones mentioned in the aforementioned weaknesses.", "limitations": "The authors have adequately addressed the limitations of the paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This research represents a pioneering investigation into the nonconvex conditional stochastic optimization (CSO) within the federated learning framework. It presents three algorithms and conducts a comprehensive analysis of their sample and communication complexities. Notably, the Acc-FCSG-M matches the lower bound of sample complexity of single-machine CSO.", "weaknesses": "There are two areas where the paper could be enhanced. Firstly, the introduction should provide a stronger motivation for federated conditional stochastic optimization (CSO). Currently, the paper separately emphasizes the significance of CSO and federated learning, but fails to incorporate specific real-world problems that would necessitate the use of federated CSO. Secondly, the paper should emphasize its technical contribution by briefly highlighting the unique aspects in the proof that are absent in prior works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "I really appreciate your time and comments. I would follow your suggestion and add more details in final version.\n\nThe technical contribution in the proof that are absent in prior works is listed here:\n\nFirst, we are the first work to do federated conditional stochastic optimization and we not only need to consider the complicated structure of the CSO problem but also need to handle the local update and global update separately.  Second, in the case of Acc-FCSG-M, a direct combination of variance reduction mechanisms like STORM with FL conditional stochastic optimization cannot work because the specially calibrated step sizes in STORM. We use different theoretical analyses and our approach involves a more lenient step size condition. At the same time, our results is tighter and does not include logarithmic terms, which are present in the outcomes of STORM-like methodologies [1] [2] .\n\n[1] Cutkosky, Ashok, and Francesco Orabona. \"Momentum-based variance reduction in non-convex sgd.\" Advances in neural information processing systems 32 (2019).\n\n[2] Khanduri, Prashant, et al. \"Stem: A stochastic two-sided momentum algorithm achieving near-optimal sample and communication complexities for federated learning.\" Advances in Neural Information Processing Systems 34 (2021): 6050-6061. https://proceedings.neurips.cc/paper/2021/file/3016a447172f3045b65f5fc83e04b554-Paper.pdf\n\n"}]

Author rebuttal:

Meta-review: Dear authors,

Thank you for submitting your paper focused on federated conditional stochastic optimization. The reviewers appreciated many of your new contributions, including combining the idea from STORM and extending it into a federated setting with a conditional stochastic flavor. Despite this, the reviewers also pointed out several weaknesses in your work, many of which you successfully addressed during the rebuttal phase. Some of the reviewers increased their scores and agreed with acceptance only based on the promise that the final camera-ready version would include the amendments you promised and demonstrated. I would also encourage you to highlight your contributions in the final camera-ready version to demonstrate further that your paper is not just a trivial extension of STORM+CSO for FL.

Best,
AC
Decision: Accept (poster)