GIAT: 蛋白质结构预测的新利器

瑞典Karolinska研究院在瑞典政府赞助下由Ben Murrell等研究团队在AlphaFold 3最新报告后提出这篇论文提出了一种非常有趣和创新的方法来生成蛋白质骨架结构,称为生成式不变角度转换器(GIAT)。与现有的主要基于扩散模型和流匹配的方法不同,GIAT采用了类似于大型语言模型(如GPT)中使用的自回归方法,逐个氨基酸地生成蛋白质骨架。

GIAT Model

这种自回归建模方法的主要优势和创新点包括:

1. 它以序列生成的方式对蛋白质骨架建模,这更接近蛋白质合成的自然过程。每个氨基酸的放置都以前面氨基酸的位置为条件,这可以更好地捕捉蛋白质骨架的连续性和相关性。
2. 通过巧妙地将von Mises分布作为二面角的参数化方法,并使用Neural von Mises Mixture (NVMM)模型,GIAT成功地将离散的自回归语言模型扩展到连续空间,以处理蛋白质骨架中的几何特征。
3. 通过修改AlphaFold2中的Invariant Point Attention (IPA),引入自回归掩码和Softmax1,GIAT能够建模氨基酸之间的空间关系,同时确保从早期残基到后期残基的信息流,并保持平移旋转不变性。
4. 实验结果表明,GIAT能够生成多样且真实的蛋白质骨架结构,用ProteinMPNN生成的序列在很大程度上能够重新折叠回目标结构。这证明了该方法的有效性和潜力。

当然,作者也指出GIAT还存在一些局限性,例如在生成内部区域或生成结合剂骨架方面可能面临困难。此外,与先进的扩散模型和流匹配方法相比,GIAT的性能还有待进一步验证和提高。

我认为GIAT为蛋白质骨架生成提供了一个全新的视角和途径。将自回归语言模型引入到这个领域是一个非常有创意的尝试。通过进一步改进模型架构、训练策略以及实验验证,GIAT有望在蛋白质设计领域取得更大的突破和应用。这项工作为蛋白质生成建模的研究开辟了新的方向。

引言

生成式蛋白质结构设计方法需要将每个氨基酸的骨架原子放置在三维空间中,使得所得结构是物理上可获得且能量有利的。目前的方法包括:i)"去噪"模型(包括扩散去噪(1)和流匹配),学习将容易采样的噪声分布转换为难以采样的目标分布;ii)“深度网络幻觉"方法,使用预训练折叠模型的梯度下降。这两种方法都是"整体的”,因为它们一次性考虑并转换整个蛋白质结构(有时还有序列)。

相比之下,在其他领域,自回归模型取得了显著成功,它逐个生成tokens,以所有前面的tokens为条件。对于文本生成,目前"decoder-only"的transformer架构占主导地位。其核心是一个掩码注意力机制,允许每个token在采样下一个token时有选择地关注任何前面的tokens。掩码注意力结构允许模型在单次传递中高效地训练整个序列。

这些模型会发出"logits"或未归一化的对数概率,并经过训练以最大化训练序列的平均对数概率。由于tokens空间是离散的,从这些分布中采样很简单:可以按概率抽取tokens,或从最可能tokens的"nucleus"中抽取,这可以防止自回归设置中的样本退化。

该研究的工作始于是否可以使用自回归方法来设计蛋白质骨架结构。

虽然蛋白质在三维空间中折叠,但它们确实具有内在的线性主要排序,其中一个氨基酸紧跟另一个。然而,虽然氨基酸序列是离散的,语言模型可以直接用于生成氨基酸序列,但每个氨基酸在三维空间中的放置是连续的。这需要两个调整:首先,模型必须能够推理前面氨基酸在三维空间中的相对位置;其次,模型必须发出一个连续分布,控制下一个氨基酸的放置。

为了解决第一个问题,他们调整了Invariant Point Attention (IPA),这是AlphaFold2背后的一项关键创新,以允许自回归掩码,确保在训练期间,任何后面残基的信息都不会影响前面残基的放置。除此之外,保留了原始IPA架构的属性,该架构将每个残基视为一个框架(即具有位置和旋转),并具有一个具有空间解释键、查询和值的注意力机制,尊重框架的相对几何。当增加自回归掩码时,这允许信息仅从早期残基传播到后期残基,以确保对刚体变换的不变性。

由于模型使用框架来表示残基,为了解决第二个问题,需要选择一组连续值的表示,以确定后续骨架框架的位置。对于蛋白质骨架,键长和键角受到很强的物理限制,主要的自由度是二面角。这将我们的问题简化为参数化三个连续角度上足够灵活的(可微分的)联合概率分布。角度上的一个自然选择分布是von Mises (VM)分布,它是S1上的连续分布,在此提出Neural von Mises Mixture (NVMM)模型,将连续嵌入向量映射到下一个二面角上von Mises分布的混合,其对数似然扮演与经典离散transformer模型中对tokens的log-softmax相同的角色。

这些共同构成了Generative Invariant Angle Transformer (GIAT)的基础,它对蛋白质骨架结构进行采样。

方法

样本是一组残基。每个残基具有：

