Amirhossein-Zare/causal-long-pfn

CausalLongPFN pretrained weights

This repository contains inference-only pretrained weights for CausalLongPFN: a prior-fitted network for time-series causal inference in longitudinal treatment-response data and zero-shot in-context counterfactual outcome prediction.

CausalLongPFN predicts history-conditional potential outcomes from longitudinal treatment-response time series. Given support trajectories from a new domain, a query history, and a planned future treatment sequence, the frozen model returns a Gaussian-mixture predictive distribution over future outcomes without target-domain gradient updates, propensity-model fitting, or adversarial balancing.

Paper

Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction
Amirhossein Zare, Amirhessam Zare, Herlock Rahimi, Reza Salarikia, Mohammad Kashkooli

• arXiv: 2606.05797
• DOI: 10.48550/arXiv.2606.05797
• Code: https://github.com/Amirhossein-Zare/causal-long-pfn

Citation

@misc{zare2026causallongitudinalpriorfittednetworks,
  title={Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction},
  author={Amirhossein Zare and Amirhessam Zare and Herlock Rahimi and Reza Salarikia and Mohammad Kashkooli},
  year={2026},
  eprint={2606.05797},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2606.05797},
}

Files

• causal-long-pfn-v1-step10000.safetensors — inference-only CausalLongPFN weights in safetensors format
• config.json — architecture, interface, paper, and source-code metadata
• causal_long_pfn_train_config.yaml — compact summary of training and prior settings
• manifest.json — checksum, tensor count, parameter count, and release metadata
• load_model.py — minimal local loading helper
• requirements.txt — minimal dependency list for loading

Release metadata

Field	Value
Training step	10,000
Number of tensors	153
Number of parameters	8,138,384
Safetensors SHA256	437de0c4b95bb89ed66d56ee42516e2ea74306d2964e122884118ec1b40614b1
Tensor fingerprint	0a8cbf027d21d423

The released file contains inference-only model weights in safetensors format. Optimizer, scheduler, and RNG states are not included.

Installation

Install the project code and loading dependencies:

pip install git+https://github.com/Amirhossein-Zare/causal-long-pfn.git
pip install safetensors huggingface_hub

Loading

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from clpfn.models.causal_long_pfn import CausalLongPFN

weights_path = hf_hub_download(
    repo_id="Amirhossein-Zare/causal-long-pfn",
    filename="causal-long-pfn-v1-step10000.safetensors",
)

state_dict = load_file(weights_path, device="cpu")

model = CausalLongPFN()
missing, unexpected = model.load_state_dict(state_dict, strict=False)
if missing or unexpected:
    raise RuntimeError(
        f"State dict mismatch: missing={missing}, unexpected={unexpected}"
    )

model.eval()

You can also clone or download this model repository and run:

python load_model.py

Model interface

The released model uses the fixed CausalLongPFN interface from the paper:

Quantity	Value
Maximum observed time points	60
Maximum rollout horizon	5
Maximum sequence length	65
Support trajectories	3–500
Time-varying covariate channels	up to 10
Outcome channels	1
Static covariates	5
Discrete treatment actions	4
GMM components	5

The model expects the same batch dictionary format used by the source repository evaluation code.

Intended use

These weights are intended for research on:

• time-series causal inference in longitudinal treatment-response data
• longitudinal causal inference
• history-conditional potential-outcome prediction
• counterfactual outcome prediction under planned treatment sequences
• dynamic treatment-response modeling
• zero-shot in-context prediction with prior-fitted networks
• amortized causal inference from synthetic temporal-SCM pretraining

The weights are especially useful for reproducing or extending the CausalLongPFN experiments and for evaluating frozen PFN-style predictors on compatible longitudinal treatment-response time series.

Limitations

As a model for time-series causal inference in longitudinal treatment-response data, CausalLongPFN does not remove the assumptions required for causal interpretation of observational data. Counterfactual validity still depends on consistency, positivity, sequential exchangeability, adequate treatment overlap, and the target domain being reasonably covered by the synthetic temporal-SCM prior.

The model should be treated as a research artifact for causal sequence modeling and hypothesis generation, not as a standalone clinical decision system.