STEMDiff：基于小波增强扩散模型的物理信息嵌入STEM图像生成及其在原子柱检测中的应用

本文推荐了一种创新的STEM图像生成模型STEMDiff，该模型通过离散小波变换（DWT）架构和物理信息嵌入策略，将晶体结构衍生的二值标签转化为高度逼真的扫描透射电子显微镜（STEM）图像。该研究解决了扩散模型固有的高频偏差问题，在保留实验噪声特征的同时实现了与实验数据近乎无异的图像生成（FID分数0.02，较现有方法提升17倍）。基于合成数据训练的完全卷积网络（FCN）在WSe2和石墨烯等材料的原子柱检测中展现出卓越精度，为材料科学领域提供了无需人工标注的高通量数据解决方案。

　　

1 引言

原子尺度材料理解对量子计算、异相催化、能源存储和水处理等领域的功能材料设计具有基础性意义。扫描透射电子显微镜（STEM）已成为实现亚埃级分辨率原子排列直接可视化的重要工具，然而从复杂图像中提取原子位置、缺陷、应变场和成分变化等定量信息日益需要机器学习（ML）方法处理现代电子显微镜数据集的规模和复杂性。

部署ML进行STEM图像分析的关键瓶颈在于高质量标注训练数据的稀缺性。实验STEM图像的获取耗时且受专用仪器限制，而手动标注会引入人为偏差并严重限制数据集规模。传统依赖物理模拟后手动添加实验伪影的方法不仅劳动密集型，且难以捕捉真实实验图像的统计复杂性，包括探测器噪声、漂移引起的畸变、探针不稳定性和样品污染等伪影。

近期尝试使用生成对抗网络（GANs）桥接模拟与实验差距显示出前景，但GANs存在训练不稳定性、模式崩溃和有限可解释性等问题。扩散模型通过渐进去噪过程学习数据的基础噪声分布，提供更好的训练稳定性和模式覆盖，但常用的U-Net架构在扩散模型中作为去噪模块存在固有高频偏差，即倾向于优先处理低频图像组件而不足够捕捉实验STEM图像中关键的高频细节。

2 方法论

2.1 STEMDiff架构设计与理论框架

STEMDiff是一个基于条件扩散模型的高精度图像生成框架，其核心架构整合三个核心组件：前向扩散过程、反向去噪过程和条件引导机制。该框架采用去噪扩散概率模型（DDPM）作为主干，通过从晶体结构文件衍生的原子坐标投影生成二值标签图，其中原子柱区域标记为1，背景区域标记为0。

前向扩散过程建立概率马尔可夫链，通过N个时间步将原始模拟STEM图像y₀转化为纯噪声。每个转移遵循精确校准的高斯分布，方差系统性增加。反向去噪过程的核心在于从标准高斯噪声y_N重建回连贯STEM图像，由结构二值标签图x引导。每个重建步骤作为条件高斯变换运行，优化目标专注于精确的噪声预测。

2.2 噪声保留块的高频保持机制

研究团队识别出传统扩散架构在科学成像应用中的固有局限性——优先重建低频信息而牺牲关键高频细节。为解决这一根本限制，引入了基于多分辨率小波分析的噪声保留块（NRBs）。这些专用架构组件实施结构化四阶段过程：多分辨率分解将特征表示通过离散小波变换分解为四个互补频率子表示；高频组件隔离选择性地保留包含对角高频信息的x_hh子表示；逆小波变换重建这些关键元素；自适应集成通过修改的跳跃连接将重建的高频元素系统集成到后续处理层中。

这种基于小波的方法改变了频率信息在网络中的流动方式，使STEMDiff能够保存和重建表征真实实验STEM观察的复杂高频细节。通过用NRBs替换整个架构中的传统跳跃连接，在保持计算效率的同时实现了显著改进的保真度。

2.3 原始模拟图像增强协议

STEMDiff的开发需要实施复杂的预处理方法来桥接模拟与实验现实之间的真实感差距。综合增强协议包含以下顺序步骤：像素级引入物理动机的泊松噪声模式，准确复制实验电子检测系统中固有的量子统计波动；应用一系列空间几何操作，包括水平剪切、垂直压缩和抖动；集成校准的高斯噪声分布以模拟真实成像环境中遇到的不可预测变化；使用傅里叶滤波从实验STEM图像中提取并整合低频背景组件，以复制由样品污染引起的局部超亮区域。

3 结果与讨论

3.1 通过STEMDiff生成实验真实STEM图像

STEMDiff展示了在桥接STEM中计算模拟与实验现实之间持续差距方面的卓越能力。以单层二硒化钨（WSe₂）为例，研究团队对原始模拟图像、STEMDiff生成输出和实验STEM图像三个不同图像类别进行了比较分析，检查它们的空间和频域特性。

