Starlight-apk blog

1.1 研究背景#

工业异常检测是智能制造质量控制的关键环节。近年来，随着深度学习技术的发展，基于卷积神经网络和视觉Transformer的方法在异常检测任务上取得了显著进展。然而，现有方法仍存在三个主要局限。

第一，模态缺失的鲁棒性问题。 实际工业产线中，不同传感器可能因故障、遮挡或成本考虑而发生部分模态缺失。例如，深度相机可能因光照条件不稳定而失效，红外传感器可能因温度漂移而数据异常。多数现有方法假设模态完整，在模态缺失时性能急剧下降。

第二，多子类正常分布的建模困难。 “正常”在工业场景中往往包含多种子类——同一产品可能存在不同规格、不同颜色、不同版本的合格品。传统方法通常假设正常样本服从单一高斯分布，面临多子类时会将合格的差异误判为异常。

第三，异常感知生成的需求。 单纯的异常检测二分类结果不足以指导实际修复。工业场景需要模型在检测到异常后，能够在异常位置生成合理的修复内容，以辅助人工决策或自动修补。

1.2 主要贡献#

针对上述挑战，本文提出统一的异常检测与生成框架，主要贡献如下：

1. 动态模态路由机制：通过可学习的注意力权重对现有模态进行自适应融合，实现对任意模态组合的鲁棒编码
2. 实时推理性能：正常样本推理时间15.2ms（>60 FPS），满足工业产线实时检测需求
3. 可学习原型记忆模块：存储多个原型向量以显式建模多子类正常分布，并通过对比学习使正常样本靠近其最近原型
4. 布局引导扩散模型：以异常区域的掩码为条件，实现异常区域的精确可控生成

2.1 异常检测#

基于重构的方法通过自编码器学习正常样本的分布，将重构误差大的区域判定为异常。基于嵌入的方法提取预训练特征并计算与正常样本库的距离，代表性工作包括：

• PaDiM[1]：概率分布建模，使用多元高斯分布估计正常特征分布
• PatchCore[2]：核心集采样，保留最具代表性的正常特征

这些方法在MVTec AD数据集上取得了优异性能。然而，它们均假设模态完整且所有正常样本来自同一分布。

2.2 多模态异常检测#

近年来，研究者开始探索融合图像和深度图的多模态异常检测：

• MMV Imputer：通过跨模态注意力补全缺失模态，但依赖大量完整模态数据训练
• 3D-STPM：融合RGB和深度特征进行蒸馏，但在模态缺失时无法工作
• CFA[3]：多模态特征对齐，通过对比学习增强模态间一致性

本文方法与上述工作的核心区别在于：我们显式设计了针对任意模态组合的鲁棒编码机制，而非假设模态完整。

2.3 扩散模型在异常检测中的应用#

扩散模型在图像生成任务中展现了强大的能力：

• DiffAD[4]：将扩散模型用于异常模拟，通过合成伪异常样本辅助训练
• DDPM[5]：经典扩散概率模型，为后续工作奠定基础

然而，现有工作主要关注生成质量而非异常区域的精确定位与可控生成。本文是较早将布局引导引入扩散模型的工作之一，实现异常区域的精确可控修复。

3.1 问题定义与符号#

设多模态输入为 $\mathcal{X} = \{x_1, x_2, \ldots, x_M\}$，其中 $M$ 为模态总数，$x_i$ 为第 $i$ 个模态的特征向量。部分模态可能缺失，缺失的模态记为 $x_i = \text{None}$。定义模态指示向量 $\mathbf{m} \in \{0,1\}^M$，其中 $m_i=1$ 表示第 $i$ 个模态存在，$m_i=0$ 表示缺失。

训练阶段仅使用正常样本，测试阶段输入可能包含异常样本。任务为：(1) 判断样本是否异常；(2) 定位异常区域；(3) 在异常位置生成修复内容。

3.2 动态模态路由机制#

为解决任意模态组合下的鲁棒编码问题，我们提出动态模态路由机制。对于每个现有模态，首先通过独立的编码器 $E_i$ 提取特征：

其中 $E_i$ 为参数共享的模态编码器，$\mathbf{z}_i \in \mathbb{R}^d$ 为第 $i$ 个模态的嵌入向量。

随后，可学习路由模块通过注意力机制计算各模态的融合权重：

其中 $\mathbf{q} \in \mathbb{R}^d$ 为可学习的查询向量，$\sqrt{d}$ 为缩放因子防止梯度消失。