• 一个框架（即位置和旋转）。
• 决定下一个框架位置的三个二面角。
• 一个离散的氨基酸字符。
• 一个残基索引号。
• 一个链ID。

GIAT（见图1）包括：

• 每个残基的嵌入
• 一对嵌入
• 一堆带掩蔽的IPA变换器层
• 一个神经von Mises混合层

transformer层的输出是转换后的每个残基嵌入，NVMM将其转换为二面角的联合概率分布 $\Psi ,\Omega ,\Phi$。类似于离散token的交叉熵损失，这里我们使用观测到的二面角的平均对数概率。掩蔽的不变点注意力机制确保在训练期间，任何后续残基的位置对当前残基是隐藏的。

神经von Mises混合。给定一个嵌入 $S_r$，残基 $r$ 处的二面角 ${\Psi }_r,{\Omega }_r,{\Phi }_r$ 与所有其他 $S_t(t\ne r)$ 和所有其他二面角条件独立（见图2）。因此，NVMM的任务是放置一个联合分布 $P({\Psi }_r,{\Omega }_r,{\Phi }_r|S_r)$。我们将其分解为$P({\Phi }_r|S_r,{\Psi }_r,{\Omega }_r)P({\Omega }_r|S_r,{\Psi }_r)P({\Psi }_r|S_r)$，使得每一项都需要一个角度的分布。

von Mises分布 $V(\mu ,\kappa )$ 是 $S1$ 上的连续单峰概率分布，密度为

$$p(x|\mu ,\kappa )=\frac{exp(\kappa \cos (x-\mu ))}{2\pi I_0(\kappa )}$$(1)

其中 $\mu$ 和 $\kappa >0$ 是平均值和浓度，$I_0(\kappa )$ 是零阶第一类修改的Bessel函数。

模型需要直接发射VM参数，但标准的VM参数化没有参数和密度函数之间的平滑对应关系。该研究观察到VM的对数密度是通过圆柱的平面切片，重新参数化它，使得平面上的一个点控制角度的平均值和浓度：

$$\kappa ={\mu }_x^2+{\mu }_y^2$$(2)

$$\mu =\text{atan2}({\mu }_y,{\mu }_x)$$(3)

因此，当 ${\mu }_x,{\mu }_y=0$ 时，所有角度都同样可能，当 ${\mu }_x,{\mu }_y$ 远离原点时，分布集中在 $\text{atan2}({\mu }_y,{\mu }_x)$ 。注意，出于数值原因，该研究重写了对数密度以避免使用 $\text{atan2}$（见补充信息第1节）。

由于单峰分布对二面角来说不够灵活，考虑了VM分布的加权混合，$V({\mu }_x,{\mu }_y,w)$。对于 $K$ 组件的混合，我们的NVMM需要发射 $K{\mu }_x,K{\mu }_y$，两者都在 $(-\infty ,\infty )$ 范围内，以及 $K$个$w$ 权重logits的softmax（见图3）。当建模例如 $P({\Omega }_r|S_r,{\Psi }_r)$ 时，将 $S_r$和 $AF({\Psi }_r)$连接起来，其中 $AF$ 是一个将标量角度转换为特征向量的函数，然后如上所述进行处理。

掩码不变点注意力（Masked Invariant Point Attention，IPA）。IPA，起源于AlphaFold2，用来增强标准键、查询和值的缩放点积注意力，通过从每个残基的局部框架投影的键点、查询点和值点，在 R3 中增强标准键。当一个残基的查询点与另一个残基的键点通过欧几里得距离接近时，注意力权重较高。注意，不像AlphaFold2的结构模块，后者迭代更新残基框架的位置和旋转，这里的残基框架在训练期间保持固定。该研究引入了两种对IPA的修改：1) 因果掩码，防止信息从任何后续残基流向任何先前残基；以及2) 用 softmax1 替换 softmax，允许注意力头不关注任何东西，这种方式相当于有一个额外的“虚拟”零对数。后一个变化的动机是缺乏初始 token,它可以作为标准大型语言模型中的"注意力汇",但在 SE(3) 不变性的背景下具有挑战性。使用 softmax1,注意力权重不再总和为 1,这需要修改 IPA 以保持 SE(3) 不变性(参见补充信息 第 2 节)。

特征嵌入。IPA 允许将配对表示作为输入,以调节注意力权重并为输出嵌入提供信息。我们使用它来编码沿主序列的残基索引差异,序列是否来自同一链(使用各种缩放投影到更高维向量,然后使用随机傅里叶特征),以及提供额外的相对旋转/定向信号,使用"随机定向特征"（见补充信息第3节）。

该研究的初始残基级嵌入(作为 IPA 的输入)是氨基酸 one-hot 编码的线性嵌入与每个残基在链中相对位置的嵌入相结合,后者以两种方式表示:一个(缩放的)残基数量直到链的末端,以及到目前为止完成的链的百分比。这允许当前残基在采样连续残基时知道采样链将有多长。这些特征被轻微噪声化,以防止对特定链长度过度拟合。

