Files
DeepHealth/MODEL_ARCHITECTURES.md

2.1 KiB

Model architectures

DeepHealth uses one codebase for both supported history-block architectures. Select the architecture explicitly when starting a training run:

model_architecture History block Checkpoint fingerprint
transformer_ffn_v1 Temporal attention + SwiGLU FFN blocks.*.mlp.w1/w2/w3 and blocks.*.ln2
traj_mixer_v5 Temporal attention + TrajMixer blocks.*.mlp.intra_*, gate_proj, and output_proj

transformer_ffn_v1 is the CLI default; pass traj_mixer_v5 explicitly for TrajMixer runs.

Training

Next-step example:

python train_next_step.py --model_architecture traj_mixer_v5 --n_layer 12

All-future example:

python train_all_future.py --model_architecture transformer_ffn_v1 --n_layer 12

New runs are separated by architecture:

runs/
  transformer_ffn_v1/
    <run_name>/
  traj_mixer_v5/
    <run_name>/

Use --runs_root to place this structure under a different root. Existing run directories are not moved or renamed.

Each generated train_config.json records model_architecture, total parameter count, and trainable parameter count.

Both training entry points use the single --n_layer option to set the number of history backbone blocks. The same value is passed to DeepHealth.n_layer and saved as n_layer in train_config.json; it must be at least 1.

Architecture validation

Evaluation resolves the architecture before constructing the model and always loads weights with strict=True.

  • Every config must include an explicit model_architecture marker.
  • Checkpoint fingerprints are used to validate that the selected architecture matches the stored weights.
  • A config marker that conflicts with the checkpoint fingerprint raises an error instead of silently choosing one architecture.
  • Unsupported historical TrajMixer markers such as traj_mixer_v2, traj_mixer_v3, and traj_mixer_v4 are rejected.
  • Checkpoints and configs created before architecture markers were introduced are intentionally unsupported.

Project code should use the architecture factory rather than instantiate a history block directly.