AF3 TriangleAttention类解读

TriangleAttention类介绍

AlphaFold3 中的 TriangleAttention 类，用于对蛋白质的几何和语义信息进行建模。

继承类的设计

• TriangleAttentionStartingNode:
  • 对应 AlphaFold 算法中的起始节点（starting=True）。
• TriangleAttentionEndingNode:
  • 对应 AlphaFold 算法中的终止节点（starting=False）。

通过 partialmethod 复用 TriangleAttention 的构造函数，保持了代码的简洁性。

---

代码用途与背景

• 核心思想： 这个模块的核心是通过三角形几何建模蛋白质的关系网络（pair representation）。通过注意力机制，TriangleAttention 能够捕获残基之间的相对位置和几何关系。

• 起始与终止节点： 起始节点和终止节点分别对应蛋白质几何关系的不同方向（如从行到列或列到行的视角）。

• 应用场景： 这个模块是 AlphaFold 中更大的模块（例如 EvoFormer 块）的基础部分。它通过三角形几何捕获复杂的蛋白质序列-结构关系。

---

核心方法总结

1. _chunk: 内存优化的块计算。
2. forward: 包含三角形注意力的前向传播。
3. 偏置生成:
   • mask_bias：屏蔽无效位置。
   • triangle_bias：表示几何相关性。

• TriangleAttentionStartingNode 处理源节点（行方向），捕获“从源到目标”的交互。捕获“从源到目标”的交互。
• TriangleAttentionEndingNode 处理目标节点（列方向），捕获“从目标到源”的交互。

注意： 算法13 和14的代码差异，通过参数starting参数控制是否转置。

源代码：

class TriangleAttention(nn.Module):def __init__(self, c_in, c_hidden, no_heads, starting=True, inf=1e9):"""Args:c_in: