Spectral Feedback for Test-Time Alignment of Protein Diffusion Models¶

会议: NeurIPS 2026
arXiv: 2609.30456
OpenReview: https://openreview.net/forum?id=PhVaSdnP4b
领域: 计算生物
关键词: 蛋白质逆折叠, 离散扩散模型, 测试时对齐, 稀疏傅里叶变换, 反馈编辑

一句话总结¶

针对离散扩散模型在蛋白质逆折叠中单向生成的不可逆缺陷，提出利用生物上位性导致的稀疏傅里叶结构学习编辑集合价值函数，以测试时谱反馈循环迭代重掩码并重采样，大幅提升生成蛋白质的稳定性。

研究背景与动机¶

蛋白质逆折叠（Inverse Protein Folding）的核心任务是根据给定的目标蛋白质骨架几何结构设计能够正确折叠的氨基酸序列。近年来，以 ProteinMPNN 架构为代表的离散扩散模型（Discrete Diffusion Models）已成为该任务的主流生成框架。然而在实际蛋白质工程中，科学家往往不仅要求序列折叠回目标骨架，还希望其满足诸如热力学稳定性（ΔΔG）等高级生化性质。现有对齐方法要么依赖强化学习微调（如 DRAKES），受限于高昂的重训练开销和分布偏移风险；要么采用单向推理引导（如 Token-level Guidance 或树搜索），将采样过程视为单向轨迹，一旦生成过程中确定了不良残基，便缺乏可回溯修正的机制。

利用离散扩散模型的重掩码（re-masking）与重采样构建编辑反馈循环，类比于连续图像扩散中的加噪编辑，提供了一种黑盒修正策略。然而其核心矛盾在于：编辑哪些位置构成了一个组合优化难题。修改一个氨基酸残基的影响高度依赖于同时被编辑的其他残基，即编辑效果具有强烈的非独立交互性；若在 $2^L$ 的候选子集空间中进行穷举搜索，在计算上是完全不可行的。以往方法为了规避组合爆炸，往往被迫假设各位置独立（如基于单位置重要性或随机掩码），从而削弱了整体对齐效果。

生物序列中广泛存在的上位性（Epistasis）表明，多个氨基酸协同突变的联合效应通常由极少数强交互组主导，而非均匀分散在所有可能子集中。本文的核心切入点正是将这一生物物理洞察映射至编辑集合价值函数，发现其在布尔傅里叶基底下的展开同样呈现高度稀疏性。核心 idea：将蛋白质离散扩散模型的测试时对齐转化为编辑子集价值函数的组合优化问题，利用生物上位性带来的傅里叶稀疏性，通过亚指数样本复杂度的稀疏谱恢复与整数规划精确求解最优编辑位置，实现模型无关的迭代自修正闭环。

方法详解¶

整体框架¶

Spectral Feedback 采用外层反馈循环与内层谱优化相结合的两阶段范式。对于给定的目标骨架 $y$ 和当前氨基酸序列 $x$，算法不改动底层扩散模型的任何权重，而是将模型视为黑盒条件采样器。整体流水线主要包括：在当前序列上对编辑子集进行伯努利掩码采样并评估短步预期奖励、利用 ProxySPEX 拟合编辑集合价值函数的低阶稀疏傅里叶近似、将谱表示转换为带势约束的 0-1 整数线性规划并求解全局最优编辑子集 $S^*$、以及对最优子集对应残基施加重掩码并执行逆向扩散步完成局部重采样。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始全掩码序列<br/>逆向扩散生成候选序列 x"] --> B["编辑子集随机采样与单步价值评估<br/>评估集 rewards"]
    B --> C["ProxySPEX 稀疏傅里叶恢复<br/>拟合低阶谱系数"]
    C --> D["Möbius 基整数线性规划<br/>求解最优编辑子集 S*"]
    D --> E["局部重掩码与扩散重采样<br/>更新序列 x"]
    E -->|未达到迭代预算 N-1| B
    E -->|完成对齐| F["输出高稳定性蛋白质序列"]

关键设计¶

1. 价值函数设计与编辑子集采样：在不展开全反向轨迹的前提下准确感知局部编辑收益 直接对每个候选编辑子集运行完整扩散反向轨迹计算真实奖励在计算上难以承受。为此，算法针对不同的底层采样机制定义了单步逆向扩散价值函数：当对基础扩散模型或 RL 微调模型进行常规采样时，采用平均价值函数 $f_{avg}(S, x_t, y) = \frac{1}{n} \sum_{i=1}^n r(z_i)$，衡量重采样后生成蛋白质的预期稳定性；当底层采用 Best-of-N 等贪心策略时，则采用极大值价值函数 $f_{max}(S, x_t, y) = \max \{r(z_1), \dots, r(z_n)\}$，避免因非最优样本的方差拉低潜力子集的估值。在此基础上，算法通过参数为 $\gamma$ 的独立伯努利试验并行抽取 $D$ 个候选编辑掩码，并借助仅更新单步反向扩散的状态转移概率进行快速模拟，在极低计算开销下构造用于谱估计的样本集。

