기술 요약: 일반 트랜스포머는 놀라울 정도로 강력한 링크 예측기이다

문제 정의

링크 예측은 그래프 머신러닝의 근본적인 과제이지만, 현재의 최첨단 접근 방식들은 상당한 확장성 및 일반화 문제를 안고 있다. 표준 그래프 신경망(GNN), 특히 메시 전달 신경망(MPNN)은 노드 중심의 집계에 의해 제한되며, 이는 종종 대칭적인 로컬 이웃을 가진 구조적으로 구별되는 노드 쌍을 구분하는 데 실패한다. 이를 극복하기 위해 기존 파이프라인은 흔히 명시적인 구조적 휴리스틱(예: 공통 이웃, 거리 인코딩), 글로벌 그래프 특징, 또는 학습 가능한 노드별 ID 임베딩에 의존한다. 이러한 접근 방식은 다음과 같은 몇 가지 결정적인 한계를 초래한다:

1. 확장성: 글로벌 캐시 또는 전체 그래프 패스에 의존하는 방법은 웹 규모의 그래프와 동적 업데이트 상황에서 어려움을 겪는다.
2. 배포 제약: ID 기반 임베딩은 새로운 노드가 추가될 때 비용이 많이 드는 재학습 또는 재계산을 요구하며, 이는 고정된 예산 내의 미니 배치 효율성 필요성을 위반한다.
3. 사전 지식에 대한 과도한 의지: 많은 모델이 일반화 가능한 패턴을 학습하기보다는 단순한 구조적 단서를 활용하여 성공하며, 하드웨어 효율성을 저해하는 복잡하고 수작업으로 제작된 구조적 인코딩(PE/SE)을 필요로 하는 경우가 많다.

그래프 트랜스포머(GT)는 전방향 어텐션(all-to-all attention)을 통해 이론적인 대안을 제시하지만, 엄격한 배포 제약이 필요한 링크 예측 작업에 부적합한 복잡한 구조적 인코딩(예: 스펙트럼 신호, 랜덤 워크)으로 인해 상당한 오버헤드가 발생한다.

방법론: PENCIL

저자들은 PENCIL(Plain ENCoder for Inferring Links)을 제안한다. 이는 수작업으로 제작된 사전 지식, 노드 ID, 또는 글로벌 전처리 없이 작동하도록 설계된 트랜스포머 기반 링크 예측기이다. PENCIL은 새로운 인코딩 방식과 아키텍처 설계를 통해 그래프 링크 예측에 적응된 표준 인코더 전용 BERT 스타일 아키텍처를 활용한다.

1. 그래프 인코딩 방식

PENCIL은 전체 그래프를 처리하는 대신, 후보 링크 $(v_{src}, v_{dst})$를 중심으로 샘플링된 고정 예산의 로컬 서브그래프 상에서 작동한다.

• 토큰화: 후보 쌍에 대해, 서브그래프 노드들은 $v_{src}$가 인덱스 0으로, $v_{dst}$가 인덱스 1로 고정된 인덱스를 할당받는다. 나머지 노드들은 무작위 인덱스를 할당받는다.
• 입력 표현: 각 토큰은 세 가지 결합된 구성 요소로 이루어진다:
  1. 샘플링된 순서 내에서의 노드 식별을 위한 원-핫 식별자(one-hot identifier).
  2. 추출된 서브그래프 내의 노드의 인접 지표 행(모델이 입력으로부터 인접 행렬을 직접 재구성할 수 있도록 함).
  3. 컨텍스트 노드와 태스크 노드(엔드포인트)를 구분하는 역할 플래그(role flag).
• 태스크 토큰: 두 개의 추가 태스크 토큰이 엔드포인트의 식별자와 인접 정보를 복사하되, 역할 플래그를 반전시켜 이들이 쿼리 대상임을 나타내며 추가된다.

2. 아키텍처 설계

PENCIL은 양방향 셀프 어텐션을 갖춘 표준 트랜스포머 인코더를 사용하며, 특정 잔차 메커니즘(residual mechanism)과 결합된다:

• 입력 투영(Input Projection): 토큰화된 입력과 선택적 노드 특징은 공유된 은닉 차원으로 투영된다. 저자들은 조건이 잘 갖춰진 초기 벡터를 보장하기 위해 투영 행렬에 직교 초기화(orthogonal initialization)를 사용한다.
• 인접성 재구성: 서브그래프 인접 행렬 $\tilde{A}$는 별도의 입력으로 제공되지 않고, 토큰 인코딩에서 인접 지표 행을 슬라이싱하여 재구성된다.
• 곱셈적 잔차(Multiplicative Residual): 복잡한 인코딩 없이 구조적 귀납 편향(inductive bias)을 주입하기 위해, PENCIL은 표준 트랜스포머 출력에 곱셈적 잔차 항을 더한다. 각 레이어 $k$에 대해:
  $$H^{(k)} = Z^{(k)} + P_k(\tilde{A}Z^{(k)})$$
  여기서 $Z^{(k)}$는 셀프 어텐션 블록의 출력이고, $P_k$는 학습 가능한 투영이다. 이를 통해 모델은 밀집 셀프 어텐션과 병행하여 재구성된 인접 행렬에 대한 명시적인 1-홉 메시지 전달을 수행할 수 있다.
