🤖 machine learning

Calibrating Generative Models to Feature Distributions with MMD Finetuning

本文介绍了 kCGM，这是一种通过最小化最大均值差异（MMD）和 KL 正则化来校准生成模型以匹配目标特征分布的方法，有效地在抗生素、蛋白质和 DNA 生成等多种任务中，在保持样本有效性的同时提升了特征对齐度。

原作者： Nathaniel L. Diamant, Brian L. Trippe

发布于 2026-06-19

📖 1 分钟阅读☕ 轻松阅读

原作者： Nathaniel L. Diamant, Brian L. Trippe

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

想象一下，你拥有一位大师级厨师，他极其擅长烹饪各种各样的菜肴。这位厨师从海量的食谱库（预训练模型）中汲取了经验，几乎每次都能做出看起来美味且合理的佳肴。然而，如果你要求这位厨师做一种特定类型的菜，比如“正宗的意大利面”，他可能会遇到困难。他可能会做出看起来很棒、味道也不错的意面，但却缺少定义“真正意大利面”的特定质地、酱汁浓稠度或香料比例。

在人工智能领域，这是一个常见的问题。生成式模型（如这位厨师）可以创造出单个看起来真实的个体，但如果你观察它们生产出的整批产品，其特征分布（整体风格、尺寸或化学组成）往往会偏离你真正想要的目标。

这篇论文介绍了一种名为 kCGM（核校准生成模型，kernel Calibrating Generative Models）的新方法，旨在解决这种偏移问题，同时又不破坏厨师原有的天赋。

问题所在：为什么“直接模仿”行不通

如果你想让你的 AI 厨师做出更好的抗生素（一种药物），一个简单的方案是直接给他们看一份由 174 种真实抗生素组成的名单，并说：“多做一些这类东西。”

论文指出，这就像是要求学生逐页背诵一本教科书。

过拟合（Overfitting）： AI 可能会仅仅死记硬背这 174 个特定分子，从而停止了创造力。它变成了一台复印机，而不是一个生成器。
关注点错误： 它可能会模仿错误的东西。它可能会学习模仿分子的形状，却丢失了使其成为药物的化学特性。
破坏模型： 在实验中，当研究人员尝试这种“直接模仿”（直接微调）的方法来处理抗生素时，AI 开始生成“无效”分子——即在现实世界中根本无法存在的化学结构。它为了匹配名单而牺牲了质量。

解决方案：kCGM（“试味”校准法）

kCGM 并不强迫 AI 去死记硬背目标名单，而是像一个聪明的品鉴师，通过对比 AI 作品的整体风味特征与目标之间的差异来进行工作。

以下是它的运作方式，使用了几个类比：

1. 特征映射（“特征菜单”）
你不需要向 AI 展示实际的分子。你只需要告诉它哪些特征是重要的。

对于药物，你可能关心“亲脂性”（它有多油）。
对于蛋白质，你可能关心“螺旋”与“折叠”形状的比例。
对于 DNA，你可能关心序列在特定细胞类型中的活性。
这些特征就像是一份特征菜单。AI 不需要看到实际的目标分子，它只需要看到目标群体在平均意义上的“菜单”长什么样。

2. MMD（“距离测量尺”）
论文使用了一个名为最大均值差异（Maximum Mean Discrepancy, MMD）的数学工具。你可以把它想象成一把测量两团点云之间距离的尺子。

想象 AI 当前的输出是一团蓝色的点。
目标抗生素是另一团红色的点。
MMD 测量的是这两团点在形状和分布上的距离，而不仅仅是它们的中心点。
创新之处： 之前的研究方法（如 CGM）仅试图匹配云团的中心（平均值）。而 kCGM 试图匹配整个云团的形状。它确保 AI 不仅仅是在制造“平均水平”的抗生素，而是制造出一种多样化的组合，看起来完全符合真实的分布。

3. 分数函数（“推力”）
由于 AI 并不总是能“看到”这些特征背后的数学逻辑（有些特征是黑盒，例如复杂的化学指纹），kCGM 使用了一个分数函数估计器（score-function estimator）。

想象一位蒙着眼睛的艺术家正在尝试临摹目标。
他看不见目标，但他会得到一个“分数”，告诉他他的画作在多大程度上接近目标的风格。
kCGM 根据这个分数计算出一个“推力”，轻轻地将 AI 的下一次尝试推向目标分布，而不会强迫它去死记硬背特定的实例。

