REMED -- 一种高效嵌入微调的检索增强生成方法

在医疗这样的高风险领域，大模型“一本正经地胡说八道”可是真要命。如何让LLMs的答案更可靠，一直是行业焦点。这项发表在IJCNN 2024的研究，来自华南师范大学、上海人工智能实验室和上海交通大学，他们提出的REMED框架，或许找到了一个不错的解。

摘要

大语言模型（LLMs）在医疗等领域的“幻觉”问题，催生了检索增强生成（RAG）方法。但RAG的效果，很大程度上取决于检索质量。为此，我们设计了一个专门针对医学文档检索的框架——REMED，并在IJCNN 2024上分享了这套方案。REMED整合了数据集构建、高效的嵌入微调模型EM-FT、检索增强生成以及人工评估。其中，EM-FT通过高效的嵌入微调，能对预训练模型的医学句子表示进行端到端优化，大幅提升医学检索的性能。我们用对比学习作为损失函数，让模型能更精准地捕捉查询和文档的相似性。在两个数据集上的验证结果显示，相比直接用嵌入模型检索，我们的方法在MMD上召回率和精确度提升了3.2%到6.0%，在MPD上更是提升了14.4%到42.6%。

研究背景

GPT-4这类大模型的能力确实惊艳，但它们的“幻觉”问题也让人头疼。尤其是在医疗、法律这些领域，一个微小的错误都可能引发严重后果，信息处理必须慎之又慎。

检索增强生成（RAG）是目前应对幻觉问题的主流方案之一。它的思路很直接：把大模型和外部知识库结合起来，让生成过程有据可查。我们在这项工作中，自主收集了两个数据集：一个是药品信息数据集MMD，一个是医学论文数据集MPD。通过将生成内容与可靠知识库进行比对，生成结果的可靠性有了显著提升。

研究方法

数据集

数据集是研究的基石。我们构建了两个专用的医学数据集：

• Medical Menu dataset (MMD)
  ：这是一个高质量的医学信息检索基准，数据源自权威的“WHO Medicine”，涵盖了“国家药典”中的所有药物信息，共计20多万条记录。

• Medical Paper dataset (MPD)
  ：我们从美国国家生物技术信息中心（NCBI）随机抽样了1,000篇论文，构建了一个样本医学文档数据集。

数据预处理

两个监督数据集的构建流程是重头戏。

上图展示了MMD的结构化细节。简单来说，就是将原始药品数据转化为结构化文本，每个样本包含药品名称、生产厂家、类别和属性等，标签则来自真实世界的标注。

MPD的处理流程则稍有不同：先用论文标题或随机段落作为LLM输入，生成相关查询；再计算每个查询和段落的BM25分数，并据此分配标签。

REMED框架

REMED框架整合了数据集构建、高效的嵌入微调模型EM-FT、检索增强生成以及人工评估四个环节。

我们的方法用监督数据集训练嵌入模型，并引入对比学习来提高性能，成本却很低。在密集检索领域，M3e和E5等嵌入模型是关键。在此基础上，我们加入了一个新模块——Gate Linear Unit (GLU)。EM-FT整合了整个嵌入模型和GLU模块。值得注意的是，微调过程中嵌入模型参数被冻结，训练只集中在GLU模块上。这样做的好处是，既能确保医学数据实时更新，又能防止私有数据泄露。

损失函数

我们设计了一个对比损失作为监督信号，来优化嵌入空间。目标是让与查询相关的文档在向量空间中更接近，不相关的则更远。

通过将对比学习纳入EM-FT的微调，模型在捕捉查询和文档相关性方面变得更高效，能有效拉近正向文本，推远负向文本。

EM-FT模型

EM-FT模型架构集成了两个核心组件：嵌入骨干（Embedding Backbone）和可训练的EM头部（Trainable EM Head）。这个紧凑的EM头部设计包含层归一化（LayerNorm）和两个带激活函数的线性层。这里我们提供了两种激活函数选择：GELU和Swish，灵感来自最新的LLM设计进展。

人类评估

我们用同一个问题“奥利司他胶囊的适应症是什么？”来对比M3e模型和我们EM-FT模型的LLM生成效果。结果一目了然：M3e检索结果不准确，导致LLM给出了错误答案；而我们的方法检索准确，LLM自然也能正确回答。

研究结论

我们在MMD和MPD两个数据集上，用多个指标进行了全面测评，重点关注前K=10个推荐结果。

从表1可以得出几个关键结论：

• EM-FT（w/L2）方法在损失函数中引入L2正则化，显著提升了模型性能，召回率提高了约6%。这验证了正则化能有效避免过拟合，增强泛化能力。
• EM-FT（使用GELU激活函数）优于EM-FT(s)（使用Swish激活函数）。Swish可能导致过于复杂的嵌入向量，引发过拟合；而GELU则产生了更优质的向量表示。
• 我们的模型表现优于M3e-FPFT（全参数微调）。M3e-FPFT的时间成本增加了10倍，但召回率和精确率却分别下降了17.2%和28.8%。这说明我们的模型在计算效率上更有优势。

从表2可以看到：

• EM-FT和EM-FT(s)相比基准模型都有显著提升。其中，用标题生成查询的EM-FT（Title）和用随机段落生成查询的EM-FT（RP），在召回率、精确率和mAP上分别大幅提升了12.2%-14.4%、30.5%-42.6%和34%-42.6%。
• 使用LLM辅助查询生成确实有益。我们通过Title和RP两种方法验证了性能改善，确认了方法的普适性。

综合来看，在医学文档检索方面，引入对比学习和用EM-FT进行向量微调，能显著提升模型性能，同时在计算效率上也表现突出。

总结

这项研究提出了REMED医学文档检索框架，涵盖数据构建、嵌入模型微调、检索增强生成和人工评估。EM-FT作为一种高效的微调方法，能端到端优化大规模预训练模型的医学句子表示。我们收集的MMD和MPD两个监督数据集，也为EM-FT模型微调提供了有力支撑。

REMED框架的关键优势在于确保了数据集的隐私和安全性。我们证明了，提升检索模型前K个结果的质量，能直接提高LLM回答的准确性。当然，目前的工作也有局限，比如二元标注系统（0和1）。未来，将标注细化为0到4的评分，可能会带来更精确的结果。这也正是我们下一步的重点方向。

[1] L. Huang, W. Yu, W. Ma, W. Zhong, Z. Feng, H. Wang, Q. Chen, W. Peng, X. Feng, B. Qin et al., “A survey on hallucination in large language models: Principles, taxonomy, challenges, and open questions,” arXiv preprint arXiv:2311.05232, 2023.

[2] P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal,H. Kuttler, M. Lewis, W.-t. Yih, T. Rocktaschel et al., “Retrieval augmented generation for knowledge-intensive nlp tasks,” Advances in Neural Information Processing Systems, vol. 33, pp. 9459–9474, 2020.

[3] Y. Wang, Q. Sun, and S. He, “M3e: Moka massive mixed embedding model,” 2023.

[4] L. Wang, N. Yang, X. Huang, B. Jiao, L. Yang, D. Jiang, R. Majumder, and F. Wei, “Text embeddings by weakly-supervised contrastive pretraining,” arXiv preprint arXiv:2212.03533, 2022.