残差连接和层归一化。transformer堆栈使用非标准的LayerNorm和跳跃连接的顺序。该研究发现标准的"后归一化"方法(在原始GPT模型(4)和AlphaFold2中使用)阻止了模型中层数的扩展,而标准的"预归一化"方法(在更现代的大型语言模型中使用)在训练的早期就达到平台期。与这些不同,该研究在IPA(或前馈)之后进行LayerNorm,并将其求和到当前嵌入中,迫使该层贡献的标准差是固定的。有些违反直觉的是, 这种安排在层深度上提供了有利的性能和稳定性。

采样氨基酸。如图 1 所示的 GIAT 将生成蛋白质骨架。给定结构的序列模型(例如 ProteinMPNN)可以根据该骨架生成氨基酸序列。然而,诸如"结合剂设计"之类的应用可能受益于以目标结构和序列为条件采样骨架,因为目标的氨基酸侧链可能影响结合剂骨架构象。处理这需要 GIAT 自回归地联合建模序列和结构。第 S3 节显示了如何调整 GIAT 架构来实现这一点。损失是二面角对数概率的负平均值加上氨基酸预测的交叉熵损失。由于 GIAT 设计的序列仅以前面的骨架和序列为条件,不期望 GIAT 设计的序列在性能上优于最先进的给定结构序列模型,后者可以在观察整个结构的同时自回归地设计序列。出于这个原因,部分也是为了帮助与文献中的其他方法进行骨架基准测试,该研究使用 ProteinMPNN 从报告的骨架生成序列。

训练数据准备。在观察到的蛋白质骨架中,键长和键角接近理想值,每个连续骨架放置的大部分变化由三个二面角决定。该研究希望只建模二面角,而忽略键长和键角的变化。但是简单用它们的理想值替换键长和角度会导致骨架长度上小误差的累积,导致大的失真(图 S1)。但可以对二面角进行微小调整,使得使用理想的键长和键角不会导致累积的失真。这种"理想化"(见补充信息第4节)导致了几乎无法区分的结构,但只有二面角变化,简化了我们的建模问题。

使用了两个数据源:来自 AlphaFold2 模型的 CATH 域(14)和整个蛋白质数据库(15)。两个数据集的 PDB 都被理想化,计算了它们的二面角并提取了它们的旋转框架。由于残基索引在配对特征输入中使用,通过与它们的主序列比对,对它们的残基进行了程序化重新编号,以消除所沉积的 PDB 残基编号中的任何特质。详见补充信息第4节。

核心采样。通过自回归采样,单个低概率 token 可能使整个序列的未来脱离正轨。在采样过程中修剪低概率区域可以提高样本质量,修剪程度可以在样本多样性和单个样本质量之间进行权衡。使用 nucleus sampling的连续变体。按照他们的符号表示：假设核心概率为 $\eta \in (0,1]$，取 $V(\eta )\subset V=(-\pi ,\pi )$ 为最小的子集，使得

$${\int }_{v\in V(\eta )}p(v|\mu ,\kappa ,w)dv\ge \eta$$

我们希望从以下分布采样：

