기술 요약: 동적 그래프 트랜스포머에서의 주의 분산

문제 정의

연속 시간 동적 그래프 (CTDG) 학습은 가변 길이의 역사적 상호작용 시퀀스를 처리하기 위해 트랜스포머 기반 아키텍처를 점점 더 많이 채택하고 있습니다. 이러한 모델들이 표준 벤치마크에서 지배적인 패러다임이 되었음에도 불구하고, 훈련 및 테스트 윈도우 간에 상당한 시간적 분포 변화를 보이는 데이터셋에서는 성능이 정체됩니다. 기존 연구는 주로 시퀀스 구성 (예: 이웃 공발생 인코딩, 시간적 패치링) 과 특징 설계 개선에 초점을 맞추어 왔으나, 이러한 접근법은 US Legis., UN Trade, UN Vote 와 같은 높은 변화 (high-shift) 데이터셋에서의 성능 저하를 해결하지 못합니다.

본 논문은 이러한 한계가 정보 가용성 부족이나 모델 용량 부족 때문이 아님을 지적합니다. 대신, 핵심 실패 모드는 주의 분산입니다: 시간적 분포 변화 하에서 표준 소프트맥스 주의 메커니즘은 중요한 역사적 토큰에 확률 질량을 집중시키지 못하여, 예측 신호와 노이즈를 신뢰성 있게 구별할 수 없는 과도하게 평평한 주의 분포를 초래합니다.

방법론

1. 주의 분산 진단

저자들은 먼저 성능 저하가 최대 평균 불일치 (MMD) 로 측정된 시간적 분포 변화의 크기와 강하게 상관관계가 있음을 입증합니다. 원인을 국소화하기 위해 중요 노드 개념을 도입합니다. 이는 구조적으로 중심적인 (로컬 상호작용 구조에서 여러 노드에 연결된) 또는 시간적으로 안정적인 (반복적이고 안정적인 관계에 관여하는) 역사적 이웃들을 의미합니다.

통제된 제거 실험 (ablation studies) 을 통해 저자들은 다음을 입증합니다:

• 정보 가용성: 중요한 노드들은 비례하여 예측 정보를 많이 담고 있습니다. 이를 마스킹하면 특히 높은 변화 데이터셋에서 성능이 크게 저하됩니다.
• 병목 현상: 모든 중요한 노드가 입력에 존재함에도 불구하고, 기존 트랜스포머 모델은 높은 변화 데이터셋에서 낮은 성능을 보입니다. 이는 실패가 정보 검색이 아닌 주의 할당에 있음을 나타냅니다.
• 메커니즘: 시간적 변화는 쿼리 - 키 점수 ($QK^\top$) 의 상대적 대비를 감소시켜, 이후 소프트맥스가 더 평평하고 분산된 분포를 생성하게 합니다. 모델은 가장 정보량이 많은 토큰에 불충분한 가중치를 부여합니다.

2. 제안된 해결책: 차분 주의 (Differential Attention)

주의 분산을 해결하기 위해 본 논문은 표준 셀프 어텐션을 **차분 주의 (Differential Attention)**로 대체할 것을 제안합니다. 언어 모델링을 위해 처음 제안된 이 메커니즘은 동일한 입력에서 유도된 두 개의 소프트맥스 주의 맵 간의 차이를 계산합니다:
$$\text{DiffAttn}(Z) = (A_1 - \lambda A_2)V$$
여기서 $A_1$과 $A_2$는 두 개의 쿼리 - 키 쌍에서 계산된 주의 맵이며, $\lambda$는 학습 가능한 스칼라입니다.

• 기능: 뺄셈은 두 맵에서 유사하게 나타나는 "공통 모드" 주의 (노이즈나 일반적인 패턴을 나타내는 신호) 를 억제하고 독특한 토큰 수준의 신호를 증폭시킵니다.
• 목표: 이는 역사적 토큰 간의 대비를 복원하여 모델이 가장 예측력이 높고 중요한 노드에 확률 질량을 집중하도록 강제합니다.

