Use matched prior controls and isolate evaluation randomness

README.md

@@ -1,37 +1,44 @@
 # ParticleGAN
 
-**GANs don't collapse when z can move too.**
+**Learnable latent particles for studying GAN mode coverage and stability.**
 
 ![100 Gaussians with Particle Prior](100gaussians.gif)
 
 ## The Problem
 
-Traditional GANs suffer from **mode collapse**: the generator learns to produce only a subset of the data distribution, ignoring other valid modes. This happens because G must warp a *fixed* prior (usually a Gaussian) to match the data. All geometric stress concentrates in G, causing the learned manifold to fold and tear.
+GANs can suffer from **mode collapse**: the generator produces only a subset of the data distribution. This project explores whether optimizing a finite latent particle cloud alongside the generator improves coverage on small, highly multimodal benchmarks.
 
 ## The Insight
 
 **What if the prior could move too?**
 
-Instead of forcing G to do all the work, we introduce learnable "particles" in latent space. These particles move during training to match the structure of the data, absorbing geometric stress alongside G. The result: stable convergence even on highly multimodal distributions.
+We introduce learnable "particles" in latent space. Both the generator and these latent vectors are optimized during training. The experiments examine how that extra flexibility interacts with discriminator regularization, optimizer dynamics, and sample quality. The results are empirical observations on these benchmarks, not a guarantee against collapse.
 
-### Without Particles: Mode Collapse
+### Historical Gaussian example
 
 ![100 Gaussians without Particle Prior](100gaussians_no_particles.gif)
 
-*Same architecture, same hyperparameters, but with a fixed Gaussian prior — the generator collapses to a subset of modes.*
+*Historical visualization from the older Gaussian example. Its architecture and training recipe differ from the particle example above, so these GIFs are not a matched prior comparison.*
 
-## Results
+## Evidence and controls
 
-| Problem | Fixed Gaussian Prior | Particle Prior |
-|---------|---------------------|----------------|
-| 5 modes (text) | collapse | **converges** |
-| 100 modes (2D grid) | collapse | **converges** |
+The historical [regularizer study](FINDINGS.md) compares discriminator penalties within the particle model. It does not establish that a fixed Gaussian prior necessarily collapses. The current examples share one training loop and matched defaults; the only training change for the Gaussian controls is removing the learned prior and its regularizer.
+
+For a reproducible three-way comparison, run:
+
+```bash
+python experiments/compare_priors.py --study-dir runs/prior_comparison --run --device cuda:0
+```
+
+This runs learned particles, a frozen Gaussian table, and fresh Gaussian noise on paired seeds 23001–23003. It records configs, source revision, final samples, coverage, transport distances, and per-mode radial and covariance shape diagnostics. See [prior controls and interpretation](docs/prior-controls.md) and [reproducing the project](docs/reproducing.md).
+
+The completed [nine-run matched comparison](reports/prior-comparison/README.md) reached 100/100 high-quality modes on every learned-prior seed, with a mean high-quality fraction of 98.6%, versus 8.1% for the frozen table and 6.4% for fresh Gaussian noise. This establishes a concentration advantage under this recipe. The report also shows remaining tail and covariance distortion, finite output support, and transport-metric tradeoffs; it does not establish complete Gaussian calibration or a general guarantee against collapse.
 
 ## How It Works
 
 1. **Particle Prior**: Instead of sampling z ~ N(0, I), we maintain a set of learnable latent vectors (particles). During training, we sample from this discrete set.
 
-2. **Joint Optimization**: Particles are optimized alongside G and D. They naturally spread out to cover the data modes.
+2. **Joint Optimization**: Particles are optimized alongside G and D. Their positions can adapt to the data modes.
 
 3. **VICReg Regularization**: We apply variance-covariance regularization to prevent particles from collapsing to a single point, while allowing arbitrary topology (clusters, gaps, etc.).
 
@@ -60,7 +67,7 @@ The main benchmark. 100 Gaussian modes arranged on a 10×10 grid. This is a stre
 python examples/100gaussians.py
 ```
 
-**With particle prior**: All 100 modes are captured — 100/100 modes with ~99% of samples within 3σ of a center after 7k steps.
+The historical particle study reports runs with 100/100 modes and approximately 99% of samples within 3σ of a center after 7k steps. Coverage alone does not establish that the within-mode distribution is correct; the trainer also records shape and transport metrics.
 
 The default recipe is RpGAN (relativistic, logistic) + a one-sided cap gradient penalty on D (`relu(‖∇ₓD‖ − 1)²` on reals and fakes, coeff 1.0), Fourier-feature D, EMA evaluation, Adam β1=0, base LR 6e-4 with a delayed cosine anneal. The cap won a 420-run bake-off against the zero-centered R1/R2 penalty, which is still available with `--reg_arm a_r1r2 --reg_coeff 0.02`. See [FINDINGS.md](FINDINGS.md) for the study and [docs/convergence-tips.md](docs/convergence-tips.md) for the transferable reasoning behind each ingredient.
 
@@ -69,14 +76,14 @@ The default recipe is RpGAN (relativistic, logistic) + a one-sided cap gradient
 python examples/100gaussians_no_particle_prior.py
 ```
 
