[论文解读] Graph Masked Autoencoders with Transformers

Sixiao Zhang, Hongxu Chen+4 more|arXiv (Cornell University)|Feb 17, 2022

Advanced Graph Neural Networks被引用 11

一句话总结

该论文提出图掩码自编码器（GMAE），一种基于Transformer的自监督模型，采用掩码机制与非对称编码器-解码器架构，高效学习图表示。通过掩码高达70%的节点并重建其特征，GMAE在图分类与节点分类任务中达到最先进性能，同时相比传统Transformer将内存消耗降低最多20%，从而支持更大规模图与更深模型的训练。

ABSTRACT

Recently, transformers have shown promising performance in learning graph representations. However, there are still some challenges when applying transformers to real-world scenarios due to the fact that deep transformers are hard to train from scratch and the quadratic memory consumption w.r.t. the number of nodes. In this paper, we propose Graph Masked Autoencoders (GMAEs), a self-supervised transformer-based model for learning graph representations. To address the above two challenges, we adopt the masking mechanism and the asymmetric encoder-decoder design. Specifically, GMAE takes partially masked graphs as input, and reconstructs the features of the masked nodes. The encoder and decoder are asymmetric, where the encoder is a deep transformer and the decoder is a shallow transformer. The masking mechanism and the asymmetric design make GMAE a memory-efficient model compared with conventional transformers. We show that, when serving as a conventional self-supervised graph representation model, GMAE achieves state-of-the-art performance on both the graph classification task and the node classification task under common downstream evaluation protocols. We also show that, compared with training in an end-to-end manner from scratch, we can achieve comparable performance after pre-training and fine-tuning using GMAE while simplifying the training process.

研究动机与目标

• 解决在图上训练深度Transformer所面临的挑战，包括高内存消耗与从零开始训练的困难。
• 在图表示学习中超越图卷积网络（GCNs）的局限性，提升全局信息捕获能力，后者在深层架构中易出现过平滑问题。
• 开发一种自监督预训练框架，简化下游微调过程，同时保持与端到端训练相当的性能。
• 通过节点掩码与非对称设计，降低Transformer在图设置下的二次方内存复杂度。
• 实现对多样化图基准数据集具有良好泛化能力的图Transformer有效预训练。

提出的方法

• GMAE采用掩码自编码目标，其中大量节点（高达70%）被掩码，模型重建被掩码节点的原始特征。
• 模型采用非对称架构：深层12层Transformer编码器处理掩码图，浅层2层Transformer解码器重建被掩码节点的特征。
• 掩码机制减少了自注意力计算中的有效节点数，将内存使用量从O(n²)降低至O(0.49n²)，当掩码比率为70%时。
• 在图数据上进行自监督预训练可实现迁移学习，使下游任务微调时可使用更高的初始学习率并减少训练轮数。
• 模型通过被掩码节点特征的重建损失进行训练，同时优化节点级别与图级别表示的质量。
• 通过实验调优掩码比率与解码器深度等超参数，以在性能与效率之间取得平衡。

实验结果

研究问题

• RQ1 采用掩码自编码方法与非对称Transformer是否能在不牺牲性能的前提下，降低图表示学习中的内存消耗？
• RQ2 掩码比率如何影响GMAE模型在多样化图数据集上的下游性能？
• RQ3 与对称Transformer相比，使用深层编码器与浅层解码器是否能提升性能与训练效率？
• RQ4 GMAE预训练能否在简化训练流程的同时，实现与端到端训练相当的性能？
• RQ5 在准确率与内存效率方面，GMAE与SOTA模型如Graphormer相比表现如何？

主要发现

• 在标准自监督评估协议下，GMAE在7个基准图数据集（PROTEINS, BZR, COX2, ER_MD, NCI1, MUTAG, DHFR）上，于图分类与节点分类任务中均达到最先进性能。
• 当掩码比率为70%时，GMAE相比传统Transformer将内存消耗降低约20%，尽管理论上自注意力计算的节省最高可达50%。
• 模型在掩码比约为0.4或0.7时表现最佳，表明当一半节点被掩码时仍具鲁棒性，显著降低计算负载。
• 1–2层的浅层解码器优于更深的解码器，验证了非对称编码器-解码器设计在效率与性能方面的有效性。
• 使用GMAE进行预训练可使微调阶段采用更高的学习率与更少的训练轮数，性能与端到端训练相当，同时简化了训练流程。
• 在部分基准测试中，GMAE优于Graphormer，且无需高度定制化的训练方案，展现出强大的泛化能力与易用性。