TTE-Flash：通过先思后嵌入令牌加速基于推理的多模态表示

# TTE-Flash：通过思考-然后-嵌入令牌加速基于推理的多模态表示

来源：https://arxiv.org/html/2605.16638
1] Meta AI
吴闲
张江帆
包文涛
Chaitanya Ahuja
Shlok Kumar Mishra
韩超宇
杨高
范霞
齐果
赵丹
樊翔军
肖军
(2026年5月15日)

###### 摘要

近期研究表明，通用多模态嵌入（Universal Multimodal Embedding, UME）受益于链式思维（Chain-of-Thought, CoT）推理。在这种范式中，生成模型为多模态查询生成显式推理轨迹，最终表示则从同时关注查询和推理的嵌入令牌中提取。尽管有效，但生成显式CoT轨迹的计算开销常常令人望而却步。在这项工作中，我们提出用潜在的思考令牌替代显式CoT，这些令牌被解释为潜变量，可以生成作为观测变量的显式CoT轨迹。通过使用CoT生成损失优化思考令牌，并使用对比损失优化后续的嵌入令牌，我们以恒定的推理成本生成了高性能、具备推理能力的表示。我们的研究探讨了两个关键的架构设计：1) 如何从同一个LLM骨干网络中提取思考令牌和嵌入令牌？2) 如何将令牌作为两个相关的任务进行训练？我们引入了TTE-Flash-2B，一个具备推理能力的多模态表示模型，在MMEB-v2基准测试上超越了其显式CoT对应模型，同时生成的潜在思考令牌在文本和视觉上都具有可解释性。此外，在15个视频数据集上的零样本评估揭示了随着思考令牌数量增加而出现的缩放行为，并激发了一项基于任务需求的自适应思考预算分配的初步研究。

## 1 引言

通用多模态嵌入（UME）旨在学习跨文本和视觉模态的指令跟随表示。最近的进展表明，在表示学习之前整合链式思维（CoT）推理可以显著提升表示质量，这被称为"思考-然后-嵌入"（TTE）框架（cui2025think; lan2025ume）。由于多模态内容和指令通常信息密集且抽象，CoT有助于挖掘潜在信息并引导表示学习过程，这反映了其在复杂推理和代码生成任务中的成功（wei2022chain）。然而，尽管有效，TTE模型的实际部署却因实时生成显式CoT轨迹的过高计算成本而受到阻碍。CoT固有的延迟瓶颈并非UME独有；通过潜在推理作为在生成任务中实现恒定推理速度的手段，这一方向已被广泛探索（hao2024training; zhu2025scaling; geiping2025scaling）。在这种范式中，LLM在提示编码和目标解码之间生成固定数量的中间隐藏状态——代表着连续空间中的潜在思考。我们将这种方法引入TTE，并定义了由LLM骨干网络发射的两类连续令牌：**思考令牌**和**嵌入令牌**。具体来说，模型处理多模态输入以在通过嵌入令牌进行表示学习之前生成预定数量的思考令牌，所有操作都在统一的骨干网络内完成。在此框架下，我们研究以下两个研究问题。

**Q1. 如何从统一的LLM骨干网络中提取思考表示和嵌入表示？** 我们寻求一种能够在高并发性和表示保真度之间取得平衡的解决方案。为此，我们评估了两种互补的范式：**循环架构**（zhu2025scaling; hao2024training）和**基于寄存器的机制**（wen2024efficient; darcet2023vision）。循环架构广泛用于潜在推理，它递归地将LLM的连续输出重新注入作为后续输入。虽然这确保了表示深度，但该过程本质上是自回归的且受限于内存。相比之下，常用于视觉Transformer中的**寄存器**（darcet2023vision）——扮演着类似于UME中[CLS]令牌的池化角色，所有输入在单次预填充传递中处理。尽管这种方法受限于计算且效率显著更高，但寄存器缺乏上下文依赖性，可能产生表达力较弱的表示。我们在MMEB上的实验结果表明（meng2025vlm2vec），对于8令牌模型，寄存器的吞吐量是循环方法的两倍，但在检索准确率上相对下降了3%。然而，我们表明，在每个Transformer层引入寄存器可以缩小这一差距，同时保持前向效率。

