[论文解读] Lipschitz Normalization for Self-Attention Layers with Application to Graph Neural Networks

George Dasoulas, Kevin Scaman+1 more|arXiv (Cornell University)|Mar 8, 2021

Advanced Graph Neural Networks参考文献 35被引用 5

一句话总结

本文提出 LipschitzNorm，一种无需参数的归一化技术，通过在自注意力层中强制实现利普希茨连续性，防止深度注意力模型中的梯度爆炸。将其应用于图神经网络（GNNs）时，LipschitzNorm 使深层 GAT 和图 Transformer 架构能够实现稳定训练，在具有长距离依赖关系的节点分类任务上达到最先进性能，尤其在 30 层时，相比残差网络，准确率最高提升 6%。

ABSTRACT

Attention based neural networks are state of the art in a large range of applications. However, their performance tends to degrade when the number of layers increases. In this work, we show that enforcing Lipschitz continuity by normalizing the attention scores can significantly improve the performance of deep attention models. First, we show that, for deep graph attention networks (GAT), gradient explosion appears during training, leading to poor performance of gradient-based training algorithms. To address this issue, we derive a theoretical analysis of the Lipschitz continuity of attention modules and introduce LipschitzNorm, a simple and parameter-free normalization for self-attention mechanisms that enforces the model to be Lipschitz continuous. We then apply LipschitzNorm to GAT and Graph Transformers and show that their performance is substantially improved in the deep setting (10 to 30 layers). More specifically, we show that a deep GAT model with LipschitzNorm achieves state of the art results for node label prediction tasks that exhibit long-range dependencies, while showing consistent improvements over their unnormalized counterparts in benchmark node classification tasks.

研究动机与目标

• 为解决深度注意力模型（尤其是图神经网络 GNNs）在训练过程中因梯度爆炸导致的性能退化问题。
• 理论分析标准自注意力机制的利普希茨连续性，识别其在深层架构中的不稳定性。
• 提出一种简单、无需参数的归一化方法——LipschitzNorm，用于在注意力层中强制实现利普希茨连续性。
• 通过实证验证，LipschitzNorm 能够稳定训练并提升深层 GNN 的性能，尤其在具有长距离依赖关系的任务中。
• 证明 LipschitzNorm 在深层 GNN 设置下优于现有归一化技术（如 PairNorm 和 LayerNorm）

提出的方法

• LipschitzNorm 通过缩放查询和键向量，使注意力机制实现利普希茨连续性，且其界为 1。
• 该方法利用注意力权重矩阵的谱范数推导出归一化因子，确保输出变化相对于输入变化受到有界限制。
• 它直接应用于点积注意力计算，通过如下方式修改注意力分数：\text{attn} = \frac{\text{softmax}(\text{QK}^T / \sqrt{d})}{\|\text{QK}^T / \sqrt{d}\|_2} \cdot \text{value}。
• 该归一化方法无需额外可学习参数，轻量化且与现有 GNN 架构完全兼容。
• 理论分析证明，LipschitzNorm 确保注意力层为 1-利普希茨连续，从而在反向传播中稳定梯度流动。
• 该方法被应用于 GAT 和图 Transformer，其有效性在多个基准数据集上随深度增加而得到评估。

实验结果

研究问题

• RQ1 在自注意力层中强制实现利普希茨连续性是否能防止深度注意力模型中的梯度爆炸？
• RQ2 LipschitzNorm 是否能稳定深层图神经网络的训练并提升其性能，特别是在具有长距离依赖关系的模型中？
• RQ3 在深层 GNN 架构中，LipschitzNorm 与 PairNorm 和 LayerNorm 等现有归一化技术相比表现如何？
• RQ4 LipschitzNorm 是否能与残差连接有效结合，从而在深层 GNN 中进一步提升性能？
• RQ5 LipschitzNorm 是否能使更深的 GNN（如 30 层）在节点分类任务上实现最先进性能？

主要发现

• LipschitzNorm 实现了 GAT 和图 Transformer 模型在高达 30 层的稳定训练，而未归一化的模型在此深度下无法收敛。
• 在 Cora 数据集上，30 层的 GAT-Lip 模型达到 69.4% 的测试准确率，相比 GAT-res（63.5%）高出 5.9 个百分点。
• 在 PubMed 数据集上，30 层的 GAT-Lip 模型达到 67.2% 的准确率，相比未归一化的 GAT（28.2%）提升 15.8%。
• 在 Ogbn-proteins 数据集上，LipschitzNorm 在 30 层内保持了高水平的 ROC-AUC 性能，而未归一化的模型则完全崩溃。
• LipschitzNorm 在所有四个基准数据集（Cora、PubMed、Ogbn-arxiv、Ogbn-proteins）上均优于 PairNorm 和 LayerNorm，在深层设置下表现出一致的性能提升。
• 将 LipschitzNorm 与残差连接结合可带来微小的进一步性能提升，表明其在深层架构中具有良好的兼容性与鲁棒性。