Towards Minimizing Feature Drift in Model Merging: Layer-wise Task Vector Fusion for Adaptive Knowledge Integration¶

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

LOT Merging (Layer-wise Optimal Task Vector Merging) is a training-free model merging method that explicitly minimizes feature drift—the divergence in intermediate representations between task-specific experts and the merged model—by formulating per-layer optimal task vector computation as a convex quadratic optimization with closed-form solutions. It consolidates multiple fine-tuned models into one unified model using only 16–64 unlabeled exemplars per task, with no gradient-based training.

Problem¶

Existing multi-task model merging methods suffer from a performance gap versus the individual-model upper bound. Parameter-level methods (Task Arithmetic, Ties-Merging, etc.) use heuristic weighting/masking without directly targeting representational fidelity; loss-based methods (AdaMerging, Surgery) close the gap but require costly secondary training. Neither family explicitly minimizes the root cause of degradation: feature drift, the layer-wise divergence between the merged model's representations and those of the task-specific experts.

Method¶

LOT Merging computes a layer-wise optimal task vector T⋆ by minimizing the squared feature error between the merged model and each expert, summed over tasks. The objective is a convex quadratic program with closed-form solutions for all three operation types found in transformers:

Linear layers (matrix multiplication): T⋆ = (Σ XₖᵀXₖ)† · Σ XₖᵀXₖTₖ — a feature-covariance-weighted sum of task vectors, equivalent to projecting each task vector onto the subspace spanned by that task's input features.
Normalization scaling (Hadamard product): per-dimension variance-weighted average of task vectors.
Bias/shift (element-wise addition): simple arithmetic mean of task vectors.

A small calibration set (16–64 samples per task) is run through a single forward pass to collect layer-wise input features Xₖˡ. No backpropagation is needed. The final merged model is Wpre + λT⋆, with λ ≈ 1.2–1.5. The method operates on task vectors (parameter deltas from pretraining), not raw weights, thereby preserving pretrained knowledge when the feature matrix is rank-deficient.

Key Contributions¶

Identification of feature drift (cosine distance between expert and merged model representations) as the primary correlate of post-merge performance degradation, empirically validated across eight vision tasks.
Formulation of layer-wise optimal task vector computation as a convex quadratic problem with analytical closed-form solutions for all three transformer operation types (linear, normalization scaling, bias).
Theoretical analysis showing LOT Merging reduces to task-specific projection in the orthogonal-feature case (zero drift) and variance-weighted averaging in the collinear case (minimal drift).
Demonstration that operating on task vectors rather than raw parameters avoids catastrophic forgetting under low-exemplar regimes, contrasting with RegMean-style direct weight merging.

Results¶

ViT-B/32, 8 vision tasks: LOT Merging achieves 82.7% average accuracy, +4.4% over the second-best training-free method (CAT Merging at 78.3%); best on 6 of 8 tasks.
ViT-L/14, 8 vision tasks: 90.5% average accuracy, +0.9% over second-best; best on 7 of 8 tasks.
BLIP, 6 vision-language tasks (COCO Caption, Flickr30k, Textcaps, OKVQA, TextVQA, ScienceQA): best performance on 5 of 6 tasks; e.g., ScienceQA accuracy 48.24% vs. CAT Merging 46.36%, TextVQA 20.82% vs. CAT Merging 19.69%.
vs. RegMean: +10.9% (ViT-B/32) and +6.8% (ViT-L/14) average accuracy.
Performance stabilizes with as few as 16 exemplars per task; competitive with 64.

Limitations¶

Requires a small exemplar set (≥16 samples per task); not applicable in strictly data-free settings.
Extending closed-form solutions to layers with non-linear operations (e.g., exponential functions in softmax attention) requires additional adaptation not addressed in this work.
Layer-wise optimization treats layers independently; cross-layer interactions in feature drift are not explicitly modeled.
Experiments are limited to CLIP ViT and BLIP architectures; generalization to other VLM families (e.g., decoder-only multimodal LLMs) is not demonstrated.

Relevance to Vision-Language Models¶

LOT Merging directly addresses the challenge of consolidating multiple VLM fine-tunes (demonstrated on BLIP across captioning and VQA tasks) into a single model without any retraining, which is practically significant given the proliferation of task-specific VLM checkpoints. The closed-form approach is architecture-agnostic within the transformer family and scales to ViT-L/14-scale encoders, making it applicable to the vision encoder components common in modern VLMs. For VLM researchers, this provides a principled, low-data alternative to training-based merge methods when deploying models across heterogeneous multimodal tasks. It also advances understanding of how feature representations degrade during merging—a concern increasingly relevant as VLMs are fine-tuned on diverse downstream tasks (captioning, VQA, grounding) and practitioners seek unified deployable models.