2. 基于 ProxySPEX 的稀疏傅里叶谱恢复：破解高阶残基协同效应的组合爆炸 将编辑集合价值函数 $f: 2^{[L]} \to \mathbb{R}$ 投影到布尔正交傅里叶基底： $$f(S) = \sum_{T \subseteq [L]} (-1)^{|S \cap T|} F(T)$$ 生物物理中的能量相互作用决定了傅里叶系数 $F(T)$ 是高度稀疏且能量集中于低阶交互项的（实测 $R^2 \ge 0.7\sim 0.84$，且主导项阶数 $|T| \le 3$）。由于传统的 LASSO 或高阶穷举法在特征维度随阶数指数爆炸，本文采用基于梯度提升树（GBT）的 ProxySPEX 算法，以亚指数级的样本复杂度高效捕获主导傅里叶支撑集 $\mathcal{A} = \{T \subseteq [L] : \hat{F}(T) \neq 0, |T| \le d\}$ 及其系数 $\hat{F}(T)$。该设计彻底摆脱了传统启发式编辑对残基独立性的僵化假设，精确刻画了多个残基间的协同上位性效应。

3. Möbius 基变换与 0-1 整数线性规划求解：将连续谱函数无损转化为离散约束优化 即便获得了稀疏傅里叶近似函数 $\hat{f}(S)$，在基数约束 $|S| \le k$ 下直接寻找全局最优编辑子集仍是 NP 难问题。本文将傅里叶基底通过变量代换精确变换至 Möbius 包含基： $$\hat{M}(T) = (-2)^{|T|} \sum_{S \in \mathcal{A}, S \supseteq T} \hat{F}(S), \quad \hat{f}(S) = \sum_{R \in \mathcal{A}^+, R \subseteq S} \hat{M}(R)$$ 并将集合指示变量转化为二进制决策变量 $x \in \{0,1\}^L$ 与单项式连乘变量 $y_R = \prod_{i \in R} x_i$。通过引入向下闭包 $\mathcal{A}^+$ 的子集激活约束（$y_R \le y_Q, \forall Q \subset R$）与去激活约束（$\sum_{i \in R} y_{\{i\}} \le |R| - 1 + y_R$），将原问题重写为紧凑的整数线性规划（ILP），变量总数仅为 $s \cdot 2^d + n$（其中 $s = |\mathcal{A}|$ 为稀疏度，$d$ 为最高阶数）。借助 Gurobi 分支割平面法，算法可在毫秒级时间内精确锁定最优编辑集合 $S^*$。

4. 局部逆向扩散重采样：利用部分加噪先验加速生成收敛 选定最优编辑子集 $S^*$ 后，算法构造局部掩码序列 $\tilde{x}_i = \text{MASK}$（若 $i \in S^*$）而保留非编辑残基不变。由于大部分上下文结构已知，重采样过程无需从完全无信息的 $t=T$ 噪声先验启动，而是利用部分加噪序列直接进行条件逆向扩散。这一设计不仅大幅削减了逆向步数和推理延迟，同时严格保留了骨架约束区域的已有优良构象，避免全局重抽样带来的构象崩塌。

损失函数 / 训练策略¶

Spectral Feedback 纯粹作为推理时（Test-Time）对齐算法运行，底层离散扩散模型与奖励模型（如基于 Megascale 训练的 $\Delta\Delta G$ 预测器）的参数完全冻结。内层 ProxySPEX 采用浅层决策树集成（叶子节点数 30~50，学习率 0.01，$\lambda=0.0001$）进行超参数交叉验证；主循环迭代预算设为 $N-1$ 次，典型设置下仅需 $2\sim 5$ 次反馈迭代即可实现超越深层采样的收敛效果。

实验关键数据¶

主实验¶

在 Megascale 数据集的多项蛋白质逆折叠骨架上，以 $\Delta\Delta G$ 作为稳定性对齐与评估预言机（$\Delta\Delta G > 0$ 表示比野生型更稳定），同时借助 ESMFold 结构预测评估自洽性 RMSD（scRMSD，要求 $< 2\text{Å}$），并用 ProtGPT2 对数似然评估序列自然度。

