面向更平坦极小点的梯度混合与外推学习：一种提升领域泛化的新方法

《Artificial Intelligence》：Learning Gradient-based Mixup with Extrapolation toward Flatter Minima for Domain Generalization

在领域泛化（DG）中，如何解决训练与测试数据间的分布偏移是核心挑战。为解决现有方法对源域过拟合、潜在未覆盖区域有限的问题，作者开展了名为FGMix的梯度混合研究。研究通过设计一种梯度兼容性权重的策略，并结合内插和外推进行混合，有效引导模型向更平坦的极小点学习。结果表明，该方法在DomainBed基准测试上实现了最先进的性能，显著提升了模型在未见域上的泛化能力。

　　

机器学习系统通常建立在训练数据和测试数据来自同一分布（即独立同分布，i.i.d.）的假设之上。然而，现实世界的应用场景中，这一假设常常不成立。例如，当训练数据和测试数据由不同的设备采集，或在不同的环境条件下获取时，就会发生分布偏移。面对这种偏移，模型如果学习了虚假的相关性（如纹理或背景）来进行预测，其在测试数据上的表现就可能急剧下降。领域泛化（Domain Generalization, DG）正是为了应对这一挑战而提出的。它利用来自多个源域的数据来训练一个模型，旨在让这个模型能够很好地泛化到训练时未见的、全新的目标域上。现有的DG方法，虽然各有侧重，但大多仍然难以摆脱对源域过拟合的困境。究其原因，部分在于特征空间中预期区域的覆盖范围有限。想象一下，如果目标域位于源域构成的区域之外，那么仅在源域上训练良好的分类器，在面对目标域时可能就表现不佳了。这正是当前DG领域亟待解决的一个核心难题。

为了应对上述挑战，来自新加坡南洋理工大学的研究团队在《Artificial Intelligence》期刊上发表了题为“Learning Gradient-based Mixup with Extrapolation toward Flater Minima for Domain Generalization”的研究论文。他们创新性地提出了一种名为平坦感知梯度混合的方法。这项研究的核心思路在于，通过数据的内插和外推来进行混合，以覆盖潜在的未见区域，从而扩展模型能够探索的特征空间。更重要的是，他们设计了一种精巧的策略来生成每个样本在混合时的权重，这个策略的核心是基于梯度。梯度是模型学习方向的关键指示器。研究团队认为，与同一组合中其他样本梯度兼容性更高的样本，往往携带了更多跨域不变的信息，因此在混合时应被赋予更大的权重。通过这种方式生成的混合数据，能够吸收更多包含不变特征的样本信息。为了让混合不仅限于内插（即在源域样本构成的凸包内部），还能进行外推（扩展到凸包外部），他们对归一化后的权重进行了缩放和平移操作，移除了权重的非负约束。最终，模型被鼓励去学习一个更平坦的损失曲面。损失曲面的平坦程度与模型的泛化能力密切相关，一个平坦的极小点意味着模型参数在发生微小扰动时，其性能变化不大，因而通常具有更好的泛化性。通过在DomainBed这一权威的领域泛化基准测试上进行大量实验，该方法（FGMix）验证了其各项设计的有效性，并展示出优于其他现有DG算法的性能。

为开展这项研究，作者们主要运用了几项关键技术方法。首先是潜在空间数据混合，研究在特征提取器生成的潜在表征空间中对同一类别、来自不同源域的样本进行线性组合，而非在原始输入或浅层特征上进行，这有助于在更高语义层次进行有效的数据增强。其次是基于梯度的权重生成策略，这是FGMix的核心创新，它通过一个可学习的兼容性评分函数 A_ω(·,·) 来计算任意两个样本梯度之间的相关性，并据此为每个样本分配混合权重，该权重随后经过Softmax归一化及可学习的缩放因子λ的调整，以同时支持内插和外推。第三是面向平坦极小点的策略学习，研究通过优化权重生成策略（即A_ω和λ），使得基于混合数据产生的损失曲面在分类器参数?的邻域内更加平坦，这通过蒙特卡洛采样近似损失差异来实现。最后，作为一种温和的正则化以防止过度外推，研究引入了对抗性训练约束，通过一个判别器来确保生成的混合数据分布不会过度偏离一个基于源域数据估计的先验分布（高斯分布）。

3.1. 潜在空间数据增强

