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 gradientfree 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 multimodel scenarios, while maintaining their performance comparable to architecture-specific methods.1

heading: 1 Introduction
text: Task embeddings serve as a valuable meta-learning tool for various downstream tasks such as multitask learning (Vu et al., 2020; Li et al., 2022; Zhou et al., 2022; Wu et al., 2023), model editing (Ilharco et al., 2023; Ortiz-Jimenez et al., 2023), and interpretability (Gueta et al., 2023). 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 (Vu et al., 2020; Zhou et al., 2022; Ilharco et al., 2023). 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
1Code is available at https://github.com/xnyuwg/fute.
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
8324
than directly improving the embedding methods themselves. Although the performance of FUTE relies on the employed task embedding methods, this aspect represents a strength, allowing FUTE to leverage future advancements in task embedding techniques without significant modifications. Our main contributions are summarized as follows: • We introduce FUTE, a framework capable of learn-
ing unified task embeddings from diverse models, including language models of different architectures, and LLMs with various prompts, within a single vector space. • We decouple the concept of task embedding into data task embedding and model task embedding. • Experimental results show that FUTE, while being more versatile, retains a performance to be comparable to existing model-specific methods.

heading: 2 Related Work
text: Task embedding, initially introduced as a task vector (Achille et al., 2019), 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, Vu et al. (2020) 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 FineTuning (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 (AlvarezMelis and Fusi, 2020; Tan et al., 2021) proposed direct computation of task similarity via Wasserstein Distance, bypassing the need for embedding similarity measures. Additional research efforts, such as (Lv et al., 2023; Huang et al., 2023a,b) 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 (Hendel et al., 2023) 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.

heading: 3 Methodology

heading: 3.1 Problem Setup
text: Existing task embedding learning methods employ a singular embedding to represent both dataset and model task information. In this paper, we refine the concept into two distinct forms: Dataset Task Embedding (DTE) and Model Task Embedding (MTE). The learning of them is independent and parallel.
Formally, we define a model as M and a dataset as D = {(xi, yi)}ni=1, where xi denotes the input and yi denotes the label. Task embedding is denoted as e, with DTE specified as eD and MTE as eM . Typically, the task embedding learning method requires a base model, denoted as Θ. The task embedding learning method f(·, ·) is formulated as:
e = f ( {(xi, yi)}ni=1,Θ ) , (1)
where the function f(·, ·) employs the base model Θ and a set of data instances {(xi, yi)}ni=1 to compute the corresponding task embedding e. In this paper, f(·, ·) encapsulates the complete pipeline of a task embedding learning method, integrating all steps from initial data to the final embedding output. It is important to note that to ensure all task embeddings reside within the same vector space, it is crucial that the function f(·, ·), responsible for learning task embeddings, utilizes a consistent base model Θ. This requirement ensures the consistency of the embedding process across different tasks.

heading: 3.2 Background
text: In this paper, we concentrate on extending two task embedding learning methods: TaskEmb (Vu et al., 2020) and TuPaTE (Zhou et al., 2022). TaskEmb employs the Fisher Information Matrix (FIM), which captures the curvature of the loss surface, to learn task embeddings. TuPaTE, on the other hand, introduces Parameter-Efficient FineTuning (PEFT) methods to learn task embedding.
TaskEmb TaskEmb (Vu et al., 2020) employs the FIM with respect to the parameters of a language model to compute task embeddings. The FIM is derived from the expected covariance of the gradients concerning the parameter θ of a language model ΘM where θ ∈ ΘM as:
Fθ = E x,y∼Pθ(x,y)
∇θ logPθ(y|x)∇θ logPθ(y|x)T . (2)
Specifically, TaskEmb computes the empirical Fisher on a fine-tuned model ΘM using the labels:
Fθ = 1
n
n∑
i=1
[ ∇θ logPθ(yi|xi)∇θ logPθ(yi|xi)T ] . (3)
Subsequently, the diagonal entries are extracted and averaged across the entire dataset, and then concatenated to form the task embedding e. The complete pipeline of TaskEmb is denoted as e = fTaskEmb({(xi, yi)}ni=1,ΘM ). TuPaTE TuPaTE demonstrates the application of three PEFT methods: Prompt Tuning (Liu et al., 2022b), Bias Tuning (Ben Zaken et al., 2022), and Low-Rank Tuning (Hu et al., 2022). In this paper, our focus is specifically on TuPaTE with Prompt Tuning. However, FUTE can be easily adapted to other methods in a similar manner.
Specifically, with a language model ΘM , TuPaTE with Prompt Tuning incorporates additional attention prefix matrices Kp and Vp as parameters. For each layer, these augmented vectors are concatenated with the existing key K and value V to form K ′ = [Kp,K] and V ′ = [Vp, V ], where [·, ·] denotes concatenation. Subsequently, the attention mechanism employs the modified K ′ and V ′ for computation. Following PEFT using the prompt Kp and Vp in the language model ΘM on dataset {(xi, yi)}ni=1, these prompt vectors are averaged across all layers and concatenated to form the task embedding e. The entire process of TuPaTE is denoted as e = fTuPaTE({(xi, yi)}ni=1,ΘM ).

heading: 3.3 Model Task Embedding
text: Model Task Embedding (MTE) is designed to encapsulate the task-specific capabilities of a model within an embedding. Previous methods require that task embeddings be derived from the same base model ΘM . Furthermore, these methods are closely tied to datasets, constraining the embeddings to models trained on specific datasets, as illustrated in Figure 2A. The dependency on datasets restricts applicability in scenarios where datasets, such as those for LLMs, are unavailable or inaccessible due to privacy concerns. To overcome these limitations, we introduce a model-agnostic and dataset-agnostic approach for extracting MTE. This method offers advantages in scenarios involving varied models and in contexts where training data is inaccessible.
MTE Framework FUTE eliminates the dependency on a specific dataset D for understanding the task information of a model, as depicted in Figure 2C, by utilizing an unsupervised text data U = {x′i}mi=1, unrelated to specific tasks, to substitute the conventional task-specific dataset D = {(xi, yi)}ni=1. By applying the model ΘM
on U = {x′i}mi=1, we first generate predictions {ŷMi }mi=1 according to:
ŷMi = ΘM (x ′ i). (4)
This approach is based on the premise that different task-specific models, such as sentiment analysis or question answering, will yield divergent outcomes even for identical text inputs. Here, ΘM (x′i) captures task-related information inherent to ΘM .
Additionally, FUTE introduces a surrogate base model for task embedding learning, denoted as ΘT , as a substitute for ΘM . This substitution preserves the uniformity of ΘT throughout the task embedding learning process, eliminating the need for the same ΘM . The T5-base (Raffel et al., 2020) is employed as ΘT in our experiments. We pair each x′i with its corresponding task-related ŷ M i to feed into the task embedding learning method f(·, ·), deriving MTE eM as:
eM = f({(x′i, ŷMi )}mi=1,ΘT ). (5)
This mechanism allows for the adaptation of the MTE learning process to various models, with the model output {ŷMi }mi=1 reflecting the nature of the task of ΘM , be it a probability, class, or sentence.
MTE for Prompt-based LLMs Previous methods typically rely on task-specific model parameters and datasets, which are not applicable or accessible for LLMs, especially those closed-source LLMs. Conversely, FUTE can be easily extended to accommodate prompt-based LLMs. In essence, a prompt in conjunction with an LLM can be treated as a unified model from which its MTE can be derived. Prompts can vary from simple instructions
to more complex In-Context Learning examples (ICL) and Chain of Thought (CoT) prompts.
We use the same MTE framework proposed above, and the only difference is that the input to Eq. (4) includes a prompt in addition to the original input x′, as shown in Figure 2D. Even though the LLM here, ΘM , may be significantly larger in size compared to the MTE learning model, ΘT , an LLM guided by a prompt, which is limited to performing a specific task, can be managed by ΘT .
This method holds particular promise for LLMs, where different prompts can trigger varied taskspecific performances, thereby offering a nuanced understanding of the task capability of LLMs without being limited to specific fine-tuned parameters.
Sourcing Unsupervised Data For sourcing unsupervised text data U , we opt for the CrossFit dataset (Ye et al., 2021). CrossFit encompasses a dataset collection of 160 diverse NLP tasks, offering a rich variety of text styles and structures. This diversity stands in contrast to large, unsupervised text corpora such as Wikipedia, which predominantly feature formal and neutral tones, and consistent style. Such uniformity in text data can skew the performance of models towards, for example, neutral predictions in tasks such as sentiment analysis. This bias may require a much larger volume of data to develop a comprehensive understanding of the model’s capabilities. Given this context, the broad spectrum of text characteristics found in CrossFit makes it an ideal choice for U in our study, especially when considering the efficiency of MTE learning.
Discussion of MTE Viewed from an alternative perspective, this MTE learning approach resembles the principles of knowledge distillation, as described in (Hinton et al., 2015), by utilizing extensive unsupervised data. It effectively facilitates the transfer of knowledge from the original model, ΘM , to the surrogate base model, ΘT .
Theoretically, FUTE can learn MTEs for any models with accessible outputs {ŷMi }mi=1, which potentially encompass not only language models and LLMs but also traditional neural networks like CNN (Fukushima, 1980) and LSTM (Hochreiter and Schmidhuber, 1997), as well as classical models such as decision trees and Support Vector Machines (Hearst et al., 1998).

heading: 3.4 Dataset Task Embedding
text: Dataset Task Embedding (DTE) aims to encapsulate the task information of a dataset within an embedding. In this paper, DTE is conceptualized similarly to the task embedding methods presented in prior studies. The process of learning DTE closely aligns with that of MTE utilizing ΘT . Specifically, as illustrated in Figure 2B, given a dataset D = {(xi, yi)}ni=1, the computation of DTE eD is revised as follows:
eD = f ( {(xi, yi)}ni=1,ΘT ) . (6)
This reformulation facilitates a more flexible and adaptable approach for generating task embeddings, by leveraging ΘT without being bounded by the constraints of architectural and initialization uniformity imposed by ΘM .

heading: 4 Experiments
text: We conducted two sets of experiments focused on smaller language models and LLMs.

heading: 4.1 Language Model Experiments

heading: 4.1.1 Experimental Setup
text: To assess the effectiveness of FUTE with language models, we follow (Vu et al., 2020; Zhou et al., 2022) to evaluate the transferability across 11 Classification or Regression (CR) and 11 Question Answering (QA) datasets. Vu et al. (2020); Zhou et al. (2022) conducted a transfer learning experiment, where initially, a language model is fine-tuned on a source dataset, followed by further fine-tuning on a target dataset. Based on the similarity between task embeddings learned from BERT-base (Devlin et al., 2019), a source dataset is selected for a given
target dataset, promising to achieve optimal transfer learning gain. Evaluating the selection process highlights the effectiveness of task embedding.
Our approach to the experiment diverges from the original method by shifting the focus from selecting the most appropriate source dataset to identifying the most compatible models for a target dataset. This is achieved by matching the MTEs of candidate models with the DTE of the target dataset based on their embedding similarity. More experimental setup details are given below:
Datasets CR datasets comprise nine datasets from the GLUE benchmark (Wang et al., 2019), in addition to SNLI (Bowman et al., 2015) and SciTail (Khot et al., 2018). QA datasets contain eleven datasets from the MultiQA repository (Talmor and Berant, 2019). See Appendix A.1 for more details.
Unsupervised Dataset We randomly sampled 100 text entries from each task within CrossFit (Ye et al., 2021) to ensure broad text styles and structures. We removed any datasets that overlap with the CR or QA datasets mentioned above, discarding any associated labels. This process resulted in the compilation of 104,162 text entries for U . In our experimental setup, we fine-tune ΘT on the unsupervised data for one epoch.
Metric We follow (Vu et al., 2020; Zhou et al., 2022) and employ two metrics: (1) the average rank ρ of the candidate source model that yields the highest transfer gain; (2) the Normalized Discounted Cumulative Gain (NDCG) (Järvelin and Kekäläinen, 2002), which assesses the overall quality of the entire ranking.
We present two extensions for TaskEmb and TuPaTE with FUTE: (1) FUTE+TaskEmb: It combines FUTE with the TaskEmb method, utilizing Eq. (5) and (6), in conjunction with the TaskEmb method fTaskEmb(·, ·). (2) FUTE+TuPaTE: It applies FUTE on TuPaTE, employing Eq. (5) and (6) alongside TuPaTE method fTuPaTE(·, ·).

heading: 4.1.2 Baselines
text: In language model experiment, we compare FUTE with the following baselines: (1) DataSize (Vu et al., 2020) prioritises dataset based on the size of dataset. (2) CurveGrad (Bingel and Søgaard, 2017; Vu et al., 2020) analyzes the gradients of the loss curve associated with the language model, and is adapted by (Vu et al., 2020) for transferability prediction. (3) TextEmb (Vu et al., 2020)
computes an average of the text representations extracted from the last layer of the language model across the entire dataset. (4) TaskEmb (Vu et al., 2020), as detailed in Section 3.2, leverages Finisher information to derive task embedding. (5) TuPaTE (Zhou et al., 2022), outlined in Section 3.2, utilize PEFT with the language model, extracting the tuned prompt as task embedding.
All these baselines are unavailable in the multimodel scenario. Dataset-centric methods, like DataSize, inherently cannot differentiate between models trained on the same dataset. Model-specific approaches, CurveGrad, TextEmb, TaskEmb, and TuPaTE, all require a single model initialization, limiting their adaptability across diverse models. CurveGrad computes the gradients from the same model. TextEmb extracts the text representations from the same model. For example, similarly, text embedding derived from BERT-base, BERT-large, and GPT2 do not reside in the same vector space. However, FUTE can be potentially adapted to extend text embedding methods as well. TaskEmb and TuPaTE, as mentioned above, are also unavailable for multiple models.

heading: 4.1.3 Transferability Results
text: The Transferability results are presented in Table 1. Upon comparison, FUTE exhibits competitive performance relative to the original methods. Specifically, when integrated with TaskEmb, FUTE achieves better results in the QA all-class setting compared to the original TaskEmb method, albeit with slightly less optimised outcomes in other settings. Similarly, the combination of FUTE with TuPaTE demonstrates improved performance in the all-class setting, while yielding very similar outcomes in the in-class setting. Overall, the dis-
crepancies between FUTE and the original methods are generally confined to less than a 2% difference.

heading: 4.1.4 Cross-Domain Results
text: FUTE utilizes consistent unsupervised data U for MTEs across all models. This approach diverges from previous methodologies, where each task embedding is tied to a specific dataset. We hypothesise that task embeddings generated by previous methods may encapsulate domain-specific information due to their dataset-centric nature. To test this hypothesis, we conduct experiments using three datasets from the finance and medical domains: FinancialPhraseBank (Malo et al., 2014), Medical Question Pairs (MQP) (McCreery et al., 2020), and MedQA (Jin et al., 2021). We randomly split 20% of the training set to serve as the validation set. These datasets exhibit text characteristics distinct from the general text content presented in the datasets used in Section 4.1.3.
The comparative results, displayed in Table 2, reveal that TaskEmb maintains consistent performance across this cross-domain experiment, closely mirroring the results obtained from FUTE. TaskEmb leverages the parameters of the entire language model, potentially embedding within it a wider range of domain information. Conversely, TuPaTE extracts task embeddings from the PEFT parameters, which may be more narrowly tailored to specific domain characteristics. This difference suggests that TaskEmb’s embeddings could carry broader information due to its reliance on the comprehensive model parameters.
Remarkably, FUTE using TuPaTE significantly outperforms the original method. This improvement can be attributed to the utilization of consistent unsupervised data U across all models, thereby
reducing the domain influence from the data source. The findings from Table 2 indicate the promising potential of FUTE to enhance the adaptability and effectiveness in cross-domain applications.

heading: 4.1.5 Visualization
text: In this section, we present a visualization in Figure 3 by extending TuPaTE to three language models, BERT-base (Devlin et al., 2019), GPT2 (Radford et al., 2019), and T5-base (Raffel et al., 2020), based on with CR and QA datasets. A notable observation from the visualization is the emergence of two main clusters based on task information, highlighting the significant distinctions of MTEs between task categories. MTEs associated with the same language model, particularly for QA tasks, demonstrate a tendency to cluster together, which indicates that the inherent characteristics of a language model also influence the task embedding space. An interesting detail is the presence of three blue points within the CR task cluster, representing MTEs associated with the BoolQ dataset (Clark et al., 2019). This positioning matches BoolQ’s characteristics with CR tasks due to its requirement for boolean answer generation for the question, which closely resembles CR tasks. This positioning matches with the characteristics of BoolQ, which is formatted to predict a boolean answer for each question and closely resembles CR tasks. See Appendix A.2 for more details.

heading: 4.2 LLMs Experiments

heading: 4.2.1 Experimental Setup
text: In this section, we assess the effectiveness of FUTE for LLMs by focusing on the task of zero-shot prompt instruction selection. The objective is to identify the most suitable prompt instruction that enables an LLM to achieve optimal performance on a given dataset. We utilize FUTE to generate MTEs of LLMs associated with prompts using Eq. (5), and compute DTE for the target dataset using Eq. (6). We then compute the similarity between MTEs and DTEs for prompt selection.
Datasets We utilize two categories of tasks for evaluation: Sentiment Analysis (SA) and Natural Language Inference (NLI). For SA, we employ: Rotten Tomatoes (Pang and Lee, 2005), IMDB (Maas et al., 2011), and SST-2 (Socher et al., 2013). For NLI, we use: WANLI (Liu et al., 2022a), CB (De Marneffe et al., 2019), and ANLI (Round 3) (Nie et al., 2020). We randomly sample 1,000 data from the test set of each dataset for evaluation.
Unsupervised Dataset We employ CrossFit (Ye et al., 2021) as the source of unsupervised data U and randomly sample 10,000 diverse text entries from U in Section 4.1. Again, we ensure no overlap between U and datasets D, and exclusively use text content, discarding any associated labels.
LLMs We employ Llama2-13B (Touvron et al., 2023) in 8 bit, Llama2-70B (Touvron et al., 2023) in 4 bit, and Mixtral 8x7B (Jiang et al., 2024) in 8 bit. Mixtral 8x7B is a Sparse Mixture of Experts (SMoE) architecture, featuring 8 experts with each token processed by two experts.
Prompts We incorporate a diverse set of 13 zeroshot prompts for each category of tasks, which comprises 10 vanilla instruction prompts sourced from PromptSource (Bach et al., 2022) and 3 Chain of Thought (CoT) prompts. Each prompt is structured to include instructions only, without any ICL examples. See Appendix A.3 for more details.
Metric We utilize three metrics: (1) the performance of the selected prompt; (2) performance rate,
computed as the ratio of the achieved performance with the selected prompt to the maximum performance observed across all prompts. A performance rate of 100% indicates perfect selection; (3) NDCG (Järvelin and Kekäläinen, 2002) evaluates the overall quality of prompt ranking, where relevance is determined by performance. An NDCG of 100% represents an ideal ranking. The final results are averaged across datasets.

heading: 4.2.2 Baselines
text: We compare FUTE against the following baselines: (1) MI (Sorensen et al., 2022) selects the prompt that maximizes the mutual information between the input and the LLM output. (2) LocalE (Lu et al., 2022) calculates the local entropy of the prediction probabilities for each entry. (3) GlobalE (Lu et al., 2022) computes the global entropy across all prediction categories in a dataset. ZPS (Liao et al., 2022) leverages an ensemble of prompts to generate pseudo labels and then predicts the performance of each prompt on these pseudo labels. Based on the ensembling technique, it has three variants: (4) ZPS-Log using the sum of log probabilities, (5) ZPS-Prob using the sum of probabilities, and (6) ZPS-Vote using majority vote. (7) Self-Select (Ramji and Kyimpopkin, 2023) introduces a meta-prompt that queries LLM for selecting the most effective prompt. (8) SPELL (Gonen et al., 2023) predicts prompt performance based on
the perplexity of the prompt. Here, we also use two extensions with FUTE: FUTE+TaskEmb with fTaskEmb(·, ·) and FUTE+TuPaTE with fTuPaTE(·, ·).

heading: 4.2.3 Overall Results
text: Table 3 presents the comparative results. No single baseline consistently outperforms the others across all scenarios. Focusing on performance and rate performance, MI excels on the Llama2-13B for the NLI task. LocalE stands out on the Mixtral for the SA task. FUTE+TaskEmb achieves superior performance on Llama2-13B and Llama2-70B for the SA task. However, given that performance and rate performance only focus on top-1 selection, we delve deeper into the NDCG scores for a comprehensive quality assessment.
GlobalE obtains the highest NDCG on Llama270B and Mixtral for the SA task. Its methodology, based on model entropy and confidence, appears less effective for the more challenging NLI task, likely due to the increased uncertainty even for correct prediction in LLM responses. SPELL achieves the best NDCG on Llama2-13B for both SA and NLI task, but underperforms on Llama2-70B and Mixtral. Given SPELL’s reliance on perplexity, this outcome may suggest that larger models possess greater resilience in delivering accurate predictions, even in scenarios of poor perplexity. FUTE outperforms others on Llama2-70B and
Mixtral for NLI tasks, demonstrating the effectiveness of our framework when handling difficult tasks. For the Llama2 model, FUTE+TaskEmb outperforms FUTE+TuPaTE on the SA task, but underperforms on NLI task. This disparity may stem from the inherent complexity of the NLI task, which tends to produce more unconfident and potentially noisier predictions from LLMs. TaskEmb leverages the parameters of the entire model, learning a broad knowledge but also risking overfitting to noisy data. TuPaTE utilizing the PEFT parameter, has a more constrained capacity that may better isolate relevant, non-noisy information. This specialization likely accounts for TaskEmb’s superior performance in SA tasks and vice versa. On Mixtral, which is based on SMoE architecture, FUTE+TaskEmb demonstrates superior results. The diverse experts within Mixtral may require more information storage capacity, a requirement that TaskEmb meets. In summary, while each baseline excels in certain tasks/models, FUTE emerges as a stable approach, achieving top-3 results across various scenarios.

heading: 4.2.4 Visualization
text: Figure 4 illustrates the visualization of learned MTEs for LLMs utilizing diverse prompts. This visualization reveals two principal clusters formed based on the prompts, highlighting the significant role prompts play in differentiating LLMs performance across task categories. While prompts across various tasks significantly influence MTEs,
the distinction becomes less pronounced within a single task category, where specific LLMs have a more substantial impact on the embedding outcomes. Especially within the NLI task category, MTEs are further divided into roughly three clusters, each corresponding to different LLMs. Mixtral MTEs also generally cluster towards the centre. See Appendix A.4 for a more detailed visualization, where prompts also slightly influence MTEs within the same task.

heading: 4.2.5 Efficiency Analysis
text: FUTE extends the existing task embedding method and utilizes embedding similarity to select prompts, leading to a distinctive computational complexity profile. The process requires evaluating a total of kp prompts across kD datasets. The computational complexity of FUTE is represented as O((kp + kD)co), where co represents the computation cost associated with the task embedding method. In contrast, the computational complexity for other baselines evaluated in this study is expressed as O(kpkDcb), where co represents the computation cost associated with baselines. The difference in complexity becomes significant with increasing values of kp and kD, where the product kpkD will exceed the sum (kp + kD). Thus, FUTE presents a potentially more efficient approach to prompt selection, especially as the number of prompts and datasets grows. However, the existing task embedding methods used in FUTE require finetuning, thereby co > cb. With a relatively small number of prompts (kp) and datasets (kD), FUTE does not demonstrate superior running efficiency in comparison to baselines. Nonetheless, co is more closely related to the inherent efficiency of the task embedding method itself, whereas FUTE focuses on extending these methods.

heading: 5 Conclusion and Future Work
text: 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.
Limitations
FUTE is designed to extend, rather than enhance, current task embedding learning methods. Consequently, it inherits potential limitations from the methodologies it builds on. However, we argue that this is more of a feature of FUTE rather than a limitation. By not altering current procedures, our framework leverages the strengths of existing task embedding methods, allowing for seamless integration of future advancements in the field. This adaptability means that as stronger task embedding methods emerge in the future, our framework will also exhibit improved performance without requiring modifications. We also acknowledge that any alteration to the surrogate base model requires the re-learning of task embeddings, which could be computationally intensive. However, we consider this issue intrinsic to the existing task embedding learning methods. If future developments in task embedding learning methods eliminate dependency on the base model, FUTE can be adjusted by simply omitting the surrogate base model as well.
Ethics Statement
FUTE aims to address challenges related to the interpretability of task embeddings across diverse models, especially for prompt-guided LLMs. While this framework holds promise for enhancing the utility and comparability of task embeddings across different models, it also raises some ethical concerns. For example, learning task embeddings requires the use of unsupervised data, which may not adequately encapsulate the essential attributes of a particular task, potentially resulting in biased task embeddings. Moreover, the reliance on a surrogate T5 model introduces potential ethical issue related to transparency. If the surrogate model fails to accurately replicate the behavior of the target model, it could lead to misleading results.

heading: Acknowledgements
text: This work was supported in part by the UK Engineering and Physical Sciences Research Council (EPSRC) through a Turing AI Fellowship (grant no. EP/V020579/1, EP/V020579/2) and a New Horizons grant (grant no. EP/X019063/1).

heading: A Appendix
text: A.1 Datasets in Language Model Experiment Text classification or regression (CR) contain: SNLI (Bowman et al., 2015), MNLI (Williams et al., 2018), QQP (Shankar Iyer and Csernai, 2017), QNLI (Wang et al., 2019), SST-2 (Socher et al., 2013), SciTail (Khot et al., 2018), CoLA (Warstadt et al., 2019), STS-B (Cer et al., 2017), MRPC (Dolan and Brockett, 2005), RTE (Dagan et al., 2006), WNLI (Levesque et al., 2011).
Question answering (QA) datasets contain: SQuAD-2 (Rajpurkar et al., 2018), NewsQA (Trischler et al., 2017), HotpotQA (Yang et al., 2018), SQuAD-1 (Rajpurkar et al., 2016), DuoRCp (Saha et al., 2018), DuoRC-s (Saha et al., 2018), DROP (Dua et al., 2019), WikiHop (Welbl et al., 2018), BoolQ (Clark et al., 2019), ComQA (Abujabal et al., 2019), CQ (Bao et al., 2016).
A.2 Visualization of Language Model Task Embedding
Figure A1 presents the visualization of task embeddings for language models learned by FUTE with TuPaTE. This visualization highlights the emergence of distinct clusters that reflect task-specific information, revealing that task embeddings are influenced by both language models and trained datasets:
Datasets Beyond the primary task clusters, some datasets demonstrate a pronounced impact on the formation of task embeddings. For example, models trained on SNLI and SST-2 exhibit a tendency to cluster together within the same task. A particular observation is the grouping of three blue points, representing task embeddings associated with the BoolQ, within the CR task cluster. This matches the format of BoolQ, a boolean answer for a question, which aligns more closely with CR tasks.
Language Model MTEs associated with the same language model, particularly BERT and
GPT2, demonstrate a propensity to cluster together. This phenomenon indicates that the inherent characteristics of a language model also influence the task embedding space.
A.3 Prompt Template
Table A1 and Table A2 list used prompts in SA and NLI in Section 4.2. Each contains 10 vanilla prompts and 3 CoT prompts. A.4 Visualization of LLM Task Embedding
with Prompts
Figure A2 demonstrates a visualization where prompt numbers are detailed in Table A1 and Table A2. Prompt task categories are differentiated by colors, and distinct shapes represent various LLMs. This visualization reveals that task embeddings are influenced by both the LLMs and the prompts used as follows:
Influence of Prompt This visualization demonstrates two influences: inter-task and intra-task influences. For inter-task influences, task embeddings are clearly separated into two main categories based on the task, indicating that prompts have a substantial impact on LLM performance across different tasks. For intra-task influences, within the same task, prompts also slightly influence task embeddings. For example, in the NLI task, all MTEs of prompt#9 are closer to the SA cluster than MTEs from the same LLM. Conversely, MTEs of prompt#10 diverge from the SA cluster. Also, the Chain of Thought (CoT) prompts#11, 12, and 13 form their own cluster, distinct from the main LLMbased clusters. A similar phenomenon is observed within the SA task, where prompts#7 and 8 diverge from the NLI cluster.
Influence of LLM Task embeddings tend to cluster by LLM, particularly for the NLI task, where embeddings from the same LLM cluster together. Mixtral’s task embeddings also generally cluster towards the center.
Figure A1: T-SNE visualization of language model MTEs using FUTE with TuPaTE. Different task categories are represented by various colors, while different models or datasets are indicated by distinct shapes.
No. Prompt
1 Does the movie review below make someone want to watch it? {text} 2 {text} How does the viewer feel about the movie? 3 {text} Did the reviewer enjoy the movie? 4 Does it seem like the reviewer who wrote this review liked the movie? {text} 5 Is the movie review below positive? {text} 6 {text} Did the reviewer find this movie good or bad? 7 {text} Is this review positive or negative? 8 {text} This is definitely not a positive or negative review 9 {text} Did the reviewer enjoy the movie? 10 {text} Did the reviewer find this movie good or bad? 11 Does the movie review below make someone want to watch it? {text} Reason step by step before answering. 12 {text} How does the viewer feel about the movie? Reason step by step before answering. 13 {text} Did the reviewer enjoy the movie?
Reason step by step before answering.
Table A1: Prompts used for SA
No. Prompt
1 Given that {premise} Does it follow that hypothesis yes, no, or maybe? 2 {premise} Using only the above description and what you know about the world, "{hypothesis}" is definitely correct, incorrect, or uncertain? 3 {premise} Question: {hypothesis} true, false, or maybe? 4 Take the following as truth: premise Then the following statement: ¨hypothesisïs true, false, or maybe? 5 {premise} Question: Does this imply that "{hypothesis}"? Yes, no, or maybe? 6 Given {premise} Should we assume that "{hypothesis}" is true? Yes, no, or maybe? 7 {premise} Based on the previous passage, is it true that "{hypothesis}"? Yes, no, or maybe? 8 {premise} Are we justified in saying that "{hypothesis}"? yes, no, or maybe? 9 Suppose it’s true that {premise} Then, is "{hypothesis}" always, sometimes, or never true? 10 Given {premise} Is it guaranteed true that "{hypothesis}"? yes, no, or maybe? 11 Given that {premise} Does it follow that hypothesis yes, no, or maybe? Reason step by step before answering. 12 {premise} Using only the above description and what you know about the world,
"{hypothesis}" is definitely correct, incorrect, or uncertain? Reason step by step before answering.
13 {premise} Question: {hypothesis} true, false, or maybe? Reason step by step before answering.
Table A2: Prompts used for NLI
Figure A2: T-SNE visualization of LLMs MTE using FUTE with TuPaTE. Prompts in different task categories are represented by colors, while different LLMs are indicated by distinct shapes. The numbers correspond to the prompt number as listed in Table A1 and Table A2.

heading: Future Work
text: We leave the exploration to other models in the future. Limitations
FUTE is designed to extend, rather than enhance, current task embedding learning methods. Consequently, it inherits potential limitations from the methodologies it builds on. However, we argue that this is more of a feature of FUTE rather than a limitation. By not altering current procedures, our framework leverages the strengths of existing task embedding methods, allowing for seamless integration of future advancements in the field. This adaptability means that as stronger task embedding methods emerge in the future, our framework will also exhibit improved performance without requiring modifications. We also acknowledge that any alteration to the surrogate base model requires the re-learning of task embeddings, which could be computationally intensive. However, we consider this issue intrinsic to the existing task embedding learning methods. If future developments in task embedding learning methods eliminate dependency on the base model, FUTE can be adjusted by simply omitting the surrogate base model as well