当前位置：首页 > news >正文

TiDAR架构：融合自回归与扩散模型的语言生成新范式

news 2026/5/1 3:38:38

1. TiDAR架构设计解析

在语言模型领域，自回归（AR）模型和扩散模型（dLM）长期处于对立状态。AR模型因其链式条件概率建模能力而享有质量优势，但受限于序列生成特性；扩散模型虽能并行生成，却常因质量妥协而难以实用化。TiDAR通过架构层面的创新，实现了鱼与熊掌的兼得。

1.1 核心设计理念

TiDAR的核心思想可概括为"扩散思考，自回归表达"（Think in Diffusion, Talk in Autoregression）。其突破点在于：

计算密度最大化：现代GPU在解码过程中存在大量空闲计算资源。当模型处于内存带宽限制区时，增加少量token处理几乎不增加延迟。TiDAR利用这一特性，在单次前向传播中同时处理三类token：
1. 前缀token（已确认的上下文）
2. 上一步生成的候选token（需验证）
3. 为下一步预生成的草稿token
双模式统一训练：通过特殊的注意力掩码设计（图3），模型同时学习两种分布：
- 自回归模式：传统链式条件概率 $p_{AR}(x_i|x_{<i})$
- 扩散模式：基于边际分布的并行预测 $p_{Diff}(x_i|\tilde{x})$

关键技术细节：训练时扩散部分全部使用掩码token，这带来三个优势：(1)损失信号更密集；(2)AR与扩散损失项数量一致，便于平衡；(3)支持单步扩散推理，提升效率。

1.2 注意力掩码机制

TiDAR的注意力掩码设计是其成功的关键（图3左）。与Block Diffusion不同，TiDAR仅保留最后一个块为双向注意力，其余前缀部分保持严格因果注意力。这种混合模式带来两大优势：

精确似然计算：支持与传统AR相同的链式条件概率计算，使拒绝采样质量有保障
训练信号增强：前缀部分可计算next-token预测损失，相比纯扩散训练数据利用率更高

实际实现中，掩码模式遵循以下规则：

def get_mask(seq_len, block_size): # 前缀部分：严格因果注意力 prefix_mask = torch.tril(torch.ones(seq_len, seq_len)) # 扩散块：块内双向注意力 block_mask = torch.ones(block_size, block_size) return combine_masks(prefix_mask, block_mask)

1.3 推理过程详解

TiDAR的推理流程（图2）体现了其精妙的设计：

并行草稿生成：通过扩散模式同时预测多个候选token（Thinking阶段）
自回归验证：用AR模式对候选token进行拒绝采样（Talking阶段）
前瞻性预生成：基于所有可能的拒绝结果，并行预生成下一批候选token

整个过程在单次前向传播中完成，KV缓存管理遵循以下原则：

因果注意力计算的token：保留其KV缓存
被拒绝的token：及时清除对应缓存
预生成token：不占用额外缓存空间

实测表明，在H100 GPU上，当总token数（前缀+候选+预生成）≤128时，前向传播延迟几乎不变，验证了"免费token槽"理论的正确性。

2. 关键技术实现

2.1 训练策略优化

TiDAR采用独特的全掩码训练策略，其损失函数设计为：

$$ \mathcal{L}{TiDAR}(\theta) = \frac{1}{1+\alpha}\left(\sum{i=1}^{S-1}\frac{\alpha}{S-1}\mathcal{L}{AR}(x_i,x{i+1};\theta) + \sum_{i=1}^{S-1}\frac{1}{S-1}\mathcal{L}_{Diff}([mask],x_i;\theta)\right) $$

其中$\alpha$控制两种损失的权重（默认取1）。这种设计解决了传统扩散训练的三大痛点：

掩码策略简化：无需复杂的噪声调度
损失平衡：AR和扩散损失的项数相同
训练-推理一致：都采用全掩码方式

实际训练中，我们观察到：

在1.5B模型上，约20B tokens后两种损失收敛
更大的block size（16 vs 4）需要更长训练时间但带来更高吞吐
混合精度训练（BF16）对稳定性无显著影响

2.2 内存效率优化

TiDAR通过以下技术创新实现高效内存利用：

结构化KV缓存：
- 前缀部分：标准AR缓存，可复用
- 候选token：按接受情况动态更新
- 预生成token：不占用持久化缓存
掩码预计算：

max_seq_len = 4096 block_size = 16 # 预初始化全局掩码 global_mask = build_hybrid_mask(max_seq_len, block_size) # 推理时动态切片 def get_slice_mask(prefix_len): return global_mask[:prefix_len+block_size, :prefix_len+block_size]

零重计算设计：
- 被拒绝token的KV缓存直接丢弃
- 接受的token立即用于后续生成
- 无需像传统扩散模型那样多步去噪

2.3 计算加速技巧

在实际部署中，我们总结了以下加速经验：

Flash Attention集成：
- 利用Flash Attention 2的块状注意力计算
- 对混合掩码模式进行特化优化
- 实测速度提升达1.8倍
批处理策略：
- 将多个请求的候选token对齐处理
- 动态padding最小化计算浪费
- 在batch size=8时达到最佳吞吐
硬件感知优化：
- 根据H100的SM单元数量调整并行度
- 使用TMA（Tensor Memory Accelerator）加速大矩阵运算
- 通过CUDA Graph消除内核启动开销

3. 性能评估与分析

3.1 质量对比实验

我们在代码生成（HumanEval、MBPP）和数学推理（GSM8K）任务上进行了系统评估（表2）。关键发现：

1.5B模型：
- TiDAR在HumanEval上达到43.29%，超过基线Qwen2.5 1.5B（35.98%）
- 平均每个前向传播生成7.45个token
- 质量接近AR模型，吞吐接近纯扩散模型
8B模型：
- "信任扩散"模式在数学任务表现更优（GSM8K 80.44% vs AR 79.83%）
- 平均吞吐达8.25 token/NF
- 与EAGLE-3相比，吞吐提升1.7倍

特别值得注意的是，TiDAR在MBPP+上的表现（61.11% @1.5B，80.95% @8B）显著优于同类模型，表明其特别适合需要长程连贯性的任务。

3.2 效率基准测试

图4的吞吐-质量曲线揭示了几个重要结论：

计算密度优势：
- TiDAR 1.5B：4.71倍AR吞吐（5.91 token/s vs 1.25 token/s）
- TiDAR 8B：5.91倍AR吞吐（7.83 token/s vs 1.32 token/s）
与竞品对比：
- 相比Block Diffusion：质量相当情况下吞吐高1.8倍
- 相比EAGLE-3：在GSM8K上快2.1倍
规模扩展性：
- 从1.5B到8B，吞吐增益保持线性
- 大模型更能发挥并行优势