最终融合嵌入为：

该机制具有以下性质：当模态完整时，各模态按注意力权重融合；当模态缺失时，权重重新分配，不受缺失影响。

3.3 原型记忆模块#

为建模包含多个子类的正常分布，我们设计可学习的原型记忆模块。该模块维护 $K$ 个原型向量 $\mathbf{p}_1, \mathbf{p}_2, \ldots, \mathbf{p}_K \in \mathbb{R}^d$，每个原型代表一个正常子类。

对于融合嵌入 $\mathbf{z}$，模型计算其与所有原型的欧氏距离：

其中 $d_k$ 为第 $k$ 个原型的距离，值越小表示越接近该子类。

样本的异常分数定义为到最近原型的距离：

异常分数越大，样本越可能异常。训练阶段，模型通过对比学习使正常样本靠近其最近原型、远离其他原型：

其中 $d_{\text{min}} = \min_k d_k$，$\tau$ 为温度超参数（默认0.1）。

同时，为防止所有原型坍缩到同一区域，我们引入聚类正则项：

其中 $\delta$ 为最小间隔超参数（默认1.0），该损失鼓励原型之间保持距离。

总损失为：

其中 $\lambda$ 为平衡系数（默认0.1）。

3.4 模型自训练的核心训练循环#

就像我刚才说的一样，下面给出的代码，展现的是一段使模型可以自己训练的核心训练循环：

1
import torch
2
import torch.nn as nn
3
import torch.optim as optim
4

5
class DynamicRouter(nn.Module):
6
    """动态路由模块 - 与3.2节公式(1)(2)(3)完全等价"""
7
    def __init__(self, d_model):
8
        super().__init__()
9
        self.query = nn.Parameter(torch.randn(d_model))
10
        self.scale = d_model ** 0.5
11

12
    def forward(self, embeddings, mask):
13
        """
14
        embeddings: [batch, M, d] - 各模态特征
15
        mask: [batch, M] - 模态存在指示
16
        """
17
        # 公式(2): 计算注意力权重
18
        scores = torch.einsum('bmd,d->bm', embeddings, self.query) / self.scale
19
        scores = scores.masked_fill(~mask, -1e9)
20
        weights = torch.softmax(scores, dim=-1)  # [batch, M]
21

22
        # 公式(3): 加权融合
23
        fused = torch.einsum('bm,bmd->bd', weights, embeddings)
24
        return fused, weights
25

26
class PrototypeMemory(nn.Module):
27
    """原型记忆模块 - 与3.3节公式(4)-(7)完全等价"""
28
    def __init__(self, d_model, num_prototypes, delta=1.0):
29
        super().__init__()
30
        self.prototypes = nn.Parameter(torch.randn(num_prototypes, d_model))
31
        self.delta = delta
32

33
    def forward(self, z):
34
        """
35
        z: [batch, d] - 融合嵌入
36
        返回: min_dist [batch], distances [batch, K]
37
        """
38
        # 公式(4): 平方欧氏距离
39
        distances = torch.cdist(z, self.prototypes, p=2) ** 2  # [batch, K]
40
        min_dist, min_idx = distances.min(dim=1)
41
        return min_dist, distances
42

43
    def clustering_loss(self):
44
        """公式(7): 防止原型坍缩"""
45
        pairwise_dist = torch.cdist(self.prototypes, self.prototypes, p=2) ** 2
46
        mask = ~torch.eye(pairwise_dist.size(0), device=pairwise_dist.device).bool()
47
        violation = torch.relu(self.delta - pairwise_dist[mask])
48
        return violation.mean()
49

50
def train(model, dataloader, config):
51
    """核心训练循环"""
52
    optimizer = optim.Adam(model.parameters(), lr=config.lr)
53

54
    for epoch in range(config.max_epochs):
55
        epoch_loss = 0.0
56
        for batch in dataloader:
57
            # 模拟模态缺失（训练增强）
58
            features = []
59
            mask = []
60
            for i, mod_feat in enumerate(batch['modalities']):
61
                if torch.rand(1).item() < config.dropout_prob:
62
                    # 随机丢弃模态
63
                    features.append(torch.zeros_like(mod_feat))
64
                    mask.append(False)
65
                else:
66
                    features.append(mod_feat)
67
                    mask.append(True)
68

69
            embeddings = torch.stack(features, dim=1)  # [B, M, d]
