Khung Kiến Trúc Truyền Thông Điệp Tổng Quát MPNN (Gilmer et al. ICML 2017)

Mục tiêu kiến thức

• Gilmer et al. (2017) thống nhất hầu hết các mô hình GNN (GCN, GAT, GraphSAGE) vào một khung tổng quát gọi là Message Passing Neural Networks (MPNN) gồm 3 giai đoạn tại mỗi bước t:
• Tạo thông điệp (Message Function M_t):

m_ij^(t) = M_t(h_i^(t), h_j^(t), e_ij) (Thông điệp gửi từ đỉnh láng giềng j tới đỉnh i phụ thuộc vào trạng thái 2 đỉnh và đặc trưng cạnh e_ij).

• Gom cụm thông điệp (Aggregation Function bigoplus):

m_i^(t) = bigoplus_{j ∈ Ni} m_ij^(t) (Phải là một toán tử bất biến hoán vị như ∑ hoặc mean).

• Cập nhật trạng thái (Update Function U_t):

h_i^(t+1) = U_t(h_i^(t), m_i^(t)) (Thường dùng mạng MLP hoặc tế bào GRU).

Yêu cầu

Xây dựng lớp MPNNLayer(torch.nn.Module):

class MPNNLayer(nn.Module):
    def __init__(self, node_dim: int, edge_dim: int, out_dim: int, aggr: str = "sum"):
        pass

    def forward(self, h: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor) -> torch.Tensor:
        # h: (N, node_dim)
        # edge_index: (2, E) where edge_index[0] = source, edge_index[1] = target
        # edge_attr: (E, edge_dim)
        # Returns: (N, out_dim)
        pass

• Message: M([h_src, h_dst, edge_attr]) qua một Linear layer (node_dim * 2 + edge_dim -> out_dim).
• Aggregate: Dùng scatter_add hoặc lặp để gom thông điệp vào đỉnh đích target = edge_index[1].
• Update: U([h_dst, m_aggregated]) qua Linear layer (node_dim + out_dim -> out_dim).

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.