#

RoPER is work by Georges Harik (@gharik), and this implementation is based on his original code.

Rotary Positional Embeddings with Relative distance (RoPER)

Rotary Positional Embeddings (RoPE) includes relative positions in attention score calculation. However, the embeddings themselves do not get any positional information , except what it can get implicitly from causal attention.

RoPER adds relative positional information explicitly to value embeddings. Specifically, it adds the relative positions of the tokens it paid attention to. We use same rotary positional embeddings to rotate the values in attention, Then, after taking the weighted sum, we rotate the final in the opposite direction. Which is equivalent to rotating each of the values (before attention) relative to the current position.

Here's the training code for training a transformer model with RoPER on an arithmetic addition where we can see significant improvement over RoPE.

Relative distances in embeddings

For any head, let $a_{n, i}$ be the attention from position $n$ to position $i$ , and $v_{i}$ be the value embeddings at position $i$ . Let's denote individual features as $v_{i}^{(1)}, v_{i}^{(2)}, \dots$ .

Normally, we would take the weight sum of value embeddings

$o_{n}^{(j)} = i \sum a_{n, i} v_{i}^{(j)}$

This doesn't explicitly add any distance information about the positions $i$ to final result $o_{n}^{(j)}$ .

RoPER pairs features like RoPE and transform them. For a pair $v_{m}^{(1)}$ and $v_{m}^{(2)}$ it transforms them by $R o PE (v_{m}^{(1)}, v_{m}^{(2)}, m)$ . Let us donate the transformed features with $\overset{v}{^}_{m}^{(1)}, \overset{v}{^}_{m}^{(2)}$ . Then it rotates the weighted sum $\overset{o}{^}_{n}^{(j)}$ in the the reverse direction with $R o PE (\overset{o}{^}_{n}^{(1)}, \overset{o}{^}_{n}^{(2)}, - n)$ . Note the $- n$ .

Note that,

R o PE (x_{m}^{(1)}, x_{m}^{(2)}, m) = (c o s m θ s i n m θ - s i n m θ c o s m θ) (x_{m}^{(1)} x_{m}^{(2)}) = (x_{m}^{(1)} c o s m θ - x_{m}^{(2)} s i n m θ x_{m}^{(2)} c o s m θ + x_{m}^{(1)} s i n m θ)

Final output after with the transformations is,

R o PE (\overset{o}{^}_{n}^{(1)}, \overset{o}{^}_{n}^{(2)}, - n) (\overset{o}{^}_{n}^{(1)} c o s n θ + \overset{o}{^}_{n}^{(2)} s i n n θ \overset{o}{^}_{n}^{(2)} c o s n θ - \overset{o}{^}_{n}^{(1)} s i n n θ) =

Note that $s i n (- n θ) = - s i n n θ$ .

Let's expand the first term $\overset{o}{^}_{n}^{(1)} c o s n θ + \overset{o}{^}_{n}^{(2)} s i n n θ$ ,

\overset{o}{^}_{n}^{(1)} c o s n θ + \overset{o}{^}_{n}^{(2)} s i n n θ i \sum a_{n, i} \overset{v}{^}_{i}^{(1)} c o s n θ + i \sum a_{n, i} \overset{v}{^}_{i}^{(2)} s i n n θ i \sum a_{n, i} (v_{i}^{(1)} c o s i θ - v_{i}^{(2)} s i n i θ) c o s n θ i \sum a_{n, i} (v_{i}^{(2)} c o s i θ + v_{i}^{(1)} s i n i θ) s i n m θ i \sum a_{n, i} v_{i}^{(1)} (c o s i θ c o s n θ + s i n i θ s i n n θ) i \sum a_{n, i} v_{i}^{(2)} (c o s i θ s i n n θ - s i n i θ c o s n θ) i \sum a_{n, i} v_{i}^{(1)} c o s (i - n) θ - i \sum a_{n, i} v_{i}^{(2)} s i n (i - n) θ i \sum a_{n, i} v_{i}^{(1)} c o s (i - n) θ - i \sum a_{n, i} v_{i}^{(2)} s i n (i - n) θ = = + = + = =

Simiarly we can show the second term is equal to,

$i \sum a_{n, i} v_{i}^{(1)} c o s (i - n) θ + i \sum a_{n, i} v_{i}^{(2)} s i n (i - n) θ$

Which gives,

R o PE (\overset{o}{^}_{n}^{(1)}, \overset{o}{^}_{n}^{(2)}, - n) (\sum_{i} a_{n, i} v_{i}^{(1)} c o s (i - n) θ - \sum_{i} a_{n, i} v_{i}^{(2)} s i n (i - n) θ \sum_{i} a_{n, i} v_{i}^{(1)} c o s (i - n) θ + \sum_{i} a_{n, i} v_{i}^{(2)} s i n (i - n) θ) i \sum a_{n, i} R o PE (v_{i}^{(1)}, v_{i}^{(1)}, (i - n) θ) = =

