# 3D-RPE: Enhancing Long-Context Modeling Through 3D Rotary Position Encoding

Xindian Ma1, Wenyuan Liu1, Peng Zhang1 , Nan Xu2
1 College of Intelligence and Computing, Tianjin University, Tianjin, China
2 Beijing Wenge Technology Co. Ltd.
{xindianma, lwy2020<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
Corresponding Author: Peng Zhang

###### Abstract

Inspired by the Bloch Sphere representation, we propose a novel rotary
position encoding on a three-dimensional sphere, named 3D Rotary Position
Encoding (3D-RPE). 3D-RPE is an advanced version of the widely used 2D Rotary
Position Encoding (RoPE), with two major advantages for modeling long
contexts: controllable long-term decay and improved position resolution. For
controllable long-term decay, 3D-RPE allows for the regulation of long-term
decay within the chunk size, ensuring the modeling of relative positional
information between tokens at a distant relative position. For enhanced
position resolution, 3D-RPE can mitigate the degradation of position
resolution caused by position interpolation on RoPE. We have conducted
experiments on long-context Natural Language Understanding (NLU) and long-
sequence Language Modeling (LM) tasks. From the experimental results, 3D-RPE
achieved performance improvements over RoPE, especially in long-context NLU
tasks.

Figure 1: 2D Rotary Position Encoding (RoPE) vs. 3D Rotary Position Encoding
(3D-RPE).

##  1 Introduction

Rotary Position Encoding (RoPE) [23] is essential in Transformer-based Large
Language Models (LLMs), such as the LLaMA models [24]. RoPE merges the
advantages of absolute and relative positional encoding by using a rotation
mechanism to represent each position. Despite its widespread use in LLMs [24,
27, 7], RoPE has notable limitations when extending LLMs with a predefined
context window. The long-term decay problem of RoPE limits the model’s ability
to extend positions outward in long-context tasks. Although the long-context
modeling capability of LLMs can be extended through position interpolation, as
more positions are inserted, RoPE encounters the challenge of decreased
position resolution.

We propose a novel position encoding mechanism for transformer architecture,
called 3D Rotary Position Encoding (3D-RPE), to address challenges in long-
context modeling faced by LLMs using RoPE. Inspired by the Bloch Sphere
representation, 3D-RPE applies rotary position encoding on a three-dimensional
spherical surface, as illustrated in Figure 1(b). In contrast, RoPE employs
rotation on a two-dimensional circular path, as depicted in Figure 1(a). RoPE
suffers from long-term decay, as shown in Figure 1(c), implying that as the
relative distance increases, the relative upper bound on token correlations at
modeled relative positions will continuously decrease. 3D-RPE addresses this
issue by segmenting a long sequence into chunks and setting rotation angles
within and between the chunks to construct position encoding. As shown in
Figure 1(d), 3D-RPE is able to control this relative upper bound through two
relative positional dimensions, namely within and between chunks. Compared to
Figure 1(c), this method improves the upper bound on correlations between long
relative distances and alleviates the issue of long-term decay.

Position Interpolation (PI) methods [4, 18] based on RoPE are often employed
to extend LLMs for modeling contexts that exceed the pre-training length.
These techniques scale the position encoding during inference, allowing the
originally out-of-range position encoding to fall within the trained position
interval after interpolation. However, as the interpolation factor increases,
PI experiences a substantial decline in positional resolution among tokens,
detrimentally affecting long-context modeling performance. As illustrated in
Figure 1(e), extending the pre-training length $L_{p}$ to $L$ using linear PI
[4] leads to reduced positional resolution with increasing $L$. 3D-RPE employs
a 3D rotating sphere for position encoding, which supports higher positional
resolution compared to the 2D circular rotation. Similarly, through linear PI
extension, 3D-RPE achieves a positional resolution superior to
$\frac{L_{p}}{L}$ (See Figure 1(f)). This benefit has been theoretically
proven (Refer to Theorem 1 in Section 3.2.2) and corroborated by experimental
results (Refer to Table 4 in Section 5.2).

We conducted experiments on long-sequence Language Modeling (LM) and long-
context Natural Language Understanding (NLU) tasks. Our experimental results
highlight the promising performance of the 3D-RPE method, especially in tasks
requiring long-context language understanding.

Our major contributions of this paper are as follows:

  * •

A position encoding method on a 3D sphere, 3D-RPE, is provided, which can
enhance the long-context modeling capability of LLMs by replacing RoPE.

  * •

It is proved that 3D-RPE has two benefits, controllable long-term decay and
mitigating the reduction in positional resolution caused by position
interpolation.

  * •

LLMs combine with 3D-RPE have achieved significant performance improvements in
long-context NLU tasks.

The structure of this paper is as follows. Section 2 covers the preliminaries
of 3D-RPE, Bloch Sphere, and RoPE. Section 3 explains the construction of
3D-RPE on a 3D rotating sphere and highlights its benefits over RoPE. Section
4 reviews related work. In Section 5, we validated the advantages of our
method through experiments. Section 6 concludes with a discussion on 3D-RPE’s
impact.

##  2 Preliminaries

The analysis of 3D-RPE relies on these concepts and results from the filed of
Bloch Sphere and RoPE. We offer an introduction to Bloch Sphere in Section 2.1
and RoPE [23] in Section 2.2.

###  2.1 Bloch Sphere

Bloch Sphere (BS) offers a geometric depiction of a quantum mechanical
system’s pure state, limited to two levels. The state vector $|\phi\rangle$ is
mathematically expressed as

$\displaystyle\ket{\phi}=e^{\mathrm{i}\theta}(\cos{\frac{\varphi}{2}}\ket{0}+\sin{\frac{\varphi}{2}}e^{i\theta_{1}}\ket{1})$
(1)

where $|0\rangle$ and $|1\rangle$ are Dirac’s notations. ${\theta}$,
${\theta}_{1}$ and $\varphi$ are rotation angles. In our work, $\theta$
encodes the relative positions of tokens within chunks, $\varphi$ encodes the
relative positions of tokens across chunks, and $\theta_{1}$ is equal to $0$.
Some other concepts about BS are showed in Supplementary Materials A.

###  2.2 Rotary Position Embedding

Rotary Position Embedding (RoPE) is a commonly used relative position encoding
technique in LLMs, such as LLaMA [24], GPT-J [27], Vicuna [7] and etc. RoPE is
a 2-dimensional space rotary encoding, which is denoted as follows:

$\displaystyle
RoPE(\bm{h}_{m},m)=e^{\mathrm{i}m\theta}\bm{h}_{m}\leavevmode\nobreak\
,\leavevmode\nobreak\ RoPE(\bm{h}_{n},n)=e^{\mathrm{i}n\theta}\bm{h}_{n}$ (2)

$\bm{h}_{m}$ and $\bm{h}_{n}$ are hidden vectors from the query and key for a
specific attention head in transformer. For ease of differentiation,
$\bm{h}_{m}$ and $\bm{h}_{n}$ can be refined later as $\bm{q}_{m}$ and
$\bm{k}_{n}$, $\mathrm{i}$ is the imaginary unit, $\theta$ is the rotary angle
in RoPE. $m$ and $n$ are indexes about positions. Then, the inner product is
employed to define the self-attention score before softmax computing:

$\displaystyle s(m-n,\bm{q}_{m},\bm{k}_{n})$ $\displaystyle=\langle
RoPE(\bm{q}_{m},m),RoPE(\bm{k}_{n},n)\rangle$ (3)
$\displaystyle=Re[\sum_{l=0}^{{d/2}-1}\bm{q}_{[2l:{2l+1}]}\bm{k}_{[2l:{2l+1}]}e^{\mathrm{i}(m-n)\theta_{l}}]$

Eq (3) is unary function respect to the relative position $(m-n)$,
representing the relative position between tokens and modeling the relative
positional information. Here, $Re[\cdot]$ denotes the calculation of the real
part of a complex number. In our study, the 3D-RPE self-attention score is a
binary function containing the relative position $(m-n)$.

##  3 Method

Section 3.1 introduces the new position encoding on a 3D sphere, 3D-RPE.
Section 3.2 focuses on analyzing two benefits of 3D-RPE, namely controllable
long-term decay and enhanced position resolution.

Figure 2:  Visualization of the 3D Rotary Position Encoding (3D-RPE). The
context size is $L$, and the chunk size is $c$. The vectors
${[\bm{h}_{j,m}^{1},\bm{h}_{j,m}^{2}]}^{T}$ and
${[-\bm{h}_{j,m}^{2},\bm{h}_{j,m}^{1}]}^{T}$ form an orthogonal basis,
corresponding to the $\ket{1}$ and $\ket{0}$ states in Eq. (1). The components
$\bm{h}_{j,m}^{1}$ and $\bm{h}_{j,m}^{2}$ represent the first and second
dimensions of the state vector $\bm{h}_{j,m}$, which is the $m_{th}$ token in
the $j_{th}$ chunk.

###  3.1 3D Rotary Position Encoding

For a long sequence of length $L$ and a chunk size set to $c$, where $c$ is
smaller than the pre-training length of LLM, the sequence can be divided into
$\lceil L/c\rceil$ chunks. Here, $\lceil.\rceil$ represents the ceiling
function, rounding up to the nearest integer (see Figure 2). The state vector
$\bm{h}_{j,m}$ comes from either Query or Key. Here, $j\in[0,{\lceil
L/c\rceil}-1]$ represents the positional index of the chunk, and $m\in[0,c-1]$
indicates the positional index of the token within the chunk. This is used to
calculate the new state vector $\widetilde{\bm{h}}_{j,m}$ by rotating the
Bloch Sphere. Specifically, two rotation angles, $\theta$ and $\varphi$ are
defined, with $\theta$ governing the position encoding within the chunk’s
internal tokens, and $\varphi$ governing the position encoding between the
chunks. Our position encoding method is called 3D Rotary Position Encoding, or
3D-RPE. The formal definition of 3D-RPE is provided as follows. The
computational process of 3D-RPE in practice is provided in Supplementary
Materials B.1.

######  Definition 1 (3D Rotary Position Encoding).

Let $\bm{h}_{j,m}\in\mathbb{R}^{d}$ be a state vector of an attention head
without position encoding, where $d$ is the dimension of the vector, which is
an even number. 3D-RPE encodes $\bm{h}_{j,m}$ into the vector
$\widetilde{\bm{h}}_{j,m}$, which can be formalized as:

$\widetilde{\bm{h}}_{j,m}=e^{\mathrm{i}m\theta}(\cos{\varphi_{j}}\bm{h}_{j,m}^{\perp}+\sin{\varphi_{j}}\bm{h}_{j,m})$
(4)

$\mathrm{i}$ is the imaginary unit. $\bm{h}_{j,m}^{\perp}$ equals to
${[-\bm{h}_{j,m}^{2},\bm{h}_{j,m}^{1}]}^{T}$, where
$\bm{h}_{j,m}^{1}\in\mathbb{R}^{d/2}$ and
$\bm{h}_{j,m}^{2}\in\mathbb{R}^{d/2}$ is the first and second halves of the
state vector $\bm{h}_{j,m}$.

In transformer-based LLMs, after applying position encoding to the state
vectors from Query and Key, it is essential to compute their attention scores.
For the sake of clarity and formalization, we denote the position encoding of
the state vector from Query as 3d-PE$(\bm{q},i,m)$ and from Key as
3d-PE$(\bm{k},j,n)$, where $i$ and $j$ range from $0$ to ${\lceil
L/c\rceil}-1$, and $m$ and $n$ range from $0$ to $c-1$. The self-attention
score can be obtained through the conjugate symmetric inner product of
$\bm{q}_{i,m}$ and $\bm{k}_{j,n}$, which are the state vectors from Query and
Key,

$s(\bm{q}_{i,m},\bm{k}_{j,n},\varphi_{i}-\varphi_{j},m-n)=Re[e^{\mathrm{i}(\varphi_{i}-\varphi_{j})}\sum\limits_{l=0}^{d/2-1}e^{\mathrm{i}(m-n)\theta_{l}}(\bm{q}_{l}\bm{k}_{l}+\bm{q}_{d/2+l}\bm{k}_{d/2+l})]$
(5)

where $l\in[0,{\frac{d}{2}-1}]$, $\varphi_{i}=base^{-i}$ and
$\varphi_{j}=base^{-j}$. Let ${\\{\bm{q},\bm{k}\\}}_{l}$ denote the $l$-th
components of ${\\{\bm{q},\bm{k}\\}}$. In experiments using the LLaMA2 models,
the $base$ is generally set to $10,000$. The self-attention score computed
after applying 3d-PE is a function of both the relative position between
chunks ($\varphi_{i}-\varphi_{j}$) and the relative position ($m-n$).

Consequently, the self-attention score relying on 3d-PE is influenced by the
relative positions at both the chunk and token levels. It is important to
highlight that when $\bm{q}_{i,m}$ and $\bm{k}_{j,n}$ reside within the same
chunk (i.e., $i=j$), Eq. (5) simplifies to the standard RoPE formulation as
depicted in Eq. (3). For a detailed derivation and computation process of Eq.
(5), as well as the complete formulation of Eq. (4), please refer to
Supplementary Materials B.2.

###  3.2 Benefits of 3D-RPE

In this section, we delve into two benefits offered by 3D-RPE: the ability to
control long-term decay and mitigate the reduction in positional resolution
caused by position interpolation.

####  3.2.1 Controllable Long-term Decay

3D-RPE has the property of controllable long-term decay. Like RoPE, taking the
absolute value $s$ in Eq (5) and applying the Abel transformation, we derive
the upper bound of the correlation coefficients related to term dependencies
as follows:

$\displaystyle|s(\bm{q}_{i,m},\bm{k}_{j,n},\varphi_{i}-\varphi_{j},m-n)|$
$\displaystyle\leq|e^{\mathrm{i}(\varphi_{i}-\varphi_{j})}||\sum\limits_{l=0}^{d/2-1}E_{l+1}(h_{l+1}-h_{l})|$
(6)
$\displaystyle\leq(\max_{l}|h_{l+1}-h_{l}|)\sum\limits_{l=0}^{d/2-1}|E_{l+1}|$

where $E_{l}=\sum_{k=0}^{l-1}e^{\mathrm{i}(m-n)\theta_{k}}$ and $E_{0}=0$. For
RoPE [23], the relative upper bound $E_{rope}$ is given by
$\frac{1}{d/2}\sum_{j=1}^{d/2}|S_{j}|$, where
$S_{j}=\sum_{t=0}^{j-1}e^{i(m-n)\theta_{t}}$ (see the section 3.4.3 of RoPE
[23]). By setting $\theta_{t}=10000^{\frac{-2t}{d}}$, the value decays as the
relative position $(m-n)$ increases. For the upper bound $E_{3d-rpe}$ of
3D-RPE, it is formalized as follows:

$E_{3d-rpe}=\frac{1}{d/2}\sum_{j=1}^{d/2}|E_{l}|$ (7)

The domains of the relative position $(m-n)$ differ between $E_{3d-rpe}$ and
$E_{rope}$. In $E_{rope}$, $(m-n)$ is in the range $[0,L-1]$, while in
$E_{3d-rpe}$, it is in $[0,c-1]$. The relative positions between tokens
exceeding the chunk size $c$ are constructed collaboratively using positional
encoding within and across chunks. The Relative Position Matrix $\bm{A}$ using
3D-RPE is shown in Figure 3.

To illustrate the advantage of controllable long-term decay, we present the
results in Figure 1(c) and Figure 1(d). As shown in Figure 1(c), when the
relative position $(m-n)$ exceeds approximately $1000$, $E_{rope}$ begins to
significantly decrease to below $5$. This limitation of $E_{rope}\leq 5$ poses
challenges for RoPE in modeling attention scores between tokens with longer
relative distances (greater than $4000$). In contrast, as shown in Figure
1(d), 3D-RPE employs both $(m-n)$ and $(\varphi_{i}-\varphi_{j})$, setting
$c=1000$ to keep $(m-n)$ within $1000$, thereby preventing decay over longer
distances. This method ensures $E_{3d-rpe}$ stays at or above $5$ for all
relative positions.

Figure 3: Visualization of the Relative Position Matrix $\bm{A}$ employing
3D-RPE, with chunk size $c$=$4$, and sequence size $L$=$12$. The matrix
elements $A_{i,j}$ represents the relative position between the $i_{th}$ query
vector $\bm{q}$ and the $j_{th}$ key vector $\bm{k}$.

####  3.2.2 Enhanced Positional Resolution

Position Interpolation (PI) [4] has been introduced to scale down the position
indices to align with the original window size, resulting in enhanced outcomes
for context extension. However, as the extension length and interpolation
increase, PI can lead to a reduction in relative positional resolution. 3D-RPE
can be used alongside PI for long-context extensions. Compared to RoPE
combined with PI, 3D-RPE has the advantage of mitigating the reduction in
positional resolution caused by positional interpolation, as demonstrated in
Theorem 1.

######  Theorem 1 (Enhanced Position Resolution).

For a pre-trained language model with a length of $L_{p}$ and an extension
length requirement of $L$, employing linear position interpolation extension
methods $\mathcal{I}$ based on Rotary Position Encoding (RoPE) can elevate the
relative positional resolution from $\mathcal{E}_{rope}$ to
$\mathcal{E}_{rope}^{\prime}$. Let $\mathcal{E}_{3d-rpe}^{\prime}$ denote the
relative positional encoding resolution achieved by the method $\mathcal{I}$
based on 3D-RPE, with chunk size $c\geq 3$, there is:

$\mathcal{E}_{3d-rpe}^{\prime}>\mathcal{E}_{rope}^{\prime}$ (8)

The Proof of Theorem 1 is provided in Supplementary Materials C.

To empirically validate the superior performance of this benefit in a
training-free setting, it has been observed that methods combining RoPE with
interpolation lead to a significant increase in Perplexity as the modeling
length increases in language modeling tasks. Conversely, the increase in
Perplexity is substantially smaller when employing 3D-RPE with linear
interpolation (Refer to Table 4 in Section 5). This phenomenon indicates that
this benefit has led to an improvement in the performance of long sequence
language modeling.

##  4 Related Work

This section provides an overview of the extensive literature on position
encoding in Transformers [26] and discusses context extending capabilities
based on RoPE.

Position Encoding (PE): PE is important for Transformer-based language models.
Earlier studies [22, 21, 28, 23] have focused on enhancing the original
absolute position encoding to develop better relative position encoding,
thereby improving the text modeling capabilities of language models. These
works [22, 21, 28] utilized trainable position vector encoding to directly
incorporate positional information into context representations. Although
effective, these methods typically add positional information to contextual
representations, making them unsuitable for linear self-attention
architectures. RoFormer [23] introduced relative position information by
rotating context representations, known as RoPE. Transformers utilizing RoPE
have become a prevalent backbone in various LLM designs [24, 8, 27, 16]. Our
proposed 3D-RPE differs from the two-dimensional space of RoPE by modeling the
relative position information of tokens through rotation on the Bloch Sphere.

Long-context LLMs based on RoPE: To enhance the contextual capabilities of
Large Language Models (LLMs) using RoPE, several positional encoding
interpolation techniques have been developed. These include Linear Position
Interpolation (LPI) [4], Neural Tangent Kernel (NTK) [17], and Yet Another
Recurrent Network (YaRN) [18] interpolation. Position Sequence Tuning (PoSE)
[31] has notably increased sequence lengths to $128k$ by amalgamating these
positional interpolation strategies. Additionally, LongLora [5] introduced the
shift-short attention mechanism, allowing for effective emulation of full
attention and extending sequences up to $100k$, leveraging the LLMa-2-7B model
and LoRA’s fine-tuning approach [12]. 3D-RPE further strengthens the
positional relationships between distant tokens by capturing inter-chunk
positional information and is compatible with existing fine-tuning techniques
like LoRA to bolster long-context representation. The Dual Chunk Attention
(DCA) [2] method, which enhances the use of pre-trained integer-based
parameters, splits query and key sequences into chunks and uses three
specialized matrices to capture the relative positions within and between
these chunks. This method enhances the model’s ability to process longer
sequences, but it is unable to model the relative positions within distant
chunks. In our work, we employ rotating positional encoding to link attention
across different chunks.

##  5 Experiments

We evaluate the method of position encoding, 3D-RPE, on LLaMA2 [24] models
(specifically, LLaMA-2-7B and LLaMA-2-7B-chat), which have a $4k$ pre-training
context, and LLaMA-3-8B-Instruct 111https://github.com/meta-llama/llama3,
which has an $8k$ pre-training context. Our experiments aim to explore the
following aspects: 1) The effect of 3D-RPE on long-context generation can be
assessed using Perplexity. 2) The impact of 3D-RPE on long-context
understanding and generation tasks, can be reflected by the accuracy of long
sequence natural language tasks, e.g., multiply documents QA. 3) Ablation
studies to confirm the advantages of 3D-RPE in position interpolation.

