🤖 machine learning

Tunable MAGMAX: Preference-Aware Model Merging for Continual Learning

This paper introduces Tunable MAGMAX, a preference-aware model merging framework for continual learning that automatically constructs preference vectors to control task-specific performance, enabling effective adaptation of merged models to diverse deployment environments without manual specification.

Original authors: Kei Hiroshima, Kento Uchida, Shinichi Shirakawa

Published 2026-05-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kei Hiroshima, Kento Uchida, Shinichi Shirakawa

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 have a master chef who has learned to cook five completely different cuisines: Italian, Japanese, Mexican, Indian, and French. In the world of Artificial Intelligence, this chef is a "Continual Learning" model. The problem is, when the chef learns to make perfect Sushi, they might start forgetting how to make a perfect Pizza. This is called "catastrophic forgetting."

Usually, to fix this, researchers try to create one "Super Chef" who is just okay at all five cuisines, aiming for the highest average score. But in the real world, that doesn't always work. A restaurant in Tokyo might only care about Sushi and don't care about Pizza at all. A restaurant in Mexico might need the opposite. They don't want a "Super Chef"; they want a chef who is a Sushi Master but is just "good enough" at the other dishes.

This paper introduces a new tool called Tunable MAGMAX to solve exactly that problem. Here is how it works, using simple analogies:

1. The Problem with the Old Way (MAGMAX)

The previous method, called MAGMAX, was like a "Talent Show Judge." It looked at the chef's skills for every dish and picked the single best move for each step of cooking. If the chef's best move for chopping onions came from their Italian training, it kept that. If the best move for seasoning came from their Japanese training, it kept that.

The result? A great "average" chef. But, the paper notes a flaw: the last cuisine the chef learned often dominated the final recipe. If the chef learned French last, the final dish might taste too much like French food, even if the restaurant wanted Italian. You couldn't easily tell the chef, "Hey, I need you to be 80% Italian and only 20% French."

2. The New Solution: Tunable MAGMAX

The authors propose Tunable MAGMAX. Think of this as giving the restaurant owner a Remote Control with a slider for each cuisine.

The Preference Vector (The Remote): This is a list of numbers that tells the system how much of each "cuisine skill" to keep. If you want a Sushi-focused chef, you turn the "Japanese" slider up and the "Italian" slider down.
The Merging Process: Instead of just picking the single best move for every step, Tunable MAGMAX looks at the "skill budget." If you set the Japanese slider to high, the system ensures that a large number of the chef's moves come from their Japanese training. If you set the Italian slider to low, it only keeps a few Italian moves.

This allows you to build a custom model for a specific environment (like a specific city or hospital) without having to retrain the whole chef from scratch.

3. How Do We Set the Remote? (The "Magic" Part)

You might ask, "How do I know what numbers to put on the sliders? Do I have to guess?"

The paper says no. They created a way to automatically set the remote based on a small sample of the new environment.

The Analogy: Imagine you are opening a new restaurant in a rainy city. You don't need to cook for the whole world; you just need to know what the locals eat.
The Method: The system takes a tiny sample of data from your new location (the "meta dataset"). It then compares this sample to the training data it already has.
- If you have labels (names of dishes): It checks, "Does the new restaurant serve more Tacos or more Sushi?" It calculates the similarity and sets the sliders accordingly.
- If you only have pictures (no names): It uses a mathematical tool called "Optimal Transport" to look at the "shape" of the data. It asks, "Do the pictures of food here look more like the Italian training set or the Japanese one?"

Once it calculates this similarity, it automatically builds the perfect "Preference Vector" (the remote settings) for that specific environment.

4. The Results

The researchers tested this on standard AI benchmarks (like CIFAR-100 and ImageNet-R), which are like standardized cooking competitions.

Control: They showed that by changing the sliders, they could make the model perform incredibly well on specific tasks (like the first or last dish learned) while letting other tasks drop slightly.
Adaptability: When they simulated different "target environments" (mixing different types of data), their automatic method consistently created better chefs than the old methods. It could build a model that was perfectly tuned for a "Rainy City" or a "Sunny City" just by looking at a small sample of data from that city.

Summary

In short, Tunable MAGMAX is a way to take a general AI model that has learned many things and custom-tailor it for a specific job without retraining it. It uses a "preference vector" to decide how much of each learned skill to keep, and it can figure out the right settings automatically by looking at a small sample of the new environment. It's like having a master chef who can instantly adjust their entire menu to perfectly suit the specific tastes of any new neighborhood they open in.

Technical Summary: Tunable MAGMAX

Problem Statement

