Transformers를 이용한 그래프 분류

Hugging Face는 Transformers 라이브러리를 사용하여 그래프 분류를 구현하는 기술 가이드를 제공했습니다. 현재 라이브러리에서 이 작업에 사용할 수 있는 주요 모델은 Microsoft의 Graphormer이며, 이를 통해 사용자는 그래프 구조 데이터를 대상으로 이진 분류, 다중 클래스 분류 또는 회귀와 같은 작업에 트랜스포머 기반 아키텍처를 적용할 수 있습니다.

그래프 데이터 포맷팅 및 로딩

그래프 데이터셋은 주로 jsonl 형식으로 저장되며, 각 그래프는 다음과 같은 필수 및 선택 필드를 포함하는 사전 형태로 표현됩니다:

  • edge_index: 엣지의 노드 인덱스를 나타내는 두 개의 평행한 정수 리스트를 포함하는 리스트입니다 (예: [[1, 1, 3], [2, 3, 1]]는 1–2, 1–3, 3–1 엣지를 의미합니다).
  • num_nodes: 그래프 내 전체 노드 수를 나타내는 정수이며, 고립된 노드도 포함됩니다.
  • y: 예측 대상 값으로, 다중 클래스 분류의 경우 정수, 회귀의 경우 실수, 다중 작업 분류의 경우 이진 레이블 리스트가 될 수 있습니다.
  • node_feat (Optional): 노드 인덱스 순서대로 각 노드의 특성을 담은 정수 리스트들의 리스트입니다.
  • edge_attr (Optional): edge_index 순서를 따르는 각 엣지의 속성을 담은 정수 리스트들의 리스트입니다.

datasets 라이브러리를 사용하여 Hub에서 직접 이러한 데이터셋을 로드할 수 있습니다. 예를 들어 Open Graph Benchmark(OGB)의 ogbg-molhiv 데이터셋이 있습니다.

Graphormer 전처리

그래프 트랜스포머 프레임워크는 학습을 돕는 특성을 생성하기 위해 특정 전처리가 필요합니다. Graphormer의 경우 노드 간 최단 경로 행렬과 입·출 차수 정보를 생성하는 것이 포함됩니다.

전처리는 두 가지 방법으로 처리할 수 있습니다:

  1. Static Preprocessing: dataset.map(preprocess_item, batched=False)를 사용하여 학습 전에 데이터를 전처리합니다.
  2. On-the-fly Processing: GraphormerDataCollator 내부에서 on_the_fly_processing=True로 설정하여 학습 중에 데이터를 처리합니다. 이는 대규모 그래프에 대해 메모리 효율이 높습니다.

모델 구현 및 파인튜닝

그래프 분류는 GraphormerForGraphClassification 클래스를 사용하여 구현됩니다. 워크플로는 사전 학습된 체크포인트를 로드하고 분류 헤드를 특정 다운스트림 작업에 맞게 조정하는 과정을 포함합니다.

로딩 및 설정

from_pretrained를 통해 모델을 로드할 때, 사용자는 특정 작업에 맞게 num_classes를 지정할 수 있습니다(예: 이진 분류의 경우 num_classes=2). ignore_mismatched_sizes=True로 설정하면 라이브러리는 사전 학습된 체크포인트의 원래 디코더 헤드를 교체하고 사용자 정의 분류 헤드를 생성합니다.

학습 워크플로우

학습 과정은 Trainer API를 통해 관리됩니다. 주요 기술적 고려 사항은 다음과 같습니다:

  • Data Collator: 개별 그래프를 배치로 변환하기 위해 Trainer에 GraphormerDataCollator를 전달해야 합니다.
  • Memory Management: 모델 크기 때문에 per_device_train_batch_size와 gradient_accumulation_steps를 조정하여 메모리 부족(OOM) 오류를 방지하는 것이 중요합니다.
  • Hardware: CPU에서도 학습이 가능하지만, 가이드에 따르면 Intel Core i7에서 20 epoch을 수행하는 데 약 하루가 소요되므로 실용적인 사용을 위해 GPU 가속을 강력히 권장합니다.

Sources