**Q2. 思考令牌和嵌入令牌应如何训练？** 潜在推理中的一个主要挑战是有效监督中间隐藏状态。先前的研究探索了从显式CoT过渡到隐式CoT的课程学习（deng2024explicit）、从基于显式CoT的教师模型中蒸馏隐藏状态（deng2023implicit），以及利用最终任务损失作为间接监督（yue2025hybrid）。在这项工作中，我们提出了一种更简单而有效的替代训练形式：我们将思考令牌视为生成显式CoT作为观测变量的潜变量。如图2所示，我们训练一个预训练的LLM发射思考令牌，该令牌随后作为解码显式CoT的信息瓶颈。对MMEB的域外评估表明，这种方法使思考令牌能够产生高保真度的CoT轨迹，这些轨迹与多模态指令和真实情况高度一致。在思考之后，后续的嵌入令牌通过使用对比学习的标准多向量检索框架来训练以提取多模态表示（santhanam2022colbertv2; faysse2024colpali）。我们研究了思考和嵌入目标之间的相互作用。初步实验表明，重叠思考令牌和嵌入令牌会导致两个任务的性能下降。这表明，尽管推理和表示学习基于相同的多模态输入，但它们也需要专用的参数专家。受此启发，我们开发了一种用于参数共享和解耦的专门架构，更好地反映了顺序的"思考-然后-嵌入"依赖关系。这项研究的成果是TTE-Flash-2B，一个通过"思考-然后-嵌入"令牌实现恒定时间推理的感知推理UME模型。在MMEB-v2基准测试上，TTE-Flash-2B超越了使用显式CoT的基线推理型UME，同时效率提高了70倍。我们展示了潜在的思考令牌既可以通过解码为CoT进行文本解释，也可以通过附加的图像生成头进行视觉解释。此外，在15个视频数据集上的零样本评估揭示了随着思考令牌数量增加而出现的任务特定缩放行为。这进一步激发了一项基于任务难度的自适应思考预算分配的初步研究。

参见图注 (a) 循环架构自回归地生成下一个隐藏状态用于推理和嵌入。
参见图注 (b) 寄存器是特殊令牌，与提示一起编码为单次预填充传递。
图1：循环方法与基于寄存器的方法在"思考-然后-嵌入"令牌中的比较。

## 2 相关工作

**通用多模态嵌入。** 通用多模态嵌入（UME）利用统一的MLLM骨干网络将多模态输入映射到共享空间（lin2024mm; gu2026unime; gu2025breaking）。与依赖分离编码器和后期对齐的CLIP风格双塔架构不同（radford2021learning; zhai2023sigmoid），UME采用以LLM为中心的融合，以促进生成丰富且具备推理能力的表示（jiang2024e5; zhang2024gme; lin2024mm）。为了增强指令跟随表示，像TTE（cui2025think）、UME-R1（lan2025ume）和MMEmb-R1（wang2026mmemb）等推理增强框架将显式CoT集成到表示学习之前的UME中。与我们工作同期，PLUME（he2026plume）在UME中引入了潜在推理，以在MMEB-v2上实现卓越的准确率-效率权衡（meng2025vlm2vec）。

**潜在推理。** 潜在推理将CoT重新构建为隐藏状态计算。模型不将推理步骤解码为显式令牌，而是生成中间隐藏状态作为潜在思维向量。例如，Coconut（hao2024training）通过递归生成最后的隐藏状态并将其反馈回模型来引入潜在推理。CoLaR（tan2025think）将多个CoT令牌压缩成一个潜在表示。LaDiR（kang2025ladir）应用扩散进行潜在推理。goyal2023think 在生成之前引入特殊的暂停令牌。为了训练潜在推理状态，常见策略包括将显式CoT迁移到潜在推理的课程学习（hao2024training; deng2024explicit）、从显式CoT中蒸馏（deng2023implicit; shen2025codi）、使用端到端损失间接监督潜在状态或使用强化学习（yue2025hybrid）。与这些训练策略不同，我们将思维向量视为在相同预训练LLM骨干网络内生成显式CoT的潜在信息瓶颈。

