p-Bridge — flagship conditional checkpoint

Weights for The Noise Plane: One Diffusion Bridge Across Exponent and Dispersion. Code: Teriri1999/p-Bridge.

p-Bridge indexes a diffusion bridge by two physical quantities — the power-variance exponent p and the cumulative measurement weight ν — instead of an abstract noise schedule. This is the flagship conditional model: a single network spanning the whole Tweedie axis p ∈ [1, 2], told the exponent at inference rather than retrained for it.

Usage

git clone https://github.com/Teriri1999/p-Bridge.git && cd p-Bridge
pip install -r requirements.txt

hf download Teriri1999/p-Bridge pbridge_cond_severe_50k.pt --local-dir weights

python -m eval.run_eval \
    --ckpt weights/pbridge_cond_severe_50k.pt \
    --nfe 5 --eval_p 1.0 1.5 2.0 \
    --out_json results/eval.json

--eval_p is the exponent the test data is corrupted with; the model is told it at inference. Results are quoted at NFE = 5 throughout.

The noise plane

Noise is located by a point (p, ν) on a plane rather than by a scalar level: the variance law Var(X | μ) = μ^p / ν separates how strongly the fluctuation depends on the signal (p) from how much measurement went into it (ν). Gaussian noise, Poisson shot noise, photon-counting detectors and SAR speckle are columns of one family rather than four problems.

Because ν carries p-dependent units, runs are calibrated by a severity η := CV(μ_ref) rather than by a shared ν. This checkpoint is trained at η = 1.0 (severe) with μ_ref = 0.464864; milder noise is handled by entering the same bridge part-way along, not by a different checkpoint.

Training configuration

torchrun --nproc_per_node=4 train.py \
    --p 1.5 --p_range 1.0 2.0 --p_quantum 0.05 \
    --eta 1.0 --T 100 --schedule log \
    --image_size 256 --model_channels 64 --cond_x1 \
    --nu_weight_alpha 0.0 \
    --lambda_rec 10 --lambda_mom 1 --lambda_cons 5 \
    --batch_size 2 --grad_accum 3 --total_steps 50000 \
    --lr_G 1e-4 --lr_schedule cosine --lr_warmup 1000 --lr_min 1e-7 \
    --ema_decay 0.999 --grad_clip 1.0 \
    --dataset_length 16000 --num_workers 4 \
    --extra_roots data/div2k/DIV2K_train_HR --seed 0

License

MIT. The network backbone is derived from OpenAI guided-diffusion, also MIT.