引言

过去几年，我们常听到“大语言模型”（LLM）这个词，比如大家熟悉的 ChatGPT。这些模型能写文章、回答问题、甚至帮人写代码，它们背后的主力就是一种叫 自回归语言模型 的技术。这种方式虽然强大但是缺点在于：

1.生成速度较慢，因为每次生成一个token需要结合历史推理一次

2.生成内容容易前后不一致，因为decoder只参考了单向（前向）序列。

这就引出了一个新方向——扩散语言模型（Diffusion Language Models）。它的灵感来自于扩散模型在图像领域的成功，比如大家熟知的 Stable Diffusion，能一步步把“噪声”变成一张清晰的图。套用到文字里，就是通过逐步去噪”的方式，把一堆随机符号慢慢打磨成一段有意义的文字。而 Dream7B 就是这样一款新型的扩散语言模型。它不再单调地“逐词接龙”，而是尝试一次生成大块内容，再慢慢修正，就像画家不断润色画布上的草图。理论上，这种方法既能加快速度，也有助于保持整体的一致性。

环境搭建

项目网址：DreamLM/Dream: Dream 7B, a large diffusion language model

模型权重：Dream-org/Dream-v0-Instruct-7B · Hugging Face

博主使用NVIDIA L40 48G进行推理实验，据观察推理过程中至少需要20G显存。博主尝试过3080直接炸缸了 卡成PPT...

原理介绍

模型推理的逻辑是：

1. 根据输入可生成的最大token长度填充整个序列。（此时全序列为待生成的<MASK>状态）。
2. 设定timesteps进行调度，此模型中调度策略为线性，即每个时间步生成相同数量的token，所以理论上每步生成的token数应该等于
3. 生成循环每次先检查哪些位置是MASK，这些是待生成位置，然后计算所有待生成位置的logits（不需要保存已生成位置的概率分布以节省空间）。
4. 根据预定义的参数选择生成位置（哪些词应该在这一步被生成）。模型中设置了4种选择模式：
   （1）origin：随机采样，缺点是可能先生成了不确定的位置，整体文本的质量会降低。
   （2）maskgit_plus：按top-1概率采样。在所有候选的位置找到logits概率的最大值并且排序，选出前k个（k为当前时间步需要生成的token数）。优点是优先生成了确定的token，保证了整体序列的稳定性，缺点是在过早的时间就锁定了简单的结果，可能没有充分利用上下文。（本来在生成一段时间，把整段都理解透彻了这个位置可能会输出另一个token,但是在早期就已经被选了。）
   （3）topk-margin：按top-1和top-2的差值进行采样。以所有候选位置logits中top-1和top-2的差值作为置信度排序，选取前k个。优点：相比(2)更加鲁棒，因为考虑了不确定性，如果一个词的第一名和第二名差的特别大那说明基本没悬念了，就选它。缺点：依然可能存在和上述一样的问题，过早揭晓。
   （4）entropy：按分布熵。对所有候选分布做softmax然后计算熵进行排序，选熵最小的k个先生成。优点：度量了全局的不确定性，比(3)中方法更稳定。缺点：计算开销大，对每个候选都做了softmax和sum工作。

一个小细节

模型推理出来的logits应用到实际采样中要整体右移一格，在代码中体现为：

logits = torch.cat([logits_raw[:, :1, :], logits_raw[:, :-1, :]], dim=1)

因为Dream-7B 的骨架其实继承自 Qwen / GPT 系列，这些都是 自回归语言模型。这些模型训练时的目标都是预测当前位置下一位置的token分布，所以位置i的输出应该是位置i+1的分布。所以我们在使用模型推理时需要整体右移一格，将输入和输出位置对齐。

推理过程

以下是我使用的推理代码，并且导出每一步的历史查看序列变化的过程。

import torch
import random
import numpy as np
from transformers import AutoModel, AutoTokenizer
from utils import transtxt



def set_seed(seed):
    """设置所有随机种子以确保可复现性"""
    random.seed(seed)                # Python的random库
    np.random.seed(seed)            # Numpy库
    torch.manual_seed(seed)         # PyTorch CPU
    torch.cuda.manual_seed(seed)    # PyTorch GPU单卡
    torch.cuda.manual_seed_all(seed)# PyTorch GPU多卡
    # 设置固定生成的随机数，这样可以保证每次运行生成的数据都是固定的
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# 设置随机种子
set_seed(42)  # 可以修改为任意整数

