Variational Task Vector Composition

We appreciate the reviewer’s feedback and for highlighting the clear structure and robust empirical validation of our work. Below, we address the questions and concerns.

W1: The gated sampling mechanism seems limited in contribution. Similar ideas may exist elsewhere, and the paper doesn’t explain what’s uniquely challenging or new about using it in task vector composition.

We agree that gating mechanisms have been explored in other works. However, applying such mechanisms to task vector composition introduces several unique challenges that our method addresses:

• First, task vectors operate in a high-dimensional parameter space, where conventional Monte Carlo sampling leads to high variance and instability[1]. Our gated sampling mechanism avoids this by introducing a deterministic, continuous gating function, which allows gradient-based optimization while maintaining sparsity and stability.
• Second, unlike existing gating approaches, our mechanism jointly considers token-level uncertainty and coefficient magnitude—a design tailored to the uncertainty-aware and sparse nature of task composition.
• Finally, we show that this design improves both efficiency and generalization, especially when combined with Spike-and-Slab priors (see Figure 4 and Appendix C.3), which to our knowledge has not been explored in prior task arithmetic literature. We will clarify these contributions and motivations more explicitly in Section 4.3.

W2/Q1: Task negation

We agree that evaluating task negation is important for demonstrating the scalability and flexibility of our framework. To address this, we have added new experiments on task unlearning. Specifically, we introduce an asymmetric gated sampling mechanism that sets a lower threshold for the coefficients of target (to-be-forgotten) tasks, while preserving important components of control tasks. This allows for effective suppression of target task influence without compromising control task performance. The results on the ViT-B/32 backbone are summarized below:

These results show that our method achieves stronger forgetting on the target task while maintaining comparable performance on the control tasks. This demonstrates the method’s applicability to both task addition and task negation scenarios.

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

When task vectors contain irrelevant vectors or blocks for the current task, the spike-and-slab prior reduces redundancy by pushing unimportant coefficients to exactly 0 through probabilistic modeling. This achieves structured sparsity, reduces redundancy, avoids negative transfer, and lowers computational and storage costs.

