DGMamba: Domain Generalization via Generalized State Space Model论文笔记

DGMamba: Domain Generalization via Generalized State Space Model

paper: https://arxiv.org/abs/2404.07794
code: https://github.com/longshaocong/DGMamba(暂时没有可用代码)

摘要

• 领域泛化(DG)旨在解决各种场景下的分布转移问题。
• 现有的方法基于卷积神经网络(CNN)或视觉transformer(ViT), 它们存在接受域有限或二次复杂度的问题。
• Mamba作为一种新兴的状态空间模型(SSM)，具有优越的线性复杂度和全局感受野。尽管如此，由于隐藏状态问题和不适当的扫描机制，它很难应用于DG来解决分布变化。
• 作者提出了一种新的DG框架(DGMamba)，它具有对未见领域的强泛化性，同时具有全局感受野和高效线性复杂度的优点。
• DGMamba包含两个核心组件: 隐藏状态抑制(HSS)和语义感知补丁细化(SPR)。HSS是为了减轻输出预测期间与特定领域相关的隐藏状态的影响。SPR是为了让模型更多地关注对象而不是上下文，由免先验扫描(PFS)和域上下文交换(DGI)。
• PFS的目标是打乱图像中的非语义补丁，从图像中创建更灵活、更有效的序列。DGI的目的是通过融合域之间的补丁，结合不匹配的非语义和语义信息来正则化Mamba。

动机

• 目标是增强Mamba模型对未知领域的通用性。
• 作者观察到隐藏状态作为Mamba的重要组成部分，通过记录序列数据中的历史信息，促进全局感受野，在长期相关性建模中发挥着重要作用。然而，当处理包含来自不同领域的不同特定信息的看不见的图像时，这种隐藏状态可能会产生不良效果。如上图(a)所示，特定领域的信息可能会在传播过程中在隐藏状态中累积甚至放大，从而降低泛化性能。
• 如何有效地将2D图像扫描成适合DG中Mamba的1D序列数据仍然是一个悬而未决的问题，因为图像的像素或斑块不表现出序列数据中存在的必要因果关系。尽管最近的工作已经探索了视觉任务的各种扫描策略，但这种简单的一维遍历策略可能会导致生成的序列数据中出现意外的特定领域信息(上图c), 从而削弱了Mamba解决分布变化的能力。此外，这些固定扫描策略很大程度上忽略了与领域无关的扫描，并且对各种不同的场景高度敏感，因此很难应用于DG。

DGMamba设计

DGMamba是一种新颖的基于状态空间模型的领域泛化框架，它擅长对未见领域的强泛化性，同时具有全局感受野和高效线性复杂度的优势。

隐藏状态抑制(HSS)

HSS通过在输出预测期间选择地抑制相应的隐藏状态来消除隐藏状态中包含的非语义信息的有害影响。通过减少SSM层中的非语义信息，DGMamba学习域不变特征。

为了抑制隐藏状态中传达的意外的特定领域信息，最初的任务是识别包含不利因素的隐藏状态。在隐藏状态的传播过程中，与真实标签显示出更强相关性的隐藏状态应该被保留。因此，他需要SSM中(SSM相关可以查看Mamba)$\bar{\mathbf{A}}$中较大的传播系数，而关联较少的隐藏状态需要$\bar{\mathbf{A}}$中相对较小的洗漱。因此, $\Delta A$的值用于确定哪些隐藏状态将受到抑制。从数学上讲，隐藏状态抑制的策略为
$$\begin{array}{rl}{y}_t& =\bar{\mathbf{C}}{h}_t\\ \overline{mathbfC}& =\mathbf{C}\odot \mathbf{M}\\ \mathbf{M}& =(\Delta >\alpha )+(1-(\Delta A>\alpha )\odot \Delta A)\end{array}$$
其中, $\alpha \in [0,0.5]$表示确定是否应抑制隐藏状态的阈值。这样，系数参数$\Delta A<=\alpha$的隐藏状态将被$\Delta A$抑制，而其余隐藏状态保持不变。

语义感知补丁细化(SPR)

SPR的作用是鼓励模型更多地关注对象而不是上下文，它包含免先验扫描(PES)和域上下文交换(DCI)。