3. DiffDyG 구현

저자들은 이 해결책을 참조 CTDG 트랜스포머인 DiffDyG에서 구체화합니다.

• 아키텍처: 표준 입력 구성 요소 (RoPE, 이웃 공발생 인코딩, 공간 거리 인코딩) 는 유지하되, 셀프 어텐션 레이어를 차분 어텐션으로 대체합니다.
• 효율성: 차분 주의의 향상된 시간적 노이즈 제거 능력으로 인해, DiffDyG 는 더 작은 임베딩 차원 (베이스라인의 50 대비 36) 과 더 적은 파라미터로 우수한 성능을 달성하며, FlashAttention 과 Grouped Query Attention 과 결합될 때 더 빠른 훈련 시간을 보입니다.

주요 기여

1. 진단: 본 논문은 시간적 분포 변화 하에서 트랜스포머 기반 CTDG 모델의 공통된 실패 모드로서 주의 분산을 식별합니다. 입력에는 중요한 예측 신호가 존재하지만 평평해진 주의 분포로 인해 무시됨을 입증합니다.
2. 검증: 저자들은 구조적으로 세 가지 다른 베이스라인 (DyGFormer, TIDFormer, TCL) 에서 이 진단을 검증합니다. 주의 메커니즘만 차분 어텐션으로 교체해도 일관된 성능 향상을 얻으며, 특히 높은 변화 데이터셋에서 두드러집니다. 직접 측정을 통해 이 수정이 주의 엔트로피를 감소시키고 중요한 노드에 할당된 주의 질량을 증가시킨 것을 확인했습니다.
3. 최고 성능 (SOTA): DiffDyG는 세 가지 부정적 샘플링 프로토콜 (무작위, 역사적, 귀납적) 하에서 9 개 벤치마크에서 최첨단 결과를 달성합니다. 이전 방법들이 정체되었던 가장 도전적인 변화 데이터셋에서 성능 격차를 크게 좁힙니다.

결과

• 성능 향상: 높은 변화 데이터셋에서 DiffDyG 는 기존 최상위 베이스라인보다 상당한 개선을 보입니다:
  • US Legis.: 71.1% $\to$ 87.5% AP
  • UN Trade: 66.5% $\to$ 99.0% AP
  • UN Vote: 69.0% $\to$ 88.8% AP
• 메커니즘 검증: 주의 수준 측정은 차분 어텐션이 표준 어텐션에 비해 일관되게 엔트로피를 감소시켜 (더 날카로운 집중을 의미) 중요한 노드에 할당된 주의 질량의 비율을 증가시킴을 보여줍니다.
• 효율성: DiffDyG 는 더 적은 파라미터를 사용하면서도 동등한 트랜스포머 기반 모델 (예: DyGFormer, TIDFormer) 보다 1.5 배에서 2.5 배 빠르게 훈련됩니다.

중요성과 주장

본 논문은 어렵고 변화가 큰 데이터셋에서의 CTDG 학습을 위한 향후 방향은 단순히 더 나은 역사적 시퀀스 구성이 아니라, 이러한 시퀀스에 대한 더 나은 주의 할당 메커니즘이 필요하다고 주장합니다. 이 연구는 입력 인코딩의 아키텍처 혁신이 해결하지 못했던 성능 병목 현상을 주의 메커니즘의 표적된 변경으로 해결할 수 있음을 보여줍니다.

저자들은 분석이 주로 경험적임을 명시하며 겸손한 입장을 견지합니다. 소프트맥스 어텐션이 언제 분산되는지 또는 차분 어텐션이 언제 도움이 될 것인지에 대한 공식적인 이론적 특성을 제공하지는 않는다고 인정합니다. 또한, "중요 노드"에 대한 정의는 진단용 탐침으로 기능할 뿐 고유한 최적 정의는 아니며, 평가는 현재 표준 벤치마크의 동적 링크 예측으로 제한됩니다. 향후 연구는 이러한 발견을 다른 동적 그래프 작업과 더 큰 산업용 그래프로 확장할 것을 제안합니다.