Score-Entropy-Discrete-Diffusion模型架构揭秘：从small.yaml配置到Transformer实现

Score-Entropy-Discrete-Diffusion模型架构揭秘：从small.yaml配置到Transformer实现

【免费下载链接】Score-Entropy-Discrete-Diffusion[ICML 2024 Best Paper] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution (https://arxiv.org/abs/2310.16834)项目地址: https://gitcode.com/gh_mirrors/sc/Score-Entropy-Discrete-Diffusion

Score-Entropy-Discrete-Diffusion（SEDD）是ICML 2024最佳论文提出的离散扩散模型，通过估计数据分布比例实现高效生成。本文将从配置文件解析到Transformer核心实现，全面揭秘这一创新模型的架构设计。

核心配置解析：small.yaml与medium.yaml对比

SEDD模型提供了两种预配置方案，通过configs/model/目录下的YAML文件定义网络结构参数：

small.yaml轻量级配置

name: small type: ddit hidden_size: 768 # 隐藏层维度 cond_dim: 128 # 条件编码维度 length: 1024 # 序列长度 n_blocks: 12 # Transformer块数量 n_heads: 12 # 注意力头数 scale_by_sigma: True # 是否使用sigma缩放 dropout: 0.1 # Dropout比率

medium.yaml增强配置

name: medium type: ddit hidden_size: 1024 # 提升至1024维隐藏层 cond_dim: 128 # 保持条件维度一致 length: 1024 n_blocks: 24 # 增加至24个Transformer块 n_heads: 16 # 提升至16个注意力头 scale_by_sigma: True dropout: 0.1

两种配置的核心差异体现在模型容量上，medium版本通过增加隐藏层维度、Transformer块数和注意力头数，提供更强的特征学习能力，适合复杂数据生成任务。

Transformer核心实现：从DDiTBlock到SEDD类

1. 底层构建模块

模型的核心构建块DDiTBlock定义在model/transformer.py中，每个块包含：

• 自适应层归一化：通过adaLN_modulation实现条件信号对归一化参数的调制
• Flash注意力机制：使用flash_attn_varlen_qkvpacked_func实现高效注意力计算
• 门控残差连接：结合bias_dropout_scale_fn实现稳定训练

2. 时间步与条件嵌入

• TimestepEmbedder：将扩散时间步σ转换为高维嵌入，使用正弦位置编码+MLP架构
• LabelEmbedder：处理类别条件信息，支持无分类器引导生成

3. 完整模型架构

SEDD主类实现了从输入序列到输出分布的完整流程：

1. 词汇嵌入：通过EmbeddingLayer将离散 tokens 转换为向量表示
2. 条件编码：融合时间步σ和类别条件信息
3. Transformer堆叠：通过DDiTBlock列表构建深度网络
4. 输出层：使用DDitFinalLayer生成最终分布

关键创新点在于引入分数熵估计和自适应调制机制，使模型能直接估计数据分布的比例，避免传统扩散模型的采样效率问题。

模型训练与推理流程

训练配置

主训练入口位于run_train.py，通过加载configs/config.yaml中的全局参数控制训练过程，包括：

• 优化器设置（学习率、权重衰减）
• 数据加载参数（批次大小、序列长度）
• 扩散过程超参数（噪声调度、采样步数）

推理采样

提供两种采样脚本：

• run_sample.py：无条件生成
• run_sample_cond.py：条件生成，支持类别引导

总结：SEDD模型的技术优势

1. 架构效率：采用Flash注意力和融合操作，显著提升训练速度
2. 配置灵活性：通过YAML文件轻松调整模型容量，适应不同任务需求
3. 生成质量：分数熵估计方法改善了离散数据的生成多样性和保真度

SEDD模型通过创新的架构设计和训练策略，为离散扩散模型领域提供了新的技术范式，其代码实现既保持了学术创新性，又兼顾了工程实用性。

【免费下载链接】Score-Entropy-Discrete-Diffusion[ICML 2024 Best Paper] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution (https://arxiv.org/abs/2310.16834)项目地址: https://gitcode.com/gh_mirrors/sc/Score-Entropy-Discrete-Diffusion