diff --git a/.gitignore b/.gitignore index 6c0eec71..4336d66b 100644 --- a/.gitignore +++ b/.gitignore @@ -15,4 +15,5 @@ logs html/ diagrams/ .comet.config -settings.md \ No newline at end of file +settings.md +labml_app.log \ No newline at end of file diff --git a/.labml.yaml b/.labml.yaml index 1290b7bf..2e384051 100644 --- a/.labml.yaml +++ b/.labml.yaml @@ -19,3 +19,4 @@ indicators: name: optim.* options: comet: false +web_api: http://localhost:5005/api/v1/track? diff --git a/docs/experiments/arithmetic_dataset.html b/docs/experiments/arithmetic_dataset.html new file mode 100644 index 00000000..c35bbd9a --- /dev/null +++ b/docs/experiments/arithmetic_dataset.html @@ -0,0 +1,900 @@ + + + + + + + + + + + + + + + + + + + + + + + Arithmetic Dataset + + + + + + + + + + +
+
+
+
+

+ home + experiments +

+

+ + + Github + + Twitter +

+
+
+
+
+
+ # +
+

This is based on code by Georges Harik (@gharik).

+ +
+
+ 
11import random
+12import string
+13from typing import List
+14
+15import torch
+16from labml.logger import Text
+17from torch.utils.data import DataLoader, Dataset
+18
+19from labml import monit, logger, tracker
+20from labml.configs import option
+21from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs, transpose_batch
+
+
+
+
+
+ # +
+

Arithmetic Dataset

+

This creates arithmetic addition problems and solutions with workings. We've only implemented addition so far.

+

It's based on a character level tokenization.

+ +
+
+ 
24class ArithmeticDataset(Dataset):
+
+
+
+
+
+ # +
+
  • seq_len is the sequence length of generated math problems. We fill as many problems as possible upto this length :max_digits: is the maximum number of digits in the operand integers :n_sequences: is the number of sequences per epoch
+ +
+
+ 
34    def __init__(self, seq_len: int, max_digits: int, n_sequences: int):
+
+
+
+
+
+ # +
+ +
+
+ 
41        self.n_sequences = n_sequences
+42        self.max_digits = max_digits
+43        self.seq_len = seq_len
+
+
+
+
+
+ # +
+

Token id to string

+ +
+
+ 
45        self.itos = list(string.digits + 'xe =\n?+;')
+
+
+
+
+
+ # +
+

Character to token id

+ +
+
+ 
47        self.stoi = {c: i for i, c in enumerate(self.itos)}
+
+
+
+
+
+ # +
+

Generates an integer with n_digit + number of digits

+ +
+
+ 
49    @staticmethod
+50    def make_int(n_digits: int):
+
+
+
+
+
+ # +
+ +
+
+ 
54        res = 0
+55        for i in range(n_digits):
+56            d = random.randrange(1, 11) if i == 0 else random.randrange(0, 11)
+57            res = res * 10 + d
+58
+59        return res
+
+
+
+
+
+ # +
+

Generates the workings for x + y +. For example for 11+29 + it generates 1e0+9e0+0e0=10e0 1e0+2e0+1e0=4e0 +.

+ +
+
+ 
61    @staticmethod
+62    def get_add_explanation(x: int, y: int):
+
+
+
+
+
+ # +
+ +
+
+ 
69        carry = 0
+70        e = 0
+71        explanation = []
+72        while x > 0 or y > 0 or carry > 0:
+73            rx, ry = x % 10, y % 10
+74            total = rx + ry + carry
+75            explanation.append(f"{rx}e{e}+{ry}e{e}+{carry}e{e}=={total}e{e}")
+76            x, y, carry = x // 10, y // 10, total // 10
+77            e += 1
+78
+79        return ' '.join(explanation)
+
+
+
+
+
+ # +
+

Make a problem with a pre_explanation or not

+

Creates an arithmetic addition problem with workings and answer.

+ +
+
+ 
82    def make_add_problem(self):
+
+
+
+
+
+ # +
+ +
+
+ 
86        x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+87        y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+88
+89        explanation = self.get_add_explanation(x, y)
+90        return f"x={x}+{y}; {explanation} x=={x + y}\n"
+
+
+
+
+
+ # +
+

Get arithmetic problem and answer. This is used for evaluation.

+ +
+
+ 
92    def get_qa(self):
+
+
+
+
+
+ # +
+ +
+
+ 
96        x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+97        y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+98
+99        return f'x={x}+{y};', f'{x + y}'
+
+
+
+
+
+ # +
+

Generate multiple problems and pack them into a sequence.

+ +
+
+ 
101    def get_packed_math_input(self):
+
+
+
+
+
+ # +
+ +
+
+ 
105        s_enc = []
+106        while len(s_enc) <= self.seq_len:
+107            s_part = self.make_add_problem()
+108            s_part_enc = self.encode('?' + s_part)
+109            s_enc = s_enc + s_part_enc
+110        return s_enc
+
+
+
+
+
+ # +
+

Encode a given string

+ +
+
+ 
112    def encode(self, s: str):
+
+
+
+
+
+ # +
+ +
+
+ 
116        return [self.stoi[c] for c in s]
+
+
+
+
+
+ # +
+

Decode a list of token ids

+ +
+
+ 
118    def decode(self, arr: List[int]):
+
+
+
+
+
+ # +
+ +
+
+ 
122        return ''.join([self.itos[c] for c in arr])
+
+
+
+
+
+ # +
+

Get a input and target pair for auto-regressive modelling

+ +
+
+ 
124    def __getitem__(self, idx: int):
+
+
+
+
+
+ # +
+ +
+
+ 
128        s = torch.tensor(self.get_packed_math_input())
+129        return s[:self.seq_len], s[1:self.seq_len + 1]
+
+
+
+
+
+ # +
+

Number of sequences per epoch

+ +
+
+ 
131    def __len__(self):
+
+
+
+
+
+ # +
+ +
+
+ 
135        return self.n_sequences
+
+
+
+
+
+ # +
+

Arithmetic Task Experiment Configurations

+ +
+
+ 
138class ArithmeticAutoregression(NLPAutoRegressionConfigs):
+
+
+
+
+
+ # +
+

Maximum number of digits per operand integer

+ +
+
+ 
143    max_digits: int = 4
+
+
+
+
+
+ # +
+

Number of training sequences per epoch

+ +
+
+ 
145    train_sequences_per_epoch: int = 2 ** 12
+
+
+
+
+
+ # +
+

Training data loader

+ +
+
+ 
147    train_loader: DataLoader = 'arithmetic_train_loader'
+
+
+
+
+
+ # +
+

Number of problems in evaluation

+ +
+
+ 
149    n_tests: int = 64
+
+
+
+
+
+ # +
+

No need of a validation dataset

+ +
+
+ 
151    validator = None
+
+
+
+
+
+ # +
+

Number of times to run evaluations per epoch

+ +
+
+ 
153    inner_iterations = 4
+
+
+
+
+
+ # +
+

Number of tokens in the vocabulary

+ +
+
+ 
155    n_tokens = len(ArithmeticDataset(1, 1, 1).itos)
+
+
+
+
+
+ # +
+

Evaluation

+

We use the sampling function to evaluate the model on a set of problems

+ +
+
+ 
157    @torch.no_grad()
+158    def sample(self):
+
+
+
+
+
+ # +
+

Skip in the first epoch

+ +
+
+ 
166        if self.training_loop.idx < 1:
+167            return
+
+
+
+
+
+ # +
+

Create a dataset to generate problems

+ +
+
+ 
170        dataset = ArithmeticDataset(self.seq_len, self.max_digits, 1)
+
+
+
+
+
+ # +
+

Get a set of problems and answers

+ +
+
+ 
172        qa = [dataset.get_qa() for _ in range(self.n_tests)]
+
+
+
+
+
+ # +
+

Collect the problems only

+ +
+
+ 
174        questions = [p[0] for p in qa]
+
+
+
+
+
+ # +
+

Create a tensor with only the initial token

+ +
+
+ 
177        data = torch.tensor([[dataset.stoi[p[0]] for p in questions]])
+
+
+
+
+
+ # +
+

Move to device

+ +
+
+ 
179        data = data.to(self.device)
+
+
+
+
+
+ # +
+

Number of sequences that have completed

+ +
+
+ 
182        finished = torch.zeros((len(questions),)).bool().to(self.device)
+
+
+
+
+
+ # +
+

Token id of the new line character - this marks end of the answer

+ +
+
+ 
184        new_line = dataset.stoi['\n']
+
+
+
+
+
+ # +
+

Sampled results

+ +
+
+ 
187        results = [p[0] for p in questions]
+
+
+
+
+
+ # +
+

Sample upto sequence length

+ +
+
+ 
190        for i in monit.iterate('Sample', self.seq_len - 1):
+
+
+
+
+
+ # +
+

If all the sequences have completed we skip this

+ +
+
+ 
192            if finished.sum() == len(finished):
+193                continue
+
+
+
+
+
+ # +
+

Get the model output

+ +
+
+ 
196            output, *_ = self.model(data)
+
+
+
+
+
+ # +
+

Get the model prediction (greedy)

+ +
+
+ 
198            output = output[-1].argmax(dim=-1)
+
+
+
+
+
+ # +
+

Find which sequences have finished

+ +
+
+ 
201            finished = finished | (output == new_line)
+
+
+
+
+
+ # +
+

Skip if all have finished

+ +
+
+ 
203            if finished.sum() == len(finished):
+204                continue
+
+
+
+
+
+ # +
+

Override with the question

+ +
+
+ 
207            for j, p in enumerate(questions):
+208                if len(p) > i + 1:
+209                    output[j] = dataset.stoi[p[i + 1]]
+
+
+
+
+
+ # +
+

Add the next token to the input

+ +
+
+ 
212            data = torch.cat([data, output[None, :]], dim=0)
+
+
+
+
+
+ # +
+

Get the sampled results

+ +
+
+ 
215            for j, c in enumerate(output):
+216                results[j] += dataset.itos[c]
+
+
+
+
+
+ # +
+

Discard everything after the answer in the results

+ +
+
+ 
219        results = [r.split('\n')[0] for r in results]
+
+
+
+
+
+ # +
+

Log a sample

+ +
+
+ 
222        res_sample = results[0].split(';')
+223        logger.log([(res_sample[0], Text.key), (';', Text.subtle), (';'.join(res_sample[1:]), Text.none)])
+
+
+
+
+
+ # +
+

Get the answers

+ +
+
+ 
226        results = [r.split('x==')[-1] for r in results]
+
+
+
+
+
+ # +
+

Count the number of correct answers

+ +
+
+ 
229        correct = 0
+230        for r, _qa in zip(results, qa):
+231            if r == _qa[1]:
+232                correct += 1
+
+
+
+
+
+ # +
+

Log the score

+ +
+
+ 
235        tracker.save('score', correct / len(results))
+
+
+
+
+
+ # +
+

Training data loader

+ +
+
+ 
238@option(ArithmeticAutoregression.train_loader)
+239def arithmetic_train_loader(c: ArithmeticAutoregression):
+
+
+
+
+
+ # +
+ +
+
+ 
243    return DataLoader(ArithmeticDataset(c.seq_len, c.max_digits, c.train_sequences_per_epoch),
+244                      batch_size=c.batch_size,
+245                      collate_fn=transpose_batch,
+246                      num_workers=4)
+
+
+
+
+
+ # +
+

Code to test generated problems

+ +
+
+ 
249def _test():
+
+
+
+
+
+ # +
+ +
+
+ 
253    dataset = ArithmeticDataset(256, 8, 10)
+254
+255    print(dataset.decode(dataset.get_packed_math_input()))
+
+
+
+
+
+ # +
+

+ +
+
+ 
259if __name__ == '__main__':
+260    _test()
+
+
+
+ Trending Research Papers + labml.ai +
+
+ + + + \ No newline at end of file diff --git a/docs/experiments/nlp_autoregression.html b/docs/experiments/nlp_autoregression.html index e35de1a6..5b0ffb0b 100644 --- a/docs/experiments/nlp_autoregression.html +++ b/docs/experiments/nlp_autoregression.html @@ -83,7 +83,7 @@

Auto-regressive NLP model trainer

19from labml.logger import Text 20from labml_helpers.datasets.text import TextDataset, SequentialDataLoader, SequentialUnBatchedDataset, TextFileDataset 21from labml_helpers.device import DeviceConfigs -22from labml_helpers.metrics.accuracy import Accuracy +22from labml_helpers.metrics.accuracy import Accuracy, AccuracyMovingAvg 23from labml_helpers.module import Module 24from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex 25from labml_nn.optimizers.configs import OptimizerConfigs @@ -293,7 +293,7 @@

