Learning to Recover Task Experts from a Multi-Task Merged Model

Conference: ECCV2026
Paper: ECCV Original Paper
Code: https://github.com/BAIKLAB/ReTeX
Area: Model Compression
Keywords: model merging, expert recovery, low-rank parameter offsets, subspace task identification, unseen-task generalization

TL;DR

ReTeX treats model-merging interference as a learnable parameter offset, identifies tasks through SVD subspaces, and recovers experts on demand, reaching 92.2% average accuracy on eight ViT-B/32 tasks without storing full expert checkpoints at inference.

Background & Motivation

Fine-tuning the same pretrained model on different datasets produces experts with distinct capabilities. Deploying every expert requires multiple large checkpoints, whereas averaging their weights introduces interference between task updates. Task Arithmetic and TIES-Merging improve the merged weights, but still produce a fixed parameter set that cannot restore specialized behavior for the current input.

Dynamic merging allows inputs to determine expert composition and therefore relaxes this constraint. However, many approaches still read expert parameters or task components at inference, or require an additional trained router. This paper takes a different perspective: instead of searching only for a shared compromise, it treats the merged model as a common starting point and learns its differences from the available experts. The objective is expert behavior recovery, not mathematical inversion of merging or unsupervised reconstruction from a checkpoint alone.

Existing expert weights provide offline recovery targets; deployment retains the merged weights, recovery modules, and task signatures. Core Idea: infer the required task capability from input features, then generate a low-rank parameter offset that moves the merged model toward the corresponding expert, rather than storing and invoking the entire expert collection.

Method

Overall Architecture

The inputs are an already merged model and the original task experts available offline. ReTeX builds an SVD signature bank for task identification and trains an offset generator for parameter reconstruction. At inference, it extracts input features, identifies the task, generates an offset, adds that offset to the merged weights, and predicts with the recovered model. Simple expert weight averaging is the default starting point; sophisticated merging is not a prerequisite.

These offline preparations use different supervision: the signature bank needs a small set of original training inputs, whereas offset training uses expert parameters alone. Streams from known tasks use hard task IDs. The paper separately introduces soft recovery for unseen-task evaluation, so these settings should not be described as one unchanged inference protocol.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Input["Input and merged model"] --> Identify["Subspace task identification"]
    Memory["Offline reference inputs<br/>SVD signature bank"] --> Identify
    Identify -->|Hard task ID| Recover["Low-rank offset recovery"]
    Identify -->|Unseen-task setting| Soft["Soft mixture recovery"]
    Soft --> Recover
    Experts["Offline expert-parameter supervision"] --> Recover
    Recover --> Output["Recover weights and predict"]

Key Designs

1. Subspace task identification: replace router training with distance to task feature spaces

For each task, the merged model extracts intermediate features from reference inputs taken from the original training set. The features are centered by their empirical mean, and SVD is applied to the resulting matrix. Keeping the leading right singular vectors defines a low-dimensional affine subspace. The default is 64 reference inputs per task; deployment stores the mean and basis, without requiring those raw inputs or complete experts for task identification.

A test feature is centered using a candidate task's mean and projected onto its subspace. A smaller unexplained residual indicates a closer match. The following notation restates the textual definitions accompanying Eqs. (8)-(9), where \(V_{t,k}\) is an orthonormal basis, \(\mu_t\) is the mean, and \(z\) is the test feature:

\[ \mathcal{R}_t(z)=\left\|(I-V_{t,k}V_{t,k}^{\top})(z-\mu_t)\right\|_2, \qquad \hat{t}=\arg\min_t\mathcal{R}_t(z). \]

Here, router-free means that no routing network is trained, not that task selection disappears or that the whole system requires no task data. The signature dimension \(k\) and the parameter-offset rank \(r\) are different hyperparameters. The main text does not establish the default \(k\) for all settings clearly enough to assign a value here.

2. Soft mixture recovery: avoid forcing unseen tasks into one existing expert

For an input from a task absent during merging, even the closest known task may be only partially relevant. In the generalization experiment, projection residuals are converted to soft mixture weights, for example with a softmax over negative residuals. This vector replaces the one-hot task indicator, allowing an input to draw on several known capabilities rather than selecting one nearest expert.

Training only on one-hot indicators does not fully describe the paper's OOD protocol. Recovery training additionally samples mixture vectors from a Dirichlet distribution and uses the corresponding convex combinations of seen expert parameters as targets. Thus, adaptive interpolation includes explicit soft-target training in parameter space; it should not be portrayed as entirely unprepared zero-shot emergence. The text does not specify the Dirichlet concentration or softmax temperature, so no values are invented here.

3. Low-rank offset recovery: share one factor and generate only the task-dependent part

