🤖 machine learning

CoAction: Cross-task Correlation-aware Pareto Set Learning

The paper proposes CoAction, a Cross-task Correlation-aware Pareto Set Learning framework that utilizes a task-aware Transformer to simultaneously solve multiple multi-objective optimization problems by leveraging inter-task correlations for improved efficiency and performance.

Original authors: Xinyue Chen, Yingxuan Liang, Yiqin Huang, Chikai Shang, Hai-Lin Liu, Fangqing Gu

Published 2026-06-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xinyue Chen, Yingxuan Liang, Yiqin Huang, Chikai Shang, Hai-Lin Liu, Fangqing Gu

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

Imagine you are a master chef trying to perfect a menu. In the world of Multi-Objective Optimization, you aren't just trying to make one dish taste good; you are trying to balance conflicting goals, like making a dish that is both spicy and sweet, or cheap and high-quality. There is no single "perfect" dish; instead, there is a whole spectrum of "best possible" compromises, known as the Pareto Set.

Traditionally, if you wanted to learn how to make the perfect spicy-sweet balance for a burger, you'd train a robot chef on burgers. If you then wanted to learn the perfect balance for a pizza, you'd have to fire the burger robot, build a whole new robot, and train it from scratch on pizza. This is slow, expensive, and ignores the fact that cooking a burger and a pizza share some similar skills (like chopping onions or managing heat).

This paper introduces CoAction, a new way to train a single "super-chef" robot that can learn to make the perfect compromises for many different dishes (tasks) at the same time.

Here is how it works, broken down into simple concepts:

1. The Problem: One Robot Per Job

Existing methods (called PSL) are like having a separate robot for every single problem. If you have 7 different engineering design problems (like designing a bridge, a rocket, or a car), you need 7 different models. This wastes time and money, and it misses the chance to let the "rocket robot" learn from the "car robot."

2. The Solution: The "Name Tag" System (Task Embedding)

CoAction solves this by giving every task a unique digital name tag (called a task embedding).

How it works: Imagine you have one giant kitchen. Instead of building a new kitchen for every recipe, you give every recipe a specific colored apron.
The Magic: The robot wears the "Rocket Apron" when making rockets and the "Bridge Apron" when making bridges. Because the robot knows which apron it's wearing, it knows which rules to follow. But because it's the same robot, it can also learn general cooking skills (like how to handle heat) that apply to both rockets and bridges. This is Knowledge Sharing.

3. The Brain: The Transformer (The "Super-Listener")

To make sure the robot understands how these different tasks relate to each other, CoAction uses a special brain architecture called a Transformer (the same technology behind modern AI chatbots).

The Analogy: Think of a traditional robot as a person who listens to one instruction at a time. The Transformer is like a conductor in an orchestra who can hear every instrument (every task) simultaneously.
The Benefit: It uses a mechanism called Self-Attention. This allows the robot to look at the "Rocket" task and say, "Hey, this part of the rocket design is actually very similar to that part of the bridge design." It connects the dots between different problems to learn faster and better.

4. The Results: Faster and Smarter

The authors tested this "Super-Chef" on 7 different problems, ranging from math puzzles to real-world engineering challenges (like designing a rocket injector).

Speed: Because they trained one model for all 7 problems instead of 7 separate models, they saved about 27% of the time.
Quality: In most cases, the single "Super-Chef" actually made better compromises than the individual robots did. It found a wider variety of solutions (covering more of the "menu") and was very stable.
The "Hard" Test: The biggest win was on a very complex 3-goal problem (the rocket injector). The old method struggled to find good solutions, but the CoAction robot, thanks to its ability to listen to all tasks at once, found a much better, more complete set of solutions.

Summary

CoAction is like upgrading from having 7 separate, isolated apprentices to having one highly skilled master chef who wears different hats for different jobs. By using a special "name tag" system and a brain that can listen to all jobs at once, this system learns faster, saves money, and produces better results than trying to solve each problem alone.

Technical Summary: CoAction – Cross-task Correlation-aware Pareto Set Learning

1. Problem Statement

Multi-objective optimization (MOO) seeks to find a set of trade-off solutions (the Pareto set, PS) for conflicting objectives. While traditional evolutionary algorithms (e.g., NSGA-II, MOEA/D) generate finite solution sets per run, they struggle with real-time decision-making where preferences change dynamically. Pareto Set Learning (PSL) addresses this by training neural networks to map preference vectors to Pareto optimal solutions in real-time.

However, existing PSL methods suffer from a fundamental limitation: they are designed to solve a single MOO problem at a time. In scenarios involving multiple related optimization problems (Multi-Task Learning, MTL), this paradigm requires training a separate model for each task. This approach incurs linearly scaling computational costs and fails to exploit the latent structural similarities and shared knowledge across related tasks. The paper identifies the need for a unified framework that can simultaneously learn Pareto sets for multiple tasks while modeling inter-task correlations to improve efficiency and generalization.

2. Methodology: The CoAction Framework