Trainer configurations

- 
80    accuracy = Accuracy()
+ 
80    accuracy = AccuracyMovingAvg()
diff --git a/docs/normalization/deep_norm/experiment.html b/docs/normalization/deep_norm/experiment.html index 9f7e0095..8ccf394a 100644 --- a/docs/normalization/deep_norm/experiment.html +++ b/docs/normalization/deep_norm/experiment.html @@ -70,7 +70,8 @@ #

DeepNorm Experiment

-

Open In Colab View Run

+

Open In Colab View Run Open In Comet

+ 
15import copy
diff --git a/docs/sitemap.xml b/docs/sitemap.xml
index 42cf8562..8f6077e4 100644
--- a/docs/sitemap.xml
+++ b/docs/sitemap.xml
@@ -204,7 +204,7 @@
 
     
       https://nn.labml.ai/normalization/deep_norm/index.html
-      2022-04-23T16:30:00+00:00
+      2022-05-18T16:30:00+00:00
       1.00
     
     
@@ -244,6 +244,13 @@
     
     
 
+    
+      https://nn.labml.ai/experiments/arithmetic_dataset.html
+      2022-06-02T16:30:00+00:00
+      1.00
+    
+    
+
     
       https://nn.labml.ai/experiments/index.html
       2020-12-26T16:30:00+00:00
@@ -603,14 +610,35 @@
 
     
       https://nn.labml.ai/transformers/rope/index.html
-      2022-04-05T16:30:00+00:00
+      2022-05-31T16:30:00+00:00
+      1.00
+    
+    
+
+    
+      https://nn.labml.ai/transformers/rope/value_pe/arithmetic_experiment.html
+      2022-06-02T16:30:00+00:00
+      1.00
+    
+    
+
+    
+      https://nn.labml.ai/transformers/rope/value_pe/index.html
+      2022-06-02T16:30:00+00:00
+      1.00
+    
+    
+
+    
+      https://nn.labml.ai/transformers/rope/value_pe/experiment.html
+      2022-05-31T16:30:00+00:00
       1.00
     
     
 
     
       https://nn.labml.ai/transformers/rope/experiment.html
-      2022-03-12T16:30:00+00:00
+      2022-05-31T16:30:00+00:00
       1.00
     
     
diff --git a/docs/transformers/rope/experiment.html b/docs/transformers/rope/experiment.html
index 156b2092..5ce9d242 100644
--- a/docs/transformers/rope/experiment.html
+++ b/docs/transformers/rope/experiment.html
@@ -92,7 +92,7 @@ 
Rotary PE attention

         

             
21def _rotary_pe_mha(c: TransformerConfigs):
 22    from labml_nn.transformers.rope import RotaryPEMultiHeadAttention
-23    return RotaryPEMultiHeadAttention(c.n_heads, c.d_model)

+23    return RotaryPEMultiHeadAttention(c.n_heads, c.d_model, 1.)
@@ -157,7 +157,7 @@

Rotary PE attention

- 
46    experiment.create(name="rotary_pe_transformer")
+ 
46    experiment.create(name="rotary_pe_transformer", writers={'screen'})
diff --git a/docs/transformers/rope/index.html b/docs/transformers/rope/index.html index 06512a47..6fbaa978 100644 --- a/docs/transformers/rope/index.html +++ b/docs/transformers/rope/index.html @@ -90,19 +90,19 @@

Rotary Positional Embeddings (RoPE)

#

RoPE module

-

Rotary encoding transforms pairs of features by rotating in the 2D plane. That is, it organizes the features as pairs. Each pair can be considered a coordinate in a 2D plane, and the encoding will rotate it by an angle depending on the position of the token.

+

Rotary encoding transforms pairs of features by rotating in the 2D plane. That is, it organizes the features as pairs. Each pair can be considered a coordinate in a 2D plane, and the encoding will rotate it by an angle depending on the position of the token.

For a pair of features

-

Let and be two features of the key or query of any head at position . Or for simplicity assume has only two features. Then the transformation is,

-

where is a constant angle. The other pairs of features are transformed similarly.

+

Let and be two features of the key or query of any head at position . Or for simplicity assume has only two features. Then the transformation is,

+

where is a constant angle. The other pairs of features are transformed similarly.

Attention is relative

-

For a pair of features, dot-product attention score between two positions and would be

-

This shows that for dot-production attention the rotary encodings gives relative attention.

+

For a pair of features, dot-product attention score between two positions and would be

+

This shows that for dot-production attention the rotary encodings gives relative attention.

For all features

The features are grouped into pairs and handled as above. They use a different for each pair.

-

The paper suggests using for the pairs of features.

-

We pair feature with feature . So for position we transform

-

to

- +

The paper suggests using for the pairs of features.

+

We pair feature with feature . So for position we transform

+

to

+ 
32class RotaryPositionalEmbeddings(nn.Module):
@@ -114,13 +114,13 @@

For all features

#
- 
118    def __init__(self, d: int, base: int = 10_000):
+ 
119    def __init__(self, d: int, base: int = 10_000):
@@ -131,33 +131,37 @@

For all features

- 
123        super().__init__()
+ 
124        super().__init__()
+125
+126        self.base = base
+127        self.d = d
+128        self.cos_cached = None
+129        self.sin_cached = None
-
+
#
-

+

Cache and values

- 
125        self.theta = nn.Parameter(1. / (base ** (torch.arange(0, d, 2).float() / d)), requires_grad=False)
+ 
131    def _build_cache(self, x: torch.Tensor):
-
+
#
-
  • x - is the Tensor at the head of a key or a query with shape [seq_len, batch_size, n_heads, d] -
+

Return if cache is already built

- 
127    def forward(self, x: torch.Tensor):
+ 
136        if self.cos_cached is not None and x.shape[0] <= self.cos_cached.shape[0]:
+137            return
@@ -165,11 +169,11 @@

For all features

#
-

Extract the shape

+

Get sequence length

- 
132        seq_len, batch_size, n_heads, d = x.shape
+ 
140        seq_len = x.shape[0]
@@ -177,11 +181,11 @@

For all features

#
-

+

- 
135        d_2 = d // 2
+ 
143        theta = 1. / (self.base ** (torch.arange(0, self.d, 2).float() / self.d)).to(x.device)
@@ -194,7 +198,7 @@

For all features

- 
138        seq_idx = torch.arange(seq_len, device=x.device).type_as(self.theta)
+ 
146        seq_idx = torch.arange(seq_len, device=x.device).float().to(x.device)
@@ -202,11 +206,11 @@

For all features

#
-

Calculate the product of position index and

+

Calculate the product of position index and

- 
141        idx_theta = torch.einsum('n,d->nd', seq_idx, self.theta)
+ 
149        idx_theta = torch.einsum('n,d->nd', seq_idx, theta)
@@ -214,11 +218,11 @@

For all features

#
-

Concatenate so that for row we have

+

Concatenate so that for row we have

- 
145        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
+ 
153        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
@@ -226,11 +230,12 @@

For all features

#
-

Calculate

+

Cache them

- 
148        neg_half_x = torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
+ 
156        self.cos_cached = idx_theta2.cos()[:, None, None, :]
+157        self.sin_cached = idx_theta2.sin()[:, None, None, :]
@@ -238,12 +243,10 @@

For all features

#
-

Calculate

-

for

- +
- 
160        rx = (x * idx_theta2.cos()[:, None, None, :]) + (neg_half_x * idx_theta2.sin()[:, None, None, :])
+ 
159    def _neg_half(self, x: torch.Tensor):
@@ -251,35 +254,37 @@

For all features

#
-

+

- 
163        return rx
+ 
161        d_2 = self.d // 2
-
+
#
-

Multi-head attention with rotary positional embeddings

-

We override multi-head attention from original transformer.

+

Calculate

- 
166class RotaryPEMultiHeadAttention(MultiHeadAttention):
+ 
164        return torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
-
+
#
- +
  • x + is the Tensor at the head of a key or a query with shape [seq_len, batch_size, n_heads, d] +
+
- 
173    def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
+ 
166    def forward(self, x: torch.Tensor):
@@ -287,12 +292,11 @@

Multi-head attention with rotary positional embeddings

#
-

The linear transformations do not need a bias since we explicitly include it when calculating scores. However having a bias for value - might make sense.

+

Cache and values

- 
177        super().__init__(heads, d_model, dropout_prob, bias=False)
+ 
171        self._build_cache(x)
@@ -300,24 +304,23 @@

Multi-head attention with rotary positional embeddings

#
-

Rotary positional embedding layers

+

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

- 
180        self.query_rotary_pe = RotaryPositionalEmbeddings(self.d_k)
-181        self.key_rotary_pe = RotaryPositionalEmbeddings(self.d_k)
+ 
174        x_rope, x_pass = x[..., :self.d], x[..., self.d:]
-
+
#
-

Calculate scores between queries and keys

+

Calculate

- 
183    def get_scores(self, query: torch.Tensor, key: torch.Tensor):
+ 
178        neg_half_x = self._neg_half(x_rope)
@@ -325,43 +328,119 @@

Calculate scores between queries and keys

#
-

Calculate dot-product with RoPE

+

Calculate

+

for

- 
189        return torch.einsum('ibhd,jbhd->ijbh', self.query_rotary_pe(query), self.key_rotary_pe(key))
+ 
190        x_rope = (x_rope * self.cos_cached[:x.shape[0]]) + (neg_half_x * self.sin_cached[:x.shape[0]])
-
+
#
-

Testing RoPE with a simple example

+

- 
192def _test_rotary():
+ 
193        return torch.cat((x_rope, x_pass), dim=-1)
-
+
#
+

Multi-head attention with rotary positional embeddings

+

We override multi-head attention from original transformer.

+ +
+
+ 
196class RotaryPEMultiHeadAttention(MultiHeadAttention):
+
+
+
+
+
+ # +
+ +
+
+ 
203    def __init__(self, heads: int, d_model: int, rope_percentage: float = 0.5, dropout_prob: float = 0.0):
+204        super().__init__(heads, d_model, dropout_prob)
+
+
+
+
+
+ # +
+

Rotary positional embedding layers

+ +
+
+ 
207        d_rope = int(self.d_k * rope_percentage)
+208        self.query_rotary_pe = RotaryPositionalEmbeddings(d_rope)
+209        self.key_rotary_pe = RotaryPositionalEmbeddings(d_rope)
+
+
+
+
+
+ # +
+

Calculate scores between queries and keys

+ +
+
+ 
211    def get_scores(self, query: torch.Tensor, key: torch.Tensor):
+
+
+
+
+
+ # +
+

Calculate dot-product with RoPE

+ +
+
+ 
217        return torch.einsum('ibhd,jbhd->ijbh', self.query_rotary_pe(query), self.key_rotary_pe(key))
+
+
+
+
+
+ # +
+

Testing RoPE with a simple example

+ +
+
+ 
220def _test_rotary():
+
+
+
+
+
+ # +
- 
196    x = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7], [7, 8, 9, 10]], dtype=torch.float)
-197    x = x[:, None, None, :]
-198    inspect(x)
-199
-200    rotary_pe = RotaryPositionalEmbeddings(3)
-201    inspect(rotary_pe(x))
-202
-203
-204if __name__ == '__main__':
-205    _test_rotary()
+ 
224    x = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7], [7, 8, 9, 10]], dtype=torch.float)
+225    x = x[:, None, None, :]
+226    inspect(x)
+227
+228    rotary_pe = RotaryPositionalEmbeddings(3)
+229    inspect(rotary_pe(x))
+230
+231
+232if __name__ == '__main__':
+233    _test_rotary()
diff --git a/docs/transformers/rope/value_pe/arithmetic_experiment.html b/docs/transformers/rope/value_pe/arithmetic_experiment.html new file mode 100644 index 00000000..bd52b4dc --- /dev/null +++ b/docs/transformers/rope/value_pe/arithmetic_experiment.html @@ -0,0 +1,403 @@ + + + + + + + + + + + + + + + + + + + + + + + Rotary Positional Embeddings with Relative distance (RoPER) Experiment + + + + + + + + + + +
+
+
+
+

+ home + transformers + rope + value_pe +

+

+ + + Github + + Twitter +

+
+
+
+
+
+ # +
+

Rotary Positional Embeddings with Relative distance (RoPER) Experiment

+ +
+
+ 
11from labml import experiment
+12from labml.configs import calculate
+13from labml_nn.experiments.arithmetic_dataset import ArithmeticAutoregression
+14from labml_nn.transformers import TransformerConfigs
+15from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs
+
+
+
+
+
+ # +
+

