Sweeping Heterogeneity with Smart MoPs: Mixture of Prompts for LLM Task Adaptation

We thank the reviewer for the detailed and thoughtful review. Below, we respond to the points raised one by one. We hope our responses will resolve any further concerns, and we are available for other questions.

-Regarding including all available prompts: Our baseline from [Xu 2023] indeed includes all available prompts in both training and inference. Our experiments and analysis show that when all prompts are used for diverse tasks (and cross-clients in federated learning), the training interference between different tasks and clients results in significant performance degradation, mitigated by our proposed gating function and expert prompts. %Our gating function is

-Regarding comparison with other PEFT methods: We acknowledge several PEFT methods in current practice. We argue that prompt tuning is more widely adopted for instruction tuning on LLMs due to its straightforward implementation, which does not require any modifications to the LLM model. Moreover, prompt tuning is more parameter efficient compared to the adapter family. Therefore, to be as efficient and realistic as possible in our paper, we focus exclusively on the prompt tuning framework; however, we acknowledge that such experimental comparison is interesting. Finally, prompt tuning is more suitable for our gating function as it is more natural to select combinations of subsets of prompts depending on the current input rather than combining different adapters.

-Regarding other designs of a gating function: We thank the reviewer for the exciting idea of different designs of the gating function. Directly using kNN to select from prompts still needs embedding the current input question from another embedding network, which might result in more computation and memory cost. In our method, the embedding for our gating function is directly reused from the intermediate layer embedding from the original LLM. Further, as all our prompts are randomly initialized during training, our gating function additionally needs to discover the set of skills across diverse tasks and guide the specialization of expert prompts towards such skills. This can only be done using a trainable gating network instead of a fixed clustering scheme, as in kNN.

-Regarding the ordering of prompts: As our prompts are randomly initialized, our method does not assume any ordering of prompts before training. During training, we treat prompts as additional tokens with different position encoding.

-Regarding evaluation metrics: Thank you for suggesting other metrics. We understand the importance of a unified evaluation metric across different tasks and datasets. That is why we chose to use PPL for our selected set of tasks, as it is particularly suitable for long text generation. However, we are open to incorporating more metrics in future drafts, and we appreciate your input.

-Regarding [Vu 2022]: We argue that [Vu 2022], similar to [Wang 2023], focuses on different scenarios and cannot be directly applicable to the settings described in our paper. In [Vu 2022], authors also assume the availability of task/domain labels during prompt tuning and prompt transfer for test clients, which, as mentioned above, is considered a significant limitation. Further [Vu 2022] considers two approaches: generic and targeted Soft Prompt Transfer (SPOT). The generic SPOT case is essentially the same as our prompt tuning baseline that we compare against in the paper, where the prompts are directly used for test clients. For targeted SPOT, it is required first to perform prompt tuning on test clients and then use similarity scores of the resulting prompt tokens to select source prompts. However, this violates the zero-shot assumption we have for the test clients, in which we do not have any ground truth label for any prompt tuning; thus, such a baseline cannot be used in our ablation studies.

-Regarding [Wang 2022] and [Asai 2022]: We thank the reviewer for referencing other works proposing a mixture of expert frameworks. As [Wang 2022] only considers single-task PEFT training, the functionality of their gating network and adapters is different from our gating function and expert prompts. In particular, our gating function needs to discover the set of skills across diverse tasks and guide the specialization of expert prompts toward such skills. During inference, our gating function, given an unseen test client, needs to analyze the input question and correctly select the combination of relevant skills/expert prompts. In summary, we require our expert prompts to be highly specialized in different skills, which is in direct contrast with [Wang 2022], which utilizes a consistency regularization loss to encourage the similarity between expert adapters. Accordingly, our gating function cannot be a random routing network as in [Wang 2022], as it needs to analyze the current task and dynamically select relevant specialized skills/expert prompts.

For the case of [Asai 2022], as we have mentioned in the related work section, [Asai 2022] assumes task labels for source prompt tuning, which are not available in our scenarios. Therefore, the gating function in [Asai 2022] does not need to discover the skills and guides the training of specialized experts in highly mixed task distribution without any task/domain labels.

-Regarding ablation studies: We appreciate the reviewer's insightful comments. Our method uses LLM intermediate embedding instead of a separate embedding for the gating function to avoid the additional cost of training and inference incurred by having a different embedding network (such as another LLM or transformer model) for the gating function. We ensure a similar training/inference cost with a prompt tuning baseline by directly reusing the intermediate layer embedding. At the same time, to further reduce the computation costs (for both our method and baseline), we strategically add prompts only to the middle layers of the model, thus avoiding backpropagation of the full model during training. Therefore, We set $L_{mid} = 10$ for an LLama-7B model with $L = 20$. Such choice of $L_{mid} = 10$ is to balance two conflicting requirements: (1) we want to inject prompts as late as possible to reduce back propagation cost during training and increase depth/capacity for the embedding network of the gating function, and (2) prompts should be injected early to have more capacity in influencing the pretrained LLM network.

