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},
}

Released checkpoint

Field	Value
Pretraining seed	42
Optimizer updates	10,000
Synthetic training episodes processed	2,561,536
Effective batch size	256
Trainable parameters	6,533,903
State-dict tensors	125
State-dict scalar elements	6,554,768
Safetensors SHA256	528ac8894c944e35e3c3414ca1e3aad24aaa50536a428214dbaa56fa3847633e
Tensor fingerprint	d827b87eeb174e0f
Config SHA256	7a4cc9cb03cac6e59c88af3ab2d9210e1305f40388779fa20e386ebe22810a30

The public checkpoint contains inference-only model weights. Optimizer, scheduler, gradient-scaler, and RNG states are not included.

Files

• causal-long-pfn-seed42-step10000.safetensors — seed-42 inference weights in safetensors format
• config.json — model constructor configuration, interface, and release metadata
• causal_long_pfn_train_config.yaml — checkpoint-derived training and prior configuration
• manifest.json — checksums, fingerprints, parameter counts, and provenance metadata
• load_model.py — strict local loading helper
• requirements.txt — minimal loading dependencies

Architecture

Quantity	Value
History encoder layers	4
Context/PFN layers	4
Hidden width	256
Attention heads	8
Feed-forward width	1024
Dropout	0.1
GMM components	5

A causal Transformer encodes each longitudinal history. Labeled support anchors and the query are then processed by a permutation-invariant context encoder. The query representation parameterizes a five-component Gaussian mixture for the normalized next outcome.

Pretraining setup

The model was pretrained entirely on synthetic episodes sampled from a broad prior over temporal structural causal models. The prior exposes the model to nonlinear lagged dynamics, treatment-confounder feedback, latent unit heterogeneity, delayed and cumulative treatment effects, and structured dynamical motifs.

Each support trajectory contributes 4 labeled anchors. Training mixes factual continuations and intervention-conditioned queries:

• factual continuation probability: 0.50
• intervention-conditioned query probability: 0.50
• future-support-covariate masking probability: 0.35

Model interface

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

For multi-step prediction, the learned one-step predictor is recursively applied under the supplied treatment sequence. The predictive mixture mean is inserted as the intermediate outcome while future query covariates remain hidden.

Installation

Install the project source and loading dependencies:

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

Loading from Hugging Face

import json

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from clpfn.models.causal_long_pfn import CausalLongPFN

repo_id = "Amirhossein-Zare/causal-long-pfn"

config_path = hf_hub_download(
    repo_id=repo_id,
    filename="config.json",
)

with open(config_path, "r", encoding="utf-8") as handle:
    release_config = json.load(handle)

weights_path = hf_hub_download(
    repo_id=repo_id,
    filename=release_config["weights_file"],
)

state_dict = load_file(weights_path, device="cpu")
model = CausalLongPFN.from_config(release_config["model"])
model.load_state_dict(state_dict, strict=True)
model.eval()

A local clone or download of this model repository can also be checked with:

python load_model.py

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.