We inherit RoPE experiment and use it for arithmetic addition task.

+

We add the option to change attention to use Rotary Positional Embeddings with Relative distance (RoPER) below.

+ +
+
+ 
18class Configs(RoPEConfigs, ArithmeticAutoregression):
+
+
+
+
+
+ # +
+ +
+
+ 
26    pass
+
+
+
+
+
+ # +
+

Use Rotary Positional Embeddings with Relative distance (RoPER) in attention.

+ +
+
+ 
29def _rotary_value_pe_mha(c: TransformerConfigs):
+
+
+
+
+
+ # +
+ +
+
+ 
33    from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
+34    return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)
+
+
+
+
+
+ # +
+

Configuration options

+ +
+
+ 
38calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
+39calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
+40calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)
+
+
+
+
+
+ # +
+ +
+
+ 
43def main():
+
+
+
+
+
+ # +
+

Create experiment

+ +
+
+ 
45    experiment.create(name="roper_addition", comment="rotary value 7", writers={'screen', 'labml', 'comet'})
+
+
+
+
+
+ # +
+

Create configs

+ +
+
+ 
47    conf = Configs()
+
+
+
+
+
+ # +
+

Override configurations

+ +
+
+ 
49    experiment.configs(conf, {
+50        'max_digits': 7,
+
+
+
+
+
+ # +
+

No fixed positional embeddings

+ +
+
+ 
53        'transformer.src_embed': 'no_pos',
+54        'transformer.tgt_embed': 'no_pos',
+
+
+
+
+
+ # +
+

Encoder with RoPER attention

+ +
+
+ 
57        'transformer.encoder_attn': 'rotary_value',
+
+
+
+
+
+ # +
+

Encoder with RoPE attention 'transformer.encoder_attn': 'rotary',

+ +
+
+ 

+        

+    

+    

+        

+            

+                #
+            

+            
 

+
+        

+        

+            
62        'model': 'rotary_pe_transformer',

+        

+    

+    

+        

+            

+                #
+            

+            
Use a context size of  

+
+        

+        

+            
65        'seq_len': 512,

+        

+    

+    

+        

+            

+                #
+            

+            
Train for 32 epochs 

+
+        

+        

+            
67        'epochs': 20,

+        

+    

+    

+        

+            

+                #
+            

+            
Batch size  

+
+        

+        

+            
69        'batch_size': 16,

+        

+    

+    

+        

+            

+                #
+            

+            
Model size 

+
+        

+        

+            
72        'd_model': 128,
+73        'transformer.ffn.d_ff': 512,
+74        'transformer.n_heads': 4,
+75        'transformer.dropout': 0.0,

+        

+    

+    

+        

+            

+                #
+            

+            
Use Adam optimizer 

+
+        

+        

+            
78        'optimizer.optimizer': 'Adam',
+79        'optimizer.learning_rate': 2.5e-4,
+80    })

+        

+    

+    

+        

+            

+                #
+            

+            
Set models for saving and loading 

+
+        

+        

+            
83    experiment.add_pytorch_models({'model': conf.model})

+        

+    

+    

+        

+            

+                #
+            

+            
Start the experiment 

+
+        

+        

+            
86    with experiment.start():

+        

+    

+    

+        

+            

+                #
+            

+            
Run training 

+
+        

+        

+            
88        conf.run()

+        

+    

+    

+        

+            

+                #
+            

+            
 

+
+        

+        

+            
92if __name__ == '__main__':
+93    main()

+        

+    

+    

+        Trending Research Papers
+        labml.ai
+    

+

+
+
+
+
\ No newline at end of file
diff --git a/docs/transformers/rope/value_pe/experiment.html b/docs/transformers/rope/value_pe/experiment.html
new file mode 100644
index 00000000..109a7605
--- /dev/null
+++ b/docs/transformers/rope/value_pe/experiment.html
@@ -0,0 +1,454 @@
+
+
+
+    
+    
+    
+
+    
+    
+    
+    
+    
+    
+
+    
+    
+    
+    
+    
+    
+    
+
+    Rotary Positional Embeddings (RoPE) Experiment
+    
+    
+    
+    
+
+    
+    
+    
+
+
+

+    

+    

+        

+            

+                home
+                transformers
+                rope
+                value_pe
+            

+            

+
+                
+                    Github
+                
+                    Twitter
+            

+        

+    

+    

+        

+            

+                #
+            

+            
Rotary Positional Embeddings (RoPE) Experiment

+
This is an annotated PyTorch experiment to train a transformer model with Rotary Positional Embeddings (RoPE).

+
View Run

+
+        

+        

+            
14from labml import experiment
+15from labml.configs import calculate
+16from labml_nn.transformers import TransformerConfigs
+17from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs

+        

+    

+    

+        

+            

+                #
+            

+            
Rotary PE attention

+
+        

+        

+            

+        

+    

+    

+        

+            

+                #
+            

+            
+        

+        

+            
22class Configs(RoPEConfigs):  # , ArithmeticAutoregression):
+23    pass

+        

+    

+    

+        

+            

+                #
+            

+            
+        

+        

+            
26def _rotary_value_pe_mha(c: TransformerConfigs):
+27    from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
+28    return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)

+        

+    

+    

+        

+            

+                #
+            

+            
Configuration options 

+
+        

+        

+            
32calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
+33calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
+34calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)

+        

+    

+    

+        

+            

+                #
+            

+            
+        

+        

+            
37def main():

+        

+    

+    

+        

+            

+                #
+            

+            
Create experiment 

+
+        

+        

+            
39    experiment.create(name="rotary_shakespeare", comment="rotary value", writers={'screen', 'labml'})

+        

+    

+    

+        

+            

+                #
+            

+            
Create configs 

+
+        

+        

+            
41    conf = Configs()

+        

+    

+    

+        

+            

+                #
+            

+            
Override configurations 

+
+        

+        

+            
43    experiment.configs(conf, {

+        

+    

+    

+        

+            

+                #
+            

+            
No fixed positional embeddings 

+
+        

+        

+            
45        'transformer.src_embed': 'no_pos',
+46        'transformer.tgt_embed': 'no_pos',

+        

+    

+    

+        

+            

+                #
+            

+            
Encoder with RoPE 

+
+        

+        

+            
49        'transformer.encoder_attn': 'rotary_value',

+        

+    

+    

+        

+            

+                #
+            

+            
'transformer.encoder_attn': 'rotary', 

+
+        

+        

+            

+        

+    

+    

+        

+            

+                #
+            

+            
 

+
+        

+        

+            
53        'model': 'rotary_pe_transformer',

+        

+    

+    

+        

+            

+                #
+            

+            
Use character level tokenizer 

+
+        

+        

+            
56        'tokenizer': 'character',

+        

+    

+    

+        

+            

+                #
+            

+            
Prompt separator is blank 

+
+        

+        

+            
58        'prompt_separator': '',

+        

+    

+    

+        

+            

+                #
+            

+            
Starting prompt for sampling 

+
+        

+        

+            
60        'prompt': 'It is ',

+        

+    

+    

+        

+            

+                #
+            

+            
Use Tiny Shakespeare dataset 

+
+        

+        

+            
62        'text': 'tiny_shakespeare',

+        

+    

+    

+        

+            

+                #
+            

+            
Use a context size of  

+
+        

+        

+            
65        'seq_len': 512,

+        

+    

+    

+        

+            

+                #
+            

+            
Train for 32 epochs 

+
+        

+        

+            
67        'epochs': 24,

+        

+    

+    

+        

+            

+                #
+            

+            
Batch size  

+
+        

+        

+            
69        'batch_size': 16,

+        

+    

+    

+        

+            

+                #
+            

+            
Switch between training and validation for  times per epoch 

+
+        

+        

+            
72        'inner_iterations': 4,

+        

+    

+    

+        

+            

+                #
+            

+            
Model size 

+
+        

+        

+            
75        'd_model': 128,
+76        'transformer.ffn.d_ff': 512,
+77        'transformer.n_heads': 4,
+78        'transformer.dropout': 0.0,

+        

+    

+    

+        

+            

+                #
+            

+            
Use Adam optimizer 

+
+        

+        

+            
81        'optimizer.optimizer': 'Adam',
+82        'optimizer.learning_rate': 2.5e-4,
+83
+84        'dataloader_shuffle_with_replacement': True
+85    })

+        

+    

+    

+        

+            

+                #
+            

+            
Set models for saving and loading 

+
+        

+        

+            
88    experiment.add_pytorch_models({'model': conf.model})

+        

+    

+    

+        

+            

+                #
+            

+            
Start the experiment 

+
+        

+        

+            
91    with experiment.start():

+        

+    

+    

+        

+            

+                #
+            

+            
Run training 

+
+        

+        

+            
93        conf.run()

+        

+    

+    

+        

+            

+                #
+            

+            
 

+
+        

+        

+            
97if __name__ == '__main__':
+98    main()

+        

+    

+    

+        Trending Research Papers
+        labml.ai
+    

+

+
+
+
+
\ No newline at end of file
diff --git a/docs/transformers/rope/value_pe/index.html b/docs/transformers/rope/value_pe/index.html
new file mode 100644
index 00000000..fa4f0e44
--- /dev/null
+++ b/docs/transformers/rope/value_pe/index.html
@@ -0,0 +1,510 @@
+
+
+
+    
+    
+    
+
+    
+    
+    
+    
+    
+    
+
+    
+    
+    
+    
+    
+    
+    
+
+    Rotary Positional Embeddings with Relative distance (RoPER)
+    
+    
+    
+    
+
+    
+    
+    
+
+
+

+    

+    

+        

+            

+                home
+                transformers
+                rope
+                value_pe
+            

+            

+
+                
+                    Github
+                
+                    Twitter
+            

+        

+    

+    

+        

+            

+                #
+            

+            
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  be the attention from position  to position , and  be the value embeddings at position . Let's denote individual features as .

+
Normally, we would take the weight sum of value embeddings

+

+
This doesn't explicitly add any distance information about the positions  to final result .

+
RoPER pairs features like RoPE and transform them. For a pair  and  it transforms them by  . Let us donate the transformed features with . Then it rotates the weighted sum  in the the reverse direction with  . Note the .

+
Note that,

+
Final output after with the transformations is,

+
Note that .

+
Let's expand the first term ,

+
Simiarly we can show the second term is equal to,

+

+
Which gives,

+
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  and  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  

+
+        

+        

+            
144        neg_half_x = self._neg_half(x_rope)

+        

+    

+    

+        

+            

+                #
+            

+            
Calculate

+
for  

+
+        

+        

+            
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 . This gives a tensor of shape [seq_len, seq_len, batch_size, heads]
+. 

+
+        

+        

+            
213        scores = self.get_scores(query, key)

+        

+    

+    

+        

+            

+                #
+            

+            
Scale scores  

+
+        

+        

+            
216        scores *= self.scale

+        

+    

+    

+        

+            

+                #
+            

+            
Apply mask 

+
+        

+        

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

+        

+    

+    

+        

+            

+                #
+            

+            
 attention along the key sequence dimension  

+
+        

+        

+            
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  

+
+        

+        