###  5.1 Experimental Settings

In this section, we elaborate on the experimental setup by introducing two
types of tasks (i.e., long-context language understanding and long sequence
language modeling) and detailing three aspects of the configuration (i.e.,
training parameters, training and evaluation datasets, and baseline models).

Training Setting: For long-context Natural Language Understanding (NLU) tasks,
we have fine-tuned LLaMA-2-7B-chat and LLaMA-3-8B-Instruct. The context length
for these models has been extended from $4k$ to $16k$ and from $8k$ to $16k$,
respectively. The fine-tuning method follows the fine-tuning strategy of
LongChat [13]. The training step is $3,000$. For the long-sequence Language
Modeling (LM) tasks, we have fine-tuned LLaMA-2-7B to support extended context
length of $32k$ tokens. The training step is $1,000$. We set the per-device
batch size as $1$, and gradient accumulation step as $8$, which means that the
batch size is $8$. We train the model with the next token prediction objective
with LoRA [12].

We employed the AdamW optimizer [15] with ${\beta}_{1}=0.9$ and
${\beta}_{2}=0.95$ for all fine-tuned models. Chunk size is set to $3k$. The
learning rate was set to $2\times 10^{-5}$, and a linear learning rate warmup
was applied. Training was conducted on a single 4xA800 GPU machine using
FlashAttention-2 [10].

Datasets: In the context of long-context NLU tasks, we employ the
LongAlpaca-12k dataset, which contains 9,000 LongQA and 3,000 short QA entries
[6], and the LongAlpace-16k-length dataset222https://github.com/dvlab-
research/LongLoRA/. To evaluate the performance of 3D-RPE for long-context
extension, we use the LongBench [3], which includes $13$ English tasks, $5$
Chinese tasks and $2$ code tasks, with most tasks having an average context
length of $5k$ to $15k$ tokens. We focus on the English and code tasks to
evaluate our method, 3D-RPE. Additionally, the LEval [1] evaluation set, which
also consists of long-context datasets, is used to verify the effectiveness of
3D-RPE. The five datasets annotated from scratch in LEval, namely Coursera,
QuALiTY, CodeU, GSM,and TOEFL, are utilized.

For long-sequence LM tasks, we use the RedPajama-Data [9] for fine-tuning
training. The dataset is a large-scale pre-training dataset (the size reaches
1.2 trillion tokens) designed to provide high-quality training data for
language models, and contains multiple data sources (i.e., github, arxiv,
book, c4 and Wikipedia, etc.). We sample $20,000$ samples from these data
sources for training. For evaluation, we utilize the PG19 book corpus dataset
[20], which includes 100 documents, and the Arxiv Math Proof-pile dataset
(test split). Additionally, all methods evaluate perplexity by using a sliding
window following [19].

