深度拆解AnomalyDiffusion：用扩散模型破解工业缺陷检测的“数据荒”，每一步原理都讲透！

前沿做工业视觉、缺陷检测的朋友大概都有过这样的崩溃时刻老板让你做一个AI质检模型正常产品的图片能堆成山可缺陷样本呢每种缺陷可能就3-5张甚至只有1张。AI模型就像一个没见过世面的学生只看几张缺陷图根本学不会“识别缺陷”——要么把正常产品误判成缺陷要么漏检真正的问题甚至连缺陷在哪都找不到。更头疼的是想人工标注缺陷又费时间又费钱工厂流水线根本耗不起。为了解决这个“数据荒”难题学术界和工业界试了无数方法GAN生成假缺陷、裁剪粘贴拼缺陷、无监督模型硬凑……但要么假得离谱要么可控性差要么下游检测效果拉胯。直到这篇发表于AAAI 2024的论文《AnomalyDiffusion: Controllable Anomaly Generation with Spatial Anomaly Embedding and Adaptive Attention Re-weighting》横空出世才算真正打通了“小样本缺陷生成→下游质检”的全链路。今天这篇深度博客不玩虚的从“工业痛点→模型原理→核心创新→实验细节→落地价值”一步步把AnomalyDiffusion讲透不管你是做算法研发、工业落地还是单纯好奇扩散模型怎么用在工业场景都能看懂、能用得上论文地址https://arxiv.org/pdf/2312.05767代码地址https://github.com/sjtuplayer/anomalydiffusion先明确核心问题工业缺陷检测的“死结”到底在哪在聊模型之前我们先搞清楚为什么工业缺陷检测这么难核心就一个“缺”字——缺缺陷样本、缺缺陷mask、缺可控的生成方法。我们先梳理下工业质检的三大核心任务以及每个任务面临的困境缺陷检测AD判断产品“有没有缺陷”。难点在于正常样本太多缺陷样本太少模型容易“过拟合”到正常样本把轻微的正常波动当成缺陷或者把不明显的缺陷当成正常。缺陷定位AL判断“缺陷在哪”。难点在于不仅需要缺陷样本还需要标注缺陷的位置也就是mask而mask标注比单纯标注“有/无缺陷”更耗时真实场景中mask数据少得可怜。缺陷分类AC判断“是哪种缺陷”。难点在于同一种产品可能有多种缺陷比如胶囊可能有挤压、破损、变形每种缺陷的样本都极少模型无法学习到不同缺陷的特征差异。为了缓解这个问题之前的方法主要分为三类但每类都有致命缺点第一类简单拼接法Crop-Paste把现有缺陷块裁剪下来随便粘贴到正常产品上。优点是简单缺点是假得离谱——缺陷和背景纹理不融合边缘生硬AI一眼就能看穿根本没法用来训练下游模型。第二类GAN生成法SDGAN、Defect-GAN用GAN模型学习缺陷特征生成假缺陷。但GAN的缺点很明显需要大量缺陷样本训练小样本场景下容易“模式崩溃”生成的缺陷千篇一律而且生成的缺陷和mask对不齐定位任务直接废了。第三类改进型生成法DFMGAN目前的SOTA方法用StyleGAN2做迁移学习生成质量比GAN好很多但依然有两个硬伤一是生成的缺陷不够真实细节缺失二是无法精准对齐mask经常出现“缺陷超出mask范围”或“mask内没填满缺陷”的情况。一句话总结之前的方法要么解决不了“小样本”要么解决不了“真实性”要么解决不了“对齐性”。而AnomalyDiffusion一次性解决了这三个问题。左图我们的模型能够生成大量的异常数据这为下游的**异常检测AD、定位AL和分类AC**任务提供了支持而以往的方法由于异常数据有限主要依赖于无监督学习或少量样本监督学习。右图展示了我们模型以及现有异常生成方法在榛子裂纹和胶囊挤压方面的生成异常结果其中我们的结果最为真实准确。。AnomalyDiffusion核心逻辑用扩散模型让“少样本”也能生成“高质量缺陷”AnomalyDiffusion的核心思路很简单但执行细节很精妙基于预训练的Latent Diffusion ModelLDM用“文本反演”学习缺陷特征再通过两大创新机制实现“可控、真实、对齐”的缺陷生成。先科普一个关键前提为什么选LDM而不是GANLDM潜在扩散模型是目前最火的生成模型之一它的优势的是“能利用大规模预训练数据的先验知识”——比如LAION-400M4亿张图像-文本对训练出的LDM已经学会了各种物体的纹理、结构特征。而工业缺陷本质上也是“图像的局部异常”LDM可以借助这种先验只用少量缺陷样本就能快速学习到缺陷的特征。更重要的是LDM支持“文本引导生成”可以通过文本嵌入text embedding控制生成内容——这也是AnomalyDiffusion的核心基础。AnomalyDiffusion的整体流程可以分为4步每一步都针对性解决一个痛点用“masked文本反演”学习“缺陷长啥样”缺陷嵌入用“空间编码器”学习“缺陷在哪”空间嵌入用“自适应注意力重加权”确保缺陷和mask精准对齐用“文本反演生成mask”解决mask数据稀缺问题。接下来我们逐一拆解这4步重点讲透两大核心创新——Spatial Anomaly Embedding空间缺陷嵌入和Adaptive Attention Re-weighting自适应注意力重加权这也是AnomalyDiffusion能吊打其他方法的关键。我们“异常扩散”模型的整体框架1空间异常嵌入 e由异常嵌入eae_{a}ea​一种用于表示异常外观类型的学得文本嵌入和空间嵌入ese_{s}es​从输入的异常掩码 m 中编码得到以指示异常位置组成作为指导异常生成过程的文本条件2)自适应注意力权重调整机制根据去噪后的图像x^0\hat{x}_0x^0​与输入的正常样本yyy之间的差异来计算权重图wmw_{m}wm​并通过权重图wmw_{m}wm​对交叉注意力图mcm_{c}mc​进行自适应调整权重以帮助模型在去噪过程中更专注于那些不太明显的异常区域。创新点1Spatial Anomaly EmbeddingSAE——把“缺陷长啥样”和“缺陷在哪”彻底拆开之前的文本反演方法比如Textual Inversion有一个致命问题会把“缺陷外观”和“缺陷位置”绑在一起。举个例子如果训练样本里的缺陷都在产品的左上角那么文本反演学到的嵌入不仅包含“缺陷是裂纹”的外观信息还包含“缺陷在左上角”的位置信息。这样一来生成缺陷时无论你给的mask在哪个位置缺陷都会生成在左上角——这在工业场景中完全不可用因为缺陷可能出现在产品的任何位置。AnomalyDiffusion提出的SAE就是为了解决这个问题把文本嵌入拆成两部分分别负责“外观”和“位置”实现彻底解耦。3.1.1 缺陷嵌入Anomaly Embedding, eₐ——只学“缺陷长啥样”缺陷嵌入的作用是专门学习缺陷的外观特征比如裂纹的纹理、凹陷的形状、变形的弧度完全不涉及位置信息。为了让它“专注”于缺陷外观作者做了一个关键操作——masked textual inversion掩码文本反演训练时用mask把正常区域和背景都遮住只让模型看到缺陷区域。这样一来模型就不会被正常背景干扰只能学习缺陷的外观特征。具体来说缺陷嵌入eₐ由k个token组成论文中k8通过“掩码扩散损失”masked diffusion loss进行优化损失函数如下Ldif∥m⊙(ϵ−ϵθ(zt,t,{ea,es}))∥22\mathcal{L}_{dif }\left\| m \odot\left(\epsilon-\epsilon_{\theta}\left(z_{t}, t,\left\{e_{a}, e_{s}\right\}\right)\right)\right\| _{2}^{2}Ldif​∥m⊙(ϵ−ϵθ​(zt​,t,{ea​,es​}))∥22​其中m是缺陷maskε是随机噪声zₜ是t时刻的带噪潜在编码εₜ是模型预测的噪声。这个损失的作用是让模型只在mask区域内学习缺陷的噪声预测也就是缺陷特征确保eₐ只包含外观信息。3.1.2 空间嵌入Spatial Embedding, eₛ——只学“缺陷在哪”空间嵌入的作用是专门学习缺陷的位置信息完全不涉及外观特征。它是从输入的缺陷mask中编码得到的确保生成的缺陷能精准出现在mask指定的位置。具体实现很巧妙作者设计了一个空间编码器E流程如下把缺陷mask输入到ResNet-50中提取不同层的图像特征低层特征负责细节位置高层特征负责全局位置用Feature Pyramid NetworksFPN把这些特征融合起来得到完整的位置特征用几个全连接网络把融合后的特征映射到LDM的文本嵌入空间输出包含精准位置信息的空间嵌入eₛ论文中eₛ由4个token组成。这样一来eₐ负责“生成什么样的缺陷”eₛ负责“缺陷生成在哪里”两者彻底解耦——你可以随便改mask的位置缺陷就会生成在对应的位置真正实现“可控生成”。创新点2Adaptive Attention Re-weightingAAR——解决“缺陷和mask对不齐”的痛点即使有了SAE生成的缺陷依然可能出现一个问题填不满mask或者边缘模糊和mask的边界对不齐这也是之前DFMGAN的主要缺陷。比如mask是一个不规则的区域但生成的缺陷只集中在mask的中心边缘区域是空的或者缺陷超出了mask范围跑到了正常区域。这种情况会导致下游定位任务的性能急剧下降——模型学不会“缺陷只在mask区域内”。AnomalyDiffusion提出的AAR机制就是专门解决这个问题的在扩散模型的去噪过程中自适应地给“缺陷不明显的区域”加权重让模型多关注这些区域把缺陷补全、对齐。AAR的具体实现分为两步非常直观3.2.1 计算权重图Weight Map, wₘ在每一步去噪过程中模型会生成一个临时的缺陷图像x̂₀还没完全去噪的版本。然后计算x̂₀和正常样本y在mask区域内的像素级差异——差异越小说明这个区域的缺陷越不明显可能没生成出来或者生成得很淡。然后通过一个自适应缩放softmaxASS操作给这些“差异小的区域”分配更高的权重权重图的计算公式如下wm∥m∥1⋅Softmax(f(∥m⊙y−m⊙x^0∥22))w_{m}\| m\| _{1} \cdot Softmax\left(f\left(\left\| m \odot y-m \odot \hat{x}_{0}\right\| _{2}^{2}\right)\right)wm​∥m∥1​⋅Softmax(f(∥m⊙y−m⊙x^0​∥22​))其中f(x)1/xx≠0f(x)-∞x0。简单来说mask区域内生成的缺陷越不明显权重越高越明显权重越低。3.2.2 注意力重加权Attention Re-weighting扩散模型的生成过程核心是交叉注意力cross-attentionQueryQ来自带噪的潜在编码zₜKeyK和ValueV来自文本嵌入也就是我们的SAEeaeseₐeₛea​es​。交叉注意力图m_c控制着生成的布局和效果——注意力越高生成的效果越强。AAR的核心就是用上面计算的权重图wmwₘwm​对交叉注意力图进行重加权mc′mc⊙wmm_c m_c ⊙ wₘmc′​mc​⊙wm​。这样一来权重高的区域缺陷不明显的区域会获得更多的注意力模型会重点在这些区域生成缺陷从而把mask填满、对齐而权重低的区域缺陷已经很明显的区域模型会减少注意力避免过度生成。论文中的实验证明没有AAR的模型生成的缺陷经常填不满mask加上AAR后缺陷能完美贴合mask的边界对齐精度提升了一个档次如图3所示。图3与具有自适应注意力权重调整我们的模型的模型以及不具有该功能的模型AAR 模型的对比。没有 AAR 功能的模型无法生成足以填充整个掩码的异常值。补充创新自己生成mask彻底解决mask稀缺问题除了缺陷样本少缺陷mask少也是一个大问题——真实场景中标注一个mask需要人工勾勒缺陷边界非常耗时。AnomalyDiffusion给出的解决方案很简单用文本反演学习mask嵌入emeₘem​然后用emeₘem​作为文本条件让LDM生成大量多样的缺陷mask。具体来说mask嵌入emeₘem​由k′kk′个随机token初始化论文中k′4k4k′4通过扩散损失优化损失函数如下em∗argminemEz∼E(m),ϵ,t[∥ϵ−ϵθ(zt,t,em)∥22]e_{m}^{*}\underset{e_{m}}{arg min } \mathbb{E}_{z \sim \mathcal{E}(m), \epsilon, t}\left[\left\| \epsilon-\epsilon_{\theta}\left(z_{t}, t, e_{m}\right)\right\| _{2}^{2}\right]em∗​em​argmin​Ez∼E(m),ϵ,t​[∥ϵ−ϵθ​(zt​,t,em​)∥22​]训练好emeₘem​后就可以用它引导LDM生成各种形状、大小的缺陷mask——比如裂纹mask、凹陷mask、变形mask彻底解决mask稀缺的问题让生成的缺陷更具多样性。实验细节拉满AnomalyDiffusion到底有多能打作者在工业缺陷检测的标准数据集MVTec AD上做了大量实验对比了目前所有主流的缺陷生成方法DiffAug、CDC、Crop-Paste、SDGAN、Defect-GAN、DFMGAN等从“生成质量”“下游任务性能”“消融实验”三个维度全面验证了模型的有效性。先说明实验设置避免大家看数据时一脸懵数据集MVTec AD包含15类工业产品瓶子、电缆、胶囊、地毯等每类产品有多种缺陷类型训练设置每类缺陷只取1/3的样本作为训练集极少样本生成1000张异常图像- mask对用于下游任务训练评价指标生成质量ISInception Score越高越好衡量生成质量和多样性、IC-LPIPS Intra-cluster LPIPS越高越好衡量生成多样性下游任务AUROC越高越好检测/定位精度、AP越高越好定位精度、F1-max越高越好综合精度、分类准确率越高越好。表 1在 MVTec 数据集上的 IS 和 IC-LPIPS 比较。我们的模型生成了质量最高且种类最丰富的异常数据实现了最佳的 IS 和 IC-LPIPS 值。粗体和斜体分别代表最优和次优结果。4.1 生成质量全面碾压所有方法论文中的表1核心对比表显示AnomalyDiffusion在IS和IC-LPIPS两个指标上均取得了最好的成绩IS平均值1.80第二名DFMGAN是1.72IC-LPIPS平均值0.32第二名DFMGAN是0.20。这意味着AnomalyDiffusion生成的缺陷不仅更真实IS高而且多样性更强IC-LPIPS高——不会出现“生成的裂纹都长一个样”的情况。从定性结果图4、图9来看差距更明显Crop-Paste缺陷和背景割裂一看就是拼接的CDC生成的缺陷很怪异比如胶囊挤压缺陷形状完全不符合实际SDGAN/Defect-GAN经常生成不出完整的缺陷比如药片裂纹只生成了一小段DFMGAN虽然比前几种好但缺陷和mask对不齐比如金属螺母的弯曲缺陷超出了mask范围AnomalyDiffusion缺陷纹理自然、边缘顺滑完美贴合mask和真实缺陷几乎一模一样。图 4在 MVTec 数据集上的生成结果对比。我们的模型生成的异常图像质量高并且与异常掩码完全对齐。图 9异常生成质量的定性比较。请注意由 DiffAug 生成的异常与训练样本位于右下角的图像完全相同。4.2 下游任务简单U-Net也能吊打SOTA生成缺陷的最终目的是为了提升下游质检任务的性能。作者用生成的异常数据训练了一个简单的U-Net不是什么复杂的SOTA模型然后对比了它和其他方法的性能缺陷定位像素级AUROC 99.1%、AP 81.4%、F1-max 76.3%全面碾压DRAEM、PRN、DFMGAN等方法表2缺陷检测图像级AUROC 99.2%、AP 99.7%、F1-max 98.7%同样排名第一表3缺陷分类平均准确率66.09%比第二名DFMGAN49.61%高出16.48个百分点断层领先表4。更夸张的是这个简单的U-Net甚至吊打了很多专门做缺陷检测的SOTA模型比如CFLOW、PatchCore、DRAEM——表5显示U-Net的AP81.4%比这些模型高出10-20个百分点。这充分说明AnomalyDiffusion生成的数据质量足够高能够有效弥补真实缺陷样本的不足让下游模型的性能实现质的飞跃。表 3图像级异常检测的比较分析。表 4基于 ResNet-18 模型对生成数据进行训练所得的异常分类准确率的比较。表 5基于我们生成的数据集训练的简单 U-Net 与现有异常检测方法包括其官方代码或预训练模型在像素级异常定位方面的比较AUROC/AP。4.3 消融实验缺一个模块都不行作者做了详细的消融实验验证了每个核心模块的必要性表6无SAE无masked loss无AARAUROC 81.3%、AP 31.1%性能惨不忍睹只有SAEAUROC 90.3%、AP 60.7%可控性提升但对齐性差SAEmasked lossAUROC 95.0%、AP 64.9%生成质量提升全模块SAEmasked lossAARAUROC 99.1%、AP 81.4%性能达到最优。除此之外作者还做了其他消融实验缺陷样本比例当训练样本比例低于30%时AP会急剧下降高于30%时性能趋于稳定说明模型在小样本场景下依然有效表7嵌入长度缺陷嵌入和空间嵌入的长度token数量对性能影响不大说明模型对超参数不敏感容易落地表8和Prompt-to-Prompt对比Prompt-to-Prompt无法生成真实的缺陷对齐性也差而AnomalyDiffusion的性能全面领先表9、图8。表 6我们所提出的空间异常嵌入SAE、掩码扩散损失MaskedL\mathit{L}L以及自适应注意力重新加权机制AAR的消融研究。表 7异常率的消融研究。表 8异常嵌入lal_{a}la​和空间嵌入lsl_{s}ls​的消融研究图 8在异常生成方面文本反转 提示到提示模型与我们的模型空间异常嵌入模型的比较结果。落地价值AnomalyDiffusion到底能解决什么实际问题聊完技术和实验我们回到工业落地场景——AnomalyDiffusion不是一篇“纸上谈兵”的论文它的每一个设计都针对工业场景的实际痛点降低数据收集成本只需要3-5张缺陷样本就能生成大量高质量的异常数据不用再花大量时间、金钱标注缺陷和mask提升质检精度生成的数据能让下游模型哪怕是简单的U-Net的检测、定位、分类精度大幅提升减少漏检、误检可控性强可以自由控制缺陷的类型和位置适配不同的产品和缺陷场景比如瓶子的破损、电缆的断裂易落地模型基于预训练LDM训练速度快300K迭代3天就能训完对硬件要求不高NVIDIA RTX 3090即可而且对超参数不敏感容易调试。对于做工业视觉的朋友来说这意味着以前因为“缺数据”做不了的质检AI现在用AnomalyDiffusion就能轻松实现以前需要复杂模型才能达到的精度现在用一个简单的U-Net就能实现。总结AnomalyDiffusion的核心价值的是什么AnomalyDiffusion的本质是用“扩散模型的预训练优势”“两大创新机制”破解了工业缺陷检测的“数据荒”难题。它的核心价值不在于“提出了一个新的扩散模型变体”而在于“把扩散模型的能力精准落地到工业场景解决了实际问题”。它的成功也给工业视觉领域带来了一个新的思路与其在“无监督检测”上硬凑不如用“小样本生成”来弥补数据不足——毕竟AI学得好不好关键还是看数据够不够、够不够好。