70
            mask = torch.tensor(mask, device=embeddings.device)  # [M]
71

72
            # 动态路由编码（公式1-3）
73
            fused, weights = model.router(embeddings, mask)
74

75
            # 原型匹配（公式4-5）
76
            min_dist, distances = model.prototype_memory(fused)
77

78
            # 对比损失（公式6）
79
            contrast_loss = -torch.log(
80
                torch.exp(-min_dist / config.tau) /
81
                torch.exp(-distances / config.tau).sum(dim=1)
82
            ).mean()
83

84
            # 聚类正则（公式7）
85
            cluster_loss = model.prototype_memory.clustering_loss()
86

87
            # 总损失（公式8）
88
            loss = contrast_loss + config.lambda_cluster * cluster_loss
89

90
            loss.backward()
91
            optimizer.step()
92
            optimizer.zero_grad()
93

94
            epoch_loss += loss.item()
95

96
        print(f"Epoch {epoch}: Loss = {epoch_loss / len(dataloader):.4f}")

其中，随机模态丢弃的作用是以概率 dropout_prob 丢弃某一种模态的特征向量，增强模型对模态缺失的鲁棒性；动态路由模块使用可学习的注意力权重对现有模态做加权融合，与3.2节公式(1)-(3)完全等价；原型记忆模块存储 $K$ 个可更新的原型向量；对比损失让正常样本靠近最近的原型、远离其他原型，聚类正则防止原型坍缩。

3.5 布局引导的扩散模型#

对于异常区域生成任务，我们引入以异常布局为条件控制的扩散模型。设查询样本的融合嵌入为 $\mathbf{z}$，原型记忆模块预测的异常区域掩码为 $\mathbf{y}$（二值掩码，1表示异常），目标生成修复内容 $\hat{\mathbf{x}}$。

我们采用条件扩散模型。前向过程逐步添加噪声，逆向过程学习去噪。条件表示为 $\mathbf{c} = [\mathbf{z}; \mathbf{y}]$，即融合嵌入与异常掩码的拼接。训练目标为：

其中 $t \sim \text{Uniform}(1, T)$ 为时间步，$\epsilon \sim \mathcal{N}(0, \mathbf{I})$ 为标准高斯噪声，$\bar{\alpha}_t$ 为噪声调度参数，$\epsilon_\theta$ 为噪声预测网络。

以下是布局引导的扩散生成采样过程的核心伪代码：

1
def diffuse_generate(z, anomaly_mask, original, model_diffusion, steps=50):
2
    """
3
    布局引导的扩散生成采样过程
4

5
    参数:
6
        z: 融合嵌入 [d]
7
        anomaly_mask: 异常区域掩码 [H, W]，1表示异常需要修复
8
        original: 原始图像 [H, W, C]
9
        model_diffusion: 预训练的扩散噪声预测模型
10
        steps: DDIM采样步数
11

12
    返回:
13
        x_final: 修复后的图像
14
    """
15
    # 从纯噪声开始
16
    x_t = torch.randn_like(original)
17

18
    for t in reversed(range(steps)):
19
        # 预测当前噪声
20
        t_tensor = torch.full((1,), t, device=x_t.device)
21
        eps_pred = model_diffusion(x_t, t_tensor, condition=z)
22

23
        # DDIM采样更新
24
        alpha_t = alpha_schedule[t]
25
        alpha_prev = alpha_schedule[t-1] if t > 0 else 1.0
26
        x_prev = (x_t - (1 - alpha_t).sqrt() * eps_pred) / alpha_t.sqrt()
27
        x_prev = x_prev * alpha_prev.sqrt() + (1 - alpha_prev).sqrt() * eps_pred
28

29
        # 核心: 布局引导 - 强制非异常区域与原始输入一致
30
        # 这一步实现了异常区域的精确可控生成
31
        x_prev = x_prev * anomaly_mask + original * (1 - anomaly_mask)
32

33
        x_t = x_prev
34

35
    return x_t

其中，x_prev * anomaly_mask + original * (1 - anomaly_mask) 实现了布局引导的核心——强制模型在非异常区域保留原始内容，使生成动作仅作用于异常区域。这保证了修复的精确性和可控性。

需要特别说明的是，扩散生成模块仅在检测到异常（即 $s > \text{阈值}$）时才被激活，以避免对正常样本的不必要计算。

3.6 与现有方法的理论对比#

本方法在实现模态缺失鲁棒性的同时，首次将实时推理性能、多子类建模和异常生成整合到统一框架中。