+            
234        x = torch.einsum("ijbh,jbhd->ibhd", attn, self.value_rotary_pe(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)

+        

+    

+    

+        Trending Research Papers
+        labml.ai
+    

+

+
+
+
+
\ No newline at end of file
diff --git a/labml_nn/experiments/algo_tasks/__init__.py b/labml_nn/experiments/algo_tasks/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/labml_nn/experiments/algo_tasks/arithmetic_addition.py b/labml_nn/experiments/algo_tasks/arithmetic_addition.py
new file mode 100644
index 00000000..8239be7a
--- /dev/null
+++ b/labml_nn/experiments/algo_tasks/arithmetic_addition.py
@@ -0,0 +1,260 @@
+"""
+---
+title: Arithmetic Dataset
+summary: >
+  This creates arithmetic problems.
+---
+
+*This is based on code by [Georges Harik (@gharik)](https://twitter.com/gharik).*
+"""
+
+import random
+import string
+from typing import List
+
+import torch
+from labml.logger import Text
+from torch.utils.data import DataLoader, Dataset
+
+from labml import monit, logger, tracker
+from labml.configs import option
+from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs, transpose_batch
+
+
+class ArithmeticAdditionDataset(Dataset):
+    """
+    ## Arithmetic Dataset
+
+    This creates arithmetic addition problems and solutions with workings.
+    We've only implemented addition so far.
+
+    It's based on a character level tokenization.
+    """
+
+    def __init__(self, seq_len: int, max_digits: int, n_sequences: int):
+        """
+        :param seq_len: is the sequence length of generated math problems.
+            We fill as many problems as possible upto this length
+        :max_digits: is the maximum number of digits in the operand integers
+        :n_sequences: is the number of sequences per epoch
+        """
+        self.n_sequences = n_sequences
+        self.max_digits = max_digits
+        self.seq_len = seq_len
+        # Token id to string
+        self.itos = list(string.digits + 'xe =\n?+;')
+        # Character to token id
+        self.stoi = {c: i for i, c in enumerate(self.itos)}
+
+    @staticmethod
+    def make_int(n_digits: int):
+        """
+        Generates an integer with `n_digit` number of digits
+        """
+        res = 0
+        for i in range(n_digits):
+            d = random.randrange(1, 11) if i == 0 else random.randrange(0, 11)
+            res = res * 10 + d
+
+        return res
+
+    @staticmethod
+    def get_add_explanation(x: int, y: int):
+        """
+        Generates the workings for `x + y`.
+        For example for `11+29` it generates
+        `1e0+9e0+0e0=10e0 1e0+2e0+1e0=4e0`.
+        """
+
+        carry = 0
+        e = 0
+        explanation = []
+        while x > 0 or y > 0 or carry > 0:
+            rx, ry = x % 10, y % 10
+            total = rx + ry + carry
+            explanation.append(f"{rx}e{e}+{ry}e{e}+{carry}e{e}=={total}e{e}")
+            x, y, carry = x // 10, y // 10, total // 10
+            e += 1
+
+        return ' '.join(explanation)
+
+    # Make a problem with a pre_explanation or not
+    def make_add_problem(self):
+        """
+        Creates an arithmetic addition problem with workings and answer.
+        """
+        x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+        y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+
+        explanation = self.get_add_explanation(x, y)
+        return f"x={x}+{y}; {explanation} x=={x + y}\n"
+
+    def get_qa(self):
+        """
+        Get arithmetic problem and answer. This is used for evaluation.
+        """
+        x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+        y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+
+        return f'?x={x}+{y};', f'{x + y}'
+
+    def get_packed_math_input(self):
+        """
+        Generate multiple problems and pack them into a sequence.
+        """
+        s_enc = []
+        while len(s_enc) <= self.seq_len:
+            s_part = self.make_add_problem()
+            s_part_enc = self.encode('?' + s_part)
+            s_enc = s_enc + s_part_enc
+        return s_enc
+
+    def encode(self, s: str):
+        """
+        Encode a given string
+        """
+        return [self.stoi[c] for c in s]
+
+    def decode(self, arr: List[int]):
+        """
+        Decode a list of token ids
+        """
+        return ''.join([self.itos[c] for c in arr])
+
+    def __getitem__(self, idx: int):
+        """
+        Get a input and target pair for auto-regressive modelling
+        """
+        s = torch.tensor(self.get_packed_math_input())
+        return s[:self.seq_len], s[1:self.seq_len + 1]
+
+    def __len__(self):
+        """
+        Number of sequences per epoch
+        """
+        return self.n_sequences
+
+
+class ArithmeticAdditionAutoregression(NLPAutoRegressionConfigs):
+    """
+    ## Arithmetic Task Experiment Configurations
+    """
+    # Maximum number of digits per operand integer
+    max_digits: int = 4
+    # Number of training sequences per epoch
+    train_sequences_per_epoch: int = 2 ** 12
+    # Training data loader
+    train_loader: DataLoader = 'arithmetic_train_loader'
+    # Number of problems in evaluation
+    n_tests: int = 64
+    # No need of a validation dataset
+    validator = None
+    # Number of times to run evaluations per epoch
+    inner_iterations = 4
+    # Number of tokens in the vocabulary
+    n_tokens = len(ArithmeticAdditionDataset(1, 1, 1).itos)
+
+    @torch.no_grad()
+    def sample(self):
+        """
+        ### Evaluation
+
+        We use the sampling function to evaluate the model on a set of problems
+        """
+
+        # Skip in the first epoch
+        if self.training_loop.idx < 1:
+            return
+
+        # Create a dataset to generate problems
+        dataset = ArithmeticAdditionDataset(self.seq_len, self.max_digits, 1)
+        # Get a set of problems and answers
+        qa = [dataset.get_qa() for _ in range(self.n_tests)]
+        # Collect the problems only
+        questions = [p[0] for p in qa]
+
+        # Create a tensor with only the initial token
+        data = torch.tensor([[dataset.stoi[p[0]] for p in questions]])
+        # Move to device
+        data = data.to(self.device)
+
+        # Number of sequences that have completed
+        finished = torch.zeros((len(questions),)).bool().to(self.device)
+        # Token id of the new line character - this marks end of the answer
+        new_line = dataset.stoi['\n']
+
+        # Sampled results
+        results = [p[0] for p in questions]
+
+        # Sample upto sequence length
+        for i in monit.iterate('Sample', self.seq_len - 1):
+            # If all the sequences have completed we skip this
+            if finished.sum() == len(finished):
+                continue
+
+            # Get the model output
+            output, *_ = self.model(data)
+            # Get the model prediction (greedy)
+            output = output[-1].argmax(dim=-1)
+
+            # Override with the question
+            for j, p in enumerate(questions):
+                if len(p) > i + 1:
+                    output[j] = dataset.stoi[p[i + 1]]
+
+            # Find which sequences have finished
+            finished = finished | (output == new_line)
+            # Skip if all have finished
+            if finished.sum() == len(finished):
+                continue
+
+            # Add the next token to the input
+            data = torch.cat([data, output[None, :]], dim=0)
+
+            # Get the sampled results
+            for j, c in enumerate(output):
+                results[j] += dataset.itos[c]
+
+        # Discard everything after the answer in the results
+        results = [r.split('\n')[0] for r in results]
+
+        # Log a sample
+        res_sample = results[0].split(';')
+        logger.log([(res_sample[0], Text.key), (';', Text.subtle), (';'.join(res_sample[1:]), Text.none)])
+
+        # Get the answers
+        results = [r.split('x==')[-1] for r in results]
+
+        # Count the number of correct answers
+        correct = 0
+        for r, _qa in zip(results, qa):
+            if r == _qa[1]:
+                correct += 1
+
+        # Log the score
+        tracker.save('score', correct / len(results))
+
+
+@option(ArithmeticAdditionAutoregression.train_loader)
+def arithmetic_train_loader(c: ArithmeticAdditionAutoregression):
+    """
+    Training data loader
+    """
+    return DataLoader(ArithmeticAdditionDataset(c.seq_len, c.max_digits, c.train_sequences_per_epoch),
+                      batch_size=c.batch_size,
+                      collate_fn=transpose_batch,
+                      num_workers=4)
+
+
+def _test():
+    """
+    Code to test generated problems
+    """
+    dataset = ArithmeticAdditionDataset(256, 8, 10)
+
+    print(dataset.decode(dataset.get_packed_math_input()))
+
+
+#
+if __name__ == '__main__':
+    _test()
diff --git a/labml_nn/experiments/algo_tasks/arithmetic_multiplication.py b/labml_nn/experiments/algo_tasks/arithmetic_multiplication.py
new file mode 100644
index 00000000..e8d9ade1
--- /dev/null
+++ b/labml_nn/experiments/algo_tasks/arithmetic_multiplication.py
@@ -0,0 +1,273 @@
+"""
+---
+title: Arithmetic Dataset
+summary: >
+  This creates arithmetic problems.
+---
+
+*This is based on code by [Georges Harik (@gharik)](https://twitter.com/gharik).*
+"""
+
+import random
+import string
+from typing import List
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from labml import monit, logger, tracker
+from labml.configs import option
+from labml.logger import Text
+from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs, transpose_batch
+
+
+class ArithmeticMultiplicationDataset(Dataset):
+    """
+    ## Arithmetic Dataset
+
+    This creates arithmetic addition problems and solutions with workings.
+    We've only implemented addition so far.
+
+    It's based on a character level tokenization.
+    """
+
+    def __init__(self, seq_len: int, max_digits: int, base: int, n_sequences: int):
+        """
+        :param seq_len: is the sequence length of generated math problems.
+            We fill as many problems as possible upto this length
+        :max_digits: is the maximum number of digits in the operand integers
+        :n_sequences: is the number of sequences per epoch
+        """
+        self.base = base
+        self.n_sequences = n_sequences
+        self.max_digits = max_digits
+        self.seq_len = seq_len
+        # Token id to string
+        self.itos = list(string.digits + 'x =\n?*;')
+        # Character to token id
+        self.stoi = {c: i for i, c in enumerate(self.itos)}
+
+    def make_int(self, n_digits: int):
+        """
+        Generates an integer with `n_digit` number of digits
+        """
+        res = 0
+        for i in range(n_digits):
+            d = random.randrange(1, self.base + 1) if i == 0 else random.randrange(0, self.base + 1)
+            res = res * self.base + d
+
+        return res
+
+    def get_add_explanation(self, x: int, y: int):
+        """
+        Generates the workings for `x + y`.
+        For example for `11+29` it generates
+        `1e0+9e0+0e0=10e0 1e0+2e0+1e0=4e0`.
+        """
+
+        explanation = []
+        while x > 0:
+            rx = x % self.base
+            explanation.append(f"{self.to_string(y * rx)}")
+            x = x // self.base
+
+        return ' '.join(explanation)
+
+    # Make a problem with a pre_explanation or not
+    def make_add_problem(self):
+        """
+        Creates an arithmetic addition problem with workings and answer.
+        """
+        x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+        y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+
+        explanation = self.get_add_explanation(x, y)
+        return f"x={self.to_string(x)}*{self.to_string(y)}; {explanation} x=={self.to_string(x * y)}\n"
+
+    def to_string(self, x: int):
+        if x == 0:
+            return '0'
+        a = []
+        while x > 0:
+            a += [f'{x % self.base}']
+            x = x // self.base
+
+        return ''.join(reversed(a))
+
+    def get_qa(self):
+        """
+        Get arithmetic problem and answer. This is used for evaluation.
+        """
+        x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+        y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
+
+        return f'?x={self.to_string(x)}*{self.to_string(y)};', f'{self.to_string(x * y)}'
+
+    def get_packed_math_input(self):
+        """
+        Generate multiple problems and pack them into a sequence.
+        """
+        s_enc = []
+        mask = []
+        while len(s_enc) <= self.seq_len:
+            s_part = self.make_add_problem()
+            s_part_enc = self.encode('?' + s_part)
+            prob, sol = s_part.split(';')
+            mask += [False] * (len(prob) + 2)
+            mask += [True] * len(sol)
+            s_enc = s_enc + s_part_enc
+        return s_enc, mask
+
+    def encode(self, s: str):
+        """
+        Encode a given string
+        """
+        return [self.stoi[c] for c in s]
+
+    def decode(self, arr: List[int]):
+        """
+        Decode a list of token ids
+        """
+        return ''.join([self.itos[c] for c in arr])
+
+    def __getitem__(self, idx: int):
+        """
+        Get a input and target pair for auto-regressive modelling
+        """
+        s, mask = self.get_packed_math_input()
+        s = torch.tensor(s)
+        mask = torch.tensor(mask)
+        target = s * mask + -1 * (~mask)
+        return s[:self.seq_len], target[1:self.seq_len + 1]
+
+    def __len__(self):
+        """
+        Number of sequences per epoch
+        """
+        return self.n_sequences
+
+
+class ArithmeticMultiplicationAutoregression(NLPAutoRegressionConfigs):
+    """
+    ## Arithmetic Task Experiment Configurations
+    """
+    # Maximum number of digits per operand integer
+    max_digits: int = 4
+    # Number of training sequences per epoch
+    train_sequences_per_epoch: int = 2 ** 12
+    # Training data loader
+    train_loader: DataLoader = 'arithmetic_train_loader'
+    # Number of problems in evaluation
+    n_tests: int = 64
+    # No need of a validation dataset
+    validator = None
+    # Number of times to run evaluations per epoch
+    inner_iterations = 4
+    # Number of tokens in the vocabulary
+    base: int = 10
+    n_tokens = len(ArithmeticMultiplicationDataset(1, 1, 1, 1).itos)
+
+    @torch.no_grad()
+    def sample(self):
+        """
+        ### Evaluation
+
+        We use the sampling function to evaluate the model on a set of problems
+        """
+
+        # Skip in the first epoch
+        if self.training_loop.idx < 1:
+            return
+
+        # Create a dataset to generate problems
+        dataset = ArithmeticMultiplicationDataset(self.seq_len, self.max_digits, self.base, 1)
+        # Get a set of problems and answers
+        qa = [dataset.get_qa() for _ in range(self.n_tests)]
+        # Collect the problems only
+        questions = [p[0] for p in qa]
+
+        # Create a tensor with only the initial token
+        data = torch.tensor([[dataset.stoi[p[0]] for p in questions]])
+        # Move to device
+        data = data.to(self.device)
+
+        # Number of sequences that have completed
+        finished = torch.zeros((len(questions),)).bool().to(self.device)
+        # Token id of the new line character - this marks end of the answer
+        new_line = dataset.stoi['\n']
+
+        # Sampled results
+        results = [p[0] for p in questions]
+
+        # Sample upto sequence length
+        for i in monit.iterate('Sample', self.seq_len - 1):
+            # If all the sequences have completed we skip this
+            if finished.sum() == len(finished):
+                continue
+
+            # Get the model output
+            output, *_ = self.model(data)
+            # Get the model prediction (greedy)
+            output = output[-1].argmax(dim=-1)
+
+            # Override with the question
+            for j, p in enumerate(questions):
+                if len(p) > i + 1:
+                    output[j] = dataset.stoi[p[i + 1]]
+
+            # Find which sequences have finished
+            finished = finished | (output == new_line)
+            # Skip if all have finished
+            if finished.sum() == len(finished):
+                continue
+
+            # Add the next token to the input
+            data = torch.cat([data, output[None, :]], dim=0)
+
+            # Get the sampled results
+            for j, c in enumerate(output):
+                results[j] += dataset.itos[c]
+
+        # Discard everything after the answer in the results
+        results = [r.split('\n')[0] for r in results]
+
+        # Log a sample
+        res_sample = results[0].split(';')
+        logger.log([(res_sample[0], Text.key), (';', Text.subtle), (';'.join(res_sample[1:]), Text.none)])
+
+        # Get the answers
+        results = [r.split('x==')[-1] for r in results]
+
+        # Count the number of correct answers
+        correct = 0
+        for r, _qa in zip(results, qa):
+            if r == _qa[1]:
+                correct += 1
+
+        # Log the score
+        tracker.save('score', correct / len(results))
+
+
+@option(ArithmeticMultiplicationAutoregression.train_loader)
+def arithmetic_train_loader(c: ArithmeticMultiplicationAutoregression):
+    """
+    Training data loader
+    """
+    return DataLoader(ArithmeticMultiplicationDataset(c.seq_len, c.max_digits, c.base, c.train_sequences_per_epoch),
+                      batch_size=c.batch_size,
+                      collate_fn=transpose_batch,
+                      num_workers=4)
+
+
+def _test():
+    """
+    Code to test generated problems
+    """
+    dataset = ArithmeticMultiplicationDataset(256, 4, 4, 10)
+
+    print(dataset.decode(dataset.get_packed_math_input()[0]))
+
+
+#
+if __name__ == '__main__':
+    _test()
diff --git a/labml_nn/experiments/algo_tasks/copy.py b/labml_nn/experiments/algo_tasks/copy.py
new file mode 100644
index 00000000..2a4d95ca
--- /dev/null
+++ b/labml_nn/experiments/algo_tasks/copy.py
@@ -0,0 +1,160 @@
+import random
+from typing import List
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from labml import tracker
+from labml.configs import option
+from labml_helpers.train_valid import BatchIndex
+from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs, transpose_batch
+
+
+class CopyPermDataset(Dataset):
+    """
+    """
+
+    def __init__(self, seq_len: int, n_sequences: int):
+        """
+        :param seq_len: is the sequence length of generated math problems.
+            We fill as many problems as possible upto this length
+        """
+        self.n_sequences = n_sequences
+        self.seq_len = seq_len
+        self.code_len = (self.seq_len + 2) // 2
+        self.letters = '01'  # 'acgt'  # string.ascii_lowercase
+        # Token id to string
+        self.itos = list(self.letters + '>')
+        # Character to token id
+        self.stoi = {c: i for i, c in enumerate(self.itos)}
+
+    def random_string(self, n_chars: int):
+        return ''.join(random.choice(self.letters) for _ in range(n_chars))
+
+    def generate_problem(self):
+        code = self.random_string(self.code_len)
+        offset = random.randrange(0, len(code))
+        perm = code[offset:] + code[:offset]
+        return code + '>' + perm
+
+    def encode(self, s: str):
+        """
+        Encode a given string
+        """
+        return [self.stoi[c] for c in s]
+
+    def decode(self, arr: List[int]):
+        """
+        Decode a list of token ids
+        """
+        return ''.join([self.itos[c] for c in arr])
+
+    def __getitem__(self, idx: int):
+        """
+        Get a input and target pair for auto-regressive modelling
+        """
+        s = torch.tensor(self.encode(self.generate_problem()))
+        return s[:self.seq_len], s[self.code_len + 1:self.seq_len + 1]
+
+    def __len__(self):
+        """
+        Number of sequences per epoch
+        """
+        return self.n_sequences
+
+
+class CopyAutoregression(NLPAutoRegressionConfigs):
+    """
+    ## Arithmetic Task Experiment Configurations
+    """
+    # Number of training sequences per epoch
+    train_sequences_per_epoch: int = 2 ** 12
+    # Training data loader
+    train_loader: DataLoader = 'copy_train_loader'
+    # Number of problems in evaluation
+    n_tests: int = 64
+    # No need of a validation dataset
+    validator = None
+    # Number of times to run evaluations per epoch
+    inner_iterations = 4
+    # Number of tokens in the vocabulary
+    n_tokens = len(CopyPermDataset(1, 1).itos)
+
+    @torch.no_grad()
+    def sample(self):
+        pass
+
+    def step(self, batch: any, batch_idx: BatchIndex):
+        """
+        ### Training or validation step
+        """
+
+        # Set training/eval mode
+        self.model.train(self.mode.is_train)
+
+        # Move data to the device
+        data, target = batch[0].to(self.device), batch[1].to(self.device)
+
+        # Update global step (number of tokens processed) when in training mode
+        if self.mode.is_train:
+            tracker.add_global_step(data.shape[0] * data.shape[1])
+
+        # Whether to capture model outputs
+        with self.mode.update(is_log_activations=batch_idx.is_last and self.is_log_model_activations):
+            # Get model outputs.
+            # It's returning a tuple for states when using RNNs.
+            # This is not implemented yet. 😜
+            output, *_ = self.model(data)
+
+        output = output[-len(target):]
+        # Calculate and log loss
+        loss = self.loss_func(output, target)
+        tracker.add("loss.", loss)
+
+        # Calculate and log accuracy
+        self.accuracy(output, target)
+        self.accuracy.track()
+
+        self.other_metrics(output, target)
+
+        # Train the model
+        if self.mode.is_train:
+            # Calculate gradients
+            loss.backward()
+            # Clip gradients
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
+            # Take optimizer step
+            self.optimizer.step()
+            # Log the model parameters and gradients on last batch of every epoch
+            if batch_idx.is_last and self.is_log_model_params_grads:
+                tracker.add('model', self.model)
+            # Clear the gradients
+            self.optimizer.zero_grad()
+
+        # Save the tracked metrics
+        tracker.save()
+
+
+@option(CopyAutoregression.train_loader)
+def copy_train_loader(c: CopyAutoregression):
+    """
+    Training data loader
+    """
+    return DataLoader(CopyPermDataset(c.seq_len, c.train_sequences_per_epoch),
+                      batch_size=c.batch_size,
+                      collate_fn=transpose_batch,
+                      num_workers=4)
+
+
+def _test():
+    """
+    Code to test generated problems
+    """
+    dataset = CopyPermDataset(32, 1)
+
+    print(dataset.generate_problem())
+
+
+#
+if __name__ == '__main__':
+    _test()
diff --git a/labml_nn/experiments/algo_tasks/copy_repeat.py b/labml_nn/experiments/algo_tasks/copy_repeat.py
new file mode 100644
index 00000000..738f6433
--- /dev/null
+++ b/labml_nn/experiments/algo_tasks/copy_repeat.py
@@ -0,0 +1,179 @@
+import random
+from typing import List
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from labml import tracker
+from labml.configs import option
+from labml_helpers.train_valid import BatchIndex
+from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs, transpose_batch
+
+
+class CopyPermRepeatDataset(Dataset):
+    """
+    """
+
+    def __init__(self, seq_len: int, substr_len: int, rnd_len: int, n_sequences: int):
+        """
+        :param seq_len: is the sequence length of generated math problems.
+            We fill as many problems as possible upto this length
+        """
+        self.rnd_len = rnd_len
+        self.substr_len = substr_len
+        self.n_sequences = n_sequences
+        self.seq_len = seq_len
+        self.letters = 'acgt'  # string.ascii_lowercase  # '01'  # 'acgt'  #
+        # Token id to string
+        self.itos = list(self.letters + '>')
+        # Character to token id
+        self.stoi = {c: i for i, c in enumerate(self.itos)}
+
+    def random_string(self, n_chars: int):
+        return ''.join(random.choice(self.letters) for _ in range(n_chars))
+
+    def generate_problem(self):
+        pure = self.random_string(self.substr_len)
+        out = pure
+        mask = [False] * len(out)
+        while len(out) <= self.seq_len:
+            s = self.random_string(random.randrange(1, self.rnd_len))
+            out += s + '>'
+            mask += [False] * (len(s) + 1)
+            pure += s
+
+            offset = random.randrange(0, len(pure) - self.substr_len)
+            copy = pure[offset:offset + self.substr_len]
+
+            out += copy
+            mask += [False] + [True] * (self.substr_len - 1)
+            pure += copy
+
+        return out, mask
+
+    def encode(self, s: str):
+        """
+        Encode a given string
+        """
+        return [self.stoi[c] for c in s]
+
+    def decode(self, arr: List[int]):
+        """
+        Decode a list of token ids
+        """
+        return ''.join([self.itos[c] for c in arr])
+
+    def __getitem__(self, idx: int):
+        """
+        Get a input and target pair for auto-regressive modelling
+        """
+        s, mask = self.generate_problem()
+        s = torch.tensor(self.encode(s))
+        mask = torch.tensor(mask)
+        target = s * mask + -1 * (~mask)
+        return s[:self.seq_len], target[1:self.seq_len + 1]
+
+    def __len__(self):
+        """
+        Number of sequences per epoch
+        """
+        return self.n_sequences
+
+
+class CopyRepeatAutoregression(NLPAutoRegressionConfigs):
+    """
+    ## Arithmetic Task Experiment Configurations
+    """
+    # Number of training sequences per epoch
+    train_sequences_per_epoch: int = 2 ** 12
+    # Training data loader
+    train_loader: DataLoader = 'copy_train_loader'
+    # Number of problems in evaluation
+    n_tests: int = 64
+    # No need of a validation dataset
+    validator = None
+    # Number of times to run evaluations per epoch
+    inner_iterations = 4
+    # Number of tokens in the vocabulary
+    n_tokens = len(CopyPermRepeatDataset(1, 1, 1, 1).itos)
+
+    substr_len: int = 16
+    rnd_len: int = 12
+
+    @torch.no_grad()
+    def sample(self):
+        pass
+
+    def step(self, batch: any, batch_idx: BatchIndex):
+        """
+        ### Training or validation step
+        """
+
+        # Set training/eval mode
+        self.model.train(self.mode.is_train)
+
+        # Move data to the device
+        data, target = batch[0].to(self.device), batch[1].to(self.device)
+
+        # Update global step (number of tokens processed) when in training mode
+        if self.mode.is_train:
+            tracker.add_global_step(data.shape[0] * data.shape[1])
+
+        # Whether to capture model outputs
+        with self.mode.update(is_log_activations=batch_idx.is_last and self.is_log_model_activations):
+            # Get model outputs.
+            # It's returning a tuple for states when using RNNs.
+            # This is not implemented yet. 😜
+            output, *_ = self.model(data)
+
+        # Calculate and log loss
+        loss = self.loss_func(output, target)
+        tracker.add("loss.", loss)
+
+        # Calculate and log accuracy
+        self.accuracy(output, target)
+        self.accuracy.track()
+
+        self.other_metrics(output, target)
+
+        # Train the model
+        if self.mode.is_train:
+            # Calculate gradients
+            loss.backward()
+            # Clip gradients
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
+            # Take optimizer step
+            self.optimizer.step()
+            # Log the model parameters and gradients on last batch of every epoch
+            if batch_idx.is_last and self.is_log_model_params_grads:
+                tracker.add('model', self.model)
+            # Clear the gradients
+            self.optimizer.zero_grad()
+
+        # Save the tracked metrics
+        tracker.save()
+
+
+@option(CopyRepeatAutoregression.train_loader)
+def copy_train_loader(c: CopyRepeatAutoregression):
+    """
+    Training data loader
+    """
+    return DataLoader(CopyPermRepeatDataset(c.seq_len, c.substr_len, c.rnd_len, c.train_sequences_per_epoch),
+                      batch_size=c.batch_size,
+                      collate_fn=transpose_batch)
+    # num_workers=4)
+
+
+def _test():
+    """
+    Code to test generated problems
+    """
+    dataset = CopyPermRepeatDataset(32, 8, 8, 1)
+
+    print(dataset.generate_problem())
+
+
+#
+if __name__ == '__main__':
+    _test()
diff --git a/labml_nn/experiments/nlp_autoregression.py b/labml_nn/experiments/nlp_autoregression.py
index fe4e67d1..cdcd0cc6 100644
--- a/labml_nn/experiments/nlp_autoregression.py
+++ b/labml_nn/experiments/nlp_autoregression.py
@@ -19,7 +19,7 @@
 from labml.logger import Text
 from labml_helpers.datasets.text import TextDataset, SequentialDataLoader, SequentialUnBatchedDataset, TextFileDataset
 from labml_helpers.device import DeviceConfigs
-from labml_helpers.metrics.accuracy import Accuracy
+from labml_helpers.metrics.accuracy import AccuracyMovingAvg
 from labml_helpers.module import Module
 from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
 from labml_nn.optimizers.configs import OptimizerConfigs
@@ -30,9 +30,9 @@ class CrossEntropyLoss(Module):
     ### Cross entropy loss
     """
 
-    def __init__(self):
+    def __init__(self, ignore_index: int = -100):
         super().__init__()
-        self.loss = nn.CrossEntropyLoss()
+        self.loss = nn.CrossEntropyLoss(ignore_index=ignore_index)
 
     def forward(self, outputs, targets):
         return self.loss(outputs.view(-1, outputs.shape[-1]), targets.view(-1))
@@ -75,9 +75,9 @@ class NLPAutoRegressionConfigs(TrainValidConfigs):
     is_save_models = True
 
     # Loss function
-    loss_func = CrossEntropyLoss()
+    loss_func = CrossEntropyLoss(ignore_index=-1)
     # Accuracy function
-    accuracy = Accuracy()
+    accuracy = AccuracyMovingAvg()
     # Model embedding size
     d_model: int = 512
     # Gradient clipping
@@ -297,10 +297,7 @@ def transpose_batch(batch):
 
     transposed_data = list(zip(*batch))
     # Stack the batch along the second dimension `dim=1`
-    src = torch.stack(transposed_data[0], dim=1)
-    tgt = torch.stack(transposed_data[1], dim=1)
-
-    return src, tgt
+    return tuple(torch.stack(d, dim=1) for d in transposed_data)
 
 
 @option(NLPAutoRegressionConfigs.train_loader)
diff --git a/labml_nn/transformers/rope/__init__.py b/labml_nn/transformers/rope/__init__.py
index ae93f859..0f4cfec2 100644
--- a/labml_nn/transformers/rope/__init__.py
+++ b/labml_nn/transformers/rope/__init__.py
@@ -115,37 +115,67 @@ class RotaryPositionalEmbeddings(nn.Module):
     \end{pmatrix} \\
     \end{align}
     """
+
     def __init__(self, d: int, base: int = 10_000):
         """
         * `d` is the number of features $d$
         * `base` is the constant used for calculating $\Theta$
         """
         super().__init__()
-        # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
-        self.theta = nn.Parameter(1. / (base ** (torch.arange(0, d, 2).float() / d)), requires_grad=False)
 
-    def forward(self, x: torch.Tensor):
+        self.base = base
+        self.d = d
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def _build_cache(self, x: torch.Tensor):
         """
-        * `x` is the Tensor at the head of a key or a query with shape `[seq_len, batch_size, n_heads, d]`
+        Cache $\cos$ and $\sin$ values
         """
-        # Extract the shape
-        seq_len, batch_size, n_heads, d = x.shape
+        # Return if cache is already built
+        if self.cos_cached is not None and x.shape[0] <= self.cos_cached.shape[0]:
+            return
 
-        # $\frac{d}{2}$
-        d_2 = d // 2
+        # Get sequence length
+        seq_len = x.shape[0]
+
+        # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
+        theta = 1. / (self.base ** (torch.arange(0, self.d, 2).float() / self.d)).to(x.device)
 
         # Create position indexes `[0, 1, ..., seq_len - 1]`
-        seq_idx = torch.arange(seq_len, device=x.device).type_as(self.theta)
+        seq_idx = torch.arange(seq_len, device=x.device).float().to(x.device)
 
         # Calculate the product of position index and $\theta_i$
-        idx_theta = torch.einsum('n,d->nd', seq_idx, self.theta)
+        idx_theta = torch.einsum('n,d->nd', seq_idx, theta)
 
         # Concatenate so that for row $m$ we have
         # $[m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}, m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}]$
         idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
 
-        # Calculate $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., -x^{(\frac{d}{2})}]$
-        neg_half_x = torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
+        # Cache them
+        self.cos_cached = idx_theta2.cos()[:, None, None, :]
+        self.sin_cached = idx_theta2.sin()[:, None, None, :]
+
+    def _neg_half(self, x: torch.Tensor):
+        # $\frac{d}{2}$
+        d_2 = self.d // 2
+
+        # Calculate $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
+        return torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
+
+    def forward(self, x: torch.Tensor):
+        """
+        * `x` is the Tensor at the head of a key or a query with shape `[seq_len, batch_size, n_heads, d]`
+        """
+        # Cache $\cos$ and $\sin$ values
+        self._build_cache(x)
+
+        # Split the features, we can choose to apply rotary embeddings only to a partial set of features.
+        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})}]$
+        neg_half_x = self._neg_half(x_rope)
 
         # Calculate
         #
