DRoPE: Directional Rotary Position Embedding for Efficient Agent Interaction Modeling

본 논문은 자율주행 시스템의 핵심인 궤적 생성(trajectory generation)에서 에이전트 상호작용을 정확하고 효율적으로 모델링하는 문제를 다룬다. 기존의 에이전트 상호작용 모델링 방법론인 scene-centric, agent-centric, query-centric 프레임워크는 각각 정확도, 계산 시간, 메모리 효율성 측면에서 장단점을 가지며, 이들 모두를 동시에 최적화하기 어렵다는 "불가능한 삼각형(impossible triangle)" 문제에 직면해 있다. 특히, query-centric 방법에서 사용되는 Relative Position Embedding (RPE)은 $O(N^2)$의 높은 공간 복잡도(space complexity)를 가지며, 이는 N개의 에이전트가 존재할 때 모든 에이전트 쌍 간의 명시적인 상대 위치를 표현하기 위해 N배의 공간 증가를 야기한다.

이러한 한계를 극복하기 위해, 논문은 자연어 처리(Natural Language Processing, NLP)에서 개발된 Rotary Position Embedding (RoPE)의 새로운 변형인 Directional Rotary Position Embedding (DRoPE)을 제안한다. RoPE는 명시적으로 상대 위치를 표현하지 않고도 효율적으로 상대 위치를 인코딩하여 $O(N)$의 공간 복잡도를 유지하지만, 각도 정보의 주기성(periodicity)을 처리하는 데 본질적인 한계가 있다. 이는 RoPE가 절대 위치(absolute position)를 기반으로 회전 변환을 수행하기 때문에, 상대 각도(relative angle)의 주기적 속성을 직접적으로 반영하기 어렵기 때문이다. 예를 들어, 두 쌍의 에이전트가 동일한 상대 각도 차이를 가지더라도, RoPE는 각도 차이를 나타내는 스칼라 $m\theta_l$이 차원 $l$에 따라 달라지는 $\theta_l = 10000^{-l/d_k}$를 포함하므로, 이 주기성을 올바르게 인코딩하지 못한다.

DRoPE는 RoPE의 2D 회전 변환에 단일(uniform) 항등 스칼라(identity scalar)를 도입함으로써 이 한계를 극복한다. RoPE의 $f_{\to}(X, m) = \text{BlockDiag}(R(m\theta_0), \dots, R(m\theta_{d_k-1}))X$와 달리, DRoPE의 각도 임베딩 함수는 $f_{\angle}(X, \theta) = \text{BlockDiag}(R(\theta), \dots, R(\theta))X$로 정의된다. 여기서 모든 $\theta_l$ 파라미터를 단일 값 $\theta$로 통일하여 사용한다. 이를 통해 $Q_i$와 $K_j$의 QK 벡터에 해당하는 $\bar{Q}_i = f_{\angle}(Q_i, \theta_i)$와 $\bar{K}_j = f_{\angle}(K_j, \theta_j)$의 내적 $\langle \bar{Q}_i, \bar{K}_j \rangle$이 $Q_i^T \text{BlockDiag}(R(\theta_j-\theta_i), \dots, R(\theta_j-\theta_i)) K_j$와 같이 오직 $Q_i$, $K_j$ 및 주기적인 상대 각도 $\theta_j - \theta_i \pmod{2\pi}$에만 의존함을 이론적으로 증명한다(Proposition 4). 이는 실제 에이전트의 헤딩(heading)과 회전 각도를 자연스럽게 정렬시켜 상대 각도 정보를 정확하게 인코딩할 수 있도록 한다.

논문은 DRoPE와 RoPE를 결합하여 에이전트의 상대 위치와 헤딩 방향을 동시에 임베딩하는 두 가지 통합 방식을 제안한다:

Head-by-head 통합(Head-by-head integration): 어텐션 헤드(attention head)를 분리하여 짝수 헤드에는 위치 정보를 위한 RoPE를, 홀수 헤드에는 각도 정보를 위한 DRoPE를 적용한다. 즉, $h \pmod{2} = 0$인 경우 $\hat{Q}_{A,i}^h = f_{\to}(Q_{A,i}^h, \text{pos}_{A,i})$ 및 $\hat{K}_{A,j}^h = f_{\to}(K_{A,j}^h, \text{pos}_{A,j})$를 사용하고, 그렇지 않은 경우 $\bar{Q}_{A,i}^h = f_{\angle}(Q_{A,i}^h, \theta_{A,i})$ 및 $\bar{K}_{A,j}^h = f_{\angle}(K_{A,j}^h, \theta_{A,j})$를 사용하여 어텐션 점수를 계산한다.

Intra-head 통합(Intra-head integration): 단일 어텐션 헤드 내에서 QK 벡터를 위치(position) 서브-벡터와 각도(angle) 서브-벡터로 분해한다. 즉, $Q_{A,i}^h = [Q_{A,i}^{h,\text{pos}}, Q_{A,i}^{h,\text{angle}}]$ 및 $K_{A,j}^h = [K_{A,j}^{h,\text{pos}}, K_{A,j}^{h,\text{angle}}]$로 분할한 후, 각각 RoPE와 DRoPE를 적용하여 두 임베딩의 내적을 합산하여 어텐션 점수를 계산한다: $\frac{\langle \hat{Q}_{A,i}^{h,\text{pos}}, \hat{K}_{A,j}^{h,\text{pos}} \rangle + \langle \bar{Q}_{A,i}^{h,\text{angle}}, \bar{K}_{A,j}^{h,\text{angle}} \rangle}{\sqrt{d_k}}$.

DRoPE는 RPE와 달리 $O(N^2)$가 아닌 $O(N)$의 공간 복잡도를 가지며, 이는 병렬 계산 설정에서 QKV 벡터를 위한 RoPE 입력의 공간 복잡도가 $O(NH(2d_k + d_v))$임을 Proposition 3으로 증명한다. 이로써 DRoPE는 궤적 생성 정확도, 시간 복잡도, 공간 복잡도를 동시에 최적화한다. 다양한 최신 궤적 생성 모델과의 경험적 비교를 통해 DRoPE가 우수한 성능을 유지하면서 공간 복잡도를 크게 줄임을 확인하였다.