Continual Learning (CL) aims to train models sequentially on multiple tasks while mitigating catastrophic forgetting (CF). While recent advances in Large Pre-trained Models (LPMs) and model merging techniques like MAGMAX (Maximum Magnitude Selection) have shown effective CL performance, existing methods primarily optimize for average performance across all tasks. They fail to address the practical necessity of constructing models that accommodate varying deployment environments or specific user preferences. For instance, a medical diagnosis system in one hospital may prioritize tumor detection, while another prioritizes cardiac abnormalities, despite both relying on the same base LPM. Current CL techniques, including MAGMAX, do not adequately provide a mechanism to flexibly adjust task-wise performance priorities. Furthermore, MAGMAX's reliance on element-wise maximum magnitude selection often results in the task vector of the last trained task dominating the merged model, potentially leading to suboptimal performance on earlier tasks.

Methodology

The paper proposes Tunable MAGMAX, a preference-aware model merging framework designed to control task-specific performance in CL. The method extends MAGMAX by introducing a preference vector that dictates the number of elements selected from each task vector during the merging process.

Core Mechanism: Tunable MAGMAX

In standard MAGMAX, the merged task vector is formed by selecting the element with the maximum absolute magnitude across all task vectors for each parameter index. Tunable MAGMAX modifies this by enforcing a budget constraint defined by a preference vector $\mathbf{n}_T = (n_1, \dots, n_T)$ , where $\sum n_t = d$ (the total number of model parameters).

Selection Process: The algorithm iterates through tasks in reverse order (from the last task $T$ to the first). For each task $t$ , it identifies unselected parameters where task $t$ 's vector has the largest magnitude among the remaining tasks.
Budget Enforcement: If the number of naturally selected elements exceeds the budget $n_t$ , the excess is randomly discarded to meet the quota. If the budget is not met after the initial pass, the algorithm performs additional alignment iterations (set to $K=2$ in the experiments) to assign remaining unselected parameters to tasks that have not yet reached their budget.
Result: This allows the merged model to be biased toward specific tasks by allocating more parameters to their corresponding task vectors, thereby controlling the trade-off in performance between early and late tasks.

Automatic Preference Construction

To eliminate the need for manual specification of the preference vector, the authors propose a similarity-based mechanism. This leverages a small amount of meta-data ( $D_{meta}$ ) from the target deployment environment to compute similarity with the training datasets ( $D_t$ ).

Feature-Based (OT): If feature sets are available, the similarity is calculated using the Optimal Transport (OT) distance between the embedded features of the training task and the meta-data.
Label-Based: If only label sets are available, similarity is computed based on the inner product of empirical label distributions.
The preference vector $n_t$ is then constructed by allocating parameters proportionally to these similarity scores.

Key Contributions

Tunable MAGMAX Framework: A novel model merging method that regulates the number of elements selected from each task vector via a preference vector, enabling flexible control over task-wise performance in continual learning.
Similarity-Based Automation: A method to automatically construct appropriate preference vectors tailored to target environments by leveraging dataset similarity (via OT distance or label distribution) between small amounts of target meta-data and training datasets.
Empirical Validation: Demonstration through CL benchmarks (CIFAR-100 and ImageNet-R) that the framework successfully constructs merged models that adapt to various target environments, outperforming or matching baseline methods in specific deployment scenarios.

Experimental Results

The authors evaluated Tunable MAGMAX on class-incremental learning settings using CIFAR-100 and ImageNet-R with varying numbers of tasks ( $T \in \{5, 20, 50\}$ ).

Controllability: By varying a scaling coefficient $\alpha$ in the preference vector, the method successfully shifted performance focus. High $\alpha$ values prioritized earlier tasks, while low $\alpha$ values prioritized later tasks, demonstrating the ability to trade off performance between tasks.
Target Environment Adaptation: In experiments simulating diverse target environments (balanced and imbalanced mixtures of tasks), Tunable MAGMAX using the Label-based similarity control consistently achieved higher or comparable average accuracy compared to baselines (MAGMAX, TIES-Merging, Model Soup, and random mixing).
Robustness: The Label-based preference control showed superior robustness across varying numbers of tasks in the target environment compared to the OT-based approach, particularly when the target environment involved a small number of tasks.

Significance and Claims

The paper claims that Tunable MAGMAX offers a practical solution for deploying continual learning models in scenarios where task priorities differ across applications or users. By enabling preference-aware control without the computational cost of retraining from scratch, the method allows practitioners with limited resources to customize merged models for specific deployment needs. The authors note that while the method effectively addresses the need for flexible task prioritization, it relies on the availability of meta-data to approximate the target environment and does not currently account for dependencies between training tasks.