问题 (Problem)

论文旨在解决当前文本嵌入模型领域的核心挑战：如何创建一个既紧凑又通用的文本嵌入模型。

现有方法存在以下痛点：

  • 通用性差：许多模型在特定任务（如语义相似度）上表现优异，但在跨任务、跨领域（如信息检索、分类、聚类等）的泛化能力上表现不佳。
  • 数据依赖严重：要构建一个覆盖多领域、多任务的通用模型，通常需要海量的、高质量的人工标注数据。这个过程不仅成本高昂、耗时费力，而且难以覆盖所有场景。
  • 模型效率问题：为了追求高性能，模型往往变得越来越大（例如参数量超过70亿），嵌入维度也越来越高（例如超过4000维），这给实际部署和应用带来了巨大的计算和存储开销。

因此，本文的核心问题是：我们能在多大程度上直接利用大型语言模型（LLM）中蕴含的丰富世界知识，来蒸馏出一个紧凑、高效且在多种任务上都表现出色的通用文本嵌入模型？

方法 (Method)

本文提出了 Gecko，一个通过两步式LLM知识蒸馏流程训练得到的文本嵌入模型。其核心是创建了一个名为 FRet (Few-shot Prompted Retrieval dataset) 的高质量合成数据集。

Figure1

FRet：两步式LLM蒸馏数据集生成

该方法不直接使用LLM生成文本内容，而是利用LLM的理解和判断能力来生成高质量的训练标签。

第一步：LLM生成多样化的任务和查询 (Diverse Query Generation)

  • 起点：从一个巨大的、无标签的网页语料库中随机抽取一个“种子段落” (pseedp_{seed})。
  • 过程：使用一个经过小样本（few-shot）提示（Prompt）的LLM，让它读取这个种子段落，并为其生成一个相关的任务描述 (tt) 和一个符合该任务的查询 (qq)。
    • 数学表示为：LLM(PQG,pseed)→(t,q)LLM(\mathbb{P}_{QG}, p_{seed}) \rightarrow (t, q)
    • 其中 PQG\mathbb{P}_{QG} 是一个包含了指令和若干示例的固定Prompt。
  • 多样性来源：
    • 网页语料库本身内容丰富，涵盖博客、新闻、百科等多种风格和主题。
    • 通过在Prompt中设计多样化的任务示例（如“问答”、“事实核查”、“寻找相似句子”等），引导LLM生成形式多样的任务和查询。

图解：如下图上半部分所示，LLM读取一个关于“Phastos…创造原子弹”的段落后，可以生成“问答”任务及对应查询“谁制造了原子弹？”，也可以生成“事实核查”任务及对应查询“Phastos创造了原子弹”。

Figure2
图：FRet数据集生成流程图

第二步：LLM挖掘正负样本 (Positive and Negative Mining)

  • 核心假设：种子段落 (pseedp_{seed}) 不一定是其生成的查询 (qq) 的最佳答案。语料库中可能存在更相关、更直接的段落。
  • 过程：
    1. 检索：对于上一步生成的每个查询 (qq)，使用一个预训练好的检索器，从网页语料库中检索出Top-N个最相关的候选段落 P=p(1),...,p(N)P={p^{(1)},...,p^{(N)}}。
    2. 重排序 (Re-ranking)：利用同一个LLM对这N个候选段落进行打分和排序。为了提高排序的鲁棒性，论文融合了两种LLM打分方法：
      • 查询似然度 (Query Likelihood, QL)：计算在给定段落 pp 的条件下，生成查询 qq 的对数似然概率，即 LLM(q∣p,PQL)LLM(q|p, \mathbb{P}_{QL})。
      • 相关性分类 (Relevance Classification, RC)：计算在给定查询 qq 和段落 pp 的条件下，生成“相关”标签的对数似然概率，即 LLM(label∣q,p,PRC)LLM(label|q,p,\mathbb{P}_{RC})。
      • 融合：使用倒数排序融合 (Reciprocal Rank Fusion, RRF) 算法将QL和RC的排序结果结合起来，得到最终的排序函数 R(q,p)R(q,p)。
    3. 重新标注：
      • 正样本 (p+p^+)：选择LLM重排序后排名第一的段落，即 p+=p1p^+ = p_1。实验发现，大约有15%的情况下，p1p_1 并不等于原始的 pseedp_{seed}，证明了重标注的必要性。
      • 难负样本 (p−p^-)：选择LLM重排序后排名最低的段落（例如pNp_N），或者从排名靠后的段落中随机采样。

通过以上两步，最终生成了包含660万个样本的FRet数据集，每个样本都包含(任务, 查询, 正样本, 负样本)四元组。

Gecko 模型训练流程

Gecko的训练分为两个阶段，基于一个12亿参数的预训练Transformer模型。

