update

toolbox/torchaudio/models/nx_clean_unet/transformer/attention.py

@@ -0,0 +1,250 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+
+class MultiHeadedAttention(nn.Module):
+    def __init__(self, n_head: int, n_feat: int, dropout_rate: float):
+        """
+        :param n_head: int. the number of heads.
+        :param n_feat: int. the number of features.
+        :param dropout_rate: float. dropout rate.
+        """
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self,
+                    query: torch.Tensor,
+                    key: torch.Tensor,
+                    value: torch.Tensor
+                    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        transform query, key and value.
+        :param query: torch.Tensor. query tensor. shape=(batch_size, time1, n_feat).
+        :param key: torch.Tensor. key tensor. shape=(batch_size, time2, n_feat).
+        :param value: torch.Tensor. value tensor. shape=(batch_size, time2, n_feat).
+        :return:
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(self,
+                          value: torch.Tensor,
+                          scores: torch.Tensor,
+                          mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+                          ) -> torch.Tensor:
+        """
+        compute attention context vector.
+        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, time2, d_k).
+        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, time1, time2).
+        :param mask: torch.Tensor. mask. shape=(batch_size, 1, time2) or
+                (batch_size, time1, time2), (0, 0, 0) means fake mask.
+        :return: torch.Tensor. transformed value. (batch_size, time1, d_model).
+                weighted by the attention score (batch_size, time1, time2).
+        """
+        n_batch = value.size(0)
+        # NOTE: When will `if mask.size(2) > 0` be True?
+        #   1. onnx(16/4) [WHY? Because we feed real cache & real mask for the
+        #           1st chunk to ease the onnx export.]
+        #   2. pytorch training
+        if mask.size(2) > 0:  # time2 > 0
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            # For last chunk, time2 might be larger than scores.size(-1)
+            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
+
+        # NOTE: When will `if mask.size(2) > 0` be False?
+        #   1. onnx(16/-1, -1/-1, 16/0)
+        #   2. jit (16/-1, -1/-1, 16/0, 16/4)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
+
+        return self.linear_out(x)  # (batch, time1, n_feat)
+
+    def forward(self,
+                query: torch.Tensor,
+                key: torch.Tensor,
+                value: torch.Tensor,
+                mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+                cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        q, k, v = self.forward_qkv(query, key, value)
+
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(
+                cache, cache.size(-1) // 2, dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE: We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+        new_cache = torch.cat((k, v), dim=-1)
+
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+class RelativeMultiHeadedAttention(nn.Module):
+
+    def __init__(self, n_head: int, n_feat: int, dropout_rate: float, max_relative_position: int = 5120):
+        """
+        :param n_head: int. the number of heads.
+        :param n_feat: int. the number of features.
+        :param dropout_rate: float. dropout rate.
+        :param max_relative_position: int. maximum relative position for relative position encoding.
+        """
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+        # Relative position encoding
+        self.max_relative_position = max_relative_position
+        self.relative_position_k = nn.Parameter(torch.randn(max_relative_position * 2 + 1, self.d_k))
+
+    def forward_qkv(self,
+                    query: torch.Tensor,
+                    key: torch.Tensor,
+                    value: torch.Tensor
+                    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        transform query, key and value.
+        :param query: torch.Tensor. query tensor. shape=(batch_size, time1, n_feat).
+        :param key: torch.Tensor. key tensor. shape=(batch_size, time2, n_feat).
+        :param value: torch.Tensor. value tensor. shape=(batch_size, time2, n_feat).
+        :return:
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(self,
+                          value: torch.Tensor,
+                          scores: torch.Tensor,
+                          mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+                          ) -> torch.Tensor:
+        """
+        compute attention context vector.
+        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, time2, d_k).
+        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, time1, time2).
+        :param mask: torch.Tensor. mask. shape=(batch_size, 1, time2) or
+                (batch_size, time1, time2), (0, 0, 0) means fake mask.
+        :return: torch.Tensor. transformed value. (batch_size, time1, d_model).
+                weighted by the attention score (batch_size, time1, time2).
+        """
+        n_batch = value.size(0)
+        if mask.size(2) > 0:  # time2 > 0
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            # For last chunk, time2 might be larger than scores.size(-1)
+            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
+
+        return self.linear_out(x)  # (batch, time1, n_feat)
+
+    def relative_position_encoding(self, length: int) -> torch.Tensor:
+        """
+        Generate relative position encoding.
+        :param length: int. length of the sequence.
+        :return: torch.Tensor. relative position encoding. shape=(length, length, d_k).
+        """
+        range_vec = torch.arange(length)
+        distance_mat = range_vec.unsqueeze(0) - range_vec.unsqueeze(1)
+        distance_mat_clipped = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position)
+        final_mat = distance_mat_clipped + self.max_relative_position
+        return final_mat
+
+    def forward(self,
+                query: torch.Tensor,
+                key: torch.Tensor,
+                value: torch.Tensor,
+                mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+                cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        q, k, v = self.forward_qkv(query, key, value)
+
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(
+                cache, cache.size(-1) // 2, dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE: We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+
+        # new_cache shape: [batch_size, self.h, time_steps, self.d_v * 2]
+        new_cache = torch.cat((k, v), dim=-1)
+
+        # Compute relative position encoding
+        length = q.size(2)
+        relative_position = self.relative_position_encoding(length)
+        relative_position_k = self.relative_position_k[relative_position.view(-1)].view(length, length, -1)
+
+        relative_position_k = relative_position_k.unsqueeze(0).unsqueeze(0)  # (1, 1, length, length, d_k)
+        relative_position_k = relative_position_k.expand(q.size(0), q.size(1), -1, -1, -1)  # (batch, head, length, length, d_k)
+
+        native_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        relative_position_scores = torch.matmul(q.unsqueeze(3), relative_position_k.transpose(-2, -1)).squeeze(3) / math.sqrt(self.d_k)
+        scores = native_scores + relative_position_scores
+
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+def main():
+    rel_attention = RelativeMultiHeadedAttention(n_head=4, n_feat=256, dropout_rate=0.1)
+
+    # x = torch.ones(size=(1, 200, 256), dtype=torch.float32)
+
+    x = torch.ones(size=(1, 1, 256), dtype=torch.float32)
+    cache = torch.ones(size=(1, 4, 199, 128), dtype=torch.float32)
+
+    xt, new_cache = rel_attention.forward(x, x, x, cache=cache)
+    print(xt.shape)
+    print(new_cache.shape)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_clean_unet/transformer/embedding.py

@@ -1,95 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import torch
-import torch.nn as nn
-
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-class RelativeMultiheadAttention(nn.Module):
-    def __init__(self, d_model, num_heads, max_len, dropout=0.1):
-        super(RelativeMultiheadAttention, self).__init__()
-        self.num_heads = num_heads
-        self.d_model = d_model
-        self.head_dim = d_model // num_heads
-        self.scale = self.head_dim ** -0.5
-
-        self.query_projection = nn.Linear(d_model, d_model)
-        self.key_projection = nn.Linear(d_model, d_model)
-        self.value_projection = nn.Linear(d_model, d_model)
-        self.output_projection = nn.Linear(d_model, d_model)
-
-        self.dropout = nn.Dropout(dropout)
-
-        # Relative position encoding
-        self.relative_positions_encoding = self.generate_relative_positions_encoding(max_len, self.head_dim)
-
-    def generate_relative_positions_encoding(self, max_len, head_dim):
-        # Generate relative positions encoding matrix
-        even_index = torch.arange(max_len)[:, None] // torch.pow(10000, torch.arange(0, head_dim, 2) / head_dim)
-        odd_index = torch.arange(max_len)[:, None] // torch.pow(10000, torch.arange(1, head_dim, 2) / head_dim)
-        even_index = torch.sin(even_index)
-        odd_index = torch.cos(odd_index)
-        pos_encoding = torch.zeros(max_len, head_dim)
-        pos_encoding[:, 0::2] = even_index
-        pos_encoding[:, 1::2] = odd_index
-        return pos_encoding
-
-    def forward(self, query, key, value, mask=None):
-        batch_size = query.size(0)
-        query_len = query.size(1)
-        key_len = key.size(1)
-
-        # Project queries, keys, and values to multiple heads
-        query = self.query_projection(query).view(batch_size, query_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key = self.key_projection(key).view(batch_size, key_len, self.num_heads, self.head_dim).transpose(1, 2)
-        value = self.value_projection(value).view(batch_size, key_len, self.num_heads, self.head_dim).transpose(1, 2)
-
-        # Apply relative position encoding
-        relative_keys = self.relative_positions_encoding[:query_len, :].unsqueeze(0).unsqueeze(0).repeat(batch_size, self.num_heads, 1, 1)
-        relative_values = self.relative_positions_encoding[:query_len, :].unsqueeze(0).unsqueeze(0).repeat(batch_size, self.num_heads, 1, 1)
-
-        # Compute attention scores
-        scores = torch.matmul(query, key.transpose(-2, -1)) * self.scale
-        scores += torch.matmul(query, relative_keys.transpose(-2, -1))
-
-        if mask is not None:
-            scores = scores.masked_fill(mask == 0, float('-inf'))
-
-        attn_weights = F.softmax(scores, dim=-1)
-        attn_weights = self.dropout(attn_weights)
-
-        # Apply attention weights to values
-        output = torch.matmul(attn_weights, value) + torch.matmul(attn_weights, relative_values)
-        output = output.transpose(1, 2).contiguous().view(batch_size, query_len, self.d_model)
-
-        # Apply output projection
-        output = self.output_projection(output)
-
-        return output
-
-
-def main():
-    # Example usage
-    batch_size = 2
-    query_len = 10
-    key_len = 10
-    d_model = 512
-    num_heads = 8
-    max_len = 100
-
-    query = torch.rand(batch_size, query_len, d_model)
-    key = torch.rand(batch_size, key_len, d_model)
-    value = torch.rand(batch_size, key_len, d_model)
-
-    attention = RelativeMultiheadAttention(d_model, num_heads, max_len)
-    output = attention(query, key, value)
-    print(output.shape)  # Output shape should be (batch_size, query_len, d_model)
-
-    return
-
-
-if __name__ == '__main__':
-    main()

toolbox/torchaudio/models/nx_clean_unet/transformer/transformer.py

@@ -5,178 +5,9 @@ from typing import Dict, Optional, Tuple, List, Union
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as f
 
 from toolbox.torchaudio.models.nx_clean_unet.transformer.mask import subsequent_chunk_mask
-
-
-class SinusoidalPositionalEncoding(nn.Module):
-    """
-    Positional Encoding
-
-    PE(pos, 2i)   = sin(pos/(10000^(2i/dmodel)))
-    PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel)))
-    """
-
-    @staticmethod
-    def demo1():
-        batch_size = 2
-        time_steps = 10
-        embedding_dim = 64
-
-        pe = SinusoidalPositionalEncoding(
-            embedding_dim=embedding_dim,
-            dropout_rate=0.1,
-        )
-
-        x = torch.randn(size=(batch_size, time_steps, embedding_dim))
-
-        x, pos_emb = pe.forward(x)
-
-        # torch.Size([2, 10, 64])
-        print(x.shape)
-        # torch.Size([1, 10, 64])
-        print(pos_emb.shape)
-        return
-
-    @staticmethod
-    def demo2():
-        batch_size = 2
-        time_steps = 10
-        embedding_dim = 64
-
-        pe = SinusoidalPositionalEncoding(
-            embedding_dim=embedding_dim,
-            dropout_rate=0.1,
-        )
-
-        x = torch.randn(size=(batch_size, time_steps, embedding_dim))
-        offset = torch.randint(low=3, high=7, size=(batch_size,))
-        x, pos_emb = pe.forward(x, offset=offset)
-
-        # tensor([3, 4])
-        print(offset)
-        # torch.Size([2, 10, 64])
-        print(x.shape)
-        # torch.Size([2, 10, 64])
-        print(pos_emb.shape)
-        return
-
-    def __init__(self,
-                 embedding_dim: int,
-                 dropout_rate: float,
-                 max_length: int = 5000,
-                 reverse: bool = False
-                 ):
-        super().__init__()
-        self.embedding_dim = embedding_dim
-        self.dropout_rate = dropout_rate
-        self.max_length = max_length
-        self.reverse = reverse
-
-        self.x_scale = math.sqrt(self.embedding_dim)
-        self.dropout = torch.nn.Dropout(p=dropout_rate)
-
-        self.pe = torch.zeros(self.max_length, self.embedding_dim)
-        position = torch.arange(0, self.max_length, dtype=torch.float32).unsqueeze(1)
-
-        div_term = torch.exp(
-            torch.arange(0, self.embedding_dim, 2, dtype=torch.float32) *
-            - (math.log(10000.0) / self.embedding_dim)
-        )
-        self.pe[:, 0::2] = torch.sin(position * div_term)
-        self.pe[:, 1::2] = torch.cos(position * div_term)
-        self.pe = self.pe.unsqueeze(0)
-
-    def forward(self,
-                x: torch.Tensor,
-                offset: Union[int, torch.Tensor] = 0
-                ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Add positional encoding.
-        :param x: torch.Tensor. Input. shape=(batch_size, time_steps, ...).
-        :param offset: int or torch.Tensor. position offset.
-        :return:
-        torch.Tensor. Encoded tensor. shape=(batch_size, time_steps, ...).
-        torch.Tensor. for compatibility to RelPositionalEncoding. shape=(1, time_steps, ...).
-        """
-        self.pe = self.pe.to(x.device)
-        pos_emb = self.position_encoding(offset, x.size(1), False)
-        x = x * self.x_scale + pos_emb
-        return self.dropout(x), self.dropout(pos_emb)
-
-    def position_encoding(self,
-                          offset: Union[int, torch.Tensor],
-                          size: int,
-                          apply_dropout: bool = True
-                          ) -> torch.Tensor:
-        """
-        For getting encoding in a streaming fashion.
-
-        Attention!!!!!
-        we apply dropout only once at the whole utterance level in a none
-        streaming way, but will call this function several times with
-        increasing input size in a streaming scenario, so the dropout will
-        be applied several times.
-
-        :param offset: int or torch.Tensor. start offset.
-        :param size: int. required size of position encoding.
-        :param apply_dropout:
-        :return: torch.Tensor. Corresponding encoding.
-        """
-        if isinstance(offset, int):
-            assert offset + size <= self.max_length
-            pos_emb = self.pe[:, offset:offset + size]
-        elif isinstance(offset, torch.Tensor) and offset.dim() == 0:  # scalar
-            assert offset + size <= self.max_length
-            pos_emb = self.pe[:, offset:offset + size]
-        else:    # for batched streaming decoding on GPU
-            # offset. shape=(batch_size,)
-            assert torch.max(offset) + size <= self.max_length
-
-            # shape=(batch_size, time_steps)
-            index = offset.unsqueeze(1) + torch.arange(0, size).to(offset.device)
-            flag = index > 0
-            # remove negative offset
-            index = index * flag
-            # shape=(batch_size, time_steps, embedding_dim)
-            pos_emb = f.embedding(index, self.pe[0])
-
-        if apply_dropout:
-            pos_emb = self.dropout(pos_emb)
-        return pos_emb
-
-
-class RelPositionalEncoding(SinusoidalPositionalEncoding):
-    """
-    Relative positional encoding module.
-
-    See : Appendix B in https://arxiv.org/abs/1901.02860
-
-    """
-    def __init__(self,
-                 embedding_dim: int,
-                 dropout_rate: float,
-                 max_length: int = 5000,
-                 ):
-        super().__init__(embedding_dim, dropout_rate, max_length, reverse=True)
-
-    def forward(self,
-                x: torch.Tensor,
-                offset: Union[int, torch.Tensor] = 0
-                ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Compute positional encoding.
-        :param x: torch.Tensor. Input. shape=(batch_size, time_steps, ...).
-        :param offset:
-        :return:
-        torch.Tensor. Encoded tensor. shape=(batch_size, time_steps, ...).
-        torch.Tensor. Positional embedding tensor. shape=(1, time_steps, ...).
-        """
-        self.pe = self.pe.to(x.device)
-        x = x * self.x_scale
-        pos_emb = self.position_encoding(offset, x.size(1), False)
-        return self.dropout(x), self.dropout(pos_emb)
+from toolbox.torchaudio.models.nx_clean_unet.transformer.attention import MultiHeadedAttention, RelativeMultiHeadedAttention
 
 
 class PositionwiseFeedForward(nn.Module):
@@ -209,151 +40,20 @@ class PositionwiseFeedForward(nn.Module):
         return self.w_2(self.dropout(self.activation(self.w_1(xs))))
 
 
-class MultiHeadedAttention(nn.Module):
-    def __init__(self, n_head: int, n_feat: int, dropout_rate: float):
-        """
-        :param n_head: int. the number of heads.
-        :param n_feat: int. the number of features.
-        :param dropout_rate: float. dropout rate.
-        """
-        super().__init__()
-        assert n_feat % n_head == 0
-        # We assume d_v always equals d_k
-        self.d_k = n_feat // n_head
-        self.h = n_head
-        self.linear_q = nn.Linear(n_feat, n_feat)
-        self.linear_k = nn.Linear(n_feat, n_feat)
-        self.linear_v = nn.Linear(n_feat, n_feat)
-        self.linear_out = nn.Linear(n_feat, n_feat)
-        self.dropout = nn.Dropout(p=dropout_rate)
-
-    def forward_qkv(self,
-                    query: torch.Tensor,
-                    key: torch.Tensor,
-                    value: torch.Tensor
-                    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        transform query, key and value.
-        :param query: torch.Tensor. query tensor. shape=(batch_size, time1, n_feat).
-        :param key: torch.Tensor. key tensor. shape=(batch_size, time2, n_feat).
-        :param value: torch.Tensor. value tensor. shape=(batch_size, time2, n_feat).
-        :return:
-        """
-        n_batch = query.size(0)
-        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
-        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
-        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
-        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
-        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
-        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
-
-        return q, k, v
-
-    def forward_attention(self,
-                          value: torch.Tensor,
-                          scores: torch.Tensor,
-                          mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
-                          ) -> torch.Tensor:
-        """
-        compute attention context vector.
-        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, time2, d_k).
-        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, time1, time2).
-        :param mask: torch.Tensor. mask. shape=(batch_size, 1, time2) or
-                (batch_size, time1, time2), (0, 0, 0) means fake mask.
-        :return: torch.Tensor. transformed value. (batch_size, time1, d_model).
-                weighted by the attention score (batch_size, time1, time2).
-        """
-        n_batch = value.size(0)
-        # NOTE: When will `if mask.size(2) > 0` be True?
-        #   1. onnx(16/4) [WHY? Because we feed real cache & real mask for the
-        #           1st chunk to ease the onnx export.]
-        #   2. pytorch training
-        if mask.size(2) > 0:  # time2 > 0
-            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
-            # For last chunk, time2 might be larger than scores.size(-1)
-            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
-            scores = scores.masked_fill(mask, -float('inf'))
-            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
-
-        # NOTE: When will `if mask.size(2) > 0` be False?
-        #   1. onnx(16/-1, -1/-1, 16/0)
-        #   2. jit (16/-1, -1/-1, 16/0, 16/4)
-        else:
-            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
-
-        p_attn = self.dropout(attn)
-        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
-        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
-
-        return self.linear_out(x)  # (batch, time1, n_feat)
-
-    def forward(self,
-                query: torch.Tensor,
-                key: torch.Tensor,
-                value: torch.Tensor,
-                mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
-                cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
-                **kwargs,
-                ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        compute scaled dot product attention.
-        :param query: torch.Tensor. query tensor. shape=(batch_size, time1, n_feat).
-        :param key: torch.Tensor. key tensor. shape=(batch_size, time2, n_feat).
-        :param value: torch.Tensor. value tensor. shape=(batch_size, time2, n_feat).
-        :param mask: torch.Tensor. mask tensor (batch_size, 1, time2) or
-                (batch_size, time1, time2).
-        :param cache: torch.Tensor. cache tensor. shape=(1, head, cache_t, d_k * 2),
-                where `cache_t == chunk_size * num_decoding_left_chunks`
-                and `head * d_k == n_feat`
-        :return:
-        torch.Tensor. output tensor. shape=(batch_size, time1, n_feat).
-        torch.Tensor. cache tensor. (1, head, cache_t + time1, d_k * 2)
-                where `cache_t == chunk_size * num_decoding_left_chunks`
-                and `head * d_k == n_feat`
-        """
-        q, k, v = self.forward_qkv(query, key, value)
-
-        # NOTE:
-        #   when export onnx model, for 1st chunk, we feed
-        #       cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
-        #       or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
-        #       In all modes, `if cache.size(0) > 0` will alwayse be `True`
-        #       and we will always do splitting and
-        #       concatnation(this will simplify onnx export). Note that
-        #       it's OK to concat & split zero-shaped tensors(see code below).
-        #   when export jit  model, for 1st chunk, we always feed
-        #       cache(0, 0, 0, 0) since jit supports dynamic if-branch.
-        # >>> a = torch.ones((1, 2, 0, 4))
-        # >>> b = torch.ones((1, 2, 3, 4))
-        # >>> c = torch.cat((a, b), dim=2)
-        # >>> torch.equal(b, c)        # True
-        # >>> d = torch.split(a, 2, dim=-1)
-        # >>> torch.equal(d[0], d[1])  # True
-        if cache.size(0) > 0:
-            key_cache, value_cache = torch.split(
-                cache, cache.size(-1) // 2, dim=-1)
-            k = torch.cat([key_cache, k], dim=2)
-            v = torch.cat([value_cache, v], dim=2)
-        # NOTE: We do cache slicing in encoder.forward_chunk, since it's
-        #   non-trivial to calculate `next_cache_start` here.
-        new_cache = torch.cat((k, v), dim=-1)
-
-        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
-        return self.forward_attention(v, scores, mask), new_cache
-
-
 class TransformerEncoderLayer(nn.Module):
     def __init__(self,
                  input_dim: int,
                  dropout_rate: float = 0.1,
                  n_heads: int = 4,
+                 max_relative_position: int = 5120
                  ):
         super().__init__()
         self.norm1 = nn.LayerNorm(input_dim, eps=1e-5)
-        self.attention = MultiHeadedAttention(
+        self.attention = RelativeMultiHeadedAttention(
             n_head=n_heads,
             n_feat=input_dim,
-            dropout_rate=dropout_rate
+            dropout_rate=dropout_rate,
+            max_relative_position=max_relative_position,
         )
 
         self.dropout1 = nn.Dropout(dropout_rate)
@@ -370,7 +70,6 @@ class TransformerEncoderLayer(nn.Module):
             self,
             x: torch.Tensor,
             mask: torch.Tensor,
-            position_embedding: torch.Tensor,
             attention_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
@@ -388,7 +87,7 @@ class TransformerEncoderLayer(nn.Module):
         xt = self.norm1(x)
 
         x_att, new_att_cache = self.attention.forward(
-            xt, xt, xt, mask=mask, cache=attention_cache, position_embedding=position_embedding
+            xt, xt, xt, mask=mask, cache=attention_cache
         )
         x = x + self.dropout1(xt)
         xt = self.norm2(x)
@@ -410,7 +109,7 @@ class TransformerEncoder(nn.Module):
                  attention_heads: int = 4,
                  num_blocks: int = 6,
                  dropout_rate: float = 0.1,
-                 max_length: int = 1024,
+                 max_relative_position: int = 1024,
                  chunk_size: int = 1,
                  num_left_chunks: int = 128,
                  ):
@@ -418,7 +117,7 @@ class TransformerEncoder(nn.Module):
         self.input_size = input_size
         self.hidden_size = hidden_size
 
-        self.max_length = max_length
+        self.max_relative_position = max_relative_position
         self.chunk_size = chunk_size
         self.num_left_chunks = num_left_chunks
 
@@ -427,17 +126,12 @@ class TransformerEncoder(nn.Module):
             out_features=self.hidden_size,
         )
 
-        self.positional_encoding = RelPositionalEncoding(
-            embedding_dim=hidden_size,
-            dropout_rate=dropout_rate,
-            max_length=max_length,
-        )
-
         self.encoder_layer_list = torch.nn.ModuleList([
             TransformerEncoderLayer(
                 input_dim=hidden_size,
                 n_heads=attention_heads,
                 dropout_rate=dropout_rate,
+                max_relative_position=max_relative_position,
             ) for _ in range(num_blocks)
         ])
 
@@ -458,10 +152,6 @@ class TransformerEncoder(nn.Module):
         xs = self.input_linear.forward(xs)
         # xs shape: [batch_size, time_steps, hidden_size]
 
-        xs, position_embedding = self.positional_encoding.forward(xs)
-        # xs shape: [batch_size, time_steps, hidden_size]
-        # position_embedding shape: [1, time_steps, hidden_size]
-
         chunk_masks = subsequent_chunk_mask(
             size=time_steps,
             chunk_size=self.chunk_size,
@@ -473,7 +163,7 @@ class TransformerEncoder(nn.Module):
         # chunk_masks shape: [batch_size, time_steps, time_steps]
 
         for encoder_layer in self.encoder_layer_list:
-            xs, _ = encoder_layer.forward(xs, chunk_masks, position_embedding)
+            xs, _ = encoder_layer.forward(xs, chunk_masks)
 
         # xs shape: [batch_size, time_steps, hidden_size]
         xs = self.output_linear.forward(xs)