feat: add option to use rel pos

Author: flavioschneider

audio_diffusion_pytorch/model.py

@@ -238,6 +238,7 @@ def get_default_model_kwargs():
         attention_heads=8,
         attention_features=64,
         attention_multiplier=2,
+        attention_use_rel_pos=False,
         resnet_groups=8,
         kernel_multiplier_downsample=2,
         use_nearest_upsample=False,

audio_diffusion_pytorch/modules.py

@@ -315,6 +315,55 @@ def __init__(
 """
 
 
+class RelativePositionBias(nn.Module):
+    def __init__(self, num_buckets: int, max_distance: int, num_heads: int):
+        super().__init__()
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.num_heads = num_heads
+        self.relative_attention_bias = nn.Embedding(num_buckets, num_heads)
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position: Tensor, num_buckets: int, max_distance: int
+    ):
+        num_buckets //= 2
+        ret = (relative_position >= 0).to(torch.long) * num_buckets
+        n = torch.abs(relative_position)
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = (
+            max_exact
+            + (
+                torch.log(n.float() / max_exact)
+                / math.log(max_distance / max_exact)
+                * (num_buckets - max_exact)
+            ).long()
+        )
+        val_if_large = torch.min(
+            val_if_large, torch.full_like(val_if_large, num_buckets - 1)
+        )
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def forward(self, num_queries: int, num_keys: int) -> Tensor:
+        i, j, device = num_queries, num_keys, self.relative_attention_bias.weight.device
+        q_pos = torch.arange(j - i, j, dtype=torch.long, device=device)
+        k_pos = torch.arange(j, dtype=torch.long, device=device)
+        rel_pos = rearrange(k_pos, "j -> 1 j") - rearrange(q_pos, "i -> i 1")
+
+        relative_position_bucket = self._relative_position_bucket(
+            rel_pos, num_buckets=self.num_buckets, max_distance=self.max_distance
+        )
+
+        bias = self.relative_attention_bias(relative_position_bucket)
+        bias = rearrange(bias, "m n h -> 1 h m n")
+        return bias
+
+
 def FeedForward(features: int, multiplier: int) -> nn.Module:
     mid_features = features * multiplier
     return nn.Sequential(
@@ -331,19 +380,33 @@ def __init__(
         *,
         head_features: int,
         num_heads: int,
+        use_rel_pos: bool,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
     ):
         super().__init__()
         self.scale = head_features ** -0.5
         self.num_heads = num_heads
+        self.use_rel_pos = use_rel_pos
         mid_features = head_features * num_heads
 
+        if use_rel_pos:
+            assert exists(rel_pos_num_buckets) and exists(rel_pos_max_distance)
+            self.rel_pos = RelativePositionBias(
+                num_buckets=rel_pos_num_buckets,
+                max_distance=rel_pos_max_distance,
+                num_heads=num_heads,
+            )
+
         self.to_out = nn.Linear(in_features=mid_features, out_features=features)
 
     def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
         # Split heads
         q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads)
         # Compute similarity matrix
-        sim = einsum("... n d, ... m d -> ... n m", q, k) * self.scale
+        sim = einsum("... n d, ... m d -> ... n m", q, k)
+        sim = (sim + self.rel_pos(*sim.shape[-2:])) if self.use_rel_pos else sim
+        sim = sim * self.scale
         # Get attention matrix with softmax
         attn = sim.softmax(dim=-1)
         # Compute values
@@ -360,6 +423,9 @@ def __init__(
         head_features: int,
         num_heads: int,
         context_features: Optional[int] = None,
+        use_rel_pos: bool,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
     ):
         super().__init__()
         self.context_features = context_features
@@ -375,7 +441,12 @@ def __init__(
             in_features=context_features, out_features=mid_features * 2, bias=False
         )
         self.attention = AttentionBase(
-            features, num_heads=num_heads, head_features=head_features
+            features,
+            num_heads=num_heads,
+            head_features=head_features,
+            use_rel_pos=use_rel_pos,
+            rel_pos_num_buckets=rel_pos_num_buckets,
+            rel_pos_max_distance=rel_pos_max_distance,
         )
 
     def forward(self, x: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
@@ -402,14 +473,22 @@ def __init__(
         num_heads: int,
         head_features: int,
         multiplier: int,
+        use_rel_pos: bool,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
         context_features: Optional[int] = None,
     ):
         super().__init__()
 
         self.use_cross_attention = exists(context_features) and context_features > 0
 
         self.attention = Attention(
-            features=features, num_heads=num_heads, head_features=head_features
+            features=features,
+            num_heads=num_heads,
+            head_features=head_features,
+            use_rel_pos=use_rel_pos,
+            rel_pos_num_buckets=rel_pos_num_buckets,
+            rel_pos_max_distance=rel_pos_max_distance,
         )
 
         if self.use_cross_attention:
@@ -418,6 +497,9 @@ def __init__(
                 num_heads=num_heads,
                 head_features=head_features,
                 context_features=context_features,
+                use_rel_pos=use_rel_pos,
+                rel_pos_num_buckets=rel_pos_num_buckets,
+                rel_pos_max_distance=rel_pos_max_distance,
             )
 
         self.feed_forward = FeedForward(features=features, multiplier=multiplier)
@@ -443,6 +525,9 @@ def __init__(
         num_heads: int,
         head_features: int,
         multiplier: int,
+        use_rel_pos: bool,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
         context_features: Optional[int] = None,
     ):
         super().__init__()
@@ -465,6 +550,9 @@ def __init__(
                     num_heads=num_heads,
                     multiplier=multiplier,
                     context_features=context_features,
+                    use_rel_pos=use_rel_pos,
+                    rel_pos_num_buckets=rel_pos_num_buckets,
+                    rel_pos_max_distance=rel_pos_max_distance,
                 )
                 for i in range(num_layers)
             ]
@@ -552,6 +640,9 @@ def __init__(
         attention_heads: Optional[int] = None,
         attention_features: Optional[int] = None,
         attention_multiplier: Optional[int] = None,
+        attention_use_rel_pos: Optional[bool] = None,
+        attention_rel_pos_max_distance: Optional[int] = None,
+        attention_rel_pos_num_buckets: Optional[int] = None,
         context_mapping_features: Optional[int] = None,
         context_embedding_features: Optional[int] = None,
     ):
@@ -588,6 +679,7 @@ def __init__(
                 exists(attention_heads)
                 and exists(attention_features)
                 and exists(attention_multiplier)
+                and exists(attention_use_rel_pos)
             )
             self.transformer = Transformer1d(
                 num_layers=num_transformer_blocks,
@@ -596,6 +688,9 @@ def __init__(
                 head_features=attention_features,
                 multiplier=attention_multiplier,
                 context_features=context_embedding_features,
+                use_rel_pos=attention_use_rel_pos,
+                rel_pos_num_buckets=attention_rel_pos_num_buckets,
+                rel_pos_max_distance=attention_rel_pos_max_distance,
             )
 
         if self.use_extract:
@@ -659,6 +754,9 @@ def __init__(
         attention_heads: Optional[int] = None,
         attention_features: Optional[int] = None,
         attention_multiplier: Optional[int] = None,
+        attention_use_rel_pos: Optional[bool] = None,
+        attention_rel_pos_max_distance: Optional[int] = None,
+        attention_rel_pos_num_buckets: Optional[int] = None,
         context_mapping_features: Optional[int] = None,
         context_embedding_features: Optional[int] = None,
     ):
@@ -689,6 +787,7 @@ def __init__(
                 exists(attention_heads)
                 and exists(attention_features)
                 and exists(attention_multiplier)
+                and exists(attention_use_rel_pos)
             )
             self.transformer = Transformer1d(
                 num_layers=num_transformer_blocks,
@@ -697,6 +796,9 @@ def __init__(
                 head_features=attention_features,
                 multiplier=attention_multiplier,
                 context_features=context_embedding_features,
+                use_rel_pos=attention_use_rel_pos,
+                rel_pos_num_buckets=attention_rel_pos_num_buckets,
+                rel_pos_max_distance=attention_rel_pos_max_distance,
             )
 
         self.upsample = Upsample1d(
@@ -756,6 +858,9 @@ def __init__(
         attention_heads: Optional[int] = None,
         attention_features: Optional[int] = None,
         attention_multiplier: Optional[int] = None,
+        attention_use_rel_pos: Optional[bool] = None,
+        attention_rel_pos_max_distance: Optional[int] = None,
+        attention_rel_pos_num_buckets: Optional[int] = None,
         context_mapping_features: Optional[int] = None,
         context_embedding_features: Optional[int] = None,
     ):
@@ -774,6 +879,7 @@ def __init__(
                 exists(attention_heads)
                 and exists(attention_features)
                 and exists(attention_multiplier)
+                and exists(attention_use_rel_pos)
             )
             self.transformer = Transformer1d(
                 num_layers=num_transformer_blocks,
@@ -782,6 +888,9 @@ def __init__(
                 head_features=attention_features,
                 multiplier=attention_multiplier,
                 context_features=context_embedding_features,
+                use_rel_pos=attention_use_rel_pos,
+                rel_pos_num_buckets=attention_rel_pos_num_buckets,
+                rel_pos_max_distance=attention_rel_pos_max_distance,
             )
 
         self.post_block = ResnetBlock1d(
@@ -844,6 +953,9 @@ def __init__(
         context_features: Optional[int] = None,
         context_channels: Optional[Sequence[int]] = None,
         context_embedding_features: Optional[int] = None,
+        attention_use_rel_pos: bool = False,
+        attention_rel_pos_max_distance: Optional[int] = None,
+        attention_rel_pos_num_buckets: Optional[int] = None,
     ):
         super().__init__()
         out_channels = default(out_channels, in_channels)
@@ -931,6 +1043,9 @@ def __init__(
                     attention_heads=attention_heads,
                     attention_features=attention_features,
                     attention_multiplier=attention_multiplier,
+                    attention_use_rel_pos=attention_use_rel_pos,
+                    attention_rel_pos_max_distance=attention_rel_pos_max_distance,
+                    attention_rel_pos_num_buckets=attention_rel_pos_num_buckets,
                 )
                 for i in range(num_layers)
             ]
@@ -945,6 +1060,9 @@ def __init__(
             attention_heads=attention_heads,
             attention_features=attention_features,
             attention_multiplier=attention_multiplier,
+            attention_use_rel_pos=attention_use_rel_pos,
+            attention_rel_pos_max_distance=attention_rel_pos_max_distance,
+            attention_rel_pos_num_buckets=attention_rel_pos_num_buckets,
         )
 
         self.upsamples = nn.ModuleList(
@@ -966,6 +1084,9 @@ def __init__(
                     attention_heads=attention_heads,
                     attention_features=attention_features,
                     attention_multiplier=attention_multiplier,
+                    attention_use_rel_pos=attention_use_rel_pos,
+                    attention_rel_pos_max_distance=attention_rel_pos_max_distance,
+                    attention_rel_pos_num_buckets=attention_rel_pos_num_buckets,
                 )
                 for i in reversed(range(num_layers))
             ]

setup.py

@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.69",
+    version="0.0.70",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",