研究在潜在空间对来自k个源域、同一类别的样本进行线性混合，生成混合潜在编码z_mix，并用其与原始数据一起训练分类器。通过允许混合权重a_i为负值，该方法能同时实现内插（所有权重介于0和1之间）和外推（部分权重为负）。训练目标是最小化原始数据损失和混合数据损失的平均值。该过程的概览如图2a所示。mix are used together with the original latent codes {Z₁, Z₂, Z₃} to train the classifier h_?. (b) Weight Generation Policy: the weight assigned to each latent code is based on the sum of its gradient’s compatibility scores w.r.t. other latent codes. A softmax and a shifting & scaling layers are applied to enable extrapolation. The compatibility score function A_ωand the scaling factor λ are learned towards flatter minima (i.e., L_flat) and matching the generated latent codes to a prior via adversarial training (i.e., L_adv).">

3.2. 基于梯度的权重生成

与以往从预设分布中随机采样权重不同，FGMix根据样本间的关系（具体是梯度的兼容性）来分配权重。梯度相似性反映了样本间在学习过程中信息共享的程度。梯度方向越一致，意味着它们包含更多共享的、跨域不变的信息。研究通过一个可学习的2层MLP网络A_ω来计算任意两个样本梯度g_i和g_j之间的兼容性得分。对于第i个样本，其总分s_i是其梯度与组合内所有其他样本梯度的兼容性得分之和。得分更高的样本被认为携带更多不变信息，因此在混合中被赋予更高的重要性。随后，对得分进行Softmax归一化得到初始权重a_i。为了支持外推，引入一个可学习的缩放因子λ，对权重进行缩放和偏移，得到最终混合权重?_i，此过程去除了权重的非负约束。权重生成策略的整体流程如图2b所示。

3.3. 学习权重生成策略

权重生成策略中的两个关键可学习部分——兼容性评分函数A_ω和缩放因子λ，其学习目标是引导模型产生一个更平坦的损失曲面。平坦性通过衡量当前分类器参数?与其γ-邻域内的参数?‘在混合数据损失上的差异来量化。优化目标是最小化该差异的平方的期望，从而使得在?附近，损失对参数变化不敏感。同时，为了约束生成的数据分布不过度偏离源域数据分布，研究引入了一个对抗性损失L_adv。一个判别器被训练来区分从源域先验分布（高斯分布）中采样的“真实”潜在编码和由FGMix生成的“伪造”混合潜在编码。权重生成策略则试图生成能够“欺骗”判别器的混合数据。最终的策略损失L_pol是平坦性损失L_flat和对抗性损失L_adv的加权和。整个FGMix模型的训练流程是联合优化基础模型参数{θ, ?}和策略参数{ω, λ}，并在训练后期（从第τ次迭代开始）才启用权重生成和混合数据训练，以确保特征提取器首先得到良好训练。

4. 实验结果

研究在DomainBed测试平台的五个真实世界基准数据集上进行了广泛的实验，包括PACS、VLCS、OfficeHome、TerraIncognita和DomainNet。实验遵循DomainBed的标准评估流程，使用ImageNet预训练的ResNet-50作为基础模型。FGMix在大多数数据集上均取得了优异的表现。与基线方法相比，包括朴素的经验风险最小化、多种混合方法以及梯度处理方法，FGMix在平均准确率上显示出竞争力。消融实验验证了梯度兼容性权重、外推能力（通过可学习λ实现）以及平坦性目标各自的有效性。定性分析表明，FGMix生成的混合权重分布确实倾向于强调那些梯度兼容性更高的样本，并且其产生的混合数据有助于探索更平坦的损失区域。此外，研究还发现将FGMix与现有的平坦性感知优化器（如SWAD）结合使用，能带来进一步的性能提升。

结论与讨论

本研究提出了一种新颖的领域泛化方法——平坦感知梯度混合。FGMix的核心贡献在于设计了一种基于梯度兼容性的、可学习的权重生成策略，用于在潜在空间进行数据混合，并同时支持内插和外推。该方法通过优化权重生成策略本身，主动塑造一个更平坦的损失曲面，从而鼓励模型收敛到泛化能力更佳的平坦极小点。理论分析和大量实验表明，这种策略能够有效缓解模型对源域的过拟合，扩展模型在特征空间中的探索范围，提高覆盖未知目标域最优解的可能性。在DomainBed基准测试上取得的先进性能，强有力地证明了FGMix的有效性。这项工作为领域泛化研究提供了一个新的视角，即不仅可以在给定的损失曲面上寻找平坦区域，还可以通过智能的数据生成来主动“平坦化”损失曲面本身。未来的工作可以探索将类似思想应用于更广泛的分布外泛化问题，或者与更强大的基础模型架构结合。