一、 模型组成与核心思想

1. 变分自编码器：负责学习晶体结构的低维、连续、有意义的隐空间表示。其编码器将复杂的晶体结构压缩成潜变量 z，解码器则尝试从 z 重建晶体结构。VAE 的引入确保了隐空间的规整性，便于后续的插值和条件控制。
2. 属性预测器：一个辅助的多任务预测网络，通常在 VAE 的瓶颈处（潜变量 z）进行训练。它预测晶体的宏观属性，如形成能、带隙、体积等。这是实现性质导向生成的关键。
3. 扩散模型：一个在 VAE 的隐空间 中操作的生成模型。它学习一个从简单噪声分布（如高斯噪声）到编码器输出的真实晶体隐变量 z 的去噪扩散过程。这是模型的核心生成引擎，负责从噪声中合成新的、合理的潜变量。

二、 模型架构详解

1. 输入与输出表示

• 晶体图表示：一个晶体被表示为一个多图（通常是有向图）。
  • 节点：代表原子。节点特征包括原子类型（One-hot 向量）、可能的位置等。
  • 边：代表原子间的连接。通过设定一个截断半径，为每个原子创建与其邻居的边。边特征包括原子间距离、可能的角度等。
  • 全局特征：如晶体的空间群、晶胞体积等。
• 输出：模型需要重建或生成完整的晶体结构，即：
  • 晶格矩阵：一个 3x3 的矩阵，定义晶胞的形状和大小（连续变量）。
  • 原子坐标：每个原子在分数坐标或笛卡尔坐标下的位置（连续变量）。
  • 原子类型：每个原子的化学元素（离散变量）。

2. 编码器

• 目标：将晶体图 G 编码为潜变量 z（一个高斯分布的均值和方差）。
• 架构：通常采用图神经网络，特别是消息传递神经网络。
  • 网络对节点、边和全局特征进行多轮信息聚合，最终通过全局池化（如求平均）得到一个全局图表示向量。
  • 这个向量通过两个不同的线性层，分别输出均值 μ 和对数方差 log σ^2。
  • 通过重参数化技巧 采样得到潜变量 z = μ + σ ⊙ ε，其中 ε ~ N(0, I)。

3. 解码器

• 目标：将潜变量 z 解码为晶体结构（晶格 L，原子类型 T，原子坐标 X）。
• 架构：解码器是 CDVAE 中最具创新性的部分之一，通常分阶段进行：
  • 第一阶段 - 生成基础框架：
    • 一个前馈网络从 z 预测晶格矩阵 L（例如，输出9个参数再重组为3x3矩阵）。
    • 另一个网络从 z 预测原子数量 和各元素类型的比例。
  • 第二阶段 - 细化原子位置和类型：
    • 基于预测的原子数量，初始化一组“ placeholder ”节点。
    • 使用另一个 GNN，以潜变量 z、初步预测的晶格信息以及初始化的节点状态为输入，进行多轮消息传递。
    • 最终，GNN 的每个节点输出：
      • 原子类型：一个 softmax 分类分布。
      • 原子坐标：在晶胞内的分数坐标。

4. 属性预测器

• 目标：从潜变量 z 回归或分类预测目标属性 y（如形成能）。
• 架构：通常是一个简单的多层感知机。它作为一个辅助任务，与 VAE 主损失联合训练，促使 z 编码与属性相关的信息。

5. 隐空间扩散模型

• 目标：学习潜变量 z 的数据分布 p(z)。
• 架构：采用标准的去噪扩散概率模型 框架，但在隐空间而非原始数据空间运行。
  • 前向过程（加噪）：在训练时，对编码器输出的真实 z 逐步添加高斯噪声，经过 T 步后变为纯噪声 z_T。这是一个固定的马尔可夫链。
  • 反向过程（去噪）：训练一个去噪网络（通常是时间条件化的 MLP 或 Transformer）来预测给定 t 时刻的噪声 z_t 中所添加的噪声 ε。该网络的训练目标是最小化预测噪声与真实添加噪声的差距。
  • 条件生成：如果需要性质导向生成，可以将目标属性 y 作为条件信息，与时间步 t 一起输入到去噪网络中，即 ε_θ(z_t, t, y)。

三、 训练步骤与完整过程

第一阶段：训练 VAE 和属性预测器

