Title: Variational Task Vector Composition
Venue: NeurIPS 2025
Accepted: yes
Keywords: Task vector

Vartiational inference

Spike-and-slab prior

Controllable posterior with gated sampling
Abstract: Task vectors capture how a model changes during fine-tuning by recording the difference between pre-trained and task-specific weights. The composition of task vectors, a key operator in task arithmetic, enables models to integrate knowledge from multiple tasks without incurring significant additional inference costs. In this paper, we propose variational task vector composition (VTVC), where composition coefficients are taken as latent variables and estimated in a Bayesian inference framework. Unlike previous methods that operate at the task level, our framework focuses on sample-specific composition. Motivated by the observation of structural redundancy in task vectors, we introduce a Spike-and-Slab prior that promotes sparsity and aims to preserve the most informative components. To further address the high variance and sampling inefficiency in sparse, high-dimensional spaces, we develop a gated sampling mechanism that constructs a controllable posterior by filtering the composition coefficients based on both uncertainty and importance. This yields a more stable and interpretable variational framework by deterministically selecting reliable task components, reducing sampling variance while improving transparency and generalization. Experimental results demonstrate that our method achieves state-of-the-art average performance across a diverse range of benchmarks, including image classification and natural language understanding. These findings highlight the practical value of our approach, offering a new, efficient, and effective framework for task vector composition.


Peer reviews:
comments: SUMMARY:
This paper proposes a novel framework for variational composition of learned task vectors, aiming to improve both performance and efficiency by reducing redundancy and introducing uncertainty-aware combination coefficients.

STRENGTHS AND WEAKNESSES:
Strengths

1. The structure and main contributions of this paper are clearly described. The authors provide sufficient background and prior knowledge, which helps the audience better understand the context and motivation behind this work.

2. The proposed methodology is conceptually sound and is empirically validated through experiments on real-world datasets.

Weaknesses

1. The contribution of the gated sampling mechanism appears limited, especially if similar techniques have already been introduced in other frameworks. The paper does not clearly explain any unique challenges or novel adaptations specific to applying this approach in the task vector setting.

2. There is a lack of experimental evaluation on the task unlearning scenario.

QUESTIONS:
1. How should we implement gated sampling for task negation?

