[논문 리뷰] Stochastic Aggregation in Graph Neural Networks

Yuanqing Wang, Theofanis Karaletsos|arXiv (Cornell University)|2021. 02. 25.

Advanced Graph Neural Networks참고 문헌 24인용 수 7

한 줄 요약

이 논문은 그래프 신경망(GNNs)에서 메시지 전파 동안 엣지 가중치에 적응형 노이즈를 주입하는 통합 프레임워크인 Stochastic Aggregation(STAG)을 제안한다. 이는 표현력 향상과 과도한 스무딩 방지에 기여하며, 인용 및 분자 그래프 벤치마크에서 성능 향상을 이룬다. 변분 추론 변종은 최소한의 계산 오버헤드로 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Graph neural networks (GNNs) manifest pathologies including over-smoothing and limited discriminating power as a result of suboptimally expressive aggregating mechanisms. We herein present a unifying framework for stochastic aggregation (STAG) in GNNs, where noise is (adaptively) injected into the aggregation process from the neighborhood to form node embeddings. We provide theoretical arguments that STAG models, with little overhead, remedy both of the aforementioned problems. In addition to fixed-noise models, we also propose probabilistic versions of STAG models and a variational inference framework to learn the noise posterior. We conduct illustrative experiments clearly targeting oversmoothing and multiset aggregation limitations. Furthermore, STAG enhances general performance of GNNs demonstrated by competitive performance in common citation and molecule graph benchmark datasets.

연구 동기 및 목표

• 정규화되지 않은 아그리게이션 메커니즘으로 인한 GNN의 과도한 스무딩과 표현력 제한 문제를 해결하기 위해.
• Dropout, DropEdge, Graph DropConnect와 같은 기존 정규화 기법들을 하나의 스토하스틱 프레임워크로 통합하기 위해.
• 노이즈 분포의 파라미터를 모델 파라미터와 함께 공동으로 학습하는 변분 추론 기반 STAG 변종을 개발하기 위해.
• 표준 GNN 벤치마크(인용 및 분자 그래프 포함)에서 STAG의 효과를 경험적으로 검증하기 위해.
• 엣지 가중치의 스토하스틱성은 유의미한 계산 비용 증가 없이 일반화 능력과 모델의 강건성을 향상시킬 수 있음을 보여주기 위해.

제안 방법

• STAG는 각 메시지 전파 단계에서 노이즈 분포에서 샘플링한 스토하스틱 엣지 가중치로 정규화되지 않은 아그리게이션을 대체한다.
• 노이즈는 연속 분포(예: Normal(1,1)) 또는 이산 분포(예: Bernoulli)를 사용해 엣지 가중치를 편향시켜 이웃 메시지의 재가중치를 가능하게 한다.
• 학습 시, 몬테카를로 샘플링을 통해 기울기를 추정하여 편향 없이 엔드 투 엔드 백프로파게이션을 가능하게 한다.
• 변분 추론(VI) 프레임워크를 도입하여 노이즈 분포의 파라미터를 변분 파라미터로 학습하며, 노드 특징과 그래프 구조에 조건부로 설정될 수 있다.
• 추론 시 예측 사후분포는 학습된 엣지 가중치 분포에 대해 마진할라이징하여 확보된다.
• 프레임워크는 파이토치 및 DGL에서 효율적으로 구현되었으며, 메시지 전파 연산에 단 한 줄의 수정만으로 구현 가능하다.

실험 결과

연구 질문

• RQ1 GNN에서 엣지 가중치에 스토하스틱 편향을 주입함으로써 정규화된 아그리게이션에 비해 표현력 향상과 과도한 스무딩 감소가 가능한가?
• RQ2 Dropout, DropEdge, Graph DropConnect와 같은 기존 정규화 방법과 비교할 때 STAG의 성능 및 일반화 능력은 어떠한가?
• RQ3 노이즈 파라미터 학습을 위한 변분 추론 프레임워크가 다양한 그래프 구조 작업에서 GNN 성능 향상에 기여할 수 있는가?
• RQ4 특히 정규화가 적용되지 않은 경우, 연속 노이즈 분포(예: Normal)가 이산 분포(예: Bernoulli)보다 STAG에서 더 뛰어난 성능을 내는가?
• RQ5 STAG는 인용 네트워크와 분자 그래프와 같은 다양한 그래프 유형으로 일반화되어 일관된 성능 향상을 보일 수 있는가?

주요 결과

• 연속 노이즈 분포(예: Normal(1,1))를 사용한 STAG는 Cora와 Citeseer에서 정규화된 기준 모델 및 이산 노이즈 변종보다 뛰어난 성능을 보이며, 여러 실행에서 일관된 향상을 보였다.
• 변분 추론 변종인 STAG_VI는 인용 및 분자 그래프 벤치마크에서 비적응형 STAG 및 Graph DropConnect보다 일관되게 뛰어난 성능을 보였다.
• 엣지 및 특징별로 노이즈 파라미터를 학습하는 가장 표현력 있는 STAG_VI 변종은 ESOL 및 FreeSolv 데이터셋에서 최신 기술 수준의 성능을 달성했다.
• V100 GPU에서 STAG는 각 전방향 추론 단계에서 5.9에서 9.3ms의 추가 시간만 소요되어 경량임을 입증했다.
• 정규화 연산은 연속 노이즈 분포와 함께 사용될 경우 성능을 떨어뜨리므로, STAG의 노이즈 주입 방식은 이러한 보정 없이도 본질적으로 안정적임을 시사한다.
• 프레임워크는 그래프 유형 간에 잘 일반화되어 있으며, 사회 네트워크 및 분자 그래프 데이터셋 모두에서 일관된 성능 향상을 보였다.