RFdiffusion：把扩散模型搬到蛋白质结构上

English version: RFdiffusion: diffusion models on protein structure

简介

RFdiffusion（Watson et al., Nature 620, 1089–1100, 2023）常被介绍成”蛋白质设计领域的 Stable Diffusion”。这个类比对了一半：它确实是一个 DDPM，确实是”给定条件、从噪声里生成结构”，代码也确实在 GitHub 上以 BSD 协议开源。

但真正值得一个 ML 工程师读它的地方，恰恰是它和图像扩散模型不一样的那些部分：数据不在欧氏空间上、没有 FID 这样的现成指标、而且它的去噪网络不是从零训起来的——是把一个已经训好的判别式结构预测模型改造成了生成模型。

这篇笔记只谈 ML 的部分：问题怎么形式化、网络怎么搭、条件怎么注入、效果怎么评、以及哪些经验可以迁移到别的领域。生物学上的意义（为什么 de novo 设计蛋白质是件大事）我不是行家，就不展开了。

一、数据长什么样：不是像素，是一堆刚体

图像扩散模型的数据是 \(H \times W \times 3\) 的张量，加噪就是逐像素加高斯噪声，干净利落。蛋白质骨架不是这样。

RFdiffusion 把每个氨基酸残基表示成一个刚体坐标系（frame）：

一个平移：Cα 原子的三维坐标 \(x \in \mathbb{R}^3\)
一个旋转：由 N–Cα–C 三个原子确定的朝向 \(R \in SO(3)\)

先解释一下这里的原子名。蛋白质是一条氨基酸链，每个氨基酸（残基）都有同样的三个主链原子，按 N → Cα → C 的顺序一节节接下去：

       R₁              R₂              R₃        <- 侧链，20 种氨基酸的区别所在
       |               |               |
 - N - Cα - C --- N - Cα - C --- N - Cα - C -
            ||              ||              ||
            O               O               O

Cα（alpha carbon，α 碳）是每个残基的中心碳原子，上面接了四样东西：主链的 N、主链的 C、一个氢、以及决定这是哪种氨基酸的侧链。也就是说，它正好是”可变部分挂在不变主链上”的那个分叉点。结构生物学里历来就用 Cα 作为一个残基位置的代表（所谓 Cα trace），比较两个结构时算的 RMSD 默认也是在 Cα 上算的。坐标单位是埃（Å，0.1 nm），相邻两个 Cα 之间大约 3.8 Å，而且几乎是个常数。

一个好用的类比：这就是三维动画里的骨骼绑定。Cα 坐标是关节的位置，N–Cα–C 构成的三角形给出关节的朝向；主链相当于一根骨头长度固定的运动链（kinematic chain）。所以”每个关节一个位置 + 一个旋转”是最自然的状态表示。

为什么是 frame 而不是直接用一堆原子坐标？一方面，每个残基有 7～24 个原子，直接在全原子坐标上做扩散太重；另一方面，主链的键长、键角几乎是刚性的，真正自由变化的就是每个残基的位置和朝向，而朝向决定了侧链指向哪边、残基之间怎么堆叠接触。这也正是 AlphaFold2 和 RoseTTAFold 内部使用的表示——这一点在下一节会很关键。

先说旋转那一半。\(SO(3)\)（special orthogonal group，三维特殊正交群）就是三维空间里所有绕原点的旋转构成的集合，具体地说，是满足下面两个条件的 3×3 矩阵：

\[R^\top R = I, \qquad \det R = +1\]

第一个条件（正交）保证它只是”转”：长度和夹角都不变，不拉伸也不扭曲。第二个条件（行列式是 +1 而不是 −1）排除掉镜像翻转——这就是名字里 “special” 的含义。这对蛋白质不是个细节：蛋白质是手性的，照镜子照出来的不是同一个分子。

虽然写出来有 9 个数，但 \(R^\top R = I\) 本身带来 6 个约束，所以 \(SO(3)\) 实际只有 3 个自由度（欧拉角、轴角、单位四元数都是它的不同参数化）。它是一个群：乘法就是矩阵相乘，逆就是转置。还有一点待会儿会用到——它是紧致的，”所有可能的朝向”是一个有限大小的空间，因此上面存在均匀分布。