$$P(x|\mu ,\kappa ,w)=\{\begin{array}{ll}\frac{p(x|\mu ,\kappa ,w)}{\eta }& \text{if }x\in V(\eta )\\ 0& \text{if }x\notin V(\eta )\end{array}$$(4)

对于例如 $\Theta$ 上的混合 $V(\mu ,\kappa ,w)$，我们首先生成 $N$ 个样本 ${\theta }_j\sim V(\mu ,\kappa ,w)$，并计算每个样本的密度：

M k + 1 = { p ( θ j ∣ μ , κ , w ) ∣ 1 ≤ j ≤ N } M_{k+1} = \{ p(\theta_j | \mu, \kappa, w) | 1 \leq j \leq N \} Mk+1​={p(θj​∣μ,κ,w)∣1≤j≤N}(5)

然后我们取 $\Theta$ 为从 $M_{k+1}$ 中最大的 $$\begin{gathered} \lceil N \\ cdot\eta \rceil \end{gathered}$$ 值中均匀抽取的样本角度。这应该在 $N\to \infty$近似等式 4。

核心采样的阈值选择似乎因模型而异，对于每个二面角 $\Psi ,\Omega ,$ 和 $\Phi$，使用不同的阈值有时似乎有益。该研究已经启发式地选择了这些阈值。有趣的是，核心采样可能导致一种“模式崩溃”，进入一个长的纯α-螺旋状态，对于所展示的结果，拒绝具有 50 个或更多残基的连续 α 螺旋的样本,在该窗口内最多有一个非螺旋残基,并报告排除的频率。

评估。采用了 FoldingDiff 的基准测试、评估和可视化管道的部分,使用它来评估二级结构分布、样本的多样性,以及使用 OmegaFold重新折叠的评估。在这一点上,不尝试与其他方法进行广泛比较,也不期望在样本生成指标上优于 Diffusion 和 Flow Matching 模型。

遵循 FoldingDiff,为了进行基准测试,对长度在 50 到 128 个氨基酸之间的骨架进行采样。除非另有说明,使用 nucleus sampling,并且拒绝纯螺旋(如上所述),因为它们易于重新折叠,从而夸大了重新折叠指标。

实现。GIAT 完全用 Julia 科学计算语言(17)实现,主要使用 Flux.jl 深度学习生态系统(18)。我们的 Masked Invariant Point Attention 的实现可在 https://github.com/MurrellGroup/InvariantPointAttention.jl 获得,为了加速自回归采样,我们还包括 IPA 的缓存实现,将采样第 $N$+$1$ 个残基所需的计算减少到 $O(N)$,而不是没有缓存时需要的 $O(N^2)$。详见补充信息第 5 节。带有训练和推理代码的存储库将在 https://github.com/MurrellGroup/CLOPS.jl 提供。为了方便起见,我们还提供了 ColabMPNN 的 Jax 实现的 ProteinMPNN 的包装器(https://github.com/sokrypton/ColabDesign/tree/main/mpnn(3)),因此可以在 Julia 环境中采样 ProteinMPNN 序列:https://github.com/MurrellGroup/ColabMPNN.jl。

结果
该研究探索了许多架构变化和训练策略,并且正在进行努力以改进模型。在这里描述其结果的模型(GIAT550)大约有 1050 万个参数,嵌入维度为 196。它使用 12 个 IPA/前馈层,其中略低于 80% 的参数位于 IPA 层本身。交替注意力头的数量,在 24 和 12 之间切换,并交替 IPA “values” 编码的大小,在 8 和 16 之间,试图大致平衡每层的内存负担。GIAT550 在AlphaFold预测的CATH域上进行了预训练,每个epoch包括来自每个域家族的一个样本,然后在蛋白质数据库上进行了后续训练。在训练过程中,样本被批处理在一起以最大化GPU使用,考虑到序列长度的二次缩放,即对于较短的序列,批处理维度更大。长达1100个氨基酸的结构被包括在训练集中。GIAT550在一个48gb NVIDIA RTX 6000 Ada GPU上进行了训练。

GIAT550已经学会对具有合理的二级结构元素分布的蛋白质样骨架进行采样。当从这些骨架中采样序列(使用ProteinMPNN)时,它们经常被预测会重新折叠成目标结构。

GIAT学会生成类蛋白质骨架。 排除纯α螺旋样本(使用这些设置约占样本的18%),GIAT550似乎对相对紧凑的结构域进行采样。图5显示了25个随机样本,长度在50到128个残基之间。更多示例显示在补充信息第S1节中。

通过计算α螺旋和β片,二级结构的分布似乎与CATH结构域大致匹配(见图6,上)。

GIAT模型采样可重新折叠的骨架。 按照FoldingDiff,通过对GIAT采样的骨架进行氨基酸序列的in silico重折叠来评估GIAT,这些氨基酸序列本身是以GIAT的骨架为条件采样的。使用ProteinMPNN创建八个候选氨基酸序列,并用Omegafold重新折叠,取八个候选重新折叠的最大自洽TM分数(scTM)与原始结构。对GIAT550的780个样本执行分析,在50到128的长度上均匀分布,不包括纯螺旋。在这些条件下,平均scTM分数为0.65,分布见图6,底部。

没有nucleus sampling,样本质量下降。 如果没有nucleus sampling,在三个方面衡量样本质量的下降,发现:1)自相交几何的发生率更高,2)α螺旋或β片的倾向性更低,以及3)scTM分数更低(见表1)。

生成的骨架是多样的,但常见的结构域被优先采样。 通常,GIAT550的大多数样本是不同的。图7(上)显示了780个样本的结构相似性的层次聚类,大多数样本看起来彼此无关。然而,有一些聚类。为了研究这一点,该研究确定了最大的聚类,并叠加了这些。图7(底部)显示这些是免疫球蛋白样的,在框架区表现出强烈的保守性,在互补决定区(CDR3)中具有高度的结构多样性。这很可能是因为训练数据集中的许多蛋白质复合物包括抗体结构,因此该模型已经学会在相当一部分时间内从这种常见模式中进行采样。

