TSP: A Two-Sided Smoothed Primal-Dual Method for Nonconvex Bilevel Optimization

We thank reviewer fDFn for your helpful comments and questions.

Theoretical Claims: I find the key contribution of this paper to be its ability to handle nonconvexity in the lower-level problem, which, to my knowledge, has been a significant challenge in bilevel optimization. As claimed by the authors, the proposed algorithm reaches a stationary point satisfying the KKT conditions for both the upper and lower levels under only a weakly convex assumption. This result appears quite strong. Given that lower bound results exist for related problems (such as in minimax optimization, which is a special case considered in this paper), [1] shows that such problems are PPAD-hard (though with additional constraints). Could the authors provide insights into why their approach successfully achieves this result despite these known hardness barriers? Understanding this aspect would further clarify the theoretical significance of the proposed method.

Thank you for your insightful comment. The main source of hardness in solving min-max problems (a special case of bilevel optimization) to the type of stationary points defined in [1] lies in the presence of constraints. In our setting, there are no explicit constraints at either the upper or lower level.

Additionally, we assume that the objective functions are coercive, which helps ensure that the iterates generated by our SPD algorithm remain within a bounded region without requiring additional projection. As a result, this assumption—together with the algorithm design—implicitly guarantees the boundedness of the loss values over the unconstrained domain.

These conditions align exactly with the discussion following Theorem 4.1 in [1], which argues that in such unconstrained settings with bounded loss value, approximate stationary points do exist and can be found in polynomial time. Therefore, our theoretical results do not contradict known hardness barriers; rather, they fall within a subclass of problems where efficient solutions remain tractable.

[1] Daskalakis, Constantinos, Stratis Skoulakis, and Manolis Zampetakis. "The complexity of constrained min-max optimization." Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing. 2021.

Experimental Designs Or Analyses: The experimental design makes sense to me. However, it would be helpful if the authors could explain how the chosen parameters in the experiments align with the theoretical claims. Providing this clarification would strengthen the connection between the theoretical guarantees and empirical results.

Thank you for the helpful comment. The parameter choices (primarily the learning rates) were selected via grid search, following standard practice in the literature. For the stochastic case, we applied a learning rate decay of $1/\sqrt{r}$, which is guided by our theoretical results. The initial learning rates were determined through grid search.

We also apologize for the typo in line 418 of the paper: the learning rate decay was incorrectly written as $1/r$; it should be $1/\sqrt{r}$. We will correct this in the revised version.

Other Comments Or Suggestions:

line 195: formulations (1) and (1)

Typo: the second reference to equation (1) should be (3).

line 215: forgotten period at the end.

Thank you for your careful reading. We will add the missing period in the revised version.

line 255 & 256: should the UL and LL get exchanged?

The terms "UL'' and "LL'' are indeed confusing in this context. We will remove them in the revised version.

We thank reviewer bXuu for your helpful comments and questions.

Claims And Evidence: *My main concern is the gap between the penalty reformulation and the original bilevel problem. It seems that both of the proposed method and the convergence analysis are based on the reformulation. *

Correct. The proposed method addresses a Moreau envelope-based reformulation of the original bilevel optimization problem. The relationship and equivalence between the original and reformulated problems are discussed between lines 194 and 199 of the paper, summarizing the findings of existing works [Gao et al., 2023] and [Liu et al., 2024a].

The equivalence presented in [Theorem A.2, Liu et al., 2024a] for the simplified case where the lower-level problem is unconstrained is further clarified below.

Original Problem:

$$\min_{x, y} f(x, y), \quad \text{s.t.} \quad y \in \mathcal{S}(x)$$

where $\mathcal{S}(x) := \arg\min_y g(x, y)$.

Moreau Envelope-based Reformulation:

$$\min_{x, y} f(x, y), \quad \text{s.t.} \quad g(x, y) - g^\star_\gamma(x) \le 0$$

• When the lower-level function $g(x, y)$ is convex or satisfies the Polyak-Łojasiewicz (PL) condition, these two formulations are equivalent.
• When $g(x, y)$ is weakly convex and $\gamma \in (1, 1/\rho)$, the Moreau envelope-based reformulation becomes equivalent to a relaxed version of the original problem:

$$\min_{x, y} f(x, y), \quad \text{s.t.} \quad y \in \mathcal{S}'(x)$$

where $\mathcal{S}'(x) := \\{ y \mid \\|\nabla_y g(x, y)\\| = 0 \\}$

Given the weak convexity of $g$, it is reasonable to aim for a stationary point, as opposed to requiring a global optimum for the lower-level problem.

Methods And Evaluation Criteria: The main idea of the method design is to reformulate the original bilevel optimization problem into a min-max optimization problem with more feasible conditions using moreau envelope and penalty technique. This makes sense as the reformulated problem is a convex-linear min-max problem is well-studied and easier to solve.

Yes, this approach is conceptually related to the primal-dual or Lagrangian methods. However, the reformulated problem is not a convex-linear min-max problem due to the weak convexity of both the upper- and lower-level objective functions. Therefore, standard techniques for solving convex-linear min-max problems are not directly applicable in this case. The main reasons are as follows:

