Graphically Structured Diffusion Models

TL;DR

This paper presents a framework for designing deep generative models with problem-specific structures, leveraging graphical models and diffusion architectures to improve scalability and accuracy in complex problem domains.

Contribution

It introduces a novel approach to tailor diffusion models to structured problems using graphical models, enhancing performance and scalability.

Findings

Improved scaling between problem size and model performance.

Effective application to problems like sorting, Sudoku, and matrix factorization.

Demonstrated benefits of explicit subcomputations and permutation invariances.

Abstract

We introduce a framework for automatically defining and learning deep generative models with problem-specific structure. We tackle problem domains that are more traditionally solved by algorithms such as sorting, constraint satisfaction for Sudoku, and matrix factorization. Concretely, we train diffusion models with an architecture tailored to the problem specification. This problem specification should contain a graphical model describing relationships between variables, and often benefits from explicit representation of subcomputations. Permutation invariances can also be exploited. Across a diverse set of experiments we improve the scaling relationship between problem dimension and our model's performance, in terms of both training time and final accuracy. Our code can be found at https://github.com/plai-group/gsdm.