📊 statistics

AugMask: Training Diffusion Models on Incomplete Tabular Data via Stochastic Augmentation and Masking

本文介绍了 AugMask，这是一个即插即用的训练框架，它通过使用随机增强将缺失值作为不确定条件上下文进行填充，同时将去噪监督限制在仅观测到的坐标上，从而使标准的基于评分的扩散模型能够有效地生成不完整的表格数据。

原作者： Jungkyu Kim, Taeyoung Park, Kibok Lee

发布于 2026-06-03

📖 1 分钟阅读☕ 轻松阅读

原作者： Jungkyu Kim, Taeyoung Park, Kibok Lee

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是对 AugMask 论文的解释，通过简单的概念和日常类比进行了拆解。

问题所在：教一位拿着残缺食谱的大厨

想象一下，你正在试图教一位顶级大厨（扩散模型/Diffusion Model）如何烹饪一道特定的菜肴（生成真实的表格数据）。这位大厨非常擅长遵循食谱，但问题在于：你给他的食谱是不完整的。有些食材缺失了，原本应该有食材的地方变成了空白或者被标记为“NaN”（非数字）。

标准的厨师（AI 模型）会被空白处搞糊涂。如果食谱写着“加入 2 杯 [空白]”，他们就没法做菜了。

旧方法（零填充/Zero-Filling）： 过去人们只是用一个占位符来填充空白，比如“0”或“水”。但这就像告诉大厨：“加入 0 杯糖。”大厨会误以为“0”是一种真实的食材，从而破坏了味道（数据分布）。
“猜谜”法（插补/Imputation）： 另一种方法是让大厨去猜缺失的食材，把它写下来，然后训练大厨去完美地实现那个猜测。但如果大厨猜错了，他就会一直试图去完善一个错误的猜测，导致厨房变得混乱且充满噪声。

解决方案：AugMask

作者提出了 AugMask，这是一种新的训练策略，它扮演了一个聪明的副厨角色。它通过将两项工作分离来解决问题：布置场景 和评分。

1. 布置场景（随机增强/Stochastic Augmentation）

副厨（一个轻量级的辅助模型）并不只是留下空白，也不用无聊的“0”来填充，而是用合理的猜测来填补空白。

转折点： 副厨不仅仅给出一个猜测，它给的是一个范围。
- 类比： 如果食谱缺少“盐”，副厨不会只说“1 茶匙”，而是说：“它可能在 0.5 到 1.5 茶匙之间。”
为什么？ 这教会了主厨，缺失的食材是具有不确定性的。主厨学会了观察全局（其他食材），并理解这个缺失的部分是一个“可能”，而不是一个“定论”。

2. 评分（掩码损失/Masking the Loss）

这是最重要的一部分。当主厨尝试重现这道菜时，老师（训练算法）只对原本存在的食材进行评分。

规则： 如果食谱原本有“2 杯面粉”，老师会检查大厨的面粉是否正确。
漏洞处理： 如果食谱原本在“盐”的位置是空白的，老师就会忽略大厨对盐的猜测。老师会说：“我不在乎你对盐猜了多少，只要确保面粉是对的就行。”
结果： 大厨学会了利用这些“猜测”作为上下文（线索）来帮助完成真实的食材，但他们不会因为猜错了猜测而受到惩罚。他们不再试图死记硬背那些猜测，而是开始学习真实数据之间的关系。

为什么有效： “不确定性惩罚”

论文使用了一些高级数学来解释为什么这有效，你可以通过一个简单的隐喻来理解：颤抖的手。

确定性猜测（不好）： 如果副厨给出一个单一且自信的猜测（例如，“正好是 1.0 茶匙”），主厨会想：“好吧，我必须完全匹配这个。”如果猜测稍有偏差，大厨就会感到困惑。
随机性猜测（好）： 如果副厨说：“在 0.5 到 1.5 之间”，主厨会意识到：“这是一个颤抖的、不确定的线索。”
惩罚机制： 数学证明，当线索是颤抖的（高不确定性）时，主厨自然会停止过度依赖它。这就像一名学生在考试中忽略了一个模糊的提示，因为他知道这个提示可能是错的。这防止了模型“陷入”错误的猜测中。

结果：一位更好的大厨

作者在许多不同的数据集上（如成年人收入数据、购物习惯等）测试了 AugMask，涵盖了不同程度的缺失信息（从缺失 10% 到缺失 90%）。

结果： 通过使用这种“布置场景、忽略猜测”的策略，标准的 AI 模型（通常无法处理缺失数据）的表现甚至优于专门为处理缺失数据而设计的模型。
核心要点： 你不需要为一个不完整的食谱建造一个特殊的、复杂的厨房。你只需要一种聪明的方法，用“不确定的猜测”来填补空白，然后告诉大厨在最后评分时忽略这些猜测即可。

一句话总结

AugMask 通过用“不确定的猜测”来填补空白以提供上下文，同时严格仅对实际存在的数据进行评分，从而教会 AI 模型如何处理缺失数据，防止其被自己的猜测所误导。

技术摘要：AugMask

问题陈述

基于评分的扩散模型已成为强大的深度生成模型，但在应用于包含缺失值的表格数据时面临显著挑战。标准的扩散架构假设输入是完全指定的，而现实世界的表格数据通常包含缺失条目（NaNs）。

现有的弥补这一差距的方法通常分为两类，两者都存在局限性：

