Multi-Task Optimization

Follow this topic Ask about this topic

Multi-task optimization studies how to improve several related tasks or objectives together, from training one model across datasets to searching for solutions across related optimization problems. Common approaches share representations or task information, balance losses, and use gradient surgery to reduce interference; multi-objective methods also use scalarization and Pareto-front search to expose trade-offs. Bayesian and evolutionary methods transfer information between tasks, sometimes using task graphs or learned task relationships. A central tension is that sharing can save data and computation but cause negative transfer when tasks differ, while conflicting objectives and mismatched loss scales make reliable balancing difficult.

Latest papers

49 papers7 papers in the last four weeks, level with the four weeks before. 0.0% of all new papers.

May 1, 2026

Near-optimal and Efficient First-Order Algorithm for Multi-Task Learning with Shared Linear Representation

Multi-task learning (MTL) has emerged as a pivotal paradigm in machine learning by leveraging shared structures across multiple related tasks. Despite its empirical success, the development of likelihood-based efficiently solvable algorithms--even for shared linear representations--remains largely underdeveloped, primarily due to the non-convex structure intrinsic to matrix factorization. This paper introduces a first-order algorithm that jointly learns a shared representation and task-specific parameters, with guaranteed efficiency. Notably, it converges in

\widetilde{\mathcal{O}}(1)

iterations and attains a near-optimal estimation error of

\widetilde{\mathcal{O}}(dk/(TN))

, improving over existing likelihood-based methods by a factor of

k

, where

d

,

k

,

T

,

N

denote input dimension, representation dimension, task count, and samples per task, respectively. Our results justify that likelihood-based first-order methods can efficiently solve the MTL problem.

Shihong Ding, Fangyu Du, Cong Fang

State Key Lab of General AI, School of Intelligence Science and Technology, Peking University, China · School of Mathematical Sciences, Peking University, China · Institute for Artificial Intelligence, Peking University, China

Shared Representation Learning Representation Learning Multi-Task Learning Multi-Task Optimization Matrix Factorization

Apr 23, 2026

Multi-Task Optimization over Networks of Tasks

Multi-task optimization is a powerful approach for solving a large number of tasks in parallel. However, existing algorithms face distinct limitations: Population-based methods scale poorly and remain underexplored for large task sets. Approaches that do scale beyond a thousand tasks are mostly MAP-Elites variants and rely on a fixed, discretized archive that disregards the topology of the task space. We introduce MONET (Multi-Task Optimization over Networks of Tasks), a multi-task optimization algorithm that models the task space as a graph: tasks are nodes, and edges connect tasks in the task parameter space. This representation enables knowledge transfer between tasks and remains tractable for high-dimensional problems while exploiting the topology of the task space. MONET combines social learning, which generates candidates from neighboring nodes via crossover, with individual learning, which refines a node's own solution independently via mutation. We evaluate MONET on four domains (archery, arm, and cartpole with 5,000 tasks each; hexapod with 2,000 tasks) and show that it matches or exceeds the performance of existing MAP-Elites-based baselines across all four domains.

Julian Hatzky, Thomas Bartz-Beielstein, A. E. Eiben +1

Vrije Universiteit Amsterdam, Amsterdam, The Netherlands · TH Köln, Cologne, Germany

Evolutionary Optimization Evolutionary Multitask Optimization Quality-Diversity Optimization Multi-Task Optimization

Apr 16, 2026

Analysis of Multitasking Pareto Optimization for Monotone Submodular Problems

Pareto optimization via evolutionary multi-objective algorithms has been shown to efficiently solve constrained monotone submodular functions. Traditionally when solving multiple problems, the algorithm is run for each problem separately. We introduce multitasking formulations of these problems that are an effective way to solve multiple related problems with a single run. In our setting the given problems share a monotone submodular function

f

