How valentyn1997/causaltransformer works
Purpose and execution
CausalTransformer is a Python research-training repository for estimating longitudinal counterfactual outcomes under time-varying treatment confounding. Its executable entry points are model-family-specific runners: runnables/train_multi.py, runnables/train_enc_dec.py, runnables/train_gnet.py, runnables/train_msm.py, and runnables/train_rmsn.py.
Configuration-driven experiment assembly
Hydra composes a root experiment configuration with a selected dataset and backbone, including saved hyperparameter variants. Central configuration is config/config.yaml; dataset selections are in config/dataset/cancer_sim.yaml, config/dataset/mimic3_synthetic.yaml, and config/dataset/mimic3_real.yaml; backbone defaults are under config/backbone/. This is the main control boundary between runners and implementations.
Longitudinal dataset layer
A shared collection abstraction coordinates cohort splits and sequence-oriented data consumed by models: src/data/dataset_collection.py. Data originates from:
• Synthetic tumor-growth trajectories: src/data/cancer_sim/cancer_simulation.py, src/data/cancer_sim/dataset.py.
• MIMIC-III semi-synthetic trajectories: src/data/mimic_iii/semi_synthetic_dataset.py.
• MIMIC-III real-world outcome data: src/data/mimic_iii/real_dataset.py.
MIMIC source extraction is loaded from local processed data through src/data/mimic_iii/load_data.py; expected external input is data/processed/all_hourly_data.h5.
Core temporal-model contract
src/models/time_varying_model.py is the likely shared abstraction for temporal causal prediction, allowing runners to train and evaluate different estimators against a common longitudinal-data interface. Shared model operations and neural primitives sit in src/models/utils.py, src/models/utils_lstm.py, and src/models/utils_transformer.py, separating estimator-specific causal logic from recurrent/Transformer mechanics.
Transformer causal estimators
The repository’s primary contribution is implemented in src/models/ct.py (Causal Transformer) and src/models/edct.py (encoder-decoder Causal Transformer). These models predict future outcomes from historical covariates and treatments; configurations support balancing through counterfactual domain confusion or gradient reversal. Architecture and training variants are selected via config/backbone/ct.yaml and config/backbone/edct.yaml.
Baseline estimator families
Comparable causal baselines are independently implemented: Counterfactual Recurrent Network (src/models/crn.py), G-Net (src/models/gnet.py), Marginal Structural Models (src/models/msm.py), and Recurrent Marginal Structural Networks (src/models/rmsn.py). Runner choice is meaningful: encoder-decoder training serves CT/EDCT/CRN, while G-Net, MSM, and RMSN have dedicated workflows. Backbone declarations include config/backbone/crn.yaml, config/backbone/gnet.yaml, and config/backbone/rmsn.yaml.
Runtime, experiment infrastructure, and outputs
The material runtime is Python with PyTorch Lightning for training lifecycle/device execution, Hydra for CLI configuration and multiruns, and MLflow for experiment tracking. GPU selection is supplied externally through CUDA_VISIBLE_DEVICES; MLflow requires an independently running tracking server. Dependencies are declared in requirements.txt, packaging metadata in setup.py, and processed external datasets are intentionally excluded except for the placeholder data/processed/.gitkeep.
Main entry points: time_varying_model.py Temporal model contract · train_gnet.py Baseline runners · ct.yaml Backbone configs · config.yaml Root Hydra config
Components 20
Experiment assembly
Hydra selects dataset and estimator variants.
- Multi-model runner · Entry point for multi-model training.
runnables/train_multi.py - Encoder-decoder runner · Trains CT, EDCT, and CRN workflows.
runnables/train_enc_dec.py - Baseline runners · Dedicated G-Net, MSM, and RMSN entry points.
runnables/train_gnet.py - Root Hydra config · Composes experiment, dataset, and backbone selections.
config/config.yaml - Dataset configs · Cancer and MIMIC-III dataset selections.
config/dataset/cancer_sim.yaml - Backbone configs · CT, EDCT, CRN, G-Net, and RMSN variants.
config/backbone/ct.yaml
Longitudinal data
Shared cohort and trajectory data layer.
- Dataset collection · Coordinates cohort splits and model-ready sequences.
src/data/dataset_collection.py - Cancer simulation · Generates tumor-growth trajectories and dataset records.
src/data/cancer_sim/cancer_simulation.py - MIMIC semi-synthetic · Builds semi-synthetic longitudinal trajectories.
src/data/mimic_iii/semi_synthetic_dataset.py - MIMIC real outcomes · Provides real-world outcome data.
src/data/mimic_iii/real_dataset.py - MIMIC source loader · Loads locally processed hourly MIMIC data.
src/data/mimic_iii/load_data.py - Processed MIMIC data · Externally supplied processed source data.
Causal estimators
Common temporal contract and estimator families.
- Temporal model contract · Shared interface for longitudinal causal prediction.
src/models/time_varying_model.py - Model primitives · Shared operations plus LSTM and Transformer mechanics.
src/models/utils.py - CT and EDCT · Transformer estimators for historical treatment and outcome prediction.
src/models/ct.py - CRN baseline
src/models/crn.py - G-Net baseline
src/models/gnet.py - MSM baseline
src/models/msm.py - RMSN baseline
src/models/rmsn.py
Training runtime
Execution, tracking, and external inputs.
- Lightning, MLflow, GPU · PyTorch Lightning executes training; MLflow records experiments.
Connections 32
- Root Hydra config → Dataset configs: selects
- Root Hydra config → Backbone configs: selects
- Multi-model runner → Root Hydra config: composes
- Encoder-decoder runner → Root Hydra config: composes
- Baseline runners → Root Hydra config: composes
- Dataset configs → Dataset collection: builds
- Cancer simulation → Dataset collection: supplies trajectories
- MIMIC semi-synthetic → Dataset collection: supplies trajectories
- MIMIC real outcomes → Dataset collection: supplies outcomes
- Processed MIMIC data → MIMIC source loader: loaded from
- MIMIC source loader → MIMIC semi-synthetic: provides source data
- MIMIC source loader → MIMIC real outcomes: provides source data
- Dataset collection → Temporal model contract: provides sequences
- Temporal model contract → CT and EDCT: implemented by
- Temporal model contract → CRN baseline: implemented by
- Temporal model contract → G-Net baseline: implemented by
- Temporal model contract → MSM baseline: implemented by
- Temporal model contract → RMSN baseline: implemented by
- Model primitives → CT and EDCT: supports
- Model primitives → CRN baseline: supports
- Backbone configs → CT and EDCT: configures
- Backbone configs → CRN baseline: configures
- Backbone configs → G-Net baseline: configures
- Backbone configs → RMSN baseline: configures
- Encoder-decoder runner → CT and EDCT: trains
- Encoder-decoder runner → CRN baseline: trains
- Baseline runners → G-Net baseline: trains
- Baseline runners → MSM baseline: trains
- Baseline runners → RMSN baseline: trains
- Multi-model runner → Lightning, MLflow, GPU: executes on
- Encoder-decoder runner → Lightning, MLflow, GPU: executes on
- Baseline runners → Lightning, MLflow, GPU: executes on