Compositional Generative Inverse Design
Authors: Tailin Wu
Takashi Maruyama
Long Wei
Tao Zhang
Yilun Du
Gianluca Iaccarino
Jure Leskovec
Venue: ICLR 2024 spotlight (2024)
Keywords: inverse design
generative design
PDE
physical simulation
compositional
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We introduce a method that uses compositional generative models to design boundaries and initial states significantly more complex than the ones seen in training for physical simulations

Abstract: 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 multi-airfoil 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-WestlakeU/cindm.

Published as a conference paper at ICLR 2024
COMPOSITIONAL GENERATIVE INVERSE DESIGN
Tailin Wu Takashi Maruyama 2* Long Weilr, Tao Zhang", Yilun Du Gianluea laccarino', Jure Leskovec Dept. of Engineering, Westlake University, NEEC Laboratories Europe,
8 Deett of Computer Science, MIT, *Dept. of Mechanical Engineering, Stanford University, Deett. of Computer Science, Stanford University
dettalli@weesfe lake. edu.cn, cn, Takashi. MookahhiMaruyyaee M @neclab.en wei long@west akee lake.
edu. cn, zhangta lake. backe.edu.ccm edu. cn, y undddêmia edu jops@stanford.eem edu, jureêcs. stanford.edu
ABSTRACT
Inverse design, where seek design input variables in order optimize underlying objective function, is an important problem that arises across fields such as mechanical engineering aerospace engineering. Inverse design typ- ically formulated as an optimization problem, with recent works leveraging op- timization across learned dynamics models. However, as models are optimized they tend to fall into adversarial modes, preventing effective sampling. We il- lustrate 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 design system compositional, enabling us to combine multiple different diffusion models rep- resenting subcomponents of our desired system to design systems with every specified component. In N-hody interaction task and challenging 2D multi- airfoil design task, we demonstrate that by composing the learned diffusion model 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 more objects for N-body dataset and discovers formation flying to minimize drag in the multi-airfol design task. Project website and code can be found at https: : / / thub. com AAA.AA44.. AT4Sciince West aeke// c ndm.
INTRODUCTION
The problem of inverse design finding a set of high-dimensional design parameters (e.g., boundary and initial conditions) for 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 al., 2019). Such inverse design problems 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 al. (2022), which addresses inverse design by first learning neural surrogate model to approx- imate the forward physical dynamics, and then performing backpropagation through the full sim- ulation trajectory to optimize the design parameters such the boundary shape. Compared with standard sampling-bar optimization methods with classical simulators, it shows comparable and sometimes better performance, establishing deep learning as 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). root cause of this is that the forward model does not have a measure of data likelihood and does
"qqual contribution. 'Corresponding author.

---