but have different knapsack constraints. We examine the case where elements within a constraint have the same cost and show that our multitasking formulations result in small Pareto fronts. This allows the population to share solutions between all problems leading to significant improvements compared to running several classical approaches independently. Using rigorous runtime analysis, we analyze the expected time until the introduced multitasking approaches obtain a

(1-1/e)

-approximation for each of the given problems. Our experimental investigations for the maximum coverage problem give further insight into the dynamics behind how the approach works and doesn't work in practice for problems where elements within a constraint also have varied costs.

Liam Wigney, Frank Neumann

Optimisation and Logistics School of Computer Science and Information Technology Adelaide University Adelaide, Australia

Submodular Optimization Evolutionary Optimization Multi-Objective Evolutionary Optimization Runtime Analysis of Evolutionary Algorithms Multi-Task Optimization Evolutionary Multitask Optimization

Apr 13, 2026

BayMOTH: Bayesian optiMizatiOn with meTa-lookahead -- a simple approacH

Bayesian optimization (BO) has for sequential optimization of expensive black-box functions demonstrated practicality and effectiveness in many real-world settings. Meta-Bayesian optimization (meta-BO) focuses on improving the sample efficiency of BO by making use of information from related tasks. Although meta-BO is sample-efficient when task structure transfers, poor alignment between meta-training and test tasks can cause suboptimal queries to be suggested during online optimization. To this end, we propose a simple meta-BO algorithm that utilizes related-task information when determined useful, falling back to lookahead otherwise, within a unified framework. We demonstrate competitiveness of our method with existing approaches on function optimization tasks, while retaining strong performance in low task-relatedness regimes where test tasks share limited structure with the meta-training set.

Rahman Ejaz, Varchas Gopalaswamy, Ricardo Luna +5

Laboratory for Laser Energetics, Rochester, NY, USA · Department of Mechanical Engineering, University of Rochester, Rochester, NY, USA · Hewlett Packard Labs, Santa Clara, CA, USA +2

Bayesian Optimization Multi-Task Optimization

Apr 10, 2026

Delve into the Applicability of Advanced Optimizers for Multi-Task Learning

Multi-Task Learning (MTL) is a foundational machine learning problem that has seen extensive development over the past decade. Recently, various optimization-based MTL approaches have been proposed to learn multiple tasks simultaneously by altering the optimization trajectory. Although these methods strive to de-conflict and re-balance tasks, we empirically identify that their effectiveness is often undermined by an overlooked factor when employing advanced optimizers: the instant-derived gradients play only a marginal role in the actual parameter updates. This discrepancy prevents MTL frameworks from fully releasing its power on learning dynamics. Furthermore, we observe that Muon-a recently emerged advanced optimizer-inherently functions as a multi-task learner, which underscores the critical importance of the gradients used for its orthogonalization. To address these issues, we propose APT (Applicability of advanced oPTimizers), a framework featuring a simple adaptive momentum mechanism designed to balance the strengths between advanced optimizers and MTL. Additionally, we introduce a light direction preservation method to facilitate Muon's orthogonalization. Extensive experiments across four mainstream MTL datasets demonstrate that APT consistently augments existing MTL approaches, yielding substantial performance improvements.

Zhipeng Zhou, Linxiao Cao, Pengcheng Wu +2

Nanyang Technological University · Hong Kong University of Science and Technology (Guangzhou) · School of Artificial Intelligence, Shanghai Jiao Tong University

Orthogonal Optimization Multi-Task Learning Multi-Task Optimization Adaptive Gradient Methods

Apr 4, 2026

Task-Guided Multi-Annotation Triplet Learning for Remote Sensing Representations

Prior multi-task triplet loss methods relied on static weights to balance supervision between various types of annotation. However, static weighting requires tuning and does not account for how tasks interact when shaping a shared representation. To address this, the proposed task-guided multi-annotation triplet loss removes this dependency by selecting triplets through a mutual-information criteria that identifies triplets most informative across tasks. This strategy modifies which samples influence the representation rather than adjusting loss magnitudes. Experiments on an aerial wildlife dataset compare the proposed task-guided selection against several triplet loss setups for shaping a representation in an effective multi-task manner. The results show improved classification and regression performance and demonstrate that task-aware triplet selection produces a more effective shared representation for downstream tasks.

