zimowang/Distance-Marching

Distance Marching checkpoints

Checkpoints corresponding to the main CIFAR-10 and ImageNet-256 B/2 results of Distance Marching. CIFAR-10 is the unmodified original training checkpoint. ImageNet is an EMA inference checkpoint exported from the original training checkpoint; all 143 EMA tensors and the training arguments were checked for exact equality to the source. Its optimizer state and unused non-EMA training weights are omitted.

Checkpoints and recorded results

Checkpoint	Dataset / model	Sampler	Steps	Step size	CFG	Recorded FID
cifar10-checkpoint-1809.pth	CIFAR-10	Sphere tracing	250	0.031	Unconditional	2.234152 (seed 0)
cifar10-checkpoint-1809.pth	CIFAR-10	Gradient descent	250	0.46	Unconditional	2.545608 (seed 0)
imagenet-b2-0400000-ema.pt	ImageNet-256, B/2	Sphere tracing	250	0.022	2.0	6.589967 (mean of seeds 0, 1, 2)
imagenet-b2-0400000-ema.pt	ImageNet-256, B/2	Gradient descent	250	0.64	2.0	6.670391 (mean of seeds 0, 1, 2)

Each evaluation uses 50,000 generated samples per seed. The CIFAR-10 sphere-tracing result corresponds to the paper's 2.23. The ImageNet means correspond to 6.59 and 6.67. The CIFAR-10 GD value is reported here directly from the saved log (the paper table displays 2.54).

For ImageNet ST, the three recorded FIDs are 6.562257, 6.597062, and 6.610580. For ImageNet GD they are 6.614579, 6.687854, and 6.708741. Exact values, recomputed means and standard deviations, and source-log paths are in evaluation_results.json. The recomputed sample standard deviations are 0.024931 and 0.049451; the manuscript reports 0.027 and 0.053, respectively.

These are historical evaluation records recovered from the original persistent volume. This release verifies checkpoint identity and transfer integrity; it does not rerun image generation or FID evaluation.

Source identity

• CIFAR-10: /pv/flow_matching/output_dir/no_ot_best/checkpoint-1809.pth, saved epoch index 1809 (1810 completed epochs).
• ImageNet B/2: /pv/EqM/results/001-EqM-B-2-Linear-dist-None/checkpoints/0400000.pt, training step 400,000. The EqM directory and model argument are historical implementation names for this Distance Marching experiment.

Both checkpoints serve their ST and GD samplers; a separate sampler-specific checkpoint is not needed. The CIFAR-10 checkpoint stores an EMA wrapper in checkpoint["model"] (including model.*, shadow_params.*, and num_updates). The ImageNet sampler uses checkpoint["ema"]. In the exported ImageNet file, both model and ema refer to those same original EMA weights so the original sampler can load it unchanged. This export is for inference, not full optimizer-state training resumption. Preserve the EMA behavior when reproducing results.

File sizes and SHA-256 values are in checkpoint_manifest.json and SHA256SUMS. The manifest also records the original ImageNet source checkpoint hash; the smaller inference export has its own hash because its serialization and included state differ.

Download and inspect

from huggingface_hub import hf_hub_download
import argparse
import torch

path = hf_hub_download(
    repo_id="zimowang/Distance-Marching",
    filename="cifar10-checkpoint-1809.pth",
)
# Original checkpoints include argparse.Namespace training configuration.
with torch.serialization.safe_globals([argparse.Namespace]):
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)

Use filename="imagenet-b2-0400000-ema.pt" for the ImageNet checkpoint.

Sampling

Use the corresponding Distance Marching implementation. The manuscript links the CIFAR-10 code and ImageNet code. These checkpoints require the custom distance predictor and EMA implementation; they are not standard Diffusers pipelines.

For the original CIFAR-10 code, run from its repository root and replace PATH_TO_CHECKPOINT with the downloaded file:

python examples/image/train.py --dataset cifar10 \
  --use_distance_marching --batch_size 128 --accum_iter 1 --epochs 3000 \
  --class_drop_prob 1.0 --cfg_scale 0.0 --use_ema --timestamp \
  --skewed_timesteps --edm_schedule --compute_fid \
  --ode_options '{"nfe": 250}' --separate_t --max_period 1000000 \
  --dm_method st --dm_criteria fixed_steps --c 4.0 --unified_scale 120.0 \
  --dm_step_size 0.031 --fid_samples 50000 --save_fid_samples \
  --seed 0 --eval_only --resume PATH_TO_CHECKPOINT --output_dir ./samples_st

For GD, use --dm_method gd --dm_step_size 0.46.

For the original ImageNet code:

torchrun --nnodes=1 --nproc_per_node=1 sample_gd.py \
  --model EqM-B/2 --ckpt PATH_TO_CHECKPOINT --global-batch-size 256 \
  --num-fid-samples 50000 --folder samples_st --ebm none \
  --sampler st --stepsize 0.022 --cfg-scale 2.0 \
  --num-sampling-steps 250 --global-seed 0 --vae ema

For GD, use --sampler gd --stepsize 0.64. The sampler restores the distance-head configuration from the checkpoint arguments and uses the frozen stabilityai/sd-vae-ft-ema VAE. The VAE and benchmark datasets are not included in this repository.

Acknowledgments

The implementations build on Meta's Flow Matching code and the EqM/SiT codebases. Their original code and the pretrained VAE retain their respective upstream licenses.