LOTFormer 저차원 최적수송 기반 이중확률 선형 어텐션

초록

LOTFormer는 쿼리와 키를 확률 측정으로 해석하고, 작은 크기의 학습 가능한 피벗 측정을 도입해 두 단계의 엔트로피 정규화 최적수송 문제를 풀어 저차원 운송 계획을 만든다. 이 운송 계획을 결합해 이중확률(doubly‑stochastic) 행렬을 얻으며, 행렬의 랭크는 r ≪ n이므로 값 연산을 O(n r) 시간에 수행한다. 실험 결과, 비전과 텍스트 벤치마크에서 기존 선형 및 이중확률 어텐션보다 높은 정확도‑효율성을 보인다.

상세 분석

본 논문은 어텐션을 두 확률 측정(쿼리와 키) 사이의 운송 계획으로 보는 최적수송(OT) 관점을 확장한다. 기존 선형 어텐션은 커널 특성맵을 이용해 행렬을 행별로 정규화하지만, 열 정규화가 없어서 몇몇 토큰에 집중(over‑focusing)되는 문제가 있다. 이를 해결하기 위해 이중확률(doubly‑stochastic) 행렬을 목표로 삼지만, Sinkhorn 같은 전통적 방법은 전체 n × n 행렬에 반복 정규화를 적용해 시간·메모리 비용이 크게 늘어난다. LOTFormer는 “피벗 측정”이라는 작은 지원 집합(크기 r) ⟨ r ≪ n ⟩을 도입한다. 쿼리 측정 µ와 피벗 측정 σ 사이, 그리고 σ와 키 측정 ν 사이에 각각 엔트로피 정규화 OT 문제를 풀어 두 개의 운송 행렬 Γ¹, Γ²를 얻는다. 이때 비용 행렬은 단순 내적 QZᵀ, KZᵀ이며, Sinkhorn 스케일링을 통해 Γ¹ = Diag(u) exp(QZᵀ/ε) Diag(v), Γ² = Diag(ũ) exp(KZᵀ/ε) Diag(ṽ) 형태가 된다. 피벗을 매개로 두 운송 행렬을 결합하면
Γ = (Γ¹)ᵀ Diag(σ)⁻¹ Γ²
가 된다. 수학적으로 Γ는 µ와 ν 사이의 운송 계획이며, 행·열 마진이 각각 µ, ν와 일치하므로 이중확률성을 만족한다. 또한 Γ는 Γ¹와 Γ²의 곱으로 표현되므로 랭크는 최대 r이다. 따라서 값 V에 대한 최종 연산은 V′ = Γ V ≈ (Γ¹)ᵀ Diag(σ)⁻¹ (Γ² V) 로, 각 단계가 O(n r) 시간에 수행된다. 이 구조는 기존 선형 어텐션의 O(n d) 복잡도에 비해 r이 작은 경우 거의 동일한 비용으로 이중확률성을 제공한다. 실험에서는 r = 4~32 범위에서 성능-효율성 트레이드오프를 조사했으며, 특히 r = 16 정도에서 가장 좋은 정확도 향상을 보였다. 또한 피벗 위치와 질량을 학습 가능한 파라미터로 두어 데이터에 맞게 적응하도록 설계했으며, 이는 기존 고정 피벗(예: k‑means) 방식보다 유연성을 높인다. 전체적으로 LOTFormer는 이중확률 어텐션을 선형 시간에 구현함으로써, 과도한 토큰 집중을 방지하면서도 대규모 시퀀스 처리에 적합한 새로운 설계 패러다임을 제시한다.