反思与优化：在蛋白质结合物设计中重新审视约束优化

背景

在当前的机器学习领域，尤其是当预训练模型（Pre-trained Models）成为主流时，许多从业者往往忽视了底层优化理论的重要性。作者坦言，虽然熟悉 Adam、AdaGrad 或 L-BFGS 等常用优化器，但对于涉及对偶空间（dual spaces）和连续函数收敛性等经典优化理论，往往感到陌生甚至觉得“过时”。

然而，在处理一个具体的科学计算问题——**从头设计蛋白质结合物（de-novo binder design）**时，作者发现“第一印象”往往不是最佳方案。这个问题本质上是一个带有约束条件的非凸优化问题：我们需要在一个概率单纯形（probability simplex）上寻找一个概率分布，以最小化某个复杂的折叠模型（如 Alphafold）输出的评分函数。

这篇文章旨在通过这一具体案例，对比两种处理约束优化的主流方法：重参数化（Reparameterization）与投影梯度下降（Projected Gradient Descent, PGD），并深入探讨在高维空间中这两种方法的差异及其对最终解的影响。

核心内容

问题设定

假设我们有一个包含 $K$ 个类别的分类概率分布，其概率由向量 $\mathbf{p}$ 表示。为了使其成为有效的概率分布，必须满足以下约束：

1. 归一化：所有概率之和为 1，即 $\sum p_i = 1$。
2. 非负性：所有概率必须大于 0，即 $p_i > 0$。

我们的目标是最小化一个非凸函数 $f(\mathbf{p})$，该函数接收概率向量并输出一个实数值。假设我们可以计算梯度 $\nabla f(\mathbf{p})$，但计算梯度和函数值本身在计算上是昂贵的。

生物学背景：在从头设计蛋白质结合物的场景中，$\mathbf{p}$ 代表氨基酸的概率分布，$f$ 是一个折叠模型（如 Alphafold）。输入到 $f$ 的是一个位置特异性评分矩阵（PSSM），每一列是一个概率向量。我们的目标是找到能产生最佳折叠结构（根据 ipSAE 等指标衡量）的序列。

第一种尝试：重参数化（Reparameterization）

面对带约束的优化问题，最直观的想法是通过重参数化消除约束，从而使用标准的无约束优化方法。

我们可以引入对数几率（logits）$\mathbf{z}$，并通过 Softmax 函数将无约束的 $\mathbf{z}$ 映射为有效的概率向量 $\mathbf{p}$： $$ \mathbf{p} = \text{softmax}(\mathbf{z}) $$

由于 Softmax 的性质，无论 $\mathbf{z}$ 取何值，输出的 $\mathbf{p}$ 始终满足归一化和非负性约束。因此，问题转化为寻找最优的 $\mathbf{z}^*$，使得 $f(\text{softmax}(\mathbf{z}))$ 最小化。随后，我们可以对 $\mathbf{z}$ 计算梯度并运行梯度下降。

这种方法在概念上简单直接，也是机器学习中常见的做法。然而，作者指出，这并非唯一的选择，且不同的变换可能会因为非凸性导致完全不同的局部最优解。

第二种尝试：概率单纯形与投影梯度下降（PGD）

让我们回到约束本身。满足归一化和非负性约束的所有 $\mathbf{p}$ 构成的空间被称为概率单纯形（Probability Simplex），记为 $\Delta^{K-1}$。

例如，当 $K=3$ 时，2-单纯形是一个垂直于向量 $(1,1,1)$ 的三角形区域。单纯形的顶点对应着确定性分布（如 $[1,0,0]$），而面上的点则表示某些类别的概率为零。

如果我们直接在单纯形上运行梯度下降 $\mathbf{p}_{new} = \mathbf{p} - \eta \nabla f(\mathbf{p})$，通常会遇到以下问题：

