Metropolis Sampling for Constrained Diffusion Models

Authors: Nic Fishman, Leo Klarner, Emile Mathieu, Michael Hutchinson, Valentin De Bortoli

NeurIPS 2023 | Conference PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental To demonstrate the practical utility and empirical performance of the proposed Metropolis diffusion models, we conduct a comprehensive evaluation on a range of synthetic and real-world tasks. In Section 5.1, we assess the scalability of our method by applying it to synthetic distributions on hypercubes and simplices of increasing dimensionality. In Section 5.2, we extend the evaluation to real-world tasks on manifolds with convex constraints by applying our method to the robotics and protein design datasets presented in Fishman et al. (2023). In Section 5.3, we additionally demonstrate that our method extends to constrained manifolds with highly non-convex boundaries a setting that is intractable with existing approaches.
Researcher Affiliation Academia Nic Fishman University of Oxford Leo Klarner University of Oxford Emile Mathieu University of Cambridge Michael Hutchinson University of Oxford Valentin De Bortoli ENS Ulm
Pseudocode Yes Algorithm 1 Metropolis approx. of RBM, Algorithm 2 Rejection approx. of RBM, Algorithm 3 Reflected Step Algorithm., Algorithm 4 Reflected Random Walk., Algorithm 5 Metropolis Random Walk.
Open Source Code Yes All source code that is needed to reproduce the results presented below is made available under https://github.com/oxcsml/score-sde/tree/metropolis, which requires a supporting package to handle the different geometries that is available under https://github.com/oxcsml/geomstats/tree/polytope.
Open Datasets Yes To this end, we derived a geospatial dataset based on wildfire incidence rates within the continental United States (see Appendix E for full details) and trained a Metropolis diffusion model constrained by the corresponding country borders on the sphere S2. Specifically, we retrieved the rasterised version of the wildfire data provided by Welty et al. (2020), converted it to a spherical geodetic coordinate system using the CARTOPY library (Met Office, 2010 2015), and drew a weighted subsample of size 1 106.
Dataset Splits Yes A quantitative comparison of the log-likelihood of a held-out test set is presented in Table 2, while a visual comparison is postponed to Appendix F.3. We selected these hyperparameters from a systematic search over learning rates (6 10 4, 2 10 4, 6 10 5, 2 10 5), learning rate schedules (cosine, log-linear), and batch sizes (128, 256, 512, 1024) on synthetic examples for the reflected and log-barrier models.
Hardware Specification Yes All models were trained on a single NVIDIA Ge Force GTX 1080 GPU.
Software Dependencies No The paper mentions software components like 'CARTOPY library', 'GEOVIEWS', and 'DATASHADER' but does not specify their version numbers or the versions of other key libraries used for model implementation (e.g., PyTorch, TensorFlow).
Experiment Setup Yes For all experiments, we use a simple 6-layer MLP with sine activations and a score rescaling function... We set T = 1, β0 = 1 10 3 and tune β1 to ensure that the forward process reaches the invariant distribution with a linear β-schedule. We use a learning rate of 2 10 4 with a cosine learning rate schedule and an ism loss with a modified loss weighting function of (1 + t), a batch size of 1024 and 8 repeats per batch.