A Near-Optimal Single-Loop Stochastic Algorithm for Convex Finite-Sum Coupled Compositional Optimization

We thank the reviewer for taking the time to read our paper and greatly appreciate your valuable feedback.

Q1: Can you tell me if the following papers are relevant? (1) Alacaoglu et al. "Forward‑reflected‑backward method with variance reduction", 2021; (2) Alacaoglu and Malitsky "Stochastic Variance Reduction for Variational Inequality Methods" 2022; (3) Zhu et al. "Accelerated Primal-Dual Algorithms for a Class of Convex-Concave Saddle-Point Problems with Non-Bilinear Coupling Term," 2020.

A: Thanks for bring these interesting works to our attention.

• Zhu et al. study the convex-concave min-max optimization problem with a non-bilinear coupling term, which is similar to the reformulation (2) in our work. However, there are key differences in our problem setting:
  • i) In our formulation (2), the problem is separable with respect to each coordinate of the dual variable $y$, whereas the problem in Zhu et al. is not;
  • ii) Zhu et al. focus on the deterministic setting, while our work focuses on the stochastic setting.
• Both Alacaoglu et al. and Alacaoglu and Malitsky study the monotone variational inequality problem with a finite-sum structure, which includes the concave min-max optimization problem with a finite-sum structure as a special case. Although our problem in (2) also has a finite-sum structure, there are two differences:
  • i) Our problem (2) is separable with respect to each coordinate of the dual variable $y$, whereas the problems considered in Alacaoglu et al. and Alacaoglu and Malitsky are not;
  • ii) Alacaoglu et al. and Alacaoglu and Malitsky assume that each summand in the finite sum of losses is directly available. In contrast, our work only requires that each summand be accessible through stochastic oracles.

We will add these discussions in the revision.

Q2: Line 60: "By the convex conjugate" do you mean that this follows from some duality framework? Can you give some references on this?

A: Sorry for the confusion. We refer to that the value of a convex and and lower semi-continuous function $f_i$ on any $u$ in its domain can be represented as $f_i(u) = \max_{y^{(i)}\in\mathcal{Y}_i}\{u^\top y^{(i)} - f_i^*(y^{(i)})\}$, where $f_i^*$ is the convex conjugate of $f_i. The formal definition can be found in Section 12 of Rockafellar (1997).

Rockafellar, R.T., 1997. Convex analysis (Vol. 28). Princeton university press.

Q3: typos: regularizatoin; line 1362: minF(x)

A: Thanks for pointing these out. We will fix them in the next version of our draft.

We thank the reviewer for taking the time to read our paper and greatly appreciate your valuable feedback.

We thank the reviewer for taking the time to read our paper and greatly appreciate your valuable feedback.

Q1: The paper provides strong theoretical results, but there is limited discussion on the sensitivity of ALEXR to hyperparameter choices. How sensitive is ALEXR’s performance to those choices?

A: ALEXR has three hyperparameters: $\eta$, $\tau$, and $\theta$. Although $\eta$ and $\tau$ need to be tuned, the sensitivity analysis below shows that ALEXR is not highly sensitive to small variations of their values. Fixing $\theta = 0.1$ seems to work well across all of our experiments and datasets. Thus, the tuning effort of ALEXR is similar to baselines SOX, MSVR, and OOA.

i) Sensitivity to step sizes

Note that in ALEXR, $1/\eta$ is the step size for the primal variable $x$, while $1/\tau$ is step size for the dual variable $y$. We sweep over different values of step sizes in the Group DRO experiments. The top-tier choices of step sizes are highlighted.

(Adult Dataset, $\alpha=0.1$)

(CelebA Dataset, $\alpha=0.1$, '-' refers to divergence)

As shown in the results above, while the step sizes require tuning for each dataset, they are not highly sensitive to small variations.

ii) Choice of $\psi_i$

Since the GDRO problem with the CVaR divergence has a non-smooth $f_i$, we report the experimental results of ALEXR with the choice quadratic choice $\psi_i(u)=\frac{1}{2}\||u\||_2^2$, as discussed in Line 199 of our paper. Next, we empirically compare the results under the quadratic choice with those under another choice $\psi_i=f_i^*$. We separately tune the step sizes for the choice $\psi_i=f_i^*$.

The results show that $\psi_i(u)=\frac{1}{2}\\||u\\||_2^2$ is indeed a better choice than $\psi_i=f_i^*$ for the non-smooth GDRO problem.

Q2: It would be helpful to present and discuss runtime performance across different algorithms, not just convergence speed in terms of oracle complexity. How does the wall-clock runtime of ALEXR compare to baselines like SOX and MSVR?

A: Next, we report the total wall-clock runtime for the GDRO experiment on the larger dataset CelebA.

The result above shows that runtime of ALEXR is comparable to the baselines while achieving a faster convergence rate.

