고차원 모듈러 어텐션: 단백질 서열을 위한 쌍방 및 삼중 상호작용 융합

초록

HOMA(Higher‑Order Modular Attention)는 기존의 쌍방(2‑D) 어텐션에 삼중(3‑D) 상호작용 경로를 추가한 새로운 어텐션 연산자이다. 겹치는 블록과 로컬 윈도우를 이용해 삼중 어텐션의 계산량을 O(L·w²)로 낮추면서도, 단백질 서열의 고차원 상호작용을 명시적으로 모델링한다. TAPE 벤치마크(Secondary Structure, Fluorescence, Stability)에서 Blockwise‑2D 기반 모델보다 일관된 성능 향상을 보였으며, 메모리와 속도 측면에서도 실용적인 비용을 유지한다.

상세 분석

본 논문은 단백질 서열‑표현 학습에서 “에피스타시스”라 불리는 다중 잔기 간의 협동 효과를 직접 모델링하고자 한다. 기존 트랜스포머의 셀프‑어텐션은 쿼리와 키 사이의 내적을 이용해 두 토큰 간의 관계만을 학습하므로, 삼중 이상의 고차원 상호작용을 암묵적으로 추정해야 한다. 이를 보완하기 위해 저자들은 HOMA라는 연산자를 설계했으며, 핵심 아이디어는 두 개의 병렬 경로를 두는 것이다. 첫 번째는 기존의 2‑D 스케일드 닷‑프로덕트 어텐션이며, 두 번째는 3‑D 삼중 어텐션이다. 삼중 어텐션에서는 쿼리 Q_i, 키 K_j, 추가 투영 U_k를 이용해 S³_{ijk}= (1/√d) Σ_c Q_{ic} K_{jc} U_{kc} 라는 스칼라 점수를 계산하고, 이를 (j,k) 쌍에 대해 소프트맥스 정규화한다. 정규화된 가중치는 V_j와 V_k의 원소별 곱 ˜V_{jk}=V_j⊙V_k와 결합돼 최종 삼중 출력 O³_i= Σ_j Σ_k A³_{ijk} ˜V_{jk} 로 합산된다.

삼중 연산은 O(L³) 복잡도를 갖기에 실용적이지 않다. 저자들은 세 가지 최적화를 적용한다. 첫째, 겹치는 블록(길이 ℓ, 스트라이드 s)으로 시퀀스를 분할해 블록 내부에서만 어텐션을 수행한다. 둘째, 각 블록 내에서 윈도우 크기 w(ℓ≫w)를 지정해 (j,k) 쌍을 쿼리 i 주변 w 범위로 제한함으로써 O(ℓ·w²)로 복잡도를 낮춘다. 셋째, 삼중 경로의 투영 행렬 W_U를 저차원(r) 팩터화해 파라미터 수를 크게 억제한다. 이러한 설계는 메모리 사용량을 O(T·ℓ·w²) 수준으로 유지하면서도, 로컬 삼중 상호작용을 충분히 포착한다.

실험에서는 TAPE 벤치마크의 세 가지 태스크에 대해 동일한 트랜스포머 백본(12층, d_model=512, 8 헤드)을 사용하고, 어텐션 모듈만을 교체해 비교했다. Baseline으로는 전역 MHSA(Pairwise‑2D), 겹치는 블록 어텐션(Blockwise‑2D), Linformer‑style 저랭크 어텐션(Linear‑2D)를 포함한다. HOMA는 Blockwise‑2D에 삼중 경로를 추가한 형태이며, 윈도우 크기 w를 3,5,7로 변형해 최적값을 탐색했다. 결과는 모든 태스크에서 일관된 개선을 보였다. 예를 들어 Secondary Structure의 CASP12 테스트에서 Blockwise‑2D가 0.6368의 정확도를 기록한 반면, HOMA(w=5)는 0.6588(≈3.5% 상대 향상)을 달성했다. Stability 회귀에서는 ρ=0.7152로, Blockwise‑2D(ρ=0.6509) 대비 약 9.9% 상승했다. 또한, 토큰당 처리 속도와 피크 메모리 사용량도 합리적인 수준을 유지했으며, 특히 w가 작을수록 효율성이 크게 향상된다.

한계점으로는 삼중 어텐션이 여전히 로컬 윈도우에 국한돼 전역 고차원 상호작용을 포착하지 못한다는 점이다. 또한, 윈도우와 블록 하이퍼파라미터 선택이 태스크마다 민감하게 작용해 자동 튜닝이 필요하다. 향후 연구에서는 동적 윈도우, 하이퍼그래프 기반 전역 삼중 연결, 혹은 멀티스케일 블록 구조를 도입해 전역 고차원 의존성을 더 효과적으로 모델링할 여지가 있다.