1. 数据准备：准备大规模的晶体结构数据集（如 Materials Project），将其全部转换为统一的晶体图表示。
2. 前向传播：
   • 将晶体图 G 输入编码器，得到 μ, log σ^2，采样得到 z。
   • 将 z 输入解码器，得到重建的晶体 (L‘, T', X')。
   • 将 z 输入属性预测器，得到预测属性 y‘。
3. 损失计算：
   • 重建损失：衡量重建晶体与原始晶体的差异。
     • 晶格损失：均方误差。
     • 坐标损失：考虑周期性边界条件的均方误差（如使用最小镜像距离）。
     • 原子类型损失：交叉熵损失。
   • KL 散度损失：约束潜变量 z 接近标准正态分布 N(0, I)，这是 VAE 的核心。
   • 属性预测损失：预测属性 y‘ 与真实属性 y 的均方误差或绝对误差。
   • 总损失：L_total = L_recon + β * L_KL + λ * L_prop，其中 β 和 λ 是超参数。
4. 反向传播与优化：使用梯度下降（如 Adam）更新编码器、解码器、属性预测器的所有参数。
5. 收敛：重复以上步骤，直到重建损失和属性预测损失稳定。训练完成后，我们就得到了一个能将任何晶体映射到规整隐空间 z，并能从中准确重建的 VAE，同时 z 也蕴含了属性信息。

第二阶段：在隐空间上训练扩散模型

1. 潜变量提取：使用训练好的编码器，为所有训练数据计算其潜变量 {z_1, z_2, ..., z_N}。这些 z 构成了扩散模型要学习的目标分布。
2. 扩散训练：
   • 从数据集中随机抽取一个潜变量 z_0。
   • 随机选择一个时间步 t（1 到 T）。
   • 根据前向过程的公式，计算加噪后的 z_t。
   • 将 z_t、时间步 t（通常通过正弦位置编码嵌入）以及（可选的）条件属性 y 输入去噪网络 ε_θ。
   • 计算去噪网络预测的噪声 ε_θ(z_t, t, y) 与真实添加的噪声 ε 之间的均方误差。
3. 优化：仅更新扩散模型（去噪网络） 的参数，VAE 的参数在此阶段冻结。
4. 收敛：重复以上步骤，直到去噪损失收敛。此时，扩散模型已经学会了如何从噪声 z_T ~ N(0, I) 出发，通过 T 步去噪，生成一个看起来像是来自真实数据分布 p(z) 的潜变量。

四、 训练后模型的使用（推理与生成）

1. 无条件生成

• 步骤：
  1. 扩散采样：从标准正态分布 N(0, I) 采样一个随机噪声 z_T。运行训练好的扩散模型的反向过程（去噪过程），经过 T 步迭代，得到一个生成的潜变量 z_gen。
  2. 解码：将 z_gen 输入到冻结的 VAE 解码器中。
  3. 后处理：解码器输出生成的晶格、原子坐标和原子类型。通常会对坐标进行归一化，并应用对称性操作（如标准化晶胞）以获得规范化的晶体结构。
• 结果：获得一个全新的、在化学上合理的虚拟晶体结构。

2. 性质导向生成（条件生成）

• 步骤：与无条件生成类似，但在扩散模型的每一步去噪中，将目标属性 y_target（例如，“形成能 < 0 eV/atom，带隙 > 2.0 eV”）作为额外条件输入到去噪网络 ε_θ(z_t, t, y_target) 中。
• 结果：生成满足或接近指定目标性质的晶体结构。这是材料逆向设计的核心应用。

3. 晶体重建与插值

• 重建：将已知晶体输入编码器得到 z，再解码回来，可以测试模型的表示能力。
• 插值：在两个已知晶体的潜变量 z_A 和 z_B 之间进行线性插值，然后将中间点的 z 解码，可以观察到晶体结构如何连续地从 A 演变为 B，这对于探索结构-性质关系非常有价值。

4. 性质预测

• 虽然属性预测器是在潜变量上训练的，精度可能不如专门训练的模型，但它可以作为一种快速预筛选工具，对生成的晶体进行初步性质评估。

总结

1. VAE 层：解决了晶体结构（混合离散-连续变量，具有对称性）的高效、规整表示难题。
2. 扩散层：在平滑的隐空间上运作，解决了高质量、多样化样本生成的难题。
3. 条件机制：将宏观属性与隐变量绑定，实现了可控的、目标驱动的材料设计。