X-CoSD: Communication-Efficient Cross-Vocabulary Collaborative Speculative Decoding

问题陈述

协作式投机采样（Collaborative Speculative Decoding, CoSD）是一种分布式推理框架，其中设备端的轻量级语言模型（SLM）起草候选 Token，而服务器端的语言大模型（LLM）对其进行验证。虽然 CoSD 加速了生成过程并减轻了计算负担，但现有方法依赖于一个关键假设：SLM 和 LLM 必须共享相同的词表。

在实际的异构场景中（例如不同的模型家族或厂商），词表往往是不同的。将 CoSD 直接扩展到这些场景需要一个直观的操作：每当 Token 被拒绝时，都需要传输整个服务器 LLM 在全量词表上的概率分布，以便设备端进行精确重采样。这造成了极高的通信瓶al 瓶颈，尤其是在无线网络中，下行链路往往是端到端延迟的主要来源。此外，现有的通信高效型 CoSD 方法（如 U-HLM）通常会牺牲解码质量（有损），或者仍然需要沉重的上行分布传输。

方法论

作者提出了 X-CoSD，这是一个专为异构 SLM–LLM 词表设计的框架，旨在保证无损解码（保留服务器 LLM 的分布）的同时最小化通信开销。该框架依赖于两个核心机制：

1. 混合重采样 (Hybrid Resampling, HR)

X-CoSD 基于 Token 级交集（Token-Level Intersection, TLI）方法构建，其中设备端 SLM 仅从两个模型共有的公共词表（$V_C$）中起草 Token。当 Token 被拒绝时，残余分布（重采样所需的概率质量）被分解为两个区域：

• 公共词表区域 ($V_C$)：两个模型都存在的 Token。
• 仅 LLM 区域 ($V_O$)：仅存在于服务器 LLM 中的 Token。

HR 不再发送完整的分布，而是利用残余分布的结构来进行分布式重采样：

1. 服务器计算公共区域（$V_C$）的总残余概率质量 ($\theta_c$) 和仅 LLM 区域 ($\theta_o$) 的概率质量 ($\theta_o$)。
2. 服务器仅向设备端传输公共区域（$V_C$）的 LLM 概率和标量 $\theta_o$。
3. 设备端以概率方式选择重采样区域：
   • 如果选择了 $V_C$，设备端在公共区域内进行本地重采样。
   • 如果选择了 $V_O$，设备端请求服务器从仅 LLM 区域进行采样。随后，服务器仅发送对应采样 Token 的反序列化字符串（而非完整的分布）。

这种方法确保了完整的 LLM 分布永远不会被传输，从而显著降低了下行负载。

2. 带设备验证的服务器重采样 (Server Resampling with Device Verification, SR-DV)

为了进一步降低开销，作者提出了增强版 X-CoSD-E，它使用了 SR-DV 技术。

• 当发生拒绝时，服务器从完整的 LLM 分布中采样少量的替换候选者（$K$ 个），并仅将这些 Token 及其概率发送给设备端。
• 设备端使用论文中推导出的特定接受概率函数 ($\alpha(x)$) 来验证这些候选者。
• 如果候选者被接受，则流程结束。如果所有 $K$ 个候选者都被拒绝（或经过 $M$ 次迭代后），系统将回退到 HR 机制以确保精确的残余重采样。
• 这种概率性方法避免了在大多数情况下传输公共区域的分布，进一步降低了延迟。

主要贡献

1. 跨词表支持：X-CoSD 是第一个专门为异构词表设计的 CoSD 框架，无需进行模型重训或在线知识蒸馏。
2. 无损保证：论文通过理论证明（定理 1 和定理 2）表明，X-CoSD 和 X-CoSD-E 都能完整保留服务器 LLM 的分布，确保生成质量与服务器 LLM 一致。
3. 通信效率：通过将分布传输限制在公共词表区域（HR）或用候选 Token 代替（SR-DV），该框架比直接扩展的方法大幅降低了下行通信负载。
4. SR-DV 机制：引入了优先进行低延迟候选者交换而非传输完整分布的回退机制。

实验结果

作者使用 Vicuna-68M (SLM) 配合 Llama-3.1-8B 和 Qwen2-7B (LLMs) 在包括机器翻译 (WMT-DeEn)、摘要生成 (XSum, CNN/DailyMail)、数学推理 (GSM8K) 和问答 (MMLU) 在内的任务中对 X-CoSD 和 X-CoSD-E 进行了评估。

• 生成质量：X-CoSD 和 X-CoSD-E 的生成质量得分均与服务器 LLM 相当（相对得分 $\approx$ 1.0），证实了其无损特性。相比之下，有损基线（Greedy/Target Resampling）和 U-HLM 显示出明显的质量下降。
• 通信负载：
  • 下行链路：与发送完整分布的“直接扩展”（Straightforward）方法相比，X-CoSD 将下行负载降低了数个数量级。X-CoSD-E 通过在许多拒绝案例中避免传输分布，进一步降低了该负载。
  • 上行链路：与某些需要上传 SLM 分布的基线不同，X-CoSD 保持了较低的上行负载，仅传输 Token ID 和概率。
• 吞吐量与延迟：
  • X-CoSD 和 X-CoSD-E 实现了比单纯服务器 LLM 更高的 Token 吞吐量，并且优于“直接扩展”基线，尤其是在下行带宽增加的情况下。
  • X-CoSD-E 在所有无损方法中表现出最低的单 Token 延迟，即使在不同的下行传输速率（100–300 Mbps）下也能保持高吞吐量。

意义

论文声称，X-CoSD 和 X-CoSD-E 代表了迈向现实世界协作式 LLM 推理部署的务实步骤。通过在不牺牲解码准确性的情况下解决词表不匹配问题，这些框架实现了异构设备端模型与强大服务器 LLM 之间高效、无损的协作。这对于以通信带宽为主要瓶颈的无线环境尤为重要，能够在不损害生成文本质量的前提下实现更快的推理。