Q3: Some hyperparameter settings are not explicitly stated (e.g., specific learning rates, batch sizes for all algorithms). A table summarizing the experimental setup would improve reproducibility.

Thank you for the suggestion! The tuned algorithm-specific hyperparameter settings for the GDRO experiments are:

The shared settings for all algorithms in the GDRO experiments are:

Due to character limitations, we will also include settings for the pAUC experiments in the next version of our draft.

We thank the reviewer for taking the time to read our paper and greatly appreciate your valuable feedback.

Q1: Assumption 4 looks strong to me: the authors assume the bounded variance for both zeroth- and first-order oracles. What are the assumption (cf Assumption 4 here) the prior work on stochastic compositional optimization made? Do they also require the bounded variance for both zeroth- and first-order oracles?

A: All prior works on stochastic compositional optimization make the same or stronger assumptions. For example, the seminal works [1, 2] assume both of the following hold: (1) The variance of the zeroth-order oracle is bounded, which is the same as ours. (2) The first-order oracle is either almost surely bounded or has a bounded 2nd moment, which are stronger than the bounded variance assumption in our work. Recent works such as SOX and MSVR all assume the same conditions as ours.

[1] Wang et al., 2017. Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. Math. Program.

[2] Wang et al., 2017. Accelerating stochastic composition optimization. JMLR.

Q2: It seems that the analysis mainly relies on the Lipschitzness of the functions, and the smoothness only helps improve the batch size requirements.

A: That is partially true. First, the smoothness of $f_i$ also allows us to derive the optimal complexity of $O(1/\epsilon)$ for strongly convex problems, which is better than $O(1/\epsilon^2)$ for non-smooth $f_i$. While it is true that the smoothness of the inner functions $g_i$ only affects the batch size scaling in iteration complexity, this is expected, as standard SGD exhibits a similar distinction between smooth and non-smooth problems.

Q3: What is the notation $u_t$ appearing in Section 3.1-3.2 and in the proofs in the appendix? Seems that the authors have refactored the writeup but are not careful to unify the notation.

A: Sorry for the confusion, but $u_t$ is not a notation inconsistency. Eq. (3) in our paper shows that if we set $\psi_i=f_i^*$, Lines 4–8 of our ALEXR algorithm can be rewritten in a form similar to the previous algorithms SOX and MSVR. To be specific, the updated $y_{t+1}^{(i)}$ can be expressed as the gradient $\nabla f_i$ evaluated at a moving-average estimator $u_{t+1}^{(i)} = \frac{\tau}{1+\tau}u_t^{(i)} + \frac{1}{1+\tau}\tilde{g}_t^{(i)}$. The form in (3) based on $u_t$ facilitates the discussion of relationship to SOX and MSVR in Section 3.2.

Q4: In Eq. (5), it seems that there should be $L(x,\bar{y}\_{t+1}) - L(x\_{t+1},y)$ instead of $L(x\_{t+1},y) - L(x,\bar{y}\_{t+1})$ on the right-hand side, based on the discussion and Lemma 5 in the appendix.

A: Thank you for pointing out this typo! We will fix it in the next version of our draft.

Q5: Could the authors explain why they need the extrapolation (Line 6 of Algorithm 1) on the dual side?

A: The extrapolation allows us to enjoy the batch size scaling in the iteration complexity by leveraging the smoothness of $g_i$. Without extrapolation (i.e., $\theta=0$), there is an additional term $\Gamma_{t+1}= \frac{1}{n} \sum_{i=1}^n (g_i(x_{t+1}) - g_i(x_t))^\top (y_*^{(i)}-y_{t+1}^{(i)})$ on the R.H.S. of (18) in Lemma 6. Directly bound this term leads to a worse rate when $g_i$ is smooth. With extrapolation, we can form a telescoping sum of $\Gamma_t$ and cancel it out. Similar idea is used in Lemma 10 for the merely convex case. We have the results in the full version of our paper and will include the discussions in the revision.

Q6: I do not quite get how the authors derive $B,S=O(1)$ from the main body. Could the authors explain their choice of $B$ and $S$?

A: $B,S=O(1)$ in Table 1 means that our algorithm only requires a constant mini-batch size to converge (c.f. impractically large batch size of $O(\epsilon^{-1})$ or $O(\epsilon^{-1})$ in previous work BSGD) and enjoys a parallel speedup using mini-batches.

Q7: For the challenge the authors discussed in Eq. (5), it is usual that the analysis of the algorithm with full updates can be adapted because here one only have randomized coordinate updates to deal with (in contrast to cyclic coordinate methods). Could the authors explain the additional technical challenges here?

A: There are two additional technical challenges here:

• First, our algorithm includes the randomized coordinate stochastic updates for the dual variable, which cause the dependence issue discussed above Theorem 2, making it difficult to directly apply standard convergence analyses used in full-update methods.
• The second challenge comes from the stochastic mirror ascent update for the sampled $y_i$, which make it more involved to derive a batch size scaled convergence rate when the inner functions are smooth.

