README.md

Status

This repository contains ongoing research on generative latent-variable models under informative missingness (MNAR), including GC-SIWAE and related experiments. Results, APIs, configurations, and documentation may change as the study evolves. The current codebase should not yet be treated as a final reproducibility release.

mnar_2026_rptu

Clean, paper-aligned implementation of Gradient Conflict-Aware Selection-Augmented Importance-Weighted Variational Learning (GC-SIWAE), baseline comparisons, and numerical theory validation.

Included experiments

• GC-SIWAE with logistic and neural selection models.
• Masked VAE, MIWAE, not-MIWAE, finite-mixture MIWAE, and DeMissVAE-style unified baseline reimplementations.
• Beijing controlled-MNAR benchmark.
• Synthetic shared, hybrid, and direct MNAR mechanisms.
• Misspecification and gradient-projection validation.
• All figures are written as PNG and PDF.
• Mean-imputation VAE and the legacy staged schedule are intentionally excluded.

Baseline note: the baseline modules use a common PyTorch benchmark implementation so all methods share architecture, data, splits, optimizer, and metrics. They are research reimplementations, not verbatim copies of each authors' released repository. For a camera-ready comparison, verify each implementation against the corresponding paper/repository and report this distinction.

Cluster setup

unzip mnar_2026_rptu.zip
cd mnar_2026_rptu
python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"

The Beijing CSV files are already included under data/raw/beijing/.

Smoke test

bash scripts/smoke_test.sh

Full Slurm run

sbatch slurm/run_all.sbatch

or interactively:

bash scripts/run_all.sh

Main outputs

results/beijing/
results/synthetic/{shared,hybrid,direct}/
results/theory/

Every benchmark produces metrics, checkpoints, training histories, GC gradient diagnostics, and PDF/PNG figures.

Where the baseline implementations are

The executable unified implementations are located at:

src/mnar_2026_rptu/baselines/
├── masked_vae.py
├── miwae.py
├── not_miwae.py
├── finite_mixture_miwae.py
└── demissvae.py

