[논문 리뷰] Lipschitz Normalization for Self-Attention Layers with Application to Graph Neural Networks

George Dasoulas, Kevin Scaman외 1명|arXiv (Cornell University)|2021. 03. 08.

Advanced Graph Neural Networks참고 문헌 35인용 수 5

한 줄 요약

이 논문은 자기주의 어텐션 레이어에서 리프시츠 연속성을 강제하여 깊은 어텐션 모델에서 기울기 폭주를 방지하는 파라미터 없는 정규화 기법인 LipschitzNorm을 소개한다. 그래프 신경망(GNN)에 적용했을 때, LipschitzNorm은 깊은 GAT 및 그래프 트랜스포머 아키텍처의 안정적인 학습을 가능하게 하며, 특히 장거리 의존성이 있는 노드 분류 작업에서 최신 기술 수준의 성능을 달성한다. 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.

연구 동기 및 목표

• 깊은 어텐션 모델, 특히 그래프 신경망(GNN)에서 학습 중 기울기 폭주로 인한 성능 저하 문제를 해결하기 위해.
• 기본 자기주의 어텐션 메커니즘의 리프시츠 연속성에 대해 이론적으로 분석하고 깊은 아키텍처에서의 불안정성을 규명하기 위해.
• 자기주의 어텐션 레이어에서 리프시츠 연속성을 강제하는 단순하고 파라미터 없는 정규화 방법인 LipschitzNorm을 제안하기 위해.
• LipschitzNorm이 깊은 GNN에서 학습을 안정화시키고 성능을 향상시키는 것을 실험적으로 검증하기 위해, 특히 장거리 의존성이 있는 작업에서의 성능 향상을 위해.
• 기존 정규화 기법들인 PairNorm 및 LayerNorm과 비교해 LipschitzNorm이 깊은 GNN 환경에서 더 우수한 성능을 보이는지 입증하기 위해.

제안 방법

• 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 및 그래프 트랜스포머 모두에 적용되었으며, 깊이가 증가하는 여러 벤치마크 데이터셋에서 성능이 평가되었다.

실험 결과

연구 질문

• RQ1 자기주의 어텐션 레이어에 리프시츠 연속성을 강제하면 깊은 어텐션 모델에서 기울기 폭주를 방지할 수 있는가?
• RQ2 LipschitzNorm은 깊은 그래프 신경망, 특히 장거리 의존성이 있는 경우에 학습을 안정화시키고 성능을 향상시킬 수 있는가?
• RQ3 LipschitzNorm은 깊은 GNN 아키텍처에서 기존 정규화 기법들인 PairNorm 및 LayerNorm과 비교해 어떻게 성능을 냈는가?
• RQ4 LipschitzNorm은 잔차 연결과 조합되어 깊은 GNN의 성능을 더욱 향상시킬 수 있는가?
• RQ5 LipschitzNorm은 30층과 같은 더 깊은 GNN 아키텍처를 통해 노드 분류 작업에서 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• LipschitzNorm은 30층에 이르는 깊은 GAT 및 그래프 트랜스포머 모델의 안정적 학습을 가능하게 하며, 정규화되지 않은 모델은 수렴에 실패한다.
• Cora 데이터셋에서 30층 GAT에 LipschitzNorm을 적용한 결과, 테스트 정확도가 69.4%에 도달하여 GAT-res(63.5%)보다 5.9%p 높게 나타났다.
• PubMed 데이터셋에서는 30층 GAT-Lip 모델이 67.2%의 정확도를 기록했으며, 정규화되지 않은 GAT(28.2%) 대비 15.8%p 향상되었다.
• Ogbn-proteins 데이터셋에서는 LipschitzNorm이 30층에 이르는 깊이에서도 높은 ROC-AUC 성능을 유지했으며, 정규화되지 않은 모델은 성능이 붕괴되었다.
• LipschitzNorm은 모든 네 가지 벤치마크 데이터셋(Cora, PubMed, Ogbn-arxiv, Ogbn-proteins)에서 PairNorm 및 LayerNorm을 모두 능가하며 일관된 성능 향상을 보였다.
• LipschitzNorm을 잔차 연결과 조합하면 약간의 추가 성능 향상이 이루어졌으며, 이는 깊은 아키텍처에서의 호환성과 강건성을 시사한다.