Diffusion Optimization Models with Trajectory Alignment

Generative models have had a profound impact on vision and language, paving the way for a new era of multimodal generative applications. While these successes have inspired researchers to explore using generative models in science and engineering to accelerate the design process and reduce the reliance on iterative optimization, challenges remain. Specifically, engineering optimization methods based on physics still outperform generative models when dealing with constrained environments where data is scarce and precision is paramount. To address these challenges, we introduce Diffusion Optimization Models (DOM) and Trajectory Alignment (TA), a learning framework that demonstrates the efficacy of aligning the sampling trajectory of diffusion models with the optimization trajectory derived from traditional physics-based methods.

This repo contains code and experiments for:

Aligning Optimization Trajectories with Diffusion Models for Constrained Design Generation
Giorgio Giannone, Akash Srivastava, Ole Winther, Faez Ahmed
Conference on Neural Information Processing Systems (NeurIPS), 2023

[paper] [code] [page]

Diffusing the Optimal Topology: A Generative Optimization Approach
Giorgio Giannone, Faez Ahmed
International Design Engineering Technical Conferences (IDETC), 2023

[paper] [code]

---

Installation

conda create -n dom python=3.8
conda activate dom

git clone https://github.com/georgosgeorgos/trajectory-alignment-diffusion/
cd trajectory-alignment-diffusion

The code has been tested on Ubuntu 20.04, Pytorch 1.13, and CUDA 11.6

---

Dataset

• 2d topology optimization dataset 64x64 - here

---

Evaluation

We use the benchmark provided in TopoDiff. Follow the instructions here to download the evaluation set.

---

Train the model

TRAIN_FLAGS="--batch_size 32 --save_interval 10000 --use_fp16 True"
MODEL_FLAGS="--image_size 64 --num_channels 128 --num_res_blocks 3 --learn_sigma True --dropout 0.3"
DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule cosine"

DATA_FLAGS="--data_dir ./dom_dataset/"


CUDA_VISIBLE_DEVICES=0 \
python scripts/image_train_intermediate_kernel.py $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS $DATA_FLAGS

---

Sample the model

#! /bin/sh
MODEL_FLAGS="--image_size 64 --num_channels 128 --num_res_blocks 3 --learn_sigma True --dropout 0.3 --use_fp16 True"
DIFFUSION_FLAGS="--diffusion_steps 1000 --timestep_respacing 100 --noise_schedule cosine"
DATA_FLAGS="--constraints_path ./dom_dataset/test_data/ --num_samples 1800"
CHECKPOINTS_FLAGS="--model_path ./dom_logdir/ema_0.9999_xxxxx.pt"


CUDA_VISIBLE_DEVICES=0 \
python scripts/sample_kernel_relaxation.py $MODEL_FLAGS $DIFFUSION_FLAGS $DATA_FLAGS $CHECKPOINTS_FLAGS

---

Acknowledgments

A lot of code and ideas were borrowed from:

• https://github.com/francoismaze/topodiff
• https://github.com/openai/improved-diffusion

Citation

@article{giannone2023aligning,
  title={Aligning Optimization Trajectories with Diffusion Models for Constrained Design Generation},
  author={Giannone, Giorgio and Srivastava, Akash and Winther, Ole and Ahmed, Faez},
  journal={arXiv preprint arXiv:2305.18470},
  year={2023}
}

@article{giannone2023diffusing,
  title={Diffusing the optimal topology: A generative optimization approach},
  author={Giannone, Giorgio and Ahmed, Faez},
  journal={arXiv preprint arXiv:2303.09760},
  year={2023}
}