Mark My Words: Repurposing LLMs for Specialized Domains via Ability Tokens

We appreciate your valuable feedback and constructive comments. We have incorporated them in the revised paper. Below, we answer the questions raised in the review.

 

More ablations: "There are no side-by-side experiments comparing the traditional prompt-tuning and the proposed tuning both adopting regression heads or regular LM heads.""Is it possible to show the results of the regular prompt tuning with regression heads on the task of Descriptor Prediction or QED Prediction alone?"

Thank you for your suggestion. We added these baselines to our rebuttal (see Table 1 in the general comment above). Specifically, we evaluate PEFT baselines w. regression head and compare them to our method (ability tokens w. regression head). This setting removes the potential influence of the prediction head and the loss function to help us better understand the effect of ability tokens. The results show that our method outperforms prompt tuning on all tasks and outperforms all PEFT methods extended with the regression head on 3 of 4 tasks. The only exception is the binding affinity task, where our method underperforms LORA w. reg head, which nevertheless uses 12x more trainable parameters than we do (~1M vs. ~86K).

Apart from the "w. regression head" baselines you suggest, we add another set of ablation studies in Table 2 of the general comment above. We replace the learned embeddings of the domain markers with the text embeddings of the domain names (e.g., inserting the actual tokens "<", "Protein", ">" into the input sentence) to examine whether the domain markers have indeed learned useful domain knowledge from the auxiliary data. Our method still achieves outstanding performance, so domain markers are effective and necessary for improving downstream performance in specialized domains.

Beyond these new experiments, our original submission contains detailed ablation studies for different components of our approach (see Section 4.2.2, Figure 3, and Appendix A.1.3), showing the performance of

• Full method: domain markers + functional tokens + prediction head
• w/o domain markers
• w/o functional tokens
• w/o prediction head

The first setting obtains the best results. We have also studied the effect of ability token length in Figure 3.

 

Difference from prompt tuning: "The proposed approach is similar to 'prompt tuning' and its related techniques. The mere adaptation is to provide a different set of tokens per domain/task, which itself is relatively trivial."

We underscore our difference from existing PEFT methods (including prompt tuning) in Section 2.1 and 3.1. In particular, our method separates task domains from task functions—we use modularized tokens to represent different domains/functions and learn them using a hierarchical training protocol. These ideas are not trivial because they lead to distinct properties compared to existing methods:

• Reusability & generalization to unseen tasks: While existing PEFT methods require learning a distinct set of parameters for every single task, regardless of whether they share domains, our learned tokens can be shared across tasks and applied to new problems without additional training. For example, as demonstrated in Section 4.2, the <SMILES> marker can be independently applied to both drug combination and the binding affinity task. Additionally, in Section 4.1, the general functional token <Translate> exhibits zero-shot transferability to unseen tasks.
• Local applicability & multiple occurrence: Prompt tuning is "global" (applies to the entire query) and therefore always appears at the beginning, whereas our markers are "local" (apply to certain substrings of the input) and therefore can appear anywhere and multiple times in the prompt.
• Taking advantage of auxiliary data: For existing PEFT methods, learning the additional parameters (e.g., soft prompts, adapters) uses only the target dataset. In contrast, our design of domain markers allows us to leverage auxiliary unsupervised data and learn domain knowledge from them, which yields better empirical results than the PEFT baselines (Table 2 & 3).
• Sample efficiency for downstream tasks: By inserting learned domain markers into the input, we amortize the sample complexity needed to learn the functional tokens and solve downstream tasks.

Apart from these desirable properties, we also show in the new experiments (Table 1 of the general comment above) that our method significantly outperforms prompt tuning, which further provides empirical evidence that these methods differ not only conceptually, but also functionally.

Thank you for your question. We will add the above discussion to the paper for better clarity.

More ablations: "Ablation experiments show that the effectiveness of domain markers is relatively limited."