Published as a conference paper at ICLR 2024
Train
Compositional Design
More time steps
Learn (x) from observed data
state
(gRo(t))
U
Design objective
More interacting bodies
Z--1
AJ (2)
state
(A7
EØ(2)
U
Energy function
Design variable
whole multiple boundaries;
state
state
(Mgg 28)
UPTTI Trajectory
Uuay
Boundary
Shared trajectory
Figure I: CinDM schematic. By composing generative models specified over subsets of inputs,
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 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 dif- ferent constraints test-time. By formulating inverse design as optimizing generative energy fune- tion trained on existing designs, 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 generative energy func- tions, where each generative model captures subset of design parameters governing the system. Each individual generative energy function ensures that designs do not locally fall out of distribu- tion, with their composition ensuring that inferred design parameters bution. Simu Simultaneously, designs from this compositional set of generative energy functions may be significantly different from the training data, as designs not constrained globally follow the observed data (Liu et al., 2022; Du et al., 2023), achieving compositional generalization in design.
roughly "locally" in distri-
We illustrate the promise of using such compositional energy functions across 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 training. Similarly, we can design systems with many more objects and more complex shapes than those seen in training.
Coneretely, we contribute the following: (1) We propose novel formulation for inverse design as an energy optimization problem. (2) We introduce Compositional Inverse Design with Diffusion Models (CinDM) method, which enables us to generalize
diittii-iitiittiionn and more complex design inputs than seen in training. (3) We present set of benchmarks for inverse design in ID and 2D. Our method generalizes 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 me- chanical engineering (Coros al., 2013), materials science (Dijkstra Luijten, 2021), nanopho- tonics (Molesky et al., 2018), robotics (Saghafi & Lavimi, 2020), chemical engineering (Bhowmik et al., 2019), and aerospace engineering (Athanasopoulos et al., 2009; Anderson & Venkatakrish- nan, 1999). Classical methods to address inverse design rely slow classical solvers. They accurate but are prohibitively inefficient (e.g., sampling-ber based methods like CEM (Rubinstein & Kroese, 2004)). Recently, deep learning-based inverse design has made promising progress. Allen
2

---

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 inner loop during inverse design, we introduce novel generative
the dynamics and use it 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 synthe- sis (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). Techni- cally, product of experts is effective kind of approaches to combine the predictive distributions of local experts (Hinton, 2002; Cohen et al., 2020; Gordon et al., 2023; Tautvaišas & 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.
3 METHOD
In this section, we detail our method of Compositional INverse design with Diffusion Models (CinDM). We first introduce the problem setup in Section 3.1. In Section 3.2, we introduce genera- tive 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; 7; be the state of a dynamical system at time and location where the dynamics is described by a partial differential equation (PDE) or an ordinary differential equation (ODE). Here y= y=(ug, B) E I consists of the initial state uy and boundary condition B, I is the design space, and we will call "boundary" for simplicity?. Given a PDE or ÓDE, a specific ry can uniquely determine specific trajectory "(o.7] (7) == (u(a, 1/ E [0, T]}, where we have written the dependence of "jo.T on ? explicitely. Let J be the design objective which evaluates the quality of the design. Typically J is a function of subset of the trajectory u[o.r] and 7 (esp. the boundary shape). The inverse design problem is to find an optimized design Ý which minimizes the design objective J:
4= arg min I (1)).1((7))7)
(1)
We see that J depends on y through two routes. On the one hand, 7 influences the future trajectory of the dynamical system, which I evaluates on. On the other hand, can directly influence I 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 (7) at discrete time steps and locations and a limited diver- sity of boundaries 7 П. We denote the above discrete version of the trajectory as (1) = (Uo. UI.... Ur) across time steps t = 0, 1..TT. Given the observed trajectories (1)) r, a straightforward method for inverse design is to use such observed trajectories to train neural surrogate model fo for forward modeling, so the trajectory can be autoregressively simulated by fo:
Ü,(7) fe(Ú,_1(7),))
Üo := Uo ? (Uo, B),
(2)
Here we use U, to represent the prediction by fo, to differentiate from the actual observed state U.- In the test time, the goal is optimize J (Ujo.7((7) 7) w.r.t. 7, which includes the autoregressive rollout with fo as an inner loop, as introduced in Allen et al. (2022). In general inverse design, the
In the case of ODE, the position x is neglected and the trajectory u(t: where y only includes the initial state uo. For more background information about PDEs, see Brandstetter et al. (2022). Sinne B the houndary in space and the initial state to can he seen as the "boundary" in time.
3

---

Published as a conference paper at ICLR 2024
trajectory length T, state dimension (i and complexity of yy may be much larger than in training, requiring significant out- -iitt-i-iittitttonnnn generalization.
3.2 GENERATIVE INVERSE DESIGN
Directly optimizing Eq. with respect to y using a learned surrogate model fo is often problematic as the optimization procedure on Y often leads sarial to the surrogate model fo. leading to poor performance, as observed in Zhao et al. (2022). A major cause of this is that fe 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.
of Up..1] that is out-of-distribution or adver-
To cireumvent this issue, we propose a generative perspective inverse design: during the inverse design process, we jointly optimize for both the design objective J and a generative objective Eg.
4 arg min [Eg(U 7) + + [U
"/U,..21
where Eg is an energy-based model (EBM) p(U 7) x (U(( (LeCun et al., 2006; Du & Mordatch, 2019) trained over the joint distribution of trajectories Ujo.r and boundaries 7, and is a hyperparameter. Both U0.27 and are jointly optimized, and the energy function Eg is minimized when both U and y are consistent with each other and serves the purpose of surrogate model fo in approximating simulator dynamics. The joint optimization optimizes all the steps of the trajectory 1j and the boundary 7 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 Eg in Eq. 3 encourages the trajectory U/ory and boundary y to be physically consistent, and the J encourages them to optimize the design objective.
To train Eg, use a diffusion objective, where learn denoising network cæ that learns to denoise all variables in design optimization 2 = Up.7 θ y supervised with the training loss
MSE = ||e cal co(vI- 1,2 / s)e e~ I).
(4)
As discussed in Liu et al. (2022), the denoising network eg corresponds to the gradient of EBM V, Ex(,, that represents the distribution over all optimization variables p(z) xx eoEo()) To opti- mize Eq. using a Langevin sampling procedure, we can initialize an optimization variable zş from Gaussian noise (O, 2). and iteratively run
---1 1 = 2g - 7 (Eo(2.) + (z)))) +5, {~~(0,,
(5)
for s = 1, ...,, This procedure is implemented with diffusion models by optimizing' 3
- = 2g 77 (ce(2,, s) + (zx)) + E. E {~N(0,,
(6)
where of and 1) correspond to a set of different noise schedules and scaling factors used in the dif- fusion process. To further improve the performance, we run additional steps of Langevin dynamics optimization given noise level following Du al. (2023).
Intuitively, the above diffusion procedure starts from a random variable 2s = (U/.. + 7s) ~ (O I). follows the demoising network e (z.,, and the gradient V=J(Zz), and step-by-step arrives final zp 1 + 7o that approximately minimizes the objective in Eq. 3.
3.3
COMPOSITIONAL GENERATIVE INVERSE DESIGN
A key challenge in inverse design that the boundary 7 or the trajectory Up..1 can be substan- tially different than seen during training. To enable generalization across such design variables, we propose to compositionally represent the design variable different energy functions Eg (Du et al., 2020b) on subsets of the design variable 2; C z. Each of the above Eg on the subset of design variable z; provides a physical consistency constraint on 2, encouraging each to be physically consistent internally. Also we make sure that different 2,, i = 1, 2, ...N overlap with each other, and overall covers z (See Fig. 1), so that the full 2 is physically consistent. Thus, test-time compositions of energy functions defined over subsets of the
U θ% using a composition of
There is also additional scaling term applied on the sample z, during the diffusion sampling procedure, which we omit below for clarity but also implement in practice.
4

---

Published as a conference paper at ICLR 2024
Algorithm Algorithm for Compositional Inverse Design with Diftusion Models (CinDM)
1: Require Compositional set of diffusion models eg (Zs, (2, 1, ..,,,, design objective I(). covariance matrix o²T, T, hyperparameters, X, S. K
2: Initialize optimization variables zg ~ (O,
optimize across diffusion steps S:
3: for 8 = S. S.....1 do
4: I/ optimizee K steps of Langevin sampling diffusion step
5: for l: = 1, ........ K do
6: { N~ o
7:
r/u a single Langevin sampling steps:
8:
2, 5,4-2. 2, 7 E (c0(21, + XV .T.))) +
9:
end for
10:
E~ S~N(0, o.
11:
scalle sample to transition to next diffusion step:
12: 2s-1 + 25 n E (eb(e), + xV=J(=z)) +
13: end for
14:
%, "UUU..
15: return 7
design variable 2; C 2 can then be composed together to generalize to new design variable z values that substantially different than those during training, but exploiting shared local structure in
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.
Generalization to more time steps. In the test time, the trajectory length T may be much longer than the trajectory length 7 seen in training. To allow generalization over longer trajectory length, the energy function over the design variable can be written in terms of a composition of energy functions over subsets of trajectories with overlapping states:
N
Eo(U = Σ EO(Ui(CI)),, =
Here 2 := U((:- itt -T + 7 subset of the design variable z := U/0.T + 7. to 21, 1) is the stride for consecutive time intervals, and we let = N .tg T
II. Generalization to more interacting bodies. Many inverse design applications require gener- alizing the trained model to more interacting bodies for a dynamical system, which far more difficult than generalizing 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 2-body to N-body generalization. Suppose that only the trajectory of 2-body interaction is given, where we have the trajectory of ((0.7) (U() u(²). ...,, U( for body E [1,2]; We can learn an
energy function Fo (T((L [o. (2(2) 7) from this trajectory. In the test time, given > 2 interacting bodies subjecting to the same pairwise interactions, the energy function for the combined trajectory Upu,7 (U:L) 10.27 (1.2)) for the bodies is then given by:
Es (U(o,,,7) = Σ Eg (0.7). 0-61
(8)
III. Generalization from part to whole for boundaries. Real-life inverse design typically involves designing shapes consisting of multiple parts that constitute an integral whole. Examples include planes that consist of wings, the body, the rudder, and many other parts. The shape of the whole may be more complex and out-of-distribution than the parts seen in training. To generalize from parts to whole, we can again compose the energy function over subsets of the design variable z. Concretely, U(6) suppose that we have trajectories [o.T] corresponding to the part y", i 1,.,...N we can learn energy functions corresponding the dynamics of each part Fa, (U(() (0.7) 7*), 1,2, An
example is that 7 represents the shape for each part of the plane, and 0.21 represents the full fluid state around the part 7/ without other parts present. In the test time, when requiring to generalize
5

---

Published as a conference paper at ICLR 2024
over whole boundary y that consisting of these N parts 24, 2 1, 1,2...N, we have
N
Eo(Uio.7)7)) = E
Eo, (U/o.r):"))
(9)
Note that here the composition, all the parts 7 share the same trajectory Up.T7 which can be intuitively understood in the example of the plane where all the parts of the plane influence the same full state of fluid around the plane. The composition of energy functions in Eq. 9 means that the full energy Exp (U 7) will be low if the trajectory UJO.T] is consistent with all the parts y"
Compositional Generative Inverse Design. Given the above composition of energy functions, we can correspondingly learn each energy function over the design variable Uiu.11 1 ① by training corresponding diffusion model over the subset of design variables 2' Our overall sampling objective given the set of energy functions (Fa(="))== is then given by
28-1 1
2g " N E (c)()).)) +V +s, E~~(0,0
(10)
s S,S 1, 1,.... Similarly to before, we can further run multiple steps of Langevin dynamics optimization at given noise level following Du al. (2023) further improve performance. We provide the overall pseudo-code of our method in the compositional setting in Algorithm 1.
Our proposed paradigm of generative inverse design in Section 3.2 (consisting of its design objective Eq. 3 and training objective Eq. 4) and our compositional inverse design method in Section 3.3, constitute our full method of Compositional INverse design with Diffusion Models (CinDM). Our approach is different from product of experts (Hinton, 2002) in that CinDM learns distribution of a subspace in training, based on which we infer distribution in much higher spaces during inference. Below, we will test our method's capability in compositional inverse design.
EXPERIMENTS
In the experiments, we aim to answer the following questions: (1) Can CinDM generalize to more complex designs in the test time using its composition capability? (2) Comparing backpropagation with surrogate models and other strong baselines, can CinDM improve on the design objective or prediction accuracy? (3) Can CinDM address high-dimensional design space? To answer these questions, we perform our experiments in three different scenarios: compositional inverse design in time dimension (Sec. 4.1), compositional inverse design generalizing to more objects (Sec. 4.2), and 2D compositional design for multiple airfoils with Navier-Stokes flow (Sec. 4.3). Each of the above experiments represents an important scenario in inverse design and has important implications in science and engineering. In each experiment, we compare CinDM with the state-of-the--e art deep learning-base baaed inverse design method proposed by Allen et al. (2022), which we term Backprop, and cross-entropy method (CEM) (Rubinstein & Kroese, 2004) which is a standard sampling-bare optimization method typically used in classical inverse design. Additionally, we compare CinDM with two inverse-design methods: neural adjoint method with the boundary loss function (NABL) and conditional invertible neural network (cINN) method (Ren et al., 2020; Ansari et al., 2022). The details and results are provided in Appendix H, in which adopt a more reasonable experimental setting to those new baselines. All the baselines and our model contain similar numbers of param- eters each comparison for fair evaluation. To evaluate the performance of each inverse design method, we feed the output of the inverse design method (i.e., the optimized initial or boundary states) to the ground-truth solver, perform rollout by the solver and feed the rollout trajectory to the design objective. We do not use the surrogate model to perform rollout since the trained surrogate models may not be faithful to the ground-truth dynamics and can overestimate the design objective. By evaluating using ground-truth solver, all inverse design methods can be evaluated fairly.
4.1 COMPOSITIONAL INVERSE DESIGN IN TIME
In this experiment, we aim to test each method's ' ability to generalize to more forward time steps than during training. This is important since in test time. the inverse design methods are typically used over longer predictions horizons than in training. We use an N-body interaction environment where each ball with radius of 0.1 is bouncing in box. The balls will exchange momentum when

---

Published as a conference paper at ICLR 2024
Table 1: Compositional Generalization Across Time. Experiment on compositional inverse de- sign in time. The confidence interval information is deligated to Table 6 in Appendix B for page constraints. Bold font denotes the best model.
2-body 24 steps 2-body 34 steps MAE
2-body 44 steps MAE
2-body 54 steps
Method
design obj
MAE
design obj
design obj
design obj
MAE
CEM. GNS (1-step)
0,2622
0.13963
0.2204
0.15378
0.2701
0.21277
0.2773
0,21706
CEM. GNS
0,2699
0,12746
0.3142
0.14637
0.3056
0,18155
0.3124
0,20266
CEM. U-Net (1-step)
0.2364
0.07720
0.2391
0.09701
0.2744
0.11885
0.2729
0.12992
CEM. U-Net
0.1762
0.03597
0,1639
0.03094
0.1816
0,03900
0.1887
0.04350
Backprop, GNS (1-step)
0,1452
0.04339
0.1497
0,03806
0.1511
0.03621
0.1851
0,04104
Backprop. GNS
0,2407
0.09788
0.2678
0.11017
0.2762
0.12395
0.2952
0.13963
Backprop. U-Net (1-step)
0.2182
0.07554
0.2445
0.08278
0.2536
0.08487
0.2751
0,10599
Backprop, U-Net
0.1228
0.01974
0.1171
0.01236
0.1143
0.00970
0.1289
0.01067
CinDM (ours)
0,1160
0.01264
0.1288
0.00917
0.1447
0.00959
0.1650
0.01064
elastically colliding with each other or with the wall. The design task is to identify the initial state (position and velocity of the balls) of the system such that the end state optimizes a certain objective (e.g., as close to a certain target as possible). This setting represents a simplified version of many real-life scenarios such as billiard, bowling, and ice hockey. Since the collisions preserve kinetic energy but modify speed and direction of each ball and multiple collisions can happen over a long time, this represents non-trivial inverse design problem with abrupt changes in the design space. During training time, we provide each method with training trajectory consisting of 24 steps, and in test time, let it roll out for a total of 24, 34, 44, and 54 steps. The design objective
minimize the last step's Euclidean distance to the center (x, y) (0.5, 0.5). For baselines, we compare with CEM (Rubinstein & Kroese, 2004) and Backprop (Allen et al., 2022). Each method uses either Graph Network Simulator (GNS, Sanchez-Gonzalez et al. (2020), state-of-the-ar art method for modeling N- body interactions) or U-Net (Ronneberger et al., 2015) as backbone architecture that either predicts I step or 23 steps in a single forward pass. For our method, we use the same U-Net backbone architecture for diffusion. To perform time composition, we superimpose N EBMs Eg (U ?) on states with overlapping time ranges: U,223), Up10,33. U120,433.... UpoLN 10,10(N-11)) 1)+333 in Eq. 7, and use Eq. 10 to perform denoising diffusion. Besides evaluating with the design objective (J), we also use the metric of mean absolute error (MAE) between the predicted trajectory and the trajectory generated by the ground-truth solver to evaluate how faithful each method's prediction is. Each design scenario run 500 times and the average performance is reported in Table We show example trajectories of our method in Fig. 2. Details for the architecture and training are provided in Appendix A. We also make comparison with a simple baseline that performs diffusion over 44 steps directly without time composition. Details and results are presented in Appendix C.
From Table 1, we see that our method is competitive in design objectives and outperforms every baseline in MAE. In the "2-body 24 steps" scenario which is the same setting as in training and without composition, our method outperforms the strongest baselines by a wide margin both on design objective and MAE, With more prediction steps, method not only performs better than any baselines in MAE but also merely is weaker than the strongest baseline in design objective. For example, our method's 's MAE outperforms the best baseline by 36.0%. 25.8%, 1.1%, and 0.3% for 24, 34, 44, and 54-step predictions, respectively, with an average of 15.8% improvement. Similarly, our method's 's design objective outperforms the best baseline by 5.5% for 24-step. This shows the two-fold advantage of our method. Firstly, even with the same backbone architecture, our diffusion method can roll out stably and accurately for much longer than the baseline, since the forward surrogate models in the basclines during design may encounter out-of-distributio and adversarial inputs which does not know how evolve properly. On the other hand, our diffusion-based method is trained to denoise and favor inputs consistent with the underlying physics. Secondly, our compositional method allows our model to generalize to longer time steps and allows for stable rollout. An example trajectory designed by our CinDM is shown in Fig. 2 (a). We see that it matches with the ground-truth simulation nicely, captures the bouncing with walls and with other balls, and the end position of the bodies tends towards the center, showing the effectiveness of our method. We also see that Backprop's 's performance are superior to the sampling-base CEM, consistent with Allen et al. (2022).
4.2 COMPOSITIONAL INVERSE DESIGN GENERALIZING TO MORE OBJECTS
7

---

Published as a conference paper at ICLR 2024
Table 2: Compositional Generalizaion Across Objects. Experiment on compositional inverse design generalizing to more objects. The confidence interval information is deligated to Table 7 in Appendix B for page constraints.
4-body 24 steps
4-body 44 steps
8-body 24 steps
8-body 44 steps
Method
design obj
MAE
design obj
MAE
design obj
MAE
design obj
MAE
CEM, GNS (1-step)
0,3173
0.23293
0.3307
0,53521
0.3323
0,38632
0.3306
0,53839
CEM, GNS
0,3314
0,25325
0.3313
0.28375
0,3314
0.25325
0.3313
0.28375
Backprop, GNS (1-step)
0.2947
0.06008
0.2933
0,30416
0.3280
0.46541
0.3317
0.72814
Backprop, GNS
0.3221
0.09871
0.3195
0.15745
0.3251
0.15917
0.3299
0.21489
CinDM (ours)
0,2034
0.03928
0.2254
0.03163
0.3062
0.09241
0.3212
0.09249
In this experiment, we test each method's performance in inverse de- sign on larger state dimensions than in training. We utilize the N- body simulation environment as in Sec. 4.1, but instead of considering longer trajectories, we test on more bodies than in training. This setting is also inspired by real-life scenar- ios where the dynamics in test time have more interacting objects than in training (e.g. astronomical simula- tion and biophysics). Specifically, all methods are trained with only 2-body interactions with 24 time steps, and tested with 4-body and 8-body inter- actions for 24 and 44 time steps us- ing Eq. 8. This is a markedly more challenging task than generalizing to more time steps since the methods need to generalize
02
6.4
0.6
6.6
1.
(8) 2-body 54 steps
(b) 4-body 44 steps
Figure 2: Example trajectories for N-body dataset with compositional inverse design in time (a) and bodies (b). The circles indicate CinDM-designed trajectory for the balls, drawn with every 2 steps and darker color indicating later states. The central star indicates the design target that the end state should be as close to as possible. "+" indicates ground-trute trajectory simulated by the solver.
much larger state space than in training. For N-body interaction, there N(N - 1)/2 pairs of 2-body interactions. The case with 44 time steps adds difficulty by testing generalization in both state size and time (composing 28 x 3 84 diffusion models for CinDM).
For the base network architecture, the U-Net in Backprop cannot generalize to more bodies due to U- Net's fixed feature dimension. Thus we only use GNS as the backbone architecture in the baselines. In contrast, while our CinDM method also uses U-Net as base architecture, it can generalize to more bodies due the compositional capability of diffusion models. The results are reported in Table 2.
From Table 2, we see that our CinDM method outperforms all baselines by a wide margin in both the design objective and MAE. On average, our method achieves an improvement of 15.6% in design objective, and an improvement of 53.4% in MAE than the best baseline. In Fig. 2 (b), we see that our method captures the interaction of the bodies with the wall and each other nicely and all bodies tend towards center at the end. The above results again demonstrate the strong compositional capability of our method: it can generalize to much larger state space than seen in training.
4.3 2D COMPOSITIONAL DESIGN FOR MULTIPLE AIRFOILS
In this experiment, we test the methods' ability to perform inverse design in high sigh-dimensional space, for multiple 2D airfoils. We train the methods using flow around single manooml--sampled shape, and in the test time, ask it to perform inverse design for one or more airfoils. The standard goal for airfoil design is to maximize the ratio between the total lift force and total drag force, thus im- proving aerodynamic performance and reducing cost. The multi-airfoil case represents important scenario in real-life engineering where the boundary shape that needs to be designed is more com- plicated and cuttof-disttibutiin than in training, but can be constructed by composing multiple parts. Moreover, when there are multiple flying agents, they may use formation flying to minimize drag. as has been observed in nature for migrating birds (Lissaman & Shollenberger, 1970; Hummel, 1995)
8

---

Published as a conference paper at ICLR 2024
Table Generalization Across Airfoils. Experiment results for multi- airfoil compositional design.
airfoil
2 airfoils
Method
design obj
lift-to-drag ratio
design obj
lift-to-drag ratio
CEM, FNO
0.0932
1.4005
0.3890
1.0914
CEM, LE-PDE
0.0794
1.4340
0.1691
1.0568
Backprop. FNO
0.0281
1.3300
0.1837
0.9722
Backprop, LE-PDE
0.1072
1.3203
0.0891
0.9866
CinDM (ours)
0.0797
2.1770
0.1986
1.4216
and adopted by humans in aerodynamics (Venkataramanan et al., 2003). For the ground-truth solver that generates a training set and performs evaluation, we use Lily-Pad (Weymouth, 2015). The fluid state Ug at each time step represented by 64 64 grid cells where each cell has three dynamic features: fluid velocity Uz. U/y, and pressure. The boundary ? is represented by a 64 x 64 x 3 tensor, where for each grid cell, it has three features: a binary mask indicating whether the cell inside boundary (denoted by 1) or in the fluid (denoted by 0), and relative position (Av, Ay) between the cell center to the closest point on the boundary. Therefore, the boundary has 64 x 64 x 3 = 12288 dimensions, making the inverse design task especially challenging.
For CinDM, we use U-Net as the back- bone architecture and train it to denoise the trajectory and boundary. In the test time, utilize Eq. to compose multiple air- foils into a formation. For both CEM and Backprob, we the state-of-the-art ar- chitecture of FNO (Li et al., 202.) and LE- PDE (Wu al., 2022a). For all methods, to improve design stability, we use the de- sign objective of J liff + drag and evaluate both this design objective and the lift-to-drag ratio. The results are Table 3. Details are provided in Appendix D.
(a) Formation flying of airfoils A and B
(b) Single flying of airfoil A
(c) Single flying of airfoil
The table shows that although CinDM has similar design objective baseline methods, it achieves a much higher lift-to- drag ratio than the baselines, especially in the compositional case of 2 airfoils. Fig. 9 and Fig. 10 show examples of the de- signed initial state and boundary for the 2- airfoil scenario, for our model and "CEM,
Figure 3: Discovered formation flying. In the 2-airfoil case, our model's designed boundary forms a "leader" and "follower" formation (a), reducing the drag by 53.6% and increases the lift-to-drag ratio by 66. % compared to each airfoil flying separately (b)(c). Col- ors represent fluid vorticity.
FNO" baseline, respectively. We see that while our CinDM can design a smooth initial state and reasonable boundaries, the baseline falls into adversarial modes. A surprising finding is that our model discovers formation flying (Fig. 3) that reduces the drag by 53.6% and increases the lift-to-drag ratio by 66. 66.1% compared to each airfoil flying separately. The above demonstrates the capability of CinDM to effectively design boundaries that are more complex than training, and achieving much better design performance.
CONCLUSION
In this work, we have introduced Compositional Inverse Design with Diffusion Models (CinDM), novel paradigm and method to perform compositional generative inverse design. By composing the trained diffusion models on subsets of the design variables and jointly optimizing the trajectory and the boundary, CinDM can generalize to design systems much more complex than the ones seen in training. We've demonstrated 2D airtil tasks, and believe that the techniques presented in this paper are general (Appendix J), and may further be applied across other settings such material, drug, and molecule design.
model's compositional inverse design capability in N-body and

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank Boai Sun and Haodong Feng for suggestions on Lily-Pad simulation. We thank the anony- mous reviewers for providing valuable feedback manuseript. We also gratefully acknowledge the support of Westlake University Research Center for Industries of the Future and Westlake Uni- versity Center for perforreffrm Computing.
The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding entities.
REFERENCES
Anurag Ajay, Seungwook Han, Yilun Du, Shaung Li, Abhi Gupta, Tommi Jaakkola, Josh Tenen- baum, Leslie Kaclbling, Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint arXiv:2309.0885 2023. 3
Kelsey Allen, Tatiana Lopez-Guevara, Kimberly Stachenfeld, Alvaro Sanchez Gonzalez, Peter Battaglia, Jessica B Hamrick, and Tobias Pfaff. Inverse design for ftuid-structure interactions us- ing graph network simulators. In Advances in Neural Information Processing Systems, volume 35,
pp. 13759-13774. Curran Associates, Inc., 2022. 1, 2, 3, 4, 6, 7,
W Kyle Anderson and V Venkatakrishnan. Aerodynamic design optimization on unstructured grids with a continuous adjoint formulation. Computers Fluids, 28(4-5):443-480, 1999. 2
Navid Ansari, Hans peter Seidel, Nima Vahidi Ferdowsi, and Vahid Babaei. Autoinverse: Uncer- tainty aware inversion of neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
https: /pwpreeveeeeeeeeeeeeeeeeeeeeeeee net forum? Potorumidd-diie d ( LAbOb. 3,6
Michael Athanasopoulos, Hassan Ugail, and Gabriela González Castro. Parametric design of aircraft geometry using partial differential equations. Advances in Engineering Software, 40(7):479 486, 2009. 1
Arghya Bhowmik, Ivano E Castelli, Juan Maria Carcia-Lastra, Peter Bjøm Jørgensen, Ole Winther, and Tejs Vegge. perspective on inverse design battery interphases using multi-scale mod- elling, experiments and generative deep learning. Energy Storage Materials, 21:446 2019. 1,2
Victor Blomqvist. Pymunk Iutorials, 2007. URL http: /ttp:///wwwww///////////////. pywppk.mmkkkkmm http://pwmmkkk..//////eeeeeeeeee tutorials. html. 13
Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations, 2022. URL https: httttttt//p/eeeeeee //tpp//pppeee..eee//e/e/ee
net/Eorum? idw.siii3xyxee....
Can Chen, Yingxue Zhang, Xue Liu, and Mark Coates. Bidirectional learning for offline model- based biological sequence design, 2023. URL https: http:///pppeeeieeeeee net forum??rm id= luEG3 9LM5-
Samuel Cohen, Rendani Mbuvha, Tshilidzi Marwala, and Mare Deisenroth. Healing products of gaussian process experts. In International Conference Machine Learning, pp. 2068-2077. PMLR, 2020. 3
Stelian Coros, Bernhard Thomaszewski, Gioacchino Noris, Shinjiro Sueda, Moira Forberg, Robert W Sumner, Wojciech Matusik, and Bernd Bickel. Computational design of mechanical characters. ACM Transactions on Graphics (TOG), 32(4):1-12, 2013.
Marjolein Dijkstra and Erik Luijten. From predictive modelling to machine learning and reverse engineering of colloidal self-assembly Nature materials, 2066):72--77,, 2021. 2
Tao Du, Kui Wu, Andrew Spielberg, Wojciech Matusik, Bo Zhu, and Eftychios Sifakis. Functional optimization of fluidic devices with differentiable stokes flow. ACM Transactions on Graphics (TOG), 19 2020a. 3
10

---

Published as a conference paper at ICLR 2024
Yilum Du and Igor Mordatch. Implicit generation and generalization in energy- basey models. arXiv preprint arXiv:1903.0800 2019. 4
Yilun Du, Toru Lin, and Igor Mordatch. Model based planning with energy based models. CORL, 2019.
Yilun Du, Shuang Li, and Igor Mordatch. Compositional visual generation with energy based mod- els. In Advances in Neural Information Processing Systems, 2020b. 3, 4
Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fer- gus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and meme. arXiv preprint arXiv:2302.11552, 2023. 2. 4,
Nikolaos Gkanatsios, Ayush Jain, Zhou Xian, Yunchu Zhang, Christopher Atkeson, and Katerina Fragkiadaki. Energy-based models as zero-shot planners for compositional scene rearrangement. arXiv preprint arXiv:2304. /4391, 2023. 3
Spencer L Gordon, Manav Kant, Eric Ma, Leonard J Schulman, and Andrei Staicu. Identifiability of product of experts models. arXiv preprint arXiv:2310.09339,, 2023. 3
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural
computation, 14(8):1771-1800, 2002. 3, 6
Dietrich Hummel. Formation flight as an energy-saving mechanism. Israel Journal of Ecology and
Evolution, 213)2261--77, 1995. 8
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. 18
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Huang. tutorial on energy-based learning. Predicting structured data, 1(0), 2006. 4
Shuang Li, Xavier Puig, Chris Paxton, Yilun Du, Clinton Wang. Linxi Fan, Tao Chen, De-An Huang, Ekin Akyürek, Anima Anandkumar, et al. Pre-trained language models for interactive decision- making. Advances in Neural Information Processing Systems, 35:31 199-31222, 2022. 3
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Arizzadenesheli, Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. URL https: /ttp:////peeeeeeeeeeemeee .nt forum? c PNQVvmmn.
Peter BS Lissaman and Carl A Shollenberger. Formation flight of birds. Science, 168(3934):1003- 1005, 1970. 8
Nan Liu, Shuang Li, Yilun Du, Josh Tenenbaum, and Antonio Torralba. Learning compose visual relations. Advances in Neural Information Processing Systems, 34:23166-22 14-23166-23178, 2021.
Nan Liu, Shuang Li, Yilun Du, Antonio Torralba, and Joshua B Tenenbaum. Compositional visual generation with composable diffusion models. arXiv preprint 2rXiv:2206.01714, 2022. 2, 3, 4
Sean Molesky, Zin Lin, Alexander Y Piggott, Weiliang Jin, Jelena Vucković, and Alejandro W Rodriguez. Inverse design in nanophotonics. Nature Photonics, 12(1 1766566596660 2018. 1,2
Weili Nie, Arash Vahdat, and Anima Anandkumar. Controllable and compositional generation with latent-space energy- based models, Advances Neural Information Processing Systems, 34, 2021.
Ryan Po and Gordon Wetzstein. Compositional 3d scene generation using locally conditioned dif- fusion. arXiv preprint arXiv:2303.1222 12218 2023. 3
Simiao Ren, Willie Padilla, and Jordan Malof. Benchmarking deep inverse models over time, and the aeural-adjoint method. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 38 -8. Cur- ran Associates, Inc., 2020. URL httpsss/pppccceeeng.....eeee hwwp://wcc......eoeoee.ee....o.. ccpppp.. files/ optp://00001000100000ffff4444222
3, 6,

---

Published as a conference paper at ICLR 2024
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical Image Computing and Computer-Assiste Intervention- MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceed- ings, Part III 18, pp. 234-241. Springer, 2015. 7, 13, 17, 18
Reuven Y Rubinstein and Dirk P Kroese. The cross-entropy method: a unified approach to com- binatorial optimization, Monte- Carlo simulation, and machine learning, volume 133. Springer, 2004. 2, 6, 7,
Mohammad Saghafi and Roham Lavimi. Optimal design of nose and tail of an autonomous under- water vehicle hull to reduce drag force using numerical simulation. Proceedings of the Institution of Mechanical Engineers, Part M: Journal of Engineering for the Maritime Environment, 234(1):
76-88, 2020. 1,
Alvaro Sanzaez- Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning simulate complex physies with graph networks. In International Confer- ence on Machine Learning, pp. 8459-8466 PMLR, 2020. 7, 13
Pete Shinners. Pygame: A set of Python modules designed for writing video games, 2000. URL
https:/ /twp::///wwww..... . pygame org/. 13
Saulius Tautvaišas and Julius Žilinskas. Heteroscedastic bayesian optimization using generalized product of experts. Journal of Global Optimization, pp. 1-21, 2023. 3
Brandon Trabucco, Aviral Kumar, Xinyang Geng, and Sergey Levine. Design-bench: Benchmarks for data-driven offline model-based optimization, 2021. URL http:s///pppeeeee..eee.e
Forum?id-cQP 3 3
Julen Urain, Anqi Li, Puze Liu, Carlo D'Eramo, and Jan Peters. Composable energy policies for
reactive motion generation and reinforcement learning. arXiv preprint arXiv:2105.04966, 2021.
Sriram Venkataramanan, Atilla Dogan, and William Blake. Vortex effect modelling in aircraft for- mation flight. In AIAA atmospheric flight mechanics conference and exhibit, pp. 5385, 2003.
Zihao Wang, Lin Gui, Jeffrey Negrea, and Victor Veitch. Concept algebra for text-controlled vision models. arXiv preprint arXiv:2302.03699; 2023.
Gabriel D Weymouth. Lily pad: Towards real-time interactive computational fluid dynamics. arXiv preprint arXiv: 1XXv:1510000686 2015. 9, 17, 18
Tailin Wu, Takashi Maruyama, and Jure Leskovec. Learning accelerate partial differential equa- tions via latent global evolution. Advances in Neural Information Processing Systems, 35:2240-
2253, 2022a. 3, 9
Tailin Wu, Megan Tjandrasuwita, Zhengxuan Wu, Xuelin Yang, Kevin Liu, Rok Sosie, and Jure Leskovee. Zeroc: A neuro-symbolic model for zero-shot concept recognition and acquisition at inference time. Advances Neural Information Processing Systems, 35:9828-999 -9800, 2022b. 3
Mengjiao Yang, Yilun Du, Bo Dai, Dale Schuurmans, Joshua Tenenbaum, and Pieter Abbeel. Probabilistic adaptation of text-to-video models. arXiv preprint arXiv:2306.0187. 2023a. 3
Zhutian Yang, Jiayuan Mao, Yilun Du, Jiajun Wu, Joshua B Tenenbaum, Tomás Lozano-Péree and Leslie Pack Kaelbling. Compositional diffusion-based continuous constraint solvers. arXiv 2023b.
preprint arXiv:2309.0090000
Xuan Zhang, Limei Wang. Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, et al. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08425 2023. 3
Qingqing Zhao, David B. Lindell, and Gordon Wetzstein. Learning to solve PDE-constrained inverse problems with graph networks. In /CML, 2022.
1,4
12

---

Published as a conference paper at ICLR 2024
A ADDITIONAL DETAILS FOR COMPOSITIONAL INVERSE DESIGN IN TIME
This section provides additional details for Section 4.1 and Section 4.2. In both sections, we use the same dataset for training, and the model architecture and training specifics are the same for both sections.
Dataset. We use two Python packages Pymunk (Blomqvist, 2007) and Pygame (Shinners, 2000) to generate the trajectories for this N-body dataset. We use 4 walls and several bodies to define the simulation environment. The walls are shaped as a 200 x 200 rectangle, setting elasticity to 1.0 and friction to 0.0. A body is described as ball (circle) with a radius of 20, which shares the same elasticity and friction coefficient as the wall it interacts with. The body is placed randomly within the boundaries and its initial velocity is determined using a uniform distribution e ~ U01000, 100)- We performed 2000 simulations, for balls, balls, and 8 balls in each simulation. Each simulation has time step of 1/60 seconds, consisting of 1000 steps in total. During these simulations, we record the positions and velocities of each particle in two dimensions at each time step to generate 3 datasets with shape of [N,,, NN,,,N,, Nyl. NN, means number of simulations, Ne means number of time steps, N₂ is number of bodies, Ny means number of features. The input of one piece of data shaped as [B, 1, NN, x Ny). B is batch size, for example, [32, 1. 8] for 2 bodies conditioning on only one step. Before training the model, the final data will be normalized by dividing it by 200 and setting the time resolution to four simulation time steps.
Model structure. The U-Net (Ronneberger et al., 2015) consists of three modules: the down- sampling encoder, the middle module, and the upsampling decoder. The downsampling encoder comprises 4 layers, each including three residual modules and downsampling convolutions. The middle module contains 3 residual modules, while the upsampling decoder includes four layers, each with 3 residual modules and upsampling. We mainly utilize one-dimensional convolutions in each residual module and incorporate attention mechanisms. The input shape of our model is de- fined as [batch size, n sleps, n features), and the output shape follows the same structure. The GNS Saancez--Gozzale et al., 2020) model consists of three main components. First, it builds an undirected graph based on the current state. Then, it encodes nodes and edges on the constructed graph, using message passing to propagate information. Finally, it decodes the predicted accelera- tion and utilizes semi-implicit Euler integration to update the next state. In our implementation of GNS, each body represents a node with three main attributes: current speed, distance from the wall, and particle type. We employ the standard k-d tree search algorithm to locate adjacent bodies within a connection radius, which is set as 0.2 twice the body radius. The attribute of an edge is the vector distance between the two connected bodies. More details are in Table 4.
Training. We utilize the MSE (mean squared error) as the loss function in our training process. Our model is trained for approximately 60 hours on a single Tesla V100 GPU, with a batch size of 32, employing the Adam optimizer for I million iterations. For the first 600,000 steps, the learning rate is set to le-4. After that, the learning rate is decayed by 0.5 every 40,000 steps for the remaining 400,000 iterations. More details provided in Table 5.
To perform inverse design, we mainly trained the following models: U-Net, conditioned on step and capable of rolling out 23 steps; U-Net (single step), conditioned on step and limited to rolling out only 1 step; GNS, conditioned on step and able to roll out 23 steps: GNS (single step), condi- tioned on step and restricted to rolling out only I step; and the diffusion model. Simultancously, we conducted a comparison to assess the efficacy of time compose by training a diffusion model with 44 steps directly for inverse design, eliminating the requirement for time compose. The results and analysis are shown in Appendix C. Throughout the training process, we maintained consistency in the selection of optimizers, datasets, and training steps for these models.
Inverse design. The center point is defined as the target point, and our objective is to minimize the mean squared error (MSE) between the position of the trajectory's 's last step and the target point. To compare our CinDM method, we utilize U-Net and GNS as forward models separately. We then use CEM (Rubinstein & Kroese, 2004) and Backprop (Allen et al., 2022) for inverse design with con- ditioned state (æo, Noo UaO, Uyo) used as input, and multiple trajectories of different bodies rolled out. While the CEM algorithm does not require gradient information, we define a parameterized Gaussian distribution and sample several conditions from it to input into the forward model for pre- diction. After the calculation of loss between the prediction and target, the best gest-performing samples selected to update the parameterized Gaussian distribution. Through multiple iterations, we can
13

---

Published as a conference paper at ICLR 2024
Table 4: Hyperparameters of model architecture for N-body task.
Hyperparameter name
23-steps
1-step
Hyperparameters for U-Net architecture:
Channel Expansion Factor
1,2,4,8,
2, 1,
Number of downsampling layers
4
Number of upsampling layers
4
4
Input channels
8
8
Number of residual blocks for each layer
3
Batch size
32
32
Input shape
[32, 24,
[32, 2, 8
Output shape
32, 24, 8
[32, 2,
Hyperparameters for GNS architecture:
Input steps
Prediction steps
23
Number of particle types
1
Connection radius
0.2
0.2
Maximum number of edges per node
6
Number of node features
8
8
Number of edge features
3
Message propagation layers
5
5
Latent size
64
64
Output size
46
2
Hyperparameters for the U-Net in our CinDM:
Diffusion Noise Schedule
cosine
cosine
Diffusion Step
1000
1000
Channel Expansion Factor
(1,2,4,88
(1, 1,2,1111
Number of downsampling layers
4
4
Number of upsampling layers
4
Input channels
8
8
Number of residual blocks for each layer
3
3
Batch size
32
32
Input shape
[32, 24, 8
[32 2, 8]
Output shape
32, 24,
[32, 2, 8
sample favorable conditions from the optimized distribution to predict trajectories with low loss val- ues. Backpropagation heavily relies on gradient information. calculates the gradient of the loss concerning the conditions and updates the conditions using gradient descent, ultimately designing conditions that result promising output.
During training, we can only predict finite number of time steps based on conditional states, but the system evolves over an infinite number of time steps starting from an initial state in real-world physical processes. To address this, we need to combine time intervals while training single model capable of predicting longer trajectories despite having limited number of training steps. For the forward model, whether using U-Net GNS, we rely an intermediate time step derived from the last prediction as the condition for the subsequent prediction. We iteratively forecast additional time steps based on single initial condition in this manner. As for the forward model (single step), we employ an autoregressive approach using the last step of the previous prediction to predict more steps.
14

---

Published as a conference paper at ICLR 2024
Table Hyperparameters of training for N-body task.
Hyperparameter name
23-steps
1-tepp
Hyperparameters for U-Net training:
Loss function
MSE
MSE
Number of examples for training dataset
3 x 100
3 x 10
Total number of training steps
106
106
Batch size
32
32
Initial learning rate
x 10
1 X 10 10-4
Number of training steps with a fixed learning rate
10
6 10
Learning rate adjustment strategy
StepLR Adam 10 0.95
StepLR
Optimizer
Adam 1 10 0.95
Number of steps for saving checkpoint Exponential Moving Average decay rate
Hyperparameters for GNS training:
Loss function
MSE
MSE
Number of examples for training dataset
x 10s
3 x 105 1 10 32
Total number of training steps
106
Batch size
32
Initial learning rate
x 10
1 x 10 44
Number of training steps with fixed learning rate
6 10a
6 10
Learning rate adjustment strategy
StepLR Adam
StepLR
Optimizer
Adam
Number of steps for saving checkpoint Exponential Moving Average decay rate
1 x 104
1 x 10
0.95
0.95
Hyperparameters for our CinDM training:
Loss function
MSE
MSE
Number of examples for training dataset
x 10
3 x 10 1 10
Total number of training steps
106 32
Batch size
32
Initial learning rate
x 10 10-4 4
1 10-4
Number of training steps with fixed learning rate
6 x 10
6 105
Learning rate adjustment strategy
StepLR Adam
StepLR
Optimizer
Adam
Number of steps for saving checkpoint
x 10
1 10a
Exponential Moving Average decay rate
0.95
0.95
B FULL RESULTS FOR COMPOSITIONAL INVERSE DESIGN OF THE N-BODY TASK
Here provide the full statistical results including the 95% confidence interval (for 500 instances) for N-body experiments, including compostional inverse design in time and more objects. Specif- ically, Table shows detailed results for Table in Section 4.1; and Table 7 extends Table 2 in Section 4.2.
Table 6: Compositional Generalization Across Time. The confidence interval information is pro- vided in addition to Table 1.
2-hedy 24
2-beb: 44
Mathod
design obj
MA
doiign dbi
MAE
dosign dbj
MAE
doisn obj
MAE
CEM, GNS (1 (sseep)
0 1.2622 e 0/00MO
B. 1333655 0.0000
02204 0.0072
0.15378
0.00 123
122701 00000
0.21277
0101264
22775 00000
0,21706 U 25c
0.00662
0.0064
0.14637
0.000577
0.1R155
0.06600
0.0002
0200065
0.00679
U-No (1-stpp)
0.00m5
0.07220
0.003
01
0/007
8A8899
0100395
0.0002
0.00096
0.00042
U e
0109400
0199651
0.0005
0/4350 a 0104667
Radspppp GNS (1-eap))
01 1482222 0.0060
BU08319
0.9008
0.1497 0.000)
0,88800 1 0.00064
61511 L.0002
0.09621
L.005222
0.1854 0.0002
0/04044 1 0.000RS
Backgro, GNS
0 240407
= 000067
B.W778 010015
U.3078
0.0072
0.11017
40000000
0 27262 00007
0.12395
0110667
U22522
000073
0.13M63 + 00022
Buckprop steg)
0-00078
Badipppp, UlN
0.1228 0.0040
0.00223
0.000004
0.00970 0.00066
0.0043
CInDM (ours)
0.1160 0.0019
0.01264 0.00057
0.1258 0.0000
8,80917 0,00000
0.1447 a 0.0000
0,80955 e 010116
0.1650 0U0045
8,01064 0.0017
15

---

Published as a conference paper at ICLR 2024
Table 7: Compositional Generalizaion Across Objects. The confidence interval information is provided in addition to Table 2.
stops
4-hody 44 stogs
S-body 24
stops
Methed
CBM, GNS (1-sep)
0:3173 Dri30
0.23223 a Buttenn
33307 + OOKEE
0.5333 0:00967
0.333 UIRSS
0.38632 910033
0.330
0.53339 e (1
CEM, GNS
0.3314200002 0.0023
0,25325 a 0.00266
0 03335 0,023
0.28375 y 0.00336
0.3514 0J0023
0,25325 0,00699
0.3313 = 0L0023
0.28375 a 0.00336
Backppop (1stegg)
0.2967
0.3317
0.72814
Btadppupp
0,00511
CinDM (ours)
0.2834 0,0032
0,83928 : 0J00361
0.2254 0,0044
0.13663 0,00251
0.3062 2 0U0021
0,8241 0,00200
0,3312 0.0023
0,09249 0.00276
C ADDITIONAL BASELINE FOR TIME COMPOSITION OF THE N-BODY TASK
We also make comparison with simple baseline that performs diffusion over 44 steps directly without time composition. We designed this baseline to verify the effectiveness of our time- compositional approach. This baseline takes the same architecture as CinDM but with 44 time steps instead of 24 time steps, thus has almost twice of number of parameters in CinDM. The results are displayed in Table 8, which indicates that this sample baseline is outperformed by our CinDM. Its reason may be the difficulty in capturing dynamics across 44 time steps simultancously using single model, due to the presence of long-range dependencies. In such cases, 2t-s diffusion model proves to be more suitable. Hence, when dealing with designs that involve larger number of time steps, employing time composition more effective approach, with lower cost and better performance.
Table 8: Compositional Generalization Across Time. Comparison to a baseline that directly
diffuses 44 steps without time composition.
Methods
#paraaetters(Milllon))
design obj
2-body steps
MAE
design obj
4-body steps
MAE
Our method
20.76 20.76M
0.13 + 0.0087
0.00695 + 0.00067
0.2281 = 0.0145
0.03195 0.00705
Directly diffuse steps 44.92M
0.2779 0.0197
0.00810 0.00200
0.2986 0.01481
0.05166
16

---

Published as a conference paper at ICLR 2024
D ADDITIONAL DETAILS FOR COMPOSITIONAL INVERSE DESIGN OF 2D AIRFOILS
D.1 DETAILS FOR THE MAIN EXPERIMENT
Dataset. We use Lily-Pad (Weymouth, 2015) as our data generator (Fig. 5). We generate 30,000 ellipse bodies and NACA airfoil boundary bodies and perform fluid simulations around each body. The bodies are sampled by randomizing location, thickness, and rotation between respective ranges. Each body is represented by 40 two-dimensional points composing its boundary. The spatial reso- lution is 64 x 64 and each cell is equipped with temporal pressure and velocities in both horizontal and vertical directions. Each trajectory consists of 100 times steps. To generate training trajectories, sliding time window over the 100 time steps. Each time window contains state data of T = 6 time steps with a stride of 4. So each original trajectory amounts to 25 training trajectories, and we get 750,000 training samples in total.
Model architecture. We use U-Net (Ronneherger et al., 2015) as our backhone for denoising from a random state sampled from a prior distribution. Without considering the mini-batch size dimension, the input includes a tensor of shape (37' + 3) x 64 x 64, which concatenates flow states (pressure, velocity of horizontal and vertical directions) of T' time steps and the boundary mask and offsets of horizontal and vertical directions along the channel dimension, and additionally the current diffusion step s. The output tensor shares the same shape with the input except s. The model architecture is illustrated in Fig. 4. The hyperparameters in our model architecture are shown in Table 9.
Training. We utilize the MSE (mean squared error) between prediction and a Gaussian noise as the loss function during training. We take batch size of 48 and run for 700,000 iterations. The learning rate is initialized as 1 x 10-4, Training details are provided in Table 10.
Evaluation of design results. In inference, we set À in Eq. 3 as 0.0002. We find that this À could get the best design result. More discussion on the selection of is presented in Appendix method and each airfoil design task (one airfoil or two airfoils), we conduct 10 batches of design, and each batch contains 20 examples. After we get the designed boundaries, we input them into Lily-Pad and run the simulation. To make the simulation more accurate and convincing, we use 128 x 128 resolution of the flow field, instead of 64 x 64 as in the generation of training data. Then the calculated horizontal and vertical flow force to compute and lift-to-drag ratio. In each batch, we choose the best-designed boundary (or pair of boundaries in two airfoil scenarios) and then we report average values regarding the two metrics over 10 batches.
For each
two metrics: -lift drag
Table 9: Hyperparameters used in 2D diffusion model architecture
Number of downsampling blocks
4
Number of upsampling blocks
4
Input channels
21
Number of residual blocks for each layer
2
Batch size
48
Input shape
[48, 21, 64, 64]
Output shape
48, 21, 64, 64
D.2 SURROGATE MODEL FOR FORCE PREDICTION
Model architecture. In the 2D compositional inverse design of multiple airfoils, we propose a neural surrogate model g. to approximate the mapping from the state U, and boundary y to the lift and drag forces, so that the design objective T differentiable to the design variable Upo +% The input of our model is a tensor comprising pressure, boundary mask, and offsets (both horizontal and vertical directions) of shape 4 x 64 x 64 for a given time step. The output is the predicted drag and lift forces of dimension 2. Boundary masks indicate the inner part (+1) and outside part (0) of a closed boundary. Offsets measure the signed deviation of the center of each square 64 x 64 grid from the boundary in horizontal and vertical direction respectively, where the deviation of a given point is defined as its distance to the nearest point on a boundary. If two or more boundaries appear in sample, the input mask (resp. offsets) is given by the summation
17

---

Published as a conference paper at ICLR 2024
64
0,
U1
64
U-Net
Ud
B3-1
3xT
3xT
3
Input
Output
Repeat S times
Figure 4: Diffusion model architecture of 2D inverse design.
Figure 5: Example of Lily-Pad simulation.
of masks (resp. offsets) of all the boundaries. Notice that since the input boundaries assumed not to be overlapped, the summed mask and offset are still valid. The model architecture half of U-Net Ronneberger al. (2015) where only take the down-sampling part to embed the input features to 512-dimmnsionna representation; then we use a linear transformation to output forces.
Dataset. We use Lily-Pad (Weymouth, 2015) to generate simulation data with 1, 2, boundaries to train and evaluate the surrogate model. Boundaries are a mixture of ellipses and NACA airfoils. We generate 10,000 trajectories for the training dataset and 1,000 trajectories for the test dataset. Each trajectory consists of 100 time steps. We use pressure as features and lift and drag forces labels. Thus we have 3 million training samples and 300 thousand testing samples in total.
airfoil
Training. We use MSE (mean squared error) loss between the ground truth and predicted forces to train the surrogate model. The optimizer is Adam (Kingma & Ba, 2014). The batch size is 128. The model is trained for 20 epochs. The learning rate starts from x 10 and multiplies a factor of 0.1 every five epochs. The error 0.04, smaller than 5% of the average force in the training dataset.
18

---

Published as a conference paper at ICLR 2024
Table 10: Hyperparameters used in 2D diffusion model training.
Loss function
MSE
Number of examples for training dataset
3 x 10
Total number of training steps
10a
Batch size
48
Initial learning rate
X 10
Number training steps with a fixed learning rate
10
Learning rate adjustment strategy
StepLR
Optimizer
Adam
Number of saving checkpoint
700
Exponential Moving Average decay rate
0.995
E VISUALIZATION OF N-BODY INVERSE DESIGN.
Examples of N-body design results are provided in this section. Figure shows the results of using the backpropagation algorithm and CinDM design 2-body 54-time step trajectories. The results of designing 2-body 54-time steps trajectories using CEM and CinDM are provided in Figure 7. Figure are the results of designing44-time 44 time steps trajectories using CEM, backpropagation, and CinDM.
19

---

Published as a conference paper at ICLR 2024
(c)
y
(e)
Figure 6: 54 time steps trajectories of bodies after performing inverse design using the back- propagation algorithm. Figures (a). (b), (c), and (d) represent the trajectory graphs obtained using GNS, GNS (single step), U-Net, and U-Net (single step) the forward models, respectively. And (e) is the result of CinDM. The legend of this figure is consistent with Figure 2.
20

---

Published as a conference paper at ICLR 2024
(a)
(d)
(c)
Figure 7: 54-step trajectories of 2 bodies after performing inverse design using the CEM al- gorithm. The trajectory graphs (a), (b), (c), and (d) depict the outcomes using different forward models such as GNS, GNS (single step), U-Net, and U-Net (single step) respectively. Additionally, figure (e) demonstrates the result generated by CinDM. The legend of this figure is consistent with Figure 2.
21

---

Published as a conference paper at ICLR 2024
(c)
(e)
Figure 8: 44-time steps trajectories of 4 bodies after performing inverse design using CEM. Figures (a), (b), (c), and (d) represent the trajectory graphs obtained using GNS, GNS (single step), U-Net. and U-Net (single step) the forward models, respectively. And (c) is the result of CinDM. The legend of this figure is consistent with Figure 2.
22

---

Published as a conference paper at ICLR 2024
F VISUALIZATION RESULTS OF 2D INVERSE DESIGN BY OUR CINDM
We show the compositional design results of our method in 2D airfoil generation in Figure 9.
Curpei i hertorta vilate
Dumzal vitulvvact
Dammaee
2ourgin herazndd vduale
D4T2E2 vattallvee
hetioraal vilads
2 wttolviac
Damale
hertanna valata
DuTzE4 vetcclvvac
Domple
Coumpns heniziid aduate
Dendes vattalveesly
Dundi
hertioaal vild's
Sarzed wtolwwes
Danale
Figure 9: Compositional design results of our method in 2D airfoil generation. Each row repre- sents an example. We show the heatmap of velocity in horizontal and vertical direction and pressure in the initial time step, inside which we plot the generated airfoil boundaries.
23

---

Published as a conference paper at ICLR 2024
G
SOME VISUALIZATION RESULTS OF 2D INVERSE DESIGN BASELINE.
We show some 2D design results of our baseline model in Fig. 10.
40
30
30
20
10
10
10-
0
20
20
40
80
40
G0
(a)
60
60
60
50
50
50
40
40
30
30
(b)
50
50
40
40
40
20
20
20
20
20
20
10
10
10
60
60
50
50
40
40
40
30
30
30
20
20
10
10
10
10
(d)
30
20
20
20
10
0
0
20
40
60
40
60
60
(e)
Figure 10: Design results of FNO with CEM in 2D airfoil generation. Fach row the heatmap of optimized velocities in the horizontal and vertical direction and optimized pressure in the initial time step, inside which we plot the generated airfoil boundaries.
24

---

Published as a conference paper at ICLR 2024
Table 11: Comparison to NABL and cINN for N-body time composition inverse design task.
2-body 24 steps
2-body 34
steps MAF
2-body 44 steps
2-body 54 steps
Method
design ohj
MAE
design obj
design obj
MAE
design ohj
MAE
NABL, U-Net (1-step)
0.1174
0.01650
0.1425
0.01511
0.1788
0.01185
0,2606
0.02042
cINN
0,3235
0.11704
0.3085
0,18015
0,3478
081222
0,3372
0.19296
CinDM (ours)
0.1143
0.01202
0.1251
0.00763
0.1326
0.00695
0.1533
0.00870
Table 12: Comparison to NABL and cINN for 2D airfoils inverse design task.
airfoil
2 airfoils lift-to-drag ratio t
Method
# parameters (Million)
design obj
lift-to- drag ratio t design obj
NABL, FNO
3.29
0.0323
1.3169
0,3071
0.9541
NABL, LE-PDE
3.13
0. 1010
1,3104
0.0891
0.9860
cINN
3.07
1.1745
0.7556
CinDM (ours)
3.11
0.0797
2.177
0.1986
1.4216
H COMPARISON TO ADDITIONAL WORKS
Besides comparison results of baselines shown in the main text, we further evaluated additional two baselines: neural adjoint method+ + boundary loss function (NABL) and conditional invertible neural network (cINN) method for both N-body and airfoils design experiments.
We implement NABL on top of baselines FNO and L.E-PDE in the airfoil design task and U-net in tcompositionalostinaal taskamed as "NABL,, FNO", "NABL,, LE-PDE" and "NABL., U-net" respec- tively. These new NABL baselines additionally use the boundary loss defined by the mean value and 95% significance radius of the training dataset. cINN does not apply to compositional design because the input scale for the invertible neural network function is fixed. Therefore, for the time composition task, trained 4 cINN models, each for one of the time steps: 24, 34, 44, and 54. These models differ only in the input size. The input x to cINN is a vector of size 2 x 4 x T, where 2 is the number of objects, 4 is the number of features and T is the number of time steps. The condi- tion y is set to O, the minimal distance to the target point. For cINN for 2D airfoil design, we adopt 2D coordinates of 40 boundary pointsarewhich is spanned 80-dimensionallensionaal vector, as the input, since the invertible constraint on the cINN model hardly accepts image-like inputs adopted in the main experiments. Therefore we evaluate cINN only in the single airfoil design task. The condition is set as the minimal value of drag lift drag in the training trajectories. In both tasks, the random variable z has a dimension of dim(x)- dim(y). It is drawn from Gaussian distribution and then input to the INN for inference. We also adjust the hyperparameters, such as hidden size and a number of reversible blocks, to make the number of parameters in cINN close to ours for fair comparison.
The results of NABL and cINN are shown in Table 11 and Table 12. We can see that CinDM significantly outperforms the new baselines in both experiments. Even compared to the original baselines (who contains contain "Backprop-") without the boundary loss function, as shown in Table 1 and Table 3, the NABL baselines in both tasks do not show the improvement in the objective for out-of-distribution data. These results show that our method generalizes to out-of-distribution while the original and new baselines struggle outperforms cINN by a large margin in both the time composition and airfoil design tasks. Despite the quantities, we also find that airfoil boundaries generated by cINN have little variation in shape, and the orientation is not as desired, which could incur high drag force in simulation. These results may be caused by the limitation of the model architecture of cINN, which utilizes fully connected layers as building blocks, and thus has an obvious disadvantage in capturing inductive bias of spatial- temporal features. We think it is necessary to extend cINN to convolutional networks when cINN is applied to such high-resolution design problems. However, this appears challenging when the invertible requirement is imposed. In summary, our method outperforms both NABL and cINN in method could be used for flexible compositional design. We use only
generalize the out-of-distributioo.. CinDM also
both tasks. Furthermore, one trained model to generate samples lying in a much larger state space than in training during inference, which is unique advantage of our method beyond these baselines.
25

---

Published as a conference paper at ICLR 2024
I PERFORMANCE SENSITIVITY TO HYPERPARAMETERS, INITIALIZATION AND SAMPLING STEPS.
This section evaluate the effects of different X, initialization and sampling steps on performance of CinDM.
I.1 INFLUENCE OF THE HYPERPARAMETER
Table 13: Effect of X in N-body time composition inverse design.
steps
sleps
steps
2-bods steps
design
ab)
design abg
0:5554
0.0185
0,003855
0.00183
0:001
0U0411 0O001C
0.0069
0,00083
0.00435 = 0.00100
0.1734
0.0046
0,01173
0.1340 =
= 0.002166
0 2.00049
0.3529
0.0300
0.0001 0.00554
0.1133
0.0172
0.0034
0.1385
0,00948 = 0:00299
0.3422
0.00763 0.001122
0.1663
0,00850
0.0000
0.00015
0.1326
0.01171
0.11400 0.00355
0.1670
0.0073
0.01596
0.01055 = 0:00337
U.1573
0.00113
a
0.01026
0.0063
0.01679
0.001
0.001499
0.1634
0.00779
0.1789
0.01
2
0 0.0079
0.02654
0:00193
0.1779 0.021
0.0707 0.00444
02113
0.01447 000130
0,2821 0.0197
0.21153 0.01057
0.2210
0.09155 0.00256
0.0232
0,02100
Table 14: Effect of X in 2D inverse design.
À
obj
lift/drag 1.015=0.2008
0.05
0.7628+0.189)
0.02
0.3849#0.0632
1.079440.
0.01
0.2292+0.0400
1.286+0. 1402
0.005
0.2061+0.0388 +0.0388
1.13
0.002
0.217
1.242940. 1243
0.001
0.2277=0.0451
1,2608+0. 1469
0.0005
0.2465+0.0473
1.110210.177
0.0002
0.1986#0.0431
1.4216+0.1600
0.0001
0.27140.0577
.196240.1284
0.35
0.30
0.25
design obj
0.20
0.10
design olbj steps
0.05
design obj of 34 stops
design oibj steps
design obj 54
0.00
-2
1o910d
Figure 11: Design objective of differrnt À in N-body time composition inverse design.
To evaluate influence the hyperparameter composition and 2D airfoils design task for wide range of X. The results are shown in Table 13, Table 14, Fig 11, Fig 12, and Fig 13, where Table 13 corresponds to Fig 11 and Fig 12 while Table Table 14 corresponds
Eq. 3, we perform inference in both N-body time
Fig 13. Our method demonstrates robustness and consistent performance across a wide range of lambda values. However, if is set too small (50.0001 in the 2D airfoil task, <0.01 in the N-body task), the design results will be subpar because there is minimal objective
26

---

Published as a conference paper at ICLR 2024
0.25
MAE steps
MAE of 34 steps
steps
0.20
MAE 54
0.15
MAE
0.10
0.00
4
3
1
log10A
Figure 12: MAE of different À in N-body time composition inverse design.
lift/drag
lifg/drag design objective
design objective
1.4
1.2
1.0
0.8
0.4
0.2
0.0
4.0
-3.5
3.0
2.5
-2.0
JogıoA
Figure 13: Performance of different À in 2D airfoil inverse design.
guidance incorporated. On the other hand, if A set too large (2 0.01 the 2D airfoil task, or 2 1.0 in the N-body task), there is higher likelihood of entering poor likelihood region, and the preservation physical consistency is compromised. In practical terms, X can be set between 0.01 and 1.0 for the N-body task, and between 0.0002 and 0.02 for the 2D airfoil task. In paper, we choose based on the best evaluation performance, namely we set as 0.4 for the N-body task and 0.0002 for the 2D airfoils task.
I.2 INFLUENCE OF INITIALIZATION
To analyze the sensitivity of initialization in our approach, we follow a similar methodology dis- cussed in Ren et al. (2020). We consider the "re-simulation" error r of target objective y as function of the number of samplings B, where each sampling starts from Gaussian initialization z. We use the simulator to obtain the output ò for each design D from the B design results given the target and compute the "re-simulation" emror L(ū. y). We then calculate the least emror among batch of B design results. This process is repeated for several batches, and the mean least error r is obtained by averaging
these batches.
Table and Fig present the results for the N-body inverse design task. We consider values of B ranging from 10 100, with N 10 batches. The target is set to be 0, which represents the
27

---

Published as a conference paper at ICLR 2024
0.120
24 steps
0.115
34 steps
44 steps
54 steps
0.110
0.105
0.100
0.095
0.090
0.085
0.080
20
100
B
Figure 14: "Re-simulation" error rB of different B in N-hody inverse design.
design
design
3.5
design of airfoils
design of airfoils
0.8
3.0
2.5
objective
0.6
2.0
1.5
0.4
1.0
0.2
0.5
0.0
0.0
30
6
B
Figure 15: Design performance (lift-to-drag) of different B in 2D airfoil inverse design.
distance fixed target point. The results show that r' gradually decreases B increases in the 24-step design, indicating that the design space is well explored and most solutions can be retrieved even with small number of samplings. This demonstrates the efficiency our method in generating designs. Moreover, similar observations can be made when time composition is performed in 34, 44, and 54 steps, indicating the effectiveness of our time composition approach in capturing long-tin time range physical dependencies and enabling efficient generation in a larger design space.
In the 2D inverse design task, the target y is slightly different. Here, we aim minimize the model output (drag lift force). Hence, we adopt the "re-simulation" performance metric, which the lift/drag ratio, as opposed to the "re-simulation" " error used in the N-body task, evaluate sensitivity to initialization. For each B, the lift/drag ratio is chosen the highest value among the simulation results batch of designed boundaries (or boundary pairs for the airfoils design). Any invalid design results, such as overlapping airfoil pairs in the 2-airfoil design, removed from the B results before computing the maximal lift/drag ratio. The reported numbers are obtained by averaging over N 10 batches for each B.
Table 16 and Fig 15 present the results for the 2D airfoils design task. In the airfoil design column, observe that the lift/drag ratio relatively low for B 10, indicating that the design space not
28

---

Published as a conference paper at ICLR 2024
Table 15: Influence of initialization. rB with respect B for N-body inverse design task. Each number is an average over 10 batches.
B
2-body 24 steps
2-body 34 steps
2-body 44 steps
2-body 54 steps
10
0.10122654
0.1022556
0.10542078
0.11227837
20
0. 1.10051144
1.10122902
0.10261874
0.10554917
30
8.09950846
0.10106587
0,10220513
0.10408381
40
0.09928784
0. 0.10066015
0.10173534
0.10409425
50
0.09794939
0.10023642
0.10168899
0.10462530
60
0,09876589
0.09997466
0.10105932
0.10257294
70
0.09858151
0.09979441
0.10124100
0.10179855
80
0.09809845
0.09972977
0.10060663
0.10203485
90
0.09808731
0,09941968
0.10108861
0.10120515
100
0.09734109
0.09912691
0.10056177
0.10135190
Table 16: Influence of initialization. Design performance (lift-t--drag) with respect to for 2D inverse design task. Each number an average over 10 batches.
B
airfoil
2 airfoils
10
1.4505
0.8246
20
2.2725
0.7178
30
2.2049
1.3862
40
2.6506
1.5781
50
2.1355
1.6055
sufficiently explored due its high dimensionality (64 64 in our boundary mask and offsets representation). For B > 20, the lift/drag performance remains steady. In the 2 airfoils design column, the lift/drag ratio increases roughly with B. This attributed to the higher dimensional and more complex design space compared to the single airfoil design task. The stringent constraints on boundary pairs, such as non-overlapping. lead to the presence of complex infeasible regions in the design space. Random initialization may lead to these infeasible regions, resulting in invalid design results. The rate of increase lift/drag ratio becomes slower when 2 30, indicating that majority of solutions have been explored. Despite the training data only containing single airfoil boundary, which lies in a relatively lower dimensional and simpler design space, our model demonstrates strong ability to generalize and efficiently generate designs for this challenging body compositional design problem.
1.3 INFLUENCE OF THE NUMBER OF SAMPLING STEPS IN INFERENCE
design ob)
0.22
design obj
0.18
0.16
0.12
200
400
600
800
1000
sample steps
Figure 16: Design objective of diffierent sampling steps in N-body inverse design.
29

---

Published as a conference paper at ICLR 2024
0.0150
0.0125
MAE
0.0100
0.0075
0.0060
200
400
600
800 1000
sample_steps
Figure 17: MAE of different sampling steps in N-body inverse design.
Fig 16 and Fig 17 illustrate the outcomes of inverse design carried out by CinDM. It is apparent that with an increase in the number of sampling time steps, the design objective gradually decreases. In contrast, the MAE fluctuates within a small range, occasionally rising. This phenomenon can be examined as follows: as the number of sampling steps increases, the participation of the design objective in the diffusion process intensifies. As a result, the designs improve and align more closely decrease in the design objective. However, when
with the design objective, ultimately leading the number of sampling steps increases, the MAE also increases. This is because, with small number of sampling steps, the initial velocities of some designed samples are very small, causing the diffusion of trajectories to be concentrated within narrow range. Consequently, both the true highly concentrated, resulting in a small calculated MAE.
trajectory and the diffused trajectory By analyzing the sensitivity of the design objective and MAE to different sampling steps, we can conclude that CInDM can achieve desired design results that align with design objectives and phys- ical constraints by appropriately selecting a sampling step size during the inverse design process.
J BROADER IMPACTS AND LIMITATIONS
Our method, CinDM, extends the scope of design exploration and enables efficient design and con- trol of complex systems. Its application across various scientific and engineering fields has pro- found implications. In materials science, utilizing the diffusion model for inverse design facilitates the customization of material microstructures and properties. In biomedicine, it enables the struc- tural design drug molecular systems and optimizes production processes. Furthermore, in the acrospace sector, integrating the diffusion model with inverse design can lead to the development of more diverse shapes and structures, thereby significantly enhancing design efficiency and quality.
CinDM combines the advantages of diffusion models, allowing us to generate more diverse and sophisticated design samples. However, some limitations need to be addressed at present. In terms of design quality and exploration space, need to strike balance between different objectives to avoid getting stuck local optima, especially when dealing with complex, nonlinear systems in the real world. We also need to ensure that the designed samples adhere complex multi- scale physical constraints. Furthermore, achieving interpretability in the samples designed by deep learning models challenging for inverse design applications. From cost perspective, training diffusion models requires large datasets and intensive computational resources. The complexity of calculations also hinders the speed of our model design.
Moving forward, we intend to incorporate more physical prior knowledge into the model, leverage multi-modal data for training, employ more efficient sampling methods to enhance training effi- ciency, improve interpretability, and generalize the model to multiple scales.
30


Reviews:
[{"review_id": "zc1Ke535MC", "reviewer": "Reviewer_4Mfs", "summary": "The authors address the complex task of inverse design with what I believe is a rather novel approach. \nA first line of work optimizes over the forward process using an optimization procedure (CEM, gradient based optimization), this suffers from falling into adversarial local optima and potentially poor likelihood of the generated solution.\nTo fight such a behavior, the others propose to optimize over a linear combination of an EBM, accounting for the generation of likely condition, and the design objective.\nIn addition, the authors propose to estimate the EBM in a compositional fashion to simplify learning. \nThe proposed framework is tested through two main sets of experiments: N-body problem and airfold optimization", "questions": "1. Can the authors describe the role of $\\alpha$ line 12, Alg.1 ?\n2. Can the authors comment on the choice of the energy function for the airfold design ? How do we compare to training data ?\n3. Can the authors comment on how to balance $\\lambda$ during the optimization ? Could the optimization end up in a poor likelihood region ?\n3Bis. Can other forward / optimization steps be considered for such a task ?\n4. What is the influence of the number of optimization steps  ?\n5. For the airfold design: what is the relationship between the initial objective function and the reported ratio  ? Which quantity is actually at stake here ? \n6. Can the authors think of any limitation when applying a compositional energy approach ? For instance is it computationally efficient to learn “smaller models” vs one big EBM ?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The authors approach is very interesting.\nThe paper is straightforward and aims at directly addressing the problem it uses. \nIt is clear and fairly well-written. The experiments provided by the authors seem to confirm the validity of the proposed method.", "weaknesses": "I personally found the experiments slightly harder to read compared to the rest of the paper. For other remarks see questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XfgfF8ScQD", "reviewer": "Reviewer_7GzX", "summary": "This works investigates inverse desigs in dynamic systems. The authors look into inverse design while avoiding adversarial samples in order to improve efficiency. The authors proposed a new formulation for inverse design by energy optimization, and introduced the Compositional Inverse Design with Diffusion Models (CinDM), which is able to branch out and generate further designs than observed.", "questions": "Additionally, I wonder whether changing its value will lead to different results in the experiments.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The generative optimization structure containing both the energy-based model and the design objective is quite unique and novel. It enables the optimization problem for design to be more readily approached via the joint learning procedure.\n2. The experiments conducted in Section 4 are complete which explains well the questions raised at the beginning of the section.\nOverall, the ability shown in the work to generalize is quite impressive and seems promising with potential to be applied to more applications.", "weaknesses": "1. This is more of a question. On the joint optimization, it is trying to minimize the energy component which is calculated from the trajectories and the boundary, and minimizing the design objective as well. It is proposed to achieve this by optimizing the design and the trajectory at the same time. In the joint optimization formulation as in Eqn.(3), the design objective function is weighted by $\\lambda$. I am curious how this hyperparameter is estimated/configured, and how sensitive the optimization results are to the change in $\\lambda$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f25irp84mI", "reviewer": "Reviewer_obHR", "summary": "The paper presents a new approach to inverse design by optimizing over the energy function learned by a diffusion model combined with the target function instead of backpropagating through (surrogate) dynamics. The \"compositional\" comes from the fact that the energy functions/diffusion models are learned over overlapping slices across time and state space. Experiments are performed on fluid dynamics and n-body dynamics tasks.", "questions": "- Did you try training a single shared network across the overlapping chunks? I was kind of expecting something like this (maybe with different degrees of subsampling to give long and short range dynamics)", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Originality:\n\n- the paper adopts or re-invents various tricks I've seen across the literature (unrolling across time steps and jointly diffusing, using a diffusion model as a smoothed ODE effectively) but does so in a clever combination\n- novelty: I'm not aware of any similar work, although conditional policydiffusion or codefusion might come close, and adding noise to FNO etc. is standard practice\n- clarity: overall clear presentation, especially on hyperparameters (kudos!), some questions (see below)\n- significance: difficult to judge in a still rather niche topic, but I think the general idea (learning sliced energy models to perform inverse design on) has promise to have high impact", "weaknesses": "- maybe I missed it, but page 7, I don't think $M$ is ever defined. How exactly do you train $M$ beyond the range of timesteps in training?\n- I would question the compositionality of the method and call it a \"piecewise\" or \"mixture\" approach? Given that you simply partition the spaces required into overlapping pieces (unless I misunderstood something)\n- Were the numbers of parameters matched for the different baselines? Given that you a partitioned energy functions, there might be potential for unfairness here?\n\n---\nedit after rebuttal: all of the concerns and questions have been addressed (see discussion),  on top of an additional request for signficance testing being promised for the camera ready. hence raising my score from 6 to 8", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eGQp7RPlXv", "reviewer": "Reviewer_WWyF", "summary": "The paper introduces a novel approach to tackle inverse design problems commonly found in engineering fields by optimizing energy functions within a diffusion model, instead of using traditional optimization techniques over learned dynamic models. This method addresses the challenge of adversarial modes encountered in the optimization process, improving design performance. By employing a compositional design system, the paper illustrates the potential to merge multiple diffusion models representing different subsystems, enhancing the complexity and specificity of the design.  The introduced Compositional Inverse Design with Diffusion Models (CinDM) method is highlighted for its capability to address out-of-distribution and more intricate design inputs beyond the training data, demonstrating promising advancements in the field of inverse design.", "questions": "- How other neural inversion methods perform in the context of the proposed experiments?\n- How sensitive is this inversion method to hyperparameters and initialization?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- A novel approach in addressing an interesting problem in the field of neural inverse design.\n- The results seem to be well enough, showcasing the benefits of this novel approach. However, further evaluation with the state-of-the-art counterpart neural inversion methods is required.", "weaknesses": "The manuscript lacks discussion on some pivotal related works, notably the contributions by Ren et al. [1] and Ansari et al [2]. Ren et al. elucidated various neural inversion methodologies and assessed their efficacy and accuracy encompassing Neural Adjoint, Tandem, Invertible Neural Networks, among other techniques. Additionally, they proposed a regularization scheme to mitigate the occurrence of out-of-distribution solutions.\n\nOn the other hand, Ansari et al. put forth a method wherein uncertainty information is integrated during the neural network inversion process. They asserted a multitude of benefits for this tactic, such as avoiding out-of-distribution solutions and erroneous local minima, alongside diminishing the model's susceptibility to initialization conditions.\n\nThese inversion methods should be mentioned and where possible compared with the proposed approach. In the cases where a comparison is not possible, sufficient explanation is required.\n\n[1]  Ren, Simiao, Willie Padilla, and Jordan Malof. \"Benchmarking deep inverse models over time, and the neural-adjoint method.\" Advances in Neural Information Processing Systems 33 (2020): 38-48.\n\n[2] Ansari, Navid, et al. \"Autoinverse: Uncertainty aware inversion of neural networks.\" Advances in Neural Information Processing Systems 35 (2022): 8675-8686.\n\n- A proper discussion over the limitations is missing.\n- The code and dataset is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)