Baseline models: For long-context NLU tasks, the fine-tuned models, including
LongAlpace-16k [5], LongChat-32k [14] LongLlama [25] and ChatGLM [11] are used
as the baseline models. Models of fine-tuning free in language modeling tasks
are also used in long-context NLU tasks.

In long sequence LM tasks, the methods of LongLoRA [5], StreamingLLM [29],
Positional Interpolation(PI) [4] and the NTK-Aware Scale RoPE(NTK) [17] are
selected as the baselines, all based on the LLaMA-2-7B-base model. Among these
baseline models, PI, NTK and StreamingLLM are fine-tuning-free methods. The
fine-tuned models include LongLoRA and Activation Beacon [30]. In Ablation
experiments, interpolation methods without training are used as baseline
models, which are PI and NTK.

Table 1: Comparison between open-source based models on long-context NLU tasks. Our model, 3D-RPE-LlaMA2-7B-Chat is fine-tuning on LLaMA-2-7b-chat, which is extended from $4k$ to $16k$ context lengths. Baseline models can be categorized into two groups: those that necessitate fine-tuning during training (such as LongAlpaca [5] and LongLLaMA [25]), and those that do not require it (including PI, NTK, StreamingLLM [29], and ChunkLLaMA-$16k$ [2]). The experimental results for each specific evaluation set in Supplementary Material D.2. Methods | Single-Doc QA | Multi-Doc QA | Summarization | Few-shot | Code
---|---|---|---|---|---
LLaMA-2-7B-chat | 24.90 | 22.60 | 24.70 | 60.01 | 48.10
LLaMA-2-7B-chat-PI | 18.98 | 17.16 | 25.03 | 49.43 | 52.73
LLaMA-2-7B-chat-NTK | 23.21 | 23.34 | 24.40 | 59.29 | 49.28
StreamingLLM | 21.47 | 22.22 | 22.20 | 50.05 | 48.00
ChunkLLaMA-$16k$ | 24.04 | 22.98 | 21.52 | 46.31 | 49.73
LongChat-$32k$ | 31.58 | 23.50 | 26.70 | 64.02 | 54.10
LongAlpaca-$16k$ | 28.70 | 28.10 | 27.80 | 63.70 | 56.00
LongLLaMA | 30.12 | 16.37 | 24.19 | 60.31 | 66.05
Vicuna-v1.5-7B-$16k$ | 28.01 | 18.63 | 26.01 | 66.20 | 47.30
ChatGLM3-6B-$32k$ | 40.30 | 46.60 | 29.50 | 68.10 | 56.20
3D-RPE-LLaMA2-7B-Chat | 47.40 | 60.10 | 28.99 | 73.16 | 76.50
Table 2: Comparison with open-source models, LLaMA-2-7B-chat, LLaMA3-8B-Instruct, on 5 closed-ended-ended tasks with various input length from LEval [1]. The evaluation metric “EM,” which represents the exact match score, is adopted. * indicates the model is train-free. Models | Coursera | QuALiTY | CodeU | GSM | TOEFL
---|---|---|---|---|---
LLaMA-2-7B-Chat | 29.21 | 37.62 | 1.11 | 19.00 | 51.67
LongChat-7B-16K | 29.74 | 33.66 | 3.33 | 10.00 | 47.95
LLaMA2-7B-NTK | 32.71 | 33.16 | 0.00 | 19.00 | 52.78
Vicuna1.5-7B-16k | 38.66 | 39.60 | 5.55 | 19.00 | 55.39
3D-RPE-LLaMA2-7B-Chat(ours) | 39.38 | 38.11 | 2.22 | 21.01 | 57.99
LLaMA3-8B-Instruct* | 51.45 | 64.34 | 4.44 | 76.00 | 82.89
3D-RPE-LLaMA3-8B-Instruct* | 51.89 | 61.38 | 4.44 | 80.00 | 82.89
Table 3: Perplexity evaluation on different extending methods. We conduct evaluation on the Proof-pile and PG-19 test datasets, varying evaluation context window size from $8k$ to $100k$. Methods | PG-19 | Proof-Pile
---|---|---
$8k$ | $16k$ | $32k$ | $100k$ | $8k$ | $16k$ | $32k$ | $100k$
LLaMA2-7B-Base | 131.09 | $>10^{2}$ | $>10^{2}$ | OOM | 16.79 | $>10^{2}$ | $>10^{2}$ | OOM
LLama2-7B-PI | 11.32 | 19.5 | $>10^{2}$ | OOM | 3.86 | 5.94 | 33.7 | OOM
LLama2-7B-NTK | 10.28 | 11.5 | 37.8 | OOM | 3.98 | 5.94 | 33.7 | OOM
StreamingLLM | 9.23 | 9.25 | 9.24 | 9.32 | 3.47 | 3.51 | 3.50 | 3.55
LongLoRA-32k | 7.33 | 7.16 | 7.04 | – | 2.78 | 2.61 | 2.50 | –
LongLoRA-100k | 7.57 | 7.33 | 7.16 | 7.04 | 2.78 | 2.60 | 2.58 | 2.52
LongChat-32k | 8.92 | 8.85 | 8.81 | OOM | 2.98 | 2.70 | 2.65 | OOM
Activation Beacon | 8.52 | 8.54 | 8.56 | 8.68 | 3.45 | 3.42 | 3.39 | 3.35
3D-RPE-LLaMA2-7B | 7.03 | 7.10 | 8.09 | 8.12 | 2.72 | 2.93 | 2.89 | 3.05
Table 4: Results are evaluated in Perplexity on PG19 validation split. ’*’ denotes train-free, implementing 3D-RPE directly on the LLaMA2-7B-Base model without additional fine-tuning, utilizing a chunk size of $3k$. The context length of $8k$ is extended directly with 3D-RPE. Achieving $16k$ and $32k$ is accomplished through linear positional interpolation with chunks based on the $8k$ context length.  Models | $4k$ | $8k$ | $16k$ | $32k$
---|---|---|---|---
LLaMA2-7B-PI | 7.94 | 9.19 | 15.11 | $>10^{2}$
LLaMA2-7B-NTK | 7.87 | 11.98 | 26.12 | 58.91
LLaMA2-7B-Yarn | 7.87 | 8.06 | 9.82 | 11.74
3D-RPE-LLaMA2-7B* | 7.87 | 7.90 | 7.71 | 9.34

###  5.2 Long-Context Natural Language Understanding

In this task, the LongBench [3] evaluation set was initially utilized. Five
categories of tasks were included: single-document QA (3 tasks), multi-
document QA (3 tasks), summarization (3 tasks), few-shot learning (3 tasks),
and code completion (2 tasks). The average score for each type is reported in
Table 1. The evaluation metrics followed those specified in LongBench [3],
which differ across tasks and are detailed in Supplementary Material D.1. The
results in Table 1 highlight our model’s significant performance advantages
over baseline models in four tasks, both for models without training and those
with fine-tuning. In summarization tasks, our model also achieved performance
comparable to ChatGLM3-6B-$32k$. These experimental outcomes indicate that our
model enhances the correlation between tokens with distant relative positions
in long contexts through 3D-RPE, resulting in improved performance.

Subsequently, the LEval Benchmark [1] was employed. Table 2 reveals that our
model, 3D-RPE-LLaMA2-7B-Chat, outperformed LLaMA2-7B-NTK and
LongChat-7B-$16K$. Although it did not surpass Vicuna1.5-7B-$16K$ in Quality
and CodeU tasks, it excelled in the Coursera, GSM, and TOEFL tasks.
Additionally, we conducted experiments on LLaMA3-8B-Instruct using a 16k
context window with 3D-RPE. The 3D-RPE-LLaMA3-8B-Instruct* showed performance
improvements in the Coursera and GSM tasks. While 3D-RPE did not enhance
performance in the CodeU, TOEFL, and QuALiTY tasks, there was no significant
performance decline either. These experimental results demonstrate the
effectiveness of the 3D-RPE method.

###  5.3 Long-Sequence Language Modeling

In Table 3, we present the perplexity scores for our model, 3D-RPE-LLaMA-2-7B
and baseline models on the proof-pile and PG19 test datasets. 3D-RPE-
LLaMA-2-7B was fine-tuned from the LLaMA2-7B-Base model using a dataset with a
$32k$ context window. To evaluate performance, we set sequence lengths of
$8k$, $16k$, and $32k$. We extended our model’s sequence length from $32k$ to
$100k$ using the position extending method from PoSE [31]. The results
indicate that our method outperforms train-free sequence extending models.
Compared to fine-tuned models, our model shows better performance at $8k$ and
$16k$ sequence lengths. This suggests that the new positional encoding,
3D-RPE, improves or maintains modeling performance for larger context windows
($32k$) compared to smaller ones ($8k$ and $16k$). For the $32k$ and $100k$
tasks, although our model did not surpass LongLoRA-$32k$ and LongLoRA-$100k$,
it did outperform LongChat-$32k$ and Activation Beacon.

Notably, our model can further extend from $32k$ to $100k$ without
significantly increasing perplexity values, in combination with other train-
free extension methods. However, due to its specific attention mechanism, the
LongLoRA models cannot be extended beyond their predefined context windows in
a train-free manner. For instance, LongLoRA-$32k$ cannot be further extended
to $100k$.

###  5.4 Ablation Study

In this section, we conduct ablation studies in this section to explore how
3D-RPE affects the linear interpolation method. We compare position
interpolation methods (PI, NTK, and Yarn) with the method that combines 3D-RPE
with position interpolation on the LLaMA-2-7B-Base model in a train-free
manner. The experimental results can be found in Table 2. The 3D-RPE-
LLaMA2-7B* model with linearly positional interpolation from $8k$ to $16k$ and
$32k$, the 3D-RPE approach yields improved results by mitigating the decrease
in positional resolution caused by interpolation methods. These results are
consistent with the findings of Theorem 1 in Section 3.2.2 presented in this
paper.

##  6 Conclusion and Future Work

In this paper, we present a novel rotary position encoding method called 3D
Rotary Position Encoding (3D-RPE). Compared to RoPE, we have theoretically
proved that 3D-RPE possesses two key advantages: controllable long-term decay
and enhanced interpolation resolution. Experimentally, 3D-RPE has demonstrated
outstanding performance in long-context Natural Language Understanding.

In the future, 3D-RPE holds promise as a foundational positional encoding
strategy for LLMs, especially in the aspect of modeling long contexts.
Moreover, given that 3D-RPE encapsulates positional encoding within a three-
dimensional framework, it has the potential to integrate with visual data,
thereby facilitating an in-depth exploration of its efficacy in synchronizing
graphical and textual semantic information.