预微调 (Pre-finetuning)

  • 目的：让模型接触大量不同类型的文本，学习通用的文本表示能力。

  • 数据：大规模的社区问答对（如论坛问答）和网页标题-正文对。

  • 查询向量的生成(q_i): 它的特殊之处在于，它不仅考虑了查询本身，还融入了任务描述（task）。

    qi=mean_pool∣t∣+∣qi∣[M(t⊕qi)∈R(∣t∣+∣qi∣)×d]∈Rdq_i = mean\_ pool_{|t|+|q_i|} \left[ \mathcal{M}(t \oplus q_i) \in \mathbb{R}^{(|t|+|q_i|) \times d} \right] \in \mathbb{R}^d

    • t⊕qit \oplus q_i: 这一步是文本拼接。
      • tt 代表任务描述，例如“问答”或“事实核查”。
      • qiq_i 是用户的实际查询内容，例如“谁发明了灯泡”。
      • ⊕\oplus 符号表示将这两个字符串连接在一起，形成一个新的输入，如：“任务：问答 | 查询：谁发明了灯泡”。这样做是为了让模型知道它当前需要执行哪种类型的任务。
    • M(...)\mathcal{M}(...): M\mathcal{M} 代表一个预训练的语言模型（例如Transformer）。
      • 这个模型会读取拼接后的文本，并为其中的每一个词元（token） 都生成一个上下文相关的向量。
      • 输出结果是一个矩阵，维度是 (词元数量)×d(词元数量) \times d。其中 dd 是预设的向量维度（比如768维），词元数量是 ∣t∣+∣qi∣|t|+|q_i|。
    • $mean_pool[…]: 这是平均池化操作。
      • 它会取模型输出的那个矩阵，然后将所有词元的向量相加再取平均值。
      • 这就好比把一句话里每个词的“含义向量”都融合起来，得到一个代表整句话核心含义的“平均向量”。
    • 最终结果 (qi∈Rdq_i \in R^d): 经过平均池化后，原来代表每个词的向量矩阵就被压缩成了一个单一的、维度为 dd 的向量 qiq_i。这个向量就是最终代表整个“任务+查询”的文本嵌入。

  • 段落向量的生成(p_i): 这个公式是为段落（passage），比如一篇文章或一个回答，生成向量。它的流程和上面类似，但更简单一些。

    pi=mean_pool∣pi∣[M(pi)∈R∣pi∣×d]∈Rdp_i = mean\_ pool_{|p_i|}[\mathcal{M}(p_i) \in \mathbb{R}^{|p_i| \times d}] \in \mathbb{R}^d

  • 目标函数：使用带有批内负样本（in-batch negatives）的对比学习目标函数。对于批次中的一个查询 qiq_i，其对应的段落 pip_i 是正样本，批次内所有其他的段落 pjp_j (j≠i) 都被视为负样本。

    Lpre=−1B∑i=1Blog⁡esim(qi,pi)/τ∑j=1Besim(qi,pj)/τL_{pre} = -\frac{1}{B} \sum_{i=1}^{B} \log \frac{e^{sim(q_i, p_i)/\tau}}{\sum_{j=1}^{B} e^{sim(q_i, p_j)/\tau}}

    其中:

    • sim(x,y)=x⊤y∣∣x∣∣⋅∣∣y∣∣sim(x,y) = \frac{x^{\top}y}{||x|| \cdot ||y||} 是余弦相似度函数
    • τ\tau 是温度超参数。

微调 (Fine-tuning)

  • 目的：利用高质量的FRet和公开数据集，对模型进行多任务、指令化的微调。
  • 数据：一个统一的混合数据集，包含FRet以及多个学术数据集（如Natural Questions, SNLI, MNLI等），所有数据都被统一格式化为(任务, 查询, 正样本, 负样本)。
  • 目标函数：一个更强的批内对比学习损失函数，它不仅包含来自其他样本的正段落作为负样本，还创新性地引入了“同塔负样本 (same-tower negatives)”，即将批次内其他的查询也作为负样本。

    Lmain=1B∑i=1B[−log⁡esim(qi,pi+)/τ∑j=1B(esim(qi,pj+)/τ+1[j≠i]esim(qi,qj)/τ)+esim(qi,pi−)/τ]L_{main} = \frac{1}{B} \sum_{i=1}^{B} \left [ -\log \frac{e^{sim(q_i, p_i^+)/\tau}}{\sum_{j=1}^{B} (e^{sim(q_i, p_j^+)/\tau} + \mathbb{1_{[j \ne i]}} e^{sim(q_i, q_j)/\tau}) + e^{sim(q_i, p_i^-)/\tau}} \right ]

    • pi+p_i^+ 是 qiq_i 的正样本。
    • pi−p_i^- 是 qiq_i 的难负样本。
    • pj+p_j^+ (j≠i) 是批内其他样本的正样本。
    • qjq_j (j≠i) 是同塔负样本，对于对称性任务（如语义相似度）尤其有效。
  • 多维度支持：采用了套娃表示学习 (Matryoshka Representation Learning, MRL) 损失，使同一个模型能够输出不同维度的嵌入向量（如768维和256维），增加了模型的灵活性。

Baseline

