LLM-Embedder

1. 目标

训出一个统一的embedding模型LLM-Embedder，旨在全面支持LLM在各种场景中的检索增强

2. 模型的四个关键检索能力

• knowledge：解决knowledge-intensive任务
• memory：解决long-context modeling
• example：解决in-context learning（上下文学习）
• tool：解决tool learning

3. 要解决的问题

• 嵌入模型必须优化其对LLM的最终检索增强影响，而不是仅仅关注中间检索结果
• 不同的检索任务旨在捕捉不同的语义关系，它们的影响可能受到相互干扰

4. base model

是在BAAI/bge-base-en的基础上训练的

5. 训练数据

• Question Answering
  • MSMARCO（labeled）
  • Natural Questions（labeled）
• Conversational Search
  • QReCC（labeled）
• Tool Learning
  • ToolLLM（labeled）
• Instruction Tuning
  • FLAN（Non-labeled）
  • UPRISE（Non-labeled）
• Generation
  • Multi-Session Chat（Non-labeled）
  • Books3（Non-labeled）
  • ArXiv（Non-labeled）
  • CodeParrot（Non-labeled）

6. 数据格式

training

{
  "query": str,
  "pos": List[str],
  "neg": List[str],
  "pos_index": Optional[List[int]],         # 正样本在语料库里面的索引，如果没有全局语料库，则忽略
  "neg_index": Optional[List[int]],         # 负样本在语料库里面的索引，如果没有全局语料库，则忽略 
  "teacher_scores": Optional[List[float]],  # 一个LM或者一个reranker的分数，被用来进行蒸馏
  "answers": Optional[List[str]],           # List of answers for the query, used for LM scoring.
}

evaluation

{
  "query": str,
  "pos_index": Optional[List[int]],         # Indices of the positives w.r.t. corpus (retrieval) / w.r.t. keys (rerank). When there is no positives pre-defined (e.g. NQ), just ignore this field.
  "answers": Optional[List[str]],           # List of answers for computing NQ metrics.
  "key": Optional[List[str]],               # Collated retrieval results for the query / candidates to rank when there are no positives and negatives.
  "key_index": Optional[List[int]],         # Key indices w.r.t. the corpus when reranking and no positives & negatives.
}

7. 训练方法

（主要是Reward from LLM，Instruction-based Fine-Tuning，Homogeneous In-Batch Negative Sampling和Stabilized Distillation四个关键技术）

• Reward from LLM

  • 除了数据集中的hard labels，这里还将LLM的输出作为reward，这些reward可以作为soft labels

  • (1)式中C为一个检索到的候选，也就是一个中间结果，O是期望的LLM的输出，一个候选的reward被表示为$r_{C|O}$。$o_i$表示期望输出的第i个token，$O_{:i-1}$表示期望输出的前i-1个token

  • LLM reward的使用范围：问答（Question Answering），指令调整（Instruction Tuning），生成（Generation）；在会话搜索（conversational search） 和 工具学习数据集（tool learning datasets）上不适用，因为在这些情况下对LLM的输出没有明确的期望

  • 使用对比学习（contrastive learning）的方法去抓取hard labels所反应的语义关系，用知识蒸馏（knowledge distillation）的方法从LLM生成的soft reward中学习

• 对比学习

  • 与一般的对比学习差不多，但为了提高embedding在不同应用场景中的辨别能力在对比学习框架中采用了几个关键设计：
    • Instruction-based Fine-Tuning：给每一个任务设计了一个对应的instruction，计算embedding的时候将instruction和query合在一起计算：$e_q\leftarrow encode([I_t,q])$
    • Homogeneous In-Batch Negative Sampling：使用了cross-device sharing，在这个方法里负样本一共有$B\times K\times N-1$个负样本，B是batch_size，K是GPU的数量，N是所有正样本和hard negative样本的数量。但是不同任务的数据会互相影响导致模型分辨能力下降，所以为了解决这一限制，引入了一种用于组织训练数据的正则化策略，其中来自同一任务的数据实例被分组为连续的小批量，该策略使大多数批内负样本来自同一数据集（即同质负样本），从而增强嵌入对每个特定任务的判别能力