@@ -157,10 +187,10 @@ def forward(self, x: torch.Tensor):
         # \end{align}
         #
         # for $i \in {1, 2, ..., \frac{d}{2}}$
-        rx = (x * idx_theta2.cos()[:, None, None, :]) + (neg_half_x * idx_theta2.sin()[:, None, None, :])
+        x_rope = (x_rope * self.cos_cached[:x.shape[0]]) + (neg_half_x * self.sin_cached[:x.shape[0]])
 
         #
-        return rx
+        return torch.cat((x_rope, x_pass), dim=-1)
 
 
 class RotaryPEMultiHeadAttention(MultiHeadAttention):
@@ -170,15 +200,13 @@ class RotaryPEMultiHeadAttention(MultiHeadAttention):
     We override [multi-head attention from original transformer](../mha.html).
     """
 
-    def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
-        # The linear transformations do not need a bias since we
-        # explicitly include it when calculating scores.
-        # However having a bias for `value` might make sense.
-        super().__init__(heads, d_model, dropout_prob, bias=False)
+    def __init__(self, heads: int, d_model: int, rope_percentage: float = 0.5, dropout_prob: float = 0.0):
+        super().__init__(heads, d_model, dropout_prob)
 
         # Rotary positional embedding layers
-        self.query_rotary_pe = RotaryPositionalEmbeddings(self.d_k)
-        self.key_rotary_pe = RotaryPositionalEmbeddings(self.d_k)
+        d_rope = int(self.d_k * rope_percentage)
+        self.query_rotary_pe = RotaryPositionalEmbeddings(d_rope)
+        self.key_rotary_pe = RotaryPositionalEmbeddings(d_rope)
 
     def get_scores(self, query: torch.Tensor, key: torch.Tensor):
         """