Gecko与当前主流的文本嵌入模型进行了对比，可以分为两类：

  • 更大规模的模型 (≥7B 参数, >3k 维度)：
    • gritlm-8x7b, e5-mistral-7b-instruct, gritlm-7b, text-embedding-3-large (OpenAI)
  • 相似规模的模型 (≤5B 参数, ≤1k 维度)：
    • gtr-t5-xxl, gtr-t5-xl, instructor-xl, text-embedding-3-large-256 (OpenAI)

实验结果显示，Gecko-1B-768 (12亿参数，768维) 的平均性能（66.31）不仅全面超越了所有相似规模的模型，甚至能够与比它大7倍、维度高5倍的模型（如gritlm-7b的66.76）相媲美。Gecko-1B-256 (256维) 的性能也远超同维度的其他模型。

数据集 (Datasets)

  • 训练数据：
    • 预微调：社区问答网站数据、网页标题-正文数据。
    • 微调：
      • FRet (本文贡献)：基于网页语料库生成的660万合成数据。
      • 公开学术数据：Natural Questions, HotpotQA, FEVER (检索与问答), SNLI, MNLI (自然语言推断), MedMCQA (医疗问答), MIRACL (多语言检索), 以及多个Huggingface上的分类数据集。
  • 评测数据：
    • MTEB (Massive Text Embedding Benchmark)：一个包含56个数据集的大规模评测基准，覆盖7大类任务：分类、聚类、成对分类、重排序、检索、语义文本相似度(STS)和摘要。
    • MIRACL：一个包含18种语言的多语言检索评测基准。

可复现性 (Reproducibility)

  • 代码：论文中未提及开源代码。
  • 算力/模型：
    • 模型依赖：训练基于一个12亿参数的Transformer模型，微调这样一个模型需要强大的计算资源（如多个高端GPU或TPU）。
    • 数据生成依赖：FRet数据集的生成依赖于一个未指定但性能强大的内部大语言模型（推测为Google自家的模型）。外部研究者难以获取同样性能的LLM。
    • 语料库依赖：生成FRet所用的网页语料库是内部数据，未公开。
  • 结论：由于核心的数据集生成方法依赖于非公开的LLM和数据，且没有提供源代码，该工作的外部可复现性很低。

可改进的几个点

  1. 降低LLM依赖：方法的核心在于强大的LLM。未来可以探索使用更小、更易于获取的开源LLM（如Llama 3, Mistral等）是否能生成类似质量的数据集，从而让该方法更具通用性和可复现性。
  2. 数据生成效率：目前“生成-检索-重排序”的两步流程计算成本非常高。可以研究能否设计出更高效的单步流程，例如让LLM直接生成高质量的(查询, 正样本, 负样本)三元组。
  3. 任务多样性的系统化生成：当前任务多样性依赖于Prompt中的示例。可以探索更系统化的方法，比如先让LLM生成一个包含上百种任务类型的清单，再基于这个清单去生成数据，以确保任务的广度和均衡性。
  4. 更优的难负样本挖掘策略：论文中使用了排序最低的样本或随机采样作为难负样本。可以探索更高级的策略，例如挖掘那些与查询在语义上高度相似但在关键事实上相悖的样本，为模型提供更具挑战性的学习信号。
  5. FRet的多语言扩展：目前的FRet数据集是纯英文的。尽管它对多语言模型也有帮助，但如果能直接用LLM生成多种语言的FRet数据，有望在非英语任务上取得更大的性能提升。

可以被引用的一些结论

  • LLM蒸馏是创建高效嵌入模型的有效途径：通过从LLM中蒸馏知识，一个12亿参数的紧凑模型（Gecko）可以在性能上媲美甚至超越大7倍（>70亿参数）的巨型模型，证明了“小模型+高质量合成数据”路线的巨大潜力。
  • LLM重标注对于合成数据质量至关重要：原始用于生成查询的“种子段落”往往不是最佳正样本。利用LLM对检索到的候选集进行重排序并重新选择正样本，能显著提升模型性能。这一“重选”操作发生在约15%的数据上。
  • 纯合成数据具有强大的零样本泛化能力：仅在合成的FRet数据集上训练的模型，在MTEB这个对它来说是纯零样本（zero-shot）的评测基准上，表现依然非常强劲，甚至超过了许多在人工标注数据上训练的基线模型。
  • 任务多样性是通用嵌入模型的关键：与在单一合成任务（如纯问答）上训练相比，在一个由LLM生成的多任务混合数据集上训练，能显著提升模型在各类下游任务上的泛化能力。
  • 低维嵌入也能达到顶尖性能：有效的知识蒸馏使得低维度（如256维）的Gecko模型能够超越许多高维度（如768维）的强大对手，这对于资源受限的应用场景意义重大。
  • 融合多种LLM打分策略可提升数据标注的鲁棒性：通过RRF融合查询似然度和相关性分类两种LLM打分方式，可以获得一个在不同任务上都表现稳健的重排序器，这对于生成覆盖多种任务的高质量数据集至关重要。