AFFLUENT

Reference implementation of AFFLUENT (quAntized Factor identiFiabLe caUsal Effect variatioNal auToencoder), introduced in:

Quantized Factor Identifiable Causal Effect Variational Autoencoder, CIKM '25.

AFFLUENT extends CEVAE with quantized factor identifiability regularizers (axis-aligned discontinuities + latent-space straightening) to mitigate unidentifiability and bias toward proxy variables in proxy-based causal effect estimation.

Installation

git clone <this-repo>
cd <this-repo>
pip install -r requirements.txt

Requires Python 3.8+. A CUDA-capable GPU is recommended but not required.

Project structure

.
├── models/
│   ├── affluent_demand.py    # CEVAE/AFFLUENT for the synthetic Demand task
│   ├── affluent_ihdp.py      # CEVAE/AFFLUENT for IHDP / Twins
│   ├── civvae.py             # CIV.VAE baseline (IHDP)
│   └── nmmr.py               # NMMR baseline (Demand)
├── experiments/
│   ├── run_demand.py         # main Demand comparison (Figure 5a/b)
│   ├── run_ihdp.py           # main IHDP comparison (Table 1)
│   ├── ablation_demand.py    # Demand ablation (Table 2)
│   ├── gridsearch_demand.py  # Demand hyperparameter sweep
│   ├── figure_heatmap_demand.py  # Figure 5c/d cross-influence heatmaps
│   └── figure_boxplot_demand.py  # aggregates Table 2 Demand numbers
└── utils/
    ├── utils_data_generation_ihdp.py
    └── utils_synthetic_demand.py

Data

• Demand (synthetic): generated on the fly by utils_synthetic_demand.py. No download required.
• IHDP: download ihdp_npci_1-1000.train.npz and ihdp_npci_1-1000.test.npz from the npci package (also mirrored at https://www.fredjo.com/) and place them in ../Mydata/IHDP/ relative to the repo root.
• Twins: download from AMLab-Amsterdam/CEVAE.

Scripts that save results assume the directory ../Mydata/Quantization/Synthetic/{Demand, Demand_optimum, Demand_result, Demand_plot}/ exists. Create these once before running:

mkdir -p ../Mydata/Quantization/Synthetic/{Demand,Demand_optimum,Demand_result,Demand_plot}
mkdir -p ../Mydata/IHDP

Reproducing results

Run from the project root (so ../Mydata/ resolves correctly):

# Demand (Figure 5a/b, MSE on test set)
python -m experiments.run_demand

# IHDP / Twins (Table 1)
python -m experiments.run_ihdp

# Demand ablation (Table 2)
python -m experiments.ablation_demand
python -m experiments.figure_boxplot_demand   # aggregates the saved npys

# Demand cross-influence heatmap (Figure 5c/d)
python -m experiments.figure_heatmap_demand

# Demand hyperparameter sweep (26,244 combinations — long run)
python -m experiments.gridsearch_demand

Set CUDA_VISIBLE_DEVICES via the shell to choose a GPU, e.g. CUDA_VISIBLE_DEVICES=0 python -m experiments.run_demand. If unset, PyTorch picks device 0 (or CPU if none is available).

What's not included

• Twins training script (paper Tables 1 & 2). The affluent_ihdp.py model can be reused, but a run_twins.py driver is not provided.
• TEDVAE + QFI and CIV.VAE + QFI plug-in scripts (paper Table 1). The bare CIV.VAE baseline lives in models/civvae.py; the QFI plug-in integration is not packaged.

Citation

@inproceedings{affluent2025,
author = {Song, Sujeong and Sohn, Junghyo and Kang, Eunsong and Suk, Heung-Il},
title = {Quantized Factor Identifiable Causal Effect Variational Autoencoder},
year = {2025},
booktitle = {Proceedings of the 34th ACM International Conference on Information and Knowledge Management},
}

License

MIT (see LICENSE).