Information-Theoretic Diffusion Models for Image Inpainting / 基于信息论的图像修复扩散模型

Channel-Aware Loss, Conditional Independence Regularization, Posterior-Optimal Sampling, Universal Non-Gaussian Diffusion

Abstract

This project implements an information-theoretic framework for image inpainting using diffusion models. The core contributions include:

1. Channel-Aware Loss: A loss function derived from the channel-induced distortion measure $\rho(z, y) = -\log p_{Z|X}(z|y)$, which explicitly encodes the probabilistic structure of the forward diffusion process.
2. Conditional Independence Regularization: Minimizing $I(X_0; \hat{X}_0 | X_t)$ to enforce the denoised output's statistical independence from the ground truth given the noisy observation.
3. Posterior-Optimal Sampling: A sampling strategy that minimizes channel-induced distortion during inference for improved inpainting quality.
4. Universal Non-Gaussian Diffusion: Extension from additive Gaussian noise to arbitrary discrete memoryless channels (erasure, impulse, Poisson).

Installation

Pip

pip install -r requirements.txt

Conda

conda env create -f environment.yml
conda activate inpainting

Docker

docker build -t inpainting-diffusion .
docker run --gpus all --rm -it inpainting-diffusion train.py --config configs/cifar10_experiment.yaml

Requirements

• Python >= 3.9
• PyTorch >= 2.0.0
• CUDA (recommended for training)

Project Structure

inpainting_experiment/
├── configs/                    # Experiment configuration files (YAML)
│   ├── default.yaml            # Default CelebA configuration
│   ├── cifar10_experiment.yaml # CIFAR-10 experiments
│   ├── mse_baseline.yaml       # MSE baseline
│   ├── hybrid_loss.yaml        # Hybrid loss (MSE + channel-aware)
│   ├── channel_aware.yaml      # Channel-aware loss
│   ├── non_gaussian_erasure.yaml  # Non-Gaussian erasure channel
│   ├── celeba_hq.yaml             # CelebA-HQ 256x256 config
│   ├── imagenet.yaml              # ImageNet config
│   ├── ablation_mi_0.yaml      # Ablation: lambda_mi = 0
│   ├── ablation_mi_01.yaml     # Ablation: lambda_mi = 0.1
│   └── conference_*.yaml       # Conference submission configs
├── data/
│   └── dataset.py              # Dataset loading and mask generation
├── models/
│   ├── ddpm.py                 # DDPM scheduler, channel-aware loss, CI regularization
│   ├── unet.py                 # UNet denoising architecture
│   ├── inpainting_sampler.py   # Inpainting samplers (resample, posterior-optimal, RePaint, DDIM)
│   ├── ddnm_inpainting.py      # DDNM zero-shot inpainting
│   ├── palette_inpainting.py   # Palette conditional diffusion inpainting
│   └── non_gaussian_diffusion.py  # Non-Gaussian diffusion framework
├── utils/
│   ├── config.py               # Configuration management
│   ├── visualization.py        # Visualization utilities
│   └── seed.py                 # Random seed management for reproducibility
├── train.py                    # Training script (Gaussian DDPM)
├── train_non_gaussian.py       # Training script (non-Gaussian diffusion)
├── evaluate.py                 # Evaluation script (FID, LPIPS, PSNR, SSIM)
├── evaluate_non_gaussian.py    # Evaluation for non-Gaussian models
├── benchmark_efficiency.py     # Inference efficiency benchmarking
├── quick_test.py               # Quick sanity check
├── run_tests.py                # Run all tests (unittest wrapper)
├── run_ablation.py             # Run ablation studies
├── requirements.txt            # Python dependencies
├── environment.yml             # Conda environment
├── Dockerfile                  # Container image definition
└── docker-compose.yml          # Container orchestration

Quick Start

Running Tests

# Run all unit tests
python run_tests.py

Running Ablation Studies

# Run all ablation experiments defined in the script
python run_ablation.py

# Dry run (print commands only)
python run_ablation.py --dry-run

Training

# Train with channel-aware loss (default)
python train.py --config configs/default.yaml --output_dir outputs/channel_aware

# Train with MSE baseline
python train.py --config configs/mse_baseline.yaml --output_dir outputs/mse_baseline

# Train with hybrid loss
python train.py --config configs/hybrid_loss.yaml --output_dir outputs/hybrid_loss

# Train non-Gaussian diffusion (erasure channel)
python train_non_gaussian.py --config configs/non_gaussian_erasure.yaml \
    --noise_type erasure --output_dir outputs/non_gaussian_erasure

# Quick demo on CIFAR-10
python train.py --config configs/demo_experiment.yaml --output_dir outputs/demo

Evaluation

# Evaluate with posterior-optimal sampling
python evaluate.py --config configs/default.yaml \
    --checkpoint outputs/channel_aware/checkpoints/best_model.pt \
    --sampling_method posterior_optimal --save_images

# Compare sampling methods
for method in posterior_optimal resample repaint ddim ddnm palette; do
    python evaluate.py --config configs/default.yaml \
        --checkpoint outputs/channel_aware/checkpoints/best_model.pt \
        --sampling_method $method \
        --output_dir eval_results/${method}
done

Efficiency Benchmarking

python benchmark_efficiency.py --config configs/default.yaml \
    --checkpoint outputs/channel_aware/checkpoints/best_model.pt

Experiment Results

Sampling Method Comparison (CelebA, 256x256, 50% mask)

Ablation: MI Regularization Coefficient

Configuration Guide

All experiments are configured via YAML files in configs/. Key parameters:

• training.loss_type: mse | channel_aware | hybrid
• training.lambda_mi: MI regularization coefficient (0 to disable)
• training.seed: Random seed for reproducibility (default: 42)
• testing.sampling_method: posterior_optimal | resample | repaint | ddim
• model.in_channels: 7 (3 image + 1 mask + 3 masked_image)

Citation

@inproceedings{infotheoretic_inpainting2026,
  title={Information-Theoretic Diffusion Models for Image Inpainting},
  author={TBD},
  booktitle={TBD},
  year={2026}
}

ACM MM Support Materials

• Curated support-material documentation: docs/support_materials/README.md
• Curated reproducible artifacts and manifests: artifacts/
• Restructure tooling:
  • python tools/restructure/restructure_for_acmmm.py --dry-run
  • python tools/restructure/restructure_for_acmmm.py
• Build support package:
  • python tools/restructure/build_support_package.py

CVPR 2027 基准流水线（baselines_runner/）

512×512 统一协议下的全基线对比实验框架（2026-08-19 起）：

• 官方源码基线：9 个扩散基线 + LaMa/MAT 前馈参照，全部 vendored 于 external/（commit 固定，见 external/manifest.json），经 baselines_runner/wrappers/ 子进程 wrapper 接入统一契约；
• 统一裁判：所有方法输出由 baselines_runner/judge.py 同一套 FID/LPIPS/PSNR/SSIM/L1 实现打分；
• 数据集：ImageNet val 1k / FFHQ 1k / CelebA-HQ 2k / Places2 2k / DIV2K 100（全部 512²，分配见 docs/datasets.md）；
• Mask：PConv 协议六档不规则 mask（assets/masks/PROVENANCE.md）。

bash baselines_runner/run_pipeline.sh                      # 全流水线（串行、断点续跑）
python baselines_runner/run_benchmark.py --matrix baselines_runner/matrix/imagenet512.yaml --dry-run
python -m pytest tests/ -q                                 # 79 项契约/单元测试

详细文档：docs/baselines.md（12 个 wrapper 的官方来源/NFE/mask 语义/ patch 索引）、docs/datasets.md（数据集分配与协议）。