2. If we have irrelevant task vectors, how does a spike-and-slab prior help reduce such redundancy?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a method for the integration of task-specific knowledge to be generalized across tasks. The authors formulate a probabilistic model for the task vector composition task and propose a spike-and-slab variational inference framework for efficient training and sample-specific inference. The authors also propose a differentiable and deterministic gated sampling algorithm for training their model. The experiments show that their method improves the classification accuracy across datasets. The authors also empirically discussed the benefits of the gated sampling approach.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: The paper is clearly written and presents a good motivation through Fig 2 that there are overlaps in task-specific knowledges and a lot of redundant features. The proposed method is intuitive and easy to follow. The experiments show improved overall results for image classification across different datasets over the competitive approach. They further show accuracy improvement due to a sample-specific gated sampling approach, which was the main methodological novelty.\nWeakness:\nI don’t have much experience in the field of knowledge transfer through the task vector arithmetic. I enjoyed reading the paper and the main competitive approach aTLAS. I think the main limitation of this work is the limited novelty of the method over aTLAS. The method is a sparse summation of the task vectors instead of a non-sparse summation in aTLAS. Methodologically, I am not sure how much this change is novel in the ML literature.\nThe notations are confusing. Is z the activations from a neural network? If so, then how can it be added to the parameters in \\Theta_0 in eq 12?\nI liked the deterministic gated sampling algorithm, but it also brings additional parameters to be trained. How is this then beneficial to use in knowledge transfer, where the objective is to fine-une pretrained weights. If one has additional compute budget, then wouldn’t it be easier and accurate to train separate models (classifiers in this case) on separate data?\n\nQUESTIONS:\nWhat is the pretrained model in Fig 2a? \nWhy does MNIST show a very distinct representation but gain a lot from the knowledge integration in Table 1?\nIn the experiments, I understand that the image is the input and the class is the output. Do you finetune the pretrained weights for all the tasks together with the ELBO loss? If so, then coming back to my earlier point, please provide how many trainable parameters you are adding.\nAt the beginning of the experiments, you mentioned, “We select these datasets to evaluate the effectiveness and scalability of our approach on diverse visual and task scenarios.” I think it is incorrect to say that the approach was evaluated on diverse tasks, since only image classification was chosen.\nPlease clarify the difference between a deterministic and stochastic approach, and Task-level vs sample-specific?\nAre the w/o Gated results in Table 2 from the same model and inference approach as the Table 1 final row? If so, then why are the numbers different?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a variational approach for task vector composition. Given a pre-trained model and target tasks, the main idea consists of sampling the weights’ difference to be added to the pre-trained weights to deliver a stronger set of weights that adapt to multiple downstream tasks. As a baseline, a simple Gaussian parameterization to the posterior is introduced, and then a spike-and-slab prior to enforce sparsity in the parameters to be updated is introduced, to enable task-specific parameter selection. Finally, the paper introduces a gating approach to dealing with the high-dimensionality of the weights and the problems it induces in Monte-Carlo sampling. The experimental setup follows prior work, showing results across 8 datasets, and 3 different ViT models.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n\nThe paper extends the prior work of aTLAS ([42] Zhang et al “Knowledge Composition using Task Vectors with Learned Anisotropic Scaling”, NeurIPS’24) by reformulating the task arithmetic in a variational manner. Using this approach, the task vector can be obtained directly in a sample-specific manner, showing that this approach results in better performance than when using task-specific vectors. \n\nThe main idea of the paper is interesting, well-motivated and developed. Although it requires some extra effort to match the pieces, the paper follows a proper story with a clear idea, that is, to have a sampling-based approach, uncertainty-aware, to producing task-specific vectors that enable the same neural network to deliver strong results without the need of retraining in downstream tasks. \n\n\nWeaknesses:\n\nThe first shortcoming of the paper is its writing, with loose narrative and confusing notation. For instance, the authors refer first to z as the latent representation of a CLIP model, to use it immediately after to refer to the task vector coefficients. The paper needs to be polished with proper notation and definitions presented forehand. For me to understand what the authors were trying to accomplished, I had to go read first the aTLAS paper. The paper should be a bit more self-contained.\n\nWhile the contribution is neat, it is not innovative; applying a Bayesian approach to an existing pipeline is a recurrent idea, and considering that it in this case it comes with a marginal advantage over existing alternatives, it is hard to ponder the contribution of the paper.\n\nI am a bit concerned in regards to the experiments. I find the experimental setup a bit ill-defined. In particular, I miss the following details:\n-\tThe training details: while the authors comment the datasets they use, they barely refer to the “standard benchmarks”. They depart from the pre-trained ViT models and my understanding is that they try to learn the network for parameter sampling based on the corresponding datasets. What it is not clear to me is if they trained a single model using all training partitions from the corresponding datasets, or if they trained a single model for each dataset using the corresponding trainining partitions. In the paper it is suggested at some point that gathering a strong set of task vectors to form a single model can result in strong results, however it doesn’t seem to be the case in this scenario. May the authors comment on that? It would be good to see a one-leave-out set of experiments where possible, considering that the model merely consists of a CLIP vision encoder. \n\n-\tA comparison against prior work like aTLAS is missing. \n\n-\tTable 1 shows that having a sample-specific approach clearly outperforms having a task-specific approach. However, the difference between the deterministic and probabilistic approaches is quite marginal, which is concerning. \n\n\n-\tThe same concern shows in Figure 4: the difference between using a simple Gaussian prior w.r.t. using a Spike-and-Slab prior seems marginal to me, putting under question what the contribution of the latter is in the whole framework. Is it indeed a necessary prior? \n\n-\tWhat is the computational complexity of having a per-sample task vector against having task-specific weights? I understand that the former approach will incur in multiple inference steps that are not necessary in the latter approach. A tradeoff or balance between gain and complexity should be also included in the experimental section.\n\nQUESTIONS:\nMy questions are all included in the Strengths and Weaknesses section.\n\nLIMITATIONS:\nThe authors mention the high GPU demand as a limitation of their approach, although I would like to see, as I expressed before, a proper breakdown on that limitation. Other than that limitations are properly addressed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a variational Bayesian framework for sample-specific task vector composition, improving over traditional task-level methods in model merging. Experiments show consistent performance gains across diverse datasets. The key contributions are: \n1. Variational inference formulation with amortized inference for per-sample coefficient estimation. \n2. A sparse Spike-and-Slab prior to eliminate redundant task components and improve interpretability. \n3. A gated sampling mechanism to deterministically select reliable components based on uncertainty and importance, enhancing inference stability.\n\nSTRENGTHS AND WEAKNESSES:\n## Strength\n**Quality**: The submission is mostly technically sound and the methods are appropriate. This is a complete piece of work. Authors analyzed both strengths and weaknesses fairly in the paper. Claims are well-supported. \n\n**Clarity**: This paper is well-organized. The authors put some important mathematical derivation in the Appendix which is very helpful for understanding of the methodology. \n\n**Significance**: Once the questions below are sufficiently addressed, I foresee the result in this paper will be very impactful for the community as this paper provides a new formulation of task vector addition and it’s easy for researchers to build on top of it. Besides, for scalability, the authors proposed a controllable deterministic solution that enables gradient based training that should be easily adapted into more modern deep learning frameworks. \n\n**Originality**: This paper provides a novel perspective of the task arithmetic from variational inference. It is very different in nature from all other existing papers to choose task vector coefficients and the use of variational inference is well-motivated by the necessity to include uncertainties. This paper further improves the interpretability of task vector coefficients.\n\n## Weaknesses\n**Clarity & Quality**: Some key components of task arithmetic are not clearly explained. Since this paper provides a sampled based task vector composition framework, which is very different from the mainstream literature, it’s better for authors to clearly and mathematically define the formulation of sample based task vector and analyze the difference in a more straightforward way, potentially not by just text. The weaknesses in the clarity also make it hard to justify the soundness of methodology, thus hurting the quality judgement of the paper. See questions section below for more details.\n\nQUESTIONS:\n1. In line 112 and Eq.(1) you listed the formulation of the classic task vector composition. Given a task vector pool $\\mathcal{T}$ and composition coefficients $\\mathbf{z}$, for your sample-level composition, what is the formulation of the merged model? \n2. Some common notations are not clearly defined in the context of task arithmetics, making it difficult to check the soundness of work. When you say input $\\mathbf{x}$ and labels $\\mathbf{y}$, is this for the target task’s input and labels or all the tasks that are used to generate the task vector pool $\\mathcal{T}$? Similarly, in equation (12), what is $\\mathcal{D}_t$? If $\\mathcal{D}_t$ is the target task, is it fair to compare formulation in equation (12) where you are allowed to optimize the coefficient specifically for a task in the task pool? In task addition, isn’t it like we are using a single copy of the task coefficients and evaluating the performance of that single merged model on all tasks for evaluation?\n3. Can you explain the motivation of Eq. 11 more? What can you represent $\\mathbf{z}$ as deterministic functions of input $\\mathbf{x}$? Seems like $\\mu, \\sigma, f$ are all parameterized by $\\Phi$, how are these related? Also, isn’t $\\mathbf{z} = f_\\Phi(x)$ very rare to happen if $\\mathbf{z}$ is continuous?\n4. Can you clarify the definition of gated ratio? This makes Figure 4 hard to read, and in Figure 4, how to read “while activating fewer task coefficients than the Gaussian prior” from the figure? The accuracy comparison is easy to understand. \n5. Table 1: how is deterministic-based sample specific different from probabilistic-based sample specific row? Is probabilistic-based task-level row also [42] or something else? It’s better to emphasize which row is your method. Another suggestion is to mention [42] is one of the SOTA methods so you use it as the baseline.\n\nLIMITATIONS:\nYes, in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces \"Variational Task Vector Composition\" (VTvC), a framework for model merging. The core idea is to reframe task vector composition as a Bayesian inference problem, where the combination coefficients are treated as sample-specific latent variables. The method employs a variational inference (VI) approach to estimate these coefficients. To promote sparsity and select informative components, the authors introduce a Spike-and-Slab prior and a deterministic \"gated sampling\" mechanism. The stated goal is to improve the stability, interpretability, and generalization of the merged model. The framework is evaluated on eight image classification datasets, with its primary comparison being against the aTLAS method.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1.\tOriginality: The paper introduces an interesting perspective by applying variational inference to the task vector composition problem. Framing the combination coefficients as sample-specific latent variables and using a Bayesian framework to estimate them is a departure from conventional deterministic or heuristic approaches in model merging.\n2.\tPrincipled Approach to Sparsity: The use of a Spike-and-Slab prior to induce sparsity is a well-grounded, principled technique. This is a more elegant and systematic approach than simple thresholding and is a clear strength from a methodological standpoint.\n\nWeaknesses:\n1.\tQuality & Reproducibility: The quality of the experimental evaluation is severely compromised by major, unexplained discrepancies in the baseline results.\n     o\tThe paper reports that its primary baseline, aTLAS [1], achieves 77.60% average accuracy on ViT-B/32. However, the original aTLAS paper reports 84.98% for the identical setup.\n     o\tEven more concerning is the 10+ percentage point drop in the zero-shot performance of the base model itself (37.99% in this paper vs. 48.14% in the original aTLAS paper).\n     o\tThese significant and unaddressed discrepancies cast serious doubt on the validity of the entire experimental framework and render the paper's claims of outperformance unconvincing.\n2.\tSignificance & Scope of Evaluation: The paper's contribution is significantly diminished by its insufficient experimental scope and its failure to compare against a sufficient set of relevant, contemporary baselines.\n     o\tLimited Task Diversity: The evaluation is confined to a single task type: task addition for image classification. To demonstrate the general applicability of a model merging framework, it is essential to evaluate it on a broader range of tasks. For example, experiments on task negation, a standard benchmark in task arithmetic, are missing. Furthermore, the evaluation is limited to the vision domain; demonstrating performance on NLP tasks, even on smaller models like RoBERTa, would greatly strengthen the paper's claims of generality.\n     o\tLack of State-of-the-Art Baselines: The comparison is almost exclusively limited to a single method (aTLAS). methods like AdaMerging [2], and data-free approaches like WUDI-merging[3] are completely ignored. This lack of a comprehensive comparison makes it impossible to assess the actual significance and standing of VTvC in the current landscape.\n3.\tQuality & Clarity on Practicality: The paper avoids a critical discussion of its own practical limitations.\n     o\tThe claim of having \"no additional inference costs\" is misleading, as this is standard for most task-vector merging methods.\n     o\tThe paper fails to acknowledge or quantify the prohibitive merging cost of VTvC, which requires a forward pass through an inference network for every sample. This makes the method far more computationally expensive than its peers.\n     o\tDue to this high cost, the method is not scalable to large language models (LLMs). This severe limitation is not addressed, raising questions about the practical utility and significance of the proposed framework.\n[1] Zhang, Frederic Z., et al. \"Knowledge composition using task vectors with learned anisotropic scaling.\" Advances in Neural Information Processing Systems 37 (2024): 67319-67354.\n[2] Yang E, Wang Z, Shen L, et al. Adamerging: Adaptive model merging for multi-task learning[J]. arXiv preprint arXiv:2310.02575, 2023.\n[3] Cheng, Runxi, et al. \"Whoever started the interference should end it: Guiding data-free model merging via task vectors.\" arXiv preprint arXiv:2503.08099 (2025).\n\nQUESTIONS:\n1.\tBaseline Performance Discrepancy: Can you provide a detailed explanation for the substantial performance drop of the aTLAS baseline (77.60% vs. 84.98% in the original paper) and the zero-shot base model (37.99% vs. 48.14%)? My assessment would significantly improve if you can demonstrate that your experimental setup is fair and the reported improvements are based on a correctly reproduced and competitive baseline.\n2.\tComprehensive Comparison and Trade-off Analysis: To establish the significance of VTvC, it is crucial to compare it against other relevant methods like AdaMerging, and a data-free method like WUDI-merging. Could you provide a discussion and, ideally, empirical results that analyze the trade-offs between your method and these baselines across multiple dimensions: (a) final accuracy, (b) merging cost (time/compute), (c) scalability to larger models, and (d) data dependency?\n3.\tBroadening the Experimental Scope: The current evaluation is limited to task addition on CV datasets. To demonstrate the generality of your method, would it be possible to provide results on other standard task arithmetic benchmarks, such as task negation? Additionally, could you show its applicability to the NLP domain, for instance, by conducting merging experiments on small models like RoBERTa-base? This would greatly strengthen the paper's empirical contribution.\n4.\tJustification for Sample-Specific Complexity: The proposed sample-specific approach introduces significant computational overhead during the merging phase. Could you elaborate on the specific scenarios or types of tasks where this high cost is justified and where simpler, task-level merging methods would fail? A clear justification for this complexity would strengthen the paper's contribution.\n\nLIMITATIONS:\nThe authors include a \"Limitations\" section in their conclusion. However, it does not adequately address the most critical limitations of the work. The paper mentions high GPU memory requirements but frames it as a general issue for combining many task vectors. It fails to explicitly state that its own sample-specific mechanism is the primary driver of this high cost and the core reason for its poor scalability compared to other methods. The discussion should be more transparent about the trade-offs, specifically acknowledging that the proposed method is likely not feasible for large-scale models (e.g., LLMs) and quantifying its merging cost relative to more efficient alternatives.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-07-02"}