Lipschitz Normalization for Self-Attention Layers with Application to Graph Neural Networks

2021/03/08 by George Dasoulas, Dasoulas, George, Kevin Scaman +3 · 7 citations
Computer Science · #Advanced Graph Neural Networks #Advanced Neural Network Applications #Adversarial Robustness in Machine Learning #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML)

paper · pdf · doi:10.48550/arxiv.2103.04886

openalex publication_date 2021/03/08 · openalex created_date 2021/03/15 · openalex updated_date 2026/07/28

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.

Citations

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• Invariant and Equivariant Graph Networks
• Detecting Functionality-Specific Vulnerabilities via Retrieving Individual Functionality-Equivalent APIs in Open-Source Repositories
• Open Graph Benchmark: Datasets for Machine Learning on Graphs
• Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification

Cited by

• ODE-ViT: Plug & Play Attention Layer from the Generalization of the ViT as an Ordinary Differential Equation
• Federated Attention: A Distributed Paradigm for Collaborative LLM Inference over Edge Networks
• Latent Refinement Decoding: Enhancing Diffusion-Based Language Models by Refining Belief States
• Pool Me Wisely: On the Effect of Pooling in Transformer-Based Models
• On the Convergence and Size Transferability of Continuous-depth Graph Neural Networks
• Understanding Sensitivity of Differential Attention through the Lens of Adversarial Robustness
• Pay Attention to Attention Distribution: A New Local Lipschitz Bound for Transformers

Related

• The Lipschitz Constant of Self-Attention
• Are GATs Out of Balance?
• On the Regularity of Attention
• A Regularized Attention Mechanism for Graph Attention Networks
• Understanding Attention and Generalization in Graph Neural Networks