## References

  * [1] Chenxin An, Shansan Gong, Ming Zhong, Mukai Li, Jun Zhang, Lingpeng Kong, and Xipeng Qiu.  L-eval: Instituting standardized evaluation for long context language models.  arXiv preprint arXiv:2307.11088, 2023.
  * [2] Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, Chang Zhou, and Lingpeng Kong.  Training-free long-context scaling of large language models, 2024.
  * [3] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li.  Longbench: A bilingual, multitask benchmark for long context understanding.  arXiv preprint arXiv:2308.14508, 2023.
  * [4] Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian.  Extending context window of large language models via positional interpolation.  arXiv preprint arXiv:2306.15595, 2023.
  * [5] Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia.  Longlora: Efficient fine-tuning of long-context large language models.  arXiv:2309.12307, 2023.
  * [6] Yukang Chen, Shaozuo Yu, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia.  Long alpaca: Long-context instruction-following models.  https://github.com/dvlab-research/LongLoRA, 2023.
  * [7] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E Gonzalez, et al.  Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, march 2023.  URL https://lmsys. org/blog/2023-03-30-vicuna, 3(5), 2023.
  * [8] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al.  Palm: Scaling language modeling with pathways.  arXiv preprint arXiv:2204.02311, 2022.
  * [9] Together Computer.  Redpajama: An open source recipe to reproduce llama training dataset.  https://github.com/togethercomputer/RedPajama-Data, 2023.
  * [10] Tri Dao.  FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.  CoRR, 2023.
  * [11] Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang, and Jie Tang.  Glm: General language model pretraining with autoregressive blank infilling.  In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 320–335, 2022.
  * [12] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  LoRA: Low-rank adaptation of large language models.  In International Conference on Learning Representations, 2022.
  * [13] Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang.  How long can context length of open-source llms truly promise?  In NeurIPS 2023 Workshop on Instruction Tuning and Instruction Following, 2023.
  * [14] Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph E. Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang.  How long can open-source llms truly promise on context length?  arXiv preprint arXiv:2306.04537, June 2023.
  * [15] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In ICLR, 2019, 2019.
  * [16] Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese, and Caiming Xiong.  Codegen: An open large language model for code with multi-turn program synthesis.  arXiv preprint arXiv:2203.13474, 2022.
  * [17] Bowen Peng and Jeffrey Quesnelle.  Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation.  https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have, 2023\.
  * [18] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole.  Yarn: Efficient context window extension of large language models.  arXiv preprint arXiv:2203.13474, 2023.
  * [19] Oriol Press, Noah A Smith, and Michael Lewis.  Train short, test long: Attention with linear biases enables input length extrapolation.  In ICLR, 2022, 2022.
  * [20] Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap.  Compressive transformers for long-range sequence modelling.  In International Conference on Learning Representations, 2020.
  * [21] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  Journal of Machine Learning Research, 21(140):1–67, 2020.
  * [22] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani.  Self-attention with relative position representations.  In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 464–468, 2018.
  * [23] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding.  Neurocomputing, 568:127063, 2024.
  * [24] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.  Llama 2: Open foundation and fine-tuned chat models.  arXiv preprint arXiv:2307.09288, 2023.
  * [25] Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Miłos.  Focused transformer: Contrastive training for context scaling.  arXiv preprint arXiv:2307.03170, 2023.
  * [26] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [27] Ben Wang and Aran Komatsuzaki.  Gpt-j-6b: A 6 billion parameter autoregressive language model.  GitHub, 2021.
  * [28] Benyou Wang, Donghao Zhao, Christina Lioma, Qiuchi Li, Peng Zhang, and Jakob Grue Simonsen.  Encoding word order in complex embeddings.  In International Conference on Learning Representations, 2020.
  * [29] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis.  Efficient streaming language models with attention sinks.  arXiv, 2023.
  * [30] Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou.  Soaring from 4k to 400k: Extending llm’s context with activation beacon, 2024.
  * [31] Dawei Zhu, Nan Yang, Liang Wang, Yifan Song, Wenhao Wu, Furu Wei, and Sujian Li.  Pose: Efficient context window extension of llms via positional skip-wise training.  arXiv preprint arXiv:2309.10400, 2023.

##  Appendix A Bloch Sphere

Figure 4: A diagram of Bloch Sphere.

Bloch Sphere: 3D Rotary Position Encoding (3D-RPE), proposed by us,
corresponds to a Bloch Sphere. In this section, we mainly introduce the basic
concept of Bloch Sphere, which corresponds to Eq. (1) in this paper.

The Bloch Sphere is a geometric tool to used to represent qubits, typically
depicted in a three-dimensional polar coordinate system as a point on the
Sphere (see Figure 4). A single quantum state is represented by the following
equation in linear algebra:

$\ket{\phi}=\alpha\ket{0}+\beta\ket{1}$ (9)

where $\alpha$ and $\beta$ are complex numbers, i.e.,
$\alpha,\beta\in\mathbb{C}$. According to Euler’s formula in complex analysis,
the coefficients $\alpha$ and $\beta$ can be reexpressed as:

$\displaystyle\alpha=a+b\mathrm{i}=r_{0}(cos(\theta_{\alpha})+\mathrm{i}sin(\theta_{\alpha}))=r_{0}e^{\mathrm{i}\theta_{\alpha}}$
(10)
$\displaystyle\beta=c+d\mathrm{i}=r_{1}(cos(\theta_{\beta})+\mathrm{i}sin(\theta_{\beta}))=r_{1}e^{\mathrm{i}\theta_{\beta}}$

By substituting Eq. (10) into Eq. (9), the quantum state representation is
denoted as:

$\displaystyle\ket{\phi}$
$\displaystyle=r_{0}(cos(\theta_{\alpha})+\mathrm{i}sin(\theta_{\alpha}))\ket{0}+r_{1}(cos(\theta_{\beta})+\mathrm{i}sin(\theta_{\beta}))\ket{1}$
(11)
$\displaystyle=r_{0}e^{\mathrm{i}\theta_{\alpha}}\ket{0}+r_{1}e^{\mathrm{i}\theta_{\beta}}\ket{1}$
$\displaystyle=e^{\mathrm{i}\theta_{\alpha}}(r_{0}\ket{0}+r_{1}e^{\mathrm{i}(\theta_{\beta}-\theta_{\alpha})}\ket{1})$

$\theta_{\alpha}$ is the global phase.

Considering the normalization condition $|\alpha|^{2}+|\beta|^{2}=1$, we have:

$|r_{0}|^{2}+|r_{1}e^{\mathrm{i}(\theta_{\beta}-\theta_{\alpha})}|^{2}=r_{0}^{2}+r_{1}^{2}|e^{\mathrm{i}(\theta_{\beta}-\theta_{\alpha})}|^{2}=r_{0}^{2}+r_{1}^{2}=1$
(12)

Given $r_{0}=cos\frac{\varphi}{2}$, $r_{1}=sin\frac{\varphi}{2}$,
$\theta_{\alpha}=\theta$ and $\theta_{\beta}-\theta_{\alpha}=\theta_{1}$, the
state $\ket{\phi}$ can be expressed as:

$\ket{\phi}=e^{\mathrm{i}\theta}(cos\frac{\varphi}{2}\ket{0}+sin\frac{\varphi}{2}e^{\mathrm{i}\theta_{1}}\ket{1})$
(13)

Therefore, the Eq. (1) of this paper is given out.

To adapt to the original 2D rotation position encoding (RoPE) of pre-trained
LLMs, such as LLaMA models, the global phase $\theta$ is used to model the
relative positions between tokens within a chunk, while the rotation angle
$\frac{\varphi}{2}$ is used to model the relative positions between tokens
across chunks.

##  Appendix B Supplementary Material for the Method Section

In this section, we mainly introduce the specific implementation of our
positional encoding method (3D-RPE), and the formula derivation details of
attention score calculation (Eq. (5) of this paper) not detailed in this
paper.

###  B.1 Implement of 3D-RPE

In Section 3.1, we give the general form of 3D Rotary Position Encoding
(3D-RPE):

$\widetilde{\bm{h}}_{j,m}=e^{\mathrm{i}m\theta}(\cos{\varphi_{j}}\bm{h}_{j,m}^{\perp}+\sin{\varphi_{j}}\bm{h}_{j,m})$

$\cos{\varphi_{j}}$ and $\sin{\varphi_{j}}$ are scalar quantities in
$\mathbb{R}$. $\bm{h}_{j,m}^{\perp}$ and $\bm{h}_{j,m}$ are shown below:

$\bm{h}_{j,m}=\begin{bmatrix}h^{0}\\\ h^{1}\\\ \vdots\\\ h^{d/2-2}\\\
h^{d/2-1}\\\ h^{d/2}\\\ h^{d/2+1}\\\ \vdots\\\ h^{d-2}\\\ h^{d-1}\\\
\end{bmatrix}\quad\quad\quad\quad\quad\quad\bm{h}_{j,m}^{\perp}=\begin{bmatrix}-h^{d/2}\\\
-h^{d/2+1}\\\ \vdots\\\ -h^{d-2}\\\ -h^{d-1}\\\ h^{0}\\\ h^{1}\\\ \vdots\\\
h^{d/2-2}\\\ h^{d/2-1}\\\ \end{bmatrix}$ (14)

In the concrete implementation, analogous to RoPE, for each two-dimensional
subspace $\mathbb{R}^{2}$ of $\mathbb{R}^{d}$, we assign angles
$\theta_{l}=base^{-2l/d}$ that vary from high to low frequencies. An
equivalent rotation matrix $\mathcal{R}_{m}^{\theta}$ is utilized to
substitute for $e^{\mathrm{i}m\theta}$:

$\mathcal{R}^{\theta}_{m}={\left[\begin{array}[]{cccccccc}\cos
m\theta_{0}&0&\cdots&0&-\sin m\theta_{0}&0&\cdots&0\\\ 0&\cos
m\theta_{1}&\cdots&0&0&-\sin m\theta_{1}&\cdots&0\\\
\vdots&\vdots&\ddots&\vdots&\vdots&\vdots&\ddots&\vdots\\\ 0&0&\cdots&\cos
m\theta_{d/2-1}&0&0&\cdots&-\sin m\theta_{d/2-1}\\\ \sin
m\theta_{0}&0&\cdots&0&\cos m\theta_{0}&0&\cdots&0\\\ 0&\sin
m\theta_{1}&\cdots&0&0&\cos m\theta_{1}&\cdots&0\\\
\vdots&\vdots&\ddots&\vdots&\vdots&\vdots&\ddots&\vdots\\\ 0&0&\cdots&\sin
m\theta_{d/2-1}&0&0&\cdots&\cos m\theta_{d/2-1}\\\ \end{array}\right]}$ (15)

Therefore, Eq.(4) of this paper can be transformed to

$\widetilde{\bm{h}}_{j,m}=\mathcal{R}^{\theta}_{m}(\cos{\varphi_{j}}\bm{h}_{j,m}^{\perp}+\sin{\varphi_{j}}\bm{h}_{j,m})$

where $\mathcal{R}^{\theta}_{m}$ is a design form equivalent to the rotation
matrix in RoPE, mainly re-mapped to correspond to specific application
implementations and calculation derivations in LLMs. In the specific
implementation, after the rotary position encoding of LLMs, the long sequence
is chunked based on the chunk size $c$. Then, the rotation $\varphi_{j}$ is
set on each chunk, $j$ is the position of chunk.

###  B.2 Derivation of Attention for 3D-RPE

The formula derivation details of attention score calculation(Eq. (5)) is as
follows.

Since $\bm{h}^{\perp}=e^{\mathrm{i}\frac{\pi}{2}}\bm{h}=\mathrm{i}\bm{h}$, we
could get:

$\displaystyle\widetilde{\bm{h}}_{j,m}$
$\displaystyle=e^{\mathrm{i}m\theta}(\mathrm{i}\cos{\varphi_{j}}\bm{h}_{j,m}+\sin{\varphi_{j}}\bm{h}_{j,m})$
(16)
$\displaystyle=e^{\mathrm{i}m\theta}(\mathrm{i}\sin{(\frac{\pi}{2}-\varphi_{j}})\bm{h}_{j,m}+\cos{(\frac{\pi}{2}-\varphi_{j}})\bm{h}_{j,m})$
$\displaystyle=e^{\mathrm{i}m\theta}e^{\mathrm{i}\frac{\pi}{2}-\varphi_{j}}\bm{h}_{j,m}$

Let $\bm{q}_{i,m}$=3d-PE$(\bm{q},i,m)$, $\bm{k}_{j,n}$=3d-PE$(\bm{k},j,n)$.
Taking the real part of the inner product of $\bm{q}_{i,m}$ and $\bm{k}_{j,n}$
yields:

$s(\bm{q}_{i,m},\bm{k}_{j,n})=Re[e^{\mathrm{i}(\varphi_{i}-\varphi_{j})}\sum\limits_{l=0}^{d/2-1}e^{\mathrm{i}(m-n)\theta_{l}}(\bm{q}_{l}\bm{k}_{l}+\bm{q}_{d/2+l}\bm{k}_{d/2+l})]$
(17)

which is a function related to both $m-n$ and $\varphi_{i}-\varphi_{j}$.

##  Appendix C 3D Rotary Position Encoding Resolution Enhancement

In this section, before proving Theorem 1, we first provide the definitions of
positional resolution for RoPE, as well as the positional resolution after
positional interpolation.

######  Definition 2 (Positional Interpolation Resolution).

Let $\bm{q}_{m+1}$ and $\bm{k}_{m}$ be query state and key state of the $m$-th
and ${(m+1)}$-th hidden states after RoPE. Given a pre-training length
$L_{p}$, the attention score $a$ is:

$a(\bm{q}_{m+1},\bm{k}_{m})=\bm{q}\bm{k}^{T}e^{\mathrm{i}\theta}$ (18)

The Resolution $\mathcal{E}_{rope}$ corresponding to the initial length
$L_{p}$ is $\mathcal{E}_{rope}=1$. After employing linear interpolation with
length $L\geq L_{p}$, the attention score is:

$a(\bm{q}_{m+1},\bm{k}_{m})=\bm{q}\bm{k}^{T}e^{\mathrm{i}\frac{L_{p}}{L}\theta}$
(19)

Note that the Resolution $\mathcal{E}_{rope}$ turns to
$\mathcal{E}^{\prime}_{rope}=L_{p}/L\leq 1$ and decreases as $L$ increases.

As the resolution decreases, the magnitude of the rotation of attention score
becomes smaller, reflecting the extent of positional difference becomes
smaller. Now we give the following theorem, explaining how 3D-RPE mitigates
the resolution decreasing in detail.

######  Theorem 2 (Chunk Position Encoding Resolution Enhancement).

For a pre-trained language model with a pre-training length $L_{p}$ and an
extension length requirement of $L$, employing linear position interpolation
extension methods $\mathcal{I}$ based on Rotary Position Encoding (RoPE) can
elevate the relative positional resolution from $\mathcal{E}_{rope}$ to
$\mathcal{E}_{rope}^{\prime}$, Let $\mathcal{E}_{3d-rpe}^{\prime}$ denote the
relative positional encoding resolution achieved by the method $\mathcal{I}$
based on 3D-RPE, with chunk size $c>=3$, there is:

$\mathcal{E}_{3d-rpe}^{\prime}>\mathcal{E}_{rope}^{\prime}$ (20)

###### Proof.

For 3D-RPE, let the chunk size and chunk number be denoted as $c$ and
$n=\lceil L_{p}/c\rceil$ respectively. Prior to interpolation, the indices
within a chunk range from $[0,1,\cdots,c-1]$. Linear interpolation involves
evenly distributing the excess $L-L_{p}$ tokens across $n$ chunks. This
results in new indices within the chunk, range from
$[0,1,2,\cdots,c^{\prime}-1]$, where $c^{\prime}=\lceil L/n\rceil\leq L_{p}$.
So the attention score of $\bm{q}_{i,m+1}$ and $\bm{k}_{i,m}$ based on 3D-RPE
after interpolation is:

$\displaystyle a_{3d-rpe}$
$\displaystyle=\bm{q}\bm{k}^{T}e^{\mathrm{i}\theta}e^{\mathrm{i}(\varphi_{i}-\varphi_{i})}$
$\displaystyle=\bm{q}\bm{k}^{T}e^{\mathrm{i}\theta}$

The resolution of relative position for 3D-RPE is:

$\mathcal{E}_{3d-rpe}^{\prime}=1$

For special cases $\bm{q}_{(i+1,0)}$ and $\bm{k}_{(i,c^{\prime}-1)}$:

$\mathcal{E}_{3d-rpe}^{\prime}\geq
c^{\prime}-1+\frac{(\varphi_{i+1}-\varphi_{i})}{\theta}>c^{\prime}-2\geq 1$
(21)

where $(\varphi_{i+1}-\varphi_{i})/\theta\geq-1/10000>-1$. As long as
$c^{\prime}\geq 3$, there is $\mathcal{E}_{3d-rpe}^{\prime}\geq
1>\mathcal{E}_{rope}^{\prime}=L_{p}/L$. Under normal case, the chunk size $c$
is not set to a very small number, hence $c^{\prime}\geq 3$ is certainly
established; moreover, for different interpolation lengths $L$, we need to
configure a varying number of chunks $n$, such that $c^{\prime}=\lceil
L/n\rceil\leq L_{p}$. ∎

##  Appendix D Experimental Supplementary Materials

###  D.1 Evaluation Metrics

This section mainly presents the utilization of evaluation metrics for a total
of 16 tasks from the LongBench.

Dataset | Metric
---|---
Narrative QA | F1_Score
Qsper | F1_Score
MultiFieldQA-En | F1_Score
Hotpot QA | F1_Score
2WikiM QA | F1_Score
Musique | F1_Score
GovReport | Rouge_Score
QMSum | Rouge_Score
MultiNews | Rouge_Score
Trec | Classification_Score
Trivia QA | F1_Score
SAMsum | Rouge_Score
PassageRetrieval-En | Retrieval_Score
Passage Count | Count_Score
Lcc | Code_Sim_Score
RepoBench-P | Code_Sim_Score

###  D.2 Details of Experimental Results

This section mainly presents the performance of all tasks corresponding to
each type of experiment in LongBench. These experimental results are reported
in Table 5.

Table 5: Comparison of Experimental Performance on Different Datasets for
Various Tasks in LongBench, Using Baseline Models Provided by LongBench.
3D-RPE-LLaMA2-7B is our model.

Single-Document QA |  Narrative QA |  Qasper |  MultiFieldQA-En
---|---|---|---
LLaMA2-7B-Chat-$4k$ |  18.7 |  19.2 |  36.8
LongChat-v1.5-7B-$32k$ |  16.9 |  27.7 |  41.4
InternLM-7B-$8k$ |  12.1 |  16.7 |  23.4
Vicuna-v1.5-7B-$16k$ |  19.4 |  26.1 |  38.5
LongLora-$16k$ |  19.8 |  29.1 |  37.1
3D-RPE-LLaMA2-7B(our) |  40.56 |  41.35 |  60.3

Multi-Document QA |  Hotpot QA |  2WikiM QA |  Musique
---|---|---|---
LLaMA2-7B-chat-$4k$ |  25.4 |  32.8 |  9.4
LongChat-v1.5-7B-$32k$ |  31.5 |  20.6 |  9.7
InternLM-7B-$8k$ |  28.7 |  22.8 |  9.0
Vicuna-v1.5-7B-$16k$ |  25.3 |  20.8 |  9.8
LongLora-$16k$ |  37.01 |  30.26 |  17.14
3D-RPE-LLaMA2-7B(our) |  62.49 |  58.80 |  59.01

Summarization |  GovReport |  QMSum |  MultiNews
---|---|---|---
LLaMA2-7B-chat-$4k$ |  27.3 |  20.8 |  25.8
LongChat-v1.5-7B-$32k$ |  30.8 |  22.7 |  26.4
InternLM-7B-$8k$ |  9.7 |  15.9 |  22.8
Vicuna-v1.5-7B-$16k$ |  27.9 |  22.8 |  27.2
LongLora-$16k$ |  31.53 |  24.13 |  27.74
3D-RPE-LLaMA2-7B(our) |  32.01 |  25.3 |  29.68

Few-shot Learning |  Trec |  Trivia QA |  SAMSum
---|---|---|---
LLaMA2-7B-chat-$4k$ |  61.5 |  77.8 |  40.7
LongChat-v1.5-7B-$32k$ |  63.5 |  82.3 |  34.2
InternLM-7B-$8k$ |  52.0 |  77.8 |  21.2
Vicuna-v1.5-7B-$16k$ |  71.5 |  86.2 |  40.8
LongLora-$16k$ |  63.5 |  85.69 |  41.88
3D-RPE-LLaMA2-7B-$16k$(our) |  89.50 |  90.00 |  40.00

Synthetic Tasks |  Passage Count |  PassageRetrival-En
---|---|---
LLaMA2-7B-chat-$4k$ |  2.1 |  9.8
LongChat-v1.5-7B-$32k$ |  1.0 |  30.5
InternLM-7B-$8k$ |  3.0 |  6.0
Vicuna-v1.5-7B-$16k$ |  6.5 |  4.5
LongLora-$16k$ |  3.61 |  29.75
3D-RPE-LLaMA2-7B-16k(our) |  4.0 |  14.5

Code Completion |  Lcc |  RepoBench-P
---|---|---
LLaMA2-7B-chat-$4k$ |  52.4 |  43.8
LongChat-v1.5-7B-$32k$ |  53.0 |  55.3
InternLM-7B-$8k$ |  44.1 |  28.8
Vicuna-v1.5-7B-$16k$ |  51.0 |  43.5
LongLora-$16k$ |  57.61 |  54.45
3D-RPE-LLaMA2-7B-$16k$(our) |  79.10 |  73.90


# Exploiting latent representation of sparse semantic layers for improved
short-term motion prediction with Capsule Networks

Albert Dulian and John C. Murray The authors are with the Department of
Computer Science and Technology, The University of Hull, Kingston Upon Hull,
United Kingdom<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

As urban environments manifest high levels of complexity it is of vital
importance that safety systems embedded within autonomous vehicles (AVs) are
able to accurately anticipate short-term future motion of nearby agents. This
problem can be further understood as generating a sequence of coordinates
describing the future motion of the tracked agent. Various proposed approaches
demonstrate significant benefits of using a rasterised top-down image of the
road, with a combination of Convolutional Neural Networks (CNNs), for
extraction of relevant features that define the road structure (eg. driveable
areas, lanes, walkways). In contrast, this paper explores use of Capsule
Networks (CapsNets) in the context of learning a hierarchical representation
of sparse semantic layers corresponding to small regions of the High-
Definition (HD) map. Each region of the map is dismantled into separate
geometrical layers that are extracted with respect to the agent’s current
position. By using an architecture based on CapsNets the model is able to
retain hierarchical relationships between detected features within images
whilst also preventing loss of spatial data often caused by the pooling
operation. We train and evaluate our model on publicly available dataset
nuTonomy scenes and compare it to recently published methods. We show that our
model achieves significant improvement over recently published works on
deterministic prediction, whilst drastically reducing the overall size of the
network.

##  I INTRODUCTION

Autonomous vehicles (AVs) are expected to provide a safe and robust
transportation solution in diverse and often highly uncertain surroundings.
One of the key aspects required to achieve this technology is providing AVs
with an ability to forecast short-term (e.g. 4 seconds) movements of other
vehicles within close proximity. This can allow for further reasoning with
regards to risk-free path planning and manoeuvre execution. Nonetheless,
short-term trajectory prediction is not a trivial task due to the inherently
difficult derivation of representative context cues with regards to agents and
the environment.

Recent techniques for motion prediction have focused on improving the accuracy
of predicted trajectories by incorporating rendered images of High-Definition
(HD) maps in a top-down manner where various road attributes are specifically
coded with different colours [4, 6, 21, 22]. The rasterised map is then used
as an input to a Convolutional Neural Network (CNN) [17] to encode global
spatial context of the scene and combined with an agent’s observed motion data
for further inference.