Meilun Zhou, Alina Zare

Department of Electrical and Computer Engineering, University of Florida, Gainesville, FL, USA, 32611 · Artificial Intelligence and Informatics Research Institute, University of Florida, Gainesville, FL, USA, 32611

Representation Learning Multi-Task Learning Geospatial Representation Learning Metric Learning Multi-Task Optimization

Apr 4, 2026

TransGP: Task-Conditioned Transformer-Guided Genetic Programming for Multitask Dynamic Flexible Job Shop Scheduling

Hyper-heuristics have become a popular approach for solving dynamic flexible job shop scheduling (DFJSS) problems. They use gradient-free optimization techniques like Genetic Programming (GP) to evolve non-differentiable heuristics. However, conventional GP methods tend to converge slowly because they rely solely on evolutionary search to find good heuristics. Existing multitask GP methods can solve multiple tasks simultaneously and speed up the search by transferring knowledge across similar tasks. But they mostly exchange heuristic building blocks without truly generating heuristics conditioned on task information. In this paper, we aim to accelerate convergence and enable task-specific heuristic generation by incorporating a task-conditioned Transformer model. The Transformer works in two ways. First, it learns the distribution of elite heuristics, biasing the search toward promising regions of the heuristic space. Second, through conditional generation, it produces heuristics tailored to specific tasks, allowing the model to handle multiple scheduling tasks at once and improving overall optimization efficiency. Based on these ideas, we propose TransGP, a Task-Conditioned Transformer-Guided GP framework. This evolutionary paradigm integrates generative modeling with GP, enabling efficient multitask heuristic learning and knowledge transfer. We evaluate TransGP on a range of DFJSS scenarios. Experimental results show that TransGP consistently outperforms multitask GP baselines, widely used handcrafted heuristics, and the pure Transformer model, achieving faster convergence, superior solution quality, and enhanced robustness.

Meng Xu, Jiao Liu, Hua Yu +1

Job Shop Scheduling Automated Heuristic Design Genetic Programming Combinatorial Optimization Multi-Task Optimization

Apr 4, 2026

Finding Sets of Pareto Sets in Real-World Scenarios -- A Multitask Multiobjective Perspective

Recently, evolutionary multitasking has been employed to generate a ``set of Pareto sets" (SOS) for machine learning models, addressing diverse task settings across heterogeneous environments. This involves creating a repository of compact, specialized solution models that are collectively tailored to each specific task setting and environment, enabling users to select the most suitable model based on particular specifications and preferences. In this paper, we further demonstrate the versatility and applicability of the SOS concept across diverse domains, focusing on three real-world problems: engineering design problems, inventory management problems, and hyperparameter optimization problems. Additionally, as evolutionary multitasking has proven effective in generating the SOS, we investigate the performance of current evolutionary multitasking methods on these real-world problems. Subsequently, we present visualizations of the generated SOS in both decision and objective spaces, complemented by the development of a measurement to gauge the similarity between different Pareto sets corresponding to diverse tasks. Finally, we show that by systematically examining the shifts in Pareto optimal designs across different task settings though the SOS solutions, users can gain deeper understandings on the dynamic interplay between design solutions and their performance in different settings or contexts.

Jiao Liu, Yew Soon Ong, Melvin Wong

College of Computing & Data Science, Nanyang Technological University, Singapore

Evolutionary Multitask Optimization Multi-Objective Optimization Multi-Task Optimization

Mar 24, 2026

Double Coupling Architecture and Training Method for Optimization Problems of Differential Algebraic Equations with Parameters