再把平移接上。一个”平移 + 旋转”的组合就是三维空间里的一次刚体运动，这些运动构成的群记作 \(SE(3)\)（special Euclidean group，三维特殊欧氏群），它的元素可以写成一个 4×4 矩阵：

\[\begin{pmatrix} R & x \\ 0 & 1 \end{pmatrix}, \qquad R \in SO(3),\; x \in \mathbb{R}^3\]

一共 6 个自由度（3 个转 + 3 个平移）。

所以一条长度为 \(L\) 的链，每个残基一个 frame，状态空间就是 \(L\) 份 \(SE(3)\) 的笛卡尔积 \(SE(3)^L\)，而不是 \(\mathbb{R}^{3L}\)。一个 150 残基的蛋白质，是这个 900 维空间里的一个点。

关键区别在于：\(SE(3)^L\) 是一个弯曲的流形（李群），不是向量空间。你没法做加法——两个旋转矩阵按元素相加，结果根本不是旋转矩阵；”走一小步”要靠流形上的指数/对数映射。这件事带来两个直接后果：

第一，噪声分布要跟着流形走。 平移部分可以照搬三维高斯噪声；但旋转部分你不能对旋转矩阵的九个元素加高斯噪声——加完就不是旋转矩阵了。论文的做法是在 \(SO(3)\) 流形上做布朗运动（即 IGSO(3) 分布，它是高斯分布在旋转流形上的对应物）。注意两边的终点还不一样：\(SO(3)\) 是紧致的，所以加噪到底时旋转收敛到所有朝向上的均匀分布；而平移那一半还是落回我们熟悉的高斯。前向加噪过程一共 200 步。

第二，网络必须是等变的。 把整个蛋白质在空间里旋转平移一下，它还是同一个蛋白质，网络的输出也应该跟着旋转平移。这就是为什么 RFdiffusion 依赖 SE(3)-equivariant transformer 层（装环境时那个要单独 conda 安装的 SE(3)-Transformers 就是干这个的）。图像领域可以靠数据增强糊弄过去的对称性，在这里是硬约束。

一个有用的直觉：噪声分布的设计，是由数据流形的几何结构决定的，不是由方便程度决定的。 当你的数据是角度、旋转、单纯形上的概率分布、或者离散 token 的时候，”加高斯噪声”这一步就需要重新想。

二、去噪网络：微调一个预训练的结构预测模型

这是整篇论文我认为最重要的一个决策。

RFdiffusion 的去噪网络不是新设计的 U-Net，而是直接拿 RoseTTAFold（和 AlphaFold2 同期的结构预测网络）改的，架构改动极小，然后用预训练权重初始化、在去噪任务上微调。

论文做了一个很干脆的消融：同样的训练时长，从预训练权重微调 vs 从随机初始化训练，结果是前者”far more successful”（见原文 Fig. 2F）。

为什么这件事有意思？因为 RoseTTAFold 原本是个判别模型：输入序列、输出结构。而 RFdiffusion 要的是生成模型。两者的训练目标完全不同，但共享同一件知识：什么样的三维构型才像一个真实的蛋白质。二级结构怎么堆叠、疏水核心怎么填充、链怎么折回来——这些”先验”已经被压在 RoseTTAFold 的权重里了，生成模型没必要从头再学一遍。

这和 CV/NLP 里”拿预训练 backbone 当初始化”是同一个套路，但跨越的是判别↔生成的边界，比常见的迁移学习更激进一些。如果你手头有一个在某个领域训得很好的预测模型，而又想做这个领域的生成，这个思路值得先试。

三、预测干净结构而不是噪声，而且 loss 不做对齐

两个细节，都和”扩散模型的标准做法”不一样：

预测目标。 图像扩散里主流是预测噪声 \(\epsilon\)（或 v-prediction）。RFdiffusion 每一步直接预测最终的干净结构 \(x_0\)，然后朝这个预测的方向走一步、再加一点噪声，作为下一步的输入。这个选择和它复用 RoseTTAFold 是配套的——结构预测网络天生就是输出一个完整结构的，让它预测”噪声”反而别扭。

损失函数。 训练用的是预测 frame 和真实结构之间的 MSE，而且不做对齐（without alignment）。这一点论文专门说明了，因为结构预测领域的标准损失是 FAPE（frame-aligned point error），FAPE 对全局参考系是不变的。但在扩散模型里，这个不变性反而有害：

你希望相邻时间步之间的全局坐标系是连续的。如果 loss 允许模型每一步都把整个分子”随便摆在哪”，去噪轨迹就会在全局刚体变换上乱跳，无法收敛成一条稳定的轨迹。

换句话说：在结构预测里对称性是要利用的，在扩散轨迹里对称性是要打破的（至少在时间维度上要锚定）。这种”同一个不变性，在不同任务里一个是特性一个是 bug”的情况，挺值得记住。

四、Self-conditioning

RFdiffusion 用了 self-conditioning：第 \(t\) 步去噪时，不仅输入当前的带噪结构，还输入上一步自己预测出来的 \(x_0\)。论文说这一项带来了显著的性能提升。

这个技巧来自 Chen et al. 的 Analog Bits，作者也指出它和 AlphaFold2 的 recycling 机制非常像——都是把网络自己的中间输出再喂回去，让后续计算在一个更好的起点上进行。

对工程实现来说，self-conditioning 的代价是训练时要随机地跑两次前向（一次不带条件产生预测，一次带上这个预测算 loss），换来的是推理轨迹的连贯性。性价比通常很高，而且几乎不改架构。

五、条件注入：蛋白质领域的 “ControlNet 时刻”

无条件生成一个蛋白质没什么用，真正的需求都是带约束的。RFdiffusion 的条件化手段相当丰富，基本可以和图像领域一一对应：

RFdiffusion	图像领域的类比	做法
Motif scaffolding	Inpainting	固定一段已知的功能片段（坐标+序列），让模型生成其余部分把它”托住”
Binder design + hotspots	带区域提示的条件生成	在复合物数据上微调，额外输入”目标分子上希望结合的残基”
Symmetry	硬约束投影	每一步去噪后显式地重新对称化，把输出投影回对称子空间
Secondary structure / block adjacency	Layout / ControlNet	额外输入拓扑描述（哪几段是螺旋、哪些段互相接触），控制生成的折叠类型
Guiding potentials	Classifier guidance	在采样时加一个外部势能项（吸引/排斥、回转半径等）去推轨迹
Partial diffusion (`partial_T`)	SDEdit / img2img	把已有结构只加到中间时间步再去噪，围绕已知骨架做多样化

几点观察：

对称性是用投影实现的，不是学出来的。 给定点群和链长，先为一个单体生成随机起始 frame，再复制 \(n-1\) 份按点群摆好；每一步去噪之后再显式重新对称化。这是把硬约束交给几何、而不是指望网络自己学会——在约束明确的场景下，这几乎总是更可靠的工程选择。
不同的条件走不同的路径：有的是改输入（motif、hotspot、二级结构），有的是改采样过程（potentials、对称化），有的只是改初始化（partial diffusion）。设计条件生成系统时，这三个口子都值得考虑，而不是一上来就想着”再微调一版模型”。

六、Pipeline：生成之后，用另一个模型来验收

这是整个工作流里最有 ML 味道的一环。RFdiffusion 只生成骨架（backbone），不产生氨基酸序列。完整流程是三段式：

RFdiffusion          ProteinMPNN              AlphaFold2
  骨架生成     ->      逆折叠/序列设计    ->     重新预测结构 -> 比对
 (SE(3) 扩散)      (每个骨架采样 ~8 条序列)     (pAE、RMSD 过滤)

RFdiffusion 在条件约束下生成三维骨架。
ProteinMPNN 做”逆折叠”：给定骨架，设计出可能折叠成这个形状的氨基酸序列，通常每个骨架采样 8 条。
AlphaFold2 拿这些序列重新预测结构，然后和原始设计比对。

论文里 in silico 成功的判据是三条同时满足：AF2 的 pAE < 5、整体骨架 RMSD < 2 Å、功能位点骨架 RMSD < 1 Å。