-The baseline demonstrates classic mode collapse — the generator covers only a fraction of the modes.
+This entrypoint uses the same architecture, losses, learning rates, schedule, and EMA as the particle example, with fresh Gaussian noise. Use `--prior frozen_gaussian` for a finite frozen-table control. The outcome depends on the recipe and seed; the baseline does not assume collapse.
 
 ## Installation
 
 ```bash
 git clone https://github.com/255BITS/ParticleGAN.git
 cd ParticleGAN
-pip install torch matplotlib numpy
+python -m pip install -e '.[dev]'
 ```
 
 ## Project Structure
@@ -119,12 +126,12 @@ from lib.gan_loss import GANLoss
 
 loss_fn = GANLoss(loss_type='hinge', mode='vanilla')
 d_loss = loss_fn.d_loss(d_real, d_fake)
-g_loss = loss_fn.g_loss(d_real, d_fake)
+g_loss = loss_fn.g_loss(d_fake)
 ```
 
 ### VICRegLikeLoss (`lib/vicreg_loss.py`)
 
-Prevents particle collapse while allowing flexible topology:
+Penalizes low marginal variance and cross-dimension covariance while allowing flexible topology:
 
 ```python
 from lib.vicreg_loss import VICRegLikeLoss

docs/prior-controls.md

@@ -0,0 +1,80 @@
+# Matched prior controls
+
+The 100-Gaussian example and its Gaussian counterpart now call the same training
+function. Architecture, discriminator loss and penalty, batch size, generator
+learning rate, discriminator learning rate, run length, schedule, and generator
+EMA are identical. The prior controls are:
+
+| Config / CLI value | Training and metric samples | Learned prior parameters |
+| --- | --- | --- |
+| `particles` | Uniform draws from a learned finite table | 20,000 × 4 = 80,000 by default |
+| `frozen_gaussian` | Uniform draws from one fixed Gaussian table | None |
+| `fresh_gaussian` | New independent Gaussian noise on every call | None |
+| `gaussian` | Compatibility alias for `frozen_gaussian` | None |
+
+The particle arm also optimizes its table at 10× the generator learning rate,
+applies VICReg, and averages its learned positions for EMA readout. The two
+Gaussian controls have no prior optimizer or VICReg gradient. Generator and
+discriminator parameter counts match; total trainable parameter counts differ.
+This isolates the learned-prior intervention, not a fixed total capacity budget.
+
+Fresh Gaussian sampling keeps a fixed reference buffer exclusively for plots
+requested with `fixed_first_n=True`. Ordinary `sample()` calls, coverage, sliced
+W1, and final metrics all draw new Gaussian noise. Re-seeding an evaluation
+generator makes the evaluation latent values repeat across checkpoints without
+restricting the training distribution to a table. `sample()` returns `None`
+for fresh-sample indices because those samples are not particle rows.
+
+The table variants can produce at most 20,000 distinct outputs through this
+deterministic generator. The trainer records `unique_samples` among the 100,000
+final draws. That limit matters when extrapolating toy-benchmark results to
+continuous or higher-dimensional generative modeling.
+
+## Reproduce a comparison
+
+```bash
+python experiments/compare_priors.py \
+  --study-dir runs/prior_comparison --run --device cuda:0
+```
+
+The default comparison is nine runs: each of the three priors on seeds 23001,
+23002, and 23003, with 7,000 steps, Rp logistic loss, Fourier-2 discriminator,
+`b_cap` coefficient 1.0, generator LR 6e-4, discriminator LR multiplier 1.5,
+Adam beta1=0, delayed cosine annealing, and EMA 0.995. These seeds are separate
+from the historical study's standard seeds. They are specified in advance;
+report all of them, including failed runs, instead of selecting the best seed.
+
+Omit `--run` to generate configs and a manifest for an external scheduler. The
+manifest includes resolved config paths, source revision, whether the checkout
+was dirty, Python and dependency versions, and the PyTorch CUDA build version.
+Use a clean committed checkout for reported runs. Generate into a new directory
+for each study; `--collect` collects completed runs from an existing manifest
+without rewriting its provenance. A small `--total-steps` value is useful for
+smoke tests, but is not evidence about convergence under the full recipe.
+
+Each run saves its exact final sample cloud, checkpoint, evaluation time series,
+and summary. `comparison.json` gathers final metrics for all seed/prior pairs.
+Use coverage together with high-quality fraction, transport distance, histogram
+balance, radial core spread, and both covariance eigenvalue ratios. A radial
+median alone can hide anisotropy or tails. Three seeds provide a modest
+replication check, not a broad guarantee of stability.
+
+## RNG isolation and historical results
+
+Training now owns separate generators for real data, latent samples, and
+penalty interpolation. Evaluation has its own re-seeded generator. Changing
+evaluation frequency no longer changes the learned parameters; a bounded CPU
+regression exercises the actual trainer for all three priors and compares every
+saved model tensor and final sample array across evaluation intervals.
+
+This correction changes trajectories relative to historical code, which drew
+training latents and diagnostic samples from the same global stream. New
+summaries and comparison manifests identify the RNG scheme as
+`separate_data_latent_penalty_v1`. Historical GIFs and results retain their
+original meaning; they should not be relabeled as results of the corrected
+matched comparison. In particular, the older no-particle GIF also used a
+different architecture and stabilization recipe and does not isolate the prior.
+
+Checkpoints include `prior_kind`; reconstruct a prior with
+`lib.particle_prior.make_prior(prior_kind, num_particles=..., z_dim=...)` before
+loading its state to preserve fresh-noise versus finite-table semantics.