The authors propose CoAction (Cross-task Correlation-aware Pareto Set Learning), a framework designed to learn Pareto optimal solution sets across multiple optimization problems within a single model. The methodology integrates three core components:

2.1 Task Embedding Mechanism

To enable a single model to distinguish between different tasks while sharing parameters, CoAction employs a task embedding scheme. Unlike learnable embedding vectors that scale poorly with the number of tasks and may fail to generalize to unseen task counts, CoAction utilizes sinusoidal positional encoding inspired by Transformer architectures.

Each task $t$ is assigned a unique $d$ -dimensional embedding vector $e_t$ .
The embedding is constructed using sine and cosine functions of the task index $t$ , ensuring that tasks with similar indices have similar embeddings (smoothness) while remaining distinguishable (injectivity).
This embedding is concatenated with the preference vector $\lambda(\theta)$ and zero-padding to form the network input, allowing the model to generate task-specific decision variables.

2.2 Transformer Backbone

The framework utilizes a Transformer encoder as its backbone architecture.

Self-Attention: The self-attention mechanism is employed to capture complex global dependencies among tasks. This allows the model to dynamically weigh the influence of different tasks during training, facilitating the modeling of inter-task relationships.
Architecture: The network consists of a linear projection layer, positional encoding, a single-layer encoder block (4 attention heads, 128 hidden dimension), a weighted pooling module, and task-specific output heads.
Parameter Sharing: The backbone parameters are shared across all tasks, while task-specific output heads allow for necessary adjustments, mitigating catastrophic forgetting.

2.3 Multitask Collaborative Training

The training process involves a collaborative algorithm (Algorithm 1) where:

Tasks are sampled randomly during each iteration.
Preference vectors are sampled uniformly from the positive orthant of the unit hypersphere.
The model is optimized using a Hypervolume Maximization-based loss function (PSL-HV1). This loss function maximizes a surrogate form of the Hypervolume (HV) indicator, encouraging the model to approximate the entire Pareto front.
Gradient clipping is applied to ensure training stability.

3. Key Contributions

The paper outlines the following primary contributions:

Unified Framework: The proposal of the CoAction framework, which leverages task embeddings, shared network parameters, and a Transformer backbone to model inter-task relationships. This enables a single model to learn Pareto sets for multiple problems simultaneously.
Efficiency and Expressiveness: The framework reduces computational costs by eliminating the need for separate models per task while enhancing model expressiveness through the self-attention mechanism.
Comprehensive Evaluation: Extensive experiments conducted on multitask test suites comprising both benchmark problems (ZDT, VLMOP) and real-world engineering applications (Truss, Hatch Cover, Rocket Injector designs).

4. Experimental Results

The authors evaluated CoAction against a single-task PSL baseline and an MLP-based multitask baseline using three metrics: Hypervolume (HV), Range, and Sparsity.

Performance vs. Single-Task Baseline:
- CoAction achieved competitive or superior performance in HV and Range on most problems (e.g., ZDT1, ZDT2, RE21, RE24).
- On the tri-objective RE37 problem, while HV differences were marginal, the framework demonstrated robustness.
- Computational Efficiency: CoAction reduced training time by 27% (270.13s vs. 370.25s) compared to the single-task baseline, despite training on seven tasks simultaneously.
- Stability: The framework exhibited low variance in HV across runs, indicating stable convergence and a lack of significant catastrophic forgetting.
Ablation Study (Transformer vs. MLP):
- The Transformer backbone significantly outperformed the MLP backbone, particularly on the complex tri-objective RE37 problem.
- On RE37, the Transformer achieved a 12.6% higher HV and a 314% higher Range compared to the MLP.
- The self-attention mechanism was credited for capturing global dependencies, leading to better coverage of complex Pareto fronts.
Extended Evaluation:
- On the bbob-biobj test suite, CoAction demonstrated robust generalizability, achieving consistently high HV values across five different task pairings.
Statistical Analysis:
- Wilcoxon signed-rank tests confirmed statistically significant improvements in HV and Range for CoAction over the single-task baseline on several problems (ZDT1, ZDT2, RE24).
- The Transformer backbone showed significant advantages over the MLP backbone on RE37 and ZDT1/2.

5. Significance and Claims

The paper claims that CoAction addresses the scalability and efficiency bottlenecks of existing PSL methods in multi-task settings. By integrating task embeddings with a Transformer encoder, the framework successfully:

Leverages Cross-Task Knowledge: It exploits shared structures across related optimization problems, improving generalization and solution diversity.
Reduces Computational Overhead: It eliminates the linear scaling of computational costs associated with training separate models for each task.
Handles Complex Geometries: The self-attention mechanism proves particularly effective for problems with complex, high-dimensional Pareto front geometries (e.g., RE37), where traditional MLP backbones struggle to maintain solution distribution.

The authors conclude that CoAction provides a principled and effective approach for multi-objective multi-task optimization, validating its utility through both benchmark and real-world engineering applications.