Triplet-Block Diffusion RWKV

TL;DR 将 RWKV 的线性推理效率与离散扩散模型结合，通过三元组块布局实现并行双向解码，在保持精度同时将吞吐量平均提升 1.6 倍。

问题

因果 Transformer 自回归解码因严格顺序依赖与二次注意力成本 ((O(L^2)))，在长文本实时生成场景中效率低下，业界正寻求兼具并行生成与线性复杂度的替代架构。

现有方法的局限

• 离散扩散模型：可并行去噪生成多 token，但依赖双向注意力进行完形填空式预测；强行用于因果模型会破坏单向遮蔽，导致训练-推理不一致。
• 线性时间因果模型（如 RWKV、Mamba）：将复杂度降为 (O(L))，但仍保持从左到右的串行解码，无法利用扩散的并行优势。
• 直接组合矛盾：扩散需要同时看到上下文预测被掩码 token，而因果模型仅允许看到历史信息，二者在注意力模式上根本对立。

为何该问题重要且困难

1. 技术挑战：必须在保持因果模型高效推理特性的前提下，引入伪双向机制支持扩散，同时避免额外计算开销和精度损失。核心难点在于设计一种布局方法，让单向模型以最小修改获得全局感受野。
2. 业界关注度：实时交互类应用（如对话系统、代码补全）对首 token 延迟和整体吞吐极敏感，传统自回归模型已成为瓶颈。线性注意力路线与扩散路线若能统一，有望将并行解码与线性复杂度两大优势结合，释放巨大工程价值。

类比：如同 GitHub Copilot 的代码补全，既要即时呈现多行建议（并行），又需在长上下文文件中保持低延迟（线性复杂度），这正是 B³D-RWKV 试图解决的双重约束。

核心洞察

• **Triplet-block 布局实现因果模型中的伪双向访问**：B³D-RWKV 通过将序列划分为块，并利用块间跳跃注意力（triplet-block）在严格单向的 RWKV 中模拟了扩散模型所需的双向条件建模。这一设计解决了因果架构与离散扩散的根本冲突，使得扩散机制首次能在 O(L) 线性复杂度的递归模型上高效运行，避免了传统双向 Transformer 或 MLM 在扩散生成中的高成本。它证明了通过结构化布局而非架构改动，即可为线性注意力模型注入更丰富的全局上下文，为后续高效扩散语言模型的设计提供了工程模板。
• **块级迭代去噪将解码并行度从 token 级提升到块级**：B³D-RWKV 采用在一次去噪迭代中并行生成整个块的所有 token，并通过少量迭代（如 8 步）完成解码，实现了平均 1.6 倍吞吐加速。这种块级并行策略结合了离散扩散的生成效率和 RWKV 的 O(L) 推理特性，避免了块内注意力额外开销。它展示了在不牺牲模型精度的前提下，如何通过改变采样流程显著提升推理吞吐，尤其适合长文本生成、批量推理等对延迟不敏感但吞吐要求高的工程场景，为 LLM 部署提供了新的效率优化方向。

方法

输入表示与 Triplet-Block 布局

B³D-RWKV 的输入是一段文本序列，通过 triplet-block 布局 打破因果模型的单向限制。该布局将序列动态划分为三个连续的子块：前文块（context）、目标块（target）、后文块（suffix）。目标块在训练时被掩码（mask），模型需根据两侧上下文预测被掩码的 token。

这种设计在线性循环模型 RWKV 的基础上实现了 伪双向访问（pseudo-bidirectional access）：模型仍然按因果顺序逐 token 计算隐状态，但通过特殊的块划分和数据流重排，使得对目标块的预测能同时依赖左侧（前文）和右侧（后文）的信息。具体实现中，后文块的信息通过反向传递的隐状态或条件嵌入注入，避免了标准 Transformer 中双向注意力带来的二次复杂度，保持 O(L) 推理开销。

Block-wise 迭代去噪推理

推理阶段采用 离散扩散去噪 流程，以并行方式逐块生成文本。主要步骤包括：

• 掩码采样规则：每个块内统一随机掩码比例，强制掩码 EOS 和 PAD token，以确保序列边界稳定；利用词表复用（vocabulary slot reuse）减少计算浪费。
• 逐步去噪：从完全掩码的目标块开始，模型迭代地预测每个位置的 token 分布，并根据置信度阈值（commit threshold）逐步提交高置信度的 token，直到整块生成完毕。
• 提交机制：每轮迭代中，对于预测概率超过阈值的 token 固定下来，剩余位置继续掩码，直到所有位置完成去噪。

该过程天然支持并行解码，一个块内的所有 token 可在单次前向传播中同时生成，而非自回归式的一个接一个输出，从而大幅提升解码吞吐量。

与同类方法的差异

与 D3PM、MDLM 等基于 Transformer 的双向离散扩散模型不同，B³D-RWKV 首次在线性时间复杂度的因果架构上实现双向扩散，无需引入额外注意力层，从结构上统一了因果模型的高效性与扩散模型的并行生成优势。

实验

实验设计

• 在 8 任务套件 上评估 B³D-RWKV-7.2B 的准确率，覆盖常识推理等典型语言理解任务（论文未列出具体数据集名称）。
• 对比 因果 Transformer 与 离散扩散模型 基线，测试自回归生成与 block 级并行扩散的解码吞吐量。
• 使用 triplet-block 布局 实现伪双向访问，推理采用 block 级迭代去噪采样。