Q8: Suggestions on citing prior works.

A: Thank you! We will add these references in the next version of our draft.

We thank the reviewer for taking the time to read our paper and greatly appreciate your valuable feedback.

Q1: For the lower bound part, there seems to be no clear definition of the oracle, regarding the layer-wise structure, you need to access $\nabla f_i$ and $\nabla g_i$ separately, also you further need to compute the proximal operator, I expect you can explicitly state the definition of the oracle, with the corresponding input and query output, similar to Zhang & Lan (2020) as you cited.

A: The oracle complexity in lower bound corresponds to the total number of calls of $g_i(x;\zeta^{(i)})$ and $\nabla g_i(x;\tilde{\zeta}^{(i)})$ required by any algorithm within the abstract scheme to achieve an $\epsilon$-accurate solution. In addition, we assume the proximal mapping of $f_i^*$ can be computed when $f_i$ is non-smooth similar to Zhang & Lan (2020). Nevertheless, the number of proximal mappings of $f_i^*$ is bounded by the number of calls of $g_i(x;\zeta^{(i)})$. We will include this clarification in the next version of our draft.

Q2: The lower bound result $\Omega(\max(S/\mu,n)\epsilon^{-1})$ looks weird to me, I think lower bound should be problem-intrinsic and depends on the function behavior, so the dependence on $n$ and $\mu$ (intrinsic problem parameters) looks fine to me; but $S$ should be a user-specified parameter, if $\mu$ is small enough, then you can choose $S$ to have a varing lower bound, which looks weird to me.

A: Thanks for your insightful question. While the lower bound naturally depends on important problem parameters such as $n$ and $\mu$, it is also tied to an abstract update scheme that describes how the space of the iterate evolves across the iterations. Please note that $S$ is a parameter of our abstract scheme presented in pg.6, which covers the typical setting $S=1$ for lower bound analysis of randomized coordinate algorithms (e.g., Ch. 5.1.4 in [1]). We intended to be general so that the abstract scheme covers our ALEXR and some baselines listed in Table 1. Our lower bound $\Omega(\max(S/\mu,n)\epsilon^{-1})$ characterizes the different behaviors of the abstract scheme in two different configurations of the problem-dependent parameters $n$, $\mu$:

• Case I ($\frac{n}{S}>\frac{1}{\mu}$): In this case, the problem's hardness primarily arises from the high dimension of the dual variable $y$. Here, the block-coordinate update of $y$ makes the term $\frac{n}{S\epsilon}$ dominate the iteration complexity, resulting in an oracle complexity bound of $\Omega(n\epsilon^{-1})$. Increasing $S$ makes more coordinates of $y$ are updated per iteration, thereby reducing the number of required iterations to find an $\epsilon$-accurate solution of the whole problem.
• Case II ($\frac{n}{S}\leq \frac{1}{\mu}$): In this case, the main hardness of the problem is from the ill-conditioning of the stochastic proximal update of primal variable $x$. Thus, increasing $S$ (i.e., updating more coordinates of the dual variable $y$) does not improve the leading term $\frac{1}{\mu\epsilon}$ of the iteration complexity, leading to a higher oracle complexity.

Q3: The algorithm for nonsmooth case works for simple $f$ only to compute the proximal mapping. Even though authors provided some examples for illustration, but the restriction may still hinder the applicability of the algorithm in general nonsmooth functions, compared to existing works.

A: We agree with the reviewer. However, existing works on cFCCO that only require computing the gradient of $f_i$, either focus on smooth problems (e.g., SOX and MSVR) or require a gigantic batch size to converge (e.g., BSGD). All of them have suboptimal complexities. Our work is the first to study the more challenging non-smooth cFCCO problem and achieve the optimal rates. In cases where the the proximal mapping of $f_i^*$ is not easily computed, our algorithm can be combined with an approach that solves the proximal sub-problems inexactly (e.g., [2]). Then, the convergence guarantee can be still established.

Q4: What is the potential obstacles if you extend your cFCCO problem from finite-sum to stochastic problem ($\frac{1}{n}\sum_{i=1}^n f_i$ to $\mathbb{E}_i f_i$), or you just intentionally choose the finite-sum problem in the study?

A: If there are an infinite number of $f_i$, our reformulation and algorithm cannot be applied directly, because the dual variable $y$ becomes infinite-dimensional. Hu et al. (2020) considered a more general setting that covers the stochastic problem, but their algorithm requires a large batch size to converge (in both finite-sum and stochastic settings). Given the broad applicability of the finite-sum structure in machine learning, we believe this is an important class to study.

Refs:

[1] Lan, 2020. First-order and Stochastic Optimization Methods for Machine Learning.

[2] Wang et al., 2017. Inexact proximal stochastic gradient method for convex composite optimization.