Scaling Graphically Structured Diffusion Models

Christian Dietrich Weilbach, William Harvey, Hamed Shirzad, Frank Wood

Keywords: diffusion models, graphical models, graph neural networks, bayesian deep learning, sparse attention, scaling

TL;DR: We scale Graphically Structured Diffusion Models and provide a new benchmark for scaling.

Abstract: Applications of the recently introduced graphically structured diffusion model (GSDM) family show that sparsifying the transformer attention mechanism within a diffusion model and meta-training on a variety of conditioning tasks can yield an efficiently learnable diffusion model artifact that is capable of flexible, in the sense of observing different subsets of variables at test-time, amortized conditioning in probabilistic graphical models. While extremely promising in terms of applicability and utility, implementations of GSDMs prior to this work were not scalable beyond toy graphical model sizes. We overcome this limitation by describing and and solving two scaling issues related to GSDMs; one engineering and one methodological. We additionally propose a new benchmark problem of weight inference for a convolutional neural network applied to $14\times14$ MNIST.

Submission Number: 96