That is, the weighted average of values rotated relative to current position.

Here's an experiment that uses RoPER on an arthmetic addition task.

118from typing import Optional
119
120import torch
121
122from labml_nn.transformers.rope import RotaryPositionalEmbeddings, RotaryPEMultiHeadAttention

#

RoPE module that rotates in the opposite direction

This inherits from RoPE rotation implementation and changes the direction.

125class ReverseRotaryPositionalEmbeddings(RotaryPositionalEmbeddings):

#

x is the Tensor at the head of a key or a query with shape [seq_len, batch_size, n_heads, d]

132    def forward(self, x: torch.Tensor):

#

Cache $c o s$ and $s i n$ values

137        self._build_cache(x)

#

Split the features, we can choose to apply rotary embeddings only to a partial set of features.

140        x_rope, x_pass = x[..., :self.d], x[..., self.d:]

#

Calculate $[- x^{(\frac{d}{2} + 1)}, - x^{(\frac{d}{2} + 2)}, ..., - x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$

144        neg_half_x = self._neg_half(x_rope)

#

Calculate

(x_{m}^{(i)} c o s - m θ_{i} - x_{m}^{(i + \frac{d}{2})} s i n - m θ_{i} x_{m}^{(i + \frac{d}{2})} c o s - m θ_{i} + x_{m}^{(i)} s i n - m θ_{i}) = (x_{m}^{(i)} c o s m θ_{i} + x_{m}^{(i + \frac{d}{2})} s i n m θ_{i} x_{m}^{(i + \frac{d}{2})} c o s m θ_{i} - x_{m}^{(i)} s i n m θ_{i})

for $i \in 1, 2, ..., \frac{d}{2}$

160        x_rope = (x_rope * self.cos_cached[:x.shape[0]]) - (neg_half_x * self.sin_cached[:x.shape[0]])

#

163        return torch.cat((x_rope, x_pass), dim=-1)

#

Multi-head attention with rotary positional embeddings

We override multi-head attention from original transformer.

166class RotaryValuePEMultiHeadAttention(RotaryPEMultiHeadAttention):

#

173    def __init__(self, heads: int, d_model: int,
174                 rope_percentage: float = 0.5, rope_value_percentage: float = 0.5,
175                 dropout_prob: float = 0.0):
176        super().__init__(heads, d_model, rope_percentage, dropout_prob)

#

Rotary positional embedding layers

179        d_rope_value = int(self.d_k * rope_value_percentage)
180
181        self.value_rotary_pe = RotaryPositionalEmbeddings(d_rope_value)
182        self.value_reverse_rotary_pe = ReverseRotaryPositionalEmbeddings(d_rope_value)

#

query , key and value are the tensors that store collection of query, key and value vectors. They have shape [seq_len, batch_size, d_model] .

mask has shape [seq_len, seq_len, batch_size] and mask[i, j, b] indicates whether for batch b , query at position i has access to key-value at position j .

184    def forward(self, *,
185                query: torch.Tensor,
186                key: torch.Tensor,
187                value: torch.Tensor,
188                mask: Optional[torch.Tensor] = None):

#

query , key and value have shape [seq_len, batch_size, d_model]

200        seq_len, batch_size, _ = query.shape
201
202        if mask is not None:
203            mask = self.prepare_mask(mask, query.shape, key.shape)

#

Prepare query , key and value for attention computation. These will then have shape [seq_len, batch_size, heads, d_k] .

207        query = self.query(query)
208        key = self.key(key)
209        value = self.value(value)

#

Compute attention scores $Q K^{⊤}$ . This gives a tensor of shape [seq_len, seq_len, batch_size, heads] .

213        scores = self.get_scores(query, key)

#

Scale scores $\frac{Q K ^{⊤}}{d _{k}}$

216        scores *= self.scale

#

Apply mask

219        if mask is not None:
220            scores = scores.masked_fill(mask == 0, float('-inf'))

#

$so f t ma x$ attention along the key sequence dimension $se q so f t ma x (\frac{Q K ^{⊤}}{d _{k}})$

224        attn = self.softmax(scores)

#

Apply dropout

227        attn = self.dropout(attn)

#

Rotate value embeddings before taking the weighted sum so that they contain positional information

230        value = self.value_rotary_pe(value)

#

Multiply by values $se q so f t ma x (\frac{Q K ^{⊤}}{d _{k}}) V$

234        x = torch.einsum("ijbh,jbhd->ibhd", attn, value)

#

Rotate in the opposite direction so that each embedding hold the relative positions

237        x = self.value_reverse_rotary_pe(x)

#

Save attentions for any other calculations

240        self.attn = attn.detach()

#

Concatenate multiple heads

243        x = x.reshape(seq_len, batch_size, -1)

#

Output layer

246        return self.output(x)