diff --git a/labml_nn/transformers/rope/experiment.py b/labml_nn/transformers/rope/experiment.py
index 24c80f9c..9f63e6e4 100644
--- a/labml_nn/transformers/rope/experiment.py
+++ b/labml_nn/transformers/rope/experiment.py
@@ -20,7 +20,7 @@
 # ### Rotary PE attention
 def _rotary_pe_mha(c: TransformerConfigs):
     from labml_nn.transformers.rope import RotaryPEMultiHeadAttention
-    return RotaryPEMultiHeadAttention(c.n_heads, c.d_model)
+    return RotaryPEMultiHeadAttention(c.n_heads, c.d_model, 1.)
 
 
 # Configuration options
@@ -43,7 +43,7 @@ def _model(c: Configs):
 
 def main():
     # Create experiment
-    experiment.create(name="rotary_pe_transformer")
+    experiment.create(name="rotary_pe_transformer", writers={'screen'})
     # Create configs
     conf = Configs()
     # Override configurations
diff --git a/labml_nn/transformers/rope/value_pe/__init__.py b/labml_nn/transformers/rope/value_pe/__init__.py
new file mode 100644
index 00000000..414c25b7
--- /dev/null
+++ b/labml_nn/transformers/rope/value_pe/__init__.py
@@ -0,0 +1,246 @@
+"""
+---
+title: Rotary Positional Embeddings with Relative distance (RoPER)
+summary: >
+  This is an implementation of RoPER which adds relative distance information to embeddings on
+  top of RoPE introduced in RoFormer: Enhanced Transformer with Rotary Position Embedding
+---
+
+*RoPER is work by [Georges Harik (@gharik)](https://twitter.com/gharik),
+and this implementation is based on his original code.*
+
+# Rotary Positional Embeddings with Relative distance (RoPER)
+
+[Rotary Positional Embeddings (RoPE)](https://papers.labml.ai/paper/2104.09864) 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](https://papers.labml.ai/paper/2c364684b15b11ecac827bce58715ee7).
+
+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](experiment.html) 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^{(1)}_i, v^{(2)}_i, \dots$.
+
+Normally, we would take the weight sum of value embeddings
+
+$$o^{(j)}_n = \sum_i a_{n,i} v^{(j)}_i$$
+
+This doesn't explicitly add any distance information about the positions $i$ to final
+result $o^{(j)}_n$.
+
+RoPER pairs features like RoPE and transform them.
+For a pair $v^{(1)}_m$ and $v^{(2)}_m$ it transforms them by
+ $RoPE\big(v^{(1)}_m, v^{(2)}_m, m\big)$.
+Let us donate the transformed features with $\hat{v}^{(1)}_m, \hat{v}^{(2)}_m$.
+Then it rotates the weighted sum $\hat{o}^{(j)}_n$ in the the reverse direction with
+ $RoPE\big(\hat{o}^{(1)}_n, \hat{o}^{(2)}_n, -n\big)$.
+*Note the *$-n$.
+
+Note that,
+
+\begin{align}
+RoPE\big(x^{(1)}_m, x^{(2)}_m, m\big) &=
+\begin{pmatrix}
+\cos m \theta & - \sin m \theta \\
+\sin m \theta & \cos m \theta
+\end{pmatrix}
+\begin{pmatrix}
+x^{(1)}_m \\
+x^{(2)}_m \\
+\end{pmatrix} \\
+&=
+\begin{pmatrix}
+x^{(1)}_m \cos m\theta - x^{(2)}_m \sin m \theta \\
+x^{(2)}_m \cos m\theta + x^{(1)}_m \sin m \theta \\
+\end{pmatrix} \\
+\end{align}
+
+Final output after with the transformations is,
+
+\begin{align}
+RoPE\big(\hat{o}^{(1)}_n, \hat{o}^{(2)}_n, -n\big) &= \\
+\begin{pmatrix}
+\hat{o}^{(1)}_n \cos n\theta + \hat{o}^{(2)}_n \sin n \theta \\
+\hat{o}^{(2)}_n \cos n\theta - \hat{o}^{(1)}_n \sin n \theta \\
+\end{pmatrix} \\
+\end{align}
+
+*Note that *$\sin (-n \theta) = -\sin n \theta$.
+
+Let's expand the first term $\hat{o}^{(1)}_n \cos n\theta + \hat{o}^{(2)}_n \sin n \theta$,
+
+\begin{align}
+\hat{o}^{(1)}_n \cos n\theta + \hat{o}^{(2)}_n \sin n \theta &= \\
+\sum_i a_{n,i} \hat{v}^{(1)}_i \cos n\theta + \sum_i a_{n,i} \hat{v}^{(2)}_i \sin n \theta &= \\
+
+\sum_i a_{n,i} \Big( v^{(1)}_i \cos i\theta - v^{(2)}_i \sin i \theta \Big) \cos n\theta &+ \\
+\sum_i a_{n,i} \Big( v^{(2)}_i \cos i\theta + v^{(1)}_i \sin i \theta \Big) \sin m \theta &= \\
+
+\sum_i a_{n,i} v^{(1)}_i \Big( \cos i\theta \cos n\theta + \sin i \theta \sin n \theta \Big) &+ \\
+\sum_i a_{n,i} v^{(2)}_i \Big( \cos i\theta \sin n\theta - \sin i \theta \cos n \theta \Big) &= \\
+
+\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta - \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta &= \\
+
+\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta - \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta
+\end{align}
+
+Simiarly we can show the second term is equal to,
+
+$$\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta + \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta$$
+
+Which gives,
+
+\begin{align}
+RoPE\big(\hat{o}^{(1)}_n, \hat{o}^{(2)}_n, -n\big) &= \\
+\begin{pmatrix}
+\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta - \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta \\
+\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta + \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta \\
+\end{pmatrix} &= \\
+\sum_i a_{n,i} RoPE \big (v^{(1)}_i, v^{(1)}_i, (i - n) \theta \big)
+\end{align}
+
+That is, the weighted average of values rotated relative to current position.
+
+[Here's an experiment](arithmetic_experiment.html) that uses RoPER on an arthmetic addition task.
+"""
+
+from typing import Optional
+
+import torch
+
+from labml_nn.transformers.rope import RotaryPositionalEmbeddings, RotaryPEMultiHeadAttention
+
+
+class ReverseRotaryPositionalEmbeddings(RotaryPositionalEmbeddings):
+    """
+    ## RoPE module that rotates in the opposite direction
+
+    This inherits from [RoPE rotation implementation](../index.html) and changes the direction.
+    """
+
+    def forward(self, x: torch.Tensor):
+        """
+        * `x` is the Tensor at the head of a key or a query with shape `[seq_len, batch_size, n_heads, d]`
+        """
+        # Cache $\cos$ and $\sin$ values
+        self._build_cache(x)
+
+        # Split the features, we can choose to apply rotary embeddings only to a partial set of features.
+        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})}]$
+        neg_half_x = self._neg_half(x_rope)
+
+        # Calculate
+        #
+        # \begin{align}
+        # \begin{pmatrix}
+        # x^{(i)}_m \cos -m \theta_i - x^{(i + \frac{d}{2})}_m \sin -m \theta_i \\
+        # x^{(i + \frac{d}{2})}_m \cos -m\theta_i + x^{(i)}_m \sin -m \theta_i \\
+        # \end{pmatrix} = \\
+        # \begin{pmatrix}
+        # x^{(i)}_m \cos m \theta_i + x^{(i + \frac{d}{2})}_m \sin m \theta_i \\
+        # x^{(i + \frac{d}{2})}_m \cos m\theta_i - x^{(i)}_m \sin m \theta_i \\
+        # \end{pmatrix} \\
+        # \end{align}
+        #
+        # for $i \in {1, 2, ..., \frac{d}{2}}$
+        x_rope = (x_rope * self.cos_cached[:x.shape[0]]) - (neg_half_x * self.sin_cached[:x.shape[0]])
+
+        #
+        return torch.cat((x_rope, x_pass), dim=-1)
+
+
+class RotaryValuePEMultiHeadAttention(RotaryPEMultiHeadAttention):
+    """
+    ## Multi-head attention with rotary positional embeddings
+
+    We override [multi-head attention from original transformer](../mha.html).
+    """
+
+    def __init__(self, heads: int, d_model: int,
+                 rope_percentage: float = 0.5, rope_value_percentage: float = 0.5,
+                 dropout_prob: float = 0.0):
+        super().__init__(heads, d_model, rope_percentage, dropout_prob)
+
+        # Rotary positional embedding layers
+        d_rope_value = int(self.d_k * rope_value_percentage)
+
+        self.value_rotary_pe = RotaryPositionalEmbeddings(d_rope_value)
+        self.value_reverse_rotary_pe = ReverseRotaryPositionalEmbeddings(d_rope_value)
+
+    def forward(self, *,
+                query: torch.Tensor,
+                key: torch.Tensor,
+                value: torch.Tensor,
+                mask: Optional[torch.Tensor] = None):
+        """
+        `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`.
+        """
+
+        # `query`, `key` and `value`  have shape `[seq_len, batch_size, d_model]`
+        seq_len, batch_size, _ = query.shape
+
+        if mask is not None:
+            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]`.
+        query = self.query(query)
+        key = self.key(key)
+        value = self.value(value)
+
+        # Compute attention scores $Q K^\top$.
+        # This gives a tensor of shape `[seq_len, seq_len, batch_size, heads]`.
+        scores = self.get_scores(query, key)
+
+        # Scale scores $\frac{Q K^\top}{\sqrt{d_k}}$
+        scores *= self.scale
+
+        # Apply mask
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+
+        # $softmax$ attention along the key sequence dimension
+        # $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
+        attn = self.softmax(scores)
+
+        # Apply dropout
+        attn = self.dropout(attn)
+
+        # Rotate value embeddings before taking the weighted sum so that they contain positional information
+        value = self.value_rotary_pe(value)
+
+        # Multiply by values
+        # $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
+        x = torch.einsum("ijbh,jbhd->ibhd", attn, self.value_rotary_pe(value))
+
+        # Rotate in the opposite direction so that each embedding hold the relative positions
+        x = self.value_reverse_rotary_pe(x)
+
+        # Save attentions for any other calculations
+        self.attn = attn.detach()
+
+        # Concatenate multiple heads
+        x = x.reshape(seq_len, batch_size, -1)
+
+        # Output layer
+        return self.output(x)
diff --git a/labml_nn/transformers/rope/value_pe/experiments/__init__.py b/labml_nn/transformers/rope/value_pe/experiments/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/labml_nn/transformers/rope/value_pe/experiments/arithmetic_addition.py b/labml_nn/transformers/rope/value_pe/experiments/arithmetic_addition.py
new file mode 100644
index 00000000..9ef756fe
--- /dev/null
+++ b/labml_nn/transformers/rope/value_pe/experiments/arithmetic_addition.py
@@ -0,0 +1,93 @@
+"""
+---
+title: Rotary Positional Embeddings with Relative distance (RoPER) Experiment
+summary: This experiment trains a transformer model with Rotary Positional Embeddings with
+ Relative Distance (RoPER) on the arithmetic addition task.
+---
+
+# Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) Experiment
+"""
+
+from labml import experiment
+from labml.configs import calculate
+from labml_nn.experiments.algo_tasks.arithmetic_addition import ArithmeticAdditionAutoregression
+from labml_nn.transformers import TransformerConfigs
+from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs
+
+
+class Configs(RoPEConfigs, ArithmeticAdditionAutoregression):
+    """
+    We inherit [RoPE experiment](../experiment.html) and use it for
+    [arithmetic addition task](../../experiments/arithmetic_dataset.html).
+
+    We add the option to change attention to use Rotary Positional Embeddings with Relative distance (RoPER)
+    below.
+    """
+    pass
+
+
+def _rotary_value_pe_mha(c: TransformerConfigs):
+    """
+    Use Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) in attention.
+    """
+    from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
+    return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)
+
+
+# Configuration options
+calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)
+
+
+def main():
+    # Create experiment
+    experiment.create(name="roper_addition", comment="rotary value 7", writers={'screen', 'labml', 'comet'})
+    # Create configs
+    conf = Configs()
+    # Override configurations
+    experiment.configs(conf, {
+        'max_digits': 7,
+
+        # No fixed positional embeddings
+        'transformer.src_embed': 'no_pos',
+        'transformer.tgt_embed': 'no_pos',
+
+        # Encoder with RoPER attention
+        'transformer.encoder_attn': 'rotary_value',
+        # Encoder with RoPE attention
+        # 'transformer.encoder_attn': 'rotary',
+
+        #
+        'model': 'rotary_pe_transformer',
+
+        # Use a context size of $256$
+        'seq_len': 512,
+        # Train for 32 epochs
+        'epochs': 20,
+        # Batch size $4$
+        'batch_size': 16,
+
+        # Model size
+        'd_model': 128,
+        'transformer.ffn.d_ff': 512,
+        'transformer.n_heads': 4,
+        'transformer.dropout': 0.0,
+
+        # Use [Adam optimizer](../../optimizers/noam.html)
+        'optimizer.optimizer': 'Adam',
+        'optimizer.learning_rate': 2.5e-4,
+    })
+
+    # Set models for saving and loading
+    experiment.add_pytorch_models({'model': conf.model})
+
+    # Start the experiment
+    with experiment.start():
+        # Run training
+        conf.run()
+
+
+#
+if __name__ == '__main__':
+    main()
diff --git a/labml_nn/transformers/rope/value_pe/experiments/arithmetic_multiplication.py b/labml_nn/transformers/rope/value_pe/experiments/arithmetic_multiplication.py
new file mode 100644
index 00000000..f60f5016
--- /dev/null
+++ b/labml_nn/transformers/rope/value_pe/experiments/arithmetic_multiplication.py
@@ -0,0 +1,94 @@
+"""
+---
+title: Rotary Positional Embeddings with Relative distance (RoPER) Experiment
+summary: This experiment trains a transformer model with Rotary Positional Embeddings with
+ Relative Distance (RoPER) on the arithmetic addition task.
+---
+
+# Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) Experiment
+"""
+
+from labml import experiment
+from labml.configs import calculate
+from labml_nn.experiments.algo_tasks.arithmetic_multiplication import ArithmeticMultiplicationAutoregression
+from labml_nn.transformers import TransformerConfigs
+from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs
+
+
+class Configs(RoPEConfigs, ArithmeticMultiplicationAutoregression):
+    """
+    We inherit [RoPE experiment](../experiment.html) and use it for
+    [arithmetic addition task](../../experiments/arithmetic_dataset.html).
+
+    We add the option to change attention to use Rotary Positional Embeddings with Relative distance (RoPER)
+    below.
+    """
+    pass
+
+
+def _rotary_value_pe_mha(c: TransformerConfigs):
+    """
+    Use Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) in attention.
+    """
+    from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
+    return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)
+
+
+# Configuration options
+calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)
+
+
+def main():
+    # Create experiment
+    experiment.create(name="roper_mult", comment="4", writers={'screen', 'labml'})
+    # Create configs
+    conf = Configs()
+    # Override configurations
+    experiment.configs(conf, {
+        'max_digits': 8,
+        'base': 4,
+
+        # No fixed positional embeddings
+        'transformer.src_embed': 'no_pos',
+        'transformer.tgt_embed': 'no_pos',
+
+        # Encoder with RoPER attention
+        # 'transformer.encoder_attn': 'rotary_value',
+        # Encoder with RoPE attention
+        'transformer.encoder_attn': 'rotary',
+
+        #
+        'model': 'rotary_pe_transformer',
+
+        # Use a context size of $256$
+        'seq_len': 512,
+        # Train for 32 epochs
+        'epochs': 20,
+        # Batch size $4$
+        'batch_size': 16,
+
+        # Model size
+        'd_model': 128,
+        'transformer.ffn.d_ff': 512,
+        'transformer.n_heads': 4,
+        'transformer.dropout': 0.0,
+
+        # Use [Adam optimizer](../../optimizers/noam.html)
+        'optimizer.optimizer': 'Adam',
+        'optimizer.learning_rate': 2.5e-4,
+    })
+
+    # Set models for saving and loading
+    experiment.add_pytorch_models({'model': conf.model})
+
+    # Start the experiment
+    with experiment.start():
+        # Run training
+        conf.run()
+
+
+#
+if __name__ == '__main__':
+    main()
diff --git a/labml_nn/transformers/rope/value_pe/experiments/copy_experiment.py b/labml_nn/transformers/rope/value_pe/experiments/copy_experiment.py
new file mode 100644
index 00000000..c2065624
--- /dev/null
+++ b/labml_nn/transformers/rope/value_pe/experiments/copy_experiment.py
@@ -0,0 +1,92 @@
+"""
+---
+title: Rotary Positional Embeddings with Relative distance (RoPER) Experiment
+summary: This experiment trains a transformer model with Rotary Positional Embeddings with
+ Relative Distance (RoPER) on the arithmetic addition task.
+---
+
+# Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) Experiment
+"""
+
+from labml import experiment
+from labml.configs import calculate
+from labml_nn.experiments.algo_tasks.copy import CopyAutoregression
+from labml_nn.transformers import TransformerConfigs
+from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs
+
+
+class Configs(RoPEConfigs, CopyAutoregression):
+    """
+    We inherit [RoPE experiment](../experiment.html) and use it for
+    [arithmetic addition task](../../experiments/arithmetic_dataset.html).
+
+    We add the option to change attention to use Rotary Positional Embeddings with Relative distance (RoPER)
+    below.
+    """
+    pass
+
+
+def _rotary_value_pe_mha(c: TransformerConfigs):
+    """
+    Use Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) in attention.
+    """
+    from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
+    return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)
+
+
+# Configuration options
+calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)
+
+
+def main():
+    # Create experiment
+    experiment.create(name="roper_copy", comment="rotary value 01", writers={'screen', 'labml'})
+    # Create configs
+    conf = Configs()
+    # Override configurations
+    experiment.configs(conf, {
+        # No fixed positional embeddings
+        'transformer.src_embed': 'no_pos',
+        'transformer.tgt_embed': 'no_pos',
+
+        # Encoder with RoPER attention
+        'transformer.encoder_attn': 'rotary_value',
+        # Encoder with RoPE attention
+        # 'transformer.encoder_attn': 'rotary',
+
+        #
+        'model': 'rotary_pe_transformer',
+
+        # Use a context size of $256$
+        'seq_len': 128,
+        # Train for 32 epochs
+        'epochs': 20,
+        # Batch size $4$
+        'batch_size': 16,
+
+        # Model size
+        'd_model': 128,
+        'transformer.ffn.d_ff': 512,
+        'transformer.n_heads': 4,
+        'transformer.n_layers': 3,
+        'transformer.dropout': 0.0,
+
+        # Use [Adam optimizer](../../optimizers/noam.html)
+        'optimizer.optimizer': 'Adam',
+        'optimizer.learning_rate': 2.5e-4,
+    })
+
+    # Set models for saving and loading
+    experiment.add_pytorch_models({'model': conf.model})
+
+    # Start the experiment
+    with experiment.start():
+        # Run training
+        conf.run()
+
+
+#
+if __name__ == '__main__':
+    main()
diff --git a/labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py b/labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
new file mode 100644
index 00000000..a7da19aa
--- /dev/null
+++ b/labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
@@ -0,0 +1,94 @@
+"""
+---
+title: Rotary Positional Embeddings with Relative distance (RoPER) Experiment
+summary: This experiment trains a transformer model with Rotary Positional Embeddings with
+ Relative Distance (RoPER) on the arithmetic addition task.
+---
+
+# Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) Experiment
+"""
+
+from labml import experiment
+from labml.configs import calculate
+from labml_nn.experiments.algo_tasks.copy_repeat import CopyRepeatAutoregression
+from labml_nn.transformers import TransformerConfigs
+from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs
+
+
+class Configs(RoPEConfigs, CopyRepeatAutoregression):
+    """
+    We inherit [RoPE experiment](../experiment.html) and use it for
+    [arithmetic addition task](../../experiments/arithmetic_dataset.html).
+
+    We add the option to change attention to use Rotary Positional Embeddings with Relative distance (RoPER)
+    below.
+    """
+    pass
+
+
+def _rotary_value_pe_mha(c: TransformerConfigs):
+    """
+    Use Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) in attention.
+    """
+    from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
+    return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)
+
+
+# Configuration options
+calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)
+
+
+def main():
+    # Create experiment
+    experiment.create(name="roper_copy_repeat", comment="rotary acgt", writers={'screen', 'sqlite'})
+    # Create configs
+    conf = Configs()
+    # Override configurations
+    experiment.configs(conf, {
+        'substr_len': 16,
+
+        # No fixed positional embeddings
+        'transformer.src_embed': 'no_pos',
+        'transformer.tgt_embed': 'no_pos',
+
+        # Encoder with RoPER attention
+        'transformer.encoder_attn': 'rotary_value',
+        # Encoder with RoPE attention
+        # 'transformer.encoder_attn': 'rotary',
+
+        #
+        'model': 'rotary_pe_transformer',
+
+        # Use a context size of $256$
+        'seq_len': 512,
+        # Train for 32 epochs
+        'epochs': 16,
+        # Batch size $4$
+        'batch_size': 16,
+
+        # Model size
+        'd_model': 128,
+        'transformer.ffn.d_ff': 512,
+        'transformer.n_heads': 4,
+        'transformer.n_layers': 3,
+        'transformer.dropout': 0.0,
+
+        # Use [Adam optimizer](../../optimizers/noam.html)
+        'optimizer.optimizer': 'Adam',
+        'optimizer.learning_rate': 2.5e-4,
+    })
+
+    # Set models for saving and loading
+    experiment.add_pytorch_models({'model': conf.model})
+
+    # Start the experiment
+    with experiment.start():
+        # Run training
+        conf.run()
+
+
+#
+if __name__ == '__main__':
+    main()
diff --git a/labml_nn/transformers/rope/value_pe/experiments/experiment.py b/labml_nn/transformers/rope/value_pe/experiments/experiment.py
new file mode 100644
index 00000000..fccc582c
--- /dev/null
+++ b/labml_nn/transformers/rope/value_pe/experiments/experiment.py
@@ -0,0 +1,98 @@
+"""
+---
+title: Rotary Positional Embeddings (RoPE) Experiment
+summary: This experiment trains a transformer model with Rotary Positional Embeddings (RoPE) on tiny Shakespeare dataset.
+---
+
+# Rotary Positional Embeddings (RoPE) Experiment
+
+This is an annotated PyTorch experiment to train a transformer model with Rotary Positional Embeddings (RoPE).
+
+[![View Run](https://img.shields.io/badge/labml-experiment-brightgreen)](https://app.labml.ai/run/1cf508e693be11ecacc98de8b38a61fe)
+"""
+
+from labml import experiment
+from labml.configs import calculate
+from labml_nn.transformers import TransformerConfigs
+from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs
+
+
+# ### Rotary PE attention
+
+class Configs(RoPEConfigs):
+    pass
+
+
+def _rotary_value_pe_mha(c: TransformerConfigs):
+    from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
+    return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)
+
+
+# Configuration options
+calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
+calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)
+
+
+def main():
+    # Create experiment
+    experiment.create(name="rotary_shakespeare", comment="rotary value", writers={'screen', 'labml'})
+    # Create configs
+    conf = Configs()
+    # Override configurations
+    experiment.configs(conf, {
+        # No fixed positional embeddings
+        'transformer.src_embed': 'no_pos',
+        'transformer.tgt_embed': 'no_pos',
+
+        # Encoder with RoPE
+        'transformer.encoder_attn': 'rotary_value',
+        # 'transformer.encoder_attn': 'rotary',
+
+        #
+        'model': 'rotary_pe_transformer',
+
+        # Use character level tokenizer
+        'tokenizer': 'character',
+        # Prompt separator is blank
+        'prompt_separator': '',
+        # Starting prompt for sampling
+        'prompt': 'It is ',
+        # Use Tiny Shakespeare dataset
+        'text': 'tiny_shakespeare',
+
+        # Use a context size of $256$
+        'seq_len': 512,
+        # Train for 32 epochs
+        'epochs': 24,
+        # Batch size $4$
+        'batch_size': 16,
+        # Switch between training and validation for $10$ times
+        # per epoch
+        'inner_iterations': 4,
+
+        # Model size
+        'd_model': 128,
+        'transformer.ffn.d_ff': 512,
+        'transformer.n_heads': 4,
+        'transformer.dropout': 0.0,
+
+        # Use [Adam optimizer](../../optimizers/noam.html)
+        'optimizer.optimizer': 'Adam',
+        'optimizer.learning_rate': 2.5e-4,
+
+        'dataloader_shuffle_with_replacement': True
+    })
+
+    # Set models for saving and loading
+    experiment.add_pytorch_models({'model': conf.model})
+
+    # Start the experiment
+    with experiment.start():
+        # Run training
+        conf.run()
+
+
+#
+if __name__ == '__main__':
+    main()
diff --git a/labml_nn/transformers/rope/value_pe/experiments/results.ipynb b/labml_nn/transformers/rope/value_pe/experiments/results.ipynb
new file mode 100644
index 00000000..1ea15295
--- /dev/null
+++ b/labml_nn/transformers/rope/value_pe/experiments/results.ipynb
@@ -0,0 +1,268 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "a779760d-52b3-41e2-98d2-adb98b8f2a93",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from labml import analytics\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "4db66c78-1949-40e3-b96d-6c797a5f061a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ind_rope = analytics.runs(*analytics.get_experiment_runs('rope_copy_repeat'))\n",
+    "ind_roper = analytics.runs(*analytics.get_experiment_runs('roper_copy_repeat'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "89bb8bfe-b2be-4542-a4e5-28daab816ec1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "analytics.set_preferred_db('sqlite')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "65ee55a7-5100-40c9-b80c-0ec32d317186",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "
\n",
+       ""
+      ],
+      "text/plain": [
+       "alt.VConcatChart(...)"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "analytics.distribution((ind_rope.loss_train + ind_roper.loss_train)[16_000_000:],\n",
+    "                       color_scheme='blueorange',\n",
+    "                       width=800, height=600, height_minimap=50, levels=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "b5ef8769-c4a6-4e88-8485-57a9a02da497",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "d_rope, _ = analytics.indicator_data((ind_rope.loss_train))\n",
+    "step_rope = np.stack([d[:, 0] for d in d_rope]).mean(axis=0)\n",
+    "d_rope = np.stack([d[:, 5] for d in d_rope]).mean(axis=0)\n",
+    "\n",
+    "d_roper, _ = analytics.indicator_data((ind_roper.loss_train))\n",
+    "step_roper = np.stack([d[:, 0] for d in d_roper]).mean(axis=0)\n",
+    "d_roper = np.stack([d[:, 5] for d in d_roper]).mean(axis=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "e737046d-ffb8-47f0-b742-d7241914061a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "
\n",
+       ""
+      ],
+      "text/plain": [
+       "alt.VConcatChart(...)"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "offset = 50\n",
+    "analytics.distribution([d_rope[offset:], d_roper[offset:]], step_rope[offset:], width=800, height=600, height_minimap=50)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "5dd17505-e567-43a0-8a2f-e092ae954183",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "d_rope, _ = analytics.indicator_data((ind_rope.loss_train))\n",
+    "d_rope = np.stack([d[-5:, 5] for d in d_rope]).mean(axis=1)\n",
+    "\n",
+    "d_roper, _ = analytics.indicator_data((ind_roper.loss_train))\n",
+    "d_roper = np.stack([d[-5:, 5] for d in d_roper]).mean(axis=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "34c8d915-06bd-484f-906c-75abef12afe0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(array([0.3234671 , 0.34104063, 0.31573243, 0.32310195, 0.31422086,\n",
+       "        0.36494735, 0.35428971, 0.31074833, 0.32917885, 0.32994516]),\n",
+       " array([0.32325992, 0.31913995, 0.3186015 , 0.32333733, 0.31854622,\n",
+       "        0.32299762, 0.31425539, 0.32346853, 0.32077246, 0.31096461]))"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "d_rope, d_roper"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f8b2379b-5465-4d4e-bedd-7c9cdde994ae",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/run_roper.sh b/run_roper.sh
new file mode 100755
index 00000000..1127f552
--- /dev/null
+++ b/run_roper.sh
@@ -0,0 +1,7 @@
+PYTHONPATH="${PYTHONPATH}:$(pwd):$(pwd)/src" python labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
+PYTHONPATH="${PYTHONPATH}:$(pwd):$(pwd)/src" python labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
+PYTHONPATH="${PYTHONPATH}:$(pwd):$(pwd)/src" python labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
+PYTHONPATH="${PYTHONPATH}:$(pwd):$(pwd)/src" python labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
+PYTHONPATH="${PYTHONPATH}:$(pwd):$(pwd)/src" python labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
+PYTHONPATH="${PYTHONPATH}:$(pwd):$(pwd)/src" python labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
+PYTHONPATH="${PYTHONPATH}:$(pwd):$(pwd)/src" python labml_nn/transformers/rope/value_pe/experiments/copy_repeat.py
diff --git a/setup.py b/setup.py
index cc48eadc..2c01dc17 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name='labml-nn',
-    version='0.4.122',
+    version='0.4.123',
     author="Varuna Jayasiri, Nipun Wijerathne",
     author_email="vpjayasiri@gmail.com, hnipun@gmail.com",
     description="🧑‍🏫 Implementations/tutorials of deep learning papers with side-by-side notes 📝; including transformers (original, xl, switch, feedback, vit), optimizers (adam, radam, adabelief), gans(dcgan, cyclegan, stylegan2), 🎮 reinforcement learning (ppo, dqn), capsnet, distillation, etc. 🧠",
@@ -20,8 +20,8 @@
                                                'labml_helpers', 'labml_helpers.*',
                                                'test',
                                                'test.*')),
-    install_requires=['labml>=0.4.147',
-                      'labml-helpers>=0.4.84',
+    install_requires=['labml>=0.4.151',
+                      'labml-helpers>=0.4.87',
                       'torch',
                       'torchtext',
                       'torchvision',