One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States

One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States

论文信息: arXiv:2603.08429 | 2026年3月9日 | Temple University
作者: Bo Jiang

1. 问题与动机

当前基于 RAG（检索增强生成）的 LLM Agent 系统普遍采用双模型管线：LLM 负责理解对话上下文并生成搜索 query 文本，然后由一个独立的 embedding 模型将 query 编码为稠密向量用于检索。这种架构存在两个根本性问题：

冗余计算：LLM 在自回归生成 query 的过程中，其隐藏层状态已经完整编码了对话上下文和检索意图的语义理解。将这些理解「丢弃」后再让另一个模型从头编码生成的文本，是一种信息浪费。

工程复杂度与延迟：部署两个大模型（LLM + Embedding Model）增加了基础设施复杂度、显存占用和推理延迟。在对话式检索场景中，embedding 编码步骤的延迟（论文实测 43.5ms/query）虽不算极端，但在实时交互中依然显著。

论文的核心洞察是：LLM 生成 query 时产生的隐藏状态（hidden states）本身就蕴含了足够的检索语义信息，只需学习一个轻量投影头将其映射到 embedding 空间，即可完全消除推理时对独立 embedding 模型的依赖。

2. 方法概述

方法核心是在 LLM 之上附加一个轻量级投影头（Projection Head），将 LLM 自回归解码过程中产生的隐藏状态序列直接映射为与目标 embedding 模型兼容的稠密向量。训练阶段采用知识蒸馏策略，以 embedding 模型作为 teacher，通过三种互补的损失函数训练投影头。推理时，LLM 正常生成 query 文本的同时「免费」产出检索向量，无需额外调用 embedding 模型。

整体流程：

1. LLM 自回归生成 query tokens t₁, t₂, ..., tₙ，同时收集每步的最后一层隐藏状态 h₁, h₂, ..., hₙ
2. 投影头将变长隐藏状态序列 H ∈ ℝⁿˣᵈʰ 映射为定长 embedding 向量 f(H) ∈ ℝᵈ
3. 使用 f(H) 代替 E(q) 进行向量检索

3. 核心技术贡献

贡献一：隐藏状态提取与投影头架构

投影头包含四个组件：

• 输入投影：线性层将 LLM 隐藏维度 dₕ=4096 映射到内部维度 dₘ=1024
• Transformer 编码器：2层、8头自注意力，添加可学习位置编码，允许跨 token 信息聚合
• 均值池化：对非 padding 位置做 mean pooling 得到定长表示
• 输出投影 + L2 归一化：映射到目标 embedding 维度 d=1024 并归一化

整个投影头仅约 25M 参数，相对于 8B 的 LLM 可忽略不计。关键设计选择是使用 Transformer 编码器而非简单线性层——因为 LLM 的隐藏状态是因果注意力（causal）的产物，每个位置只看到之前的 token，而 embedding 模型通常使用双向注意力。Transformer 编码器的自注意力机制提供了「事后双向化」的能力，弥补了因果与双向表示之间的结构性鸿沟。

贡献二：三重损失函数设计

1. 对齐损失（Alignment Loss）：最小化投影向量与 teacher embedding 的余弦距离，提供逐 query 的绝对位置监督

- ℒ_align = 1 - (1/B) Σ cos(f(H⁽ʲ⁾), y⁽ʲ⁾)

1. 对比损失（Contrastive Loss）：InfoNCE 损失，确保每个投影向量在 batch 内与自己的 teacher embedding 最相似，保持判别性结构

- 温度参数 τ = 0.05

1. 排序蒸馏损失（Rank Distillation Loss）：对 teacher 检索的 top-K=128 候选文档，通过 KL 散度将 teacher 的排序偏好迁移给 student

- ℒ_rank = KL(softmax(s/τᵣ) || softmax(ŝ/τᵣ))

组合权重：λₐ = 0.5, λ_c = 0.5, λᵣ = 0.5。论文通过系统消融证明三者各有不可替代的作用，且排序蒸馏不能单独使用（会导致完全崩溃），必须以对齐损失建立的几何基础为前提。

贡献三：系统性消融与训练策略发现

论文提供了12种配置的完整消融实验和统计显著性检验。最重要的发现是训练策略（recipe）对性能的贡献超过任何单一损失组件：80 epoch + 学习率 2×10⁻⁴ 的配置比 30 epoch + 3×10⁻⁴ 提升了 Recall@10 从 0.595 到 0.607，而学习率稍高至 5×10⁻⁴ 则导致完全崩溃。这表明 LLM 因果隐藏空间到双向 embedding 空间的映射涉及微妙的几何关系，需要缓慢而充分的收敛。

4. 实验结果分析

实验设置：

