Add switchable DIFF V1 attention

train_next_step.py

@@ -16,6 +16,7 @@ from torch.optim import AdamW
 from torch.utils.data import DataLoader, RandomSampler
 from tqdm.auto import tqdm
 
+from attention_types import DEFAULT_ATTENTION_TYPE, SUPPORTED_ATTENTION_TYPES
 from dataset import HealthDataset, collate_fn
 from losses import build_loss
 from model_architectures import (
@@ -88,6 +89,12 @@ def parse_args() -> argparse.Namespace:
         default=DEFAULT_MODEL_ARCHITECTURE,
         choices=SUPPORTED_MODEL_ARCHITECTURES,
     )
+    parser.add_argument(
+        "--attention_type",
+        type=str,
+        default=DEFAULT_ATTENTION_TYPE,
+        choices=SUPPORTED_ATTENTION_TYPES,
+    )
 
     parser.add_argument("--t_min", type=float, default=0.0027378507871321013)
     parser.add_argument("--max_exp_input", type=float, default=60.0)
@@ -151,6 +158,7 @@ def build_model(
         dist_mode="exponential",
         dropout=args.dropout,
         model_architecture=args.model_architecture,
+        attention_type=args.attention_type,
     )
 
 
@@ -373,6 +381,7 @@ def build_metadata(
         "collate_fn": "next_step_collate_fn",
         "model_class": "DeepHealth",
         "model_architecture": args.model_architecture,
+        "attention_type": args.attention_type,
         "model_target_mode": "next_token",
         "target_mode": "delphi2m",
         "event_stream_version": "disease_death_only_v1",
@@ -425,6 +434,7 @@ def main() -> None:
     logger.info(f"Starting next-step training run: {run_name}")
     logger.info(f"Device: {device}")
     logger.info(f"Model architecture: {args.model_architecture}")
+    logger.info(f"Attention type: {args.attention_type}")
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
     logger.info("Continuous value scaling: RobustScale (required)")
     logger.info("time_mode=absolute, readout=token, target_mode=delphi2m")