1. 破坏归一化：如果梯度不与单纯形平面垂直，直接步进会导致概率之和不再为 1。
2. 破坏非负性：即使归一化保持正确（例如通过减去梯度均值来修正方向），步进后的点仍可能落在单纯形外部，导致某些概率变为负数。

为了解决这个问题，我们引入投影梯度下降（Projected Gradient Descent, PGD）。其核心思想是：先执行一步标准的梯度下降，然后将结果投影回可行域（即单纯形）上。

$$ \mathbf{p}{new} = \Pi{\Delta} (\mathbf{p} - \eta \nabla f(\mathbf{p})) $$

其中 $\Pi_{\Delta}$ 是将点投影到单纯形上的算子。投影操作旨在找到距离当前梯度步进点最近的、满足约束的单纯形上的点。虽然计算投影在数学上并不 trivial，但存在高效的算法（如 ICML 2008 年论文所述），其复杂度通常远低于梯度计算本身。

高维空间中的陷阱：稀疏性

在低维（如 2D）情况下，PGD 看起来非常直观。但在高维空间中，情况变得复杂且反直觉：

1. 脱离单纯形的概率增加：随着单纯形维度 $K$ 的增加，单纯形内部靠近边缘的点比例急剧增加。随机方向的一步梯度下降极大概率会跳出单纯形。
2. 投影到低维面：一旦跳出，投影操作会将点拉回到单纯形的边界（即至少有一个坐标为 0 的面）。
3. 解的稀疏化：随着优化的进行，点被投影到越来越低维的面上，导致解向量 $\mathbf{p}$ 变得越来越稀疏（即越来越多的概率值变为 0）。

模拟数据显示，从顶点出发，在 2D 中回到单纯形内部的概率是 $1/2$，在 3D 中约为 $1/6$，而在更高维度中几乎为 0。这意味着，在高维约束优化中，PGD 倾向于产生稀疏解。

镜像下降（Mirror Descent）的引入

作者最后提到，Softmax 重参数化之所以常见，是因为它隐含地使用了某种几何结构。为了更严谨地选择变换方法，我们需要引入**镜像（Mirror）**的概念。镜像下降是一种利用非欧几里得几何结构（通过凸函数定义的势函数）来指导优化路径的方法，它比简单的重参数化或投影提供了更丰富的理论框架来解释为何选择特定的变换函数。

关键要点

• 直觉陷阱：在处理带约束优化问题时，直接重参数化（如使用 Softmax）是常见且自然的，但它改变了目标函数的几何结构，可能导致不同的局部最优解。
• 投影梯度下降（PGD）：PGD 是一种概念上更简单的约束优化方法，通过“步进-投影”循环来确保解始终在可行域内。
• 高维稀疏性效应：在高维单纯形上进行 PGD 时，由于几何特性，优化过程会自然地将解推向边界，导致最终解具有高度的稀疏性（许多概率值为 0）。这在某些应用场景下可能是有益的，但在其他场景下可能是偏差来源。
• 计算效率：尽管投影操作在数学上复杂，但存在高效的算法，其计算成本远低于梯度评估，因此在实际应用中是可行的。
• 理论深度：简单的重参数化背后涉及复杂的几何选择（如镜像下降中的势函数），理解这些底层原理有助于在特定问题（如蛋白质设计）中选择更合适的优化策略。

意义与影响

这篇文章对于从事科学机器学习（Scientific ML）和组合优化问题的研究人员具有重要参考价值。

1. 方法论启示：它提醒研究者，不要盲目依赖标准的无约束优化技巧（如 Softmax + SGD）。在约束空间（如概率分布、混合权重、注意力机制）中，优化算法的选择会深刻影响解的性质（如稀疏性）。
2. 蛋白质设计应用：在从头设计蛋白质结合物时，解的稀疏性可能意味着模型倾向于选择少数几种氨基酸组合，这可能对应于更稳定或更特异的折叠结构。理解 PGD 的这一特性有助于更好地解释模型