关键发现

• 准确率可比：B³D-RWKV 在 8 项任务上与现有模型达到同等水平，未因结构改变而损失质量。
• 解码速度显著提升：平均 1.6 倍加速，来自并行 block 扩散对顺序生成瓶颈的突破。
• O(L) 推理效率：模型保持线性复杂度，长序列场景下优势进一步放大。

与基线的深度对比

• 传统自回归模型每一步都需等待前文，B³D-RWKV 通过 block 内并行去噪，将部分解码并行化，突破了顺序依赖。
• 相比纯扩散模型，triplet-block 布局利用 RWKV 的线性注意力实现高效双向上下文建模，不再受限于平方级计算。
• 这一设计为 线性时间序列模型 与 扩散生成 的融合提供了工程上可落地的方案，尤其适合对延迟敏感的在线服务场景。

行业影响

落地场景

B³D-RWKV 将扩散模型的并行双向生成能力与 RWKV 的 O(L) 推理效率相结合，1.6× 解码吞吐提升使其天然适合高并发、低延迟的文本生成场景。典型落地包括：

• 实时对话系统：智能客服、语音助手等要求毫秒级首 token 延迟的场景，可借助块级并行去噪大幅减少用户等待时间。
• 批量内容生产：商品描述生成、新闻摘要、报告撰写等需要大量长文本输出的任务，利用扩散的并行解码特性可实现流水线式批量推理，提升系统吞吐。
• 长文档处理：法律合同审查、论文解读等超长上下文任务，线性复杂度 RWKV 背板避免了传统 Transformer 的二次方注意力开销，使长文本推理成本可控。

商业价值

• 降本：减少 GPU 实例占用时间。推理阶段 O(L) 复杂度结合块级扩散，相比同等规模的因果 Transformer，同等硬件下可支撑更高并发，直接降低云服务或私有化部署的算力成本。
• 增收：在广告文案实时生成、电商直播弹幕自动回复等高吞吐场景中，更低的延迟能直接提升互动转化率；更快的推理也为按量计费的 API 服务提供了更大的利润空间。
• 体验提升：并行去噪机制打破了自回归模型逐 token 串行生成的体验瓶颈，在聊天、代码补全等交互式应用中，响应速度的大幅提升可以显著改善用户粘性。

与现有产品/工作流的接口

• 模型替换：B³D-RWKV 可直接替换现有的 Causal Transformer 解码器，尤其适合已采用 RWKV 的生产管线。它使用三元组块布局 (triplet-block layout) 实现伪双向注意力，在微调阶段仅需将训练目标从因果语言建模切换为扩散去噪任务即可。
• 推理框架适配：块级迭代去噪流程要求推理引擎支持多 token 并行提交与掩码调度，这一模式与现有批处理 (batching) 框架高度兼容。可在 vLLM、TGI 等框架中实现自定义采样器，或直接接入 RWKV 的 C++/CUDA 推理运行时。
• 工作流集成：作为纯文本模型，输出可直接对接下游的 LLM 编排工具（如 LangChain、Semantic Kernel），其更高的吞吐意味着在 multi-agent 或检索增强生成 (RAG) 流水线中可作为快速响应的生成节点，降低整条链路的端到端延迟。

具体落地用例

1. 电商平台智能商品描述生成：面对百万级 SKU，平台需要为每个商品生成段落级描述。使用 B³D-RWKV 替代 GPT 风格模型，批量并行扩散生成可将单商品耗时压缩 40%，同时保持可读性，大幅降低每日百万请求的 GPU 集群规模。
2. 全球社交平台多语言内容审核与摘要：平台需对用户贴文进行实时审核并生成摘要。部署 B³D-RWKV 作为后端生成模型，其 O(L) 复杂度和块并行解码能力在分析长文本时能稳定保持在 10 ms 以内延迟，满足高吞吐的 feed 流场景，相比使用 LLaMA 等模型节省超过 60% 的推理预算。

局限

• **普适性仅限单一 backbone 且成本为 3 倍序列长**：论文声称 triplet-block 布局具有结构普适性，但只在一个 backbone（RWKV-7.2B）上做了验证，未在其它线性注意力模型（如 Mamba、RetNet）上测试。同时，该布局会将物理序列长度扩展为原来的 3 倍（因需构造三个块：过去、当前、未来），这在实际长上下文场景中会抵消部分 O(L) 效率优势，且文中未讨论如何缓解该开销。
• **扩散采样步骤与大规模对齐的缺失**：实验仅使用小规模 SFT 数据，无 RLHF 对齐，且未在对话、指令遵循等复杂场景评估。离散扩散模型通常需多步迭代去噪，其推理质量对采样步数敏感，但论文未深入探讨步数、噪声调度与生成质量的 trade-off；此外，未与其它扩散 LM（如 MDLM、SEDD）在相同规模下直接对比，难以判断该方法的相对竞争力。
• **训练与推理效率分析不足**：论文主要强调解码吞吐提升 1.6×，但未报告训练吞吐、显存占用或超参数搜索成本。triplet-block 布局需在训练时构造额外块，可能引入额外的计算和通信开销，但文中并未给出训练效率量化。同时，推理阶段的 commit threshold 和采样步骤选择缺乏消融实验支撑，方法的实用性在大模型实际部署中仍有待检验。