Add switchable DIFF V1 attention

models.py

@@ -11,6 +11,7 @@ from backbones import (
     TokenAutoDiscretization,
     build_backbone_block,
 )
+from attention_types import resolve_attention_type
 from model_architectures import resolve_model_architecture
 from targets import PAD_IDX
 
@@ -221,6 +222,7 @@ class DeepHealth(nn.Module):
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
         model_architecture: str | None = None,
+        attention_type: str | None = None,
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -242,6 +244,7 @@ class DeepHealth(nn.Module):
         if n_layer < 1:
             raise ValueError(f"n_layer must be >= 1, got {n_layer}")
         model_architecture = resolve_model_architecture(model_architecture)
+        attention_type = resolve_attention_type(attention_type)
         self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
         self.gender_embedding = nn.Embedding(
             2, n_embd)  # Assuming binary gender
@@ -261,6 +264,7 @@ class DeepHealth(nn.Module):
         self.dist_mode = dist_mode
         self.extra_pool_reduce = extra_pool_reduce
         self.model_architecture = model_architecture
+        self.attention_type = attention_type
         self.n_layer = n_layer
         self.n_embd = n_embd
         self.vocab_size = vocab_size
@@ -282,7 +286,9 @@ class DeepHealth(nn.Module):
                     use_time_rope=False,
                     use_rbf_bias=False,
                     mlp_dropout=dropout,
-                ) for _ in range(n_layer)
+                    attention_type=attention_type,
+                    depth=layer_index,
+                ) for layer_index in range(n_layer)
             ])
             self.rope = None
             self.rbf = None
@@ -296,7 +302,9 @@ class DeepHealth(nn.Module):
                     use_time_rope=True,
                     use_rbf_bias=True,
                     mlp_dropout=dropout,
-                ) for _ in range(n_layer)
+                    attention_type=attention_type,
+                    depth=layer_index,
+                ) for layer_index in range(n_layer)
             ])
             self.rope = TimeRoPE(n_embd // n_head)
             self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)