Simulation and modeling are essential in product development, integrated into the design and manufacturing process to enhance efficiency and quality. They are typically represented as complex nonlinear differential algebraic equations. The growing diversity of product requirements demands multi-task optimization, a key challenge in simulation modeling research. A dual physics-informed neural network architecture has been proposed to decouple constraints and objective functions in parametric differential algebraic equation optimization problems. Theoretical analysis shows that introducing a relaxation variable with a global error bound ensures solution equivalence between the network and optimization problem. A genetic algorithm-enhanced training framework for physics-informed neural networks improves training precision and efficiency, avoiding redundant solving of differential algebraic equations. This approach enables generalization for multi-task objectives with a single, training maintaining real-time responsiveness to product requirements.

Wenqiang Yang, Wenyuan Wu, Yong Feng +1

Chongqing Key Laboratory of Automated Reasoning and Cognition, Chongqing Institute of Green and Intelligent Technology, Chinese Academy of Sciences, 266 Fangzheng Avenue, Liangjiang District, 400714, Chongqing, China. · Chongqing School, University of Chinese Academy of Sciences, 266 Fangzheng Avenue, Liangjiang District, 400714, Chongqing, China.

Constrained Optimization Physics-Informed ML Multi-Task Optimization

Mar 23, 2026

mSFT: Addressing Dataset Mixtures Overfitting Heterogeneously in Multi-task SFT

Current language model training commonly applies multi-task Supervised Fine-Tuning (SFT) using a homogeneous compute budget across all sub-datasets. This approach is fundamentally sub-optimal: heterogeneous learning dynamics cause faster-learning tasks to overfit early while slower ones remain under-fitted. To address this, we introduce mSFT, an iterative, overfitting-aware search algorithm for multi-task data mixtures. mSFT trains the model on an active mixture, identifies and excludes the earliest overfitting sub-dataset, and reverts to that specific optimal checkpoint before continuing. Extensive evaluations demonstrate that mSFT consistently outperforms 4 baselines across 10 benchmarks and 6 base models. Further analysis confirms mSFT maintains robust gains across diverse dataset sizes, task granularities, and is insensitive to its single new hyperparameter (compute budget). Notably, at low compute budget, mSFT can improve performance while lowering training FLOPs. Ultimately, mSFT establishes a practical overfitting-aware algorithm for multi-task SFT that maximizes the potential of models across diverse data mixtures.

Woosung Koh, Jeyoung Jeon, Youngjin Song +4

KAIST AI · Yonsei University · Samsung Advanced Institute of Technology

LLM Fine-Tuning Supervised Fine-Tuning Data Mixture Optimization Multi-Task Learning Multi-Task Optimization

Mar 11, 2026

UniPINN: A Unified PINN Framework for Multi-task Learning of Diverse Navier-Stokes Equations

Physics-Informed Neural Networks (PINNs) have shown promise in solving incompressible Navier-Stokes equations, yet existing approaches are predominantly designed for single-flow settings. When extended to multi-flow scenarios, these methods face three key challenges: (1) difficulty in simultaneously capturing both shared physical principles and flow-specific characteristics, (2) susceptibility to inter-task negative transfer that degrades prediction accuracy, and (3) unstable training dynamics caused by disparate loss magnitudes across heterogeneous flow regimes. To address these limitations, we propose UniPINN, a unified multi-flow PINN framework that integrates three complementary components: a shared-specialized architecture that disentangles universal physical laws from flow-specific features, a cross-flow attention mechanism that selectively reinforces relevant patterns while suppressing task-irrelevant interference, and a dynamic weight allocation strategy that adaptively balances loss contributions to stabilize multi-objective optimization. Extensive experiments on three canonical flows demonstrate that UniPINN effectively unifies multi-flow learning, achieving superior prediction accuracy and balanced performance across heterogeneous regimes while successfully mitigating negative transfer. The source code of this paper will be released on https://github.com/Event-AHU/OpenFusion

Dengdi Sun, Jie Chen, Xiao Wang +1