1. The upper-level objective function $f(x, y)$ are nonconvex rather than convex.
2. The lower-level function $g(x, y)$, even after applying the Moreau envelope, does not necessarily yield a convex constraint.

Experimental Designs Or Analyses: *Regarding the data hyper-cleaning task, as described in section 4, the lower-level objective is the cross-entropy loss function, which is convex. This problem can be solved by methods designed for bilevel problems with convex lower-level problems. This application seems not suitable. *

As noted on line 382, the lower-level objective includes a nonconvex regularization term. While the cross-entropy loss $\ell_{\text{tr}}$ is convex, the presence of the nonconvex regularizer makes the overall lower-level function weakly convex. Therefore, this data hyper-cleaning task remains a suitable example for evaluating methods designed for bilevel problems with weakly convex lower-level objectives.

All the baselines compared in the experiments are in deterministic setting, which makes the experimental results not very convincing. Moreover, it is well known that stochastic methods generalize better than deterministic methods in general. Thus, the claim 'This example further highlights ... optimization' in the representation learning task part may not necessarily hold.

Please kindly note that the numerical experiments on the representation learning task are conducted in a stochastic setting. As mentioned on line 418, a batch size of 32 is used. This setting is applied to all compared algorithms, indicating that they are all implemented stochastically. We will clarify this more explicitly in the revised version to ensure it is clear that all methods are evaluated in a stochastic manner for a fair comparison.

Other Comments Or Suggestions:

In Assumption 3.1, the weak-convexity assumption A2 is unnecessary as it is implied by the smoothness assumption A1.

Thank you for your comment. We will remove Assumption A2 in the revised version.

Using $G(x,y;\xi_i)$ for single sample stochastic estimator, and $\hat{g}(x,y)$ for mini batch stochastic estimator is rather confusing. It would help the readers to understand more easily if such notations are more consistent.

Thank you for the helpful suggestion. We will use $\nabla \hat{G}(x, y)$ to denote the mini-batch stochastic gradient estimator, so that the notation is more consistent throughout the paper.

We thank Reviewer 85Pr for your positive feedback, thoughtful comments, and constructive questions.

Questions: Is $p$ changing for different? Could the authors specify the value of $p$ we should take in the algorithm and main theorem?

In theory, $p$ is a constant and should be chosen to be on the same order as the dual variable, which is upper bounded by a constant. In practice, we can initialize it as a small value and increase it if the algorithm does not sufficiently decrease the loss, following a similar philosophy to tuning the learning rate.

How does the algorithm in this paper compare to ``SLM: A smoothed Lagrangian method for structured nonconvex constrained optimization''? Could the authors provide details?

In comparison to the referenced work, the main differences are as follows:

• (Different class of lower-level problems) The algorithms are designed for different bilevel optimization formulations. Our work focuses on a Moreau envelope-based reformulation, where the lower-level problem is weakly convex. In contrast, the referenced work assumes that the lower-level objective satisfies the PL condition.

• (Stochastic vs. deterministic setting) Our algorithm is designed for the stochastic setting, whereas the referenced method addresses only the deterministic case. This distinction introduces significant challenges, particularly in designing the update rule for the dual variable under stochastic noise.

• (Single-loop vs. double-loop structure) Our algorithm adopts a single-loop structure, which is more suitable for stochastic optimization. The referenced method, on the other hand, uses a double-loop approach.

• (New high-probability error bounds): Our theoretical analysis explicitly quantifies stochastic errors arising from the updates of the upper-level variable, lower-level variable, and auxiliary variables. This includes handling the non-negativity constraint on the dual variable. As a result, the analysis framework is significantly different from that of the referenced work.

The Moreau envelope reformulation is proposed in previous papers. Could the authors provide more details about the novelty compared to these papers?

The main novelties of this work, compared to previous papers that also study the Moreau envelope reformulation, are as follows:

-(Primal-Dual Update Strategy): To the best of our knowledge, this is the first work that applies a primal-dual algorithm to solve the Moreau envelope-based bilevel optimization problem. It is well-established in the optimization community that primal-dual methods often achieve better convergence properties than penalty-based methods. This is largely due to the dynamic update of the dual variable, whereas penalty methods require either a sufficiently large penalty parameter or a monotonically increasing schedule, which can force the learning rate to be very small. This would be particularly problematic when handling multiple lower-level problems.

-(Stochastic Problem Formulation and Algorithm Design) The proposed SPD method tackles a stochastic version of the Moreau envelope reformulation, which significantly broadens its applicability to real-world machine learning problems. In contrast, prior works have largely focused on deterministic settings. The stochastic nature of our formulation introduces substantial challenges in both algorithm design and theoretical analysis, particularly in achieving convergence guarantees. This further differentiates our approach from existing literature.

We thank reviewer Dp7u for your helpful comments and questions.

Claims And Evidence: This statement about equilibrium constraints and bilevel optimization

