Revert "Add switchable DIFF V1 attention"

This reverts commit bf5cae8758.

train_next_step.py

@@ -16,7 +16,6 @@ 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 (
@@ -89,12 +88,6 @@ 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)
@@ -158,7 +151,6 @@ def build_model(
         dist_mode="exponential",
         dropout=args.dropout,
         model_architecture=args.model_architecture,
-        attention_type=args.attention_type,
     )
 
 
@@ -381,7 +373,6 @@ 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",
@@ -434,7 +425,6 @@ 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")