-Regarding paper structure: We appreciate the feedback and will separate those sections to better present our proposed method for the camera-ready version.

We thank the reviewer for the detailed and thoughtful review. Below, we respond to the points raised, one by one. We hope our responses resolve further concerns and are available for other questions.

-Regarding uncompressed LLMs: Our method can be applied as a generic method agnostic to federated/centralized and compression ratios. However, we are motivated to focus on aggressively compressed LLM, as in real-life scenarios, users mostly have very limited computation and memory resources for LLM deployment. Unfortunately, we cannot train full-size LLM models on our GPUs to cover other scenarios due to limited computation and memory resources. We agree that, although our method achieves a consistently significant advantage compared with baselines in all pruning ratios, the absolute gain decreases when the pruning ratio goes down. From our experiments, it appears that this is because a higher compression ratio will impose more implicit prompt training interference since trainable prompts are responsible for both recovering more model capacity loss –due to higher compression– and model adaptation for downstream tasks. Yet, without additional experiments on full-size LLMs, this cannot be safely assumed for all compression ratios, and we will certainly look into this as future work.

-Regarding additional baselines [Wang 2023] [Vu 2022] [Wang 2022]: We thank the reviewer for the thoughtful comments on the baselines used in our paper. Before discussing each of these papers in detail below, we state that it is a critical point in our motivation to be as close as possible to real-life LLM instruction tuning scenarios; this leads us to consider a limited set of baselines to be applied to our scenario.

-Regarding [Wang 2023]:
We argue that our proposed method and targeted scenario are different from those in [Wang 2023]. Firstly, our paper assumes the task-type/domain labels and questions and ground truth answers during training are unavailable. This is closer to a real-life scenario as the instruction tuning data are highly mixed collections of tasks asked by the diverse users without providing any task/domain labels. Yet, [Wang 2023] assumes knowledge of task labels for both the source and target tasks when performing prompt training and prompt distillation. Second, [Wang 2023] requires labeled data from target tasks on test clients to perform prompt distillation, which might not be available in a real-life case. Instead, we focus on a more realistic zero-shot scenario, requiring our gating function to select correct prompt experts for the current test client data without any knowledge (zero-shot) of ground truth label samples or label distribution of such test client. Finally, our method can be applied to centralized and federated learning scenarios. In federated learning scenarios, our approach is shown to automatically mitigate model drift problems in federated learning without any additional cost. Such might not be the case for the proposed method in [Wang 2023].

We thank the reviewer for the review and references. Below, we respond to the points raised one by one. We hope our responses will resolve further concerns and are available for questions.

-Regarding comparison to [1]: We thank the reviewer for pointing us to this paper. We argue that our proposed method and targeted scenario differ from those described in [1]. In particular, first, we focus on LLMs' standard instruction tuning scenario, which will include diverse tasks such as open-QA, closed-QA, creative writing, etc. On the other hand, [1] only utilizes relational datasets, focusing on tuning the relation extraction/inference abilities of language models. Secondly, we assume the task-type/domain labels are unavailable with questions and ground truth answers (as in $(x,y)$ pairs during training). This is closer to real-life scenarios as the instruction tuning data are highly mixed collections of tasks asked by the diverse users without providing any task/domain labels. Yet, in [1], authors assume the knowledge of task/domain label (the relation type $r$) for their gating function throughout the training. Thus, our gating function is different in architecture and training dynamics from the one in [1], as our gating function needs to discover the relevant set of skills in diverse tasks without any task/domain labels and guides the training of specialized expert prompts. As shown in Figures 2 and 3, our gating function successfully learns to select similar skill prompts for similar tasks without knowledge of task labels or similarity between tasks. This functionality is not straightforward whether it is possible in [1].

-Regarding comparison to [2]: We again appreciate the reviewer's pointed paper. We argue that our proposed method and targeted scenario differ from [2]. In particular, for similar reasons in comparison to [1], [2] assumes knowledge of task labels for both the source and target tasks when performing prompt training and prompt distillation. This is a significant limitation since the real-life collected instruction training data mostly lacks task labels. Second, [2] requires labeled data from target tasks on test clients to perform prompt distillation, which might not be available in real-life cases. In contrast, we ask our gating function to select the correct prompt experts for the current test client data without any knowledge (zero-shot) of ground truth label samples or label distribution of such test clients. We agree that our method, as pointed out by the reviewer, can be applied to centralized and federated learning scenarios, further demonstrating its flexibility and robustness under different setups.