model_path = "Dream-org/Dream-v0-Instruct-7B"
tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True,local_files_only=True,)
model = AutoModel.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True,local_files_only=True,).to("cuda").eval()

messages = [
    {"role": "user", "content": "Give me an intersting story:"}
]
inputs = tok.apply_chat_template(messages, return_tensors="pt", return_dict=True, add_generation_prompt=True)
input_ids = inputs.input_ids.to("cuda")
attention_mask = inputs.attention_mask.to("cuda")

# --- 用于收集每步状态/日志 ---
x_init_and_steps = []       # 可选：记录初始 x（t=1）以便和logits对齐
logits_steps = []           # 记录每步 logits（采样前）

def tokens_hook(step, x, logits):
    # step 为 None 时是在进入循环前，x 是“初始画布”（已pad到max_length）
    if step is None:
        x_init_and_steps.append(x.detach().clone().cpu())
    return x

def logits_hook(step, x, logits):
    # ① 全量保存（占显存/内存较大）：
    # logits_steps.append(logits.detach().float().cpu())

    # ② 只保存被 [MASK] 位置的logits（大幅省内存，推荐）：
    mask_id = model.generation_config.mask_token_id
    mask = (x == mask_id)
    mask_logits = logits[mask]  # 形状：[num_mask_positions, vocab]
    logits_steps.append(mask_logits.detach().float().cpu())

    return logits  # 一定要原样返回

out = model.diffusion_generate(
    input_ids,
    attention_mask=attention_mask,
    max_new_tokens=20,
    steps=10,
    temperature=0.1,
    top_p=0.9,
    alg="entropy",
    alg_temp=0.0,
    # 关键开关：
    output_history=True,                 # 返回每步 x 的历史
    return_dict_in_generate=True,
    # 两个 hook：
    generation_tokens_hook_func=tokens_hook,
    generation_logits_hook_func=logits_hook,
)


# --- 读取结果 ---
# 每步更新后的 x（长度 = steps），张量形状约为 [B, max_length]

x_history = out.history  # list[Tensor]
final_x = out.sequences  # 最终 x
#print(final_x)

# 对齐关系：logits_steps[i]（采样前） -> 从 (i) 的 x 过渡到 x_history[i]（采样后）
# 若需要“采样前的初始 x”，用 x_init_and_steps[0]
print(len(x_history), "steps,", len(logits_steps), "logits steps has been collected")

torch.save({
    "x_history": [x.cpu() for x in out.history],   # list of Tensors
    "final_x": out.sequences.cpu(),
    "logits_steps": logits_steps                  # 你在 hook 中收集的logits
}, "dream_debug.pt")

# 示例：解码最终结果
'''
    final_x 是包含prompt 和生成的内容的 token序号列表
    tok.decode将上述token序号翻译成单词

'''
gen = tok.decode(final_x[0][len(input_ids[0]):].tolist()).split(tok.eos_token)[0]
print("new generated tokens_num=", len(final_x[0]) - len(input_ids[0]))
print(gen)

# for i in final_x[0].tolist():
#     print(tok.decode(i))


transtxt("dream_debug.pt",tok,"dream_x_history.txt")

将原代码自带的history转换为可视的txt文件函数transtxt如下:

def transtxt(debug_file_path,tok,out_path):
    # 加载保存的文件
    data = torch.load(debug_file_path, map_location="cpu")
    x_history = data["x_history"]   # list[Tensor]
    final_x = data["final_x"]
    logits=data["logits_steps"]
    with open(out_path, "w", encoding="utf-8") as f:
        for step, x in enumerate(x_history, start=1):
            ids = x[0].tolist()        # 取 batch=0
            text = tok.decode(ids)
            f.write(f"=== Step {step}/{len(x_history)} ===\n")
            f.write(text + "\n\n")

    print(f"所有步骤已写入 {out_path}")

下图为保存的历史文件的一部分，可以清楚地看到两个步骤之间token的生成。

总结

此语言扩散模型本质还是一个AR内核的生成模型，在结构上具有扩散的形状。这其实与文本的编码方法有关，文本作为一种天生离散的载体注定与transformer利用离散token的范式强绑定。我认为此范式最有效的应该是双向注意力，但是这种生成方法天生带来的问题就是得预先定义生成的长度，并且待填位置的前后都已经生成了内容才能发挥双向注意力的优势。有没有可能一段定长的文字如果前后内容都填完了整体质量其实已经定调了呢？只看前向而没有后向约束会不会创造力更好呢？扩散范式是否真有应用到语言模型中的价值还需要更多的实验验证。