随着制造业自动化的进步，对精确且复杂的缺陷检测技术的需求不断增长。现有的用于缺陷识别的视觉模型不足以处理现代制造业环境中缺陷的复杂性和多样性。这些模型尤其在涉及有限或失衡缺陷数据的情况下表现不佳。
在这项工作中，作者引入了MemoryMamba，这是一种新型的内存增强状态空间模型（SSM），旨在克服现有缺陷识别模型的局限性。MemoryMamba将状态空间模型与内存增强机制相结合，使系统能够在训练中保持和检索关键的缺陷特定信息。
其架构旨在捕捉依赖关系和复杂的缺陷特征，这对于有效的缺陷检测至关重要。在实验中，MemoryMamba在四个具有不同缺陷类型和复杂性的工业数据集上进行了评估。该模型一致地优于其他方法，证明了其能够适应各种缺陷识别场景的能力。
1 Introduction
深度学习技术的出现显著推动了各个行业的发展，特别是在制造业中，通过转变质量控制在缺陷识别过程中的作用。在制造业中，缺陷识别在提高生产效率、降低成本和确保产品可靠性方面起着关键作用。随着行业向更自动化、更精确的制造技术发展，对复杂且精确的缺陷识别系统的需求日益增强。这些系统不仅对于减少缺陷产品到达消费者手中的发生率至关重要，而且通过减少浪费和改进安全协议，也有助于提高生产线的整体可持续性。
卷积神经网络（CNNs）的出现显著改变了缺陷识别，提高了准确性和自动化程度[52; 64]。这些模型从适应于既定架构如VGG[1]和ResNet[40]到更复杂的配置如T-CNN[38]和GCNN[54]，显著提高了缺陷识别系统的精度和速度[26]。技术如迁移学习和分类器融合的融合进一步增强了它们在多样化制造环境中的适应性和鲁棒性。然而，这些模型在有限或高度不平衡的缺陷样本场景中经常遇到挑战，即在这些特定工业环境中普遍存在的条件。这些模型依赖于大量标注的数据集以达到高性能是一个重大限制，特别是在这些数据稀缺的环境中。
状态空间模型（SSMs）[21]最近为有效解决这些挑战提供了新的途径。Mamba模型[19]及其在计算机视觉中的变体，如VMamba[35]和VIM[69]，在提高视觉识别方面已显示出巨大潜力。VMamba通过利用跨扫描模块（CSM）降低计算复杂性，该模块在2D图像空间中进行1D选择性扫描，从而无需与视觉 Transformer （ViTs）相关的繁重计算成本即可实现全局感受野。另一方面，VIM使用双向状态空间模型以及位置嵌入来处理图像序列，这有助于捕捉到识别细微且复杂缺陷所需的全面视觉数据。此外，由于Mamba模型的硬件感知设计，其效率和可扩展性使其适合在需要实时处理的工业环境中部署。
为了解决上述问题，作者提出了MemoryMamba，这是一种专门为缺陷识别任务设计的内存增强状态空间模型。MemoryMamba的架构结合了状态空间技术与内存增强，以有效捕捉依赖关系和复杂的缺陷特征。该模型融合了粗粒度和细粒度内存网络，以更好地保留和访问先前训练样本中的关键缺陷信息。此外，作者引入了一个融合模块，以整合从这些内存网络中提取的特征，增强模型的能力。作者还提出了基于对比学习和互信息最大化的优化策略，分别用于粗粒度和细粒度内存网络。
在实验中，作者通过在四个工业缺陷识别数据集上进行全面实验，评估了MemoryMamba的有效性，这些数据集涵盖了各种缺陷类型和复杂性。此外，MemoryMamba在所有现有模型中一致表现出色。
本文的主要贡献如下：
- 在作者了解的范围内，MemoryMamba是第一个用于工业应用中缺陷识别的状态空间模型。
- MemoryMamba模型融合了一种新颖的内存增强机制，允许从历史数据中保留和高效检索关键的缺陷相关信息。
- 作者设计了粗粒度和细粒度内存网络的优化方法，并提出了一个融合模块来整合视觉特征和内存向量。此外，作者提出了基于对比学习和互信息最大化的优化策略，分别用于粗粒度和细粒度内存网络。
- 作者的实验对现有缺陷检测模型进行了全面的比较分析，证明了MemoryMamba的卓越性能。
2 Related Work
Defect Recognition
随着深度学习应用在各个行业的进步，缺陷检测技术在提高制造业产品质量和运营效率方面变得至关重要，尤其是在工业4.0的出现之下。这些技术融合了机器学习和计算机视觉，将传统的缺陷检测方法转变为前所未有的准确性和效率[2; 3; 11]。重要的进展包括由Martin-Ramiro等人开发的张量卷积神经网络（T-CNN）[38]，该网络在不过度牺牲准确性的情况下，减少了参数数量并提高了训练速度。
同样，Shi等人[43]引入了基于中心的转移特征学习与分类器适配（CTFLCA），有效地适应了不同制造环境下的多样化图像分布，并实现了高缺陷检测准确度。Zaghdoudi等人[61]为钢缺陷分类设计了一种分类器融合方法，结合了支持向量机（SVM）和随机森林（RF）以及贝叶斯规则，提高了准确性和速度。
此外，还有尝试将VGG19改进为多路径VGG19，以增强在不同数据集上的缺陷检测能力，如Apostolopoulos和Tzani的工作。相比之下，Fu等人在数据有限的情况下，有效地使用了预训练的VGG16模型与自定义CNN分类器，用于检测钢带表面的缺陷。在用于质量控制的红外成像领域，Rafiei等人利用通过结构剪枝优化的基于ResNet的模型，增强了矿物棉生产中的缺陷检测。Garcia Perez等人开发了使用CNN的自动缺陷识别（ADR）系统，提升了工业X射线分析的可靠性和速度，同时最小化了主观差异。Gamdha等人开发了Sim-ADR，使用合成数据和光线追踪训练CNN进行X射线异常检测，准确度为87%。
此外，Gao等人提出了一种基于多层次信息融合的方法，适用于小样本量，使用具有三个VGG16网络的高斯金字塔提高了识别准确度。后续工作涉及了半监督学习方法，通过伪标签增强CNN，利用有标签和无标签数据改善缺陷检测。
与基于CNN的方法不同，Wang等人 引入了图引导卷积神经网络（GCNN），它将图引导机制融入CNN中，以提高特征提取能力并管理缺陷类别内的变化。后续工作 引入了一种可变形卷积网络（DC-Net），专为混合型缺陷检测设计，在新型结构设计中实现了高准确度。Yu和Lu 提出了一种基于JLNDA的流形学习系统，用于检测晶圆图上的缺陷，使用WM-811K数据优于传统方法。针对特定工业应用的专业CNN适配包括Tao等人针对弹簧线插座以及Mentouri等人针对钢铁制造行业的在线表面缺陷监控。Yang等人开发了一种用于检测风力涡轮机叶片损伤的模型，在具有挑战性的环境中表现出卓越的性能。其他在缺陷检测领域的工作包括Konovalenko等人和Cheng和Yu[9]针对滚轧金属和钢表面缺陷使用先进的神经网络模型。Su等人在faster R-CNN框架内开发了一个互补注意力网络（CAN），以细化太阳能电池图像中的缺陷检测。Zhang等人通过结合ResNet、ECA-Net、动态 Anchor 点以及可变形注意力，增强了DETR，以实现更优越的铸造缺陷检测。
State Space Models
与Transformers 不同，状态空间模型（SSMs）持续发展，越来越多地通过解决各种数据类型中的计算效率和长距离依赖关系管理挑战，塑造序列建模的前沿 。在时间序列预测方面，Xu等人 引入了Mambaformer，这是一种结合了Mamba和Transformer的混合模型，高效地管理长距离和短距离依赖关系，并优于传统模型。Liang等人 提出了Bi-Mamba4TS模型，通过高效的计算和自适应的标记化增强长期预测，准确度上超过了当前方法。在计算机视觉领域，Chen和Ge 提出了MambaUIE，通过有效结合全局和局部特征优化状态空间模型，以用于水下图像增强，大幅降低计算需求同时保持高准确度。
对于图像恢复，Deng和Gu 提出了通道感知U型Mamba（CUMomba）模型，该模型将双状态空间模型整合到U-Net架构中，以有效编码全局上下文并保持通道相关性。在 hyperspectral 图像去噪方面，Liu等人 提出了HSIDMamba，这是一个选择性状态空间模型，它整合了先进的空谱机制，显著提高了效率和性能，比基于Transformer的方法提高了30%。此外，Wang等人 提出了InsectMamba，将状态空间模型、CNN、多头自注意力和MLP在Mix-SSM块中整合，有效地提取精确害虫分类所需的细节特征，在多个数据集上展示了卓越的性能。在3D点云分析中，Han等人 开发了Mamba3D，利用局部归一化池化和双向SSM，在准确性和可扩展性上都显著超越了Transformer模型。
在医学成像方面，有几个模型展示了SSM的应用：Wu等人 提出了H-vmunet，通过高阶2D选择性扫描（H-SS2D）和Local-SS2D模块增强特征提取。Yue和李 提出了MedMamba，利用Conv-SSM模块高效捕捉长距离依赖关系。Ruan和Xiang 开发了VM-UNet，使用视觉状态空间块增强上下文信息捕捉。此外，Wu等人 提出了UltraLight VM-UNet，这是一个高效的模型，使用新型PVM层进行并行特征处理。
在增强SSM的能力方面，He等人 提出了DenseSSM，将浅层隐藏状态整合到更深层中，以提高性能同时保持效率。Smith等人提出了ConvS5，这是一种卷积状态空间模型，在长时空序列建模方面表现出色，训练速度更快，生成样本效率更高。Fathullah等人 开发了MH-SSM，这是一个多头状态空间模型，在LibriSpeech上超越了Transformer转导器，并在整合到Stateformer中时取得了最先进的结果。
预备知识
State Space Models 2024-05-07-03-56-53
状态空间模型（SSM）为时间序列数据分析提供了一个健壮的框架，通过一系列数学表示来封装系统的动态。这些模型将时间序列表述为潜在状态和观测的函数，状态方程描述了这些潜在状态的演变，而观测方程描述了从这些状态得到的测量值。
在时间t的状态向量的演变由状态转移方程控制：
其中表示定义状态向量动态的状态转移矩阵，表示调节过程噪声影响的控制输入矩阵，而假设服从均值为零、协方差矩阵为的高斯分布。
观测模型通过以下方式将观测数据与状态向量相关联：
其中是促进从状态空间到观测数据映射的观测矩阵，而是观测噪声，通常被建模为均值为零、协方差矩阵为的高斯分布。
状态空间模型在捕捉各种系统动态方面的有效性取决于矩阵、、以及噪声过程和的精确表征。这些矩阵可能是静态的，也可能随时间变化，反映了所研究系统动态的变化。从观测估计潜在状态通常采用递归算法，如线性模型的卡尔曼滤波和非线性变体的粒子滤波。这些方法依赖于关于初始状态分布和噪声成分统计性质的假设。这一基本描述强调了状态空间模型在处理跨领域的众多应用中的适应性，它们在建模受随机干扰和观测噪声影响的动态系统中起着关键作用。
4 MemoryMamba
在本节中，作者首先详细阐述MemoryMamba的整体架构。随后，作者介绍作者的Mem-SSM模块，其中包括作者提出的大尺度和小尺度记忆编码以及融合模块。
Overall Architecture
给定一幅尺寸为 的图像 ，MemoryMamba模型首先从图1所示的Patch Embedding过程开始。输入图像 被转换为具有尺寸 的嵌入 Patch 特征 ，即：
其中 表示嵌入的 Patch 特征。在Patch Embedding之后，作者采用Mem-SSM Blocks来迭代地改进特征表示，即：
其中 和 分别是第 个Mem-SSM Block的输入和输出特征集。作者的模型由 个Mem-SSM Blocks组成。每个块操作进一步压缩空间维度并增加通道容量，有效地用空间粒度交换特征深度。最终的输出 是一个封装了图像语义内容的高维表示。
最终的输出 然后馈送到多层感知机（）进行分类，即：
其中 表示 的隐藏向量。然后，作者利用softmax函数计算类别的预测概率分布，即：
其中 表示在 个类别上的预测概率分布。最后，作者的方法可以通过最小化预测概率分布 和真实标签 之间的交叉熵损失来端到端地训练，即：
其中 表示交叉熵损失。
Mem-SSM Block
由于缺陷样本的可用性有限，作者设计了Mem-SSM模块，它集成了记忆网络，从训练样本中学习记忆信息。如图2所示，Mem-SSM模块包括记忆编码和选择性扫描，从而提取输入的健壮表示，这种表示特别敏感于准确检测缺陷所需的细微差别。
该模块对输入特征进行操作，首先通过一个线性层生成一个中间特征集，形式上可以表示为：
其中代表线性变换的可学习参数。
在初始线性变换之后，Mem-SSM模块引入了粗粒度和细粒度记忆编码模块，负责捕捉和编码分层记忆状态。记忆编码可以表示为：
其中和分别表示粗粒度和细粒度记忆状态，表示控制记忆编码过程的参数集。为了将这些编码的记忆状态与中间特征融合，引入了一个融合模块：
其中表示融合特征集，代表融合模块的参数。
融合特征随后通过2D选择性扫描模块处理，该模块选择性地强调重要的特征响应，同时抑制不相关的特征：
其中是选择性扫描的特征集，是该模块的参数。
选择性扫描模块的输出随后通过层归一化步骤进行归一化：
其中和是归一化过程中的缩放和平移参数。
归一化的特征通过残差连接与原始输入特征结合，然后进行第二次线性变换：
其中代表第二个线性层的可学习参数，这完成了Mem-SSM模块内的处理。残差连接有助于保留原始特征信息，同时允许网络自适应地学习特征集的修改。
通过重复应用Mem-SSM模块，MemoryMamba逐步细化视觉表示。
4.2.1 Coarse- & Fine-Grained Memory Encoding
粗粒度与细粒度记忆编码是作者Mem-SSM块的核心，因为它支撑了模型辨别和编码不同 Level 特征细节的能力。这种编码机制对于像缺陷检测这样的细腻任务至关重要，因为细微的区分可以决定正确的分类。
记忆编码过程从对中间特征集应用卷积操作开始，这有助于提取数据中的空间层次结构：
其中是经过卷积处理后的特征集，表示卷积层的参数。
在获得之后，作者将这些特征映射到同时体现粗粒度和细粒度记忆状态的记忆向量。这种映射通过一系列旨在保留特征图内的空间相关性同时降低维度至所需记忆大小的变换来实现：
其中和表示记忆 Query 向量，和分别是粗粒度和细粒度记忆网络的大小。
为了为这些记忆 Query 向量分配相关性，应用了一个softmax层来生成一组注意力权重，从而使模型在检索过程中能够关注最相关的记忆向量：
其中和表示记忆网络的注意力权重。
记忆编码的最后一步涉及通过用注意力权重加权原始记忆向量来聚合记忆知识，形式化如下：
其中和分别表示粗粒度和细粒度 Level 的聚合记忆向量，和分别索引各个记忆向量。随后，记忆向量和被提供给融合模块，在那里它们与融合。
4.2.2 Fusion Module
融合模块旨在将来自粗糙和细粒度记忆编码的信息与中间特征集 进行合并。融合过程首先从记忆向量与中间特征的对齐开始，即：
然后，作者计算记忆向量与特征之间的相似性：
其中 和 分别是粗糙和细粒度记忆向量与中间特征集之间的相似性得分。
相似性得分用于调节记忆特征的贡献。这是通过加权机制实现的，该机制放大相关特征并抑制不那么重要的特征：
其中 和 分别代表加权后的粗糙和细粒度记忆向量。然后，作者将 和 扩展到与 的维度相匹配。这些扩展后的向量随后被添加到 中，以形成增强的特征集 。
Memory Network Optimization
为了提升MemoryMamba架构的性能，作者为粗粒度和细粒度记忆网络采用了专门的优化策略。这些策略旨在通过利用分类损失和基于独特记忆的损失来优化记忆编码过程。
4.3.1 Coarse-Grained Memory Network Optimization
粗粒度记忆网络采用对比学习方法进行优化，该方法利用来自不同类别的 Query 记忆向量 。这种方法鼓励网络更有效地区分各类别的粗略特征，即：
其中 和 分别是第 类和第 类的 Query 记忆向量， 表示余弦相似度，而 是定义类别间最小可接受距离的边缘值。
4.3.2 Fine-Grained Memory Network Optimization
对于细粒度记忆网络，优化的重点是最大化中间特征 与其通过MLP处理后的对应记忆表示 之间的互信息。互信息的最大化确保了记忆网络捕捉到对于细粒度任务至关重要的详细且相关的特征，即：
其中 表示相似性度量（例如，点积），而 表示从记忆中抽取的一组负样本，这些样本不对应于 。
Overall Training Objective
整个训练目标结合了分类损失与记忆特定损失，以有效地训练MemoryMamba模型：
其中 和 分别是平衡对比损失和噪声对比估计损失贡献的权重因子。
5 Experiments
以下是第5节实验部分的开始部分。
Dataset
铝1、GC10 [37]、MT [25] 和 NEU [10] 数据集对于评估缺陷识别模型的性能至关重要。每个数据集都包含不同数量的类别以及在训练和测试数据之间的划分，如表1所示。为了评估模型性能，作者计算以下指标：准确度（ACC）、精确度（Prec）、召回率（Rec）和 F1 分数。使用这些指标进行的综合评估使作者能够彻底评估作者的缺陷识别模型在不同场景和数据集特性下的性能，确保它们是健壮且有效的。### 实验设置
在本研究中，作者使用了 Adam 优化器 [27] 来促进模型的学习过程。学习率设置为 ，训练过程为 10 个周期。权重衰减以 0.01 的速率实施，以正则化和防止神经元的共同适应。此外，作者在训练过程中融入了线性学习率衰减，并包含了一个 Warm up 阶段，占总训练时间的 5%。在这个阶段，学习率从零逐渐增加到设定的初始速率。批量大小设置为 64 以优化模型。训练使用的是 NVIDIA 80 GB A100 GPU。比较方法包括 ResNet [23]、DeiT [48]、Swin-Transformer（Swin [36]）和 Vmamba [35]。
Results
作者的实验结果在四个数据集上表明，即铝、GC10、MT和NEU，MemoryMamba模型的性能卓越。比较结果分别展示在表2、表3、表4和表5中。与传统的架构如ResNet和基于变换的模型如DeiT和Swin Transformers相比，MemoryMamba在准确度、精确度、召回率和F1分数上始终获得最高分。值得注意的是，在具有挑战性的缺陷检测场景中，它在NEU数据集上所有评估指标中取得了高达99%的成绩。将粗粒度和细粒度记忆编码集成显著增强了模型捕捉详细上下文信息的能力，从而提高了在多种条件下复杂视觉模式识别任务的有效性。
Ablation Study
在消融研究中，作者评估了在MemoryMamba架构中的粗粒度记忆网络（CMN）、细粒度记忆网络（FMN）以及融合模块的贡献。移除每个组件都一致导致了准确性和F1分数的下降，突显了它们各自和集体的重要性。当同时排除多个组件时，性能下降最为显著，强调了它们之间的协同效应。这些发现强调了CMN和FMN在捕捉分层特征细节以及融合模块在有效整合这些特征方面的重要作用，这对于模型在工业缺陷检测任务中的性能和鲁棒性至关重要。
Impact on Fusion Module
融合模块对MemoryMamba至关重要，直接影响其分类准确性和F1分数，如图3所示。该模块在将粗粒度和细粒度记忆编码与特征集整合方面的作用，通过在Aluminum和GC10数据集上使用不同的相似性评估方法的性能指标得到了验证。作者的发现强调，余弦相似性在Aluminum和GC10上实现了更高的性能。因此，相似性度量的选择在调整融合模块以达到最佳的缺陷检测性能方面发挥着关键作用。
Impact on Memory Networks
表6：MemoryMamba的消融研究。
图4：研究粗粒度记忆网络在不同记忆大小对铝（左）和GC10（右）数据集的表现。
记忆网络的表现对于缺陷识别系统的鲁棒性至关重要。作者的研究，如图4和图5所示，突显了记忆大小对粗粒度和细粒度记忆网络的准确率（ACC）和F1分数的影响。对于铝数据集，粗粒度记忆网络在记忆大小为4时表现出最佳性能，随着大小的增加，回报逐渐减少。在更复杂的GC10数据集上，两种类型的记忆网络随着记忆大小的变化都显示出显著的性能差异。这些观察表明，最佳的记忆大小依赖于数据集粒度和网络记忆类型的上下文依赖。
Similarity Calculation on Memory Networks
对于记忆网络来说，相似度计算方法的选择可以显著影响模型的性能。为了研究这一点，作者在粗粒度和细粒度的记忆网络上采用了三种不同的相似度度量方法：余弦相似度、L1距离和L2距离。图6和图7展示了这些方法在Aluminum和GC10数据集上的影响。作者的观察结果显示，在准确率（ACC）和F1分数方面，余弦相似度始终优于L1和L2距离。
6 Conclusion
在这项工作中，作者将状态空间模型与记忆增强结合在一起，以解决其他方法在缺陷识别系统中的局限性。作者证明了MemoryMamba的独特架构，它结合了粗粒度记忆网络和细粒度记忆网络以及一个新颖的融合模块，有效地捕捉和利用历史缺陷相关数据。这种能力使得能够增强检测到之前模型可能忽略的复杂和细微缺陷。在优化这些记忆网络时应用对比学习和互信息最大化策略进一步增强了缺陷检测过程的鲁棒性和准确性。来自四个不同工业数据集的实验结果强调了MemoryMamba相对于现有技术如CNN和Vision Transformers的优越性。
参考
[1].MemoryMamba: Memory-Augmented State Space Model for Defect Recognition.