Value Gradient Sampler (VGS)

Official implementation of Value Gradient Sampler: Learning Invariant Value Functions for Equivariant Diffusion Sampling (AISTATS 2026).

We recommend using Conda:

# Clone the repository
git clone https://github.com/swyoon/value-gradient-sampler.git
cd value-gradient-sampler

# Create the environment
conda create -n vgs python=3.10 -y && conda activate vgs  
                                                                                                
# Install PyTorch matching your CUDA driver (see https://pytorch.org/get-started/locally/)      
# Example for CUDA 12.8 (driver >= 570):                                                        
pip install torch --index-url https://download.pytorch.org/whl/cu128  
  
# Install the rest
pip install -r requirements.txt

We evaluate VGS on two families of targets (see the paper for details):

• n-body particle systems: DW-4, LJ-13, LJ-55
• Synthetic distributions: 25GMM, Funnel, Manywell

Train VGS on a target distribution with:

python train.py --config configs/<target_name>.yaml --device 0 --run <run_name>

where <target_name> is one of dw4, lj13, lj55, gmm, funnel, manywell. Configuration files, logs, and checkpoints are written to:

results/<target_name>/<run_name>/   # e.g. results/lj13/test/

Training progress can be monitored with Weights & Biases.

Any field in the YAML config can be overridden by passing --<dotted.key> <value> on the command line. For example:

# Train for 1000 iterations with batch size 512
python train.py --config configs/lj13.yaml --device 0 --run test \
    --training.n_iter 1000 --training.batch_size 512

The three LJ-13 configs correspond to the value-network ablation in Figure 2 of the paper:

Checkpoints are saved automatically and include the full training state — model weights, optimizer state, and RNG states — so resumed runs reproduce the original training dynamics exactly.

# Resume from a specific checkpoint
python train.py --config configs/<target_name>.yaml --device 0 --run <run_name> \
    --resume /path/to/checkpoint.pth

# Auto-resume from the latest checkpoint in the run's log directory
python train.py --config configs/<target_name>.yaml --device 0 --run <run_name> \
    --resume auto

Checkpoint files in the log directory:

Below is a brief summary of the evaluation metrics reported in the paper. See the paper for full definitions.

We report Total Variation Distance of Interatomic Distance and Energy (TVD-D, TVD-E) and Wasserstein Distance ($\mathcal{W}^2$), computed from 10,000 samples. The checkpoint with the lowest validation TVD-D (exp_best.pth) is saved to the log directory during training. To compute the metrics, open the evaluation notebooks in notebooks/ and set sampler_ckpt to your checkpoint path (replace <LOGDIR> in the notebook with the run's log directory, e.g., results/lj13/test).

We report Log-Partition Function Error ($|\Delta \log Z_r|, |\Delta \log Z_f|$) and Wasserstein Distance ($\mathcal{W}^2$), computed from 2,000 samples. These metrics are written to results.txt in the log directory at the end of training.

If you find this repository useful for your research, please cite our paper:

@inproceedings{hwang2025value,
  title={Value Gradient Sampler: Learning Invariant Value Functions for Equivariant Diffusion Sampling},
  author={Hwang, Himchan and Jeong, Hyeokju and Shin, Dong Kyu and Park, Che-Sang and Kweon, Sehee and Yoon, Sangwoong and Park, Frank C},
  booktitle={The 29th International Conference on Artificial Intelligence and Statistics},
  year={2025}
}