Apart from the "w. regression head" baselines you suggest, we add another set of ablation studies in Table 2 of the general comment above. We replace the learned embeddings of the domain markers with the text embeddings of the domain names (e.g., inserting the actual tokens "<", "Protein", ">" into the input sentence) to examine whether the domain markers have indeed learned useful domain knowledge from the auxiliary data. Our method still achieves outstanding performance, so domain markers are effective and necessary for improving downstream performance in specialized domains.

Besides these new experiments, our original submission contains detailed ablation studies for other components of our approach (see Section 4.2.2, Figure 3, and Appendix A.1.3), showing the performance of

• Full method: domain markers + functional tokens + prediction head
• w/o domain markers
• w/o functional tokens
• w/o prediction head

The first setting obtains the best results. We have also studied the effect of ability token length in Figure 3.

Thank you for your feedback and constructive comments, which we have incorporated in the revised paper. Below, we answer the questions raised in the review.

 

Difference from previous PEFT methods: "This method is similar to existing PEFT methods like Prompt-Tuning…What are the main differences between your approach and Prompt Tuning? It would be helpful to add a discussion about this."

We underscore our difference from existing PEFT methods (including prompt tuning) in Section 2.1 and 3.1. In particular, our method separates task domains from task functions—we use modularized tokens to represent different domains/functions and learn them using a hierarchical training protocol. This leads to:

• Reusability & generalization to unseen tasks: While existing PEFT methods require learning a distinct set of parameters for every single task, regardless of whether they share domains, our learned tokens can be shared across tasks and applied to new problems without additional training. For example, as demonstrated in Section 4.2, the <SMILES> marker can be independently applied to both drug combination and the binding affinity task. Additionally, in Section 4.1, the general functional token <Translate> exhibits zero-shot transferability to unseen tasks.
• Local applicability & multiple occurrence: Prompt tuning is "global" (applies to the entire query) and therefore always appears at the beginning, whereas our markers are "local" (apply to certain substrings of the input) and therefore can appear anywhere and multiple times in the prompt.
• Taking advantage of auxiliary data: For existing PEFT methods, learning the additional parameters (e.g., soft prompts, adapters) uses only the target dataset. In contrast, our design of domain markers allows us to leverage auxiliary unsupervised data and learn domain knowledge from them, which yields better empirical results than the PEFT baselines (Table 2 & 3).
• Sample efficiency for downstream tasks: By inserting learned domain markers into the input, we amortize the sample complexity needed to learn the functional tokens and solve downstream tasks.

Apart from these desirable properties, we also show in the new experiments (Table 1 of the general comment above) that our method significantly outperforms prompt tuning, which further provides empirical evidence that these methods differ not only conceptually but also functionally.

Thank you for the suggestion. We will add the above discussion to the paper for better clarity.

 

Difference from text instructions: "What is the difference between ability tokens with text instructions used in chat-aligned LLMs?"

The biggest difference is that our ability tokens are not represented as explicit text or actual words. Besides, our method has the following functional benefits:

• Knowledge compression: Ability tokens can compress large amounts of domain knowledge (via domain markers) or task knowledge (via functional tokens) into a few embedding parameters in a data-driven fashion. Such information is hard to summarize in a few sentences as text instructions.
• Task-agnostic: Ability tokens are more general, e.g., domain markers are task-agnostic and can be reused for new tasks in the same domain. Text instructions are always task-specific.
• Flexibility & generalization to unseen tasks: Ability tokens are composable and can generalize to unseen tasks in a zero-shot fashion, as exemplified in our multilingual experiments (Section 4.1).
• Effectiveness: Ability tokens are more effective and stable compared to hand-crafted text prompts, as shown in Table 2 & 3, where our method outperforms the hard prompt baseline by a large margin.

 

More baselines: "I think that training regression heads for numerical prediction problems may cause an unfair comparison with other adaption methods based on text generation.""Do you use the task-specific linear heads for other baseline methods, like LoRA and prompt tuning?"

