EGNN 等变图神经网络

EGNN 的全称是 E(n)-Equivariant Graph Neural Network，即对 $n$ 维欧氏空间中的平移、旋转、反射等变的图神经网络。它常用于分子建模、蛋白质结构建模、粒子系统模拟等任务。

普通 GNN 主要处理节点特征和边关系，例如分子中的原子类型、化学键、邻接关系。但在很多科学问题中，节点还带有空间坐标：

x_{i} \in R^{n}

例如三维分子中，第 $i$ 个原子的坐标为：

x_{i} = (x_{i}, y_{i}, z_{i}) \in R^{3}

这时模型不仅要理解图结构，还要正确处理几何变换。

1. 为什么需要等变性

假设一个分子整体旋转或平移后，它的物理性质不应该改变。例如分子的能量、毒性、溶解度等标量性质，与分子放在空间中的绝对位置和朝向无关。

如果把所有原子坐标做平移：

x_{i}^{'} = x_{i} + t

或者做旋转、反射：

x_{i}^{'} = Q x_{i}

其中 $Q$ 是正交矩阵，满足：

Q^{T} Q = I

那么模型应该满足两类性质：

不变性

如果预测的是分子整体性质，例如能量 $y$ ，则输出不应随坐标系变化而变化：

f (Q X + t, H, A) = f (X, H, A)

这叫 invariance，不变性。

等变性

如果预测的是坐标、速度、力等向量量，则输出应该随输入坐标一起变换：

F (Q X + t, H, A) = Q F (X, H, A) + t

或者对于力这种不受平移影响的向量：

F (Q X + t, H, A) = Q F (X, H, A)

这叫 equivariance，等变性。

简单理解：

性质	含义	例子
不变性	输入变换后，输出不变	分子能量、分类标签
等变性	输入变换后，输出按同样方式变化	坐标、速度、力

2. 普通 GNN 的消息传递

普通 GNN 的一层更新通常可以写为：

m_{i j}^{(l)} = ϕ_{e} (h_{i}^{(l)}, h_{j}^{(l)}, e_{i j})

m_{i}^{(l)} = \sum_{j \in N (i)} m_{i j}^{(l)}

h_{i}^{(l + 1)} = ϕ_{h} (h_{i}^{(l)}, m_{i}^{(l)})

其中：

符号	含义
$h_{i}^{(l)}$	第 $l$ 层中节点 $i$ 的特征
$e_{i j}$	边 $(i, j)$ 的边特征
$N (i)$	节点 $i$ 的邻居集合
$ϕ_{e}$	边消息函数，通常是 MLP
$ϕ_{h}$	节点更新函数，通常是 MLP

如果直接把坐标 $x_{i}$ 拼到节点特征里：

{\tilde{h}}_{i} = [h_{i}, x_{i}]

模型就可能学到依赖绝对坐标轴的规律。例如同一个分子旋转后，模型输出可能变化。这不符合物理系统的几何对称性。

3. EGNN 的核心思想

EGNN 的关键想法是：

消息函数只使用距离等不变量，而不是直接使用坐标方向。
坐标更新只沿着相对方向 $x_{i} - x_{j}$ 进行。
用标量函数控制相对方向的更新幅度。

这样模型既可以更新节点特征，也可以更新坐标，同时保持对欧氏变换的等变性。

设图为：

G = (V, E)

每个节点有：

h_{i}^{(l)} \in R^{d}

和坐标：

x_{i}^{(l)} \in R^{n}

EGNN 一层通常由三部分组成：

边消息计算
坐标更新
节点特征更新

4. EGNN 的数学表达式

4.1 边消息

对边 $(i, j)$ ，先计算节点间距离平方：

d_{i j}^{(l)} = ∥ x_{i}^{(l)} - x_{j}^{(l)} ∥^{2}

然后生成边消息：

m_{i j}^{(l)} = ϕ_{e} (h_{i}^{(l)}, h_{j}^{(l)}, d_{i j}^{(l)}, e_{i j})

其中 $ϕ_{e}$ 是一个 MLP，输出是标量或向量消息。

注意这里使用的是距离平方 $d_{i j}$ ，而不是原始坐标 $x_{i}$ 、 $x_{j}$ 。距离在旋转、反射和平移下都不变。

4.2 坐标更新

EGNN 使用相对坐标方向更新节点坐标：

x_{i}^{(l + 1)} = x_{i}^{(l)} + C \sum_{j \neq i} (x_{i}^{(l)} - x_{j}^{(l)}) ϕ_{x} (m_{i j}^{(l)})

其中：

符号	含义
$ϕ_{x}$	坐标更新函数，通常是输出 1 维标量的 MLP
$C$	归一化常数，常取 $\frac{1}{
$x_{i} - x_{j}$	节点间相对方向

也可以只对邻居求和：

x_{i}^{(l + 1)} = x_{i}^{(l)} + \frac{1}{| N (i) |} \sum_{j \in N (i)} (x_{i}^{(l)} - x_{j}^{(l)}) ϕ_{x} (m_{i j}^{(l)})

这里 $ϕ_{x} (m_{i j})$ 是一个标量，它只决定沿着相对方向移动多少，不会引入新的坐标方向。

4.3 节点消息聚合

把边消息聚合到节点上：

m_{i}^{(l)} = \sum_{j \in N (i)} m_{i j}^{(l)}

4.4 节点特征更新

更新节点特征：

h_{i}^{(l + 1)} = ϕ_{h} (h_{i}^{(l)}, m_{i}^{(l)})

其中 $ϕ_{h}$ 通常也是 MLP。

如果需要残差连接，也可以写成：

h_{i}^{(l + 1)} = h_{i}^{(l)} + ϕ_{h} (h_{i}^{(l)}, m_{i}^{(l)})

5. EGNN 的等变性推导

下面证明 EGNN 对平移、旋转和反射具有等变性。

设对所有坐标做欧氏变换：

x_{i}^{'} = Q x_{i} + t

其中：

Q^{T} Q = I

$Q$ 可以表示旋转或反射， $t$ 表示平移向量。

5.1 距离不变

EGNN 的边消息依赖距离平方：

d_{i j} = ∥ x_{i} - x_{j} ∥^{2}

变换后：

\begin{aligned} d_{i j}^{'} & = ∥ x_{i}^{'} - x_{j}^{'} ∥^{2} \\ = ∥ (Q x_{i} + t) - (Q x_{j} + t) ∥^{2} \\ = ∥ Q (x_{i} - x_{j}) ∥^{2} \\ = (x_{i} - x_{j})^{T} Q^{T} Q (x_{i} - x_{j}) \\ = (x_{i} - x_{j})^{T} (x_{i} - x_{j}) \\ = ∥ x_{i} - x_{j} ∥^{2} \\ = d_{i j} \end{aligned}

所以距离平方在平移、旋转和反射下都不变。

因此边消息也不变：

m_{i j}^{'} = m_{i j}

因为：

m_{i j}^{'} = ϕ_{e} (h_{i}, h_{j}, d_{i j}^{'}, e_{i j}) = ϕ_{e} (h_{i}, h_{j}, d_{i j}, e_{i j}) = m_{i j}

5.2 坐标更新等变

原始坐标更新为：

x_{i}^{+} = x_{i} + C \sum_{j \in N (i)} (x_{i} - x_{j}) ϕ_{x} (m_{i j})

其中 $x_{i}^{+}$ 表示更新后的坐标。

变换后的坐标更新为：

\begin{aligned} {x_{i}^{'}}^{+} & = x_{i}^{'} + C \sum_{j \in N (i)} (x_{i}^{'} - x_{j}^{'}) ϕ_{x} (m_{i j}^{'}) \\ = (Q x_{i} + t) + C \sum_{j \in N (i)} ((Q x_{i} + t) - (Q x_{j} + t)) ϕ_{x} (m_{i j}) \\ = Q x_{i} + t + C \sum_{j \in N (i)} Q (x_{i} - x_{j}) ϕ_{x} (m_{i j}) \\ = Q (x_{i} + C \sum_{j \in N (i)} (x_{i} - x_{j}) ϕ_{x} (m_{i j})) + t \\ = Q x_{i}^{+} + t \end{aligned}

因此坐标更新满足：

{x_{i}^{'}}^{+} = Q x_{i}^{+} + t

这说明 EGNN 的坐标输出对欧氏变换是等变的。

5.3 节点特征不变

节点特征更新依赖：

h_{i}^{+} = ϕ_{h} (h_{i}, m_{i})

而 $m_{i j}$ 不变，所以聚合消息：

m_{i} = \sum_{j \in N (i)} m_{i j}

也不变。因此节点特征更新结果不变：

{h_{i}^{'}}^{+} = h_{i}^{+}

所以 EGNN 中：

变量	变换性质
坐标 $x_{i}$	等变
节点特征 $h_{i}$	不变
图级标量输出	不变

6. EGNN 与普通 GNN 的区别

模型	是否使用坐标	是否保证几何等变	常见任务
GCN	否	否	节点分类、图分类
GAT	否	否	图结构数据建模
普通 3D GNN	是	不一定	分子性质预测
EGNN	是	是	分子建模、粒子模拟、蛋白结构

EGNN 的优势是结构相对简单，不需要使用球谐函数、张量表示等复杂几何工具，就能实现 $E (n)$ 等变性。

7. 一个简单的 PyTorch EGNN 示例

下面给出一个最小可运行的 EGNN 模型，不依赖 PyG。任务是：输入一批小分子的原子特征和三维坐标，预测一个图级标量性质。

为了演示方便，假设每个图都有相同数量的节点，并使用全连接图。

7.1 导入库

python

import torch
import torch.nn as nn

7.2 构建 EGNN 层

python

class EGNNLayer(nn.Module):
    def __init__(self, hidden_dim, edge_dim=0):
        super().__init__()

        self.edge_mlp = nn.Sequential(
            nn.Linear(hidden_dim * 2 + 1 + edge_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
        )

        self.coord_mlp = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, 1),
        )

        self.node_mlp = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )

    def forward(self, h, x, edge_index, edge_attr=None):
        """
        h: [num_nodes, hidden_dim] 节点特征
        x: [num_nodes, 3] 节点坐标
        edge_index: [2, num_edges] 边索引，第一行为源节点，第二行为目标节点
        edge_attr: [num_edges, edge_dim] 可选边特征
        """
        row, col = edge_index

        h_i = h[row]
        h_j = h[col]
        x_i = x[row]
        x_j = x[col]

        relative = x_i - x_j
        dist2 = (relative ** 2).sum(dim=-1, keepdim=True)

        if edge_attr is None:
            edge_input = torch.cat([h_i, h_j, dist2], dim=-1)
        else:
            edge_input = torch.cat([h_i, h_j, dist2, edge_attr], dim=-1)

        m_ij = self.edge_mlp(edge_input)

        coord_weight = self.coord_mlp(m_ij)
        coord_message = relative * coord_weight

        dx = torch.zeros_like(x)
        dx.index_add_(0, row, coord_message)

        degree = torch.zeros(x.size(0), 1, device=x.device)
        degree.index_add_(0, row, torch.ones_like(coord_weight))
        dx = dx / degree.clamp(min=1.0)

        x = x + dx

        m_i = torch.zeros_like(h)
        m_i.index_add_(0, row, m_ij)

        h = h + self.node_mlp(torch.cat([h, m_i], dim=-1))

        return h, x

这层中最重要的是：

python

relative = x_i - x_j
dist2 = (relative ** 2).sum(dim=-1, keepdim=True)
coord_message = relative * coord_weight

dist2 是旋转和平移不变的量，relative 是随坐标一起旋转的向量。因此坐标更新可以保持等变性。

7.3 构建完整模型

python

class EGNNModel(nn.Module):
    def __init__(self, node_dim, hidden_dim, num_layers):
        super().__init__()

        self.embedding = nn.Linear(node_dim, hidden_dim)

        self.layers = nn.ModuleList([
            EGNNLayer(hidden_dim)
            for _ in range(num_layers)
        ])

        self.readout = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, h, x, edge_index, batch):
        """
        h: [num_nodes, node_dim]
        x: [num_nodes, 3]
        edge_index: [2, num_edges]
        batch: [num_nodes]，表示每个节点属于哪个图
        """
        h = self.embedding(h)

        for layer in self.layers:
            h, x = layer(h, x, edge_index)

        num_graphs = int(batch.max().item()) + 1
        graph_h = torch.zeros(num_graphs, h.size(-1), device=h.device)
        graph_h.index_add_(0, batch, h)

        count = torch.zeros(num_graphs, 1, device=h.device)
        count.index_add_(0, batch, torch.ones(h.size(0), 1, device=h.device))
        graph_h = graph_h / count.clamp(min=1.0)

        y = self.readout(graph_h)
        return y

这里使用平均池化把节点表示变成图表示：

h_{G} = \frac{1}{| V |} \sum_{i \in V} h_{i}

然后用 MLP 输出图级标量。

7.4 构造一个简单数据集

下面构造两个小图，每个图有 4 个节点。为了简单，所有节点之间都连边，但不包含自环。

python

def fully_connected_edges(num_nodes, device):
    row = []
    col = []
    for i in range(num_nodes):
        for j in range(num_nodes):
            if i != j:
                row.append(i)
                col.append(j)
    return torch.tensor([row, col], dtype=torch.long, device=device)


device = "cuda" if torch.cuda.is_available() else "cpu"

num_graphs = 2
nodes_per_graph = 4
num_nodes = num_graphs * nodes_per_graph
node_dim = 5

h = torch.randn(num_nodes, node_dim, device=device)
x = torch.randn(num_nodes, 3, device=device)

edge_indices = []
for g in range(num_graphs):
    edge_index_g = fully_connected_edges(nodes_per_graph, device)
    edge_index_g = edge_index_g + g * nodes_per_graph
    edge_indices.append(edge_index_g)

edge_index = torch.cat(edge_indices, dim=1)

batch = torch.arange(num_graphs, device=device).repeat_interleave(nodes_per_graph)

target = torch.randn(num_graphs, 1, device=device)

7.5 训练模型

python

model = EGNNModel(
    node_dim=node_dim,
    hidden_dim=64,
    num_layers=3,
).to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

for step in range(200):
    pred = model(h, x, edge_index, batch)
    loss = loss_fn(pred, target)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 20 == 0:
        print(f"step={step}, loss={loss.item():.4f}")

这是一个教学用例。真实分子性质预测中，一般需要：

使用真实原子特征，例如原子序数、电荷、杂化状态。
使用真实边特征，例如键类型、键长、是否芳香键。
使用批处理数据加载器。
使用训练集、验证集、测试集划分。
对坐标更新幅度做更谨慎的归一化或裁剪。

8. 验证等变性的简单实验

下面用随机正交矩阵 $Q$ 和平移向量 $t$ 检查模型输出的图级标量是否基本不变。

python

def random_rotation_matrix(device):
    A = torch.randn(3, 3, device=device)
    Q, R = torch.linalg.qr(A)
    return Q


Q = random_rotation_matrix(device)
t = torch.randn(1, 3, device=device)

x_transformed = x @ Q.T + t

model.eval()
with torch.no_grad():
    y1 = model(h, x, edge_index, batch)
    y2 = model(h, x_transformed, edge_index, batch)

print("original prediction:")
print(y1)
print("transformed prediction:")
print(y2)
print("max difference:", (y1 - y2).abs().max().item())

理论上，如果模型结构和数值计算完全理想，则图级标量输出应保持不变。实际浮点计算中可能会有很小误差。

9. 小结

EGNN 的核心可以概括为：

m_{i j} = ϕ_{e} (h_{i}, h_{j}, ∥ x_{i} - x_{j} ∥^{2}, e_{i j})

x_{i}^{+} = x_{i} + \sum_{j \in N (i)} (x_{i} - x_{j}) ϕ_{x} (m_{i j})

h_{i}^{+} = ϕ_{h} (h_{i}, \sum_{j \in N (i)} m_{i j})

它通过距离构造不变量，通过相对坐标构造等变向量，因此可以自然处理三维结构数据。

对于分子和蛋白质这类对象，EGNN 的意义在于：模型不需要重新学习“旋转和平移不应该改变物理规律”这件事，而是从结构上就满足这种几何约束。

编码-解码器架构及其变种

Diffusion Model

EGNN 等变图神经网络

1. 为什么需要等变性

不变性

等变性

2. 普通 GNN 的消息传递

3. EGNN 的核心思想

4. EGNN 的数学表达式

4.1 边消息

4.2 坐标更新

4.3 节点消息聚合

4.4 节点特征更新

5. EGNN 的等变性推导

5.1 距离不变

5.2 坐标更新等变

5.3 节点特征不变

6. EGNN 与普通 GNN 的区别

7. 一个简单的 PyTorch EGNN 示例

7.1 导入库

7.2 构建 EGNN 层

7.3 构建完整模型

7.4 构造一个简单数据集

7.5 训练模型

8. 验证等变性的简单实验

9. 小结

EGNN 等变图神经网络 ​

1. 为什么需要等变性 ​

不变性 ​

等变性 ​

2. 普通 GNN 的消息传递 ​

3. EGNN 的核心思想 ​

4. EGNN 的数学表达式 ​

4.1 边消息 ​

4.2 坐标更新 ​

4.3 节点消息聚合 ​

4.4 节点特征更新 ​

5. EGNN 的等变性推导 ​

5.1 距离不变 ​

5.2 坐标更新等变 ​

5.3 节点特征不变 ​

6. EGNN 与普通 GNN 的区别 ​

7. 一个简单的 PyTorch EGNN 示例 ​

7.1 导入库 ​

7.2 构建 EGNN 层 ​

7.3 构建完整模型 ​

7.4 构造一个简单数据集 ​

7.5 训练模型 ​

8. 验证等变性的简单实验 ​

9. 小结 ​

EGNN 等变图神经网络

1. 为什么需要等变性

不变性

等变性

2. 普通 GNN 的消息传递

3. EGNN 的核心思想

4. EGNN 的数学表达式

4.1 边消息

4.2 坐标更新

4.3 节点消息聚合

4.4 节点特征更新

5. EGNN 的等变性推导

5.1 距离不变

5.2 坐标更新等变

5.3 节点特征不变

6. EGNN 与普通 GNN 的区别

7. 一个简单的 PyTorch EGNN 示例

7.1 导入库

7.2 构建 EGNN 层

7.3 构建完整模型

7.4 构造一个简单数据集

7.5 训练模型

8. 验证等变性的简单实验

9. 小结