今天和大家分享一篇2026年发布在arXiv上的工作“A-CODE: Fully Atomic Protein Co-Design with Unified Multimodal Diffusion”。这篇文章想解决的,是蛋白生成模型里一个老问题:结构和序列到底能不能真正放在同一个生成过程中一起设计,而不是先生成骨架,再用另一个模型去补序列。
目前很多性能很强的蛋白设计方法,本质上仍然是两阶段流程:第一步用扩散或流匹配生成骨架或全原子结构,第二步再用ProteinMPNN之类的逆折叠模型给结构配序列。这个流程很实用,但问题也明显:结构和序列被拆开了,模型很难直接学习它们之间的联合分布,尤其是在侧链、原子类型和非天然氨基酸这些更细的层面。
A-CODE的思路更激进:不再直接预测残基类型,而是在atom14表示里同时去噪连续的原子坐标和离散的原子名称。最后,残基身份不是模型显式输出的标签,而是由预测出的原子名称集合反推出来。换句话说,它把蛋白co-design从“残基级序列+坐标”推进到了“原子类型+原子坐标”的统一生成。
1 引言
A-CODE是一个全原子的蛋白结构-序列共设计模型。它在统一的多模态扩散框架中,同时处理两类不同模态:连续的原子坐标,以及离散的原子类型。与传统两阶段方法不同,A-CODE不先生成结构再补序列,而是在每一步降噪中同步更新原子坐标和原子类型。
在无条件蛋白设计任务中,A-CODE在designability和codesignability上都超过已有co-design与two-stage模型。相较第二强的La-Proteina,A-CODE的designability/codesignability提升约6-9%,同时推理速度实现2.5-12倍加速。
在binder设计任务中,A-CODE第一次让co-design模型接近甚至超过部分two-stage SOTA模型。相较已有co-design baseline Protpardelle,A-CODE在10个靶点上都有明显提升,在部分hard targets上success rate甚至提升约10倍。更重要的是,由于它直接预测原子,不被20种标准氨基酸的残基标签所硬约束,因此可以自然扩展到非天然氨基酸建模。
2 为什么需要co-design
如图1所示,作者把三类模型放在一起比较。AlphaFold3是结构预测模型,给定氨基酸序列后,通过扩散提取原子坐标;AF3-like模型通常把序列设计和结构生成拆开,先生成结构或骨架,再预测序列;而A-CODE则把带噪原子坐标和掩码原子一起作为输入,在同一个扩散中同步提取。
这个区别很关键。传统两阶段流程相当于假设联合分布可以拆成p(X)和p(S|X):先从结构分布里采样X,再根据这个结构补一个S。但真实蛋白不是这样拼起来的,侧链、局部几何、主链走向等因素本来就是同时互相限制的。A-CODE直接建模p(X,S),目标就是减少这种结构-序列解耦带来的误差传递。
表1也把A-CODE的定位说得很直接:MultiFlow是one-stage,但不是all-atom;Protpardelle等方法虽然是all-atom或one-stage,但不支持ncAA modeling;A-CODE同时满足all-atom、one-stage co-design和ncAA modeling三个条件。这也是作者认为它能作为统一全原子生成框架的原因。
3 方法
A-CODE最重要的设计,是把残基级预测改成原子层面的预测。作者采用atom14表示,每个氨基酸最多用14个heavy atom位置表示。骨架原子N、CA、C、O永远存在,因此固定;侧链位置则用XPA0-XPA9(输入时的掩码)这类mask tokens初始化,让模型在扩散过程中逐步预测它们到底是什么原子。
这里的原子不是单纯元素类型,而是带有化学位置语义的名称,例如CG、CD、ND1、OE1等。也就是说,模型不是只知道这里是氮或这里是氧,而是知道这个原子在氨基酸侧链拓扑里的位置。最终去掉dummy atoms(DMY ,输出占位)后,只要原子集合和某种氨基酸匹配,就能反推出残基类型。
如图2所示,整个采样过程从随机带噪坐标和一串XPA mask atom tokens开始。每一步 diffusion都同时做两件事:一边用EDM solver更新连续坐标,一边用discrete diffusion更新原子类型。最终得到干净原子坐标、原子类型和由原子名称反推出来的序列。这个机制让序列身份不再是额外解码器的输出,而是全原子生成过程自然涌现出的结果。
4 训练与采样
A-CODE同时生成原子类型和原子坐标,听起来很直接,但真正难点在侧链。侧链如果太早定型,很容易在早期骨架还没稳定时就导致碰撞,反过来干扰全局。作者因此提出side-chain denoising lag,也就是侧链生成滞后。
这个策略可以理解成先骨架,后侧链。具体做法是:骨架原子按标准噪声流程去噪,而侧链原子采用更高噪声的滞后性,论文中使用eta=1.5。也就是说,在采样早期,侧链被故意保留在更模糊的状态,让模型优先把骨架拓扑搭好,之后再逐渐细化侧链几何。
这个技巧的价值在后面的实验里非常明显。去掉side-chain lag后,A-CODE的all-atom codesignability从81.2%掉到72.6%,Cα codesignability从87.0%掉到77.6%。
这说明结构序列的联合生成的瓶颈并不只是序列能不能预测对,更在于不同模态在时间上的收敛节奏是否合理。而今年6月有一篇发布在arXiv的工作它的核心就是学习了一种模态感知的时间耦合方式,其提出了GeoCouping,A-CODE和其都在反对同步生成,都在强调生成节奏要分层或耦合;不过GeoCouping则是结构和序列之间的时间耦合,而A-CODE是骨架与侧链之间的去噪滞后,一个是模态级问题,一个是原子结构内部问题。
5 无条件蛋白生成
在无条件生成任务中,作者按照La-Proteina的设置,在长度100、200、300、400、500上各采样100条蛋白。评价重点包括codesignability、designability、secondary structure composition和diversity。
表2的结果很直接:A-CODE的all-atom codesignability达到81.2%,C-alpha codesignability达到87.0%,PMPNN@1 designability为91.6%,PMPNN@8达到97.0%。相比第二强的La-Proteina w/ tri,A-CODE在codesignability上仍然明显领先。
速度上,表3更能说明问题。在单张A100-80G上,长度500的样本,A-CODE需要14.6秒,而La-Proteina w/ tri需要182秒,MultiFlow也需要61秒。也就是说,A-CODE不是靠堆算力换性能,而是在保持高designability的同时,把采样效率也做了上去。
图3则从氨基酸组成角度看生成分布。A-CODE相对PDB的total variation distance为0.219,优于La-Proteina的0.297。更有意思的是,A-CODE对HIS、MET、TRP、CYS这类低频氨基酸的恢复更好,说明原子名称预测结合位置权重可能在隐式平衡氨基酸分布上起到了作用。
6 条件binder设计
Binder design是更能体现模型实用价值的任务。作者采用PXDesign同样的10个靶点,包括BHRF1、H1、IL17A、IL7RA、IR、PDL1、SC2RBD、TNFa、TrkA和VEGFA,并使用AF2-IG filter评估designable samples的比例。
A-CODE在这里给出了两种模式:一种是two-stage用法,即先用A-CODE生成结构,再用ProteinMPNN重新设计序列;另一种是真正的A-CODE co-design,直接使用模型生成的序列。结果显示,作为two-stage model使用时,A-CODE(PMPNN)在10个任务中有5个拿到最佳;作为co-design model时,它在全部10个靶点上都显著超过Protpardelle。
最值得注意的是hard targets。比如SC2RBD和TNFa上,A-CODE co-design甚至超过了最强two-stage版本;对比Protpardelle,则能看到近一个数量级的提升。作者想强调的也正是这一点:全原子共设计不只是概念上更统一,在binder设计这种复杂条件任务上也已经开始变成真正有竞争力的路线。
7 非天然氨基酸建模:A-CODE最有潜力的扩展点
A-CODE这篇文章最值得探究的外延,是非天然氨基酸建模。因为模型不直接预测 残基标签,而是预测原子元素,所以理论上不需要把输出空间限制死在20种标准氨基酸内。作者从PDB中筛出含ncAA的结构,得到约23k条单链和29k个binder pairs,再在这些数据上微调模型。
图4a显示,A-CODE对top 10高频ncAA的side-chain geometry建模整体不错,大多数side-chain RMSD低于1.0Å,说明模型不是只生成一个占位”,而是能恢复相当合理的侧链几何。图4b中,生成ncAA到靶点的最小距离主要集中在2-4Å,这正是生物物理相互作用常见的距离范围。
界面分析也很关键。图4c显示,32%的生成界面ncAA参与氢键,15%形成盐桥,说明模型确实会利用ncAA的特殊化学性质来稳定复合物,而不是把它们随机放在表面。图4d的例子也很直观:无条件生成设置中,两个methyllysine被放在类似组蛋白修饰的可及loop区域;binder设计设置中,hydroxycysteine出现在界面,延长侧链并额外形成两个氢键。
不过,非天然氨基酸(ncAA)的种类无穷无尽,A-CODE难道为每一种全新的化学结构都建了一个独立的原子词表吗?答案是:没有。如果来一个ncAA就加一套新原子名,词表会瞬间爆炸。 况且很多ncAA在训练集里只出现过几次,模型根本学不好。
为了解决这个痛点,作者在附录(Algorithm 2&3)中展示了一个极为巧妙的折中方案——不扩充原子词表,而是把ncAA的原子名字,强行“翻译(映射)”到已有的标准氨基酸原子体系中(总共仅保留40类atom names):
第一步:四级原子名称映射(Algorithm 3)
当模型遇到一个陌生的ncAA(例如包含:N, CA, C, O, CB, CG, CD, OE3, C10, P)时,它会按照以下“降级匹配”规则把原子名洗一遍:
保留原名:如果原子名本身就在标准词表中,直接保留(如N, CA, C, O, CB, CG, CD)。
同类映射:如果名字不对,但元素和位置能对上,就映射到同类标准原子(如OE3在标准库没有,但有OE类,于是可以映射为OE1)。
元素映射:如果位置也对不上,但元素类型在,就退一步只按元素映射(如C10→映射为某个标准C类侧链原子)。
未知标记:如果连元素都没有对应的了,直接标注为UNK(如磷原子P→UNK)。
第二步:CCD拓扑逆向识别
既然原子名都被洗成了“标准名”,模型在主流程里怎么知道它到底是谁?
-
主流程:A-CODE的主流程并不会直接识别具体的ncAA名字。它先把无法完美匹配成20种标准氨基酸的残基直接统一标记为UNK(如Algorithm 2所示)。
-
后处理:在生成结束后,模型再利用CCD进行拓扑匹配,通过看原子的连接关系,把一部分UNK重新反向识别回具体的ncAA名字。
这种处理方式,优缺点也是极为明显:它完美避开了词表爆炸以及数据稀缺的问题,最重要的问题是它不需要修改底层模型架构就能直接喂入包含ncAA的数据进行微调;但这也并不是一种严谨的化学空间建模。它本质上是把复杂ncAA简化近似成了标准氨基酸的命名体系,而在处理某些结构过于离奇或者说带有大量不合规元素时,这可能会带来一定程度的损失。
这部分结果还很初步,但意义很大。过去很多蛋白生成模型遇到ncAA,要么直接忽略,要么把它当配体处理,丢失了残基层面的连续性。A-CODE则提供了一条更自然的路径:在原子空间中统一处理天然和非天然氨基酸。
8 讨论与展望
A-CODE的核心价值,不只是几个benchmark数字更高,而是它明确推动了一个方向:蛋白co-design不应该停留在残基级标签和骨架几何的拼接上,而应该把离散的化学身份和连续的原子坐标放到同一个生成动态里处理。
从结果看,这个方向是有收益的。无条件生成里,A-CODE同时提高designability和codesignability;binder设计里,它把one-stage co-design的success rate明显拉高;ncAA中,它又展示了传统残基级模型很难自然扩展的能力。尤其是side-chain lag的成功,说明多模态生成不只是把两个loss加在一起,还要认真处理不同影响层的去噪时间尺度,这恰恰也是目前很多模型没有考虑到的点,时间耦合关系。
当然,A-CODE也不是终点。作者自己也承认,one-stage co-design和two-stage PMPNN-redesign之间仍然存在差距;ncAA评估还缺乏公认benchmark;另外,当前方法主要还是围绕蛋白和非天然氨基酸。
整体来看,A-CODE给出了一个很清晰的技术信号:未来的蛋白设计模型,很可能不会满足于“先骨架、再序列”的pipeline,而会越来越多地转向原子层面的多模态设计。只有在这个粒度上,模型才有机会同时理解侧链拓扑、局部化学、非天然残基和界面相互作用。
对做肽和环肽设计的人来说,这篇文章尤其值得注意。因为环肽、非天然氨基酸、修饰残基、小分子桥联和共价连接,本质上都要求模型理解超出20种标准氨基酸的化学结构。A-CODE现在还没有直接解决所有这些问题,但它给出的解决方式,确实比传统残基级模型更接近这些任务的需求。
简单说,A-CODE不是又一个普通蛋白生成模型,而是在尝试把“序列-结构共设计”这件事改写成“原子类型-原子坐标共生成”。这个转变如果继续走下去,很可能会直接影响后续全原子蛋白设计、非天然氨基酸蛋白设计以及更复杂的生物分子生成模型。
原文链接
https://arxiv.org/abs/2605.03360
往期文章
亲和力预测模型为什么普遍虚高——从Vina到Boltz-2,没人愿意承认的那个天花板

