Graph Neural Networks (GNNs) have emerged as a powerful representation learning framework for graph-structured data. A key limitation of conventional GNNs is their representation of each node with a singular feature vector, potentially overlooking intricate details about individual node features. Here, we propose an Attention-based Message-Passing layer for GNNs (AMPNet) that encodes individual features per node and models feature-level interactions through cross-node attention during message-passing steps. We demonstrate the abilities of AMPNet through extensive benchmarking on real-world biological systems such as fMRI brain activity recordings and spatial genomic data, improving over existing baselines by 20% on fMRI signal reconstruction, and further improving another 8% with positional embedding added. Finally, we validate the ability of AMPNet to uncover meaningful feature-level interactions through case studies on biological systems. We anticipate that our architecture will be highly applicable to graph-structured data where node entities encompass rich feature-level information.

翻译：图神经网络（GNN）已成为处理图结构数据的强大表示学习框架。传统GNN的一个关键局限在于每个节点仅用单一特征向量表示，可能忽略节点特征的细微细节。本文提出一种基于注意力机制的消息传递层（AMPNet），该层为每个节点编码独立特征，并在消息传递步骤中通过跨节点注意力建模特征层级交互。通过在真实生物系统（如fMRI脑活动记录和空间基因组数据）上的广泛基准测试，我们展示了AMPNet的能力：在fMRI信号重建任务中较现有基线方法提升20%，加入位置嵌入后额外提升8%。最后，通过生物系统案例研究验证了AMPNet揭示有意义的特征级交互的能力。我们预期该架构将高度适用于节点实体包含丰富特征层级信息的图结构数据。