通过使用abTEM和原子模拟环境（ASE）Python包生成理想化STEM模拟建立基线，匹配实验条件的精确电子光学参数。然而，这些传统模拟图像本质上缺乏关键实验特征，包括探测器噪声、漂移畸变、探针抖动、透镜像差、厚度变化和表面污染现象。STEMDiff方法通过实施物理信息条件扩散过程从根本上重新概念化这一挑战。

3.2 STEMDiff图像真实性的定量评估

为客观量化STEMDiff生成图像的实验真实性，实施了多指标统计评估框架，比较该方法与常规模拟和实际实验采集。主要定量基准采用Fréchet起始距离（FID），这是一种统计稳健的指标，通过利用预训练深度神经网络的特征提取能力在高维特征空间内量化生成与真实图像群体之间的分布相似性。

比较分析评估了三个不同图像类别与实验参考的对比：非增强计算模拟、具有人工噪声配置文件的手动增强模拟和STEMDiff生成图像。结果表明STEMDiff实现了0.02的FID分数，相比Khan等人先前报道的最先进GAN基方法（FID=0.35）提高了17倍，比非增强模拟（FID=196.04）提高了近四个数量级。

3.3 高频保持机制的光谱分析

STEMDiff实现的显著实验保真度源于其旨在克服传统扩散模型固有局限性的基本架构创新。频域详细分析阐明了STEMDiff保存定义实验STEM imagery的关键高频组件的机制。

对扩散模型行为的近期研究揭示了影响科学图像生成质量的内在高频偏差。具体而言，扩散模型表现出系统性地优先重建低频图像组件，同时在去噪过程中逐渐衰减高频细节。这种行为反映了神经网络在训练期间优先学习低频表示的基础倾向，随后损害了复杂纹理元素和精细结构细节的生成。

比较光谱分析显示，传统DDPM在去噪过程关键时间步（n=20和n=1980）表现出显著对比：虽然在早期去噪阶段（n=1980）显示均匀光谱特性，但在终端去噪阶段（n=20）显示中高频组件的显著衰减。STEMDiff架构通过DWT基NRBs直接解决这些限制，从根本上改变去噪过程中的信息流动力学。通过小波分解选择性隔离和保存高频组件，STEMDiff在整个去噪轨迹中保持关键光谱特性。

3.4 STEMDiff在全卷积网络原子柱检测中的应用

STEMDiff实用价值的最终验证在于其支持材料表征中下游分析应用的能力。为证明这一实用价值，实施了全卷积网络（FCN）用于实验STEM imagery中的原子柱检测，仅使用STEMDiff生成的合成数据训练。

原子柱的精确定位代表了材料表征中的基本分析挑战，对理解结构-性能关系具有重大影响。解决这一挑战的传统方法大致分为两个方法家族：使用几何拟合算法的实空间技术和使用模板匹配的倒易空间方法。虽然实空间方法几何直观，但通常需要大量手动参数优化，并且对实验伪影的鲁棒性有限。相反，倒易空间方法在理想化成像条件下表现出优异性能，但在面对束诱导样品修改、污染效应和其他实验变量时出现显著退化。

当代深度学习方法通过直接从数据学习复杂图像-结构映射的能力，为应对这些挑战提供了引人注目的优势。然而，它们的实际实施受到适当标注训练数据集稀缺性的严重限制，因为实验STEM图像的获取和手动标注代表了一个极其资源密集型的过程。

训练在通过物理模拟产生的STEM图像上，STEMDiff框架在推理时支持有效的结构到图像生成方法，有效将晶体学表示与密切类似于真实实验结果的图像联系起来。为证明这一能力，实施了具有架构配置的FCN来识别实验STEM图像中的单层原子位置，包括具有对应分辨率级别之间跳跃连接的对称编码器-解码器结构。该网络仅使用合成数据训练，在本研究中以单层WSe₂和石墨烯结构为例。

4 结论

STEMDiff的开发代表了解决材料表征中数据驱动基础挑战的重要进展——为电子显微镜中的ML应用创建现实、物理信息训练数据。通过将条件扩散模型与离散小波变换基噪声保留块相结合，成功克服了传统生成网络中固有的高频偏差，实现了保存定义真实STEM imagery的关键原子级细节和实验伪影。竞争性低FID分数（0.02）和KL散度指标证明STEMDiff生成图像与实验数据几乎无法区分，性能优于传统方法数个数量级。

更重要的是，该方法的实用效用通过下游原子柱检测任务得到验证，其中仅使用合成STEMDiff数据训练的FCN在包含复杂噪声模式和污染伪影的真实实验图像上实现了显著精度。这种模拟-实验差距的桥接为加速材料发现开辟了新途径，通过实现对电子显微镜数据的快速、自主分析且最小化人工干预。该方法提供了一个适用于多样化材料系统和成像模式的通用框架。STEMDiff因此为物理信息合成数据生成提供了一条新颖路径，可能从根本上改变材料表征和发现中ML模型的训练方式，并扩展至环境科学等领域。