基础模型 / 配置	对齐 $\Delta\Delta G$ ↑	评估 $\Delta\Delta G$ (avg) ↑	scRMSD (avg) ↓	scRMSD < 2Å (%) ↑	对数似然 (Naturalness) ↑
Pretrained (基线)	0.10	-0.49	1.10	91.8%	-147.8
Pretrained + Spectral	0.53 (+0.43)	0.33 (+0.82)	1.06 (-0.04)	94.2% (+2.4%)	-146.5 (+1.3)
BoN-10 (基线)	0.45	0.09	1.09	92.4%	-148.1
BoN-10 + Spectral	0.70 (+0.25)	0.86 (+0.77)	1.07 (-0.02)	96.7% (+4.3%)	-147.5 (+0.6)
DRAKES (RL 微调基线)	0.69	0.99	1.17	93.4%	-157.5
DRAKES + Spectral	0.83 (+0.14)	1.38 (+0.39)	1.53 (+0.36)	85.8% (-7.6%)	-157.7 (-0.2)

消融实验与编辑策略对比¶

对比不同编辑子集选择策略在一轮反馈迭代下的对齐增益与延迟权衡（以基数 $k=20$ 为约束）：

编辑位置选择策略	交互阶数建模	优化机制	样本扩展性 / 延迟特性	对齐收益与行为特征
Gradient-Weighting	仅一阶	基于嵌入梯度显著性得分贪心选取	极快，但忽略多残基协同	收益有限，容易陷入局部次优
Token Exclusion	仅一阶	逐残基孤立掩码评估单点收益	随序列长度线性扩展	无法捕获残基组合非线性效应
Argmax Selection	任意阶	直接选取采样库 $D$ 中最高奖励样本	依赖穷举覆盖度，上限受制于采样方差	无法跳出采样池发现稀有最优子集
Hill Climbing	高阶逐步	串行 1-flip 随机局部搜索	严格串行化，无法利用 GPU 并行批量打分	延迟极高，大规模采样下耗时爆炸
Spectral (ProxySPEX)	自适应 1~3 阶	稀疏傅里叶恢复 + Möbius ILP	高并发批采样 + 极速整数求解	在采样预算扩展下收益最高、延迟最优

关键发现¶

超越 Best-of-N 边际递减效应：在预训练模型上，Spectral Feedback 仅需 2 轮迭代即可达到 Best-of-10 的稳定性水平，5 轮迭代达到 Best-of-50 水平，且能突破 Best-of-N 随样本规模增加而饱和的收益瓶颈。
上位性结构与蛋白质构象强相关：谱分析证实不同骨架的高阶交互强度差异显著。富含 $\alpha$-螺旋的 2KVV 骨架受局部氢键主导，表现出强烈的一阶效应（LASSO 与 ProxySPEX 表现接近）；而复杂环区交织的 R6-560 骨架具有显著的高阶多残基协同（ProxySPEX 比一阶 LASSO 最终奖励高出 24%）。
DRAKES 上的结构一致性折损根源：DRAKES + Spectral 在 $\Delta\Delta G$ 上达到了全场最高的 1.38，但 scRMSD 退化至 1.53。分析表明这并非算法缺陷，而是由于 DRAKES 本身在极高奖励区域引入了偏离野生型分布的假阳性伪影，需引入多目标正则化来平衡结构稳定性。

亮点与洞察¶

将生物上位性转化为算法稀疏先验：论文敏锐地洞察到生物体突变适应度景观的低阶稀疏特性同样适用于扩散掩码价值函数，巧妙地将高维离散组合选择规约为稀疏正交基重构。
Möbius 整数线性规划映射：利用 Möbius 展开将复杂的非线性傅里叶奇偶符号变换为单调的子集覆盖指标，从而能使用成熟的商业 ILP 求解器毫秒级求出全局最优解，工程闭环极其严谨。
黑盒即插即用泛化性：不改变反向扩散的任何数学公式和网络权重，既可作用于预训练基础模型，也可无缝叠加在推理搜索（Best-of-N）或强化学习微调模型之上。

局限与展望¶

价值函数估值的随机方差：目前的价值评估基于单步反向扩散随机模拟，非确定性导致谱拟合的 $R^2$ 受限，后续可训练专用的部分掩码序列价值回归网络来提供低方差预测。
高奖励区域的预言机过拟合风险：在单一 $\Delta\Delta G$ 驱动下，算法易将序列推向牺牲折叠精度（scRMSD 升高）的病态区域，未来需结合 Pareto 多目标优化同时平衡序列自然度与构象一致性。
向更长序列与复杂模态的迁移挑战：当序列长度达到数百残基或词表显著扩充时，采样规模 $D$ 与 ILP 约束规模将进一步攀升，需探索更轻量级的分块谱分解机制。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将布尔傅里叶稀疏谱分析引入离散扩散模型的测试时对齐与重掩码决策。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 Megascale 真实骨架、三类底层模型基线、详尽的谱阶数分析与消融对比。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰，生化上位性与组合优化算法结合紧凑，图表表现力强。
价值: ⭐⭐⭐⭐⭐ 为离散扩散模型在生物计算以及通用离散序列自修正提供了一种普适且具理论支撑的新范式。