4.1 数据集与设置#

数据集：使用 MVTec AD 3D 数据集，包含 10 个类别（榛子、螺丝、饼干等），每个类别包含 RGB 图像和深度图两个模态。训练集仅包含正常样本，测试集包含正常与异常样本。

模态缺失模拟：以概率 30% 随机丢弃深度模态，模拟实际场景中的传感器失效。

基线方法：

· PaDiM[1]：概率分布建模的单模态异常检测 · PatchCore[2]：核心集采样方法 · 本方法（完整）：完整的双模态输入 · 本方法（消融）：移除原型记忆模块，使用单原型

评价指标：

· AUROC：异常检测准确率（越高越好） · FID：生成质量（越低越好，仅本方法） · 参数量：模型参数量（单位M） · 推理速度：FPS（帧每秒）

实现细节：

· 编码器：ResNet18预训练特征 · 原型数量：$K=10$ · 嵌入维度：$d=512$ · 温度系数：$\tau=0.1$ · 聚类间隔：$\delta=1.0$ · 平衡系数：$\lambda=0.1$ · 优化器：Adam，学习率 $10^{-4}$

4.2 模型架构图#

以下是具体图表（模型架构图）：

图1：本文提出的整体模型架构图。 输入层支持RGB图像和深度图两种模态（文本模态作为扩展方向）。每个模态经独立编码器提取特征后，动态路由模块通过可学习的注意力权重进行自适应融合（公式1-3）。原型记忆模块计算融合嵌入到K个原型中心的距离（公式4），最小距离作为异常分数（检测路径），同时预测异常区域掩码。当检测到异常时（虚线箭头），扩散生成模块以融合嵌入和异常掩码为条件，生成修复内容（公式9）。实线箭头表示所有训练样本（仅正常）均经过检测路径并参与损失计算（公式8）。

4.3 定量实验结果#

以下是具体图表（实验结果表）：

结果分析：

1. 模态完整时：本方法AUROC达到0.923，相比PatchCore的0.891提升了3.6%，证明了动态路由和原型记忆的有效性。
2. 模态缺失30%时：本方法AUROC为0.897，仅下降2.6%（从0.923）。相比之下，PaDiM下降12.1%（从0.852到0.731），PatchCore下降12.3%（从0.891到0.768）。这一结果证明动态路由机制对缺失模态的鲁棒性——下降幅度仅为基线的1/5。
3. 消融实验：本方法(无原型)移除原型记忆模块后，AUROC在模态完整时降至0.889（下降3.4%），在模态缺失时降至0.812（下降9.5%）。这证明了多原型机制对多子类建模的必要性。
4. 生成质量：本方法的FID为32.5，在工业场景中达到可用水平（通常FID<50可认为生成质量可接受）。
5. 参数量与速度：本方法完整版本参数量24.3M，推理速度67 FPS，在精度和效率之间取得良好平衡。

4.4 推理效率分析#

工业场景中，推理速度是重要考量指标。在单个NVIDIA Tesla T4 GPU上，本方法各阶段的平均推理时间为：

关键结论：

· 对于99%的正常样本（不触发生成模块），推理时间仅15.2ms，对应>65 FPS，满足工业实时检测需求（通常要求>30 FPS） · 对于约15%的异常样本（触发生成模块），总推理时间约202ms（~5 FPS），作为后处理辅助决策可接受 · 扩散生成模块仅在检测到异常时才激活，避免了不必要的计算开销

4.5 定性结果对比#

以下是具体图表（定性结果对比）：

图2：定性结果对比。 上行为螺丝类别，包含点状表面缺陷；下行为榛子类别，包含划痕缺陷。从左到右依次为：原始输入图像、真实异常掩码（Ground Truth）、PaDiM方法的热力图、本方法的热力图、本方法生成的修复结果。

关键观察：

1. 假阳性降低：PaDiM在螺丝样本中将正常纹理区域错误标记为异常（红色大面积区域），假阳性率约18%；本方法通过原型记忆模块将此类纹理归入正确的正常子类，假阳性率降至7%，相对降低61%
2. 定位精度提升：PaDiM对划痕缺陷的定位边界模糊，难以区分缺陷区域与正常区域；本方法通过动态路由融合多模态信息，生成边界清晰的热力图
3. 生成质量：本方法生成的修复内容在视觉上连贯，与周围正常纹理融合自然，可用于辅助人工复核或下游任务