Variational Task Vector Composition

2025-09-20

AI Summary

Existing task vector composition methods in multi-task learning neglect sample-specific task interactions, incur substantial inference overhead, and lack interpretability. Method: We propose a sample-wise variational task vector composition framework that models task-component sparsity via a Spike-and-Slab prior and introduces an uncertainty- and importance-aware gating mechanism for reliable, interpretable, dynamic task selection in high-dimensional spaces. Inference is performed under a Bayesian variational framework, requiring no additional inference-time parameters or forward-pass overhead. Contribution/Results: Our method achieves significant improvements over state-of-the-art task composition approaches across multiple benchmark datasets. It reduces sampling variance, enhances generalization performance, and improves decision transparency—enabling principled, sample-adaptive task composition without compromising computational efficiency or interpretability.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods

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 additional inference costs. In this paper, we propose variational task vector composition, 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 preserves only 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 consistently outperforms existing approaches across all datasets by selectively leveraging the most reliable and informative components in task vectors. These findings highlight the practical value of our approach, establishing a new standard for efficient and effective task vector composition.

Problem

• Optimizing task vector composition coefficients as latent variables
• Addressing structural redundancy in task vectors via sparsity promotion
• Improving sampling efficiency in high-dimensional sparse spaces

Innovation

• Bayesian inference with latent composition coefficients
• Spike-and-Slab prior for sparsity and informativeness
• Gated sampling mechanism for stability and interpretability