-Regarding any contributions on compressed LLMs: Our primary motivation for focusing on aggressively compressed LLMs is to simulate real-life scenarios where users typically have limited computation and memory resources for LLM deployment. As stated in Section 4, we directly utilize SparseGPT as our compression technique. Our contribution in this area is to demonstrate that prompt tuning is much more difficult with higher compression ratios, but our proposed mixture of prompts framework can alleviate this.

-Regarding how MoPs perform when lower complexity models are used: Although our method achieves a significant PPL loss decrease in all pruning ratios compared to the baseline, we can observe a relatively significant degradation from 85% to 90% pruning ratio. Thus, we can assume such a threshold should be between these ratios. Of course, such "phase transitions" could happen at different levels if different families of models, tasks, and datasets are used; we indicate that such a phenomenon might happen, and we report that in this paper.

-Regarding evaluation on live, real-world datasets: We acknowledge the value of testing on real-world datasets. However, we have focused on standard instruction tuning datasets due to the lack of access to such data and testing environments. This decision is based on two main reasons: (1) as these instruction tuning datasets contain task labels (which are not used during training), we can ensure the heterogeneity of task and data distribution in federated learning scenarios when we create local datasets; (2) we can also more easily visualize and analyze our gating function and expert prompts, as demonstrated in Figure 2 and 3.

We thank the reviewer for the detailed and thoughtful review. Below, we respond to the points raised one by one. We hope our responses will resolve any further concerns, and we are available for any other questions.

-Regarding implementation details: We thank the reviewer for the suggestion, and we will add more details on implementation in the future draft, which was shortened due to the page limit.

Our proposed framework is simple (i.e., only adding the gating function as an additional module) and needs minimal modification from prompt tuning baseline and no hyper-parameter tuning. Our method can be divided into $i)$ the expert prompts and $ii)$ the gating function. For our expert prompts, we are using 70 prompts.

To decrease the size/complexity of the gating function and increase the capacity of each expert, we further split prompts into seven groups, each with ten prompts: 10 prompts per task for a 7-task problem; but as the results we have shown, it is not necessary to know a priori the number of tasks, as the system adapts and utilizes the required amount of ``experts''; see, e.g., Figure 3 in the main text. Therefore, the gating function will generate expert weights for each group, and all prompts within each group are assigned the same expert weight from the gating function in the following layers.

Our gating function is a simple shallow MLP network that uses intermediate layer embedding as input and output expert weight (a seven-dimensional vector) with softmax. We inject all expert prompts in the 10th layer (the middle layer in LLAMA-7B mode). As described in the paper, this further reduces the training cost by only needing backpropagating to the middle layer.

-Regarding the number of datasets used: We understand your concerns about not using more datasets, but the computation resources available were a limitation for conducting more experiments. Nonetheless, we emphasize that the Dolly-15K and Super-Natural Instruction datasets cover many sub-tasks, are non-trivial/easy datasets, and are widely used as representative LLM instruction tuning datasets.

-Regarding analysis on mitigation of model drift: First, we note that any analysis can only be restricted to an experimental study rather than theoretical research, given the task complexity and models involved. As shown in Figure 3, our gating function learns to assign separate sets of experts for different tasks and the same group for similar tasks throughout training. In federated learning scenarios, the gating function during local training will only update the relevant subset of experts based on the task type of local training data. During synchronization, we update each expert by only collecting its updates from those clients (with training data according to task type). Therefore, we avoid averaging expert models trained on different task types, mitigating the model drift problem. We provide more analysis on the gating function on different pruning ratios in our appendix sections B and C.

-Regarding the scalability of the gating function: We thank the reviewer for this question. Regarding scalability concerning the number of tasks, as we have pointed out in Section 3.1, we intend our expert prompts not to be simple experts specialized in one task; in that case, we would trivialize our gating function to a simple task classification and only learn the one-to-one relationship between prompt and tasks. Instead, each expert prompt should represent expert knowledge of a particular skill. For each task, the gating function should learn to select/utilize a combination of relevant skills (prompts) for the current task. As shown in Figures 2 and 3, our gating function learns to choose a similar subset of prompts as skill experts for similar tasks. In this way, we decouple the number of tasks with the number of prompts or the size of the gating function, as we assume those tasks, although diverse, can all be solved with a combination of finite skills. This indicates (only experimentally) that our gating function and the collection of trained prompts will be robust with the number of tasks. Finally, we intentionally choose the Super-Instruct dataset to account for task complexity, as it contains more complex tasks than Dolly-15K. Our experimental results show that our method maintains its advantage in both datasets.