Their shared model, objectives, and trainers are in src/mnar_2026_rptu/models/ and src/mnar_2026_rptu/training/. GC-SIWAE (this paper's method) is implemented in the same models/ and training/ modules, with its gradient projection logic in src/mnar_2026_rptu/training/projection.py.

The baseline source repositories supplied by the researcher are preserved, unmodified, for method-fidelity checks at:

external/
├── miwae/
├── notMIWAE/
└── demiss-vae/

See BASELINES.md for the exact file-by-file mapping.

GC-SIWAE variants and ablations

Beyond the two primary gc_siwae_logistic / gc_siwae_neural models, the codebase supports:

• Projection ablation — gc_siwae_logistic_noproj / gc_siwae_neural_noproj run the identical architecture, data, seed, and joint-training loop as their counterparts above, but with the gradient-conflict projection (Proposition 6.10) switched off at the source (training/projection.py's enable_projection flag). Comparing these against the projected variants isolates what the projection itself contributes, rather than just the joint-training setup it's embedded in. Included by default in configs/default.yaml and configs/smoke.yaml.
• Adaptive selection weight — set adaptive_selection_weight: true in a config (see configs/adaptive_weight.yaml) to scale the selection term's weight each step by the reconstruction/selection gradient-norm ratio (GradNorm-style), clamped by grad_norm_ratio_min / grad_norm_ratio_max, instead of a fixed selection_weight throughout training. Off by default so configs/default.yaml reproduces the original fixed-weight behavior. The applied weight is logged per step as effective_selection_weight in results/*/diagnostics/gc_gradient_diagnostics.csv.

Robustness guards in training/projection.py (project_conflicting): skips the projection (rather than dividing by a near-zero norm) when ‖g_R‖ is below min_grad_norm, and detects non-finite gradients, in which case train_gc skips the optimizer step for that batch and records it in skipped_steps_total instead of propagating NaNs into the model weights.

Choosing K (importance-weighted samples)

iwae_samples (K) is currently a single shared value consumed by both training and evaluation, for both the reconstruction and selection terms (except masked_vae, which is always trained/evaluated at K=1 since a plain VAE has no importance weighting). scripts/run_k_sweep.py searches over a range of K values and reports, per model, the mean +/- std of a chosen metric across seeds at each K, picking the best:

python scripts/run_k_sweep.py \
    --config configs/default.yaml \
    --data data/processed/synthetic_shared \
    --output results/k_sweep \
    --k-min 4 --k-max 20 --k-step 2 \
    --models gc_siwae_logistic,gc_siwae_neural \
    --metric missing_reconstruction_mse --direction min

This writes k_sweep_raw_metrics.csv (every run), k_sweep_summary.csv (mean/std per model per K), best_k.json (the selected K per model), and k_sweep.png/.pdf (metric vs. K, selected K marked with a star). Use --metric selection_f1 --direction max or any other column from evaluation.metrics.evaluate() to optimize a different criterion, and --models all to also sweep K for every baseline that consumes cfg["iwae_samples"] (all except masked_vae).

Two things worth knowing before picking a wide range: (1) the gradient-conflict check in GC-SIWAE (project_conflicting's dot < 0 test) is a decision made on a K-sample gradient estimate, so very low K makes that decision noisier, independent of the metric this sweep optimizes; (2) IWAE's encoder gradient signal-to-noise ratio does not improve monotonically with K (Rainforth et al., 2018), and gR/gS here are both encoder gradients, so "K as high as possible" is not guaranteed to be best — this is exactly why the sweep selects by measured metric rather than assuming higher K wins.

Choosing selection_weight

scripts/run_weight_sweep.py sweeps a fixed selection_weight (unlike adaptive_selection_weight, which rescales it per-step -- see above) across a list of values and reports the best per model, same interface as the K sweep. It clears any conflicting selection_weight_by_model entry for the model(s) being swept, so the sweep's value always takes effect:

python scripts/run_weight_sweep.py \
    --config configs/default.yaml \
    --data data/processed/synthetic_shared \
    --output results/weight_sweep \
    --weights 0.2,0.5,0.75,1,1.5,2,3,4,5,5.5,6,6.5,7,8,10,20 \
    --models gc_siwae_logistic,gc_siwae_neural \
    --metric missing_reconstruction_mse --direction min

Findings so far (synthetic shared mechanism; full details/plots in the git history of configs/default.yaml and configs/pilot.yaml):

• There is a real, non-monotonic optimum, not "higher is always better." gc_siwae_logistic peaks at w=6.0 (fine grid: 5.5/6.0/6.5/7.0 -> 0.1728/0.1723/0.1724/0.1741 missing-reconstruction MSE at pilot scale); gc_siwae_neural's pilot-scale peak sits slightly lower, around w=5.0.
• Past roughly w=8, both the mean error and the cross-seed variance get worse -- std roughly triples by w=10. This traces to cosine_before (the raw, pre-projection g_R/g_S alignment) dropping from -0.12 (w=6) to -0.36 (w=20) in the diagnostics CSV: higher weight doesn't just make conflict more frequent, it makes it more severe, and different seeds resolve that escalating tug-of-war differently. Not a gradient-clipping artifact (clip rate stays flat ~13-18% across that whole range).
• At full scale (hidden_dim=128, 300 epochs, 2 seeds so far -- not yet the full 5), gc_siwae_logistic at w=6.0 reached 0.1010/0.0997 missing MSE, beating not_miwae (0.1114/0.1020) on both seeds, and beating its own w=1.0 result (0.1045). This is the first configuration in this project where GC-SIWAE has outperformed a baseline rather than trailing it -- worth reproducing with the full seed count before relying on it. gc_siwae_neural at its pilot-optimal w=5.0 has not yet been confirmed at full scale under a fair (matching 300-epoch) budget; its MLP selector is slower per step than logistic's, and a capped 150-epoch full-scale test underperformed logistic -- inconclusive, since it's plausible neural simply needed more of its budget to converge, not that it's worse.

selection_weight_by_model (in configs/default.yaml and configs/pilot.yaml) lets different GC-SIWAE architectures use different weights -- a single flat selection_weight can't be simultaneously optimal for both logistic and neural selectors, which is exactly what the sweep above found. It maps an architecture name (e.g. gc_siwae_logistic) to a weight; a _noproj ablation variant automatically uses its projected counterpart's weight (only the projection itself should differ between them), and any model not listed falls back to the flat selection_weight. The two current entries carry different levels of evidence -- read the comments in configs/default.yaml before citing either in a paper: gc_siwae_logistic: 6.0 is validated at full scale (2 seeds, real 300 epochs, beat not_miwae on both); gc_siwae_neural: 5.0 is pilot-scale only and not yet fairly tested at full scale.

Data/config mismatch guard

prepare_data.py writes synthetic data to data/processed/synthetic_{mechanism} based only on the mechanism name, not on which config generated it. Running it once with, e.g., configs/smoke.yaml and later with configs/pilot.yaml (or any two configs with different synthetic: settings) silently overwrites the first config's data at that same path. This isn't hypothetical -- it happened during development of this codebase and produced a batch of misleading results before being caught. run_benchmark now calls data.common.verify_synthetic_metadata before training, which compares the data's embedded shape fingerprint (n_samples, n_features, latent_dim, stored in metadata.json) against the config actually being run, and raises a clear RuntimeError naming the exact mismatch instead of training silently on stale data. If you regenerate data for a different config, either point --output at a distinct directory or re-run prepare_data.py for the config you're about to use immediately before training.

Figure style (ICML / NeurIPS camera-ready)

Every figure-producing module (evaluation/plots.py, theory/validate.py, scripts/run_k_sweep.py) applies evaluation/style.apply_style(), which fixes four things plain matplotlib defaults get wrong for a camera-ready submission:

• Embedded vector fonts, not bitmap. pdf.fonttype/ps.fonttype are forced to 42 (TrueType/CID, scalable) instead of matplotlib's default Type 3 (bitmap), which several venues' PDF/A compliance checks flag and which looks visibly blurry when a reviewer zooms in. Verified directly by inspecting a generated PDF's font subtypes (CIDFontType2/Type0, no Type3).
• A serif font matching LaTeX body text, via STIXGeneral — a Times-like font bundled with matplotlib itself, so it renders identically on any machine with no external font install required (fallbacks: Liberation Serif, Times New Roman, DejaVu Serif).
• Figures sized to actual column widths — evaluation.style.figsize() returns sizes matched to the standard ICML/NeurIPS single-column (3.25in) and double-column (6.75in) widths, at a 9/8pt font size, so text stays legible at the size the figure will actually be placed in the paper, rather than being shrunk from an arbitrary default.
• A colorblind-safe palette — the 8-color Okabe & Ito (2008) qualitative palette, replacing matplotlib's default tab10.

In-axes titles are intentionally not used (the LaTeX \caption{} carries that text instead, per standard camera-ready practice).