🤖 machine learning

A Tale of Two Problems: Multi-Task Bilevel Learning Meets Equality Constrained Multi-Objective Optimization

This paper bridges multi-task bilevel learning and equality constrained multi-objective optimization by reforming the former under relaxed convexity assumptions into the latter, for which the authors propose a novel weighted Chebyshev-penalty algorithm that achieves finite-time convergence to KKT-based Pareto stationarity and systematically explores the Pareto front.

Original authors: Zhiyao Zhang, Myeung Suk Oh, Zhen Qin, Jiaxiang Li, Xin Zhang, Jia Liu

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhiyao Zhang, Myeung Suk Oh, Zhen Qin, Jiaxiang Li, Xin Zhang, Jia Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: A Tangled Knot of Two Problems

Imagine you are trying to bake the perfect cake (the Upper Level). But to bake that cake, you first need to find the perfect recipe (the Lower Level).

In the world of machine learning, this is called Bilevel Optimization. You are constantly tweaking the recipe to make the cake better, but the recipe itself changes based on the ingredients you have.

Now, imagine you don't just want one perfect cake. You want a cake that is:

Delicious (Taste)
Healthy (Nutrition)
Cheap (Cost)
Fast to make (Speed)

These goals often fight each other. Making it healthier might make it taste worse or cost more. This is Multi-Task Learning.

The Problem:
For years, scientists could only solve this "Cake vs. Recipe" puzzle if the recipe was very simple and predictable (mathematically, "strongly convex"). But modern AI is messy and complex. The old rules break down when the recipe isn't perfectly predictable. Furthermore, no one had figured out how to solve this when you have multiple conflicting goals (Taste, Health, Cost) all at once in this messy environment.

The Paper's Solution: A Magic Transformation

The authors, Zhiyao Zhang and colleagues, say: "Let's stop trying to untangle the knot directly. Instead, let's turn the whole thing into a different kind of puzzle that we can solve."

They propose a clever trick: The Transformation.

From "Recipe Search" to "Rule Following":
Instead of asking the computer to "find the best recipe," they tell it: "Just make sure the recipe follows the basic laws of physics (mathematically, the first-order stationarity condition)."
- Analogy: Instead of searching for the perfect path through a maze, you just tell the robot, "Don't walk into walls." If it follows that rule, it's on the right track.
The New Puzzle (ECMO):
By making this switch, they turn the messy "Bilevel" problem into a new type of problem called Equality Constrained Multi-Objective Optimization (ECMO).
- Analogy: Imagine you are juggling five balls (the five goals) while standing on a tightrope (the equality constraint). You can't fall off the rope, and you want all five balls to stay in the air as high as possible.

The New Tool: The "Weighted Chebyshev" Penalty

Now that they have this new "Juggling on a Tightrope" problem, they needed a new way to solve it. Existing methods were like trying to juggle by guessing. The authors built a new tool called the WC-Penalty Algorithm.

How it works: Imagine you have a "Worst-Case Scorecard." The algorithm looks at your five balls and asks, "Which one is the lowest?" It then tries to push that lowest ball up.
The "Penalty": If you step off the tightrope (violate the rule), the algorithm hits you with a heavy penalty (a mathematical "ouch"). This forces you to stay on the rope.
The "Weight": You can tell the algorithm, "I care 90% about the red ball and 10% about the blue ball." By changing these weights, the algorithm can explore every possible balance between the goals.

What They Achieved

The paper claims three major victories:

They Defined the Rules of the Game:
Before this, no one knew exactly what "winning" looked like for this specific "Juggling on a Tightrope" problem. They created a new definition called KKT-based Pareto Stationarity.
- Simple term: They wrote the rulebook for what a "good enough" solution looks like when you can't get the perfect one.
They Built a Guaranteed Solver:
They proved mathematically that their new algorithm (WC-Penalty) will definitely find a solution within a certain number of steps. It's not just a guess; it's a guaranteed path to a solution, even in the messy, complex scenarios where old methods failed.
They Closed the Loop:
They showed that if you solve the "Juggling" problem, you have automatically solved the original "Cake and Recipe" problem.

Real-World Tests (The "Cake" Examples)

To prove their method works, they tested it on two real-world scenarios involving Large Language Models (LLMs):

Training a "Reward Model" for AI:
They tried to train an AI to judge other AIs based on five different criteria (Helpfulness, Correctness, Coherence, Complexity, Verbosity). These criteria often conflict (e.g., a very helpful answer might be too long). Their method found a better balance of these traits than previous methods.
Aligning an AI with Human Values:
They tried to fine-tune an AI (Llama) to be helpful, correct, and concise all at once. Again, their method found a better "Pareto front" (the best possible trade-offs) than existing tools.

The Bottom Line

This paper is a bridge. It connects two difficult worlds: Bilevel Learning (nested problems) and Multi-Objective Optimization (conflicting goals).

Old Way: "We can only solve this if the problem is simple and has one goal."
New Way: "We can solve this even if the problem is messy and has five conflicting goals, by turning it into a 'Juggling on a Tightrope' game and using our new penalty-based juggling technique."

They didn't just build a better juggling act; they proved mathematically that their act will never drop the balls, provided you follow their instructions.