Thank you for your suggestion. We added these baselines to our rebuttal (see Table 1 in the general comment above). Specifically, we evaluate PEFT baselines w. regression head and compare them to our method (ability tokens w. regression head). This setting removes the potential influence of the prediction head and the loss function to help us better understand the effect of ability tokens. The results show that our method outperforms prompt tuning on all tasks and outperforms all PEFT methods extended with the regression head on 3 of 4 tasks. The only exception is the binding affinity task, where our method underperforms LORA w. reg head, which nevertheless uses 12x more trainable parameters than we do (~1M vs. ~86K).

More domains: "Additionally, the experiments are mainly focused on one specific domain without testing the generalization of the approach to other domains.""Evaluations in other domains would strengthen the paper's findings."

Thank you for the suggestion. We showcase the general effectiveness of our approach across a diverse set of ten domains, encompassing eight languages, protein sequences ("MTVPDRSEIAGKWYVV…"), and chemical formulas ("CCOC@HCc1…"). We deliberately chose the latter two domains because (1) they represent non-linguistic, highly specialized domains with distributions significantly divergent from natural language distributions; (2) they are relatively well-understood with established baselines that we can compare to; and (3) they are substantially different from each other, allowing us to investigate the interactions between them. We acknowledge the potential for expanding the scope of domains in future iterations of this work.

 

Embedder initialization: "Is there any specific procedure for initializing the embedding for ability tokens, or is it done in a standard manner?"

We have outlined the initialization procedure in Section 3.2 (last paragraph on page 4). Specifically, the ability tokens are initialized using the average embedding of the LLM's original dictionary, scaled by a factor to match the norm of this average embedding with the average norm of all token embeddings in the LLM's original dictionary.

 

Ability token length: "Are the lengths of domain markers and functional tokens always the same, or can they vary depending on the task?"

The length $p$ is specified as a hyperparameter in Section 3.2 (see the 6th-to-last line on page 4), so it is variable. We explore its impact on downstream performance using the drug combination dataset (Section 4.2.2 and Figure 3).

Thank you for your feedback and constructive comments, which we have incorporated in the revised paper. Below, we answer the questions raised in the review.

 

Data contamination: "However, a major concern is that the performance seems to be primarily driven by overlap between the training data of the LLMs and the downstream tasks.""The ability tokens, which generally correspond to dataset or task names, may result in LMs memorizing information about the datasets during pretraining."

We would like to clarify a couple of points to address these misunderstandings about our work.

First, it is important to note that the ability tokens do not "correspond to dataset or task names." They exist only in the embedding space and are not represented as actual words. For example, to learn the token representing the SMILES domain, we introduce "<ability token 0>" to the tokenizer with a newly initialized embedding. Then, we prepend "<aibility token 0>" (not the word "SMILES" or any dataset name related to SMILES) to SMILES strings and learn the embedding for "<aibility token 0>" via next-token prediction. At inference time, "<ability token 0>" is inserted into the input. Similarly, when learning the binding affinity function, the functional token added to the tokenizer is "<aibility token 1>", and its embedding is not tied to the task name "binding affinity."

Second, the datasets used to train the ability tokens do not appear in pretraining. Our goal of using ability tokens is to inject additional domain/task knowledge into the LLM by leveraging auxiliary data. For instance, the dataset we use to train the <Protein> marker is from a recently published source [1], which should not have appeared in the pretraining corpora.

Lastly, the datasets used for evaluation do not appear in pretraining. For instance, we generate the QED and descriptor prediction datasets ourselves using specialized Python packages. These data have not appeared anywhere on the web, so there should not be a memorization problem.

We hope this clarifies any misconceptions. Please feel free to share if you have additional concerns. Thank you!

[1] Blanchard AE, Gounley J, Bhowmik D, et al. Language models for the prediction of SARS-CoV-2 inhibitors. The International Journal of High Performance Computing Applications. 2022; 36(5-6):587-602.

 