4. KL 正则化（“安全网”）
这一点至关重要。当你推动 AI 时，你不希望它彻底忘记如何烹饪。

kCGM 包含一个“安全网”（KL 正则化），它会说：“向目标移动，但不要离你原本的高质量风格太远。”
这可以防止 AI 为了达到目标数值而生成毫无意义的内容（无效分子）。

实验结果如何？

研究人员在三个不同的“厨房”中测试了这种“品鉴师”方法：

抗生素（小分子）：
- 结果： 当他们尝试让 AI 生成抗生素时，“直接模仿”的方法破坏了 AI，产生了许多无效的化学结构。
- kCGM： AI 生成了有效的、符合化学原理的分子，完美匹配了真实抗生素的统计分布。它在提升匹配度的同时保持了高质量。
蛋白质（折叠形状）：
- 结果： AI 之前生成的蛋白质几乎全是单一形状（α-螺旋），缺失了自然界中的多样性。
- kCGM： 它成功地引导 AI 产生多样化的蛋白质形状组合，比之前的方法更接近自然界的真实情况。
DNA（调控序列）：
- 结果： AI 在生成能激活特定细胞中特定基因的 DNA 序列方面遇到了困难。
- kCGM： 它校准了 AI，使其产生的 DNA 序列具有针对特定细胞类型的正确“活性特征”，即使用于衡量这些特征的工具非常复杂且属于“黑盒”（由另一个 AI 预测）。

核心结论

论文声称 kCGM 是微调生成模型的一种更优方式。它不是强迫模型去死记硬背一小份示例列表（这会破坏模型），而是温柔地引导模型去匹配目标群体的统计指纹。

它就像一个校准旋钮：你可以旋转它，使 AI 的输出更贴近你的目标，但“安全网”确保了 AI 永远不会失去创造有效、高质量样本的能力。即使当你关注的特征非常复杂、不可微或者来自“黑盒”工具时，它依然有效。

技术摘要：通过 MMD 微调实现生成模型的特征分布校准 (kCGM)

问题陈述
生成模型生成的样本往往在个体层面看似合理，但在关键特征的分布上存在偏差。例如，在一个广泛化学空间上预训练的模型可能会生成与特定治疗类别（如抗生素）特征分布显著不同的分子。通过在目标集上进行直接微调来纠正这一问题存在问题：这会导致对小数据集的过拟合风险，需要完整的目标样本（而这些样本可能很稀缺），并且无法指定应匹配哪些特定特征。此外，现有方法在处理离散值或黑盒特征（例如 Morgan 指纹）时往往表现不佳，因为这些特征是不可微的。本文指出，目前缺乏能够仅利用特征级监督、在不需要特征映射或采样过程可微的情况下，将生成模型校准至目标特征分布的算法。

方法论：kCGM
作者引入了核校准生成模型 (kernel Calibrating Generative Models, kCGM)，这是一种通过最小化生成特征分布与目标特征分布之间的最大均差异 (Maximum Mean Discrepancy, MMD)，并辅以与预训练模型 Kullback-Leibler (KL) 散度正则化的微调方法。

目标函数： 该方法优化一个由两部分组成的损失函数 $L_\theta$ ：
1. MMD 损失： 模型诱导特征分布 $P$ （其中 $y = h(x)$ 且 $x \sim p_\theta$ ）与经验目标分布 $Q$ （来自目标特征 $y^{targ}_{1:N}$ ）之间平方 MMD 的无偏有限样本估计量。
2. KL 正则化： 一个 $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$ 项，用于确保微调后的模型仍与预训练分布保持接近，从而防止灾难性遗忘或模式崩溃。
核函数选择： 不同于以往可能依赖于仅匹配均值的点积核方法，kCGM 使用了定义在概率分布上的度量核（仅当分布完全相同时为零）。
- 对于连续欧几里得特征，使用能量距离核 (energy-distance kernel)。
- 对于离散位向量特征（如 Morgan 指纹），采用 Tanimoto/Jaccard 核。
- 这些选择避免了引入额外的带宽超参数。
梯度估计： 由于特征映射 $h$ 和采样过程可能是不可微的，kCGM 采用了得分函数梯度 (score-function gradients)（类似于 REINFORCE 方法）。
- 留一法 (Leave-One-Out, LOO) 基线： 这是一个关键的技术贡献，作者推导出了针对 MMD 损失的新型 LOO 基线。标准的 LOO 基线是对批次中其他样本的系数求平均，但 MMD 系数在样本间是耦合的。作者推导出了一个特定的“MMD-LOO”估计量，该估计量近似于如果某个特定样本未出现在当前小批量中时的平均系数，从而在不引入偏差的情况下降低了方差。
- 算法： 该过程包括从当前模型中采样（不进行梯度追踪）、计算特征向量、结合 LOO 基线计算 MMD 和 KL 系数，并通过得分函数梯度更新模型权重。