• 知识蒸馏

  • 使用KL散度来最小化使用LLM奖励计算的候选分布与嵌入模型预测的候选分布之间的差距。对于每个query $q$都有一个候选列表$\mathcal{P}:[p_1,\cdots ,p_N]$，然后通过公式(1)获得大模型对候选的奖励，表示为$R:[r_1,\cdots ,r_N]$，为了使得大模型的rewards适合蒸馏，将每个reward转化为标准化的权重：$w_i\leftarrow softma{x}_R(r_i/\alpha )$，其中$\alpha$表示温度系数

  • 简单理解就是rewards越大（$w_i$越大）越好，query和正样本的相似度越大越好（$exp(⟨e_q,e_p⟩)/\tau$越大），query和负样本的相似度越小越好（${\sum }_{p\prime \in \mathcal{P}}exp(⟨e_q,e_p\prime ⟩)/\tau$越小）

  • 由于来自不同任务的不同训练样本，LLM的奖励幅度可能表现出高度波动，导致训练效果不好

  • 为了解决多任务场景下reward的波动问题，提出Stabilized Distillation

    • 基于LLM的rewards$R:[r_1,\cdots ,r_N]$，我们将候选从高到低排序，这样就生成了一个新的候选集，定义为$\mathbb{P}:[p_1,\cdots ,p_N]$，其中$r_i\ge r_{i+1}$。在上面的式子中$\mathbb{P}$由两部分组成：排名比$p_i$低的候选$[p_{i+1},\cdots ,p_N]$以及in-batch negative samples
    • 这个调整过的公式从两方面来稳定波动的rewards：一方面，该模型经过持续的训练以促进$p_i$ 与排名较低的同一个候选集的item$[p_{i+1},\cdots ,p_N]$进行比较， 这意味着，无论rewards的绝对价值如何，模型都能够从LLM的偏好中学习；另一方面，当排名第一的候选比其他候选获得明显更高的reward时，权重将变成one-hot，在这种情况下，提取过程将简化为对比学习的形式，排名第一的候选将被视为正样本。
    • 简单理解就是：一方面我们拿到的负样本中的$[p_{i+1},\cdots ,p_N]$部分都是在LLM给出的reward比正样本低的，这样避免了出现负样本reward比正样本要高的情况，减少了波动；另一方面当排名第一的候选比其他候选获得明显更高的reward时，式子就变成了：
      $min-log\frac{exp(⟨e_q,e_i⟩)/\tau }{{\sum }_{p\prime \in \mathbb{P}}exp(⟨e_q,e_p\prime ⟩)/\tau }$

8. LLMs的检索增强

• Knowledge Enhancement：将知识语料库里面的文档进行编码并存入向量数据库中，在许多情况下对于用户提出的问题可以直接查用于查询向量数据库，在其他情况下生成过程中的上下文可以用作查询，然后在大模型里用[knowledge, query] → answer方式得到答案

• Long-Context Modeling：当处理长上下文时，可以对整个历史进行分块、编码，并将其加载到向量数据库中，生成过程中的上下文可用于查询向量数据库中的相关chunk。在很多情况下查询到的相关chunk比如chunk_i以及其后面的chunk_i+1都会被用来增强记忆，因为后续的chunk对生成下文来说可能更为关键。[retrieved chunks, current context] → new generation.

• In-context Learning（上下文学习）：演示示例以“（(task instruction, expected output）”的形式组织，可以在向量数据库中进行编码和预存。当给出一个新任务时，该任务的指令（instruction）用于查询向量数据库，检索到的示例与任务的指令相连接，在此基础上可以进行上下文学习。[retrieved examples, instruction] → response.

  In-context Learning的示例：

  • Few shot（示例出现多个)）：6+7=13,6+6=12,5+5=10,8+9=？
  • One shot（示例出现一个)）: 5+5=10,8+9=？
  • Zero shot（示例没有出现)）: 8+9=？

• Tool Learning：该工具的功能可以用语言描述，并与其API配对：“（description，API）”，这样向量数据库就可以根据被编码的description来管理一个庞大的工具包，考虑到涉及工具使用的用户请求，可以对该用户请求进行编码并用于查询向量DB。[user request, tool’s execution result] → generation

9. 论文中四个关键技术的工程实现

• Reward from LLM
  • 通过给case选择正负样本实现，因为正负样本都通过计算LLMs的输出logits和label的交叉熵给出了分数，所以根据分数选择更好的正负样本就实现了Reward from LLM
• Instruction-based Fine-Tuning
  • 通过给query和正负样本添加instruction实现
• Homogeneous In-Batch Negative Sampling
  • 记录每种task数据的index范围，然后只在每种task的index范围内进行shuffle，然后在训练的时候将相同task的数据分到一个batch里面
• Stabilized Distillatio
  • 跟据teacher_scores（即LLM给样本的打分）对负样本进行排序，选择分数较低的一部分样本作为负样本

论文链接
项目的github地址