Towards Unified Task Embeddings Across Multiple Models: Bridging the Gap for Prompt-Based Large Language Models and Beyond (2024)

Abstract: Task embedding, a meta-learning technique that captures task-specific information, has gained popularity, especially in areas such as multi-task learning, model editing, and interpretability. However, it faces challenges with the emergence of prompt-guided Large Language Models (LLMs) operating in a gradient-free manner. Existing task embedding methods rely on fine-tuned, task-specific language models, which hinders the adaptability of task embeddings across diverse models, especially prompt-based LLMs. To hardness the potential of task embeddings in the era of LLMs, we propose a framework for unified task embeddings (FUTE), harmonizing task embeddings from various models, including smaller language models and LLMs with varied prompts, within a single vector space. Such uniformity enables comparison and analysis of similarities amongst different models, broadening the scope and utility of existing task embedding methods in multi-model scenarios, while maintaining their performance comparable to architecture-specific methods.

## introduction task embeddings serve as a valuable meta-learning tool for various downstream tasks such as multitask learning @xcite @xcite , model editing @xcite , and interpretability @xcite . current task embedding methods focus on encapsulating the task-specific information into a vector representation, which is derived from parameters of task-specific, fine-tuned models @xcite @xcite . however, the existing paradigm's reliance on the parameters of a single model restricts the comparison and adaptability of task embeddings across various models, especially hindering their applicability to prompt-guided, gradient-free large language models (llms). to overcome these challenges, we propose a model-agnostic framework for unified task embedding (fute), which enables the application of task embedding across diverse models including language models of various architectures and llms with different prompts, as illustrated in figure 1 . fute generates task embeddings by treating these models as "black boxes" and does not require knowledge of the original models' training datasets. specifically, we decouple the concept of task embedding into two distinct forms: dataset task embedding (dte) and model task embedding (mte). dte is designed to capture the task information specific to a dataset, focusing on its unique characteristics. in contrast, mte captures the model's behavior according to the task performed, independent of its training dataset. by decoupling mte and dte, our framework facilitates a more granular analysis of task characteristics, providing insights into both the data and models employed. moreover, the dtes and mtes of various models are computed using a single surrogate model, which replaces the role of the original model in existing methods. this consistent approach to computing dtes and mtes ensures that the learned task embeddings reside in the same embedding space, hence facilitating the comparison and analysis of similarities among them. building upon the aforementioned improvements, fute conceptualizes the combination of a prompt and a llm as a single model unit, enabling the computation of task embeddings for llms guided to perform specific tasks through prompts. experimental results show that, while being more adaptable and covering a wider range of models, our proposed framework performs on par with existing model-specific embedding methods. fute primarily focuses on extending task embedding methods to multi-model scenarios rather existing methods generate task embeddings by leveraging model parameters, confining to models of identical architecture and initialization. this limitation, depicted on the left, results in embeddings that reside in separate vector spaces depending on the model, making cross-model similarity comparisons impossible. our unified framework introduces a method for independently computing task embeddings, enabling the incorporation of diverse models, ranging from different neural architectures to large language models (llms) with distinct prompts, into a unified vector space. ## related work task embedding, initially introduced as a task vector @xcite , as a meta-learning method for visual classification tasks. this concept utilizes the empirical fisher information derived from model gradients to encode task-specific information into a vector representation, demonstrating that the model parameter can provide detailed task information. expanding into nlp, @xcite applied fisher information to capture task-specific details in fine-tuned language models, formalizing it as task embedding. zhou et al. (2022) explored three parameter-efficient fine-tuning (peft) methods, prompt tuning (liu et al., 2022b), bias tuning (ben zaken et al., 2022), and low-rank tuning (hu et al., 2022) to generate task embeddings. they discovered that peft parameters, being the only components adjusted during fine-tuning, inherently encapsulate task-specific information. ilharco et al. ( 2023 ) introduced a method with arithmetic operations by subtracting the original parameters from its fine-tuned parameters to extract task vectors. studies by @xcite proposed direct computation of task similarity via wasserstein distance, bypassing the need for embedding similarity measures. additional research efforts, such as @xcite have indirectly engaged with concepts akin to task embedding for multi-task learning, albeit without explicitly naming them as such. recent investigations into the task embedding of llms by @xcite extracted task vectors from parameters and activations of llms. however, similar to previous approaches, this method is confined to a single model and is unable to generate task embeddings for datasets, thereby limiting its applicability. all the aforementioned methods predominantly focused on learning task embeddings for individual models, limiting their application in scenarios involving multiple models. ## experiments we conducted two sets of experiments focused on smaller language models and llms. ## conclusion and future work this paper introduces a unified framework fute that separate task embeddings into dte and mte, by employing an additional surrogate base model and unsupervised data, without altering existing task embedding learning methodologies. fute can effectively extend current methods to learn mtes, which are positioned in a unified vector space, for various language models and llms with distinct prompts. fute is, theoretically, capable of learning mtes for any model with accessible outputs, within a single vector space. we leave the exploration to other models in the future.