Technical Summary: A Tale of Two Problems: Multi-Task Bilevel Learning Meets Equality Constrained Multi-Objective Optimization

Problem Statement
The paper addresses the challenge of Multi-Task Bilevel Learning (MTBL) under Lower-Level General Convexity (LLGC). While bilevel optimization (BLO) is widely used in machine learning (e.g., hyperparameter tuning, meta-learning, and LLM alignment), existing theoretical frameworks and algorithms predominantly rely on the Lower-Level Strong Convexity (LLSC) assumption. This assumption ensures a unique lower-level solution and a well-defined hyper-gradient, but it is often violated in modern, complex machine learning problems where the lower-level objective is merely convex.

Furthermore, existing MTBL works are largely confined to single-task settings or rely on LLSC. When extending MTBL to the LLGC setting, two critical issues arise:

The hyper-gradient becomes ill-defined due to the lack of strong convexity.
The presence of multiple, potentially conflicting upper-level objectives requires a Pareto equilibrium interpretation, complicating the optimality criteria.

Methodology
The authors propose a novel framework that reformulates the LLGC-MTBL problem into an Equality Constrained Multi-Objective Optimization (ECMO) problem. The core methodological steps are:

Reformulation to ECMO:
Instead of solving the nested bilevel problem directly, the authors observe that for a convex lower-level function $g(x, y)$ , the optimality condition $\nabla_y g(x, y) = 0$ is both necessary and sufficient. They reformulate the MTBL problem as:
$\min_{z} F(z) = [f_1(z), \dots, f_S(z)]^\top \quad \text{s.t.} \quad h_i(z) = 0, \forall i \in [q]$
where $z = [x^\top, y^\top]^\top$ and $h_i(z)$ represents the $i$ -th component of the lower-level gradient $\nabla_y g(x, y)$ .
Characterizing Pareto Stationarity for ECMO:
The paper identifies that standard Pareto stationarity definitions for unconstrained multi-objective optimization (MOO) are insufficient for ECMO due to feasibility constraints. The authors establish a new KKT-based Pareto stationarity criterion. They leverage Weighted-Chebyshev (WC) scalarization to create a one-to-one mapping between the Pareto front of the ECMO problem and the solutions of a scalarized problem. By analyzing the Karush-Kuhn-Tucker (KKT) conditions of the WC-scalarized problem, they define a rigorous convergence metric based on the norm of the KKT residual system.
Algorithm Design (WC-Penalty):
To solve the ECMO problem, the authors propose the WC-Penalty algorithm. This method:
- Reformulates the WC-scalarized problem by introducing slack variables to handle the $\ell_\infty$ -norm minimization.
- Converts the equality constraints (both the original constraints and the slack variable constraints) into penalty terms within the objective function.
- Solves the resulting single-level, constrained single-objective problem using Projected Gradient Descent (GD).
- Operates in both deterministic and stochastic settings.

Key Contributions

Theoretical Extension: This work is the first to extend bilevel optimization to the multi-task setting under the relaxed LLGC assumption.
New Problem Formulation: It introduces the ECMO problem as a distinct formulation that had not been theoretically studied prior to this work, specifically addressing the lack of finite-time convergence guarantees in existing literature.
New Convergence Metric: The authors define a novel KKT-based Pareto stationarity metric for ECMO, resolving the ambiguity of optimality conditions in constrained multi-objective settings.
Algorithm and Convergence Rate: They propose the WC-Penalty algorithm and prove a finite-time convergence rate of $O(S T^{-1/2})$ to the KKT-based Pareto stationarity, where $S$ is the number of objectives and $T$ is the total number of iterations. This is the first such guarantee for ECMO problems.
Pareto Front Exploration: By varying the preference vector over the simplex, the method systematically explores the entire Pareto front.

Experimental Results
The authors validate their approach on two multi-objective data weighting tasks:

RLHF Reward Model Training: Using the HelpSteer dataset, the algorithm optimizes data weights to balance multiple human-aligned metrics (helpfulness, correctness, etc.). Results show the WC-Penalty algorithm covers a significantly larger portion of the Pareto front compared to baselines like MOML, MoCo, and FORUM.
Multi-Objective LLM Alignment: The method is applied to fine-tuning Llama-3.2-1B-Instruct (and scaled to 3B and 8B models) to minimize multiple validation losses. The algorithm demonstrates superior performance in hypervolume and $\epsilon$ -metric comparisons against existing MTBL and bilevel baselines.
Synthetic Examples: Additional synthetic experiments confirm the algorithm's ability to recover the full Pareto front in both ECMO and LLGC-MTBL settings.

Significance and Claims
The paper claims to close the loop between foundational optimization frameworks by translating solutions from the ECMO problem directly back to the original MTBL problem. The authors emphasize that their work provides the first finite-time convergence guarantee for ECMO problems, a gap previously unaddressed in the literature. By relaxing the LLSC assumption, the proposed framework extends the applicability of bilevel learning to more complex, real-world machine learning scenarios where strong convexity cannot be assumed. The results suggest that the WC-Penalty method offers both theoretical rigor and practical efficiency in exploring Pareto optimal solutions for multi-task bilevel learning.