[논문 리뷰] Graph-level Representation Learning with Joint-Embedding Predictive Architectures

Geri Skenderi, Hang Li외 2명|arXiv (Cornell University)|2023. 09. 27.

Advanced Graph Neural NetworksComputer Science인용 수 3

한 줄 요약

이 논문은 그래프 수준의 자기지도 학습을 위한 최초의 공동 임베딩 예측 아키텍처인 Graph-JEPA를 제안한다. 데이터 증강이나 음성 샘플을 사용하지 않고도 막힌 서브그래프 임베딩을 문맥 서브그래프를 기반으로 예측하며, 계층적인 그래프 개념을 포착하는 새로운 쌍곡좌표 예측 목표함수를 통해 그래프 분류 및 회귀 작업에서 경쟁적인 성능을 달성한다.

ABSTRACT

Joint-Embedding Predictive Architectures (JEPAs) have recently emerged as a novel and powerful technique for self-supervised representation learning. They aim to learn an energy-based model by predicting the latent representation of a target signal y from the latent representation of a context signal x. JEPAs bypass the need for negative and positive samples, traditionally required by contrastive learning while avoiding the overfitting issues associated with generative pretraining. In this paper, we show that graph-level representations can be effectively modeled using this paradigm by proposing a Graph Joint-Embedding Predictive Architecture (Graph-JEPA). In particular, we employ masked modeling and focus on predicting the latent representations of masked subgraphs starting from the latent representation of a context subgraph. To endow the representations with the implicit hierarchy that is often present in graph-level concepts, we devise an alternative prediction objective that consists of predicting the coordinates of the encoded subgraphs on the unit hyperbola in the 2D plane. Through multiple experimental evaluations, we show that Graph-JEPA can learn highly semantic and expressive representations, as shown by the downstream performance in graph classification, regression, and distinguishing non-isomorphic graphs. The code is available at https://github.com/geriskenderi/graph-jepa.

연구 동기 및 목표

• 대조적 및 생성적 자기지도 학습의 한계를 해결하기 위해, 데이터 증강 및 음성 샘플에 의존하거나 분포 추정으로 인한 과적합 문제를 해결한다.
• 잠재 공간에서 작동하며 명시적 재구성 또는 대조 목표를 피하는 그래프 수준 표현 학습을 위한 JEPA 기반 프레임워크를 개발한다.
• 2차원 잠재 공간에서 단위 쌍곡선 상의 서브그래프 좌표를 예측하여 계층적인 그래프 수준 개념을 포착하는 예측 목표함수를 설계한다.
• 표준 그래프 벤치마크 데이터셋을 활용해 그래프 분류 및 회귀 작업에서 Graph-JEPA의 효과성을 평가한다.
• 위치 인코딩, 어텐션 메커니즘, 서브그래프 샘플링 전략과 같은 핵심 설계 선택 사항을 아블레이션한다.

제안 방법

• Graph-JEPA는 입력 그래프에서 추출한 문맥 서브그래프를 기반으로 랜덤하게 선택된 타겟 서브그래프의 잠재 표현을 예측하기 위해 마스킹 모델링을 사용한다.
• 서브그래프 표현은 풀링되어 전반적인 그래프 수준의 임베딩을 형성하며, 이는 후속 분류 및 회귀 작업에 활용된다.
• 새로운 예측 목표함수가 도입되며, 이는 타겟 서브그래프 임베딩을 2차원 공간에서 단위 쌍곡선 상의 좌표로 인코딩하여 계층적인 그래프 개념을 모델링한다.
• 예측 모델은 문맥 서브그래프 표현을 타겟 서브그래프의 잠재 표현으로 매핑하는 예측망을 사용하며, 공동 임베딩 철학을 적용한다.
• 상대 웨이블릿 스펙트럼 인코딩(RWSE)을 위치 인코딩으로 사용하며, 노드 수준 대비 패치 수준 RWSE를 비교하는 아블레이션 연구를 수행한다.
• 그래프 구조 데이터에 대한 인덕티브 바이어스를 제공하기 위해 하다르드 자기어텐션을 사용하며, 표준 자기어텐션과의 아블레이션을 수행한다.

Figure 1: Illustration of the three main SSL approaches: (a) Joint-Embedding Architectures learn to create similar embeddings for inputs x and y that are compatible with each other and dissimilar embeddings for inputs that are not compatible. This compatibility is implemented in practice by creating

실험 결과

연구 질문

• RQ1 JEPA 기반 접근법이 데이터 증강이나 음성 샘플 없이도 효과적으로 그래프 수준 표현을 학습할 수 있는가?
• RQ2 단위 쌍곡선 상의 서브그래프 좌표를 예측하는 것이 전통적인 잠재 재구성보다 그래프 SSL 성능을 향상시키는가?
• RQ3 구조적(METIS) 대비 무작위 서브그래프 샘플링 전략이 모델 성능과 안정성에 미치는 영향은 어떠한가?
• RQ4 위치 인코딩 유형(노드 수준 대비 패치 수준 RWSE)이 Graph-JEPA의 성능에 미치는 영향은 무엇인가?
• RQ5 거리 함수 선택(Euclidean 대비 쌍곡)이 다양한 그래프 데이터셋에서 모델의 일반화 능력에 미치는 영향은 어떠한가?

주요 결과

• Graph-JEPA는 그래프 분류 작업에서 최신 기준 성능을 달성하였으며, MUTAG에서 91.25±5.75%의 정확도와 IMDB-B에서 73.68±3.24%의 정확도를 기록하여 대비 및 생성 기반 기준을 능가했다.
• ZINC 회귀 벤치마크에서 Graph-JEPA는 평균 절대 오차(MAE) 0.434±0.01을 달성하여, 다음으로 우수한 방법(0.471±0.01)보다 유의미하게 뛰어났다.
• 쌍곡 거리 함수를 사용할 경우 MUTAG와 REDDIT-M에서 성능 향상을 보였으며, 쌍곡 거리 사용 시 89.43±5.67%의 정확도를 기록하여 유클리드 거리 사용 시 87.04±6.01%보다 높았다.
• 아블레이션 연구 결과, 노드 수준(전역) RWSE가 패치 수준(상대적) RWSE보다 일관되게 뛰어나며, 겹치는 이웃 정보로 인해 패치 수준 RWSE는 더 높은 분산을 유발한다.
• 하드르드 자기어텐션을 표준 자기어텐션으로 대체할 경우 약간 더 뛰어난 성능를 보였지만, 안정성이 떨어져 인덕티브 바이어스가 훈련의 일관성을 향상시킨다는 것을 시사한다.
• 무작위 서브그래프 샘플링은 METIS 기반 샘플링과 유사한 성능를 보였지만 일반화 능력이 떨어지며, 다양한 그래프에서 더 뛰어난 인덕티브 바이어스와 안정성을 확보하기 위해 METIS를 권장한다.

Figure 2: An overview of Graph-JEPA. We first extract non-overlapping subgraphs (patches) (a.), perform a 1-hop neighborhood expansion (b.), and encode the subgraphs with a GNN (c.). After the subgraph encoding, one is randomly picked as the context and $m$ others as the targets (d.) and they are fe