# compositional_generative_inverse_design__45e6b655.pdf Published as a conference paper at ICLR 2024 COMPOSITIONAL GENERATIVE INVERSE DESIGN Tailin Wu1 , Takashi Maruyama2 , Long Wei1 , Tao Zhang1 , Yilun Du3 , Gianluca Iaccarino4, Jure Leskovec5 1Dept. of Engineering, Westlake University, 2NEC Laboratories Europe, 3Dept. of Computer Science, MIT, 4Dept. of Mechanical Engineering, Stanford University, 5Dept. of Computer Science, Stanford University wutailin@westlake.edu.cn, Takashi.Maruyama@neclab.eu, weilong@westlake.edu.cn,zhangtao@westlake.edu.cn, yilundu@mit.edu jops@stanford.edu, jure@cs.stanford.edu Inverse design, where we seek to design input variables in order to optimize an underlying objective function, is an important problem that arises across fields such as mechanical engineering to aerospace engineering. Inverse design is typically formulated as an optimization problem, with recent works leveraging optimization across learned dynamics models. However, as models are optimized they tend to fall into adversarial modes, preventing effective sampling. We illustrate that by instead optimizing over the learned energy function captured by the diffusion model, we can avoid such adversarial examples and significantly improve design performance. We further illustrate how such a design system is compositional, enabling us to combine multiple different diffusion models representing subcomponents of our desired system to design systems with every specified component. In an N-body interaction task and a challenging 2D multiairfoil design task, we demonstrate that by composing the learned diffusion model at test time, our method allows us to design initial states and boundary shapes that are more complex than those in the training data. Our method generalizes to more objects for N-body dataset and discovers formation flying to minimize drag in the multi-airfoil design task. Project website and code can be found at https://github.com/AI4Science-Westlake U/cindm. 1 INTRODUCTION The problem of inverse design finding a set of high-dimensional design parameters (e.g., boundary and initial conditions) for a system to optimize a set of specified objectives and constraints, occurs across many engineering domains such as mechanical, materials, and aerospace engineering, with important applications such as jet engine design (Athanasopoulos et al., 2009), nanophotonic design (Molesky et al., 2018), shape design for underwater robots (Saghafi & Lavimi, 2020), and battery design (Bhowmik et al., 2019). Such inverse design problems are extremely challenging since they typically involve simulating the full trajectory of complicated physical dynamics as an inner loop, have high-dimensional design space, and require out-of-distribution test-time generalization. Recent deep learning has made promising progress for inverse design. A notable work is by Allen et al. (2022), which addresses inverse design by first learning a neural surrogate model to approximate the forward physical dynamics, and then performing backpropagation through the full simulation trajectory to optimize the design parameters such as the boundary shape. Compared with standard sampling-based optimization methods with classical simulators, it shows comparable and sometimes better performance, establishing deep learning as a viable technique for inverse design. However, an underlying issue with backpropagation with surrogate models is over-optimization as learned models have adversarial minima, excessive optimization with respect to a learned forward model leads to adversarial design parameters which lead to poor performance (Zhao et al., 2022). A root cause of this is that the forward model does not have a measure of data likelihood and does Equal contribution. Corresponding author. Published as a conference paper at ICLR 2024 More time steps More interacting bodies Learn !!(#) from observed data !! ... !" !"#$ ... !% Design objective Energy function & = (($, *) +[$,'] Trajectory Part to whole for multiple boundaries Design variable Shared trajectory Train Compositional Design Figure 1: Cin DM schematic. By composing generative models specified over subsets of inputs, we present an approach which design materials significantly more complex than those seen at training. not know which design parameters are in or out of the training distribution it has seen, allowing optimization to easily fall out-of-distribution of the design parameters seen during training. To address this issue, we view the inverse design problem from an energy optimization perspective, where constraints of the simulation model are implicitly captured through the generative energy function of a diffusion model trained with design parameters and simulator outputs. Designing parameters subject to constraints corresponds to optimizing for design parameters that minimize the energy of both the generative energy function and associated design objective functions. The generative energy function prevents design parameters from deviating and falling out of distribution. An essential aspect of inverse design is the ability to further construct new structures subjects to different constraints at test-time. By formulating inverse design as optimizing generative energy function trained on existing designs, a na ıve issue is that it constrains design parameters to be roughly those seen in the training data. We circumvent this issue by using a set of generative energy functions, where each generative model captures a subset of design parameters governing the system. Each individual generative energy function ensures that designs do not locally fall out of distribution, with their composition ensuring that inferred design parameters are roughly locally in distribution. Simultaneously, designs from this compositional set of generative energy functions may be significantly different from the training data, as designs are not constrained to globally follow the observed data (Liu et al., 2022; Du et al., 2023), achieving compositional generalization in design. We illustrate the promise of using such compositional energy functions across a variety of different settings. We illustrate that temporally composing multiple compositional energy functions, we may design sequences of outputs that are significantly longer than the ones seen in training. Similarly, we can design systems with many more objects and more complex shapes than those seen in training. Concretely, we contribute the following: (1) We propose a novel formulation for inverse design as an energy optimization problem. (2) We introduce Compositional Inverse Design with Diffusion Models (Cin DM) method, which enables us to generalize to out-of-distribution and more complex design inputs than seen in training. (3) We present a set of benchmarks for inverse design in 1D and 2D. Our method generalizes to more objects for N-body dataset and discovers formation flying to minimize drag in the multi-airfoil design task. 2 RELATED WORK Inverse Design. Inverse design plays a key role across science and engineering, including mechanical engineering (Coros et al., 2013), materials science (Dijkstra & Luijten, 2021), nanophotonics (Molesky et al., 2018), robotics (Saghafi & Lavimi, 2020), chemical engineering (Bhowmik et al., 2019), and aerospace engineering (Athanasopoulos et al., 2009; Anderson & Venkatakrishnan, 1999). Classical methods to address inverse design rely on slow classical solvers. They are accurate but are prohibitively inefficient (e.g., sampling-based methods like CEM (Rubinstein & Kroese, 2004)). Recently, deep learning-based inverse design has made promising progress. Allen Published as a conference paper at ICLR 2024 et al. (2022) introduced backpropagation through the full trajectory with surrogate models. Wu et al. (2022a) introduced backpropagation through latent dynamics to improve efficiency and accuracy. For Stokes systems, Du et al. (2020a) introduced an inverse design method under different types of boundary conditions. While the above methods typically rely on learning a surrogate model for the dynamics and use it as an inner loop during inverse design, we introduce a novel generative perspective that learns an energy function for the joint variable of trajectory and boundary. This brings the important benefit of out-of-distribution generalization and compositionality. Ren et al. (2020); Trabucco et al. (2021); Ansari et al. (2022); Chen et al. (2023) benchmarked varieties of deep learning-based methods in a wide range of inverse design tasks. Compositional Models. A large body of recent work has explored how multiple different instances of generative models can be compositionally combined for applications such as 2D images synthesis (Du et al., 2020b; Liu et al., 2021; Nie et al., 2021; Liu et al., 2022; Wu et al., 2022b; Du et al., 2023; Wang et al., 2023), 3D synthesis (Po & Wetzstein, 2023), video synthesis (Yang et al., 2023a), trajectory planning (Du et al., 2019; Urain et al., 2021; Gkanatsios et al., 2023; Yang et al., 2023b), multimodal perception (Li et al., 2022) and hierarchical decision making (Ajay et al., 2023). Technically, product of experts is an effective kind of approaches to combine the predictive distributions of local experts (Hinton, 2002; Cohen et al., 2020; Gordon et al., 2023; Tautvaiˇsas & ˇZilinskas, 2023) . To the best of our knowledge, we are the first to introduce a compositional generative perspective and method to inverse design, and show how compositional models can enable us to generalize to design spaces that are much more complex than seen at training time. In this section, we detail our method of Compositional INverse design with Diffusion Models (Cin DM). We first introduce the problem setup in Section 3.1. In Section 3.2, we introduce generative inverse design, a novel generative paradigm for solving the inverse design problem. In Section 3.3, we detail how our method allows for test-time composition of the design variables. 3.1 PROBLEM SETUP We formalize the inverse design problem using a similar setup as in Zhang et al. (2023). Concretely, let u(x, t; γ) be the state of a dynamical system at time t and location x where the dynamics is described by a partial differential equation (PDE) or an ordinary differential equation (ODE).1 Here γ = (u0, B) Γ consists of the initial state u0 and boundary condition B, Γ is the design space, and we will call γ boundary for simplicity2. Given a PDE or ODE, a specific γ can uniquely determine a specific trajectory u[0,T ](γ) := {u(x, t; γ)|t [0, T]}, where we have written the dependence of u[0,T ] on γ explicitely. Let J be the design objective which evaluates the quality of the design. Typically J is a function of a subset of the trajectory u[0,T ] and γ (esp. the boundary shape). The inverse design problem is to find an optimized design ˆγ which minimizes the design objective J : ˆγ = arg min γ J (u[0,T ](γ), γ) (1) We see that J depends on γ through two routes. On the one hand, γ influences the future trajectory of the dynamical system, which J evaluates on. On the other hand, γ can directly influence J at future times, since the design objective may be directly dependent on the boundary shape. Typically, we don t have access to the ground-truth model for the dynamical system, but instead only observe the trajectories u[0,T ](γ) at discrete time steps and locations and a limited diversity of boundaries γ Γ. We denote the above discrete version of the trajectory as U[0,T ](γ) = (U0, U1, ..., UT ) across time steps t = 0, 1, ...T. Given the observed trajectories U[0,T ](γ), γ Γ, a straightforward method for inverse design is to use such observed trajectories to train a neural surrogate model fθ for forward modeling, so the trajectory can be autoregressively simulated by fθ: ˆUt(γ) = fθ( ˆUt 1(γ), γ), ˆU0 := U0, γ = (U0, B), (2) Here we use ˆUt to represent the prediction by fθ, to differentiate from the actual observed state Ut. In the test time, the goal is to optimize J ( ˆU[0,T ](γ), γ) w.r.t. γ, which includes the autoregressive rollout with fθ as an inner loop, as introduced in Allen et al. (2022). In general inverse design, the 1In the case of ODE, the position x is neglected and the trajectory is u(t; γ), where γ only includes the initial state u0. For more background information about PDEs, see Brandstetter et al. (2022). 2Since B is the boundary in space and the initial state u0 can be seen as the boundary in time. Published as a conference paper at ICLR 2024 trajectory length T, state dimension dim(U[0,T ](γ)), and complexity of γ may be much larger than in training, requiring significant out-of-distribution generalization. 3.2 GENERATIVE INVERSE DESIGN Directly optimizing Eq. 1 with respect to γ using a learned surrogate model fθ is often problematic as the optimization procedure on γ often leads a set of U[0,T ] that is out-of-distribution or adversarial to the surrogate model fθ, leading to poor performance, as observed in Zhao et al. (2022). A major cause of this is that fθ does not have an inherent measure of uncertainty, and cannot prevent optimization from entering a design spaces γ that the model cannot guarantee its performance in. To circumvent this issue, we propose a generative perspective to inverse design: during the inverse design process, we jointly optimize for both the design objective J and a generative objective Eθ, ˆγ = arg min γ,U[0,T ] Eθ(U[0,T ], γ) + λ J (U[0,T ], γ) , (3) where Eθ is an energy-based model (EBM) p(U[0,T ], γ) e Eθ(U[0,T ],γ) (Le Cun et al., 2006; Du & Mordatch, 2019) trained over the joint distribution of trajectories U[0,T ] and boundaries γ, and λ is a hyperparameter. Both U[0,T ] and γ are jointly optimized, and the energy function Eθ is minimized when both U[0,T ] and γ are consistent with each other and serves the purpose of a surrogate model fθ in approximating simulator dynamics. The joint optimization optimizes all the steps of the trajectory U[0,T ] and the boundary γ simultaneously, which also gets rid of the time-consuming autoregressive rollout as an inner loop as in Allen et al. (2022), significantly improving inference efficiency. In addition to approximating simulator dynamics, the generative objective also serves as a measure of uncertainty. Essentially, the Eθ in Eq. 3 encourages the trajectory U[0,T ] and boundary γ to be physically consistent, and the J encourages them to optimize the design objective. To train Eθ, we use a diffusion objective, where we learn a denoising network ϵθ that learns to denoise all variables in design optimization z = U[0,T ] L γ supervised with the training loss LMSE = ϵ ϵθ( p βsϵ, s) 2 2, ϵ N(0, I). (4) As discussed in Liu et al. (2022), the denoising network ϵθ corresponds to the gradient of a EBM z Eθ(z), that represents the distribution over all optimization variables p(z) e Eθ(z). To optimize Eq. 3 using a Langevin sampling procedure, we can initialize an optimization variable z S from Gaussian noise N(0, I), and iteratively run zs 1 = zs η ( z(Eθ(zs) + λ J (zs))) + ξ, ξ N 0, σ2 s I , (5) for s = S, S 1, ..., 1. This procedure is implemented with diffusion models by optimizing3 zs 1 = zs η (ϵθ(zs, s) + λ z J (zs)) + ξ, ξ N 0, σ2 s I , (6) where σ2 s and η correspond to a set of different noise schedules and scaling factors used in the diffusion process. To further improve the performance, we run additional steps of Langevin dynamics optimization at a given noise level following Du et al. (2023). Intuitively, the above diffusion procedure starts from a random variable z S = (U[0,T ],S L γS) N(0, I), follows the denoising network ϵθ(zs, s) and the gradient z J (zs), and step-by-step arrives at a final z0 = U[0,T ],0 L γ0 that approximately minimizes the objective in Eq. 3. 3.3 COMPOSITIONAL GENERATIVE INVERSE DESIGN A key challenge in inverse design is that the boundary γ or the trajectory U[0,T ] can be substantially different than seen during training. To enable generalization across such design variables, we propose to compositionally represent the design variable z = U[0,T ] L γ, using a composition of different energy functions Eθ (Du et al., 2020b) on subsets of the design variable zi z. Each of the above Eθ on the subset of design variable zi provides a physical consistency constraint on zi, encouraging each zi to be physically consistent internally. Also we make sure that different zi, i = 1, 2, ...N overlap with each other, and overall covers z (See Fig. 1), so that the full z is physically consistent. Thus, test-time compositions of energy functions defined over subsets of the 3There is also an additional scaling term applied on the sample zs during the diffusion sampling procedure, which we omit below for clarity but also implement in practice. Published as a conference paper at ICLR 2024 Algorithm 1 Algorithm for Compositional Inverse Design with Diffusion Models (Cin DM) 1: Require Compositional set of diffusion models ϵi θ(zs, s), i = 1, 2, ...N, design objective J ( ), covariance matrix σ2 s I, hyperparameters λ, S, K 2: Initialize optimization variables z S N(0, I) // optimize across diffusion steps S: 3: for s = S, . . . , 1 do 4: // optimize K steps of Langevin sampling at diffusion step s: 5: for k = 1, . . . , K do 6: ξ N 0, σ2 s I 7: // run a single Langevin sampling steps: 8: zs zs η 1 N PN i=1 ϵi θ(zi s, s) + λ z J (zs) + ξ 9: end for 10: ξ N 0, σ2 s I 11: // scale sample to transition to next diffusion step: 12: zs 1 zs η 1 N PN i=1 ϵi θ(zi s, s) + λ z J (zs) + ξ 13: end for 14: γ, U[0,T ] = z0 15: return γ design variable zi z can then be composed together to generalize to new design variable z values that substantially different than those seen during training, but exploiting shared local structure in z. Below, we illustrate three different ways compositional inverse design can enable to generalize to design variables z that are much more complex than the ones seen during training. I. Generalization to more time steps. In the test time, the trajectory length T may be much longer than the trajectory length T tr seen in training. To allow generalization over a longer trajectory length, the energy function over the design variable can be written in terms of a composition of N energy functions over subsets of trajectories with overlapping states: Eθ(U[0,T ], γ) = i=1 Eθ(U[(i 1) tq,i tq+T tr], γ). (7) Here zi := U[(i 1) tq,i tq+T tr] L γ is a subset of the design variable z := U[0,T ] L γ. tq {1, 2, ...T 1} is the stride for consecutive time intervals, and we let T = N tq + T tr. II. Generalization to more interacting bodies. Many inverse design applications require generalizing the trained model to more interacting bodies for a dynamical system, which is far more difficult than generalizing to more time steps. Our method allows such generalization by composing the energy function of few-body interactions to more interacting bodies. Now we illustrate it with a 2-body to N-body generalization. Suppose that only the trajectory of a 2-body interaction is given, where we have the trajectory of U (i) [0,T ] = (U (i) 0 , U (i) 1 , ..., U (i) T ) for body i {1, 2}. We can learn an energy function Eθ((U (1) [0,T ], U (2) [0,T ]), γ) from this trajectory. In the test time, given N > 2 interacting bodies subjecting to the same pairwise interactions, the energy function for the combined trajectory U[0,T ] = (U (1) [0,T ], ..., U (N) [0,T ]) for the N bodies is then given by: Eθ(U[0,T ], γ) = X i