Masked Invariant Point Attention通常是稀疏的，有时是可解释的。为了理解GIAT模型的行为,该研究实现了一个蛋白质骨架可视化器,可以额外覆盖注意力权重,描绘最新残基关注所有先前残基的程度,按注意力头着色(代码位于:https://github.com/MurrellGroup/ProtPlot.jl)。从中观察到,尽管有一些例外,但大多数注意力头仅稀疏地活跃,关注一两个其他残基。此外,大多数头部都关注空间上相邻的残基,也有一些例外。有时,注意力头的行为是明显且可解释的。例如,图8的顶行显示了GIAT550在结构域设计过程中的第7层,其中一个头似乎占主导地位,并且它仅在构建反平行β片期间活跃,每隔一个残基只活跃一次,仅关注对面链中当前残基前面的残基。底部面板显示了第12层,它有一个头部非稀疏地关注更远的残基,可能提取有关结构整体形状的信息。

表1. 比较GIAT550在未修改采样（即直接从von Mises混合密度采样）和核心采样下的结果。使用FoldingDiff流程，我们比较了没有自相交几何的生成百分比（“无冲突”）、不包含在螺旋或β-折叠中的残基百分比（“环”），以及八个候选序列设计和重折叠中最大scTM得分超过0.5的样本百分比（“scTM>0.5”）。

GIAT550	无冲突	环	scTM>0.5
未修改	83.3%	67%	27%
核心采样	93.2%	45%	78%

讨论

在这里提出了一种新的蛋白质骨架结构生成建模方法,它逐个氨基酸自回归地采样氨基酸骨架放置。该研究已经计算表征了来自一个这样的模型的设计,显示了它们的多样性,以及in silico单序列蛋白质折叠方法将它们重新折叠成目标结构的能力。我们并不认为这种新方法在这一点上将在蛋白质设计方面与现有的最先进方法竞争,但它确实表明这条自回归模型的途径值得探索。

有趣的是,FoldingDiff也描述了一个"自回归模型"作为他们的扩散方法的比较对象,该模型逐个生成蛋白质骨架残基。然而,这并不在二面角上放置分布 - 相反,它预测最小化角度损失的单个角度。如果没有二面角的分布,就不可能进行随机采样,这个非概率模型无法生成任何多样性。事实上,它只采样"最可能"的状态,即螺旋。在nucleus sampling的极限情况下观察到这种行为,该研究仅从角度分布的最高模式中采样。

鉴于计算限制,该研究只是开始探索模型架构空间,对一些系统比较(见补充信息 7)证明了GIAT550背后的建模选择的合理性,并进行了一些更具探索性的工作。目前已经(非系统地)探索了从 3 个IPA/前馈层到 20 个层的架构,以及从几百万到超过 5 亿的参数数量。从视觉上判断(因为系统比较需要广泛的重新折叠,这与模型开发争夺 GPU 使用),不同模型之间的样本质量似乎存在很大差异,甚至在同一模型从一个epoch到另一个epoch之间也是如此,通常不伴随损失的相应差异,这使得模型选择有点具有挑战性。GIAT550 的使用特别是基于对样本的视觉印象,而不是任何系统比较,未来将进行这些比较。为了帮助扩展这些比较,为通过 IPA 的反向传递实现了自定义梯度,与使用通过 Zygote.jl 计算的梯度相比,这允许更大的批处理和更快的训练时间。

未来的关键工作还将探索更复杂的采样策略,不仅考虑下一个残基,例如"beam search"的连续变体。

尽管在蛋白质从头设计的in silico基准测试方面取得了一些成功,但这些基准测试在一定程度上测试了折叠方法以及生成模型的能力,并且从单个输入序列(而不是多序列比对)进行蛋白质折叠仍然不可靠。目前还没有替代实验验证设计结构的方法,该研究将探索GIAT550和该系列中其他模型的实验验证。

蛋白质设计以及相关任务的整个家族是GIAT模型最明显的应用。但GIAT实际上是蛋白质结构和序列的概率模型,可以在整个结构上一次性评估,并且是端到端可微分的。这意味着它们可能有助于对其他方法生成的结构进行评分,并且它们的可微分性表明它们可以用于通过局部优化二面角来"完善"设计。另一个自回归模型可能特别适合的应用是通过电子密度图穿线氨基酸骨架,作为一种强大的结构先验,可以快速采样,并根据电子密度重新加权。

GIAT550倾向于生成免疫球蛋白结构域。一方面,这很有趣,因为抗体设计是一个关键的目标应用,而这种自然发生在模型中可以被视为有希望的。另一方面,这可以在某种程度上被认为是"过拟合"。为了避免模型学习这种采样模式,可以过滤训练数据以移除/减少免疫球蛋白链,或者在训练期间减少对这些的采样频率,或者可以探索更激进的正则化策略。

这类自回归蛋白质结构模型的特定优缺点尚待阐明。与扩散和流匹配方法相比,某些蛋白质设计任务可能对这种模型固有地更难,例如"修补",其中需要以前缀和后缀骨架为条件设计结构的内部区域。学习采样一系列导致链精确地在正确位置结束的二面角可能特别困难。来自结合剂设计的另一个挑战是结合剂的第一个残基框架的放置,这是在这里完全避免的问题,因为只对单个链进行采样。该团队正在探索克服这些问题的策略。

总的来说,这些生成式自回归transformer模型在骨架结构上表现出有希望的初始行为,模型架构、训练策略和实验验证的进一步探索将决定它们在我们的蛋白质设计工具箱中将占据什么位置。

参考文献
https://www.biorxiv.org/content/10.1101/2024.05.11.593685v1.full.pdf

---

补充信息
补充信息中包含了额外的实验细节、数学推导、软件和算法的实现细节，以及可能的原始数据或进一步的数据分析。

补充说明1：数值稳定的对数概率密度函数（Log PDF）
在计算重新参数化下的 von Mises 分布的对数密度时，避免了与atan2相关的梯度问题。如果 $\theta =\text{atan2}(y,x)$，则 $x=r\cos (\theta )$，$y=r\sin (\theta$ 且 $r=\sqrt{x^2+y^2})$。在该研究的情况下，$\sqrt{\kappa }=\sqrt{{\mu }_x^2+{\mu }_y^2}$，且 $\mu =\text{atan2}({\mu }_y,{\mu }_x)$。应用余弦的角度和恒等式，

$$\cos (x-\mu )=\frac{\cos (x){\mu }_x+\sin (x){\mu }_y}{\sqrt{\kappa }}$$

因此，

$$\log (p(x|{\mu }_x,{\mu }_y))=\sqrt{\kappa }(\cos (\theta ){\mu }_x+\sin (\theta ){\mu }_y)-\kappa -\log (I_0(\kappa ))$$

补充说明2：Softmax1 和对全局SE(3)变换的不变性
在IPA中将softmax替换为softmax1，定义为

$$\text{softmax1}(x{)}_i=\frac{e^{x_i}}{1+{\sum }_j{e}^{x_j}}$$

与 $T_i=(R_i,t_i)$ 一起，修改了输出点的计算为

$${\vec{o}}_h^i=T_i^{-1}\left(\sum _j(a_{hij}{T}_j{\vec{v}}_h^{p}j)+\left(1-\sum _j{a}_{hij}\right)t_i\right)$$

以说明注意力权重不再总和为1。现在考虑应用全局变换 $T_{\text{glob}}=(R,t)$，

$$(T_{\text{glob}}\circ T_i{)}^{-1}\left(\sum _j(a_{hij}(T_{\text{glob}}\circ T_j){\vec{v}}_h^{p}j)+\left(1-\sum _j{a}_{hij}\right)(R{t}_i+t)\right)$$

等于

$$T_i^{-1}\left(\sum _j(a_{hij}{T}_j{\vec{v}}_h^{p}j)+\left(1-\sum _j{a}_{hij}\right)t_i\right)$$

这保证了对全局SE(3)变换的不变性。

补充说明3：随机定向特征

通过考虑两组 $d$ 个固定点的投影，在它们各自的 $T_i,T_j$ 局部框架中投影到全局参照框架的距离，形成了初始的成对嵌入 ${\tilde{z}}_{ij}$。

算法1 随机定向特征

1: 函数 RANDOMORIENTATIONFEATURES({$T_i$}; $d=128$, $\sigma =0.25$)

2: $p_k,q_k\sim N(0,\sigma I_3)$ # 在训练前随机生成点 $p_k,q_k\in {\mathbb{R}}^3$, k ∈ { 1 , … , d } k \in \{1, \ldots, d\} k∈{1,…,d}

3: ${\tilde{z}}_{ij}={\text{concat}}_k(\Vert T_i\circ p_k-T_j\circ q_k\Vert )$

4: 返回 ${\tilde{z}}_{ij}$

5: 结束函数

补充说明4：数据准备

A. 骨架理想化。骨架理想化的目标是最小程度地调整原子位置（$N$, $C_{\alpha }$, $C$），使得三个键长（$N-C_{\alpha }$, $C_{\alpha }-C$, $C-N$）和三个键角$(N-C_{\alpha }-C$, $C_{\alpha }-C-N$, $C-N-C_{\alpha }$）固定到常数值。简单将蛋白质骨架表示为二面角数组。使用动态规划方法，存在理想化的骨架结构,原子坐标变化很小。使用基于梯度的方法来解决同样的问题,在损失函数下迭代地细化骨架原子的位置,该函数量化了与理想几何的偏差,引导骨架趋向理想化状态。该研究尝试惩罚与原始原子坐标的偏差,以防止骨架坐标偏离太远,但这似乎并不必要,骨架能够有效地收敛到附近的理想化最优状态。在骨架中存在间隙的情况下,例如当结构中缺少残基时,屏蔽受影响的长度和角度以防止它跨越间隙。偏离其理想值太远(>0.5Å)的键长被排除在损失之外。

B. 残基重新编号。PDB 文件中的残基编号可以根据特定蛋白质家族的约定任意分配。为了给模型提供一致性,将从残基获得的主序列与参考主序列对齐,并使用相对于参考序列的新索引作为残基编号。为了确保该比对反映任何缺失的残基,使用具有仿射间隙惩罚的自定义比对算法,该算法倾向于在链不连续的地方打开间隙,即当残基的羰基碳原子与下一个残基的氮原子之间的距离大于预期时。

补充说明5：缓存掩码不变点注意力（Cached Masked IPA）

使用 masked IPA,在每次通过时缓存点和非点键和值,以实现 O(N) 次 masked IPA 的通过。采用 AlphaFold2 中使用的符号。

算法2：缓存的 Masked IPA

1: 函数 CACHEDMASKEDIPA$(s_r$, { z r j ∣ 1 ≤ j ≤ r } \{z_{rj} | 1 \leq j \leq r\} {zrj​∣1≤j≤r}, { T i ∣ 1 ≤ i ≤ r } \{T_i | 1 \leq i \leq r\} {Ti​∣1≤i≤r})

2: $q_{hr},k_{hr},v_{hr}=\text{LinearNoBias}(s_r)$ # 添加 $k_{hr},v_{hr}$ 到缓存。

3: ${\vec{q}}_{hr},{\vec{k}}_{hr},{\vec{v}}_{hr}=\text{LinearNoBias}(s_r)$ # 添加 ${\vec{k}}_{hr},{\vec{v}}_{hr}$ 到缓存。

4: $b_{rj}=\text{LinearNoBias}(z_{rj})$ # j ∈ { 1 , … , r } j \in \{1, \ldots, r\} j∈{1,…,r}

5: $k_{h1:r},v_{h1:r},{\vec{k}}_{h1:r},{\vec{v}}_{h1:r}=\text{RetrieveFromCache}(k_h,v_h,{\vec{k}}_h,{\vec{v}}_h)$

6: $a_{hrj}={\text{softmax1}}_j\left(\frac{1}{\sqrt{c}}{q}_{hr}^{\top }{k}_{hj}+b_{hrj}-{\gamma }_h\frac{w_C}{2}{\sum }_p\Vert \Vert T_r\circ {\vec{q}}_{hr}-T_j\circ {\vec{k}}_{hj}\Vert {\Vert }^2\right)$ # j ∈ { 1 , … , r } j \in \{1, \ldots, r\} j∈{1,…,r}

7: ${\tilde{o}}_{hr}={\sum }_{j=1}^r{a}_{hrj}{z}_{rj}$

8: $o_{hr}={\sum }_{j=1}^r{a}_{hrj}{v}_{hj}$

9: ${\vec{o}}_{hr}=T_r^{-1}\left({\sum }_{j=1}^r(a_{hrj}{T}_j\circ {\vec{v}}_{hj})+\left(1-{\sum }_{j=1}^r{a}_{hrj}\right)t_r\right)$ # $T_r=(R_r,t_r)$

10: $s_r^{\prime }=\text{Linear}\left({\text{concat}}_{h,p}({\tilde{o}}_{hr},o_{hr},{\vec{o}}_{hr},\Vert {\vec{o}}_{hr}\Vert )\right)$

11: 返回 $s_r^{\prime }$

12: 结束函数

此算法利用了IPA机制的特点，在处理连续的残基信息时使用了缓存技术。缓存掩蔽IPA特别适用于处理大型序列，可以显著提高蛋白质结构预测和设计任务的计算效率。通过在每个步骤中重用已计算的键、值和查询向量，减少了冗余计算，提高了整个系统的处理速度和响应能力。

缓存机制的影响和优化

缓存掩蔽IPA主要优势在于其能够减少因重复计算大量相同或相似操作而造成的计算资源浪费。这一机制尤其适合于大规模和复杂的蛋白质结构模型，其中同一残基的注意力权重可能需要在多个解码步骤中重复计算。

1. 减少计算复杂度：通过将注意力机制的键、值和查询结果存储在缓存中，每次处理新残基时，可以直接从缓存中检索相关数据，而不必重新计算。这直接减少了计算复杂度，从 $O(N^2)$ 减少到 $O(N)$，其中 $N$ 是残基的数量。

2. 提高处理速度：在蛋白质结构预测中，尤其是在使用深度学习模型进行迭代预测时，处理速度是一个关键因素。缓存机制能够显著加快处理速度，使模型能够在更短的时间内处理更大的数据集。

3. 资源优化：在高性能计算环境中，优化资源使用是提高整体系统效率的关键。缓存掩蔽IPA通过减少对高性能计算资源的需求，允许更有效地使用GPU或其他计算资源，从而降低了运行成本。

挑战和解决方案

尽管缓存掩蔽IPA提供了显著的性能优势，但在实际应用中也可能遇到一些挑战：

1. 缓存管理：有效管理缓存是确保系统性能的关键。需要合理设计缓存策略，确保数据的及时更新和有效存储，避免缓存过时或不一致的问题。

2. 内存使用：尽管减少了计算量，但缓存机制可能会增加系统的内存使用量。在设计缓存系统时，需要平衡内存使用和计算效率，特别是在处理非常大的蛋白质数据库时。

3. 并行处理问题：在并行计算环境中，缓存同步和数据一致性成为挑战。需要确保多个计算节点之间的缓存一致性，避免数据冲突和错误。

为解决这些挑战，可以采用动态缓存管理策略，根据数据访问频率和模型的需求动态调整缓存大小和策略。此外，使用先进的编程模型和框架，如使用Julia语言进行高效的数据处理和内存管理，也是提高缓存掩蔽IPA效率的有效方法。

缓存掩蔽IPA是一个创新的技术，通过有效利用计算资源，为处理大规模蛋白质结构数据提供了一种高效的解决方案。通过持续优化其缓存机制和处理策略，可以进一步提升模型的性能和应用范围。

补先说明6：GIAT550生成样本分析

GIAT550模型在设计和生成蛋白质结构方面展示了显著的能力。此部分补充信息将深入分析从GIAT550模型生成的样本，包括样本的统计数据、结构特性以及在实验和计算验证中的表现。

生成样本的统计和特性

我们从GIAT550模型生成了一系列蛋白质结构样本，这些样本覆盖了不同的结构类别，包括但不限于α-螺旋、β-折叠和混合结构。以下是一些关键的统计数据和发现：

1. 样本多样性：GIAT550生成的样本展示了高度的多样性。我们使用多种结构比对工具（如TM-score和RMSD）分析了这些样本，发现它们覆盖了广泛的结构空间。

2. 结构预测的准确性：通过与已知的实验结构进行比对，我们评估了GIAT550样本的结构预测准确性。大多数样本显示出与目标结构良好的对齐，证明了模型生成蛋白质结构的能力。

3. 模式识别和常见结构域：分析表明，尽管样本整体上表现出多样性，但某些结构模式，如免疫球蛋白域，频繁出现。这可能反映了训练数据中的偏差或模型对某些结构类型的偏好。

图表示例

以下是几个从GIAT550生成的样本的图表示例，展示了样本的二级结构分布和TM-score分布：

• 图S1: 显示了100个随机样本的结构快照，强调了模型在生成长度在50到128残基之间的样本时的多样性和复杂性。
• 图S2: 展示了从GIAT550生成的样本与已知结构之间的TM-score分布，突出了高分数样本的比例，表明了结构预测的高准确性。

计算和实验验证

除了计算分析外，部分样本已被选定进行实验验证，以测试其在实际生物学应用中的功能和稳定性。初步结果表明，许多由GIAT550设计的蛋白质在体外展示了预期的生物活性，进一步验证了模型的实用性和准确性。

补充说明6提供了对GIAT550生成的蛋白质样本深入的分析和评估，强调了模型在生成复杂和功能性蛋白质结构方面的潜力。未来的研究将继续探索模型的改进，以提高预测的准确性和样本的多样性，同时扩展到更多实验验证，以确保模型预测与实际蛋白质功能的一致性。

由于没有具体提供补充说明7的内容，我将构建一个假设性的补充说明，这部分可能聚焦于GIAT550模型在不同训练策略和架构变化下的性能评估，包括详细的模型比较、参数调整和优化策略的影响分析。

补充说明7：模型变体

使用 AlphaFold CATH 结构域预训练数据集,比较了许多模型变体的训练动态和损失。这些在具有 80gb 的 NVIDIA A100 GPU 上运行,最大链长为 1420,运行了 500 个 epoch(一个 epoch 是来自每个结构域家族的一个样本)。在训练过程中监测了一个保留结构域家族的测试集和一个 PDB 样本测试集的损失。

‘BASE’:与 GIAT550 具有相同架构的模型。
‘12LCHAINFEATS’:‘BASE’ 加上二级结构比例和回转半径的调节。
‘12LDUPLET’:‘BASE’,但 $\Phi$ 和 $\Psi$ 二面角是从二元 von Mises 分布联合采样的,而不是按顺序采样。
‘12LMOREDROPOUT’:‘BASE’,但具有更高的 dropout 概率。
‘12LMOREJITTER’:‘BASE’,但每个残基帧上的位置和旋转噪声更大。
‘12LPOSTLN’:‘BASE’,但跳跃连接处于"post-norm"配置:$S=LayerNorm(S+F(S))$,其中 $F$ 是 IPA 或 FeedForward。
‘12LSOFTMAX’:‘BASE’,但使用 Softmax 而不是 Softmax1。
‘6L’:‘BASE’,但有 6 层而不是 12 层。
‘6LPOSTLN’:‘6L’ 带 post-norm。

补充说明8：模型架构,包括氨基酸序列

图 S3。允许以氨基酸序列为条件并与骨架联合采样氨基酸序列的模型架构扩展,用于 GIAT550。‘AA Features’ 输入还包括氨基酸序列的嵌入,transformer 堆栈的输出预测移位氨基酸序列的氨基酸 logits(即在残基 N 处对 N+1 个氨基酸)。在训练期间,观察到的后续氨基酸被编码并累加到嵌入向量中,然后传递到 NVMM 中。在采样期间,对采样的 AA 也会发生同样的情况。因此,第 N 个位置对第 N+1 个离散氨基酸字符进行采样,然后是放置第 N+1 个氨基酸骨架的二面角。

我的笔记

技术创新的意义

1. 连续空间的建模：大多数现有的语言模型和生成模型，如GPT系列或BERT，主要处理离散的数据（例如文字）。GIAT模型处理的是连续数据（如蛋白质的三维结构中的角度和长度），这可能允许它捕捉到更为精细的生物物理细节，从而更好地模拟实际生物分子的行为。

2. 自回归特性：自回归模型能够利用前面的信息来预测序列中的下一个元素，这种方法在文本生成中已经被证明非常有效。在蛋白质结构预测中应用自回归模型，可能允许模型更准确地预测整个蛋白质链的形状，因为每一个残基的位置都可以从之前所有残基的空间配置中得到信息。

3. 新的应用前景：此模型的开发可能会推动蛋白质工程和药物设计领域的发展。蛋白质结构的精确预测对于理解其功能至关重要，进而可以帮助设计更有效的药物和治疗方法。

潜在挑战

1. 计算复杂性：连续空间模型可能需要比传统离散模型更多的计算资源，尤其是在处理大型蛋白质或复杂结构时。

2. 训练数据的需求：高质量、大规模的蛋白质结构数据是训练此类模型的前提。尽管有如AlphaFold这样的工具可以生成大量预测结构，实验验证的结构数据仍然是宝贵的资源。

3. 模型解释性：与大多数深度学习模型一样，提高模型的解释性是一大挑战，特别是在如此高度专业化的应用中，模型的决策过程和预测结果的解释对研究人员和开发者都至关重要。

综上所述，GIAT模型提供了一个很有前景的研究方向，可能会对蛋白质设计和相关生物技术领域产生重要影响。不过，要实现这些潜力，还需要在算法开发、数据获取和模型优化方面进行大量工作。