[논문 리뷰] Does Invariant Graph Learning via Environment Augmentation Learn Invariance?

Yongqiang Chen, Yatao Bian외 4명|arXiv (Cornell University)|2023. 10. 29.

Advanced Graph Neural Networks인용 수 5

한 줄 요약

이 논문은 지식 기반 환경 레이블이 없이도 불변 그래프 표현 학습을 가능하게 하는 새로운 프레임워크인 GALA를 제안한다. GALA는 대체 환경 보조 모델을 활용해 올바르게 예측된 그래프와 잘못 예측된 그래프 간의 대비 학습을 통해 불변 부분그래프를 식별한다. 변동성의 충분성과 일관성이라는 최소한의 가정 하에 GALA는 불변 부분그래프를 증명 가능하게 고립시키며, DrugOOD 포함 12개 데이터셋에서 최신 기준 성능을 달성하며, 경쟁적인 학습 효율성도 확보한다.

ABSTRACT

Invariant graph representation learning aims to learn the invariance among data from different environments for out-of-distribution generalization on graphs. As the graph environment partitions are usually expensive to obtain, augmenting the environment information has become the de facto approach. However, the usefulness of the augmented environment information has never been verified. In this work, we find that it is fundamentally impossible to learn invariant graph representations via environment augmentation without additional assumptions. Therefore, we develop a set of minimal assumptions, including variation sufficiency and variation consistency, for feasible invariant graph learning. We then propose a new framework Graph invAriant Learning Assistant (GALA). GALA incorporates an assistant model that needs to be sensitive to graph environment changes or distribution shifts. The correctness of the proxy predictions by the assistant model hence can differentiate the variations in spurious subgraphs. We show that extracting the maximally invariant subgraph to the proxy predictions provably identifies the underlying invariant subgraph for successful OOD generalization under the established minimal assumptions. Extensive experiments on datasets including DrugOOD with various graph distribution shifts confirm the effectiveness of GALA.

연구 동기 및 목표

• 지식 기반 환경 레이블이 없이도 환경 증강 기반 그래프 표현 학습이 진정으로 불변성을 학습할 수 있는지 조사하기 위해
• 기존 환경 증강 방법이 허위 부분그래프 변동성을 포착하는 데 있어 근본적인 한계를 규명하기 위해
• 지식 기반 환경 레이블이 없을 때 불변 그래프 학습이 가능해지는 데 필요한 최소한의 가정—변동성의 충분성과 일관성—을 설정하기 위해
• 예측 정확도 기반 대비 학습을 통해 최대한의 불변 부분그래프를 추출하는 데 목적이 있는 GALA 프레임워크를 제안하기 위해
• 다양한 분포 이탈이 존재하는 12개의 다양한 그래프 데이터셋에서 GALA의 효과성과 효율성을 검증하기 위해

제안 방법

• GALA는 그래프 레이블을 예측하기 위해 경험적 리스크 최소화(Empirical Risk Minimization, ERM) 방식으로 훈련된 환경 보조 모델을 도입하며, 이 모델의 예측 결과를 환경 변화의 대체 지표로 사용한다.
• 프레임워크는 보조 모델의 출력 결과에 기반해 훈련 그래프를 올바르게 예측한 그룹(긍정군)과 잘못 예측한 그룹(부정군)으로 분리한다.
• 같은 레이블을 가진 그래프이지만 예측 정확도가 다른 경우에 대해 일방적인 대비 손실을 적용함으로써 최대한의 불변 부분그래프를 학습한다.
• 모델이 올바른 예측 그룹과 잘못된 예측 그룹 모두에서 일관되게 유지되는 부분그래프에 집중하도록 유도함으로써 불변성을 보장한다.
• 데이터 특성에 따라 환경 보조 모델은 ERM 모델 또는 클러스터링 기반 예측을 활용하는 해석 가능한 GNN으로 구현될 수 있다.
• GALA는 대비 정규화 적용 이전에 주 GNN을 10 에포크 동안 ERM로 사전 훈련하며, 필요에 따라 클래스 불균형을 보완하기 위해 업샘플링을 활용한다.

Figure 1: An illustration of GALA with the task of classifying graphs according to whether there exists a “House” or “Cycle” motif. Given the training data where the “House” subgraph often co-occurs with a “Grid” and the “Cycle” subgraph often co-occurs with a “Hexagon”. An ERM trained environment a

실험 결과

연구 질문

• RQ1 지식 기반 환경 레이블이 없이도 환경 증강을 통해 불변 그래프 표현 학습을 달성할 수 있는가?
• RQ2 기존의 환경 증강 접근법이 그래프 OOD 일반화에서 허위 부분그래프 변동성을 포착하는 데 있어 근본적인 한계는 무엇인가?
• RQ3 환경 레이블이 가용하지 않을 때 불변 그래프 학습이 가능해지는 최소한의 가정은 무엇인가?
• RQ4 지식 기반 환경 레이블이 없을 경우 대체 모델의 예측 결과를 어떻게 활용해 진정한 불변 부분그래프를 식별할 수 있는가?
• RQ5 GALA는 기존 방법보다 OOD 일반화 성능을 뛰어나게 하면서도 경쟁적인 학습 효율성을 유지하는가?

주요 결과

• GALA는 DrugOOD 포함 12개의 다양한 그래프 데이터셋에서 최신 기준 성능을 달성하며, CIGA, GREA, DisC와 같은 기준 모델들보다 일관된 개선을 보였다.
• GALA는 경쟁적인 학습 시간을 유지하며, CIGA에 비해 약간의 증가 외에는 DisC 및 GIL과 비교해 훨씬 낮은 오버헤드를 보였다.
• 평균적으로 GALA는 GIL 및 DisC보다 학습 시간을 15~25% 감소시켰으며, 이는 조기 정지 조건 하에서도 느린 수렴을 보이는 두 모델 대비 유의미한 개선이었다.
• DrugOOD 데이터셋에서 클러스터링 기반 대체 예측을 사용할 경우 ERM 기반 예측보다 경험적으로 더 높은 성능을 기록했다.
• 이론적 분석을 통해 변동성의 충분성과 일관성 조건 하에 GALA가 기저의 불변 부분그래프를 증명 가능하게 식별할 수 있음을 확인했다.
• 광범위한 아블레이션 실험을 통해 환경 보조 모델의 예측 정확도가 불변 구조를 고립시키는 데 핵심적임을 입증하며, 핵심 설계 원리의 타당성을 검증했다.

(a) Correlation strengths