• 数据集：QReCC（对话式检索基准），346 个对话、2,189 个检索触发点
• 模型：Qwen3-8B（LLM）+ Qwen3-Embedding-8B（teacher embedding model），同系列配对
• 硬件：NVIDIA RTX PRO 6000

主要结果（Table 1）：

方法保留了 baseline 97% 的检索质量，同时实现了 21.8 倍的延迟降低（43.5ms → 2.0ms）。McNemar 检验 χ²=12.21, p=0.0005，差异具有统计显著性，但绝对差距较小。

逐 query 分析：在 2,189 个检索触发点中，Win/Tie/Loss = 140/1843/206，84.2% 的 query 两种方法检索结果一致。失败集中在特定对话模式中（48 个高失败率对话），而非均匀分布。

损失消融（Table 2）：

• 仅对齐损失：R@10=0.567（最强单一损失）
• 仅对比损失：R@10=0.498
• 仅排序蒸馏：R@10=0.001（完全崩溃——缺乏绝对位置锚定，梯度无有效信号）
• 对齐+对比：R@10=0.582
• 对齐+排序蒸馏：R@10=0.579
• 三者组合：R@10=0.595
• 三者组合+优化训练策略：R@10=0.607（最佳）

训练策略消融（Table 3）：

• 30 epoch, lr=3×10⁻⁴: R@10=0.595
• 50 epoch, lr=3×10⁻⁴: R@10=0.552（过拟合退化）
• 80 epoch, lr=2×10⁻⁴: R@10=0.607（最佳——慢学习率+长训练）
• 30 epoch, lr=5×10⁻⁴: R@10=0.001（崩溃）

关键洞察：增加 epoch 但不降低学习率会导致过拟合（50 epoch 比 30 epoch 更差），而降低学习率配合延长训练才能充分收敛。

权重调优（Table 4）：过度加权排序蒸馏（λᵣ=1.0）导致 R@10 从 0.591 下降到 0.516。增加候选文档数（K=128→256）效果甚微。

5. 局限与未来方向

1. 单数据集验证：仅在 QReCC 上实验，对开放域检索、多语言场景、非对话式检索的泛化能力未知。需要在 MS MARCO、BEIR 等基准上进一步验证。

1. 同系列模型假设：Qwen3-8B 与 Qwen3-Embedding-8B 共享预训练，隐藏空间可能天然更接近。跨系列（如 LLaMA + BGE）场景下映射难度可能显著增加，论文未探索。

1. 质量差距仍然显著：虽然保留了 97% 的质量，但 p=0.0005 的统计显著性意味着在大规模部署中这个差距是确定性的，而非随机波动。对于高精度要求的场景（如医疗、法律检索），3-3.6% 的下降可能不可接受。

1. 训练仍依赖 embedding 模型：投影头的训练需要 LLM 和 embedding 模型共同生成训练数据（hidden states + teacher embeddings），仅在推理时消除了 embedding 模型依赖。

1. 未探索的方向：

- 是否可以在 LLM 微调阶段端到端训练投影头（而非冻结 LLM）
- 多任务投影（同时支持检索和重排序）
- 投影头对分布外文档的泛化能力
- 更大规模 LLM（70B+）上的效果
- 与 HyDE 等方法的组合

6. 工业实践意义

架构简化与成本降低：对于大规模 RAG 系统，消除独立 embedding 模型意味着：减少一个模型的 GPU 显存占用（8B embedding 模型约需 16GB FP16）、简化服务架构（少一个微服务和 RPC 调用）、降低运维复杂度。25M 参数的投影头几乎不增加额外成本。

延迟优化显著：21.8 倍的延迟降低（43.5ms → 2.0ms）在实时对话系统中极有价值。对于需要多轮检索的 agentic 系统，累积延迟节省更加可观（如 5 次检索从 217.5ms 降至 10ms）。

部署建议：

• 适用场景：对话式检索、多轮 Agent 交互、延迟敏感型 RAG 应用，以及 GPU 资源受限需要减少模型数量的场景
• 慎用场景：高精度检索（法律/医疗）、跨系列模型部署、需要频繁更新 embedding 索引的场景
• 渐进式采用策略：可先在非关键路径上部署，与双模型管线并行运行，通过 A/B 测试验证业务指标影响

同系列模型策略的启示：论文选择 Qwen3-8B + Qwen3-Embedding-8B 的同系列配对并非偶然。对于工业实践，这暗示选择模型家族时应考虑 LLM 和 embedding 模型的「配套性」，共享预训练基础可能是投影成功的重要前提。

对 Embedding 即服务的影响：如果这一方向被广泛采用，独立 embedding 模型的角色将从推理时的必需组件退化为训练时的 teacher 信号提供者，这可能改变 embedding 模型的产品化和商业化逻辑。