More ablations: "More ablation studies are required to demonstrate the effectiveness of the ability tokens. For example, evaluating the llama-7b model with and without ability tokens for different tasks."

We have performed such ablations in the paper. In particular:

• The "hard prompt" baseline in Table 2 and 3 is llama-7b without ability tokens and regression head. This demonstrates the effect of using both ability tokens and regression head.
• The "w/o marker" baseline in Figure 3 (and Appendix A.1.3) is llama-7b with functional tokens and regression head but without domain markers. This demonstrates the effect of domain markers.
• The "w/o functional" baseline in Figure 3 (and Appendix A.1.3) is llama-7b with domain markers and regression head but without functional tokens. This demonstrates the effect of functional tokens.

For rebuttal, we add two additional sets of ablations as requested (see the general comment above).

• In Table 1, we evaluate PEFT baselines w. regression head and compare them to our method (ability tokens w. regression head). This setting removes the potential influence of the prediction head and the loss function to help us better understand the effect of ability tokens. The results show that our method outperforms prompt tuning on all tasks and outperforms all PEFT methods extended with the regression head on 3 of 4 tasks. The only exception is the binding affinity task, where our method underperforms LORA w. reg head, which nevertheless uses 12x more trainable parameters than we do (~1M vs. ~86K).
• In Table 2, we replace the learned embeddings of the domain markers with the text embeddings of the domain names (e.g., inserting the actual tokens "<", "Protein", ">" into the input sentence) to examine whether the domain markers have indeed learned useful domain knowledge from the auxiliary data. Our method still achieves outstanding performance, so domain markers are effective and necessary for improving downstream performance in specialized domains.

We appreciate your valuable feedback and constructive comments. We have incorporated them in the revised paper. Below, we answer the questions raised in the review.

 

More ablations: "The paper lacks ablation results, which is crucial to demonstrate the effectiveness of the ability tokens."

Thank you for your suggestion. We conducted two additional sets of ablation experiments to provide further understanding of the contribution of the ability tokens. The detailed results are shown in the general comment above.

• In Table 1, we evaluate PEFT baselines w. regression head and compare them to our method (ability tokens w. regression head). This setting removes the potential influence of the prediction head and the loss function to help us better understand the effect of ability tokens. The results show that our method outperforms prompt tuning on all tasks and outperforms all PEFT methods extended with the regression head on 3 of 4 tasks. The only exception is the binding affinity task, where our method slightly underperforms LORA w. reg head, which nevertheless uses 12x more trainable parameters than we do (~1M vs. ~86K).
• In Table 2, we replace the learned embeddings of the domain markers with the text embeddings of the domain names and prompt the LLM, as you suggest. Our method still achieves outstanding performance, which indicates that domain markers have indeed learned useful domain knowledge from the auxiliary unsupervised data.

Beyond these new experiments, our original submission contains detailed ablation studies for different components of our approach (see Section 4.2.2, Figure 3, and Appendix A.1.3), showing the performance of

• Full method: domain markers + functional tokens + prediction head
• w/o domain markers
• w/o functional tokens
• w/o prediction head

The first setting obtains the best results. We have also studied the effect of ability token length in Figure 3.

 

Understanding training protocol: "What is the effectiveness of the 3-stage training process and how much each stage contributed to generalization results?"

In each stage of the 3-stage training process, we learn a different type of ability token. The effect and contribution of individual stages are summarized below:

