PyKEEN 실습 정리: 지식 그래프 임베딩 TransE·DistMult·RotatE 링크 예측 비교

핵심 메시지

• PyKEEN은 지식 그래프 임베딩 모델의 학습과 평가를 하나의 파이프라인으로 묶은 파이썬 라이브러리로, 모델·데이터셋·손실 함수·네거티브 샘플러를 독립 인자로 분리해 변인을 통제한 비교를 가능하게 한다.
• 지식 그래프에는 참인 사실만 들어 있어, 헤드나 테일을 무작위로 바꿔치기해 거짓 샘플을 만드는 네거티브 샘플링이 학습의 전제가 된다.
• TransE는 관계를 평행 이동으로, DistMult는 원소별 곱의 가중치로, RotatE는 복소수 공간의 회전으로 본다. 이론적 표현력은 RotatE, DistMult, TransE 순이다.
• 실습에 쓴 Nations 데이터셋에서는 표현력이 중간인 DistMult가 가장 좋은 성능을 냈다. 국가 간 외교 관계라 대칭 관계가 많아 모델의 가정과 데이터 특성이 맞아떨어진 결과다.
• 링크 예측이 내놓는 상위 후보는 사실이 아니라 '그럴듯한 후보'다. 실제 활용에서는 전문가 검토나 실험 검증 단계로 넘겨야 한다.

쉽게 이해하기

서울대학교 산업공학과 DSBA 연구실의 온톨로지·지식 그래프 스터디 18회차는 PyKEEN 실습이었다. 발표자는 먼저 이 라이브러리가 왜 필요했는지부터 짚었다. 지식 그래프 임베딩 모델은 논문마다 구현 방식도, 하이퍼파라미터 설정도, 평가 지표도 제각각이어서 공정한 비교가 사실상 불가능했다는 것이다. 어떤 논문이 자기에게 유리한 조건으로 실험했더라도 이를 확인할 방법이 없었다. PyKEEN은 데이터셋을 나누고 학습해 평가하는 과정을 하나의 파이프라인으로 통일하고, 구성 요소를 전부 독립적인 인자로 분리해 하나씩 바꿔 가며 실험할 수 있게 만들었다. 발표에 따르면 현재 모델 40개, 데이터셋 37개, 순위 기반 지표와 분류 기반 지표 각 22개, 손실 함수 15개가 들어 있다.

학습의 뼈대는 세 모델이 공유한다. 모든 개체와 관계를 벡터로 표현하고, 그래프에 있는 트리플을 참 샘플로, 없는 트리플을 거짓 샘플로 삼아 둘 사이에 미리 정한 마진 이상의 점수 차이가 나도록 임베딩을 조정한다. 문제는 지식 그래프에 참인 사실만 들어 있다는 점이다. 그래서 헤드나 테일의 개체를 무작위로 갈아 끼워 거짓 샘플을 인위적으로 만드는 네거티브 샘플링이 필요하다. 참과 거짓의 점수 차이가 이미 마진보다 크면 손실이 0이 되어 업데이트가 일어나지 않고, 차이가 부족하면 그만큼 임베딩이 조정된다.

세 모델의 차이는 점수 함수를 어떻게 정의하느냐에 있다. 2013년 발표된 TransE는 헤드에 관계를 더하면 테일이 되어야 한다고 보는데, 양방향 관계를 표현하려면 관계 벡터가 0이 되어야 해서 사실상 관계가 사라지는 한계가 있다. 2015년 DistMult는 관계를 헤드와 테일 사이 연관성의 가중치로 보고 원소별 곱의 합을 점수로 쓰는데, 곱셈의 교환법칙 때문에 헤드와 테일을 구분하지 못하고 모든 관계를 양방향으로 취급한다. 2019년 RotatE는 더하기를 회전으로 바꿔, 양방향 관계는 180도 회전, 방향이 있는 관계는 임의의 각도로 표현한다. 앞선 두 모델의 제약이 모두 없어 이론적 표현력이 가장 크다.

실습 데이터셋은 PyKEEN에 내장된 Nations로, 1960년대 후반 국가 간 외교·경제·군사 관계를 담고 있다. 개체는 14개, 관계는 55개, 학습 트리플 1592개에 검증 199개, 테스트 201개 규모다. 세 모델 모두 100에폭, 시드 42, CPU라는 같은 조건으로 돌렸다. 학습된 임베딩은 TransE와 DistMult가 14×50, RotatE가 14×200으로 나오는데, 이는 RotatE가 더 큰 모델이어서가 아니라 복소수 임베딩의 실수부와 허수부를 펼쳐 보여주기 때문이다.

평가는 링크 예측으로 했다. 테스트 트리플의 한쪽 개체를 가리고 모델이 정답을 몇 번째 순위에 올리는지 보는 과제로, 헤드 방향과 테일 방향을 모두 평가해 402개의 순위를 얻는다. 이미 참으로 알려진 다른 트리플은 후보에서 빼는 필터링이 공정성의 조건이다. 지표는 정답 순위의 평균인 MR, 역수 평균인 MRR, 상위 K등 안에 든 비율인 Hits@K를 함께 봤다. 결과는 DistMult가 대부분의 지표에서 가장 좋았고, 표현력이 가장 크다던 RotatE가 그다음이었다. 발표자는 개체가 14개뿐이라 Hits@10 같은 지표는 아무렇게나 찍어도 점수가 잘 나오니 하나만 보지 말라고 당부했다.

주요 인사이트

• 표현력 순위와 실제 성능 순위는 일치하지 않는다. Nations처럼 대칭 관계가 많은 데이터에서는 대칭을 강제하는 DistMult의 약점이 오히려 강점이 됐고, 양방향 관계를 못 다루는 TransE가 가장 뒤처졌다.
• 정성 분석에서 TransE는 '브라질이 브라질과 국제기구에 함께 가입했다'는 식으로 자기 자신을 1순위로 올렸다. 정보량이 없는 답으로, 이 데이터셋에서 제대로 학습되지 않았다는 신호로 읽힌다.
• 모델마다 점수의 척도가 달라 서로 다른 모델의 점수를 확신도처럼 비교하면 안 된다. 점수는 같은 모델 안에서 상대적 순위를 정하는 데만 써야 한다.
• PCA 산점도에서 지리적·정치적으로 가까운 국가들이 뭉쳐 나타났지만 배치 모양은 모델마다 달랐다. 점수 함수가 다르면 만들어지는 의미 공간 자체가 달라진다는 뜻이다.
• 이 기법이 쓰이는 분야는 신약 후보 탐색, 약물 재창출, 부작용 예측, 추천 시스템, 이상거래 탐지처럼 겉보기엔 제각각이지만, 모두 관측되지 않은 링크를 순위로 제시하는 같은 문제다.

자주 묻는 질문

관련 AI 소식