Key Laboratory of Intelligent Computing & Signal Processing (ICSP), Ministry of Education, School of Artificial Intelligence, Anhui University, Hefei 230601, China · Anhui Provincial Key Laboratory of Multimodal Cognitive Computing, School of Computer Science and Technology, Anhui University, Hefei 230601, China

Physics-Informed ML Multi-Task Learning Computational Fluid Dynamics Multi-Task Optimization

Mar 5, 2026

Asymptotic Behavior of Multi--Task Learning: Implicit Regularization and Double Descent Effects

Multi--task learning seeks to improve the generalization error by leveraging the common information shared by multiple related tasks. One challenge in multi--task learning is identifying formulations capable of uncovering the common information shared between different but related tasks. This paper provides a precise asymptotic analysis of a popular multi--task formulation associated with misspecified perceptron learning models. The main contribution of this paper is to precisely determine the reasons behind the benefits gained from combining multiple related tasks. Specifically, we show that combining multiple tasks is asymptotically equivalent to a traditional formulation with additional regularization terms that help improve the generalization performance. Another contribution is to empirically study the impact of combining tasks on the generalization error. In particular, we empirically show that the combination of multiple tasks postpones the double descent phenomenon and can mitigate it asymptotically.

Ayed M. Alrashdi, Oussama Dhifallah, Houssem Sifaou

Department of Electrical Engineering, College of Engineering, University of Ha’il, P.O. Box 2440, Ha’il, 81441, Saudi Arabia · John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA 02138, USA · Department of Electrical and Electronic Engineering, King’s College London – Strand, London WC2R 2LS – London, UK

Double Descent Neural Network Generalization Multi-Task Learning Multi-Task Optimization Statistical Learning Theory Implicit Regularization

Mar 2, 2026

Can Computational Reducibility Lead to Transferable Models for Graph Combinatorial Optimization?

A key challenge in developing unified neural solvers for combinatorial optimization (CO) is the efficient generalization of models from a given set of tasks to new tasks unseen during initial training. To address this, we first establish a new GNN encoder, which uses a GCON module as a form of expressive message passing together with energy-based unsupervised loss functions. This model achieves highly competitive performance across multiple CO tasks when trained individually on each task. We then leverage knowledge from the computational reducibility literature to propose pretraining and fine-tuning strategies that transfer effectively (a) between MVC, MIS and MaxClique, and (b) in a multi-task learning setting that additionally incorporates MaxCut, MDS and graph coloring. Additionally, in a leave-one-out, multi-task learning setting, we observe that pretraining on all but one task almost always leads to faster convergence on the remaining task when fine-tuning, while avoiding negative transfer. Our findings indicate that learning common representations across multiple graph CO problems is viable through the use of expressive message passing coupled with pretraining strategies that are informed by the polynomial reducibility literature, thereby taking an important step towards enabling the development of foundational models for neural CO. We provide an open source implementation of our work at https://github.com/semihcanturk/COPT-MT.

Semih Cantürk, Thomas Sabourin, Frederik Wenkel +2

Dept. of Computer Science and Operations Research, Université de Montréal · Mila – Quebec AI Institute · Dept. of Mathematics & Statistics, Université de Montréal +2

Neural Combinatorial Optimization Transfer Learning Graph Neural Networks Multi-Task Optimization

Feb 16, 2026

DeepMTL2R: A Library for Deep Multi-task Learning to Rank

This paper presents DeepMTL2R, an open-source deep learning framework for Multi-task Learning to Rank (MTL2R), where multiple relevance criteria must be optimized simultaneously. DeepMTL2R integrates heterogeneous relevance signals into a unified, context-aware model by leveraging the self-attention mechanism of transformer architectures, enabling effective learning across diverse and potentially conflicting objectives. The framework includes 21 state-of-the-art multi-task learning algorithms and supports multi-objective optimization to identify Pareto-optimal ranking models. By capturing complex dependencies and long-range interactions among items and labels, DeepMTL2R provides a scalable and expressive solution for modern ranking systems and facilitates controlled comparisons across MTL strategies. We demonstrate its effectiveness on a publicly available dataset, report competitive performance, and visualize the resulting trade-offs among objectives. DeepMTL2R is available at https://github.com/amazon-science/DeepMTL2R.