核心贡献

公式化： 本文将特征分布匹配公式化为仅使用目标特征样本的 KL 正则化 MMD 微调，从而无需原始样本空间中的完整目标样本。
无偏估计量： 作者推导出了 MMD 目标的无偏得分函数梯度估计量，使得使用不可微特征映射和不可重参数化采样器（如自回归模型、离散扩散模型）成为可能。
泛化性： kCGM 推广了 CGM 方法。虽然 CGM 仅匹配目标特征的均值，但 kCGM 通过 MMD 匹配完整的分布。文中指出，在 kCGM 中使用点积核可以还原 CGM-relax 目标。

实验结果
作者在三个领域评估了 kCGM：小分子生成、蛋白质结构生成和调节性 DNA 生成。

小分子抗生素：
- 任务： 校准 G2PT 自回归图模型，使其匹配 174 种已知抗生素的特征分布。
- 特征： Morgan 指纹（局部子结构）、Murcko 骨架（拓扑结构）和理化描述符。
- 发现： 与 CGM 和直接微调相比，kCGM 在匹配目标特征分布（通过 MMD 衡量）与远离预训练模型之间实现了更优的权衡。至关重要的是，当直接在小型抗生素集上进行微调时，会导致高比例的化学无效分子；而 kCGM 在改善特征匹配的同时，还增加了化学有效性。kCGM 还展示了跨特征迁移能力，即仅在指纹特征上进行训练时，也能提高骨架和描述符的匹配度。
蛋白质结构多样性：
- 任务： 微调 Genie 2（一种蛋白质扩散模型），以匹配 CATH 数据集的二级结构分布（ $\alpha$ -螺旋 vs $\beta$ -折叠）。
- 发现： 预训练模型主要产生 $\alpha$ -螺旋。kCGM 与目标分布的不一致性比 CGM 减少了约 2 倍。kCGM 在 KL 权重 $\lambda$ 方面表现出平滑的权衡曲线，而 CGM 的性能对 $\lambda$ 的敏感度较低，导致优化困难。
调节性 DNA 生成：
- 任务： 微调离散扩散模型，以生成具有现实细胞类型特异性活性谱（由 AlphaGenome 预测）的 DNA 序列。
- 发现： kCGM 成功匹配了 47 种不同顺式调节主题的完整活性谱分布，表现优于仅匹配均值的 CGM。微调后的模型生成的序列具有与真实增强子相似的活性谱，而预训练模型生成的序列活性较低。

意义与主张
论文声称 kCGM 提供了一种鲁棒的替代方案，特别是在目标数据集较小或仅需保留分布的特定方面时。其主要意义在于能够：

仅通过特征级监督将生成模型对齐到用户指定的特征分布。
处理科学领域中常见但与标准基于梯度的微调不兼容的黑盒及不可微特征（如指纹、预测器输出）。
在纠正分布失配的同时，保持或提高样本质量（有效性），这是直接微调常见的失败模式。
适用于多种架构，包括自回归模型、连续空间扩散模型和离散扩散模型。

作者也指出了局限性，例如由于 MMD 假设检验中的维度灾难，该方法在处理极高维特征时可能效果不佳，以及需要可处理的采样轨迹对数概率（不包括流匹配模型）。他们建议未来的工作可以扩展到仅存在分布约束、而没有原始生成空间中完整样本的情境。

问题所在：为什么“直接模仿”行不通

解决方案：kCGM（“试味”校准法）

实验结果如何？

核心结论

技术摘要：通过 MMD 微调实现生成模型的特征分布校准 (kCGM)

类似论文