Will remove that statement.

On lines 69-71

Will remove the statement.

contributions section (lines 167-169)

Will add "approximate''.

Additionally, the reformulation (3) is incorrect.

The reformulation (3) should be:

$$\min_{x,y} f(x,y):=\mathbb{E}_{\xi\sim D\_{UL}} F(x,y;\xi) \\\\ \\quad \textrm{s.t.} g(x,y)- g^{\\star}\_{\gamma} (x,y)\le 0, \quad g^{\\star}\_{\gamma} (x,y):=\min_z \mathbb{E}\_{\xi\sim D\_{LL}} G(x,z;\xi)+\frac{1}{2\gamma}\\|z-y\\|^2.$$

Methods And Evaluation Criteria:

Response:

• Once the two quadratic proximal terms, $\frac{p}{2}\\|x - \hat{x}\\|^2$ and $\frac{p}{2}\\|y - \hat{y}\\|^2$, are introduced with auxiliary variables $\hat{x}$ and $\hat{y}$, the function $K(x, y, \hat{x}, \hat{y}, \lambda)$ becomes strongly convex. This strong convexity ensures the existence of a unique optimal solution for $x$, $y$, and $\lambda$, given any fixed $\hat{x}$ and $\hat{y}$. These optimal solutions are denoted as $x^{\star}(\hat{x}, \hat{y}; \lambda)$ or $\bar{x}^{\star}(\hat{x}, \hat{y})$, and similarly for $y$ (see equations (14) and (15) for details).

• Based on this setup, we can quantify the convergence process of the iterates $(x^r, y^r, \lambda^r)$ by measuring the distance between $x^r$ and $x^{\star}(\hat{x}^r, \hat{y}^r; \lambda^{r+1})$, and similarly between $y^r$ and $y^{\star}(\hat{x}^r, \hat{y}^r; \lambda^{r+1})$. These distances can be further upper bounded by the successive differences of the iterates, i.e., $\\|x^{r+1} - x^r\\|^2$ and $\\|y^{r+1} - y^r\\|^2$, using the standard primal error bound (Zhang & Luo, 2020).

• Therefore, the proximal terms are critical. In particular, the distances to the proximal mappings automatically shrink to zero, ensuring that the algorithm converges to the solution of the optimization problem. As such, the auxiliary variables $\hat{x}$ and $\hat{y}$ play a key role in the algorithm design and in quantifying its convergence to the KKT points of the original problem.

Experimental Designs Or Analyses: *The representation learning task in this paper lacks a clear problem ... the details of how the bilevel structure is formulated and applied remain vague. *

The representation learning task can be formulated as a bilevel optimization problem, as shown below:

$$\min_{x,\{y_i\}} \quad f(x, \{y_i\}) := \mathbb{E}\_{\xi \sim \mathcal{D}^{\text{val}}} \left[\frac{1}{K} \sum_{i=1}^K \ell(x, y_i; \xi)\right] \\\\ \text{s.t.} \quad y_i \in \arg\min_{y'_i} \mathbb{E}\_{\xi_i \sim D^{\text{tr}}_i} \ell(x, y'_i; \xi_i), \quad \text{for } i \in [K].$$

• $x$ represents the shared model parameters, typically corresponding to the common feature encoder or backbone network shared across all tasks.
• $y_i$ denotes the task-specific head parameters, i.e., the final classification layer for task $i$, which is optimized using the task-specific training data $\mathcal{D}^{\text{tr}}_i$.

• This formulation captures the core idea of our approach: learning a shared feature representation ($x$) that enables effective adaptation to multiple downstream tasks through task-specific parameters ($y\_i$).

The dataset split and the specific role of SPD in optimizing meta-learning objectives are not thoroughly explained.

We use a custom split of the MNIST dataset, where we create eight sub-datasets, each corresponding to a single digit class. Each sub-dataset contains 2,500 training samples and 1,500 validation samples.

• We use five sub-datasets for pretraining the model.
• The remaining three sub-datasets are used for meta-learning, which includes both meta-training and meta-testing.

The key role of SPD in this framework is to adjust the dual variables individually for each task-specific head. This flexibility helps to balance the optimization dynamics between the shared and task-specific components, which in turn improves the generalization capability of the learned representation across unseen tasks.

Furthermore, the experimental discussion is limited.

For bilevel optimization problems, our experimental discussion is centered on evaluating the performance of the algorithms at both the upper and lower levels. We also include addtional results w.r.t training accuracy, convergence behavior, and runtime efficiency in the appendix, showing that the proposed SPD algorithm consistently achieves superior performance across these metrics.

Weakness 1

• In the dual update, $h$ is an iterative variable that serves as a stochastic estimate of the gradient of the function $K(x, y, \hat{x}, \hat{y}; \lambda)$ .
• While the symbol $h$ also appears in the primal update, each instance is clearly distinguished by different superscripts (e.g., $f$ or $g$) and subscripts.

Weakness 2

Will remove unnecessary numbers.