기술 요약: 열대 축 어텐션을 활용한 계통수 추론

문제 정의

계통수 재구성은 다중 서열 정렬 (MSA) 로부터 진화적 관계를 추론하는 것을 목표로 하는 조합 최적화 문제이다. 표준 통계적 접근법 (예: 최대 우도법, 베이지안 추론) 과 거리 기반 방법 (예: 이웃 연결법, 균형 최소 진화법) 이 존재하지만, 이들은 상당한 도전에 직면해 있다. 거리 기반 방법은 쌍별 거리 행렬을 추정하고 이를 유효한 계통수의 공간으로 투영해야 한다. 이 공간은 매끄러운 유클리드 다양체가 아니라, 다면체 제약 (특히 4 점 조건과 초거리 부등식) 으로 정의된 구조화된 비유클리드 부분공간이다.

최근 Phyloformer 및 Phyloformer 2 와 같은 딥러닝 접근법은 표준 소프트맥스 점곱 어텐션을 사용하는 트랜스포머 아키텍처를 활용하여 MSA 로부터 거리 행렬을 직접 예측한다. 그러나 이러한 모델들은 유클리드 공간에서 작동하며 매끄럽고 밀집된 확률 분포에 의존한다. 저자들은 이러한 기하학적 불일치가 거리 측정의 고유한 조합적이고 조각별 선형적인 구조를 포착하지 못해, 결과물이 트리 제약을 위반하거나 분포 변화 (예: 서열 길이 또는 진화 모델의 변화) 하에서 일반화에 어려움을 겪을 수 있다고 주장한다.

방법론

본 논문은 모델의 귀납적 편향을 계통수 공간의 기하학과 정렬되도록 설계된 열대 축 어텐션 (Tropical Axial Attention) 신경 아키텍처를 제안한다.

1. 이론적 기반: 열대 기하학

핵심 이론적 통찰은 $n$개의 종을 가진 계통수의 공간과 열대 그라스마니안 (tropical Grassmannian) 사이의 동형성이다. 열대 대수는 표준 덧셈과 곱셈을 각각 $\max$와 덧셈으로 대체한다 (max-plus 반환).

• 트리 거리: 유효한 트리 거리는 선형 결합이 아닌 최대 연산을 포함하는 4 점 조건으로 특징지어진다.
• 초거리 (Ultrametrics): 뿌리가 있는 등거리 트리는 또한 최대 연산에 의해 지배되는 초거리 조건을 만족한다.
• 거리 공간: 저자들은 표준 유클리드 거리와 달리 덧셈적 이동에 불변하고 상대적 구조에 민감한 몫 공간 $\mathbb{R}^D / \mathbb{R}\mathbf{1}$ 위의 **힐베르트 사영 거리 (Hilbert projective metric, 열대 거리)**를 활용한다.

2. 모델 아키텍처

모델은 바닐라 소프트맥스 어텐션을 열대 축 어텐션으로 대체하여 Phyloformer 아키텍처를 확장한다:

• 입력 처리: MSA 는 원-핫 인코딩되어 임베딩된다. 순서가 없는 잎 (leaf) 쌍들을 연결하여 시퀀스 위치 전반에 걸친 쌍별 관계를 나타내는 텐서를 형성한다.
• 열대 어텐션 메커니즘:
  • 쿼리 ($Q$), 키 ($K$), 값 ($V$) 은 max-plus 반환에서 열대 선형 투영 (최대 및 덧셈 사용) 을 통해 계산된다.
  • 어텐션 점수는 열대 거리를 사용하여 계산된다: $S_{i,j} = -(\max(k_j - q_i) - \min(k_j - q_i))$.
  • 집계는 열대 행렬 곱셈을 통해 수행된다: $\text{Attn}(Q, K, V)_i = \max_j (S_{i,j} + v_j)$.
  • 이는 소프트맥스의 "부드러운" 평균화를 조각별 선형 결정 경계를 보존하는 "날카로운" 선택 메커니즘으로 대체한다.
• 축 분해: 모델은 두 축을 따라 어텐션을 적용한다: 시퀀스 축 (정렬 사이트 전반의 돌연변이 증거 학습) 과 쌍 축 (잎 쌍 간의 관계 학습 및 기하학적 트리 제약 강제).
• 출력: 포인트별 순방향 투영이 채널 차원을 축소하여 쌍별 거리를 예측한 후, 음수가 아니도록 보장하기 위해 Softplus 활성화 함수가 적용된다.