Figure 1: Disentanglement process of a map chunk at time $t-2$. The map is
turned into several geometrical layers each representing a single part of the
whole map. In addition, we draw an agent (bottom middle and right images) with
its origin corresponding to the position on the chunk of interest.

Yet, despite good results and definite advantages of incorporating map data
into the network, there are still numerous potential limitations that might
arise from these methods, effectively restricting the accuracy of the
predicted motion. First of all, the scene context is extracted globally i.e.
for each tracked agent an identical, single rasterized image is used to encode
spatial information. We argue that this approach leads to a great loss of
valuable data that can otherwise be associated with an individual agent, by
extracting small local regions of the map. These regions correspond to the
agent’s current position (Fig. 1), thus enabling the model to learn a more
representative latent space with respect to each tracked entity. We also
believe that incorporating local map regions characterised by several semantic
layers could potentially lead to the trajectory being predicted with respect
to physical boundaries of the road.

In addition, a typical approach in which previously mentioned methods encode
the image map data is by employing a variation of a CNN architecture e.g.
ResNet-50 [11] to extract more complex features with each consecutive
convolutional layer. Although over the past number of years CNNs have achieved
state-of-the-art performance in several computer vision tasks such as object
classification [32], detection [1], and semantic segmentation [31], these type
of networks still exhibit two important drawbacks [10]:

  1. 1.

Equivariance \- CNNs are not naturally equivariant to other transformations
but translation (parameter sharing). For instance, a slight rotation of an
object within the image will often cause relevant neurons to not fire and thus
fail when detecting salient features.

  2. 2.

Local Invariance \- CNNs achieve local transformation invariance e.g.
translation and rotation through pooling operation by taking maximum value
(max-pooling) from a small window of former convolution layer’s output.
Although pooling operation might not have significant drawbacks if the aim is
to solely detect whether the object is present or not, it does ignore
hierarchical relationship between features and leads to loss of large volumes
of spatial data.

To tackle these issues Hinton et al. [12] proposed a novel type of neural
network known as Capsule Network (CapsNet) that implements an idea of using
capsules (locally invariant group of neurons) to learn various properties
(e.g. pose) of the same object and encode them in an output vector whose
length corresponds to the probability of that object being present. Encoded
parameters can be further conceptualised as an object’s instantiation
parameters that enable the model to learn more robust and equivariant
representation of features with respect to change in viewpoint. Furthermore,
Sabour et al. [26] introduced the routing-by-agreement mechanism by which
capsules from lower levels decide which of their output vectors should be sent
to higher level capsules. In essence, output vectors from lower level capsules
are used to predict output of higher level capsules, predictions are then
compared with actual outputs to iteratively compute ”agreement” (cosine
similarity) between lower and higher capsules. For instance, the mere presence
of a nose or eyes (lower level capsules) should not be a sole indicator that
the face (higher level capsules) exists within an image, a hierarchical
relationship (e.g. rotation) between low and high level features should have a
high impact on the final prediction.

We are strongly motivated by the advantages offered by CapsNets and believe
that this type of network is ideal for building a well representative latent
space of map data, whose hierarchical relationship between features is of
great importance. We therefore propose an architecture where a standard CNN
backbone is replaced with Capsule layers. To the best of our knowledge this is
the first piece of work that proposes use of CapsNets in the context of short-
term motion prediction for autonomous vehicles. Through empirical evaluation
we demonstrate the benefits of using a CapsNet based network. To summarise,
our contributions are:

  * •

We demonstrate the first use of Capsule Networks with respect to the task of
predicting short-term future motion of vehicles in complex environments.

  * •

We propose the use of a local hierarchical semantic layers as opposed to a
global rasterized top-down view of the environment. This aids with creating a
more meaningful latent representation of the tracked agent.

  * •

We show that a network with CapsNet feature extractor can outperform a CNN
encoder based on e.g. ResNet architecture, whilst also significantly reducing
number of parameters within the network.

  * •

We examine and compare the performance of our proposed approach against state-
of-the-art methods on a publicly available dataset

##  II RELATED WORK

Forecasting future motion of nearby agents has been a topic of extensive study
in recent years as the domain of autonomous transportation gained exponential
interest both from numerous car manufacturers as well as various research
institutions. Some of more traditional approaches to short-term motion
prediction assume that the evolution of an object’s state through time remains
primarily governed by laws of physics and can be therefore modeled with e.g.
well known bicycle model [9]. In addition, noise reduction techniques
implementing filters, e.g. Kalman Filter [15], are often combined with the
motion model to improve prediction accuracy and account for noise in
measurements from sensors. Nevertheless, techniques such as the one presented
in [20] are greatly restricted to simple environments and to very short
prediction horizon (for instance 1 second into the future). Alternatively,
models based on Bayesian Framework such as Bayesian Networks (BNs) [14] and
Hidden Markov Models (HMMs) [23] are often employed to further account for the
probabilistic nature of the problem and capture relationships between sets of
random variables. For instance, [27] models a probability distribution over a
discrete set of maneuvers and then samples a number of possible future
trajectories with respect to detected maneuvers. For a more comprehensive
review of similar techniques we direct the reader to [19].

More recent and advanced techniques on the other hand focus on employing Deep
Learning (DL) [10] based methods in order to account for both temporal and
spatial complexity of the task. Lee et al. introduced DESIRE [18], a framework
for a multi-modal trajectory prediction based on CNNs, Conditional Variational
Auto-Encoders (CVAE) [28], and Recurrent Neural Networks (RNN) [25] encoder-
decoder to predict multiple plausible future trajectories based on the encoded
scene context, past motion histories, as well as interactions between multiple
agents. Moreover, [24] used a generative model to predict agents’ interactions
from image and LiDAR data to then forecast plausible motion from joint state
that encapsulates an encoding of all agents within the scene. Furthermore, the
use of raster top-down view of scene context was presented in [6] with a focus
on predicting multiple future paths as well as their associated probabilities
through use of CNN. Another body of work that explores motion forecasting
through probabilistic nature in multi-agent scenarios was demonstrated in
[30]. Authors modeled a joint behavior of numerous agents with a dynamic
attention-based state encoder that allowed to capture past and future
interactions for more precise multi-predictions. An interesting study was
conducted by Skrikanth et al. in [29] where instead of using a single top-down
view image of the scene, multiple intermediate representations of objects were
created by fusing images from either a stereo camera or LiDAR with depth
information. Generated images represented lanes, roads, obstacles etc, and
were used as an input to a CNN and LSTM [13] based network to obtain a future
location of the vehicle of interest on an occupancy grid map. Recently
proposed CoverNet [22] approaches the problem of motion forecasting by
formulating it as classification problem over a diverse set of trajectories.
Additionally, trajectories that were physically infeasible were excluded from
the train set, thus, during inference if the model predicted such motion it
would assign it a low probability, effectively restricting it from being taken
into consideration. In [21] Marchetti et al. presented MANTRA, a model based
on memory networks [33] that learns the association between past and future
motion and memorises most meaningful samples. Moreover, MANTRA presented
capabilities of updating its internal representation of motion samples in an
online fashion, thus allowing for continual improvement as new samples are
collected. On the other hand, the following work [8] investigated a novel
approach of transforming input features such as lanes, agents’ trajectories,
and crosswalks into a vectorized form and incorporating those into a model
based on Graph Neural Networks (GNN) [2] called VectorNet. The structure of
VectorNet was divided into local and global graphs in order to allow the model
to capture data from both individual polylines as well as from its aggregated
global representation. Experimental results indicate a significant performance
boost for predictions of up to 3 seconds into the future when compared with a
CNN model based on ResNet-18 architecture, whilst reducing number of learnable
parameters by 70%.

Figure 2: An overview of the proposed Capsule Net encoder. A single map chunk
at time $t$ is disentangled into separate geometrical layers. Each layer is
then passed separately to convolutional base, followed by lower capsules and
then to its corresponding higher capsule. Finally, outputs of all higher
capsules are concatenated and passed to final capsule to compute final
encoding.

##  III PROPOSED APPROACH

###  III-A Problem Formulation

First, we assume access to a sufficient detection and tracking module that
yields data corresponding to the current state of the tracked agent (e.g.
position, velocity) at a fixed interval e.g. 2Hz. We also assume access to a
HD map that define the following road layers; road segments, drivable areas,
lanes, and walkways.

Furthermore, we are interested in predicting a sequence of an agent’s future
positions for $\tau$ time-steps into the future which we denote as a matrix
$\mbox{\boldmath$\hat{Y}$}=[\mbox{\boldmath$\hat{y}$}_{t+1},\mbox{\boldmath$\hat{y}$}_{t+2},\cdots,\mbox{\boldmath$\hat{y}$}_{t+\tau}]$
where $\mbox{\boldmath$\hat{y}$}_{t+1}$ corresponds to a vector containing
future predictions at time $t+1$ such that
$\mbox{\boldmath$\hat{y}$}_{t+1}=(x_{t+1},y_{t+1})$. Moreover, let
$\mbox{\boldmath$S$}=[\mbox{\boldmath$s$}_{t-\rho},\cdots,\mbox{\boldmath$s$}_{t-1},\mbox{\boldmath$s$}_{t}]$
denote a standardised matrix of the tracked agent’s state features for past
$\rho$ time-steps where:

$\mbox{\boldmath$s$}_{t}=[\mbox{\boldmath$v$}_{t},\mbox{\boldmath$a$}_{t},\theta_{t}]$
(1)

is a vector of features at time-step $t$ containing velocity $(x,y)$,
acceleration $(x,y)$ and yaw angle respectively. In addition, we define a
normalised tensor
$\mbox{\boldmath$\mathsf{M}$}=[\mbox{\boldmath$\mathsf{L}$}_{t-\rho},\cdots,\mbox{\boldmath$\mathsf{L}$}_{t-1},\mbox{\boldmath$\mathsf{L}$}_{t}]$
which encapsulates local map chunks for past $\rho$ time-step. A single map
chunk $\mbox{\boldmath$\mathsf{L}$}_{t}$ at time-step $t$ is further
disentangled into separate layers (Fig. 1) such that
$\mbox{\boldmath$\mathsf{L}$}_{t}=[\mbox{\boldmath$L$}_{t,0},\mbox{\boldmath$\
L$}_{t,1},\cdots,\mbox{\boldmath$L$}_{t,n}]$ where $n$ is equal to number of
layers that describe disentangled map, and $\mbox{\boldmath$L$}_{t,i}$ is a
matrix representing a single sparse image of a layer of type $i$. We use the
same number of past time-steps $\rho$ for both $\mathsf{S}$ and $\mathsf{M}$
so that states as well as corresponding maps can be jointly encoded in a
temporal manner.

###  III-B Sparse Semantic Layers

In the literature, as previously outlined [4, 6, 21, 22], a common technique
to obtain the context of the environment is to rasterise a large portion of
its surrounding (e.g. $60m\times 60m$) into a top-down view HD map, allowing
the network to detect and encode salient features as well as enable the agent
to ’see’ what lies in front of it. In contrast, we focus on creating a more
meaningful encoding with respect to a single agent by exploiting its spatial
representation within an environment through a semantic, hierarchical view of
a local chunk of a given HD map. The following summarises the process of
creating disentangled map chunks as depicted in Fig. 1.

First, let
$\\{(x_{t},y_{t})\mid\mbox{\boldmath$p$}_{t}\in\mbox{\boldmath$P$}\\}$ denote
a set of tracked agent’s observed coordinates from time $t$ to $t-\rho$ where
each $\mbox{\boldmath$p$}_{t}\in\mbox{\boldmath$P$}$ is used to define the
origin of extraction for a map chunk $\mbox{\boldmath$\mathsf{L}$}_{t}$ at
time $t$. Moreover, as previously mentioned, we define a single map chunk
$\mbox{\boldmath$\mathsf{L}$}_{t}$ by $n$ sparse matrices (all pixels but the
ones corresponding to the given layer are set to zero)
$[\mbox{\boldmath$L$}_{t,0},\mbox{\boldmath$\
L$}_{t,1},\cdots,\mbox{\boldmath$L$}_{t,n}]$ where the
$\mbox{\boldmath$L$}_{t,i}$ corresponds to a semantic layer of type $i$. Thus,
we can further extract each local semantic layer with respect to
$\mbox{\boldmath$p$}_{t}$ as:

$\mbox{\boldmath$L$}_{t,i}=\Phi_{\mathrm{get\\_l}}(\mbox{\boldmath$p$}_{t},\lambda,\mathsf{s}),\quad\forall\
\mathsf{s}\in\mathsf{S}$ (2)

where $\lambda$ defines the extraction offset in meters. We set the
$\lambda=10$ and use a resolution of $3$ pixels per meter. The local chunk of
the layer of interest is then extracted from $x_{t}-3\lambda$ to
$x_{t}+3\lambda$ in a horizontal direction, and from $y_{t}-3\lambda$ to
$y_{t}+3\lambda$ in a vertical direction, hence, the final size of the
$\mbox{\boldmath$L$}_{t,i}$ is $20m\times 20m$. In addition, we create one
extra layer that portrays the agent by rendering it at its origin defined by
$\mbox{\boldmath$p$}_{t}$, with its initial orientation facing up. We then
compute:

$\hat{\theta}=((\pi/2)+\mathrm{sign}(-\theta)\cdot\lvert\theta\rvert)\cdot
180/\pi$ (3)

with respect to agent’s $\theta$ which defines the rotation angle by which the
agent is rotated so that its orientation is aligned with its global heading
direction. We notice no improvement when layers, as well as the agent, are
rotated so that their orientation aligns with agent’s heading facing up, and
therefore decide to leave the original orientation with respect to the global
coordinate system. Lastly, we transform extracted images (layers) from RGB to
grayscale, and upscale each image so that its final size is equal to $64\times
64$ pixels.

###  III-C Capsule Encoder’s Architecture

A typically employed approach that is used to encode the rasterised HD map is
to use some form of pre-trained CNN architecture e.g. ResNet-50 that 1) does
not provide equivariance with respect to extracted features and 2) uses
pooling operations, effectively discarding useful information as a trade off
for a small local translation invariance. With regards to the extraction of
salient map features it is crucial to retain valuable data and encode
hierarchical relationship between low and high level features. We propose to
remedy these issues through the use of a spatial encoder that is based on the
CapsNet as a feature extractor as opposed to the conventional CNN. We loosely
follow the implementation from [26], and construct our CapsNet encoder as a
four-part network as shown in Fig. 2.

Let $(k$, $s$, $o)$ denote the tuple that specifies a convolution layer’s
kernel size, stride size, and the number of output channels respectively (we
set the padding value within all convolution layers to 0). First, we define a
shallow convolutional base $\Phi_{\mathrm{base}}(\cdot)$ to extract local,
low-level features of semantic layers which is composed of a single
convolutional layer:

  * •

$(k=9,s=2,o=64)$

followed by the ELU non-linearity [5]. Next, we define the second part of the
encoder i.e. lower level (primary) capsules $\Phi_{\mathrm{lower}}(\cdot)$ to
learn parameters of more of a trivial parts of input data. Every capsule is a
$4D$ unit where each of its dimension corresponds to the scalar output value
of two consecutive convolutional operations:

  * •

$(k=9,s=2,o=32)$

  * •

$(k=2,s=2,o=16)$

hence the number of convolution layers within $\Phi_{\mathrm{lower}}(\cdot)$
layer is equal to $4\times 2$. We do not use any activation function in-
between layers but rather run capsules’ output vectors through the ’squashing’
non-linearity (Eq: 4) [26] in order to normalize the input vector $v$ so that
the magnitude of short and long vectors is squashed to almost 0 and just below
1 respectively:

$\Phi_{\mathrm{squash}}(\mbox{\boldmath$v$})=\frac{\|\mbox{\boldmath$v$}\|^{2}}{1+\|\mbox{\boldmath$v$}\|^{2}}\frac{\mbox{\boldmath$v$}}{\|\mbox{\boldmath$v$}\|}$
(4)

Since low level features of input images e.g. edges and lanes resemble strong
similarities among different types of semantic layers, it is reasonable to use
a single $\Phi_{\mathrm{lower}}(\cdot)$ layer to learn and extract their
parameters. Nonetheless, as their final representation differs significantly
we decided to implement the subsequent part of the encoder by defining $n$
higher capsule layers $\Phi_{\mathrm{higher}}(\cdot)$ per each type of the
input image. Hence, each type of semantic layer $\mbox{\boldmath$L$}_{t,i}$ is
encoded through its own respective higher capsule
$\Phi_{\mathrm{higher}}(\cdot)$ that outputs a $32D$ vector of its latent
representation. Lastly, we define the encoder’s final part as a single capsule
$\Phi_{\mathrm{final}}(\cdot)$ which outputs a $128D$ vector containing
jointed representation of all semantic layers.

###  III-D Overall Framework

We now describe the computational flow of data through the proposed network
that yields short-term future positions for the agent of interest. Although
the following represents computation process of a single agent, it is trivial
to extend the method for a multi-agent scenario. For the purpose of simplicity
and to focus on the main aspects of the paper we maintain a case of a single-
agent.

####  III-D1 Encoding of Semantic Layers

We transform each disentangled grayscale layer $\mbox{\boldmath$L$}_{t,i}$ at
time $t$ separately through the capsule encoder to compute its output
representation vector by first running the matrix through convolutional base:

$\mbox{\boldmath$\mathsf{Z}_{t,i}$}=\Phi_{\mathrm{base}}(\mbox{\boldmath$L$}_{t,i})$
(5)

where $\mathsf{Z}_{t,i}$ is the output tensor containing $64$ feature maps of
size $28\times 28$. Next, the $\mathsf{Z}_{t,i}$ is further passed into the
$\Phi_{\mathrm{lower}}(\cdot)$ layer to compute lower capsules:

$\mbox{\boldmath$Z_{t,i}$}=\Phi_{\mathrm{lower}}(\mbox{\boldmath$\mathsf{Z}_{t,i}$})$
(6)

where $Z_{t,i}$ is the squashed output matrix with $400\times 4D$ capsules
(second convolution layer outputs $16$ feature maps of size $5\times 5$).
Furthermore, we run $Z_{t,i}$ through its respective
$\Phi_{\mathrm{higher\\_i}}(\cdot)$ to get the $32D$ output vector $z_{t,i}$
whose scalar values corresponds to the layer’s $\mbox{\boldmath$L$}_{t,i}$
latent instantiation parameters:

$\mbox{\boldmath$z_{t,i}$}=\Phi_{\mathrm{higher\\_i}}(\mbox{\boldmath$Z_{t,i}$})$
(7)

We repeat this process for each semantic layer and concatenate each output to
create an input matrix for the final capsule:

$\mbox{\boldmath$z_{t}$}=\Phi_{\mathrm{final}}(\mathrm{concat}(\mbox{\boldmath$z_{t,0}$},\mbox{\boldmath$z_{t,1}$},\cdots,\mbox{\boldmath$z_{t,n}$}))$
(8)

that outputs a final $128D$ vector $z_{t}$ at time $t$.

####  III-D2 Encoding-Decoding of Future Motion

Moreover, each state vector $s_{t}$ is encoded through a fully-connected layer
with $128$ output units, followed by the ELU activation. We then concatenate
it with a corresponding $z_{t}$ to form the matrix of shape $[\rho,256]$ where
$\rho$ refers to the number of observed time-steps. The matrix is then encoded
in temporal manner through use of the LSTM layer. We set the size of the
layer’s hidden-state to $128$. Finally, we run the last output of the LSTM
through the decoder (fully-connected) layer with $2(x,y)\times\tau$ units
where $\tau$ corresponds to the future time-horizon, to get future
predictions. The output of the decoder can be reshaped to create the target
matrix
$\mbox{\boldmath$\hat{Y}$}=[\mbox{\boldmath$\hat{y}$}_{t+1},\mbox{\boldmath$\hat{y}$}_{t+2},\cdots,\mbox{\boldmath$\hat{y}$}_{t+\tau}]$
with each vector containing predicted future position $(x,y)$ of the tracked
agent, relative to its position at the last observed time-step $t$.

TABLE I: Ablation study with different backbone architectures. Each cell contains ADE/FDE error in meters for given future time horizon. In addition we examine whether additional map data, collected over observed time-steps $\rho$ provides performance increase. Number of each backbone parameters is reported in millions. Backbone | Future Time Horizon (seconds) |  #Params (Backbone)
---|---|---
1s |  2s |  3s |  4s |  5s |  6s |
ResNet-501 |  0.25/0.32 |  0.72/1.29 |  1.22/2.45 |  1.41/3.39 |  2.18/5.46 |  2.93/7.59 |  23.5m
ResNet-50ρ |  0.22/0.30 |  0.48/0.90 |  1.11/2.16 |  1.39/3.28 |  2.10/5.15 |  2.82/7.12 |  23.5m
Capsule-Encoder1 |  0.20/0.29 |  0.47/0.89 |  0.89/1.90 |  1.44/3.32 |  2.13/5.17 |  2.95/7.25 |  0.95m
Capsule-Encoderρ |  0.20/0.29 |  0.46/0.88 |  0.84/1.85 |  1.34/3.17 |  1.99/4.91 |  2.74/6.89 |  0.95m

##  IV EXPERIMENTS AND RESULTS

###  IV-A Dataset

We report results of our experiments on the publicly available self-driving
dataset nuTonomy Scenes (nuScenes) [3] designed for variety of tasks such as
detection, tracking as well as motion prediction. nuScenes provides access to
1000 scenes (approximately 20 seconds each) that were collected in Boston and
Singapore. Collected scenes provide a wide diversity with regards to weather
conditions, traffic situations and traffic density. In addition, nuScenes
contains human-annotated vectorized maps with 11 different semantic layers.
Scenes and objects within each scene (e.g. vehicles, pedestrians) were
accurately annotated at the rate of $2Hz$ and modeled as a cuboid, providing
further access to object’s position, size and yaw angle. We notice that some
parts of the HD map with respect to agents’ positions are not rendered
properly and we therefore remove these samples.

###  IV-B Training Setup and Evaluation Metrics

The proposed network is trained for $70$ epochs with Adam optimizer [16] and
set the initial learning rate to $5e-4$. We then reduce the learning rate at
epoch $5$ and $20$ by $\gamma=0.1$. We optimise the network by minimising:

$\mathcal{J}=\alpha\mathcal{J}_{1}+\beta\mathcal{J}_{2}$ (9)

where the $\mathcal{J}_{1}$ is the mean absolute error and the
$\mathcal{J}_{2}$ is the mean squared error. We notice a slight improvement
when combining both losses, and through empirical evaluation we find that
setting $\alpha=1.0$ and $\beta=1.0$ provides best results. Dataset is further
split in accordance to nuScene’s $train$ and $val$ split sets 111Available on
the official repository. Since the $test$ set has not been annotated we split
the $train$ set into train and validation sets and use the $val$ split as a
test set. For every sample, we observe $2$ seconds of agent’s past states and
make predictions for every $t\in\\{1,2,3,4,5,6\\}$ seconds of the future time
horizon. We report quantitative results by employing Average Displacement
Error (ADE) between all predictions, and Final Displacement Error (FDE)
between final prediction at time $t$. Prediction errors for both ADE and FDE
is measured in meters. In addition, we include the comparison with respect to
the number of parameters within adopted feature extractors (e.g. ResNet-50) as
well as networks as a whole.