Chaosheng Dong, Peiyao Xiao, Yijia Wang +1

Amazon USA

Learning to Rank Multi-Task Learning Multi-Task Optimization

Dec 28, 2025

Multi-Task Learning for Metal Alloy Property Prediction: An Empirical Study of Negative Transfer and Mitigation Strategies

Multi-task learning (MTL) in materials science relies on the assumption that physically related properties share learnable representations. We challenge this assumption using a 54,028-sample metal alloy dataset exhibiting extreme task-level imbalance. Our results reveal a striking dichotomy: MTL significantly degrades regression performance for resistivity and hardness but improves classification recall for amorphous-forming ability. We trace this divergence to mismatched functional forms--such as resistivity's polynomial dependence versus hardness's complex interactions--which cause severe gradient misalignment during optimization. Evaluating Deep Imbalanced Regression techniques, we find that projecting conflicting gradients (PCGrad) recovers single-task performance, while combining label distribution smoothing with gradient normalization achieves the best overall balance. Consequently, we propose a strategic framework: utilize independent models for high-precision characterization, but employ MTL for high-throughput screening where recall is paramount. These findings support a "materials property clustering" hypothesis, suggesting that distinct physical mechanisms require specialized optimization strategies to overcome negative transfer.

Sungwoo Kang

Department of Electrical and Computer Engineering, Korea University, Seoul 02841, Republic of Korea

Multi-Task Learning Materials Property Prediction Negative Transfer Multi-Task Optimization Class-Imbalanced Learning

Sep 21, 2025

Graph Coloring for Multi-Task Learning

When different objectives conflict with each other in multi-task learning, gradients begin to interfere and slow convergence, thereby potentially reducing the final model's performance. To address this, we introduce SON-GOKU, a scheduler that computes gradient interference, constructs an interference graph, and then applies greedy graph-coloring to partition tasks into groups that align well with each other. At each training step, only one group (color class) of tasks are activated, and the grouping partition is constantly recomputed as task relationships evolve throughout training. By ensuring that each mini-batch contains only tasks that pull the model in the same direction, our method improves the effectiveness of any underlying multi-task learning optimizer without additional tuning. Since tasks within these groups will update in compatible directions, multi-task learning will improve model performance rather than impede it. Empirical results on six different datasets show that this interference-aware graph-coloring approach consistently outperforms baselines and state-of-the-art multi-task optimizers. We provide extensive theory showing why grouping and sequential updates improve multi-task learning, with guarantees on descent, convergence, and the ability to accurately identify what tasks conflict or align.

Santosh Patapati, Ian Noronha

Stony Brook University AI Innovation Institute, Stony Brook, NY 11794, USA · Purdue University, West Lafayette, IN 47907, USA

Gradient Interference Multi-Task Learning Multi-Task Optimization

Jul 16, 2025

Learning Task Mixtures from Task Affinities: A Probabilistic Graphical Model for Supervised Fine-Tuning

Supervised fine-tuning performance for large language models depends strongly on how training budget is distributed across a heterogeneous set of tasks. In practice, mixtures are often fixed using simple heuristics (e.g., uniform or size-proportional sampling) that ignore task interactions, which can hurt transfer and waste budget on redundant sources. We introduce TaskPGM, a framework for learning continuous task mixtures via an energy-based model over tasks. Tasks form the nodes of a Markov random field: unary potentials capture per-task utility, and pairwise potentials encode inter-task relationships using behavioral divergences computed from predictive distributions of single-task fine-tuned models (e.g., Jensen--Shannon divergence and pointwise mutual information). Optimizing this objective yields mixtures that balance coverage against redundancy. We show that the resulting set function is weakly submodular under budget constraints, enabling approximation guarantees for discrete selection variants. Across multiple model families (LLaMA-7B, Qwen2-7B) and evaluation suites (BIG-Bench Hard), TaskPGM improves over standard mixing strategies and provides interpretable structure over task interactions.

