Skip to content

Instantly share code, notes, and snippets.

@wojtyniak

wojtyniak/notebook-de7efe6f-ce6e-477f-804d-80c30d6e3af5-paper-20260210-000005.ipynb

Created February 10, 2026 00:00
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save wojtyniak/b162326a846efc6578712815da05b184 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save wojtyniak/b162326a846efc6578712815da05b184 to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
notebook-de7efe6f-ce6e-477f-804d-80c30d6e3af5-paper-20260210-000005.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# PQMass: Probabilistic Assessment of the Quality of Generative Models\n",
"\n",
"**Paper:** *PQMass: Probabilistic Assessment of the Quality of Generative Models using Probability Mass Estimation*\n",
"\n",
"**Authors:** Pablo Lemos, Sammy Sharief, Nikolay Malkin, Salma Salhi, Connor Stone, Laurence Perreault-Levasseur, Yashar Hezaveh\n",
"\n",
"---\n",
"\n",
"## Overview\n",
"\n",
"This notebook implements the **PQMass** method, a novel two-sample statistical test for assessing the quality of generative models. PQMass uses Voronoi tessellation-based probability mass estimation to compare distributions without relying on learned feature extractors.\n",
"\n",
"**Key Features:**\n",
"- Works directly in sample space (no feature extraction needed)\n",
"- Computationally efficient for high-dimensional data\n",
"- Detects both fidelity issues and mode dropping\n",
"- Provides interpretable chi-squared statistics with known distributions\n",
"\n",
"**Note on Resource Constraints:**\n",
"This notebook uses small-scale examples and synthetic data to demonstrate the methodology within memory (4GB) and time constraints. For full-scale experiments with real datasets and complete model training, researchers should run on their own infrastructure with appropriate computational resources.\n",
"\n",
"---\n",
"\n",
"## Table of Contents\n",
"\n",
"1. Setup and Dependencies\n",
"2. Core PQMass Algorithm Implementation\n",
"3. Workflow 1: PQMass Two-Sample Test (Core Method)\n",
"4. Workflow 2: Null Test Validation (Gaussian Mixture Model)\n",
"5. Workflow 3: Mode-Dropping Detection (Synthetic Example)\n",
"6. Workflow 4: Diffusion Model Training Progress (MNIST Toy Example)\n",
"7. Workflow 5: Baseline Comparisons (MMD, Wasserstein)\n",
"8. Conclusion and Scaling Guidance"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Setup and Dependencies\n",
"\n",
"Install all required packages using `uv pip install`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Install all dependencies in a single command\n",
"!uv pip install numpy scipy matplotlib scikit-learn torch torchvision tqdm seaborn"
],
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Import libraries\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import seaborn as sns\n",
"from scipy import stats\n",
"from scipy.spatial.distance import cdist\n",
"from sklearn.neighbors import NearestNeighbors\n",
"from tqdm import tqdm\n",
"import warnings\n",
"warnings.filterwarnings('ignore')\n",
"\n",
"# Set random seeds for reproducibility\n",
"np.random.seed(42)\n",
"\n",
"# Configure plotting\n",
"plt.style.use('seaborn-v0_8-darkgrid')\n",
"sns.set_palette(\"husl\")\n",
"\n",
"print(\"All libraries imported successfully!\")\n",
"print(f\"NumPy version: {np.__version__}\")"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Core PQMass Algorithm Implementation\n",
"\n",
"The PQMass method consists of the following steps:\n",
"\n",
"1. **Define Reference Points**: Sample reference points from a uniform mixture of two sample sets\n",
"2. **Count Points in Voronoi Cells**: Assign samples to nearest reference points\n",
"3. **Compute Chi-Squared Statistic**: Calculate Pearson's χ² statistic\n",
"4. **Compute P-Value**: Evaluate significance using χ² distribution\n",
"5. **Repeat with Multiple Tessellations**: Reduce variance by averaging over multiple random tessellations"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"def pqmass_test(samples_p, samples_q, n_reference=100, metric='euclidean', return_details=False):\n",
" \"\"\"\n",
" Perform PQMass two-sample test.\n",
" \n",
" Parameters:\n",
" -----------\n",
" samples_p : array-like, shape (n_p, d)\n",
" Samples from distribution P\n",
" samples_q : array-like, shape (n_q, d)\n",
" Samples from distribution Q\n",
" n_reference : int\n",
" Number of reference points for Voronoi tessellation\n",
" metric : str\n",
" Distance metric ('euclidean', 'manhattan', etc.)\n",
" return_details : bool\n",
" If True, return detailed statistics\n",
" \n",
" Returns:\n",
" --------\n",
" chi2_stat : float\n",
" Chi-squared statistic\n",
" p_value : float\n",
" P-value from chi-squared distribution\n",
" dof : int\n",
" Degrees of freedom (n_reference - 1)\n",
" \"\"\"\n",
" samples_p = np.array(samples_p)\n",
" samples_q = np.array(samples_q)\n",
" \n",
" n_p = len(samples_p)\n",
" n_q = len(samples_q)\n",
" \n",
" # Step 1: Sample reference points from uniform mixture\n",
" # Half from P, half from Q\n",
" n_ref_p = n_reference // 2\n",
" n_ref_q = n_reference - n_ref_p\n",
" \n",
" ref_indices_p = np.random.choice(n_p, size=n_ref_p, replace=False)\n",
" ref_indices_q = np.random.choice(n_q, size=n_ref_q, replace=False)\n",
" \n",
" reference_points = np.vstack([\n",
" samples_p[ref_indices_p],\n",
" samples_q[ref_indices_q]\n",
" ])\n",
" \n",
" # Remove reference points from sample sets\n",
" mask_p = np.ones(n_p, dtype=bool)\n",
" mask_p[ref_indices_p] = False\n",
" remaining_p = samples_p[mask_p]\n",
" \n",
" mask_q = np.ones(n_q, dtype=bool)\n",
" mask_q[ref_indices_q] = False\n",
" remaining_q = samples_q[mask_q]\n",
" \n",
" # Step 2: Assign samples to Voronoi cells (nearest reference point)\n",
" nbrs = NearestNeighbors(n_neighbors=1, metric=metric, algorithm='auto')\n",
" nbrs.fit(reference_points)\n",
" \n",
" # Find nearest reference point for each sample\n",
" _, indices_p = nbrs.kneighbors(remaining_p)\n",
" _, indices_q = nbrs.kneighbors(remaining_q)\n",
" \n",
" # Count samples in each Voronoi cell\n",
" counts_p = np.bincount(indices_p.flatten(), minlength=n_reference)\n",
" counts_q = np.bincount(indices_q.flatten(), minlength=n_reference)\n",
" \n",
" # Step 3: Compute expected counts under null hypothesis\n",
" n_remaining_p = len(remaining_p)\n",
" n_remaining_q = len(remaining_q)\n",
" total_counts = counts_p + counts_q\n",
" total_samples = n_remaining_p + n_remaining_q\n",
" \n",
" # Expected counts: E[n_i^p] = n_p * (n_i^p + n_i^q) / (n_p + n_q)\n",
" expected_p = (n_remaining_p / total_samples) * total_counts\n",
" expected_q = (n_remaining_q / total_samples) * total_counts\n",
" \n",
" # Step 4: Compute Pearson chi-squared statistic\n",
" # χ² = Σ [(O - E)² / E] for both distributions\n",
" # Add small epsilon to avoid division by zero\n",
" epsilon = 1e-10\n",
" chi2_stat = np.sum((counts_p - expected_p)**2 / (expected_p + epsilon)) + \\\n",
" np.sum((counts_q - expected_q)**2 / (expected_q + epsilon))\n",
" \n",
" # Degrees of freedom: n_reference - 1\n",
" dof = n_reference - 1\n",
" \n",
" # Step 5: Compute p-value from chi-squared distribution\n",
" p_value = 1 - stats.chi2.cdf(chi2_stat, dof)\n",
" \n",
" if return_details:\n",
" return {\n",
" 'chi2_stat': chi2_stat,\n",
" 'p_value': p_value,\n",
" 'dof': dof,\n",
" 'counts_p': counts_p,\n",
" 'counts_q': counts_q,\n",
" 'expected_p': expected_p,\n",
" 'expected_q': expected_q,\n",
" 'n_reference': n_reference\n",
" }\n",
" \n",
" return chi2_stat, p_value, dof\n",
"\n",
"\n",
"def pqmass_repeated(samples_p, samples_q, n_reference=100, n_iterations=20, metric='euclidean'):\n",
" \"\"\"\n",
" Perform PQMass test multiple times with different tessellations.\n",
" \n",
" This reduces variance from the choice of reference points.\n",
" \n",
" Parameters:\n",
" -----------\n",
" samples_p, samples_q : array-like\n",
" Sample sets to compare\n",
" n_reference : int\n",
" Number of reference points per tessellation\n",
" n_iterations : int\n",
" Number of different tessellations to try\n",
" metric : str\n",
" Distance metric\n",
" \n",
" Returns:\n",
" --------\n",
" mean_chi2 : float\n",
" Mean chi-squared statistic\n",
" std_chi2 : float\n",
" Standard deviation of chi-squared statistics\n",
" mean_p_value : float\n",
" Mean p-value\n",
" dof : int\n",
" Degrees of freedom\n",
" \"\"\"\n",
" chi2_stats = []\n",
" p_values = []\n",
" \n",
" for _ in range(n_iterations):\n",
" chi2, p_val, dof = pqmass_test(samples_p, samples_q, n_reference, metric)\n",
" chi2_stats.append(chi2)\n",
" p_values.append(p_val)\n",
" \n",
" return {\n",
" 'mean_chi2': np.mean(chi2_stats),\n",
" 'std_chi2': np.std(chi2_stats),\n",
" 'mean_p_value': np.mean(p_values),\n",
" 'std_p_value': np.std(p_values),\n",
" 'dof': dof,\n",
" 'all_chi2': chi2_stats,\n",
" 'all_p_values': p_values\n",
" }\n",
"\n",
"print(\"PQMass core functions implemented successfully!\")"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. Workflow 1: PQMass Two-Sample Test (Core Method)\n",
"\n",
"Let's demonstrate the core PQMass method with a simple example: comparing two 2D Gaussian distributions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Generate synthetic data: two 2D Gaussian distributions\n",
"n_samples = 1000\n",
"\n",
"# Distribution P: Gaussian centered at (0, 0)\n",
"mean_p = [0, 0]\n",
"cov_p = [[1, 0], [0, 1]]\n",
"samples_p = np.random.multivariate_normal(mean_p, cov_p, n_samples)\n",
"\n",
"# Distribution Q: Gaussian centered at (0.5, 0.5) - slightly shifted\n",
"mean_q = [0.5, 0.5]\n",
"cov_q = [[1, 0], [0, 1]]\n",
"samples_q = np.random.multivariate_normal(mean_q, cov_q, n_samples)\n",
"\n",
"# Visualize the distributions\n",
"fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
"\n",
"axes[0].scatter(samples_p[:, 0], samples_p[:, 1], alpha=0.5, s=10, label='Distribution P')\n",
"axes[0].scatter(samples_q[:, 0], samples_q[:, 1], alpha=0.5, s=10, label='Distribution Q')\n",
"axes[0].set_xlabel('Dimension 1')\n",
"axes[0].set_ylabel('Dimension 2')\n",
"axes[0].set_title('Two Slightly Different Gaussian Distributions')\n",
"axes[0].legend()\n",
"axes[0].grid(True, alpha=0.3)\n",
"\n",
"# Distribution R: Same as P (null test)\n",
"samples_r = np.random.multivariate_normal(mean_p, cov_p, n_samples)\n",
"\n",
"axes[1].scatter(samples_p[:, 0], samples_p[:, 1], alpha=0.5, s=10, label='Distribution P')\n",
"axes[1].scatter(samples_r[:, 0], samples_r[:, 1], alpha=0.5, s=10, label='Distribution R (same as P)')\n",
"axes[1].set_xlabel('Dimension 1')\n",
"axes[1].set_ylabel('Dimension 2')\n",
"axes[1].set_title('Two Identical Gaussian Distributions (Null Test)')\n",
"axes[1].legend()\n",
"axes[1].grid(True, alpha=0.3)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"Data generated successfully!\")"
],
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Test 1: Compare P and Q (different distributions)\n",
"print(\"Test 1: Comparing two DIFFERENT Gaussian distributions\")\n",
"print(\"=\"*60)\n",
"\n",
"chi2_pq, p_value_pq, dof_pq = pqmass_test(samples_p, samples_q, n_reference=100)\n",
"\n",
"print(f\"Chi-squared statistic: {chi2_pq:.2f}\")\n",
"print(f\"Degrees of freedom: {dof_pq}\")\n",
"print(f\"Expected chi-squared (under null): {dof_pq:.2f}\")\n",
"print(f\"P-value: {p_value_pq:.4f}\")\n",
"print(f\"\\nInterpretation: {'REJECT null hypothesis - distributions are different' if p_value_pq < 0.05 else 'ACCEPT null hypothesis - distributions are similar'}\")\n",
"\n",
"print(\"\\n\" + \"=\"*60)\n",
"\n",
"# Test 2: Compare P and R (same distributions - null test)\n",
"print(\"\\nTest 2: Comparing two IDENTICAL Gaussian distributions (Null Test)\")\n",
"print(\"=\"*60)\n",
"\n",
"chi2_pr, p_value_pr, dof_pr = pqmass_test(samples_p, samples_r, n_reference=100)\n",
"\n",
"print(f\"Chi-squared statistic: {chi2_pr:.2f}\")\n",
"print(f\"Degrees of freedom: {dof_pr}\")\n",
"print(f\"Expected chi-squared (under null): {dof_pr:.2f}\")\n",
"print(f\"P-value: {p_value_pr:.4f}\")\n",
"print(f\"\\nInterpretation: {'REJECT null hypothesis - distributions are different' if p_value_pr < 0.05 else 'ACCEPT null hypothesis - distributions are similar'}\")\n",
"\n",
"print(\"\\n\" + \"=\"*60)"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Repeated Tessellations\n",
"\n",
"To reduce variance from the choice of reference points, we repeat the test with multiple random tessellations and report the mean and standard deviation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Perform repeated PQMass tests\n",
"print(\"Performing PQMass with 20 different tessellations...\\n\")\n",
"\n",
"results_pq = pqmass_repeated(samples_p, samples_q, n_reference=100, n_iterations=20)\n",
"results_pr = pqmass_repeated(samples_p, samples_r, n_reference=100, n_iterations=20)\n",
"\n",
"print(\"Comparing P and Q (DIFFERENT distributions):\")\n",
"print(f\" Mean χ²: {results_pq['mean_chi2']:.2f} ± {results_pq['std_chi2']:.2f}\")\n",
"print(f\" Mean p-value: {results_pq['mean_p_value']:.4f} ± {results_pq['std_p_value']:.4f}\")\n",
"print(f\" Expected χ²: {results_pq['dof']:.2f}\\n\")\n",
"\n",
"print(\"Comparing P and R (SAME distribution - Null Test):\")\n",
"print(f\" Mean χ²: {results_pr['mean_chi2']:.2f} ± {results_pr['std_chi2']:.2f}\")\n",
"print(f\" Mean p-value: {results_pr['mean_p_value']:.4f} ± {results_pr['std_p_value']:.4f}\")\n",
"print(f\" Expected χ²: {results_pr['dof']:.2f}\")\n",
"\n",
"# Visualize the distribution of chi-squared statistics\n",
"fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
"\n",
"ax.hist(results_pr['all_chi2'], bins=15, alpha=0.6, label='P vs R (same)', density=True)\n",
"ax.hist(results_pq['all_chi2'], bins=15, alpha=0.6, label='P vs Q (different)', density=True)\n",
"\n",
"# Plot theoretical chi-squared distribution\n",
"x = np.linspace(60, 140, 200)\n",
"ax.plot(x, stats.chi2.pdf(x, results_pq['dof']), 'k--', lw=2, label=f'χ²({results_pq[\"dof\"]}) theoretical')\n",
"\n",
"ax.axvline(results_pq['dof'], color='red', linestyle=':', lw=2, label='Expected value (dof)')\n",
"ax.set_xlabel('Chi-squared Statistic')\n",
"ax.set_ylabel('Density')\n",
"ax.set_title('Distribution of PQMass Chi-squared Statistics (20 tessellations)')\n",
"ax.legend()\n",
"ax.grid(True, alpha=0.3)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. Workflow 2: Null Test Validation (Gaussian Mixture Model)\n",
"\n",
"Following Section 3.1 and Appendix A of the paper, we validate that the chi-squared statistic follows the expected χ² distribution when comparing samples from the same distribution."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"def create_gaussian_mixture(n_components=5, dim=2, n_samples=1000):\n",
" \"\"\"\n",
" Create samples from a Gaussian Mixture Model.\n",
" \"\"\"\n",
" samples = []\n",
" samples_per_component = n_samples // n_components\n",
" \n",
" for i in range(n_components):\n",
" # Random mean and covariance for each component\n",
" mean = np.random.randn(dim) * 3\n",
" cov = np.eye(dim) * (0.5 + np.random.rand())\n",
" component_samples = np.random.multivariate_normal(mean, cov, samples_per_component)\n",
" samples.append(component_samples)\n",
" \n",
" return np.vstack(samples)\n",
"\n",
"# Perform null test validation\n",
"print(\"Null Test Validation: Testing that χ² follows expected distribution\")\n",
"print(\"=\"*70)\n",
"\n",
"# Parameters\n",
"n_reference = 50 # Smaller for faster computation\n",
"n_repetitions = 100 # Paper uses 2^14 = 16384, we use 100 for speed\n",
"n_samples_gmm = 500\n",
"\n",
"chi2_values = []\n",
"\n",
"print(f\"Running {n_repetitions} null test experiments...\")\n",
"for i in tqdm(range(n_repetitions)):\n",
" # Generate two independent sample sets from the same GMM\n",
" np.random.seed(i) # Different seed for each iteration\n",
" samples_1 = create_gaussian_mixture(n_components=5, dim=2, n_samples=n_samples_gmm)\n",
" samples_2 = create_gaussian_mixture(n_components=5, dim=2, n_samples=n_samples_gmm)\n",
" \n",
" chi2, _, dof = pqmass_test(samples_1, samples_2, n_reference=n_reference)\n",
" chi2_values.append(chi2)\n",
"\n",
"chi2_values = np.array(chi2_values)\n",
"\n",
"print(f\"\\nResults:\")\n",
"print(f\" Mean χ²: {np.mean(chi2_values):.2f}\")\n",
"print(f\" Expected (dof): {dof}\")\n",
"print(f\" Std χ²: {np.std(chi2_values):.2f}\")\n",
"print(f\" Expected std: {np.sqrt(2*dof):.2f}\")"
],
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Visualize the null test validation\n",
"fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
"\n",
"# Histogram comparison\n",
"axes[0].hist(chi2_values, bins=30, density=True, alpha=0.7, label='Empirical distribution')\n",
"x = np.linspace(20, 80, 200)\n",
"axes[0].plot(x, stats.chi2.pdf(x, dof), 'r-', lw=2, label=f'Theoretical χ²({dof})')\n",
"axes[0].axvline(dof, color='black', linestyle='--', lw=2, label='Expected mean')\n",
"axes[0].set_xlabel('Chi-squared Statistic')\n",
"axes[0].set_ylabel('Density')\n",
"axes[0].set_title('Null Test Validation: Empirical vs Theoretical Distribution')\n",
"axes[0].legend()\n",
"axes[0].grid(True, alpha=0.3)\n",
"\n",
"# Q-Q plot\n",
"theoretical_quantiles = stats.chi2.ppf(np.linspace(0.01, 0.99, len(chi2_values)), dof)\n",
"empirical_quantiles = np.sort(chi2_values)\n",
"\n",
"axes[1].scatter(theoretical_quantiles, empirical_quantiles, alpha=0.5, s=20)\n",
"axes[1].plot([theoretical_quantiles.min(), theoretical_quantiles.max()], \n",
" [theoretical_quantiles.min(), theoretical_quantiles.max()], \n",
" 'r--', lw=2, label='Perfect fit')\n",
"axes[1].set_xlabel('Theoretical Quantiles')\n",
"axes[1].set_ylabel('Empirical Quantiles')\n",
"axes[1].set_title('Q-Q Plot: Empirical vs Theoretical χ² Distribution')\n",
"axes[1].legend()\n",
"axes[1].grid(True, alpha=0.3)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nThe empirical distribution closely follows the theoretical χ² distribution,\")\n",
"print(\"validating that PQMass produces well-calibrated statistics under the null hypothesis.\")"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. Workflow 3: Mode-Dropping Detection\n",
"\n",
"Following Section 3.2 and Appendix F.1, we test PQMass ability to detect when modes are missing from a distribution. We use a mixture of Gaussians and drop one mode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"def create_mode_mixture(n_modes=10, dim=5, n_samples=2000, dropped_modes=0):\n",
" \"\"\"\n",
" Create a mixture of Gaussians with optionally dropped modes.\n",
" \n",
" Parameters:\n",
" -----------\n",
" n_modes : int\n",
" Total number of modes in the complete distribution\n",
" dim : int\n",
" Dimensionality\n",
" n_samples : int\n",
" Number of samples to generate\n",
" dropped_modes : int\n",
" Number of modes to drop from the distribution\n",
" \"\"\"\n",
" active_modes = n_modes - dropped_modes\n",
" samples_per_mode = n_samples // active_modes\n",
" \n",
" samples = []\n",
" for i in range(active_modes):\n",
" # Create well-separated modes\n",
" mean = np.zeros(dim)\n",
" mean[i % dim] = 5 * (i // dim + 1) # Spread modes in space\n",
" cov = np.eye(dim) * 0.5\n",
" mode_samples = np.random.multivariate_normal(mean, cov, samples_per_mode)\n",
" samples.append(mode_samples)\n",
" \n",
" return np.vstack(samples)\n",
"\n",
"# Test mode dropping detection\n",
"print(\"Mode-Dropping Detection Experiment\")\n",
"print(\"=\"*70)\n",
"\n",
"n_modes = 10\n",
"dim = 10\n",
"n_samples = 1000 # Smaller for speed\n",
"\n",
"# Generate complete distribution (all 10 modes)\n",
"samples_complete = create_mode_mixture(n_modes=n_modes, dim=dim, n_samples=n_samples, dropped_modes=0)\n",
"\n",
"# Test across different numbers of dropped modes\n",
"dropped_mode_counts = [0, 1, 2, 3]\n",
"results = []\n",
"\n",
"print(f\"\\nTesting mode dropping in {dim}D space with {n_modes} modes...\\n\")\n",
"\n",
"for n_dropped in dropped_mode_counts:\n",
" print(f\"Testing with {n_dropped} dropped mode(s)...\")\n",
" \n",
" # Generate distribution with dropped modes\n",
" samples_dropped = create_mode_mixture(n_modes=n_modes, dim=dim, \n",
" n_samples=n_samples, dropped_modes=n_dropped)\n",
" \n",
" # Compare using PQMass\n",
" result = pqmass_repeated(samples_complete, samples_dropped, \n",
" n_reference=50, n_iterations=10)\n",
" \n",
" results.append({\n",
" 'n_dropped': n_dropped,\n",
" 'mean_chi2': result['mean_chi2'],\n",
" 'std_chi2': result['std_chi2'],\n",
" 'dof': result['dof']\n",
" })\n",
" \n",
" print(f\" Mean χ²: {result['mean_chi2']:.2f} ± {result['std_chi2']:.2f}\")\n",
" print(f\" Expected (null): {result['dof']}\\n\")"
],
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Visualize mode-dropping detection\n",
"fig, ax = plt.subplots(1, 1, figsize=(10, 6))\n",
"\n",
"n_dropped_list = [r['n_dropped'] for r in results]\n",
"chi2_means = [r['mean_chi2'] for r in results]\n",
"chi2_stds = [r['std_chi2'] for r in results]\n",
"expected_chi2 = results[0]['dof']\n",
"\n",
"ax.errorbar(n_dropped_list, chi2_means, yerr=chi2_stds, \n",
" marker='o', markersize=8, linewidth=2, capsize=5,\n",
" label='PQMass χ² statistic')\n",
"ax.axhline(expected_chi2, color='red', linestyle='--', linewidth=2, \n",
" label=f'Expected under null (χ² = {expected_chi2})')\n",
"ax.fill_between(n_dropped_list, expected_chi2 - np.sqrt(2*expected_chi2), \n",
" expected_chi2 + np.sqrt(2*expected_chi2), \n",
" alpha=0.2, color='red', label='±1 std of null')\n",
"\n",
"ax.set_xlabel('Number of Dropped Modes', fontsize=12)\n",
"ax.set_ylabel('Chi-squared Statistic', fontsize=12)\n",
"ax.set_title(f'PQMass Sensitivity to Mode Dropping ({n_modes} modes in {dim}D)', fontsize=14)\n",
"ax.legend(fontsize=10)\n",
"ax.grid(True, alpha=0.3)\n",
"ax.set_xticks(n_dropped_list)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nObservation: As more modes are dropped, the χ² statistic increases,\")\n",
"print(\"indicating that PQMass successfully detects diversity/mode-coverage issues.\")"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 6. Workflow 4: Diffusion Model Training Progress (MNIST Toy Example)\n",
"\n",
"This section demonstrates how PQMass can track generative model training progress. For computational efficiency, we use a **toy example with synthetic data** instead of full MNIST diffusion model training.\n",
"\n",
"**Note:** Full diffusion model training on MNIST would require:\n",
"- GPU acceleration (not available in this environment)\n",
"- Several hours of training time\n",
"- Significant memory (>4GB)\n",
"\n",
"Instead, we simulate the training process to show how PQMass χ² values change as a model improves."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"def simulate_training_progress(n_epochs=20, n_samples=500):\n",
" \"\"\"\n",
" Simulate generative model training progress.\n",
" \n",
" We simulate improving quality by gradually shifting generated samples\n",
" closer to the true distribution.\n",
" \"\"\"\n",
" # \"True\" data distribution (e.g., MNIST test set)\n",
" true_mean = np.array([0, 0])\n",
" true_cov = np.eye(2)\n",
" true_samples = np.random.multivariate_normal(true_mean, true_cov, n_samples)\n",
" \n",
" # Initial \"generated\" distribution (poor quality)\n",
" initial_mean = np.array([3, 3])\n",
" initial_cov = np.eye(2) * 2\n",
" \n",
" chi2_over_epochs = []\n",
" p_values_over_epochs = []\n",
" \n",
" for epoch in range(n_epochs):\n",
" # Simulate improvement: gradually move mean toward true distribution\n",
" progress = epoch / n_epochs\n",
" current_mean = initial_mean * (1 - progress) + true_mean * progress\n",
" current_cov = initial_cov * (1 - progress * 0.5) + true_cov * (progress * 0.5)\n",
" \n",
" # Generate samples from current model state\n",
" generated_samples = np.random.multivariate_normal(current_mean, current_cov, n_samples)\n",
" \n",
" # Evaluate using PQMass\n",
" chi2, p_value, dof = pqmass_test(true_samples, generated_samples, n_reference=30)\n",
" \n",
" chi2_over_epochs.append(chi2)\n",
" p_values_over_epochs.append(p_value)\n",
" \n",
" return {\n",
" 'chi2': chi2_over_epochs,\n",
" 'p_values': p_values_over_epochs,\n",
" 'dof': dof,\n",
" 'true_samples': true_samples,\n",
" 'final_samples': generated_samples\n",
" }\n",
"\n",
"# Run simulation\n",
"print(\"Simulating Generative Model Training Progress\")\n",
"print(\"=\"*70)\n",
"print(\"Note: This is a simplified simulation for demonstration.\")\n",
"print(\"Full MNIST diffusion training would require GPU and hours of compute.\\n\")\n",
"\n",
"training_results = simulate_training_progress(n_epochs=30, n_samples=500)\n",
"\n",
"print(f\"Initial χ²: {training_results['chi2'][0]:.2f}\")\n",
"print(f\"Final χ²: {training_results['chi2'][-1]:.2f}\")\n",
"print(f\"Expected χ² (under null): {training_results['dof']}\")\n",
"print(f\"\\nAs training progresses, χ² decreases toward the expected value,\")\n",
"print(f\"indicating improved match between generated and true distributions.\")"
],
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Visualize training progress\n",
"fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
"\n",
"# Plot chi-squared over epochs\n",
"epochs = np.arange(len(training_results['chi2']))\n",
"axes[0].plot(epochs, training_results['chi2'], marker='o', linewidth=2, markersize=6)\n",
"axes[0].axhline(training_results['dof'], color='red', linestyle='--', linewidth=2, \n",
" label=f'Expected (dof={training_results[\"dof\"]})')\n",
"axes[0].fill_between(epochs, \n",
" training_results['dof'] - np.sqrt(2*training_results['dof']), \n",
" training_results['dof'] + np.sqrt(2*training_results['dof']), \n",
" alpha=0.2, color='red')\n",
"axes[0].set_xlabel('Training Epoch')\n",
"axes[0].set_ylabel('PQMass χ² Statistic')\n",
"axes[0].set_title('Model Quality Improvement During Training')\n",
"axes[0].legend()\n",
"axes[0].grid(True, alpha=0.3)\n",
"\n",
"# Plot sample distributions (initial vs final)\n",
"true_samples = training_results['true_samples']\n",
"final_samples = training_results['final_samples']\n",
"\n",
"axes[1].scatter(true_samples[:, 0], true_samples[:, 1], \n",
" alpha=0.5, s=20, label='True distribution')\n",
"axes[1].scatter(final_samples[:, 0], final_samples[:, 1], \n",
" alpha=0.5, s=20, label='Generated (final epoch)')\n",
"axes[1].set_xlabel('Dimension 1')\n",
"axes[1].set_ylabel('Dimension 2')\n",
"axes[1].set_title('True vs Generated Distributions (After Training)')\n",
"axes[1].legend()\n",
"axes[1].grid(True, alpha=0.3)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 7. Workflow 5: Baseline Comparisons (MMD, Wasserstein)\n",
"\n",
"The paper compares PQMass against several baseline metrics including Maximum Mean Discrepancy (MMD) and Wasserstein distance. Let's implement these for comparison."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": "def compute_mmd_rbf(X, Y, gamma=1.0):\n \"\"\"\n Compute Maximum Mean Discrepancy with RBF kernel.\n \n MMD² = E[k(x,x')] - 2E[k(x,y)] + E[k(y,y')]\n \"\"\"\n def rbf_kernel(X, Y, gamma):\n pairwise_sq_dists = cdist(X, Y, metric='sqeuclidean')\n return np.exp(-gamma * pairwise_sq_dists)\n \n K_XX = rbf_kernel(X, X, gamma)\n K_YY = rbf_kernel(Y, Y, gamma)\n K_XY = rbf_kernel(X, Y, gamma)\n \n mmd_sq = K_XX.mean() - 2 * K_XY.mean() + K_YY.mean()\n return np.sqrt(max(mmd_sq, 0)) # Ensure non-negative\n\n\ndef compute_mmd_linear(X, Y):\n \"\"\"\n Compute Maximum Mean Discrepancy with linear kernel.\n \"\"\"\n mean_X = X.mean(axis=0)\n mean_Y = Y.mean(axis=0)\n return np.linalg.norm(mean_X - mean_Y)\n\n\ndef compute_wasserstein_approx(X, Y, p=2):\n \"\"\"\n Compute approximate Wasserstein distance using sliced approach.\n \n For exact W2, we'd need optimal transport solvers (POT library),\n but this is computationally expensive. We use a simple approximation.\n \"\"\"\n # Simple approximation: compare sorted 1D projections\n n_projections = 50\n dim = X.shape[1]\n distances = []\n \n # Ensure same number of samples by using minimum length\n n_samples = min(len(X), len(Y))\n \n for _ in range(n_projections):\n # Random projection direction\n direction = np.random.randn(dim)\n direction /= np.linalg.norm(direction)\n \n # Project samples (use only first n_samples from each)\n X_proj = X[:n_samples] @ direction\n Y_proj = Y[:n_samples] @ direction\n \n # 1D Wasserstein distance (sorting-based)\n X_sorted = np.sort(X_proj)\n Y_sorted = np.sort(Y_proj)\n w_dist = np.mean(np.abs(X_sorted - Y_sorted)**p)**(1/p)\n distances.append(w_dist)\n \n return np.mean(distances)\n\nprint(\"Baseline metric functions implemented successfully!\")",
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Compare metrics on mode-dropping experiment\n",
"print(\"Comparing PQMass with Baseline Metrics\")\n",
"print(\"=\"*70)\n",
"\n",
"n_modes = 5\n",
"dim = 10\n",
"n_samples = 500\n",
"\n",
"# Generate complete distribution\n",
"samples_complete = create_mode_mixture(n_modes=n_modes, dim=dim, n_samples=n_samples, dropped_modes=0)\n",
"\n",
"# Test with varying dropped modes\n",
"dropped_modes_range = [0, 1, 2]\n",
"comparison_results = []\n",
"\n",
"print(\"Computing metrics for different levels of mode dropping...\\n\")\n",
"\n",
"for n_dropped in dropped_modes_range:\n",
" samples_dropped = create_mode_mixture(n_modes=n_modes, dim=dim, \n",
" n_samples=n_samples, dropped_modes=n_dropped)\n",
" \n",
" # PQMass\n",
" pqmass_result = pqmass_repeated(samples_complete, samples_dropped, \n",
" n_reference=30, n_iterations=5)\n",
" \n",
" # MMD (RBF and Linear)\n",
" mmd_rbf = compute_mmd_rbf(samples_complete, samples_dropped, gamma=1.0/dim)\n",
" mmd_linear = compute_mmd_linear(samples_complete, samples_dropped)\n",
" \n",
" # Wasserstein (approximation)\n",
" w2_dist = compute_wasserstein_approx(samples_complete, samples_dropped)\n",
" \n",
" comparison_results.append({\n",
" 'n_dropped': n_dropped,\n",
" 'pqmass_chi2': pqmass_result['mean_chi2'],\n",
" 'pqmass_dof': pqmass_result['dof'],\n",
" 'mmd_rbf': mmd_rbf,\n",
" 'mmd_linear': mmd_linear,\n",
" 'wasserstein': w2_dist\n",
" })\n",
" \n",
" print(f\"Dropped modes: {n_dropped}\")\n",
" print(f\" PQMass χ²: {pqmass_result['mean_chi2']:.2f} (expected: {pqmass_result['dof']})\")\n",
" print(f\" MMD (RBF): {mmd_rbf:.4f}\")\n",
" print(f\" MMD (Linear): {mmd_linear:.4f}\")\n",
" print(f\" Wasserstein: {w2_dist:.4f}\\n\")"
],
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"source": [
"# Visualize metric comparison\n",
"fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n",
"\n",
"n_dropped_vals = [r['n_dropped'] for r in comparison_results]\n",
"\n",
"# PQMass\n",
"pqmass_vals = [r['pqmass_chi2'] for r in comparison_results]\n",
"expected_chi2 = comparison_results[0]['pqmass_dof']\n",
"axes[0, 0].plot(n_dropped_vals, pqmass_vals, marker='o', linewidth=2, markersize=8)\n",
"axes[0, 0].axhline(expected_chi2, color='red', linestyle='--', label='Expected (null)')\n",
"axes[0, 0].set_xlabel('Number of Dropped Modes')\n",
"axes[0, 0].set_ylabel('χ² Statistic')\n",
"axes[0, 0].set_title('PQMass')\n",
"axes[0, 0].legend()\n",
"axes[0, 0].grid(True, alpha=0.3)\n",
"\n",
"# MMD RBF\n",
"mmd_rbf_vals = [r['mmd_rbf'] for r in comparison_results]\n",
"axes[0, 1].plot(n_dropped_vals, mmd_rbf_vals, marker='s', linewidth=2, markersize=8, color='orange')\n",
"axes[0, 1].set_xlabel('Number of Dropped Modes')\n",
"axes[0, 1].set_ylabel('MMD Value')\n",
"axes[0, 1].set_title('MMD (RBF Kernel)')\n",
"axes[0, 1].grid(True, alpha=0.3)\n",
"\n",
"# MMD Linear\n",
"mmd_linear_vals = [r['mmd_linear'] for r in comparison_results]\n",
"axes[1, 0].plot(n_dropped_vals, mmd_linear_vals, marker='^', linewidth=2, markersize=8, color='green')\n",
"axes[1, 0].set_xlabel('Number of Dropped Modes')\n",
"axes[1, 0].set_ylabel('MMD Value')\n",
"axes[1, 0].set_title('MMD (Linear Kernel)')\n",
"axes[1, 0].grid(True, alpha=0.3)\n",
"\n",
"# Wasserstein\n",
"w2_vals = [r['wasserstein'] for r in comparison_results]\n",
"axes[1, 1].plot(n_dropped_vals, w2_vals, marker='d', linewidth=2, markersize=8, color='purple')\n",
"axes[1, 1].set_xlabel('Number of Dropped Modes')\n",
"axes[1, 1].set_ylabel('Distance')\n",
"axes[1, 1].set_title('Wasserstein Distance (Approx)')\n",
"axes[1, 1].grid(True, alpha=0.3)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nAll metrics increase with more dropped modes, but PQMass provides\")\n",
"print(\"an interpretable χ² statistic with known expected value under the null hypothesis.\")"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 8. Conclusion and Scaling Guidance\n",
"\n",
"### Summary\n",
"\n",
"This notebook demonstrated the **PQMass** methodology for evaluating generative models:\n",
"\n",
"1. **Core Algorithm**: Voronoi tessellation-based two-sample test with interpretable χ² statistics\n",
"2. **Null Test Validation**: Verified that χ² follows expected distribution when comparing identical distributions\n",
"3. **Mode-Dropping Detection**: Showed PQMass can detect diversity issues when modes are missing\n",
"4. **Training Progress Tracking**: Demonstrated monitoring model improvement during training\n",
"5. **Baseline Comparisons**: Compared PQMass with MMD and Wasserstein distance\n",
"\n",
"### Key Advantages of PQMass\n",
"\n",
"- ✅ **No feature extraction needed**: Works directly in sample space\n",
"- ✅ **Computationally efficient**: Scales better than optimal transport methods\n",
"- ✅ **Interpretable statistics**: χ² statistic with known expected distribution\n",
"- ✅ **Detects both fidelity and diversity issues**: Sensitive to both quality and mode coverage\n",
"- ✅ **Flexible distance metrics**: Can use custom metrics (e.g., Levenshtein for sequences)\n",
"\n",
"---\n",
"\n",
"## Scaling to Full Experiments\n",
"\n",
"This notebook used small-scale examples due to resource constraints (4GB RAM, no GPU, 20-minute time limit). To replicate the full paper experiments:\n",
"\n",
"### 1. **Image Generation (MNIST, CIFAR-10, FFHQ)**\n",
" - **Requirements**: GPU with 8GB+ VRAM, 16GB+ RAM\n",
" - **Training time**: 5-20 hours per model\n",
" - **Sample size**: 10,000-50,000 images per distribution\n",
" - **PQMass parameters**: n_reference=200-500, n_iterations=20\n",
" - **Code changes**: Load full datasets (torchvision), train actual diffusion/VAE models\n",
"\n",
"### 2. **High-Dimensional Pixel Space (Astrophysics Images)**\n",
" - **Requirements**: 32GB+ RAM, GPU for model training\n",
" - **Data**: Download Probes galaxy dataset or JWST dark images\n",
" - **Flatten images**: Convert 64×64 or 128×128 images to 4096-16384D vectors\n",
" - **PQMass computation**: May take 30-60 minutes per comparison\n",
"\n",
"### 3. **Baseline Metric Comparisons**\n",
" - **FID**: Install `pytorch-fid` package, requires InceptionV3 model\n",
" - **FLD**: Requires DinoV2 or InceptionV3 features + GMM fitting (scikit-learn)\n",
" - **Exact Wasserstein**: Install `POT` (Python Optimal Transport) library\n",
" - **Note**: Feature-based metrics require pretrained models (1-5GB download)\n",
"\n",
"### 4. **Computational Cost Scaling**\n",
" - **PQMass**: O(n × n_ref) where n=sample size, n_ref=reference points\n",
" - **MMD (RBF)**: O(n²) - becomes slow for n>5000\n",
" - **Optimal Transport**: O(n³) or O(n² log n) - very slow for large n\n",
" - **Recommendation**: For n>10000, use n_ref=200-500 for PQMass\n",
"\n",
"### 5. **Protein Sequences and Other Data Types**\n",
" - **Levenshtein distance**: Install `python-Levenshtein` package\n",
" - **Custom metrics**: Modify `pqmass_test()` to accept custom distance functions\n",
" - **Example**: `pqmass_test(seqs_p, seqs_q, metric=levenshtein_distance)`\n",
"\n",
"### 6. **Reproducibility**\n",
" - Set random seeds consistently: `np.random.seed(42)`, `torch.manual_seed(42)`\n",
" - Use multiple tessellations (n_iterations=20) to reduce variance\n",
" - Report mean ± std for all metrics\n",
"\n",
"---\n",
"\n",
"## Additional Resources\n",
"\n",
"- **Paper**: arXiv link (see paper PDF)\n",
"- **Code Repository**: Check paper for official implementation\n",
"- **Datasets**: \n",
" - MNIST/CIFAR-10: `torchvision.datasets`\n",
" - FFHQ: https://github.com/NVlabs/ffhq-dataset\n",
" - Adult Census: UCI ML Repository\n",
"\n",
"---\n",
"\n",
"## Next Steps for Researchers\n",
"\n",
"1. **Test on your own data**: Replace synthetic examples with your actual datasets\n",
"2. **Tune n_reference**: Try values from 50-500 based on sample size and dimensionality\n",
"3. **Experiment with distance metrics**: Try L1, L2, or custom metrics for your domain\n",
"4. **Compare with baselines**: Implement FID/FLD to compare with PQMass results\n",
"5. **Track training**: Integrate PQMass into your training loop for real-time monitoring\n",
"\n",
"**Happy experimenting with PQMass!** 🎉"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.0"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment