가속화된 어텐션 블록: 밀도 다양체 위 관성 동역학

본 논문은 최근 트랜스포머의 어텐션 블록을 확률 밀도 공간 위의 상호작용 입자 시스템으로 해석한 연구들을 확장한다. 기존 연구에서는 토큰 업데이트를 Wasserstein‑2 거리 위의 gradient flow 로 모델링했으며, 이는 1차 미분 방정식 형태의 연속 PDE 로 표현된다. 저자는 여기서 한 단계 나아가, Nesterov 가속법을 밀도 다양체에 적용한 관성(두 번째 차원) 동역학을 도입한다. 1. **이론적 배경** - **Nesterov 가속**: 유클리드 공간에서의 Nesterov 가속은 연속 시간으로는 \(\ddot{x}(t)+3/t\,\dot{x}(t)+\nabla F(x(t))=0\) 형태의 damped Hamiltonian 시스템에 대응한다. 이는 최적화 수렴 속도를 O(1/k²) 로 끌어올린다. - **밀도 다양체**: 확률 밀도 \(\rho\) 를 정의역 \(\Omega\subset\mathbb{R}^d\) 위에 두고, 두 종류의 메트릭(워셔스테인-2와 Stein)을 사용한다. 각각의 메트릭은 이동 연산자 \(G_\rho^{-1}\) 로 정의되며, Wasserstein‑2는 \(-\nabla\cdot(\rho\nabla\Phi)\), Stein은 커널 \(k\) 를 포함한 비선형 형태이다. - **가속화 흐름**: Hamiltonian \(\mathcal{H}(\rho,\Phi)=\frac12\langle\Phi,G_\rho^{-1}\Phi\rangle+F(\rho)\) 로부터, 연속식 (7)–(9) 를 얻는다. 이는 두 개의 연립 PDE 로, 하나는 연속 방정식(질량 보존), 다른 하나는 Hamilton‑Jacobi‑type 방정식에 감쇠 항 \(\alpha(t)\Phi\) 가 추가된 형태이다. 2. **선형 어텐션과 Stein 변분 흐름** - 선형 어텐션에서는 쿼리·키·밸류 행렬을 선형 변환 \(A\) 로 표현하고, 토큰 간 상호작용 에너지를 \(F(\rho)=\frac12\iint \langle x,y\rangle_A \rho(x)\rho(y)dxdy\) 로 정의한다. - 이 경우, 가속화 흐름은 Stein 메트릭을 사용한 변분 흐름과 정확히 일치한다. 커널 \(k(x,y)=\langle x,y\rangle_A\) 로 설정하면, 흐름은 \(\partial_t\rho+\nabla\cdot(\rho\,\mathbb{E}_\rho