PES旨在对图像内的上下文补丁进行打乱，并减少对标签预测的贡献。它提供了一种有效地2D扫描机制，将2D图像遍历成1D序列数据。因此，PFS具有将模型的注意力从上下文转移到对象的能力。

为了减轻跨不同域的不同上下文信息和局部纹理细节的影响，DCI将图像的上下文补丁替换为来自不同域的上下文补丁。所提出的DCI引入了局部纹理噪声，并根据不匹配的上下文和对象的组合对模型进行了正则化。通过利用线性复杂性和异构上下文标记，DCI可以有效地学习更稳健的表示。

从领域不变角度来看，上下文和对象是两个基本元素。对象对应前景，对分类结果贡献最大，在不同场景中保持静止。上下文与特定域的信息相关，例如背景和图像样式，这些信息在不同域之间差异很大。因此，将模型的焦点转向对象可以帮助减少特定于领域的信息。

免先验扫描

尽管基于SSM的模型在视觉任务中表现出了出色的性能，但多样化和随机的上下文环境对于在DG中部署Mamba仍然至关重要。这一结论表明，仍然需要一种有效的扫描机制来应对图像像素或块之间的非因果相关性带来的挑战。合适的扫描机制应该能够打破由手动创建的图像序列引起的意外的虚假相关性。然而，现有的基于SSM的方法仅限于以固定的展开方式将图像扫描成patch。这些助管遍历策略可能会在生成的序列中产生域特定信息，从而使这些模型难以解决DG中的分布变化。

为了打破patch之间的虚假相关性并为DG任务提供有效的扫描机制，作者提出了PFS来解决Mamba中的方向敏感问题。如主图所示，PFS尝试随机打乱上下文补丁，这可能有助于展开序列中的域特定信息，同时保留对象patch。特别是，对于表示$z=z_c+z_o\in {\mathbb{R}}^{H\times W\times C}$, 其中$z_c$和$z_o$表示上下文信息和对象信息，PFS策略后的$z_{pfs}$可以表示为:
$$\begin{array}{rl}{z}_{pfs}& =z_c^s+z_o\\ z_c^s& =Shuffle(z_c)\end{array}$$
其中$z_c^s$表示在空间维度上使用$Shuffle(\cdot )$函数打乱的上下文信息。该操作可以通过生成上下文干扰或噪声，同时保持一致的对象信息，为Mamba提供展现灵活扫描方向的序列数据。因此，它减轻了手动固定扁平化策略带来的域特定信息，并打破了虚假的相关性。

域上下文交换

DG中不同域的上下文信息是异构的。PFS中的上下文patch在给定场景中受到限制，不足以提供足够多样化的上下文信息来删除域特定信息。此外，来自不同域的异构上下文patch不仅表现出不同的上下文信息，而且还包含不同的局部纹理特征。

为了充分解决异构上下文和不同局部纹理细节的不利影响，作者建议创建充足的上下文场景并通过域上下文交换(DCI)引入局部纹理噪声。DCI将图像上下文patch替换为来自不同域的补丁。该操作在反制样本上对模型进行正则化，即一个域中的语义信息和来自不同域的非语义特征的组合。这种策略进一步迫使模型专注于可概括的特征，同时丢弃文本细节或其他特定的领域特征。

上下文patch识别

为了区分文本和目标patch, 利用Grad-CAM作为度量来衡量不同域的贡献图像。由于包含对象的区域会极大地激活Grad-CAM, 而表现出上下文patch在Grad-CAM中具有较低的值。因此，图像patch根据Grad-CAM生成的激活图中的值分为context和object。具体的说，将小于$25$的Grad-CAM值的patch确定为上下文信息$z_c$, 其余的则为对象信息$z_o$

实验结果

遵循 DG中的标准协议，作者评估了我们提出的 DGMamba 的有效性，并在五个常用基准上将其与 DG 中最先进的方法进行比较：（1）PACS 包含 9991 张图像，分为 7 个类别，展示 4 种风格。 （2）VLCS涉及4个数据集，共10729张图像，分布在5个类别。 (3) OfficeHome 包含来自 4 个数据集的 65 个类别的 15588 张图像。 (4) TerraIncognita包含在 4 个不同地点拍摄的 10 种动物的 24330 张照片。 (5) DomainNet包含 586575 张图像，分为来自 6 个域的 345 个类别。