常数占位符插补（Constant Placeholder Imputation）： 在训练前将缺失值替换为固定常数（例如零、列均值）。这迫使模型将缺失条目视为有效的零，从而引入稀疏性偏差，并使预测依赖于缺失程度。此外，即使损失函数对观测到的坐标进行了掩码处理，确定性的占位符仍作为输入存在，成为与缺失性相关的伪相关特征。
自插补（Self-Imputation，例如 DiffPuter）： 在训练过程中通过期望最大化（EM）方法迭代更新缺失条目。虽然这避免了固定占位符，但它创造了一个“移动目标”优化问题，即模型是在其自身不断演变的预测基础上进行训练。这可能导致梯度噪声、非平稳动力学以及过度分散的补全结果，因为扩散目标优先考虑多样性而非点对点准确性。

核心挑战在于如何为去噪器提供稳定的数值输入，同时又不将填补后的缺失值视为地面真值（ground-truth），从而避免常数占位符带来的偏差和自插补带来的不稳定性。

方法论：AugMask

作者提出了 AugMask，这是一个即插即用的训练框架，通过将**条件化（conditioning）与监督（supervision）**分离，使缺失不感知（missing-unaware）的骨干网络能够适应不完整数据。该方法分为两个截然不同的步骤：

1. 随机条件增强（输入构建）

AugMask 不使用确定性的占位符，而是通过从上下文相关的分布 $T(\cdot | x_{obs}, m)$ 中采样缺失条目来构建一个全数值输入。

机制： 对于每个缺失特征 $j$ ，从以观测特征和缺失掩码为条件的分布中抽取一个值 $Z_j$ 。
实现： 作者利用了一种轻量级的分解方法，使用每特征独立的 LightGBM 模型。对于连续特征，他们拟合了一个条件均值回归器和一个对数方差回归器，以便从异方差高斯分布中采样。对于类别特征，他们使用了带有标签平滑的条件分类器。
作用： 这些采样值作为去噪器的不确定条件上下文（uncertain conditioning context），而非地面真值目标。这种随机性编码了缺失性固有的不确定性。

2. 仅观测值监督（损失掩码）

去噪损失仅应用于原始观测到的坐标。

目标： 模型最小化观测坐标的重建误差 $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$ ，其中 $x_{A,0}$ 是增强后的输入。
效果： 增强后的（先前缺失的）条目通过噪声输入 $x_{A,t}$ 进入模型以提供上下文，但模型永远不会因无法重建它们而受到惩罚。这防止了模型过度依赖可能不准确的补全结果。

理论证明

论文将这一训练规则与 Rao–Blackwellized 目标函数 联系起来。通过对可能的补全进行边缘化处理，作者表明该目标函数产生了一个方差加权的敏感性惩罚（variance-weighted sensitivity penalty）。

解释： 目标函数局部近似中的二阶项与增强的条件方差成正比。高条件方差会惩罚模型对该特定补全结果的敏感性，从而有效地抑制模型对不确定填充值的依赖。低方差则允许补全结果更直接地作为上下文使用。
对比 AugFull： 作者将 AugMask 与 “AugFull”（对包括增强项在内的所有坐标应用损失）进行了比较。他们证明了如果增强规则设定错误，AugFull 会引入偏差，因为模型会被迫拟合潜在错误的伪目标。AugMask 通过将补全严格视为上下文，避免了这一问题。

核心贡献

一种针对不完整表格扩散的训练原则： 引入了 AugMask，它通过随机增强（条件化）与观测值仅损失掩码（监督）的分离，实现了两者的解耦。
一种边缘化视角： 通过理论分析表明，随机增强会诱导一个方差加权的敏感性惩罚，从而正则化模型，使其防止过度依赖不确定的补全。
一种即插即用的方案： 一个轻量级框架，使标准的缺失不感知评分扩散模型无需改变架构即可在不完整数据上进行训练。

实验结果

作者在多种数据集（Adult, Default, Shoppers, Magic, Beijing, News）和缺失机制（MCAR 和 MAR）下评估了 AugMask，缺失率范围从 0.1 到 0.9。

性能： AugMask 使标准的缺失不感知扩散模型（如 CDTD, TabDiff）能够超越专门的缺失感知基线模型（如 MIWAE, MissDiff, DiffPuter, ForestDiff）。
生成质量： 使用 AugMask 训练的模型在样本级保真度（通过 $\alpha$ -Precision, $\beta$ -Recall 衡量）和特征级统计特性（Trend, Shape）方面均取得了更优的表现。
机器学习效用： 在使用 AugMask 模型生成的合成数据进行下游任务（分类和回归）训练时，其表现优于基线模型。
鲁棒性： 随着缺失率的增加，AugMask 的优势变得更加显著。在高缺失率（如 90%）的情况下，AugMask 避免了其他方法出现的严重性能退化。
敏感性分析： 该方法对增强规则的误设定具有鲁棒性。即使在扰动增强规则（例如降低随机性或引入偏差）时，AugMask 仍保持了相对于 AugFull 的性能优势，证实了将补全视为上下文而非目标的做法是关键因素。

重要性与主张

论文声称 AugMask 为解决全指定扩散骨干网络与不完整现实数据之间的不匹配问题提供了原则性的解决方案。通过将随机补全视为不确定上下文而非训练目标，该方法：

避免了由常数占位符引入的稀疏性偏差和人工模式。
避免了自插补方法中的优化不稳定性和移动目标动态。
利用增强的方差来正则化模型，使其对不确定输入不那么敏感。

作者将 AugMask 定位为一种通用策略，允许现有的、高性能的缺失不感知模型在需要处理缺失数据的场景中部署，而无需复杂的架构修改或专门的插补损失函数。该工作侧重于生成；作者指出，将其扩展到潜在扩散流水线（latent diffusion pipelines）或特定于插补的目标仍是未来的研究方向。