• The spike-and-slab prior models each combination coefficient as a mixture of point mass and Gaussian: $p(\mathbf{z})=(1-\pi)\delta_0(\mathbf{z})+\pi\mathcal{N}(0,\sigma^2)$, here $\delta_0$ is the Dirac distribution at $0$ (the "spike"), and $\mathcal{N}$ is the "slab". This means the prior prefers most coefficients to be $0$, with only a few truly useful dimensions falling in the continuous Gaussian part. As a result, dimensions irrelevant to the current sample/task are naturally pushed to $0$ during learning, effectively "turning off" redundant task vectors or their sub-blocks.
• In variational inference, each coefficient $\mathbf{z}_i^j$ has a binary indicator variable $\omega_i^j$: $\omega=1$ means keeping that dimension, $\omega=0$ means setting it to zero. The model uses factorized posteriors and jointly infers $\omega$ and $\mathbf{z}$. If a dimension contributes little to the likelihood, the optimal posterior pushes mass toward $\omega=0$ (corresponding to the prior's spike), achieving exact $0$ sparsity. The KL term in the ELBO further penalizes unnecessary non-zeros, compressing redundant dimensions. We hope this clarification helps address your concerns about the spike-and-slab prior.

References:

[1] Reichelt et al. Rethinking variational inference for probabilistic programs with stochastic support. NeurIPS’22.

[2] Ilharco et al. Editing models with task arithmetic. ICLR’23.

[3] Yadav et al. TIES-Merging: Resolving interference when merging models. NeurIPS’23.

[4] Iurada et al. Efficient Model Editing with Task-Localized Sparse Fine-tuning. ICLR’25.

[5] Zhang et al. Knowledge Composition using Task Vectors with Learned Anisotropic Scaling. NeurIPS’24.

Thank you again for helping us improve the paper and hope our response can resolve your concerns!

Thank you for the rebuttal and clarification. I will keep my current rating (4 as positive) and do not have any further questions.

Thank you for your thoughtful feedback and for noting the clarity of our research motivation and approach. In what follows, we respond to the raised points with clarifications.

W1: The paper suffers from loose narrative and confusing notation, particularly the use of z. It also relies too much on external papers like aTLAS to be understood.

We appreciate the reviewer’s feedback and agree that clarity can be improved. In the revised version:

• We will consistently reserve $\mathbf{z}$ for task vector coefficients and avoid reusing it for CLIP embeddings (which are now denoted as $h$).
• We will and a notation table and clarify key terms early in Section 3 to make the paper more self-contained.
• A brief summary of task composition and aTLAS will be included to reduce reliance on external papers. These changes aim to resolve the ambiguity and improve readability.

W2: The Bayesian angle feels like a standard extension and offers only marginal improvement, making the contribution harder to assess.

We appreciate the reviewer’s positive overall evaluation. While we adopt a Bayesian perspective, our motivation is grounded in the specific challenge of composing task vectors under uncertainty and redundancy, which is not addressed in prior works like aTLAS. Unlike previous methods that rely on deterministic, task-level composition, we introduce a structured variational framework with Spike-and-Slab priors, enabling sample-specific sparse inference and controllable coefficient selection.

This design is particularly suited for multi-task transfer with heterogeneous tasks, where a fixed combination often fails to generalize. Moreover, our deterministic gating mechanism avoids the instability of Monte Carlo sampling in high-dimensional weight space—an aspect rarely considered in prior literature. These ingredients not only bring consistent performance improvements, but also offer better sparsity-efficiency trade-offs and interpretability. We hope this clarifies that our contribution is more than a marginal extension.

W3a: Standard benchmarks.

We clarify the task vector acquisition process. We use CLIP-ViT-B/32 as the base model and perform lightweight fine-tuning separately on 8 public standard datasets (Cars, DTD, EuroSAT, GTSRB, MNIST, RESISC45, SUN397, SVHN), which are the most commonly used visual benchmarks in aTLAS and Task-Vector literature. The fine-tuning settings are consistent with aTLAS, so each dataset corresponds to one task vector.

W3b: One-leave-out set of experiments.

Regarding your question, we have conducted one-leave-out experiments on the ViT-B/32 model(Deterministic-based+task-level), with results shown in the table below:

The experimental results show that when excluding one dataset, the average accuracy of model merging on the remaining seven datasets decreased by 4.82% compared to merging eight models. This experimentally demonstrates that expanding the task vector pool can enhance model merging performance. For theoretical supplementation and proof, please refer to Ortiz-Jimenez et al.[5] and Li et al.[6].

Q1: A comparison against prior work like aTLAS is missing.

The "Deterministic-based+Task-level" results in Table 1 of our paper are actually the aTLAS results. We did not explicitly label them as "aTLAS" when writing to distinguish between task-level and sample-level effects, which may have caused confusion in understanding the comparison. We will clearly mark the aTLAS name in the final version and have added comparisons with more prior work including Task Arithmetic[1], AdaMerging[2], and WUDI-Merging[3]. The results are shown in the table below:

Q2: The difference between the deterministic and probabilistic approaches.

Thank you for pointing out the superiority of sample-specific approaches in our method. For the probabilistic approach, we can see from Table 1 that there is a 4.45% improvement at the task level and a 0.98% improvement at the sample level. However, we notice that on the most challenging dataset SUN397, the ViT-B/32 model shows a 4.96% improvement with the sample-level probabilistic method compared to the deterministic method.

We believe one reason for the smaller average improvement at the sample level is that the sample-specific method already fully utilizes each sample's features through learning a lightweight network. However, the probabilistic approach has the natural advantage of quantifying the reliability of output results, which provides the necessary foundation for our subsequent sparse priors and gated sampling approaches.

Experimental results demonstrate that probabilistic methods not only introduce reliability evaluation of results but also further improve model merging performance. Therefore, the introduction of probabilistic methods represents a valuable innovation for model merging.

Q3: The difference between using a simple Gaussian prior w.r.t. using a Spike-and-Slab prior.

We agree that compared to other experiments, the improvement of the Spike-and-Slab prior-based method over the Gaussian prior method is indeed marginal. However, on the most challenging dataset SUN397, the ViT-B/16 model achieves a 2.29% improvement under the Spike-and-Slab prior, demonstrating that this sparse prior has more potential on more challenging tasks.

To more clearly show the difference between Gaussian prior and Spike-and-Slab prior results, we present the average accuracy across three models and eight datasets, along with the gated ratio for the Spike-and-Slab prior, as shown in the table below:

From the table results, we can see that although the Spike-and-Slab prior shows relatively small accuracy improvements compared to the Gaussian prior-based method, it achieves this performance using only around 70% of activated block coefficients. In contrast, the Gaussian prior learns all block coefficients but causes slight accuracy degradation. This experimentally validates our findings mentioned in Figure 2 that task vector redundancy indeed affects model merging, and also inspires us to design the gated sampling method. Therefore, we believe the Spike-and-Slab prior is necessary.

Q4. The computational complexity and a tradeoff between gain and complexity.

Thank you for your suggestion to evaluate the tradeoff between gain and complexity. We will add this tradeoff analysis in our revision. For intuitive understanding, we compare the complexity between sample-level and task-level approaches as follows:

• Parameters: Sample-level approach adds 0.75M parameters (~0.85% of backbone); task-level approach only needs to store one set of fixed scalars $\Lambda$, with negligible parameter overhead.
• Inference FLOPs: Sample-level requires forward pass through $\Phi$ first (~0.2 GFLOPs), then one ViT forward pass (~4.5 GFLOPs), total overhead increase ~4.4%; task-level only has ViT forward pass.
• Time (A40 GPU, batch 64): ViT-B/32 standalone inference 5.6 ms/sample; sample-level composition 5.9 ms/sample, increase of 0.3 ms/sample.
• Benefits: Average accuracy on eight datasets improved by 5.74% for deterministic methods and 2.27% for probabilistic methods. The sample-level approach achieves higher accuracy gains with <5% inference cost, making it a better choice for scenarios requiring the highest cross-domain adaptability.

References:

[1] Ilharco et al. Editing models with task arithmetic. ICLR’23.

[2] Yang et al. AdaMerging: Adaptive Model Merging for Multi-Task Learning. ICLR’24.

[3] Cheng et al. Whoever Started the Interference Should End It: Guiding Data-Free Model Merging via Task Vectors. ICML’25.

[4] Ortiz-Jimenez et al. Task arithmetic in the tangent space: Improved editing of pre-trained models. NeurIPS'23.

[5] Li et al. When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers. ICLR’25.

Thank you again for helping us improve the paper and hope our response can resolve your concerns!

We appreciate your constructive review and for recognizing the clarity of presentation and empirical improvements, we are glad you enjoyed reading the paper. We now offer clarifications and address the concerns raised.

W1: The main difference from aTLAS seems to be introducing sparsity into the summation. Methodologically, it’s unclear how novel this change is in the ML literature.

While our method indeed builds on aTLAS in spirit, it introduces several non-trivial methodological departures tailored for sample-specific and uncertainty-aware task composition:

• We go beyond fixed task-level scaling and instead propose a variational Bayesian formulation that learns input-dependent coefficient distributions rather than fixed values.
• Our use of Spike-and-Slab priors enables structured sparsity, which not only improves interpretability but also reduces overfitting and improves generalization on heterogeneous tasks.
• We further design a deterministic gating mechanism to replace stochastic sampling, which addresses variance and scalability issues commonly encountered in variational inference over high-dimensional parameter spaces.
• Empirically, this framework consistently outperforms aTLAS while using fewer parameters (see Fig. 4 and App. C.3), indicating that the proposed changes bring measurable benefits beyond architectural variation. To our knowledge, these components—Bayesian sample-level task vector composition with sparse gating— have not been explored in prior work, including aTLAS. We believe this constitutes a meaningful and task-relevant extension in both the modeling and inference aspects.

W2: Notations.

Thank you for pointing out the confusion in our notation. We used $z$ in line 104 to represent the latent representation of input image $x$ from the neural network, and $\mathbf{z}$ in lines 155-156 to represent the composition coefficients in our variational framework. In Eq. 12, $\mathbf{z}(\mathbf{x}_s;\Phi)$ refers to the composition coefficients for sample $\mathbf{x}_s$, which are filtered by gating $\mathcal{G}$ (Eq. 7) and then combined with the task vector matrix $\mathcal{T}$ to produce weight increments of the same shape as $\Theta_0$, resulting in the final prediction. We will clarify the symbol definitions in our final revision to ensure a clear distinction of symbol meanings.

W3. Comparison of training separate models on separate data.

Regarding your concerns about additional computational costs, we provide a detailed discussion:

• Minimal parameter overhead: The gated sampling model has extremely small parameter requirements. The additional parameters only include $\Psi=\lbrace \psi_1, \psi_2\rbrace$ (two scalar thresholds) and a three-layer MLP with parameters $\Phi$, totaling <1% of the backbone model parameters. Compared to training separate complete models for each task, the overhead of our gating mechanism is negligible. Despite this minimal cost, gating significantly improves knowledge transfer performance — achieving 1.35% improvement over spike-and-slab prior on ViT-B/32, and 3.56% improvement when applied to the aTLAS method. Moreover, our gated sampling method achieves state-of-the-art results across three models and eight datasets, demonstrating its excellent performance.
• Dynamic knowledge composition: For your suggestion of training separate models on separate data, multi-task applications often require dynamically switching or combining knowledge across different domains. With separate models trained on individual datasets, it is nearly impossible to enable/suppress specific domain features at runtime. Additionally, when adding N new tasks in the future, you would need N large-scale training runs and storage for N complete model weights. In contrast, task arithmetic only requires fine-tuning once per new task and saving task vectors, with minimal storage growth during composition. Our gated sampling mechanism further improves inference efficiency and storage space, providing more flexible choices for task vector composition in multi-task scenarios. We hope this discussion helps you better understand the advantages of gated sampling.

Q1: Pre-trained model in Fig 2a.

The pretrained model used in Figure 2a is the CLIP-ViT-B/32 image encoder.

Q2: Explanation of MNIST.

You made a very careful observation. We believe the possible reason for this interesting phenomenon is as follows: In the t-SNE projection, MNIST clusters appear clearly separated from other datasets mainly because of the huge domain gap between its visual distribution (grayscale, low resolution, handwritten strokes) and natural images. Due to this difference, pretrained features alone perform poorly on MNIST. However, our sample-level composition automatically selects a few task vectors related to "contours/strokes" for handwritten digits during inference while suppressing redundant channels. Therefore, MNIST accuracy in Table 1 improves dramatically from 25.23% (using only $\Theta_0$) to 98.72%, showing the most significant gain.

Q3: 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.

In our training process, the pretrained weights are not jointly fine-tuned with the ELBO loss across all tasks. Our process consists of two stages:

1. Task vector creation: We independently fine-tune each dataset to obtain corresponding task vectors, then freeze the pretrained weights $\Theta_0$ and all $\tau$.
2. Composition stage: We only train the inference network $\Phi$(3-layer MLP) and two gating thresholds $\Psi$ using ELBO loss to predict sample-level composition coefficients $\mathbf{z}$ and perform gating. Trainable parameter scale: Using ViT-B/32 as an example, the backbone has approximately 87.8M parameters. The inference network uses a three-layer fully-connected structure, plus two scalar thresholds and a global scaling coefficient $\lambda$, totaling approximately 0.75M parameters, which is ~0.85% of the backbone. The additional FLOPs during inference are also <5%. Compared to training $K$ complete separate models for $K$ tasks, our storage and update overhead is significantly reduced.

Q4. Single task scenario.

Thank you for pointing out the imprecise wording in our paper. You are correct – we indeed only selected image classification experiments, and we agree that expanding our method to other task scenarios is important for validating its effectiveness. Therefore, we have added experiments on discriminative NLP tasks. The experimental results are as follows:

For discriminative language tasks, we use RoBERTa[4] as the backbone and evaluate on the 8-task GLUE benchmark[5]. We will provide detailed experimental settings and more comprehensive results in the revised version.

Q5. Clarify the difference between a deterministic and stochastic approach, and Task-level vs sample-specific.

• Deterministic vs Stochastic approaches: Deterministic methods directly provide a single point estimate with unique composition results. Stochastic methods treat coefficients as random variables, explicitly model uncertainty, and complete composition through distributions or sampling.
• Task-level vs Sample-specific approaches: Task-level methods use coefficients that only depend on the overall "task" label, shared by all samples. Sample-specific methods use coefficients that are further influenced by specific inputs, enabling fine-grained adaptation to differences within samples.

Q6. Are the w/o Gated results in Table 2 from the same model and inference approach as the Table 1 final row?

Table 2 compares method performance under sparse conditions. The w/o Gated results in Table 2 come from the model with spike-and-slab prior, while the final row in Table 1 comes from a sample-level probabilistic model under Gaussian sampling. The difference between these two methods is that the former uses a spike-and-slab prior to reduce redundancy. Detailed results can be found in Appendix C.2.

References:

[1] Ilharco et al. Editing models with task arithmetic. ICLR’23.

[2] Yadav et al. TIES-Merging: Resolving Interference When Merging Models. NeurIPS’23.

[3] Yu et al. Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch. ICML’24.

[4] Liu et al. Roberta: A robustly optimized bert pretraining approach. CoRR’19.

[5] Wang et al. Glue: A multi-task benchmark and analysis platform for natural language understanding. EMNLP Workshop BlackboxNLP’18.

Thank you again for helping us improve the paper and hope our response can resolve your concerns!

Thank you for your thoughtful feedback and for acknowledging the sound theoretical formulation involving our variational framework and Spike‑and‑Slab prior. We address your questions as follows.

Q1: Baseline Performance Discrepancy.

Thank you for mentioning the performance discrepancy, and we elaborate on the reason for you now. First of all, our experimental setup is fair as all models are compared under the same experimental conditions, including the same datasets, the same preprocessing procedures, and the same evaluation metrics.

The performance discrepancy is mainly due to the data augmentation strategies we used in the original manuscript. As we stated in our paper (lines 264-266):

"To improve computational efficiency, we precompute feature representations for all datasets and apply data augmentation to generate multiple feature versions, thereby constructing a more diverse training set."

This augmentation strategy is static and relatively weak compared to the dynamic data augmentation strategy used by the original aTLAS method (applying random transformations to each sample in real-time during training) and therefore caused accuracy decline. We adopted the static one purely for the sake of computational efficiency, and it presumably affects the performance of all models equally. To demonstrate this, in this response, we conduct experiments using the same dynamic data augmentation method as the original aTLAS paper. The results are shown in the following table:

Results show that after restoring the original augmentation strategy, the ViT-B/32, ViT-B/16, and ViT-L/14 models improved by 2.75%, 4.98%, and 2.35% respectively compared to aTLAS. This proves that, regardless of experimental settings, our method consistently performs better. We will supplement the complete experimental results under identical data augmentation conditions in the revised version of our paper. We hope this addresses your concerns about the performance discrepancy.

Q2: Comprehensive Comparison and Trade-off Analysis.

We follow your suggestion and have conducted additional experiments with related methods for a comprehensive comparison. The results are reported as follows:

(a) Final accuracy

On 8 visual benchmarks, VTVC achieves state-of-the-art performance with an average accuracy of 87.45%.

(b) Merging cost (time/compute)

The training and inference costs of the four methods (based on 1000 samples, running on a single NVIDIA A40 with a ViT-B/32 model and batch size 64) are shown in the table below:

Due to VTVC's structural design, the training stage requires learning a small number of additional parameters ${\Phi, \Psi}$. However, during inference, there is no need to traverse composition coefficients, significantly improving inference efficiency. Overall, VTVC trades a slightly higher one-time offline cost for significantly higher accuracy and sample adaptation capability, while maintaining millisecond-level online latency.

(c) Scalability

The scalability discussion for the four methods is shown in the table below:

• Larger Models
  • Task Arithmetic: Inference cost increases linearly with number of layers, grid search becomes even slower.
  • AdaMerging: Number of parameter blocks increases, linear solver matrix becomes significantly larger.
  • WUDI-merging: Computation scales linearly with parameters.
  • VTVC(Ours): No additional parameters added, inference latency still <1ms.
• More Tasks
  • Task Arithmetic: Requires complete grid search for $\lambda$ for new tasks.
  • AdaMerging: Must re-solve all block-level coefficients.
  • WUDI-merging: Direct SVD merging on $\tau_{new}$.
  • VTVC(Ours): Only requires few iterations of fine-tuning $\Phi$ to quickly absorb $\tau_{new}$.

The results in Figure 5 and Appendix C.1 show that our method VTVC achieves better gating effects and model merging performance when transferred to larger backbone networks with only minimal additional cost. For new tasks, there is no need to traverse composition coefficients again, maintaining fast inference time.

(d) Data dependency

• WUDI-merging: Completely data-free, most lightweight, but cannot optimize for specific samples.
• AdaMerging: Requires small amounts of validation data to estimate block-level coefficients, low data threshold, but lacks sample-level adaptation.
• VTVC(Ours): Requires rich validation data to learn coefficient mapping, suitable for data-rich or continuous incremental learning scenarios. Therefore, if pursuing highest accuracy and cross-domain generalization, VTVC's data investment is reasonable and necessary. In extremely data-scarce scenarios, WUDI-Merging or AdaMerging can be chosen as cost-accuracy trade-offs. We will provide complete experimental configurations, cost calculations, and validation results in the revised version, and highlight the above trade-off analysis in the main text discussion.

Q3: Broadening the Experimental Scope.

We strongly agree with your viewpoint and have promptly conducted supplementary experiments on task negation and the NLP domain. The specific results are as follows:

• Task negation

  In the task negation setting, the model aims to forget the target task as much as possible while maintaining control tasks within a certain accuracy range (>95%).

  For task negation, we introduce an asymmetric gated sampling mechanism. This mechanism sets lower thresholds for target task coefficients while protecting important components from being filtered out. The experimental results for task negation using ViT-B/32 are as follows:

  Methods	Target($\downarrow$)	Control($\uparrow$)
  Task Arithmetic(ICLR’23)	22.62	60.73
  TIES-Merging(NeurIPS’23)	20.93	61.52
  TaLoS(ICLR’25)	17.28	61.33
  aTLAS(NeurIPS’24)	19.42	61.31
  VTVC (Ours)	16.75	60.84

  Our experiments show that the asymmetric gating mechanism achieves better forgetting effects on target tasks while keeping control tasks within the required performance range. This demonstrates that our method works well not only for task addition but also for task negation.

• NLP domain

  For discriminative language tasks, we use RoBERTa as the backbone and evaluate on the 8-task GLUE benchmark. We will provide detailed experimental settings and more comprehensive results in the revised version.

  Methods	COLA	STS-2	MRPC	STS-B	QQP	QNLI	MNLI	RTE	Avg.
  Zero-shot	0.00	53.76	85.01	4.01	37.48	53.05	37.09	71.19	42.70
  Weight Averaging	0.00	59.21	85.79	46.99	45.37	63.94	48.00	71.19	52.56
  Task Arithmetic(ICLR'23)	8.35	88.26	89.57	32.84	82.03	85.40	75.54	80.43	67.80
  TIES-Merging(NeurIPS'23)	31.76	88.86	86.18	10.94	61.05	85.94	83.01	69.56	64.66
  DARE(ICML'24)	0.00	88.14	86.61	30.19	84.33	79.09	63.95	77.16	63.68
  VTVC (Ours)	80.60	90.60	88.48	74.37	80.40	83.20	74.00	83.80	81.93

  The experimental results further demonstrate that VTVC is also applicable to discriminative NLP tasks, showing cross-modal generality and robustness.

Q4: Justification for Sample-Specific Complexity.

We acknowledge that sample-level composition introduces additional computational overhead, primarily during inference-time coefficient prediction. However, our model amortizes this cost via a lightweight inference network, which outputs coefficients in a single forward pass and does not require iterative optimization. The results in [Q2] show that compared to task-level methods commonly used in existing task arithmetic work, VTVC can complete coefficient learning with minimal time during the inference stage after training, without trailing composition coefficients. Meanwhile, as illustrated in Figure 2 of our paper, sample-level methods possess the ability to capture and identify differences between samples that task-level methods cannot achieve. Experimental results demonstrate that a more fine-grained use of sample features can significantly improve model merging performance. Based on this finding, we believe the following scenarios are more suitable for deploying sample-level methods:

1. Online inference/edge deployment scenarios with mixed input sources, where sample-level methods can capture distinctive features of new data and provide more accurate predictions.
2. High-risk or continual learning scenarios (such as medical imaging, autonomous driving) that require evaluating decision confidence, where our method can smoothly absorb new domain knowledge while maintaining performance on old tasks. Based on the above discussion, we believe the additional cost of sample-level methods is controllable in many real-world complex scenarios and brings significant accuracy and generalization benefits.

Thank you again for helping us improve the paper, and hope our response can resolve your concerns!

We sincerely appreciate your constructive comments and recognition of our method’s novelty, significance, and potential scalability. Below, we provide clarifications regarding the mathematical formulations and soundness of the methodology.

Q1: 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?

A1: The merging formulation for our sample-level composition extends the task-level definition by making the composition coefficients $\mathbf{z}$ input-dependent, i.e., $\mathbf{z}(x)$. Specifically:

• Non-block case:

\theta(x) = \theta_0 + \sum_{i=1}^N \sum_{j=1}^M \mathbf{z}_{ij}(x) \cdot \tau_{ij}, $$where each task vector $\tau_i$ is divided into $M$ parameter blocks, and each block $\tau_{ij}$ has its own coefficient $\mathbf{z}_{ij}(x)$ conditioned on the input $x$. We will clarify these formulations in the revised paper (Section 3, after Eq. (1)) to explicitly distinguish the sample-level case.

Q2: Some notations related to task arithmetics are unclear, making it hard to assess correctness.

A2: Thank you for pointing out the potential confusion in notation. We clarify the notations and design choices below:

1. Meaning of $\mathbf{x}$ and $\mathbf{y}$: In our formulation, $\mathbf{x}$ and $\mathbf{y}$ refer to arbitrary input-label pairs sampled during the composition training stage. They are not tied to any specific target task, but are drawn from the joint validation set used for learning composition parameters.
2. Definition of $\mathcal{D}_t$ in Eq. (12): $\mathcal{D}_t$ denotes the joint training data used to optimize the variational components $\Phi$ and $\Psi$ during composition. Specifically, we construct $\mathcal{D}_t = \bigcup_{i=1}^8 \mathcal{D}_i^{\text{val}}$, i.e., the union of the validation splits from all 8 datasets. We uniformly sample mini-batches from this set to train the gating components under the variational objective.
3. On “optimizing coefficients for a single task”: Our method does not tune coefficients per task individually. The model learns a single pair of shared networks $\Phi$ and $\Psi$ to produce sample-level coefficients based on $\mathbf{x}$, using multi-task supervision from $\mathcal{D}_t$. During evaluation, the same shared components and task vector pool $\mathcal{T}$ are used for all tasks without any task-specific tuning. Thus, Eq. (12) does not confer any advantage to a particular task. In contrast to classical task addition, where one fixed coefficient vector is shared across all samples, our formulation dynamically generates composition coefficients per input. However, the model remains shared and fair across tasks. We will clarify these points in Section 3.3 and the caption of Eq. (12) for clarity.
4. Whether evaluation still uses "one merged model": Yes. During inference, the model parameter set consists of $\Theta_0$, task vector pool $\mathcal{T}$, and trained $\Phi, \Psi$—all three are frozen. For any test sample (regardless of which task it belongs to), the model first generates $\mathbf{z}$ through $\Phi$, obtains increment $\Delta\Theta$ through gating, and completes prediction with one forward pass. The entire process no longer updates weights, all tasks share the same network, only $\mathbf{z}$ varies with samples. We hope this clarification of the training process helps you better understand our work.

Q3: Can you explain the motivation of Eq. (11) more? What does it mean to represent $\mathbf{z}$ as a deterministic function of input $\mathbf{x}$? Since $\mu$, $\sigma$, and $f$ are all parameterized by $\Phi$, how are they related? Also, isn't $\mathbf{z} = f_\Phi(\mathbf{x})$ very unlikely to happen if $\mathbf{z}$ is continuous?

A3: We provide a comprehensive clarification for each sub-question:

1. Motivation of Eq. (11).

• Eq. (11) transforms the traditional "obtain $\mathbf{z}$ via sampling $\rightarrow$ then gate" approach into "first give $\mathbf{z}$ a deterministic point estimate $f\_\Phi(x)$ → then use gating $\Omega$ to filter by uncertainty". The core motivation is: in high-dimensional sparse scenarios, directly performing Bernoulli-Gaussian mixture sampling on each dimension has high variance, heavy computation, and difficult convergence. By folding the randomness of the sampling process into a continuous, differentiable gating function, we can achieve differentiability, stability, and sparse selection simultaneously.

2. Why can z be represented as a deterministic function $f\_\Phi(\mathbf{x})$?

• From the variational inference perspective, we freely choose the approximate posterior family. We set: $q\_{\Phi,\Psi}(\Omega,\mathbf{z}|\mathbf{x}) = q\_\Psi(\Omega|\mathbf{z})\delta\_{\mathbf{z}=f\_\Phi(\mathbf{x})}$ Using Dirac $\delta$ to represent "posterior concentrated at point $f\_\Phi(\mathbf{x})$". This is not a random event of "continuous variable hitting exactly one point", but defines a new point-mass posterior family. Randomness and uncertainty are handled by $\Omega$.
• $f\_\Phi(\mathbf{x})$ typically takes the mean $\mu_\Phi(\mathbf{x})$ of the original Gaussian posterior. This retains $\mu,\sigma$ for subsequent uncertainty measurement while avoiding high-variance sampling of $\mathbf{z}$.

3. How are $\mu, \sigma, f$ related when all parameterized by $\Phi$?

• The neural network parameterized by $\Phi$ outputs three sets of tensors:

  1. $\mu_\Phi(\mathbf{x})$ serves as the value of $f_\Phi(\mathbf{x})$, i.e., composition coefficient point estimate
  2. $\sigma_\Phi(\mathbf{x})$ used for computing gradient sensitivity + distribution deviation, participating in constructing uncertainty $U$
  3. $\pi\_\Phi(\mathbf{x})$ for Bernoulli probability if using stochastic gating

  Since all three share forward features, computational cost is low and gradient updates are consistent, avoiding multiple network calls.

4. Isn't "$\mathbf{z}=f_\Phi(\mathbf{x})$" with zero probability unreasonable?

• The Dirac $\delta$ here is a modeling choice for the posterior family, not an actual sampling event; thus there's no paradox of "zero probability yet happening".
• The KL term in ELBO can still be computed:

This can be viewed as regularization on prior probability density, still providing constraints.

Q4: Can you clarify the definition of gated ratio?

A4: We clarify as follows:

• Definition of gated ratio: In the Spike-and-Slab prior setting (Sec. 4.2), each coefficient $\mathbf{z}\_{ij}$ is associated with a binary gating variable $\omega_{ij} \in \lbrace0, 1\rbrace$. If $\omega_{ij}=1$, the coefficient is active and participates in model merging; otherwise, it is masked to zero. The gated ratio is defined as the average proportion of active coefficients: $\text{Gated Ratio} = \mathbb{E}\_{i,j}[\omega\_{ij} = 1].$ For example, a gated ratio of 0.7 means that only 70% of coefficient blocks are used during merging.

• How to interpret Figure 4:

  • In Figure 4, we plot the accuracy (blue and orange lines) and the gated ratio (green line and shaded area) across eight datasets for three backbone models. The Spike-and-Slab prior (orange line) consistently achieves higher accuracy than the Gaussian prior (blue line), while the green curve shows that it uses fewer coefficients (i.e., a lower gated ratio). The precise numbers are reported in Appendix C.2.

  We will revise the main text and caption of Figure 4 to make this explanation clearer.

Q5: The difference between deterministic-based sample specific and probabilistic-based sample specific row and writing suggestions.

A5: To address your confusion, we clarify the meaning of each row as follows:

• Deterministic vs. Probabilistic
  • Deterministic-based: Directly uses point estimation $\mathbf{z}=f_\Phi(\mathbf{x})$ without random sampling; for sample-level, it then applies gating function $\Omega$ for sparse pruning.
  • Probabilistic-based: Establishes complete variational posterior $q(\mathbf{z}|\mathbf{x})$ for $\mathbf{z}$, uses sampling/expectation during training, explicitly retains noise terms and KL regularization.
• Task-level vs. Sample-specific
  • Task-level: Same task shares one set of coefficients; "Prob-TL" is our control implementation that restricts variational inference to task level.
  • Sample-specific: Coefficients vary with input; both "Det-SS" and "Prob-SS" are sample-level methods proposed in this paper, where Prob-SS is our main method. We will clearly identify our method in the revised version and explain why we use aTLAS as the baseline.

Thank you again for helping us improve the paper and hope our response can resolve your concerns!

Thanks for author's reply. Given the clarified notation, I can go back to read the paper and other reviews carefully again to re-evaluate the paper. First, I hope the authors to include these clarification for notations again, given that your setup is very different from standard task vector merging for both merging step (i.e., training) and inference step (i.e., evaluate the merged model on different tasks) so it's not very starightforward to get these points in the current version. After the rebuttal discussion, methodology-wise confusion due to these notation problems are mostly solved and I appreciate author's effort again.

There are a few remaining questions and suggestions:

1. Sparsity-based merging is not a very novel idea for task-based merging. Although I'm not questioning the novelty of sample-based model merging, I recommend authors to describe sparsity in model merging as a separate section in related work to lead to your motivation in the context of sample-based merging. Just to name a few:

Empirical work to show the effectiveness of sparse merging:

[1] Localizing task information for improved model merging and compression.

[2] Localize-and-Stitch: Efficient Model Merging via Sparse Task Arithmetic.

Theoretical work to support the rationale of sparse merging:

[3] Efficient Model Editing with Task Vector Bases: A Theoretical Framework and Scalable Approach

Although somewhat surprisingly, the gain from using sparsity in your set-up (Gaussian vs. Spike and Slab prior) is not that big compared to the gain in task-based method. Also, the sparsity value based on Figure 4 are only moderate but in task-based method the sparsity can be pushed further down for more efficiency gain. Do you have any insights on this?

2. In Figure 4, if my understanding is correct, gated ratio is only defined for Spike and Slab Priors? If yes, can you replace the title as comparison of Accuracy under 2 priors with Spike-and-Slab Priors gated ratio. I was confused because there are two lines (blue and orange) comparing two priors and you only have one shaded area for sparsity.

3. How does sparsity improve computational cost during inference? Since unfortunately, your training cost is substantially (x10 times based on your reply to TZ6C) higher than other already on the expensive-end task-based merging methods (like Adamerge), it will be important to clarify your contribution on inference cost minimization. Also, I would imagine the method is hard to scale up with the number of tasks because of the huge increase in training time. The current benchmark is on the standard 8-task vision benchmark, but can you run the algorithm on 14 and 20 vision task setup in [4]?

[4] Task Singular Vectors: Reducing Task Interference in Model Merging

The point of model merging is a light-weight module to avoid painful multi-task training. The first task arithmetic paper avoid this issue on limited grid search of merging coefficients on validation datasets to significantly reduce the computational cost. Even the multi-task version (mix all task data, maybe just validation data, together to create a merged model like what you did here in Eq. 12), has a low inference cost but should have the optimal "merging" performance right? Then what is the major contribution of your method?

4. For me this method, compared to the task-based task addition merging, might be also prone to with-in task distribution overfitting. Note that a practical concern might be in many public benchmark, especially in NLP domain, doesn't have a validation set so they evaluate the model on held-out tasks [5,6]. Although we see merging performance gain on 8-task vision dataset, is sample-based merging method, due to the increased training parameters, robust to distribution shift?

[5] What matters for model merging at scale?

[6] MergeBench: A Benchmark for Merging Domain-Specialized LLMs

Q3: Computational Cost and Scalability Analysis

A3: To better ananlyze computational complexity, we first clarify core notation: let pretrained weights be $\theta_{pre} \in \mathbb{R}^d$, and after fine-tuning on $N$ tasks, obtain task vectors $\tau\_t = \theta\_t^{ft} - \theta\_{pre}$. Traditional task arithmetic combines task vectors through coefficients: single coefficient case $\theta_{new} = \theta\_0 + \lambda\sum\_t \tau\_t (K=1)$, multi-coefficient case $\theta\_{new} = \theta\_0 + \sum_t \lambda\_t \tau\_t (K=N)$, or block-level coefficients ($K=N\cdot M$). Key parameters include: grid search candidates per dimension n, dataset sizes $|D\_{train}|$, $|D\_{val}|$, $|D\_{test}|$, and computational costs $C\_{forward}$, $C\_{backward}$. VTVC's amortized inference network cost is denoted as $C\_{\mathbf{z},\Omega}(K)$, growing linearly with coefficient dimension K.

When adding $n'$ new tasks, traditional methods and VTVC exhibit fundamentally different scaling characteristics. Traditional grid search in single-coefficient scenarios requires no training and has high inference efficiency, but severely limits model merging effectiveness; in multi-coefficient scenarios, hyperparameter complexity is $O(n^{K+\Delta K} \cdot |D\_{val}| \cdot C\_{forward})$, where $\Delta K = n'$ (task-level) or $\Delta K = n'\cdot M$ (block-level), showing exponential explosion. In contrast, VTVC's training complexity increment is only $O(\Delta K · |D\_{train}|)$, and inference complexity increment is $O(\Delta K · |D\_{test}|)$, both showing linear growth. Specifically for scaling from 8 tasks to 14 or 20 tasks: traditional methods' hyperparameter cost increases by $n^6$ or $n^{12}$ times, while VTVC's training cost only increases by 1.75 or 2.5 times. This linear scalability enables VTVC to easily handle large-scale multi-task scenarios with good scalability.

Traditional task arithmetic faces a fundamental dilemma: single coefficients are efficient but too coarse, unable to provide optimal weights for different tasks/samples; multi-coefficients enable fine-grained control but suffer exponential computational complexity. VTVC fundamentally changes this paradigm through amortized inference: using a lightweight neural network $\lambda(\mathbf{x}) = f\_\phi(\mathbf{x})$ to directly learn sample-level, block-level coefficients, transforming hyperparameter tuning from discrete grid search $O(n^K)$ to continuous optimization $O(K)$. More importantly, VTVC not only solves efficiency problems but also brings capabilities that traditional methods cannot provide:

1. Uncertainty quantification and confidence measurement through probabilistic modeling;
2. Interpretability through sparse priors, clearly showing which tasks/modules contribute most to specific samples;
3. Avoiding feature redundancy through gating mechanisms, preventing performance degradation from excessive task vector stacking. Therefore, VTVC's core value lies in maintaining modularity and scalability while approximating multi-task joint training performance through fine-grained adaptation, providing an efficient and accurate solution for dynamic task composition in practical applications.

Q4a: Distribution Shift Robustness

A4a: To show our method's generalization performance under distribution shift, we designed rigorous leave-one-out experiments: training the model on only 7 datasets each time, then evaluating on the completely unseen 8th dataset through linear probing. This setup simulates the model's generalization ability when facing entirely new domains, and the results are shown below:

The results show VTVC achieves an average accuracy of 78.38% across all 8 held-out tasks, demonstrating clear advantages over aTLAS's 76.87% and Zero-shot's 48.22%. Notably, VTVC achieves an improvement (+3.46%) on the challenging cross-domain task SUN397, again showing that sample-specific variational merging strategies can indeed capture shared knowledge structures across tasks rather than simply memorizing specific distribution features. Although VTVC introduces additional variational inference parameters, its uncertainty-based sparse gating mechanism actually serves as a regularization term, effectively preventing overfitting to source task distributions by dynamically selecting relevant task vectors.

Q4b: NLP Domain Applicability

A4b: In NLP discriminative tasks, VTVC trains a variational inference network to dynamically predict composition coefficients and gating probabilities for task vectors for each input text sample, achieving sample-specific task knowledge fusion. This network automatically determines which tasks' knowledge the current sample needs based on textual semantic features, and selectively activates relevant task vectors through sparse gating mechanisms. Compared to traditional methods that require searching for fixed composition hyperparameters on validation sets, VTVC adaptively customizes composition strategies for each sample through end-to-end learning, effectively solving the hyperparameter optimization difficulties caused by validation data scarcity in NLP tasks. Experimental results demonstrate our method performs excellently in the NLP domain as well:

These results demonstrate that our sample-level merging method can effectively handle NLP domain tasks, maintaining good performance even in the absence of dedicated validation sets.

Thank you very much for your detailed and constructive suggestions! We are happy to further clarify the notation confusion. In the revised version, we will add a detailed notation table that clearly explains the meaning and scope of each symbol to help readers better understand our method.

Q1: Supplementing sparsity-related work and result interpretation

A1: We fully accept your suggestion and will add a separate section on sparsity methods in model merging in the related work section. This will help readers better understand our motivation for using sparsity in the context of sample-level merging. Following your suggestion, we have supplemented sparsity-related work in model merging:

Sparse Model Merging Methods. Prior work improves efficiency via structural sparsity. Wang et al. [1] identify task-specific, largely non-overlapping weight supports and use binary masks with consensus merging; He et al. [2] localize tiny skill-bearing regions and stitch them back, enabling local sparse arithmetic. From a subspace view, Zeng et al. [3] compress and de-redundantize a task-vector basis; Gargiulo et al. [4] retain leading SVD components to merge in a low-rank sparse subspace; Marczak et al. [5] equalize spectra and merge within task-specific subspaces to curb dominance. Iurada et al. [6] perform task-localized sparse fine-tuning to produce additive/subtractive, decoupled sparse task vectors. These works made great progress while leaving a shared gap: existing sparsity strategies are mostly static designs or task-level constraints, lacking adaptive sparse composition mechanisms for "sample-level" scenarios, and do not involve uncertainty modeling of merging coefficients. We treat composition coefficients as latent variables, achieving coefficient-level sparsity under sample conditions through variational inference combined with sparse priors, and introduce uncertainty-aware gating mechanisms for deterministic selection.

[1] Localizing Task Information for Improved Model Merging and Compression.

[2] Localize-and-Stitch: Efficient Model Merging via Sparse Task Arithmetic.

[3] Efficient Model Editing with Task Vector Bases: A Theoretical Framework and Scalable Approach.

[4] Task Singular Vectors: Reducing Task Interference in Model Merging.

[5] No Task Left Behind: Isotropic Model Merging with Common and Task-Specific Subspaces.

[6] Efficient Model Editing with Task-Localized Sparse Fine-tuning.

Regarding the explanation of Table 1 and Figure 4, we apologize for any confusion. We would like to further explain them in details as follows::

• Table 1 shows the comparison between task-level and sample-level merging methods, as well as the differences between deterministic and probabilistic modeling methods. This table does not include sparse prior comparisons. The purpose of this table is to demonstrate the effectiveness of our modeling at the sample-level merging and in the probabilistic formulation.
• Figure 4 shows sample-level merging scenarios, comparing the accuracy performance of Gaussian priors and Spike-and-Slab priors, while showing the gated ratio for Spike-and-Slab priors. This figure justifies the effectiveness of sparse modeling based on spike-and-slab priors.
• Table 4 comprehensively compares the performance of Gaussian and Spike-and-Slab priors across three different scales of ViT models (ViT-B/32, ViT-B/16, ViT-L/14). Figure 4 represents the results for ViT-B/32 from Table 4. From Table 4, we observe that although the overall average accuracy improvement of Spike-and-Slab priors over Gaussian priors is relatively moderate, we notice particularly significant improvements on the most challenging dataset SUN397 (e.g., 2.29% improvement under ViT-B/16 and 1.17% improvement under ViT-L/14). This indicates that Spike-and-Slab sparse priors can more effectively discover and utilize the sparse structure of task vectors in more complex and challenging task scenarios, thereby significantly enhancing the model's generalization ability. We will re-optimize the descriptions of these tables and figures in the revised version to make them clearer and easier to understand.

Q2: Figure 4 title correction

A2: Agreed. Your understanding is completely correct! The gated ratio is indeed only defined for the Spike-and-Slab prior. We follow your suggestion to modify the title of Figure 4 to: "Accuracy Comparison under Two Priors with Spike-and-Slab Prior's Gated Ratio." Again, apologies for the confusion caused, and thank you for your suggestion.

Thanks for your reply, I want to point out some more points:

• Clarification for Computational Cost Coefficient Learning
  • TIES: the paper suggested a universal merging coefficient to be 0.4. Grid search can be applied to boost the performance.
  • Localize and Stitch: The paper refers "Both Algorithms" to "TIES and Task Arithmetic", not the proposed Localize and Stitch method. Their merging coefficient is automatically and cheaply computed in the stitching step elegantly through learned masks (Algorithm 1). Let me give more early popular model merging methods where merging coefficients are learned not searched, therefore this merging coefficients training time is a very important efficiency metric.
  • Fisher Merging "Merging models with fisher-weighted averaging.": learns weights by Fisher information.
  • RegMean "Dataless knowledge fusion by merging weights of language models.": learns weights by linear regression.
  • AdaMerge: learns weight layer-wise.
  • Survey paper "Model Merging in LLMs, MLLMs, and Beyond: Methods, Theories, Applications and Opportunities" even categorized this merging coefficient learning merging method as a major direction of model merging methods in Sec 2.3.2 as it was published in 2024, with most methods more complicated than naive merging coefficients grid search. Also I'm not even talking about training loss hyperparameter search because VTVC technically needs tuning for regularization loss.
• Sample Level Adaptation:
  • I very recently realized this paper "Merging Multi-Task Models via Weight-Ensembling Mixture of Experts" (WEMOE) is also a sample-level model merging design because given frozen task vectors, the router is learned based on each input sample, and the claimed performance is even better than MTL. I wonder if this should be a more reasonable baseline and therefore, MTL should not be the oracle that can be easily ignored in most of the task-level model merging methods.
  • Can you confirm if in your framework only $\Phi, \Psi$ are learned but all other model parameters, especially $\Theta$ and $\mathcal{T}$ are frozen?
  • I understand that one of your novel contributions is uncertainty quantification due to the variational framework, I have no concerns about this.
• Question about Task Negation Setup:
  • Because you claimed your flexibility on the task negation setup, so I want to ask for clarification:
  • Task Arithmetic/TIES can arbitarily minimize the target performance by adjusting the negation coefficient at the cost of control metric loss if grid search of merging coefficient is allowed. Since I don't see an absolute advantage of VTVC over all presented baselines in your reply to ZAnD, I wonder how do you guarantee a fair comparison? How does the negation coefficient selected for all baselines?