• Stage 1: training domain markers. This stage allows us to inject general domain knowledge into the prompting process. It improves downstream performance since the "w. domain markers" setting outperforms the "w/o domain markers" setting (Figure 3). This stage also lays out the foundation for later stages because training functional tokens requires learned domain markers. It contributes the most to the generalization results, given that domain markers are task-agnostic and applicable to various in-domain tasks.
• Stage 2: training single-domain functional tokens. This stage allows the model to learn task knowledge solely from feature-label pairs (without human instructions) and compress the information into a few embedding parameters. Stage 2 boosts the efficacy of our method since the "w. functional tokens" setting outperforms the "w/o functional tokens" setting (Appendix A.1.3).
• Stage 3: training multi-domain functional tokens. This stage is similar to stage 2, but the target functions are multi-domain, so we extend the capacity of LLMs to solving multi-domain tasks like drug-protein binding affinity prediction (Section 4.2.3). Moreover, as we show in the multilingual translation experiments (Section 4.1), multi-domain tokens like <Translate> can be combined with different markers and applied to unseen tasks during training. Thus, this stage also contributes to our markup system's zero-shot generalization ability.

We will add the above discussion to the paper for better clarity.

Relevant work: "A common practice in adapting Llama to specialized domains is to add specialized tokens as words… What is the advantage of the proposed methods to this way of adding ability tokens?"

Thank you for pointing out this line of related work. The approach you referenced involves two key components: adding "<", "molecule", ">" to the input, and using adapters like LORA. We discuss our advantages relative to both components:

• Better controllability & conditioning effect. The efficacy of using pretrained tokens like "<", "molecule", ">" relies heavily on their occurrence in the pretraining corpora. However, end users do not have control over this, e.g., we are uncertain how frequently the word “molecule” occurs (if the frequency is low, then adding the tokens may not be useful) and whether these occurrences align with the target domain. For example, the domain of "molecule" is different in the sentence "2 atoms of H and 1 of oxygen form a molecule of water" (chemistry) and the sentence "the four molecules of life are proteins, carbohydrates, lipids and nucleic acids" (biology). Our method addresses this limitation by learning continuous embeddings for specialized domains explicitly using data from the target domain. This allows more fine-grained control over how we condition the LLM.
• Generalization to unseen tasks. Using LORA after adding "<", "molecule", ">" to the input still necessitates learning a distinct set of parameters for each downstream task, even though the tasks may share knowledge, as discussed in Section 3.1. On the contrary, our learned domain markers can be used for various tasks in the same domain (e.g., the <SMILES> token can be applied to both drug combination and binding affinity prediction). The learned functional tokens also exhibit zero-shot generalization to unseen domains (e.g., the <Translate> token can generalize to unseen language pairs, as we show in Section 4.1). Thus, our ability tokens are reusable, lowering the deployment costs for new tasks.
• Taking advantage of auxiliary data: For existing PEFT methods, learning the additional parameters (e.g., soft prompts, adapters) uses only the target dataset, whereas our design of domain markers leverages auxiliary unsupervised data and learns domain knowledge from them. This yields better empirical results (Table 2 & 3) and amortizes the sample complexity needed to learn the functional tokens.

To provide a more direct comparison between our method and the approach you mention, we evaluate LORA/prompt tuning with domain names inserted into the input text (see Table 2 in the general comment). The results show that our method is more effective in addition to being more controllable and generalizable.

Thank you for the references. We will add the above discussion to the paper.

 

Data efficiency: "What is the data efficiency of the proposed method, compared to prompt tuning and LORA?"

We discuss data efficiency in two scenarios:

• Case 1: no relevant domain marker or functional token learned yet. Prompt tuning and LORA learn task-specific parameters entirely from the target dataset. However, our method requires learning task-agnostic domain markers using unsupervised domain data, in addition to learning the functional token using the target dataset. This entails using more data than the PEFT methods, but in practice, the additional unsupervised data is more widely available and much easier to obtain (or generate) than labeled data.
• Case 2: there are relevant markers and functional tokens trained for previous tasks which we can use off-the-shelf. For instance, in our multilingual translation experiments, we have already learned various language markers and the functional token. Thus, we can generalize zero-shot to new translation tasks by directly inserting the learned ability tokens into the input. As we don't have to train any additional parameters, we achieve better data efficiency than the PEFT methods, which still require learning new parameters for new tasks.

We will add the above discussion to the paper.