3. 손실 함수

학습 목적 함수는 정확성과 기하학적 일관성 모두를 강제하기 위해 세 가지 구성 요소를 결합한다:

1. $\ell_1$ 손실: 환경 유클리드 공간에서의 절대 좌표별 오차를 측정하여 실제 값과의 근접성을 보장한다.
2. 열대 대칭 손실: 힐베르트 사영 거리 ($d_{tr}$) 를 사용하여 열대 기하학의 왜곡을 패널티로 부과하며, 모델이 절대값뿐만 아니라 거리의 상대적 구조를 보존하도록 장려한다.
3. 초거리 패널티: 초거리 조건 (임의의 삼중체에서 두 개의 가장 큰 거리가 같아야 함) 위반에 대해 패널티를 부과한다.

• 학습 전략: 초거리 패널티는 모델이 엄격한 트리 제약을 부과하기 전에 먼저 거친 거리 기하학을 학습할 수 있도록 첫 20 에포크 이후에 활성화된다.

주요 결과

모델은 세 가지 유형의 데이터셋에 대해 Phyloformer 및 Phyloformer 2 와 비교 평가되었다:

1. 분포 내 (In-Distribution) 시뮬레이션: 훈련 세트와 동일한 출생 - 사망 과정으로 생성된 데이터 (잎 수 $n=10$에서 $n=100$까지 변화) 에 대해, 열대 모델은 경쟁력 있는 성능을 보였으며, 일반적으로 위상학적 (Robinson-Foulds) 및 가지 길이 메트릭에서 Phyloformer 와 Phyloformer 2 사이의 결과를 달성했다.
2. 분포 외 (Out-of-Distribution, OOD) 시뮬레이션: 다른 시뮬레이션 과정 (GTR 모델을 사용한 PhyML 벤치마크) 으로 생성된 데이터에서, 열대 모델은 우수한 견고성을 보여주었다. 위상학적 정확도는 Phyloformer 2 와 유사하면서도 가지 길이에 민감한 메트릭 (가중 RF 및 KF 거리) 에서 두 기준 모델보다 현저히 우수한 성과를 거두었다.
3. 실증 데이터셋 (DS1–DS11): 참나무를 알 수 없는 실제 생물학적 데이터에서 평가는 행렬 - 트리 일관성에 초점을 맞추었다. 열대 모델은 기준 모델들보다 유도된 BME (균형 최소 진화) 트리 거리와 MAE 및 열대 거리 모두에서 훨씬 더 가까운 거리 행렬을 생성했다. 이는 열대 모델이 트리 공간의 근본적인 다면체 구조와 더 일관된 표현을 학습함을 시사한다.

중요성과 주장

본 논문은 계통 추론의 어려움이 단순히 계산적인 것이 아니라, 트리의 비유클리드적이고 조합적인 성질에서 비롯된 구조적 문제라고 주장한다.

• 기하학적 귀납적 편향: 주요 기여는 유클리드 어텐션을 열대 어텐션으로 대체하는 것이 계통 추론을 위한 자연스러운 기하학적 프레임워크를 제공한다는 것을 입증하는 것이다. 네트워크의 연산 (max-plus) 을 트리 거리의 제약 (4 점 조건) 과 정렬함으로써, 모델은 트리 위상과 본질적으로 더 일관된 표현을 학습한다.
• 분포 변화에 대한 견고성: 결과는 열대 어텐션이 특히 훈련 세트와 다른 데이터 분포에서 테스트될 때 강력한 귀납적 편향으로 작용함을 시사한다.
• 순수 정확도보다 일관성: 모델이 분포 내 설정에서 모든 메트릭에서 기준 모델을 보편적으로 능가하지는 않지만, 일관되게 거리 행렬의 수학적 정의와 더 일관된 출력을 생성한다. 이는 거리 행렬의 유효성이 중요한 하류 작업에서 특히 중요하다.

저자들은 열대 축 어텐션이 딥러닝 아키텍처와 진화 트리의 대수 기하학 사이의 간극을 메우는 원칙적인 신경 계통 추론 접근법을 제공한다고 결론지었다.