The authors first express expert parameters as an affine transformation of merged parameters and report that scaling terms concentrate near 1, motivating offset-only recovery. For a merged weight matrix of size \(a\times b\), the offset is not generated as a full matrix. Instead, it is factorized into a task-dependent \(a\times r\) factor and a shared \(r\times b\) factor. A learned task embedding enters a single-layer generator for that layer, which predicts the task-dependent factor; the shared factor is learned during recovery training.

Restating Eqs. (10)-(11), the central operation is:

\[ \hat{\theta}^{(l)}_t=\theta^{(l)}_{\mathrm{merge}}+\beta^{(l)}_t, \qquad \beta^{(l)}_t=h^{(l)}(e_t)\beta^{(l)}_s, \qquad r<\min(a,b). \]

The difference is not calculated from an expert checkpoint at inference: its structure has already been learned by the task embedding, generator, and shared factor. Nor is the method merely a generic adapter attached to layer outputs; inference uses recovered parameters. The low-rank constraint makes reconstruction approximate, while sharing a factor reduces the cost of storing a complete offset for each task.

Generator outputs can still be large. ReTeX-E therefore splits the task-dependent \(a\times r\) factor into an \(a\times r_g\) factor and a shared \(r_g\times r\) factor, with \(r_g<r\). This reduces generator-side parameters but does not guarantee faster sample-wise execution: Table 5 reports higher sample-wise latency than ReTeX. The variant changes offset parameterization rather than adding another task-identification route.

A Worked Example

Consider the eight-task vision deployment. Offline, 64 reference inputs are collected for each of SUN397, Cars, RESISC45, EuroSAT, SVHN, GTSRB, MNIST, and DTD, totaling 512 inputs for the signature bank. Recovery training separately uses the eight experts' parameters, not the labels of these 512 inputs for a parameter-reconstruction loss.

When a traffic-sign image arrives, the merged model first extracts its intermediate feature. The identifier compares projection residuals against all eight task subspaces. If GTSRB has the smallest residual, its task embedding generates layer-wise offsets, the corresponding parameters are recovered, and the image is classified. This is an illustrative execution trace, not a reported test image or a measured set of residual values.

For a batch of 64 inputs, the system first predicts every task ID and groups inputs with matching IDs. It recovers parameters once per represented task, processes each sub-batch, and restores the original output order. An eight-task batch needs at most eight recoveries rather than 64 sample-wise recoveries; the actual count depends on how many predicted tasks occur.

Loss & Training

Training fixes the base merged weights and samples task indices uniformly. The generated offsets make the reconstructed weights approximate the corresponding expert, minimizing a parameter-space reconstruction error rather than classification cross-entropy or output distillation. Learned components include task embeddings, layer-wise offset generators, and shared low-rank factors. Text extraction has lost some norm notation in Eq. (12), so this note does not invent its exponent or claim an unverified normalization convention.

The main-text configuration uses Adam for 5000 iterations, learning rate \(2\times10^{-4}\), a cosine schedule, and 600 warm-up steps. The embedding-dimension ablation fixes \(r=256\); this does not establish a universal rank for every main experiment. Original experts remain necessary during offline recovery training and are removed only from inference deployment.

Key Experimental Results

Main Results

Vision experiments use CLIP ViT-B/32, ViT-B/16, and ViT-L/14 with no per-input task ID supplied; each dataset has its own fine-tuned expert. All accuracies below are percentages, and differences are percentage points. Comparators come from the same source table, backbone, and task-count setting.

Source table and setting	ReTeX	Comparator	Comparator accuracy	Gain	Individual experts
Table 1, ViT-B/32, 8 tasks	92.2	WEMoE	90.4	+1.8	92.9
Table 1, ViT-B/32, 20 tasks	89.8	MoW-Merging	79.3	+10.5	91.4
Table 1, ViT-L/14, 20 tasks	93.6	MoW-Merging	81.8	+11.8	94.8
Table 2, ViT-B/32, 30 tasks	91.2	WEMoE	67.1	+24.1	93.1
Table 4, T5-large, 7 NLP tasks	86.6	WEMoE	83.4	+3.2	88.8

Table 1 reports 99.3% normalized accuracy for eight ViT-B/32 tasks, while Table 4 reports 97.5% normalized average accuracy for NLP. These compare accuracy against individual experts, not the fraction of parameters reconstructed or relative error reduction. Table 2 lists ReTeX at 92.0% and 89.4% for eight and twenty tasks, whereas Table 1 lists 92.2% and 89.8%. The main text does not clearly explain this discrepancy, so values retain their source-table attribution rather than being collapsed into one supposedly unique result.

Ablation Study