• 예측: 두 태스크 토큰의 최종 표현을 결합하고 스칼라 로짓으로 투영하며, 이진 교차 엔트로피 손실로 학습된다.

주요 기여

1. 성능 및 효율성

PENCIL은 다양한 벤치마크(ogbl-ppa, ogbl-citation2와 같은 OGB 데이터셋 포함)에서 최첨단 성능을 달성하는 동시에, 선도적인 ID 기반 방법들보다 훨씬 적은 파라미터(22배에서 146배 적음)를 사용한다. 또한 우수한 학습 효율성을 보여주며, 대규모 데이터셋에서 순수 GNN 아키텍처보다 6.7배에서 40배 적은 에포크(epoch) 내에 수렴한다. 특히, 여러 데이터셋에서 노드 특징 없이도 이러한 결과를 달성하여, 특징이 보편적으로 필수적이라는 가설에 도전한다.

2. 이론적 통합

논문은 PENCIL을 기성 패러다임과 연결하는 공식적인 이론적 분석을 제공한다:

• 휴리스틱 회복: 저자들은 PENCIL이 소스 조건부 MPNN(예: NBFNet)으로 퇴화될 수 있으며, 특정 파라미터 설정 하에서 고전적인 경로 기반 휴리스틱(Katz index, Personalized PageRank, Shortest Path Distance) 및 로컬 휴리스틱(Common Neighbors)을 구현할 수 있음을 증명한다.
• 표현력: $k_\phi$-$k_\rho$-$m$ 프레임워크를 통해, 저자들은 로컬 관계 풀링(local relational pooling)을 갖춘 PENCIL이 동일한 샘플링 제약 조건 하에서 SEAL(서브그래프 기반 GNN)보다 표현력이 낮지 않음을 보여준다.
• 순열 불변성(Permutation Invariance): 인코딩이 무작위 인덱스를 사용함에도 불구하고, 저자들은 결과적인 예측기가 분포에 대해 순열 불변성을 가지며 그래프 작업의 대칭 요구 사항을 충족함을 증명한다.

3. 구조적 충분성

실험 결과는 많은 표준 벤치마크에서 구조적 정보만으로도 높은 성능을 내기에 충분하다는 것을 나타낸다. 본 논문은 노드 특징이 때때로 변동성을 유발하거나, 트랜스포머의 어텐션 메커니즘이 포착하는 풍부한 구조적 신호에 비해 제한적인 한계 이득만을 제공할 수 있음을 입증한다.

실험 결과

• 벤치마크: PENCIL은 Planetoid 데이터셋(Cora, Citeseer, Pubmed)과 OGB 링크 예측 데이터셋(ogbl-citation2, ogbl-ppa, ogbl-ddi, ogbl-collab)에서 평가되었다.
• 비교: P-CEL은 정확도와 안정성(낮은 분산) 측면에서 휴리스틱 기반 GNN(예: NCN, Neo-GNN) 및 ID 기반 방법(예: MPLP+, Refined-GAE)을 능가한다.
• 휴리스틱 추정: 로컬 서브그래프로부터 쌍별 휴리스틱(예: Katz, PageRank)을 추정해야 하는 작업에서, PENCIL은 표준 GNN을 크게 앞지르며, MPNN이 오버스쿼싱(oversquashing)으로 인해 겪는 장거리 구조적 의존성 포착 능력을 입증한다.
• 깊이(Depth): 깊이가 증가함에 따라 성능 저하를 겪는 경우가 많은 MPNN과 달리, PENCIL은 깊어질수록 일관된 개선을 보이며 깊이 증가에 따른 Hits@K 지표의 향상을 보여준다.

의의 및 주장

본 논문은 단순한 설계 선택만으로도 최첨단 링크 예측 능력을 달 achieve 할 수 있음을 주장하며, 복잡한 엔지니어링 기술, 수작업으로 제작된 구조적 인코딩, 그리고 방대한 파라미터 수에 대한 기존의 의존성에 도전한다.

• 패러다임 전환: PENCIL은 단순한 로컬 서브그래프 샘플링 전략과 곱셈적 잔차를 결합한 "일반" 트랜스포머가 표준 GNN보다 더 풍부한 구조적 신호를 추출할 수 있음을 보여준다.
• 배포 가능성: 글로벌 전처리, ID 임베딩 또는 오프라인 구조적 인코딩의 필요성을 제거함으로써, PENCL은 실제 대규모 배포 시나리오에 적합한 확장 가능하고 하드웨어 효율적인 솔루션을 제공한다.
• 이론적 근거: 이 연구는 높은 표현력의 링크 예측과 실질적인 배포 사이의 간극을 메우며, 일반적인 바이디렉셔널 트랜스포머가 전형적인 그래프 트랜스포머에 필요한 강력한 위치 인코딩 없이 어떻게 효과적으로 링크 예측기로 기능할 수 있는지에 대한 첫 번째 체계적인 설명을 제공한다.

저자들은 더 큰 샘플링된 서브그래프를 처리하기 위한 하드웨어 및 시스템의 발전이 이러한 모델의 실질적인 역량을 직접적으로 확장할 수 있다고 결론지으며, 일반 트랜스포머를 그래프 머신러닝의 정당하고 이론적 근거가 있는 패러다임으로 자리매김한다.