Prateek Chanda, Saral Sureka, Parth Pratim Chatterjee +3

IIT Bombay · IBM Research · Red Hat AI Innovation +1

Supervised Fine-Tuning Multi-Task Learning Probabilistic Graphical Models Multi-Task Optimization

Mar 11, 2025

( $θ_l, θ_u$ )-Parametric Multi-Task Optimization: Joint Search in Solution and Infinite Task Spaces

Multi-task optimization is typically characterized by a fixed and finite set of tasks. The present paper relaxes this condition by considering a non-fixed and potentially infinite set of optimization tasks defined in a parameterized, continuous and bounded task space. We refer to this unique problem setting as parametric multi-task optimization (PMTO). Assuming the bounds of the task parameters to be (

\boldsymbolθ_l

,

\boldsymbolθ_u

), a novel (

\boldsymbolθ_l

,

\boldsymbolθ_u

)-PMTO algorithm is crafted to operate in two complementary modes. In an offline optimization mode, a joint search over solution and task spaces is carried out with the creation of two approximation models: (1) for mapping points in a unified solution space to the objective spaces of all tasks, which provably accelerates convergence by acting as a conduit for inter-task knowledge transfers, and (2) for probabilistically mapping tasks to their corresponding solutions, which facilitates evolutionary exploration of under-explored regions of the task space. In the online mode, the derived models enable direct optimization of any task within the bounds without the need to search from scratch. This outcome is validated on both synthetic test problems and practical case studies, with the significant real-world applicability of PMTO shown towards fast reconfiguration of robot controllers under changing task conditions. The potential of PMTO to vastly speedup the search for solutions to minimax optimization problems is also demonstrated through an example in robust engineering design.

Tingyang Wei, Jiao Liu, Abhishek Gupta +2

College of Computing and Data Science, Nanyang Technological University, Singapore · School of Mechanical Sciences, Indian Institute of Technology, Goa, India · Singapore Institute of Manufacturing Technology (SIMTech), Agency for Science, Technology and Research, Singapore +1

Evolutionary Optimization Multi-Task Optimization

Dec 11, 2024

Variational Model Merging for Pareto Front Estimation in Multitask Finetuning

Pareto fronts are useful to find good task-mixing strategies for multitask finetuning, but they are also costly to compute. To reduce costs, recent works have used existing model merging methods to help train cheap surrogate models to estimate the Pareto fronts. However, no work has yet considered designing new model-merging methods to directly, and provably, improve the quality of Pareto fronts. Here, we fill this gap by proposing a new Bayesian approach called Variational Model Merging. In this approach, existing model-merging methods are obtained as special cases of "posterior-merging" when Gaussian posteriors are used and new model-merging strategies can be derived by using non-Gaussian posteriors. Our main theoretical result is to show that more flexible posteriors necessarily yield better estimates of Pareto fronts. For instance, a Pareto front estimate obtained by merging full-Gaussian posteriors is expected to be better than that obtained by using isotropic Gaussian posteriors. We validate the theory through extensive empirical results on vision and language transformers where better Gaussian families consistently yields better or comparable Pareto fronts. Our work is a rare instance where Bayesian ideas are used to improve Pareto analysis.

Hugo Monzón Maldonado, Nico Daheim, Thomas Möllenhoff +2

RIKEN Center for AI Project, Japan · Ubiquitous Knowledge Processing Lab (UKP Lab), Department of Computer Science, Technical University of Darmstadt and National Research Center for Applied Cybersecurity ATHENE, Germany · National Research Center for Applied Cybersecurity ATHENE, Germany

Bayesian Inference Multi-Task Learning Multi-Task Optimization Language Model Merging Variational Inference Multi-Objective Optimization