docs/reproducing.md

@@ -0,0 +1,101 @@
+# Reproducing and checking experiments
+
+Use Python 3.10 or newer. The CPU test workflow checks Python 3.11 and 3.12.
+Run these commands from the repository root:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+python -m pip install -e '.[dev]'
+python -m pip check
+OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 python -m pytest -q
+```
+
+`pyproject.toml` includes the dependencies needed by the experiment scripts,
+including PyYAML and POT (`import ot`) for exact transport evaluation. A CUDA
+installation of PyTorch is useful for full studies. For CPU-only work, install
+PyTorch from its CPU wheel index before installing the project:
+
+```bash
+python -m pip install 'torch>=2.6,<3' --index-url https://download.pytorch.org/whl/cpu
+python -m pip install -e '.[dev]'
+```
+
+The tests exercise numerical gradient correctness, evaluation isolation, sparse
+gradient penalties, particle identity bookkeeping, experiment failure/resume
+behavior, and small training runs. They do not reproduce the reported full
+studies or establish convergence for a new dataset. The exact-transport test
+uses a distribution with a known answer so a missing solver is caught before
+an expensive training run reaches its final evaluation.
+
+## Record the environment
+
+The manifest defines supported dependency ranges, not a bit-for-bit numerical
+environment. Save the resolved environment and commit with every full study:
+
+```bash
+mkdir -p results/reproduction
+python -m pip freeze > results/reproduction/requirements.txt
+git rev-parse HEAD > results/reproduction/commit.txt
+python -c 'import torch; print(torch.__version__); print(torch.version.cuda)' > results/reproduction/torch.txt
+```
+
+Also retain the expanded run configurations, hardware description, summaries,
+and samples. Matching a seed is insufficient for exact reproduction across
+different PyTorch/CUDA versions and hardware. Historical tables were produced
+before the evaluation-randomness fix and should not be expected to match new
+trajectories bit for bit.
+
+## Controlled prior comparison
+
+The current comparison protocol uses the same generator, discriminator, loss,
+penalty, learning rates, and training budget for three priors:
+
+* `particles`: a learned finite table, with its prior optimizer and VICReg term;
+* `frozen_gaussian`: a fixed table initialized from a Gaussian;
+* `fresh_gaussian`: newly sampled Gaussian noise for ordinary draws.
+
+The legacy configuration name `gaussian` continues to mean the frozen table.
+It must not be interpreted as fresh continuous Gaussian sampling. The learned
+table's optimizer and regularizer are part of the learned-prior intervention;
+further ablations are needed to distinguish their individual effects.
+
+```bash
+python experiments/compare_priors.py --help
+python experiments/compare_priors.py --study-dir runs/prior_comparison --run --device cuda:0
+```
+
+Omit `--run` to emit configurations and a protocol manifest for an external
+scheduler. Use a new study directory for each generation. See
+[the prior-control protocol](prior-controls.md) for outputs and collection.
+Treat the old GIFs and
+regularizer tables as historical recipe results, not as this controlled
+three-way comparison. Use fresh seeds for confirmation after selecting a
+recipe, and retain individual seed scores rather than only an aggregate.
+
+## Interpreting shape and diversity
+
+For a deterministic generator over M fixed particles, ordinary sampling can
+emit at most M distinct outputs. Drawing 100,000 samples from 20,000 particles
+does not create 100,000 independent locations in output space. Mode coverage,
+sample plausibility, continuous diversity, and distributional fidelity are
+separate properties.
+
+The robust core ratio estimates a radial median using an isotropic-Gaussian
+conversion. A two-point cloud in each mode can pass that check while having
+zero variance along one axis. The covariance eigenvalue ratios now expose
+this failure: each eigenvalue is divided by the true variance, and the minimum
+and maximum are averaged over sufficiently populated modes. These are
+tail-sensitive population moments; read them alongside core spread, tail
+mass, and the number of audited modes. Even correct covariance does not prove
+Gaussian shape, so inspect radial and angular structure for stronger claims.
+
+Generated-sample NLL under the target rewards concentrating at mode centers.
+It is a plausibility score, not a calibration test. Likewise, a confidence
+interval containing zero is an inconclusive difference test, not proof of
+equivalence. Establish an equivalence margin and adequate replication before
+claiming two recipes perform the same.
+
+For future scaling studies, vary the particle count and separately test noise
+around particles. Adding noise changes the model distribution and requires a
+fresh comparison; the finite-support limitation is not silently removed here.