Prompt-based multi-task learning for robust text retrieval

The exponential growth of digital information necessitates the development of robust text retrieval methods since most of the methods are domain or task-specific which limits their implementation. In this case multi-task learning is a promising alternative as it helps a model to have more meaningful embeddings; however such cases require usage of task separation methods. Many studies explore multi-task learning to improve generalization but tend to focus on large models. However, in real-world, speech analytics tasks that require searching through hundreds of millions of vectors in real-time, smaller models become more appropriate. This paper presents a novel approach to enhance the robustness of multi-task text retrieval models through the use of prompts. We use contrastive learning to train encoder models both in single-task and multi-task configurations and compare their performances as well as analyze the efficiency of different prompt usage strategies including hard prompts represented by explicit natural language instructions and soft prompts of varying lengths represented by model special tokens. Experiments are conducted by applying prompts to both the query and candidate document as well as to queries only keeping the candidate without prompts to reuse pre-encoded candidates in multi-task retrieval without significant quality loss. The obtained results are compared using R@1, R@5, and MRR metrics which are most applicable for evaluating in-domain and out-of-domain search. Single-task models show better performance on in-domain training data, while multi-task models demonstrate superior performance on out-of-domain data highlighting their increased robustness to domain shifts. Applying prompts to both elements–query and document–yields better performance than applying them solely to the query. Soft prompts are found to be preferable to hard as they better adapt the model to different domains. The findings of this study can be useful for improving text retrieval models, especially in scenarios involving multi-task systems where high adaptability and performance on new data are required. Trainable prompts could be an effective tool for enhancing the flexibility of models in various applications, such as information retrieval and question-answering systems.

1. Hambarde K.A., Proença H. Information retrieval: recent advances and beyond. IEEE Access, 2023, vol. 11, pp. 76581–76604. https://doi.org/10.1109/access.2023.3295776

2. Zhang W., Xiong C., Stratos K., Overwijk A. Improving multitask retrieval by promoting task specialization. Transactions of the Association for Computational Linguistics, 2023, vol. 11, pp. 1201– 1212. https://doi.org/10.1162/tacl_a_00597

3. Muennighoff N., Tazi N., Magne L., Reimers N. MTEB: Massive Text Embedding Benchmark. Proc. of the 17th Conference of the European Chapter of the Association for Computational Linguistics, 2023, pp. 2014–2037. https://doi.org/10.18653/v1/2023.eacl-main.148

4. Thakur N., Reimers N., Rücklé A., Srivastava A., Gurevych I. BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models. Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021, pp. 105.

5. Muennighoff N., Su H., Wang L., Yang N., Wei F., Yu T., Singh A., Kiela D. Generative representational instruction tuning. arXiv, 2024, arXiv:2402.09906. https://doi.org/10.48550/arXiv.2402.09906

6. Reimers N., Gurevych I. Sentence-BERT: Sentence embeddings using siamese BERT-Networks. Sentence-BERT. Proc. of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), 2019, pp. 3982–3992. https://doi.org/10.18653/v1/d19-1410

7. Wang L., Yang N., Huang X., Jiao B., Yang L., Jiang D., Majumder R., Wei F. Text embeddings by weakly-supervised contrastive pretraining. arXiv, 2024, arXiv:2212.03533. https://doi.org/10.48550/arXiv.2212.03533

8. Chen T., Zhang M., Lu J., Bendersky M., Najork M. Out-of-Domain semantics to the rescue! Zero-shot hybrid retrieval models. Lecture Notes in Computer Science, 2022, vol. 13185, pp. 95–110. https://doi.org/10.1007/978-3-030-99736-6_7

9. Ruder S. An overview of multi-task learning in deep neural networks. arXiv, 2017, arXiv:1706.05098. https://doi.org/10.48550/arXiv.1706.05098

10. Maillard J., Karpukhin V., Petroni F., Yih W., Oğuz B., Stoyanov V., Ghosh G. Multi-task retrieval for knowledge-intensive tasks. Proc. of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing. Vol. 1, 2021, pp. 1098–1111. https://doi.org/10.18653/v1/2021.acl-long.89

11. Su H., Shi W., Kasai J., Wang Y., Hu Y., Ostendorf M., Yih W., Smith N.A., Zettlemoyer L., Yu T. One embedder, any task: Instruction-finetuned text embeddings. Findings of the Association for Computational Linguistics: ACL 2023, 2023, pp. 1102–1121. https://doi.org/10.18653/v1/2023.findings-acl.71

12. Li X.L., Liang P. Prefix-tuning: Optimizing continuous prompts for generation. Proc. of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing. Vol. 1, 2021, pp. 4582–4597. https://doi.org/10.18653/v1/2021.acl-long.353

13. Matveev Y, Makhnytkina O, Posokhov P., Matveev A., Skrylnikov S. Personalizing hybrid-based dialogue agents. Mathematics, 2022, vol. 10, no. 24, pp. 4657. https://doi.org/10.3390/math10244657

14. Posokhov P., Apanasovich K., Matveeva A., Makhnytkina O., Matveev A. Personalizing dialogue agents for Russian: Retrieve and refine. Proc. of the 31st Conference of Open Innovations Association (FRUCT), 2022, pp. 245–252. https://doi.org/10.23919/fruct54823.2022.9770895

15. Posokhov P., Matveeva A., Makhnytkina O., Matveev A., Matveev Y. Personalizing retrieval-based dialogue agents. Lecture Notes in Computer Science, 2022, vol. 13721, pp. 554–566. https://doi.org/10.1007/978-3-031-20980-2_47

16. Wang L., Yang N., Huang X., Yang L., Majumder R., Wei F. Improving text embeddings with large language models. Proc. of the 62nd Annual Meeting of the Association for Computational Linguistics. Vol. 1, 2024, pp. 11897–11916. https://doi.org/10.18653/ v1/2024.acl-long.642

17. Xu L., Xie H., Qin S.-Z.J., Tao X., Wang F.L. Parameter-efficient fine-tuning methods for pretrained language models: A critical review and assessment. arXiv, 2023, arXiv:2312.12148. https://doi.org/10.48550/arXiv.2312.12148