Calibrating Generative Models to Feature Distributions with MMD Finetuning¶
会议: NeurIPS 2026
arXiv: 2606.19496
代码: https://github.com/smithhenryd/cgm
领域: 优化/理论
关键词: 生成模型校准, 最大均值差异, 得分函数梯度, 留一法控制变量, 生物分子生成
一句话总结¶
本文提出核校准生成模型方法(kCGM),利用无偏核最大均值差异(MMD)估计器配合专为耦合核项设计的 \(O(M^2)\) 留一法梯度基线(MMD-LOO)与预训练 KL 正则,在仅依赖非可微黑盒目标特征样本的条件下校准自回归与扩散生成模型,在抗生素小分子、蛋白质二级结构和调控 DNA 序列生成上全面超越均值匹配基线 CGM 与直接监督微调。
研究背景与动机¶
在大规模宽泛数据集上预训练的生成模型虽然能够生成单个看似合理的样本,但其在关键下游特征上的群体分布往往与特定科学目标集合存在严重的分布失准(distributional miscalibration)。例如,在包含 130 万药用小分子的 GuacaMol 数据集上预训练的自回归分子图模型 G2PT,其生成的化学官能团与骨架分布与已知抗生素相去甚远;蛋白质结构扩散模型 Genie 2 受采样启发式策略影响,生成的样本几乎全部坍缩为 \(\alpha\)-螺旋结构,无法重现 CATH 天然蛋白质结构域丰富的二级结构比例分布;而在增强子 DNA 序列上训练的离散掩码扩散模型,也往往生成染色质开放活性偏低的序列,难以复现特定细胞状态下的真实调控活性分布。
面对这一分布对齐需求,现有的生成模型微调范式陷入了明显的两难困境。如果直接在小规模目标样本集(如仅含 174 个已知抗生素的数据集)上进行监督微调(Direct Finetuning),模型不仅要求必须获取原始样本空间中的完整样本,无法由用户指定只对齐哪些目标属性,还会因在小样本上逐分子模仿而严重过拟合,导致生成分子的化学有效性(chemical validity)从预训练的 87% 暴跌至 67%。另一类基于强化学习的奖励微调方法(如 RLHF、DPO)旨在最大化单样本标量奖励的期望,容易诱发模式坍缩,无法直接约束生成群体的整体特征分布形状。最接近的前人工作 CGM(Calibrating Generative Models)虽然引入了得分函数梯度以摆脱对特征提取器可微性的依赖,但它的松弛版本(CGM-relax)仅约束生成特征的均值与目标均值一致,完全忽略了高阶矩与多模态结构;即便通过多维经验累积分布函数(CDF)分箱来近似分布匹配,分箱数量也会随特征维度呈指数级爆炸,在超过三维的特征空间中完全不可行,且在二维蛋白质任务上仍残留较大的失准误差。
这一系列矛盾的核心在于:如何在不要求特征提取器或采样过程可微、且无需高维网格分箱的前提下,直接利用目标特征样本把生成模型的特征分布完整对齐到目标分布。本文的切入角度是将再生核希尔伯特空间中的最大均值差异(MMD)引入黑盒得分函数优化:一方面利用特征核(如连续空间的能量距离核与离散指纹空间的 Tanimoto 核)无偏度量两分布的高阶差异;另一方面攻克二次型 MMD 损失中多样本强耦合导致的梯度高方差难题。核心 idea:将生成模型的特征分布校准建模为带预训练 KL 正则的核最大均值差异(MMD)最小化问题(kCGM),并通过对称核展开与精确重算的 \(O(M^2)\) 耦合留一法基线(MMD-LOO)构造无偏、低方差的得分函数代理梯度,从而仅凭非可微黑盒特征样本即可高效校准自回归、连续扩散与离散扩散模型。
方法详解¶
整体框架¶
kCGM 的目标是微调预训练生成模型 \(p_{\theta_{\mathrm{pre}}}(\boldsymbol{x})\) 的参数 \(\theta\),使得从微调模型采样的生成轨迹 \(\boldsymbol{x} \sim p_\theta\) 在经过任意黑盒、非可微特征映射 \(\boldsymbol{y} = \boldsymbol{h}(\boldsymbol{x})\) 后,其诱导的特征分布 \(P_{\boldsymbol{y}}\) 能够紧密匹配由外部目标特征样本 \(\boldsymbol{y}^{\mathrm{targ}}_{1:N}\) 定义的目标分布 \(Q\)。在每一轮迭代中,算法首先在不追踪梯度的条件下从当前策略采样一批完整生成轨迹,随后并行计算每个样本相对于预训练模型的轨迹对数似然比以及黑盒特征核矩阵;接着利用留一法基线分别对 KL 惩罚项和耦合的核差异项做去中心化方差削减,组合成单样本的停止梯度权重系数;最后通过加权轨迹对数似然构造代理损失,执行无偏的得分函数梯度更新。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["预训练生成模型与目标特征样本<br/>无梯度采样M条完整生成轨迹"] --> B["非可微特征分布的核最大均值差异(MMD)建模<br/>提取黑盒特征并计算自排斥项与目标吸引项"]
B --> C["无偏得分函数梯度与耦合留一法基线(MMD-LOO)<br/>由核矩阵行和以O(M²)重算留一法中心化系数"]
C --> D["KL正则化与代理损失分步累积更新<br/>融合KL-LOO系数并沿生成轨迹分步反传梯度"]
D --> E["校准后的生成模型<br/>对齐目标特征分布且保持样本质量与有效性"]
关键设计¶
1. 非可微特征分布的核最大均值差异(MMD)建模:用无带宽超参的特征核替代均值匹配与指数分箱
针对 CGM 仅能匹配特征均值或受困于高维 CDF 分箱维度灾难的痛点,kCGM 直接在特征空间上采用再生核 \(k(\cdot, \cdot)\) 定义的平方最大均值差异 \(\operatorname{MMD}_k^2(P, Q)\) 来度量生成特征 \(\boldsymbol{y}_{1:M}\) 与目标特征 \(\boldsymbol{y}^{\mathrm{targ}}_{1:N}\) 的分布距离。给定当前批次生成的 \(M\) 个特征样本与 \(N\) 个目标特征样本,并引入可调的自排斥权重 \(\alpha > 0\)(默认 \(\alpha=1\)),其有限样本无偏估计量由模型内部自排斥项与模型-目标吸引项构成:
其中第一项惩罚生成样本在特征空间中的过度聚集以抑制模式坍缩,第二项鼓励生成特征向目标样本靠拢,\(C_{\mathrm{targ}}\) 为仅依赖目标样本的常数项。论文在理论上严格证明(Proposition A.1 与 A.2),若选用平凡的内积核 \(k(\boldsymbol{u},\boldsymbol{v}) = \boldsymbol{u}^\top \boldsymbol{v}\),上述 \(\widehat{\operatorname{MMD}}_k^2\) 在相差目标常数的意义下精确退化为仅匹配经验目标均值 \(\bar{\boldsymbol{y}}^{\mathrm{targ}}\) 的 CGM-relax 目标 \(\left\|\mathbb{E}[\boldsymbol{y}] - \mathbb{E}[\boldsymbol{y}^{\mathrm{targ}}]\right\|_2^2\),从而将 CGM 统一为 kCGM 的退化特例。为了真正捕捉全分布的高阶结构且避免引入高斯 RBF 核难以调节的带宽超参数,kCGM 针对连续欧氏特征采用无超参数的能量距离核(Energy-Distance Kernel)\(k(\boldsymbol{u},\boldsymbol{v}) = \frac{1}{2}(\|\boldsymbol{u}\|_2 + \|\boldsymbol{v}\|_2 - \|\boldsymbol{u}-\boldsymbol{v}\|_2)\),针对小分子二值指纹特征则采用化学领域标准的 Tanimoto 核 \(k(\boldsymbol{u},\boldsymbol{v}) = \frac{\boldsymbol{u}^\top \boldsymbol{v}}{\|\boldsymbol{u}\|_2^2 + \|\boldsymbol{v}\|_2^2 - \boldsymbol{u}^\top \boldsymbol{v}}\)。此外,当 \(\alpha < 1\) 时,减弱自排斥项会在离散类别空间中导出等价于单纯形投影(simplex projection)\(P_i^\star = (Q_i - \tau_\alpha)_+ / \alpha\) 的阈值锐化效应,适度向高密度主模态集中,在二维蛋白质二级结构等复杂多模态分布上进一步改善拟合精度。
2. 无偏得分函数梯度与耦合留一法基线(MMD-LOO):在 \(O(M^2)\) 时间内消除多样本耦合核的梯度方差
由于分子指纹提取、二级结构注释(Biotite)以及基因活性预测器(AlphaGenome)等特征映射 \(\boldsymbol{h}(\cdot)\) 包含不可微操作或离散采样过程,传统的重参数化梯度完全失效。kCGM 转而利用生成轨迹 \(\boldsymbol{x}_m\) 的对数概率 \(\log p_\theta(\boldsymbol{x}_m)\) 推导得分函数梯度(REINFORCE)。利用核函数对称性 \(k(\boldsymbol{y}_m, \boldsymbol{y}_{m'}) = k(\boldsymbol{y}_{m'}, \boldsymbol{y}_m)\) 与独立同分布样本的可交换性,模型内部交叉项关于参数 \(\theta\) 的梯度恰好带有系数 2,从而可将多样本耦合的期望 MMD 梯度精确分解为每个生成样本的核差异系数 \(c_m^{\mathrm{k}} = \frac{2}{M}(\alpha s_m - r_m)\) 与各自轨迹得分梯度 \(\nabla_\theta \log p_\theta(\boldsymbol{x}_m)\) 的乘积之和,其中 \(s_m = \frac{1}{M-1}\sum_{m'\neq m} k(\boldsymbol{y}_m, \boldsymbol{y}_{m'})\) 为样本 \(m\) 与批次内其余生成样本的平均核相似度,\(r_m = \frac{1}{N}\sum_{n=1}^N k(\boldsymbol{y}_m, \boldsymbol{y}_n^{\mathrm{targ}})\) 为其与目标特征集合的平均核相似度。
然而,得分函数梯度天然伴随极高的方差,必须引入控制变量基线(baseline)。在标准单样本奖励微调中,留一法基线只需简单取批次内其余样本奖励的均值;但在 MMD 目标下,其余样本 \(j \neq m\) 的系数 \(c_j^{\mathrm{k}}\) 本身通过 \(s_j\) 包含了与样本 \(\boldsymbol{x}_m\) 的核交互项 \(k(\boldsymbol{y}_j, \boldsymbol{y}_m)\),直接对 \(c_j^{\mathrm{k}}\) 取平均会因为基线依赖于 \(\boldsymbol{x}_m\) 而破坏梯度的无偏性。为此,kCGM 构造了严格不依赖 \(\boldsymbol{x}_m\) 的耦合留一法基线(MMD-LOO):先将样本 \(m\) 从批次中剔除,计算剩余子批次 \(\boldsymbol{x}_{-m}\) 中任意样本 \(j \neq m\) 的重构系数,再对这 \(M-1\) 个重构系数取平均。若朴素地对每个被剔除样本重算核矩阵,时间复杂度将高达 \(O(M^3)\);本文巧妙地利用模型特征核矩阵 \(K_{\boldsymbol{y}\boldsymbol{y}}\) 的行和关系,直接从原相似度 \(s_j\) 中代数扣除样本 \(m\) 的贡献:
通过定义中心化核系数 \(c_m^{\mathrm{k,LOO}} = c_m^{\mathrm{k}} - b_m^{\mathrm{k}}(\boldsymbol{x}_{-m}, \boldsymbol{y}^{\mathrm{targ}}_{1:N})\),由于基线 \(b_m^{\mathrm{k}}\) 完全独立于 \(\boldsymbol{x}_m\),利用恒等式 \(\mathbb{E}_{\boldsymbol{x}_m}[\nabla_\theta \log p_\theta(\boldsymbol{x}_m)] = 0\) 可严格保证 \(\mathbb{E}[b_m^{\mathrm{k}} \nabla_\theta \log p_\theta(\boldsymbol{x}_m)] = 0\)(无偏性保持不变),同时仅需 \(O(M^2)\) 的矩阵行和运算即可大幅压低梯度方差。
3. KL 正则化与代理损失分步累积更新:防退化约束与跨模型架构的常数显存反传
为了防止生成模型在迎合目标特征分布时丢失预训练模型积累的化学或生物学先验语法,kCGM 在目标中引入相对于预训练模型的轨迹级 KL 散度惩罚 \(\lambda D_{\mathrm{KL}}(p_\theta \| p_{\theta_{\mathrm{pre}}})\)。对于每个无梯度采样的轨迹 \(\boldsymbol{x}_m \sim p_{\texttt{stop-grad}(\theta)}\),先计算单样本对数似然差 \(l_m = \log p_{\texttt{stop-grad}(\theta)}(\boldsymbol{x}_m) - \log p_{\theta_{\mathrm{pre}}}(\boldsymbol{x}_m)\),由于各样本的 \(l_m\) 彼此独立,可直接采用标准留一法均值进行去中心化得到 \(c_m^{\mathrm{KL}} = \frac{1}{M}\left(l_m - \frac{1}{M-1}\sum_{j\neq m} l_j\right)\)。随后,将停止梯度的 KL 系数与核差异系数线性叠加为总样本权重 \(c_m = \lambda c_m^{\mathrm{KL}} + c_m^{\mathrm{k,LOO}}\),构造如下标量代理损失(Surrogate Loss):
根据 Proposition A.3,对该代理损失直接关于 \(\theta\) 求导所得的 \(\sum_{m=1}^M c_m \nabla_\theta \log p_\theta(\boldsymbol{x}_m)\) 恰好是原正则化 MMD 目标的严格无偏梯度估计。更重要的是,这一形式将特征分布评估与模型反向传播彻底解耦:无论是自回归 Transformer(G2PT 中按序列 token 展开 \(\sum_{t=1}^T \log p_\theta(x_t \mid x_{<t})\))、连续扩散模型(Genie 2 中将 100 步 Euler-Maruyama 采样轨迹分解为每步高斯转移密度的对数和),还是离散掩码扩散模型(DNA 任务中按余弦解掩码调度分解为 50 步被解掩码位置分类交叉熵之和),都可以先在无梯度模式下采样完整轨迹并算出标量权重 \(c_m\),随后在重放(replay)轨迹时按时间步逐块累积 \(c_m \nabla_\theta \log p_\theta(\boldsymbol{x}_m^{(t)} \mid \boldsymbol{x}_m^{(t-1)})\) 的梯度,从而无需在显存中保存长达 50 至 100 步的完整采样计算图,在单张 L40S GPU 上即可完成全量微调。
损失函数 / 训练策略¶
kCGM 的总体优化目标结合了有限样本期望核 MMD 与预训练 KL 散度正则项:
在工程训练中,为了使最优学习率能够稳定迁移到不同的正则化强度 \(\lambda\) 和不同尺度的核函数上,论文采用了两项关键的数值缩放策略:第一,在微调开始前,先从预训练模型采样 5 个批次计算其相对于目标特征集的初始平均 \(\overline{\operatorname{MMD}^2}\),将核函数重缩放为 \(k'(\boldsymbol{u},\boldsymbol{v}) = 10 k(\boldsymbol{u},\boldsymbol{v}) / \overline{\operatorname{MMD}^2}\),使预训练模型的初始 \(\operatorname{MMD}^2\) 统一固定为 10;第二,使用保比例归一化权重 \(\frac{1}{1+\lambda}\mathbb{E}[\widehat{\operatorname{MMD}}_k^2] + \frac{\lambda}{1+\lambda}\operatorname{D}_{\mathrm{KL}}(p_\theta \| p_{\theta_{\mathrm{pre}}})\) 保持不同 \(\lambda\) 下总损失的绝对量级恒定(使所有方法在 G2PT 上均能统一使用 \(10^{-5}\) 的最优学习率)。此外,对于生成的化学无效分子图,特征映射统一将其映射为零向量 \(\boldsymbol{0}\),使其自然偏离真实抗生素特征分布并受到 MMD 损失的自动惩罚,从而在不额外设计有效性奖励的前提下显著提升生成分子的化学合法率。
实验关键数据¶
主实验¶
主实验在 100 万参数的自回归分子图生成模型 G2PT-small(预训练于 130 万 GuacaMol 类药分子)上展开,目标是校准到 174 个已知小分子抗生素的特征分布。实验分别在三种互补的特征空间中进行微调与交叉评测:① Morgan 指纹(256-bit 半径 2 局部子结构,配 Tanimoto 核);② 通用 Murcko 骨架(经 ChemNet 嵌入并 PCA 降至 64 维拓扑骨架,配能量距离核);③ 11 维分子描述符(分子量、\(\mathrm{logP}\)、TPSA、QED、合成可及性等理化性质,按目标集标准差归一化,配能量距离核)。下表展示了在匹配相同预训练偏离预算(\(\operatorname{D}_{\mathrm{KL}}(p_\theta \| p_{\theta_{\mathrm{pre}}}) = 25\),基于 9,600 个生成分子线性插值,括号内为 3 个随机种子的标准误 SEM)下,kCGM 与均值匹配基线 CGM 及直接监督微调(Direct Finetuning)的全面对比:
| 微调方法与目标特征 | Morgan 指纹 \(\operatorname{MMD}^2 \downarrow\) | Murcko 骨架 \(\operatorname{MMD}^2 \downarrow\) | 理化描述符 \(\operatorname{MMD}^2 \downarrow\) | 化学有效性 (Validity) \(\uparrow\) |
|---|---|---|---|---|
| Pretrained (G2PT-small 预训练基线) | 0.036 (0.0002) | 0.24 (0.003) | 1.1 (0.00031) | 87% (0.22%) |
| Direct Finetuning (174 抗生素直接微调) | 0.007 (0.00007) | 0.085 (0.0021) | 0.036 (0.00049) | 67% (0.093%) |
| CGM Morgan (均值匹配) | 0.014 (0.00019) | 0.15 (0.015) | 0.26 (0.012) | 92% (0.21%) |
| CGM Scaffold (均值匹配) | 0.037 (0.00083) | 0.045 (0.00052) | 1.2 (0.054) | 97% (0.13%) |
| CGM Descriptor (均值匹配) | 0.024 (0.00028) | 0.12 (0.0083) | 0.12 (0.0059) | 94% (0.29%) |
| kCGM Morgan (本文) | 0.005 (0.00017) | 0.039 (0.004) | 0.15 (0.007) | 95% (0.2%) |
| kCGM Scaffold (本文) | 0.038 (0.00073) | 0.028 (0.0009) | 1.2 (0.04) | 98% (0.098%) |
| kCGM Descriptor (本文) | 0.024 (0.00028) | 0.13 (0.009) | 0.027 (0.0042) | 96% (0.2%) |
消融实验¶
论文在 G2PT Morgan 指纹任务上对核函数选择(Tanimoto 核 vs. 退化为均值匹配的内积核)和耦合留一法基线(MMD-LOO)进行了严格消融(图 4),并在连续扩散蛋白质生成模型 Genie 2(校准至 CATH 二级结构比例分布)和离散掩码扩散 DeepMEL2 调控 DNA 生成(校准至 47 类细胞状态下的 AlphaGenome 预测染色质开放峰值分布)上验证了自排斥权重 \(\alpha\) 与跨模态泛化能力:
| 实验任务与模型架构 | 对比配置 / 变体 | 核心分布匹配指标与表现 | 关键结论与机制验证 |
|---|---|---|---|
| 小分子生成 (G2PT) 自回归图模型 + 256维二值指纹 |
完整 kCGM (Tanimoto + MMD-LOO) | \(\mathrm{KL}=25\) 时 Morgan \(\operatorname{MMD}^2 = \mathbf{0.005}\),有效性 95%,独特性 >98% | 同时实现最优分布对齐、高化学合法率与零记忆化过拟合 |
| 小分子生成 (G2PT) 核函数消融 |
w/ Dot-Product Kernel (内积核 + MMD-LOO) | 退化为均值匹配,全 \(\lambda\) 范围 \(\operatorname{MMD}^2\) 显著劣于 Tanimoto 核,但优于原版 CGM (0.014) | 证明捕捉高阶分布特性至关重要;同时验证 MMD-LOO 对均值匹配同样有直接增益 |
| 小分子生成 (G2PT) 方差削减消融 |
w/o MMD-LOO Baseline (去掉留一法基线) | 全 \(\lambda\) 曲线上 \(\operatorname{MMD}^2\) 大幅恶化,高正则化下收敛严重受阻 | 耦合核得分函数梯度方差极大,MMD-LOO 是保证优化成功的核心组件 |
| 蛋白质结构生成 (Genie 2) 连续扩散 (100步) + 2D二级结构比例 |
kCGM (能量距离核, \(\alpha \in \{0.25, 0.5, 1.0\}\)) vs. 99-bin CDF CGM | 2D KDE 对称 KL 误差比 99 分箱 CGM 降低约 \(2\times\)(\(\alpha=0.5\) 与 \(0.25\) 最优) | 无需二维网格分箱即可打破预训练全 \(\alpha\)-螺旋坍缩,且随 \(\lambda\) 平滑可控 |
| 调控 DNA 生成 (DeepMEL2) 条件离散掩码扩散 (50步, 500 bp) |
kCGM (能量距离核 + 32维 PCA 白化 AlphaGenome 特征) vs. CGM | 在全部 \(\lambda \in \{10^{-1}, 10^{-2}, 10^{-3}, 10^{-4}\}\) 及 47 个细胞状态主题中的绝大多数上取得更低测试集 \(\operatorname{MMD}^2\) | 成功驱动黑盒深度学习预测器(AlphaGenome)校准条件离散扩散模型,消除低活性序列偏差 |
关键发现¶
- 特征级校准能避免小样本直接微调的灾难性退化并提升有效性:在 174 个抗生素上直接做似然微调会使化学有效性从 87% 跌至 67%,因为模型强行记忆小样本中的局部序列模式;而 kCGM 在三种特征下分别将有效性提升至 95%–98%(同时保持 \(>98\%\) 的生成分子独特性),并将目标特征 \(\operatorname{MMD}^2\) 压缩至 \(0.005\)(Morgan)、\(0.028\)(Scaffold)和 \(0.027\)(Descriptor),在所有对应指标上均大幅超越 CGM 和直接微调。
- 不同粒度特征空间的跨特征迁移具有非对称性:从表 1 可见,针对局部子结构的 Morgan 指纹微调能够很好地顺带改善骨架(\(\operatorname{MMD}^2\) 从 0.24 降至 0.039)和理化描述符(从 1.1 降至 0.15);但反过来,仅针对拓扑骨架(Scaffold)微调几乎无法改善 Morgan 指纹(0.038)或理化描述符(1.2),说明细粒度局部化学基团分布隐式约束了宏观理化性质,而纯拓扑环系骨架丢弃了原子与键级信息。
- 自排斥权重 \(\alpha < 1\) 的阈值锐化有助于多模态边缘覆盖:在 Genie 2 蛋白质二级结构任务中,\(\alpha=1.0\) 已经全面超越 99 分箱的 CGM,而将自排斥权重下调至 \(\alpha=0.5\) 或 \(\alpha=0.25\) 能进一步降低与 CATH 目标分布的对称 KL 散度,验证了适度减弱模型内部自排斥有助于推动原本高度集中的全 \(\alpha\)-螺旋分布跨越低密度区向 \(\beta\)-折叠及混合结构模态迁移。
亮点与洞察¶
- 在核视角下统一均值校准与全分布校准:作者敏锐地证明了前人工作 CGM-relax 完全等价于采用线性内积核 \(k(\boldsymbol{u},\boldsymbol{v})=\boldsymbol{u}^\top \boldsymbol{v}\) 的 kCGM(Propositions A.1–A.2)。这一理论洞察不仅解释了为什么 CGM 无法区分均值相同但方差或多模态结构不同的特征分布,而且顺理成章地通过替换为无超参数的特征核(能量距离核与 Tanimoto 核)把校准能力从一阶矩推广到了完整概率分布。
- 精巧的 \(O(M^2)\) 耦合留一法基线(MMD-LOO)设计:在 REINFORCE 框架下优化成对交互的 U-统计量(如 \(\widehat{\operatorname{MMD}}^2\))时,直接套用常规留一法均值会因样本间核项耦合而引入偏差,若逐个样本重算子批次又会带来 \(O(M^3)\) 开销。作者利用核矩阵行和的代数减法 \(s_j^{(-m)} = \frac{(M-1)s_j - k(\boldsymbol{y}_j, \boldsymbol{y}_m)}{M-2}\),在不增加额外核矩阵计算的前提下实现了严格无偏的方差削减,这一技巧可直接复用到任何基于集合级成对奖励(如多样性奖励、核对比损失)的强化学习微调任务中。
- 轨迹级重放与代理损失解耦实现跨架构通用校准:通过先无梯度采样评估黑盒特征、再按停止梯度权重 \(c_m\) 沿生成轨迹逐步重放累积 \(\nabla_\theta \log p_\theta(\boldsymbol{x}_m)\),kCGM 在单张消费级/工作站级 GPU 上即可微调 100 步连续扩散和 50 步离散扩散模型,为接入冷冻电镜(cryo-EM)、RNA 化学探测等仅有特征分布测量而无完整分子样本的科学实验约束提供了通用接口。
局限与展望¶
- 高维特征空间中核 MMD 的统计检验力衰减:论文在讨论部分坦承,双样本非参数核 MMD 统计量的检验功效会随特征维度的增加呈多项式级下降(Ramdas et al., 2015)。因此当目标特征维度极高(如原始高维序列嵌入或多通道全基因组轨迹)时,直接计算欧氏或能量距离核的效果会退化,目前需依赖在目标集上预拟合 PCA 白化降维(如 DNA 实验中降至 32 维、骨架实验中降至 64 维)来维持校准效率。
- 依赖显式可计算的采样轨迹对数似然:kCGM 的得分函数梯度要求能够精确求取整个生成采样轨迹的对数概率 \(\log p_\theta(\boldsymbol{x})\)。对于确定性 ODE 采样器或未引入随机噪声调度的流匹配(Flow Matching)模型,转移概率退化为 Dirac \(\delta\) 分布,无法直接计算轨迹级得分函数梯度,需要额外引入随机伴随匹配或随机微分方程(SDE)转换才能适用。
- 未来可向无完整目标样本的纯实验分布测量拓展:当前实验中的目标特征分布仍是从一批真实分子或序列样本上提取得到的;未来的重要方向是直接对接仅有群体统计测量、缺乏单个完整生物分子真值的湿实验场景(如单分子化学探测、蛋白质组学荧光图像分布校准)。
相关工作与启发¶
- vs CGM / CGM-relax (Smith et al., 2026):CGM 只约束生成特征的均值等于目标均值,或者通过指数级膨胀的多维 CDF 分箱近似低维分布(超过 3 维即不可行);本文的 kCGM 用无偏核 MMD 替代均值与分箱约束,证明了 CGM-relax 只是 kCGM 在内积核下的特例,并在任意维度特征上实现了更紧致的分布对齐和 \(2\times\) 更低的蛋白质结构分布误差。
- vs 直接监督微调(Direct Supervised Finetuning):直接微调在原始分子空间最大化目标小样本集的对数似然,不仅无法选择性对齐特定属性,还会在 174 个抗生素这样的小数据集上严重过拟合(有效性跌至 67%);kCGM 仅在用户选定的特征空间施加分布约束并结合预训练 KL 正则,在精准对齐目标属性的同时将化学有效性提升至 95%–98%。
- vs 奖励微调与多样性 RL(Reward Finetuning / Miao et al., 2024):传统 RL 微调最大化单样本标量奖励,而 Miao et al. (2024) 虽在扩散模型中引入 MMD 信号,但仅用于最大化生成多样性且依赖启发式的单样本奖励拆分;kCGM 则针对外部指定的目标特征分布做精确匹配,并严格推导出无偏的有限样本耦合核梯度系数与 MMD-LOO 基线。
评分¶
- 新颖性: ⭐⭐⭐⭐ (将核 MMD 分布匹配与无偏耦合留一法得分函数梯度相结合,从理论上统一并推广了 CGM,方法简洁扎实)
- 实验充分度: ⭐⭐⭐⭐ (横跨小分子自回归图生成、蛋白质连续扩散、DNA 条件离散扩散三大生物分子模态,消融与特征迁移分析详尽)
- 写作质量: ⭐⭐⭐⭐⭐ (数学推导严密清晰,附录对内积核等价性、无偏代理梯度、\(O(M^2)\) 留一法基线和单纯形投影均给出了完整证明)
- 价值: ⭐⭐⭐⭐ (为黑盒不可微科学特征与预训练生成模型之间的分布级对齐提供了即插即用、低显存开销的通用微调工具)