这个”self-consistency“指标值得单独说一下。生成模型最难的从来不是生成，而是评测——图像有 FID、有人眼，蛋白质两样都没有。RFdiffusion 的解法是用一个独立的、在不同任务上训练的模型来做裁判：如果我设计的骨架真的是”可实现的”，那么从它设计出的序列，应该能被一个独立的结构预测器折叠回原来的形状。这本质上是一个 cycle-consistency / round-trip 指标。

这个思路可以迁移，但有两个需要清醒的地方：

裁判本身是个神经网络，就有被”刷分”的风险。 优化一个可微分（或可搜索）的指标，长期看总会滑向对抗样本的方向——”AF2 觉得很好”不等于”在真实世界里成立”。论文最终靠的是实验验证：他们表达并测试了数百个设计，其中一个结合流感血凝素的设计，冷冻电镜解出的结构和设计模型几乎完全一致。体外实验才是这个 pipeline 真正的 test set。
它是一个必要非充分条件，用来做过滤（把 99% 不靠谱的候选挡掉、只把少数送去昂贵的下游验证）非常划算，用来当作最终目标函数则很危险。

做任何生成系统时都可以问自己一句：有没有一个独立的、方向相反的模型，能把我的输出”还原”回输入？如果有，你就有了一个廉价的自动评测指标。

七、推理期的旋钮

两个在实践中常用、也很有 ML 共鸣的参数：

时间步数 diffuser.T：训练用 200 步，但推理可以大幅缩短——仓库的默认值是 50，论文里提到轨迹甚至可以压到 ~15 步。和图像扩散里”训练 1000 步、推理用 DDIM 20 步”是同一个现象。
噪声尺度 noise_scale_ca / noise_scale_frame：在反向过程中调低甚至去掉额外噪声。论文报告这个做法在 23 个 benchmark 问题里有 17 个提升了成功率；代价是生成多样性下降。这几乎就是 GAN 的 truncation trick、或者 LLM 的低温采样在扩散模型上的对应物——用多样性换样本质量，在”只要有少数几个能用”的场景里是笔划算的买卖。

八、它能做什么

论文验证过的任务，按条件类型归类：

无条件 / 拓扑受限的单体设计——给定长度或折叠类型，生成全新的蛋白质骨架。
蛋白质结合物（binder）设计——给定靶点和希望结合的位点，设计一个能贴上去的新蛋白。这是目前应用最广的一类，直接关系到药物和诊断试剂。
对称寡聚体——按指定点群（环状、二面体、四面体等）生成多链组装体。
功能位点支架化（motif scaffolding）——把一个已知的功能片段（酶活性位点、金属结合位点、表位等）固定住，设计出一个把它稳定呈现出来的新骨架。

后续工作沿着同样的框架继续扩展，比如 2025 年的抗体设计版本。

小结

抛开生物学，RFdiffusion 对 ML 从业者的几条启发：

预训练的判别模型，可以当生成模型的去噪器用。 领域知识藏在权重里，不分判别还是生成；消融实验表明这比从头训练好得多。
噪声过程要尊重数据的几何。 数据在什么流形上，前向过程就该在什么流形上定义（\(SO(3)\) 上的布朗运动，而不是对矩阵元素加高斯噪声）。
同一个不变性，在不同任务里可能是特性也可能是 bug。 FAPE 的全局不变性利于结构预测，却会破坏扩散轨迹的连续性。
把硬约束交给几何，而不是交给网络去学。 对称性靠每步投影实现，比训练一个”希望它学会对称”的模型可靠得多。
条件注入有三个口子：改输入、改采样过程、改初始化。
没有 FID 的时候，找一个反向模型做 round-trip 验证。 它是很好的过滤器，但别把它当成最终目标——真正的 test set 在模型之外。

参考

论文：De novo design of protein structure and function with RFdiffusion，Nature 620, 1089–1100 (2023)
预印本（开放获取）：bioRxiv 2022.12.09.519842
代码：RosettaCommons/RFdiffusion（BSD 协议）
ProteinMPNN / RoseTTAFold
Self-conditioning 的出处：Analog Bits (Chen et al., 2022)