**联合对比-生成模型。** 双对比-生成目标的整合在用于统一理解和生成的基础模型中很常见。早期工作，如CoCA（yu2022coca），使用对比损失和字幕损失联合预训练图像-文本编码器-解码器。类似地，BLIP（li2022blip）通过结合对比、字幕和图像-文本匹配目标来优化单模态编码器和文本解码器。最近，InternVL（chen2024internvl）展示了在共享LLM骨干网络内同时进行视觉语言对比和生成训练的有效性。虽然我们的工作也在统一的LLM中利用了双重目标，但我们专注于"思考-然后-嵌入"任务，使用专用的令牌头进行推理和表示学习。

## 3 TTE-Flash

在本节中，我们详细介绍TTE-Flash架构，重点介绍思考令牌和嵌入令牌如何与统一的LLM交互、训练和使用。

### 3.1 骨干网络

如图1所示，我们采用统一的LLM骨干网络进行多模态内容编码、思考以及后续的表示学习，将所有内容视为令牌。为了简单起见，LLM注意力是完全因果的：思考步骤关注原始多模态输入编码；嵌入步骤关注多模态输入和思考步骤，与显式TTE设计相同（cui2025think）。在这种范式中，我们考虑了两种表示思考令牌和嵌入令牌的方法，并提取相应的隐藏状态作为表示。

**循环。** 如图1(a)所示，循环架构递归地将最终的提示隐藏状态作为下一个输入反馈给LLM骨干网络。该方法需要N个自回归步骤（其中N是思考令牌和嵌入令牌的总数）来生成所有表示，因此受限于内存。

**寄存器。** 如图1(b)所示，寄存器是可学习的特殊令牌，附加在输入序列的末尾。总共有N个在不同输入间共享的寄存器，该架构使模型能够在单个预填充阶段提取所有思考潜在变量和嵌入。该方法受限于计算。理论上，与循环架构相比，使用寄存器令牌节省了N个解码步骤的延迟，其中N是思考令牌和嵌入令牌的总数。在循环方法中，GPU必须重复加载提示KV缓存N次以顺序生成每个嵌入。在寄存器方法中，GPU计算并加载所有KV恰好一次。这种显著的效率提升对于扩展检索等嵌入应用至关重要。尽管我们的消融研究表明，在4令牌和8令牌设置中，寄存器模型的性能不如循环模型，但**由于其卓越的效率，我们在设计中专注于基于寄存器的方法**。在实验4.1中，我们通过在每个Transformer层添加寄存器来缩小循环和寄存器之间的性能差距，从而增加了分别用于思考和嵌入任务的专家参数总数。

参见图注
图2：思考令牌产生显式CoT生成损失。
参见图注
图3：嵌入令牌产生标准的对比损失。

### 3.2 训练

**思考损失。** 我们定义以下给定多模态输入 \(x\) 的CoT轨迹条件生成过程：
\[
\begin{aligned}
\mathbf{H}_{thi} &= f_\phi(x) \quad (1) \\
\text{CoT}_l &\sim p_\theta(\cdot \mid \text{CoT}_{<l}, \mathbf{H}_{thi}, x) \quad (2)
\end{aligned}
\]
其中 \(\mathbf{H}_{thi}\) 是第 \(i\) 个思考令牌的隐藏状态，\(f_\phi\) 是LLM骨干网络，\(p_\theta\) 是生成CoT的语言模型，\(\text{CoT}_l\) 是第 \(l\) 个生成的CoT令牌。思考损失 \(\mathcal{L}_{\text{think}}\) 是生成显式CoT轨迹的负对数似然：
\[
\mathcal{L}_{\text{think}} = -\sum_l \log p_\theta(\text{CoT}_l \mid \text{CoT}_{<l}, \mathbf{H}_{thi}, x) \quad (3)
\]

**嵌入损失。** 在思考步骤之后，嵌入令牌提取多模态表示。我们使用标准的多向量对比损失（ColBERTv2风格）：
\[
\mathcal{L}_{\text{embed}} = -\log \frac{\exp(s(q, d^+)/\tau)}{\sum_{d \in \mathcal{D}} \exp(s(q, d)/\tau)} \quad (4)
\]
其中 \(s(q, d)\) 是查询 \(q\) 和文档 \(d\) 之间的相似度分数，\(\tau\) 是温度参数，\(d^+\) 是正例文档。

**联合训练。** 总损失是思考损失和嵌入损失的加权和：
\[
\mathcal{L} = \mathcal{L}_{\text{think}} + \lambda \mathcal{L}_{\text{embed}} \quad (5)
\]
其中 \(\lambda\) 是平衡超参数。