###  IV-C Ablation Study - CNN vs Capsule based feature extractor

First, we analyse the performance of the network by employing two backbone
feature extractors i.e. 1) our proposed Capsule based encoder, and 2) a
commonly used ResNet-50 architecture. In addition, as previously discussed,
the encoder uses disentangled map chunks that are collected over observed past
$\rho$ times-steps, we therefore conduct a further study on the impact of
incorporating these map chunks over time against the use of a single map chunk
obtained at the initial time $t$. We notice that the initial, large output
size of ResNet-50 (i.e. $2048$ units) leads to either significant overfitting
or exploding gradient early in the training. To account for this, we
downsample the output through additional single fully-connected layer to
$128$. Results are presented in Table I with each cell containing the ADE/FDE
error for the corresponding time horizon. As demonstrated, the proposed
Capsule encoder whose input considers multiple map chunks over time (Capsule-
Encoderρ) outperforms other methods. Interestingly, for the first 3 seconds,
there is no significant improvement between Capsule-Encoderρ and Capsule-
Encoder1 (single map at time $t$), however, as the model is looking further
into the future, the performance difference becomes more apparent. Next, both
Capsule-Encoder and ResNet-50 demonstrate an improved performance when spatial
data for all observed time-step $\rho$ is included. Lastly, it is crucial to
emphasize that apart from significant performance improvement over ResNet-50
backbone, the Capsule encoder is substantially lighter, having less than 1
million parameters as opposed to ResNet-50 (over 23 million).

TABLE II: Comparison study between our final model vs two physics-based models and two recently proposed methods from the literature. Again, we provide a ADE/FDE error for each second of the future time horizon for up to 6 seconds. For CoverNetϵ (modes) we adjust number of modes with accordance to the $\epsilon$ value. We also included number of parameters (millions) of each model. Model | Future Time Horizon (seconds) |  #Params
---|---|---
1s |  2s |  3s |  4s |  5s |  6s |
Const. Vel. & Head. |  0.48/0.66 |  0.96/1.75 |  1.60/3.32 |  2.38/5.30 |  3.28/7.61 |  4.28/10.22 |  N/A
Physics Oracle |  0.42/0.55 |  0.77/1.35 |  1.26/2.55 |  1.89/4.18 |  2.64/6.15 |  3.50/8.44 |  N/A
CoverNet8 (64) [22] |  0.74/0.96 |  1.22/1.98 |  1.82/3.38 |  2.33/4.93 |  3.24/7.10 |  4.07/8.89 |  32.0m
CoverNet4 (415) [22] |  0.73/0.94 |  1.18/1.89 |  1.76/3.31 |  2.25/4.68 |  3.15/6.89 |  4.09/9.34 |  33.5m
CoverNet2 (2206) [22] |  0.65/0.86 |  1.13/1.90 |  1.73/3.32 |  2.23/4.63 |  3.24/7.23 |  4.23/9.74 |  40.9m
MTP [6] |  0.46/0.63 |  0.81/1.36 |  1.29/2.57 |  1.79/3.97 |  2.49/5.83 |  3.44/8.28 |  32.0m
Our Final Model |  0.20/0.29 |  0.46/0.88 |  0.84/1.85 |  1.34/3.17 |  1.99/4.91 |  2.74/6.89 |  1.2m

###  IV-D Comparison with state-of-the-art methods

Finally, we compare our results with two physics-based models as well as with
two state-of-the-art approaches that use CNN based on ResNet-50 as well as the
rasterised top-down view of the surrounding:

  * •

Constant Velocity & Heading Angle: Dynamic model explained in [22], based on
an assumption of the agent having constant velocity and heading angle across
the prediction time.

  * •

Physics Oracle: An extended version of physics-based models explained in [22].
This model computes future prediction for numerous physics-based models and
chooses the one with the minimum $L2$ distance with respect to the ground
truth data.

  * •

MTP: This is the implementation of the model presented in [6]. In the original
paper, authors focused on multi-modal prediction for fixed amount of modes.
Since our model focuses on deterministic prediction, we train the MTP to
predict a single mode. We adjust the learning settings with accordance to
details provided within the paper.

  * •

CoverNet: Another method proposed in [22] for multi-modal prediction. We
examine this model across three different settings of error tolerance
$\epsilon$ (2, 4, 8) that were publicly provided by authors (for detailed
information refer to the original publication). CoverNet was trained to
predict probability distribution over a set of diverse trajectories, we
therefore use the trajectories that were provided for each of the error
tolerance and pick top1 prediction as the final trajectory. Again, as with MTP
we also adopt training settings for CoverNet with accordance to details
outlined in the original paper.

Results of the comparison study are presented in Table II. It can be seen that
our approach based on CapsNet significantly outperforms all other methods
whilst drastically reducing number of parameters within the network. A
noticeable difference between our approach vs MTP and CoverNet might further
imply that using, for instance ResNet based backbone could not be considerably
beneficial due to e.g. large output size which, when re-training these
networks, we noticed leads to early overfitting. Furthermore, pre-trained
ResNet-50 consists of weights that were optimised for a completely different
domain [7], and therefore it is arguable whether such an approach leads to a
valuable and meaningful extraction of salient features from the global HD map.

##  V DISCUSSION

In this work we explored a novel approach to learning salient features from
local HD maps. In essence, our work investigated how a spatial encoder based
on the CapsNet architecture can be used to replace a generally employed CNN
that uses a pre-trained backbone model e.g ResNet-50 as a feature extractor.
In addition, we propose a use of local HD map chunks as opposed to global
maps, that can be further disentangled into separate, semantic geometrical
layers to encode more meaningful latent representation of a tracked agent.
Experiments on public dataset (nuScene) presents satisfactory results with
respect to prediction of a deterministic trajectory. Moreover, our model
outperforms recently proposed methods from the literature, whilst
substantially reducing size of the network. In the future, we would like to
extend this approach to multi-modal predictions in order to account for
inherently uncertain future that is associated with the discussed task.

## References

  * [1] Manoj Acharya, Tyler L Hayes, and Christopher Kanan.  Rodeo: Replay for online object detection.  arXiv preprint arXiv:2008.06439, 2020.
  * [2] Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al.  Relational inductive biases, deep learning, and graph networks.  arXiv preprint arXiv:1806.01261, 2018.
  * [3] Holger Caesar, Varun Bankiti, Alex H Lang, Sourabh Vora, Venice Erin Liong, Qiang Xu, Anush Krishnan, Yu Pan, Giancarlo Baldan, and Oscar Beijbom.  nuscenes: A multimodal dataset for autonomous driving.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11621–11631, 2020.
  * [4] Chiho Choi, Abhishek Patil, and Srikanth Malla.  Drogon: A causal reasoning framework for future trajectory forecast.  arXiv preprint arXiv:1908.00024, 2019.
  * [5] Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter.  Fast and accurate deep network learning by exponential linear units (elus).  arXiv preprint arXiv:1511.07289, 2015.
  * [6] Henggang Cui, Vladan Radosavljevic, Fang-Chieh Chou, Tsung-Han Lin, Thi Nguyen, Tzu-Kuo Huang, Jeff Schneider, and Nemanja Djuric.  Multimodal trajectory predictions for autonomous driving using deep convolutional networks.  In 2019 International Conference on Robotics and Automation (ICRA), pages 2090–2096. IEEE, 2019.
  * [7] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
  * [8] Jiyang Gao, Chen Sun, Hang Zhao, Yi Shen, Dragomir Anguelov, Congcong Li, and Cordelia Schmid.  Vectornet: Encoding hd maps and agent dynamics from vectorized representation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11525–11533, 2020.
  * [9] Thomas D Gillespie.  Fundamentals of vehicle dynamics, volume 400.  Society of automotive engineers Warrendale, PA, 1992.
  * [10] Ian Goodfellow, Yoshua Bengio, and Aaron Courville.  Deep learning.  MIT press, 2016.
  * [11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [12] Geoffrey E Hinton, Alex Krizhevsky, and Sida D Wang.  Transforming auto-encoders.  In International conference on artificial neural networks, pages 44–51. Springer, 2011.
  * [13] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  Neural computation, 9(8):1735–1780, 1997.
  * [14] Finn V Jensen et al.  An introduction to Bayesian networks, volume 210.  UCL press London, 1996.
  * [15] Rudolph Emil Kalman.  A new approach to linear filtering and prediction problems.  1960\.
  * [16] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [17] Yann LeCun, Bernhard E Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne E Hubbard, and Lawrence D Jackel.  Handwritten digit recognition with a back-propagation network.  In Advances in neural information processing systems, pages 396–404, 1990.
  * [18] Namhoon Lee, Wongun Choi, Paul Vernaza, Christopher B Choy, Philip HS Torr, and Manmohan Chandraker.  Desire: Distant future prediction in dynamic scenes with interacting agents.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 336–345, 2017.
  * [19] Stéphanie Lefèvre, Dizan Vasquez, and Christian Laugier.  A survey on motion prediction and risk assessment for intelligent vehicles.  ROBOMECH journal, 1(1):1–14, 2014.
  * [20] Chiu-Feng Lin, A Galip Ulsoy, and David J LeBlanc.  Vehicle dynamics and external disturbance estimation for vehicle path prediction.  IEEE Transactions on Control Systems Technology, 8(3):508–518, 2000\.
  * [21] Francesco Marchetti, Federico Becattini, Lorenzo Seidenari, and Alberto Del Bimbo.  Mantra: Memory augmented networks for multiple trajectory prediction.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7143–7152, 2020.
  * [22] Tung Phan-Minh, Elena Corina Grigore, Freddy A Boulton, Oscar Beijbom, and Eric M Wolff.  Covernet: Multimodal behavior prediction using trajectory sets.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14074–14083, 2020.
  * [23] Lawrence R Rabiner.  A tutorial on hidden markov models and selected applications in speech recognition.  Proceedings of the IEEE, 77(2):257–286, 1989.
  * [24] Nicholas Rhinehart, Rowan McAllister, Kris Kitani, and Sergey Levine.  Precog: Prediction conditioned on goals in visual multi-agent settings.  In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), October 2019.
  * [25] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams.  Learning representations by back-propagating errors.  nature, 323(6088):533–536, 1986.
  * [26] Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton.  Dynamic routing between capsules.  In Advances in neural information processing systems, pages 3856–3866, 2017.
  * [27] Matthias Schreier, Volker Willert, and Jürgen Adamy.  Bayesian, maneuver-based, long-term trajectory prediction and criticality assessment for driver assistance systems.  In 17th International IEEE Conference on Intelligent Transportation Systems (ITSC), pages 334–341. IEEE, 2014.
  * [28] Kihyuk Sohn, Honglak Lee, and Xinchen Yan.  Learning structured output representation using deep conditional generative models.  In Advances in neural information processing systems, pages 3483–3491, 2015.
  * [29] Shashank Srikanth, Junaid Ahmed Ansari, Sarthak Sharma, et al.  Infer: Intermediate representations for future prediction.  arXiv preprint arXiv:1903.10641, 2019.
  * [30] Charlie Tang and Russ R Salakhutdinov.  Multiple futures prediction.  In Advances in Neural Information Processing Systems, pages 15424–15434, 2019.
  * [31] Andrew Tao, Karan Sapra, and Bryan Catanzaro.  Hierarchical multi-scale attention for semantic segmentation.  arXiv preprint arXiv:2005.10821, 2020.
  * [32] Hugo Touvron, Andrea Vedaldi, Matthijs Douze, and Hervé Jégou.  Fixing the train-test resolution discrepancy: Fixefficientnet.  arXiv preprint arXiv:2003.08237, 2020.
  * [33] Jason Weston, Sumit Chopra, and Antoine Bordes.  Memory networks.  arXiv preprint arXiv:1410.3916, 2014.