The following is the efficiency analysis from Table 5, not a causal module-removal ablation. Conditions are an NVIDIA GeForce RTX 3090, CLIP ViT-B/32, eight tasks, and unknown task IDs. Timing includes task identification, offset generation, and the recovered-model forward pass. Grouped rows use \(B=64\); time is measured batch latency divided by batch size, not single-request response latency.

Config	Per-input time (s)	Peak VRAM (GB)	Average accuracy (%)
Weight Averaging	0.0004	1.3	66.3
ID-LoRA, grouped	0.0009	1.9	89.3
ID-EMR-Merging, grouped	0.004	1.9	90.5
ReTeX, sample-wise	0.05	2.5	92.2
ReTeX, grouped \(B=64\)	0.0015	2.7	92.2
ReTeX-E, sample-wise	0.07	1.8	92.0
ReTeX-E, grouped \(B=64\)	0.0015	2.0	92.0

Figure 3a varies the base merged model and recovery rank: a better starting merge helps more under tight low-rank budgets. Figure 3b fixes \(r=256\) and finds diminishing returns as task embedding dimension increases, with strong recovery once its scale is comparable to the task count. The figure supplies no directly readable coordinate table, so visual estimates are not presented as exact ablation measurements.

Key Findings

• Table 3 holds out MNIST and EuroSAT as unseen tasks and uses the remaining six vision datasets as seen tasks, without unseen expert checkpoints for recovery training. Adding ReTeX to AdaMerging raises unseen-task accuracy from 66.8% to 69.1%, a gain of 2.3 percentage points.
• The Table 4 average hides task-level variation: ReTeX scores 48.0% on WSC against 79.2% for its individual expert, only 60.6% normalized accuracy. Recovering more than 95% of expert performance therefore does not apply to every task.
• In Table 5, ReTeX-E trades a small accuracy change for lower VRAM without improving sample-wise latency. Throughput improvements mainly rely on grouping by predicted task and amortizing recovery cost.

Highlights & Insights

• Supervising parameter differences turns post-merge quality repair into learning in weight space. This is useful when experts exist but revisiting large training datasets is inconvenient, provided a small reference-input set remains available for task signatures.
• Separating identification from recovery lets subspace routing signals condition parameter generation rather than merely select stored components. The saving concerns resident expert representations, not a guarantee that deployment uses only one backbone's worth of memory.
• Convex combinations of expert parameters provide a reusable parameter-space augmentation strategy for unseen tasks. Its benefit needs evaluation on actual held-out tasks; linear interpolation alone does not establish a universal generalization guarantee.

Limitations & Future Work

• Recovery requires related experts with corresponding parameter structures and access to their weights during training. It neither directly merges heterogeneous architectures nor recovers forgotten experts from an arbitrary checkpoint without expert supervision.
• Task identification still needs per-task reference inputs and depends on task separability in merged-model feature space. The reported results do not guarantee reliable identification for closely related tasks, drifting streams, or entirely new domains.
• Grouped timing measures amortized throughput cost, which low-batch online services may not realize. Task-switch frequency, within-batch task distributions, and recovered-parameter caching deserve dedicated evaluation.
• The cached main text repeatedly refers to Appendices A and C, but the file ends after the references and contains neither appendix. Damaged equation extraction and small discrepancies between Tables 1 and 2 limit full verification of hyperparameters and normalized-score aggregation; these gaps have not been filled with guesses.

Related Work & Insights

• vs Task Arithmetic / TIES-Merging: The former combines task vectors and the latter resolves conflicts to improve static merged parameters. ReTeX instead restores task-conditioned offsets after merging and can therefore augment either method.
• vs Twin-Merging / WEMoE: These dynamic methods compose task components according to the input, whereas ReTeX generates recovery offsets through compact modules. Both are input-dependent; the main distinction is how expert information is stored and accessed.
• vs SiM / nearest-subspace classification: Subspace residuals as task signals build on existing work. ReTeX connects those signals to expert recovery, rather than originating all SVD-based routing ideas.
• vs LoRA / HyperNetworks: Neither low-rank structure nor conditional weight generation is new by itself. The specific contribution is using these mechanisms for post-merge expert reconstruction supervised by expert parameters instead of task labels.

Rating

• Novelty: 4/5. Post-merge recovery is a clear objective, while low-rank generation and subspace identification have established foundations.
• Experimental Thoroughness: 4/5. Multiple backbones, task counts, NLP, unseen tasks, and deployment costs are covered, but task-level failures and protocol details need further analysis.
• Writing Quality: 3/5. The central pipeline is clear, while cross-table discrepancies and missing appendices in the available cache reduce verifiability.
• Value: 4/5. A practical direction for compact deployment of existing experts, with benefits conditional on batching and task-identification quality.