Leveraging Discrete Function Decomposability for Scientific Design
Authors: James C Bowden
Sergey Levine
Jennifer Listgarten
Venue: ICLR 2026 Poster (2026)
Keywords: scientific
protein
design
generative model
decomposability
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: In the era of AI-driven science and engineering, we often want to design discrete
objects (e.g., circuits, proteins, materials) in silico according to user-specified
properties (e.g., that a protein binds its target). Given a property predictive model,
in silico design typically involves training a generative model over the design
space (e.g., over the set of all length-L proteins) to concentrate on designs with the
desired properties. Distributional optimization, formalized as an estimation of distribution algorithm or as reinforcement learning policy optimization, maximizes
an objective function in expectation over samples. Optimizing a distribution over
discrete-valued designs is in general challenging due to the combinatorial nature
of the design space. However, many property predictors in scientific applications
are decomposable in the sense that they can be factorized over design variables in a
way that will prove useful. For example, the active site amino acids in a catalytic
protein may need to only loosely interact with the rest of the protein for maximal catalytic activity. Current distributional optimization algorithms are unable to
make use of such structure, which could dramatically improve the optimization.
Herein, we propose and demonstrate use of a new distributional optimization algorithm, Decomposition-Aware Distributional Optimization (DADO),
that can leverage any decomposability defined by a junction tree on the design
variables. At its core, DADO employs a factorized “search distribution”—a
learned generative model—for efficient navigation of the search space, and invokes graph message passing to coordinate optimization across all variables.

Published as a conference paper at ICLR 2026
LEVERAGING DISCRETE FUNCTION DECOMPOSABILITY FOR SCIENTIFIC DESIGN
James C. Booden', Sergey Levine', Jennifer Listgarten 1.2.3 [jcbowden, sergey.levine, jennl]@berkelle edu
ABSTRACT
In the era of Al-driven science and engineering, we often want to design discrete objects (e.g., circuits, proteins, materials) in silico according user-specified properties (e.g., that protein binds its target). Given property predictive model, in silico design typically involves training generative model over the design space (e.g., over the set of all length- proteins) to concentrate on designs with the desired properties. Distributional optimization, formalized as an estimation of dis- tribution algorithm or reinforcement learning policy optimization, maximizes an objective function in expectation over samples. Optimizing a distribution over discrete-valus designs is in general challenging due to the combinatorial nature of the design space. However, many property predictors in scientific applications are decomposable in the sense that they can be factorized over design variables in way that will prove useful. For example, the active site amino acids in a catalytic protein may need only loosely interact with the rest of the protein for maxi- mal catalytic activity. Current distributional optimization algorithms are unable to make use of such structure, which could dramatically improve the optimization. Herein, we propose and demonstrate use of a new distributional optimization algo- rithm, DECOMPOSITION- AWARE DISTRIBUTIONAL OPTIMIZATION (DADO). that can leverage any decomposability defined by junction tree the design variables. At its core, DADO employs factorized "search distribution" learned generative for efficient navigation of the search space, and in- vokes graph message passing to coordinate optimization across all variables.
DESIGN IN DISCRETE STATE SPACES
The integration of Al into scientific research has opened new avenues for property-driven, in sil- ico design of discrete objects- from molecular structures like proteins, to fabricated systems like circuits wheeee where computational methods guide design with user-specified properties. For example, may seek to design an amino acid sequence for protein so that the protein binds its target.
Given a property predictive model, f(x). the simplest version of in silico design entails enumerating all possible designs, x [21, .2....., x2] X' (e.g., g. all possible amino acid sequences of length L.), evaluating each one under the predictive model, & = f(x), and choosing the design with the highest Such a setup is complicated by two primary challenges. First, in most realistic problems, the de- sign space is too large to fully enumerate. Second, if (2x has parameters estimated from data, then it is most likely not accurate over the whole space. While many works address the second problem, there has been little recent development on the first; as such, we focus herein on performing effi- cient optimization in -igheniinnii discrete design spaces. The second problem can be handled with approaches complementary to ours (e.g., Brookes et al. (2019); Trabucco et al. (2021); Uehara et al. (2024)), and is not here considered. We will focus our examples and experiments on designing amino acid sequences, but our method is general-
to design any discrete space.
Finding the highest-scoring design, x", naively requires D4 evaluations of f()) for D-amino acid alphabet (typically D 20). To develop some intuition, in this section only, let us assume that D 2 so that x is a binary sequence. First, consider the simple but unrealistic case where f() BI 3z ..... 3L L. for scalar parameters {8.). In such scenario. finding the design
'Department of Electrical Engineering and Computer Sciences, UC Berkeley. *Center for Computational Biology, UC Berkeley. *UC Berkeley-UCSt Graduate Program in Binengineering, UC San Francisco.

---

Published as a conference paper at ICLR 2026
with the highest f(x) requires considering only D L partial designs because the L components of z do not interact with each other, acting only linearly additively. A more realistic setting would allow for more complicated functional forms while maintaining some notion of linear additivity. For example, consider the form f(x) Ci (&1) + C2(52),.... Cr (āx), where Cj denotes an arbitrary function set of design variables, 32, such as (1 Ci ( In the case where sets of variables do not overlap, i.e., where can appear in only one component Cj. finding the global optimum of f(x) requires a number of evaluations of component functions that scales as Ddd, where M L is the cardinality of the largest variable set. When the sets of variables do overlap across component functions. tying them together, then the number of required evaluations becomes correspondingly higher. Note that this more general form of linear additivity admits representation of any function, possibly requiring that M = L in which case there is no linear additivity, nor consequently, decomposability. In realistic settings with many design variables and a larger alphabet (e.g., D 20), the difference in number of evaluations required between the decomposed and standard scenarios is even greater still. Importantly, in most real problems, we expect some level of decomposability. For example, in designing a protein to bind to a target, few key sequence positions may make up the binding interface, which primarily dictates the binding strength, whereas the other positions may be involved in stabilizing the protein for binding.
exp(7x
Classical message passing algorithms can leverage the aforementioned general type of structure in f(x) to exactly find the global optimizer with the lowest possible time complexity (Vlassis et al., 2004 However, for reasons discussed momentarily, we are interested in distributional optimization, wherein standard optimization problem over the design space, x' arg max fxx), is replaced by one over the parameters of a generative model, po (2), namely, 0" - arg maxg Epp(a) [f))] These formulations are equivalent in that the optimum of each is the same, assuming that po(x) has the capacity place all of its mass x². However, each formulation lends itself to different algorithms and extensions, mentioned below. Distributional optimization may employ strategies to prevent pa (a(x) the search distribution/poliey--ri prior (Brookes et al. 2019) or entropy regularizer (Ziebart et al. 2008).
collapsing to point mass, such as by using a
Our interest in the distributional optimization formulation is motivated by its extensibility. First, such a setup enables us to directly use innovations from the Estimation of Distribution Algorithm (EDA) (Brookes et al., 2020; Larrañaga & Lozanoo (2001) and policy optimization (Peters & Schaal, 2007 Peng et al.) 2019) communities. Of particular note are methods that enable combining a pre-trained, unconditional generative model, p(z), with a property predictor, p(ylæ), to execute Bayes rule and thereby obtain sampling distribution that can be used for design, p(rly E Y) (e.g.. Brookes (e.g..Brookes et al. (2019); [aanjiaag & Listgarten (2020); [eehara et al. (2024)). Second, as the key object that navigates the search space, pe (z), is generative model, we stand to benefit from advances in generative modeling. Herein, for clarity of contribution, we focus on the purest form of the EDA, without entropy regularization
prior. We leave such extensions to future work.
Contributions. We develop a distributional optimization algorithm in the form of a generalized EDA/policy optimization algorithm that can leverage any decomposability in f(x) in the form de- scribed earlier as f(z) (āt) + C2 ....).... (f..).)) ). to more efficiently navigate the design space and find high segh-performin designs quickly (Fig. [.. We call our method DECOMPOSITION- AWARE DISTRIBUTIONAL OPTIMIZATION (DADO). We first empirically investigate DADO on synthetic examples, illustrating that the anticipated optimization efficiency emerges compared to decomposition-unawaee
baselines. Then we further substantiate this efficiency on problems an- chored on real protein data, that is, by optimizing protein property predictive models. Additionally, in case studies on a few protein predictive models, we find that accuracy robust to modifications of the decomposition, suggesting that perfect priori knowledge of decomposability is not required.
As DADO requires decomposed objective of the form f(z) (ā + C2(#2). ... Cxx))), on our protein data problems, we construct these as follows. We first obtain a graph topology of which variables (protein residues) are coupled by thresholding residue distances from an AlphaFold3 structure (Abramson et al.] 2024 Brookes et al.) (022). This graph dictates the functional form of the predictive model, by way of an automatically constructed junction tree. Recall that junction trees represent arbitrarily complex relationships between random variables by transforming any graph into a tree structure, where each node contains a set of the original variables, and where these sets may be overlapping in adjacent nodes (Lauritzen & Spiegelhalter] 1988). Then we fit neural network that enforces this decomposition, to the training data. Analogous procedures in different application areas could include the following, for example. In circuit design, certain topologies of components
2

---

Published as a conference paper at ICLR 2026
b
More optimal
Less optimal
2))
x2)
Naive EDA
f(x)
pe(az
faalee R4)
DADO
fa2(sr m2)
faalsa. m4)
nd
L= 50
Qo(2n)
po(2o)
50
EDA DADO
40
(1 200)
Po(zi 20)
30
Q2(22,
po(#2 mo)
20
10
Qa 222
Po(za m2)
Qa(21, 22)
pa(z4 x2)
Training Iteration
Figure 1: Key components of DADO. DADO requires as input an objective function in its decomposed C2(2)),..... C2 tree with nodes size 1, i.e., regular tree, simplicity. Variables with edges interact to directly influence f. Some variables participate in multiple component functions, requiring coordination. b, To update the search distribution at cach iteration, naive EDAs weight entire samples, drawn from a joint distrihution over all design variables, with fx)) In contrast, DADO leverages the decomposition of fx: separately weight sample dimensions for each node. DADO uses cessage-passingg æ; interacting with its descendants. These Q; serve weights for each part the search distribution, which factorized like f. c, Example performance comparison on a synthetic problem with an exact tree decomposition over a discrete design space of size 20 (D 20, L 50), Each method drew 100 samples per iteration. We evaluated these samples with fz)), computing the per-iteration mcan and 95% confidence interval. Results shown were averaged over random seeds for the same f (details in Sec. 4 The p-value is from two-sided paired t-test that the mean at the final iteration different between methods, over the 20 seeds.
form, f(x) (21)
C. (R.). which corresponds junction tree. Here show junction
compute value functions, Q, accounting for each
may be prescribed ahead of time according production constraints domain knowledge. When designing a telescope, designers consider different arrangements of lenses-i.e. iie. topologies-and optimize the diameter, curvature, material, coating, etc. for each component.
DECOMPOSITION- -ANOOISSTTOO-AARRE DISCRETE OPTIMIZATION (DADO)
We begin by reminding the reader of the three primary steps that are iterated in standard EDA, which is unaware of any dccomposability C.g..,roooke et al. (2020)). From there, we describe how infusing this algorithm with awareness of the known decomposition of f (xz stands to make the EDA more statistically efficient, after which we formally introduce DADO.
Intuitively, one can think of the search distribution tively gets moved toward areas with high f(x) ) in expectation. Modern EDAs parameterize the search model with neural network generative models, such as Variational Autoencoder (e.g., Brookes et al. (2019). In a standard EDA, after initializing the search distribution, P&o (2r), the EDA pro- ceeds by iterating through these three steps either N times until convergence (Sec. A.2
spotlight on the design space, which itera-
1. Draw K samples from the current search distribution, {a*)
TR=1 ~ pa- (x)
2. Score each sample with the objective function to get its weight, " f(**).
3. Update search distribution parameters by weighted maximum likelihood estimation (MLE), using the weighted samples, gr+1 = arg maxç E w log log (x*).
The weighted MLE problem is typically solved with gradient descent; choosing to only fixed number of gradient steps, as we will do, can be theoretically justified through the equivalence of EDAs to -ixitttioonaaaiizaaiiiin (Brookes al. 2020
3

---

Published as a conference paper at ICLR 2026
Now that standard EDA is fresh in our minds, we can consider what it might mean to infuse it with knowledge of how f(z) is decomposed, and how this might prove useful. Before discussing the general case, let us first build intuition with simpler example. Specifically, consider f(x) = where there is no overlap in variables, Ti, between the meta-variables 21 and $2. In this setting, we can replace step 1, of sampling from one search model over all the design variables, = [+1 [2] = [ [I,, T2,..., ZL] ~ Pan (I), with instead sampling each meta-variable from its own search model separately, of ~ poy (f1) and a ~ Pag (ā2). Having done so, we can execute steps 2 and 3 in similar manner by scoring each meta-variable separately with its component function (e.g., / = (&f)) and then updating each search model independently of the other, via weighted MLE, to obtain Pop+1 (2)) and Poy+= (i2) (see Sec. A.3 for side-by-side algorithm sketches).
C1
What has this bought us? We have split one optimization problem over combinatorial space into two independent optimization problems, each over a much smaller space. For a problem where each meta-variable comprises just 15 binary design variables, we have transformed
optimization over space of size 200 s 10 to two problems each of size 215 ~ 10 For larger alphabets and sequence lengths, these differences will be larger still. Importantly, the weighted MLEE update with fixed number of samples, K, is more statistically efficient for each it is for dccomposition-unatl unaawre EDA using a single joint search model over all design variables.
search model than
The general case allows overlapping variables across component functions, requiring coordination between meta-variabl variables -e cannot divide and conquer as above; rather, we must divide, collabo- rate, and then conquer. It will be necessary to equivalently describe the decomposability of f(x) with a junction tree, defined by nodes corresponding to meta- vetialaaissless and undirected edges con- necting them (details in Sec. 1 and Sec. A.I). Given such a junction tree decomposition, we use factorized search distribution Each factor specifies the density of node's variables conditioned
directed, acyelic graphical (DAG) model matching its topology.
its parent node's variables, and will be parameterized by a neural network such as an autoregressive model. Sampling from such distribution is both computationally efficient and easy, so step 1 of the EDA remains straight- forward. However, steps 2 and 3 of the standard EDA are not so easily generalized. To update the factorized search distribution in manner that reduces the effective size of the optimization problem, we will need to generalize "max-plus" message passing, which yields a global optimum xx = arg max f(xx) (Vlassis al. 20044), to procedure integrating message passing with a fac- torized search model fitting step, thus constituting a decomposition-aweer EDA. Having provided an intuitive overview of DADO (Fig. [[. we now proceed to more formal exposition.
2.1 FORMAL EXPOSITION OF DADO
The goal of DADO is obtain a generative model, po(x), that maximizes the EDA objec- tive, arg max, Eyo (a)) [x)], while leveraging decomposabilityy in f(x) for optimization efficiency. DADO requires as input a decomposed version of ((x), which can be described by an undirected junction tree, T := (W. with nodes, NN, and edges, E. As noted in the introduction and shown in our experiments on proteins, identifying useful decomposability is feasible in practice. Given the junction tree topology, we write f(x) = ΣiEN fāā.) + E(i)Es ft, žz). where we refer to fz())) and fry (ã, 2;) respectively as node and edge component functions (Fig. [b); these are intimately related to the "epistatic landscape" of a protein property function (Sec. A. When the component functions are not known a priori, they can be parameterized and fit to labeled data.
We will begin by recalling how to do decomposition that is, to solve arg max, (x)) This problem is efficiently solved with a classical message- pesssage-passing algorithm, which coordinates local optimizations across parts of the junction tree obtain a single global optimum. Its efficiency comes from breaking optimization over all variables jointly into separate optimizations for each (smaller) meta- meta-variable, Having loaded the reader with this intuition, we then adapt these ideas to distributional optimization, yielding DADO.
exact (non-distrbbutional) optimization,
2.1.1 CLASSICAL MESSAGE-PASSEEE FOR NON- DISTRIIBTTIINA OPTIMIZATION
Although classical message message-passingg has been largely used for probabilistic inference on probabilistic graphical models (Pearl 1988, Shah, 2014), it can also be used for exact optimization of function, x² = arg max, fzx) (Vasssi et al. 2004). Given a function defined on an undirected junction tree, T. message-passing makes use of its topology find global maximum. In particular, one roots
4

---

Published as a conference paper at ICLR 2026
the undirected junction tree to obtain directed tree, inducing a hierarchy among the meta- variables from root to leaves. Each node is responsible for accumulating information from all nodes in its sub-tree and then passing this information to its parent. Consequently, the root node receives information from the entire tree, which is sufficient to set its variables in a globally optimal manner. Then, starting with the root, each parent communicates its variables' optimal settings to its children, which can in turn set their variables optimally, and so forth.
To obtain the rooted tree, T' :- (N'' ,'", from the junction tree, one keeps the same nodes, .'' N,. chooses a root node, r, and directs all edges in & outward from r. yielding directed edges, Er. Although rooting at any node will suffice, we choose r such that T' has the shortest height possible. Message-passingg finds a global optimum in two passes through the tree: one round of passing messages up from leaves to root and one round passing messages back down to the leaves.
Given T. classical message- passing first uses dynamic programming to accumulate information from the leaves up to the root. Similarly to any dynamic programming procedure, we accumulate solutions to increasingly larger intermediate sub-problems. In this case, each sub-problem is to find the value of (xx) evaluated on only a subset of meta-variables, rather than on the full set of variables in æ. Each sub-problem is tractable owing to the decomposition of the objective function into compo- nent functions for each node and edge, and by respecting the partial order of peo-llrobles induced by T'. Specifically, one computes value function. V. (āp(s)) ). for each node i E AA \\ {r}. which tells us for each setting of its parent, Fp(t). the value of the intermediate objective function defined by the edge component function, f zz). plus all component functions over the sub-tree rooted at given that all nodes maximize their respective intermediate objectives. Computing value functions constitutes the first pass through the tree, from leaves to root:
(Fp(s) = maxe, (Λ.( + (p(o), ..)+ + >
V (7:)).
eEchildren())
Notably, V.man (#p(z)) provides sufficient information about all nodes in the sub-tree rooted at to optimally choose the value of εP()) with respect to its children. Thus it follows that once all value functions have been computed, the root's assignment can be set in a globally optimal manner from its children's 's value functions.
arg maxx, (J~(E,) + Σ
eEchikdren(r) Vo((.))
Having chosen the root assignment, we then pass it down the tree as fp(a) = ( to its children, which successively pass their chosen assignments their children, all the way the leaves,
J arg max;, (I((A.) S0CO), 5.) + E eEchildrea())
Va(
resulting in a global maximizer x² of f=x) This dynamic programming "traceback" of optimal assignments back down the tree constitutes our second and final pass of messages.
Alternative notation. For convenience of our generalization to distributional optimization, we re-write the parent value functions (Tp((e) in terms of child-parent, Q (ā;, Tp(a)) and single- node, m.T (æ+) value functions:
V.Tax (āp(1) = max (Tt, #p(t),
where
(i., (p()) (H,
== ( + max (ā:)) and QTIX (ø:) = f((
Vmmax (.:).
eEchildren(s)
In contrast to the original parent value functions, Q (ăx, (p)) represents the effect of the choice of both āp()) and ū; on their edge component function plus all component functions over the sub- tree rooted at i, assuming all descendants of maximize their corresponding value functions. In- tuitively, Qu (((())
the value function on edge (p(i), prior ã, being maximized out, which will become useful if we want to, say, sample i; according to some distribution instead. Q ) is composed of two terms, one of which depends on its parent, and one of which doesn't, Q,aa( (z.). Written using the Q-functions just defined, the equivalent traceback equa- tions for selecting a globally optimal assignment, x², are simply i == arg max;, Qumax (ž,) and a == arg max;, Q.aa (ia, (p())) In other words.
x² = arg max, f(x)= = {arg max;, Q.ILK (i+)}} + farg max;, QUaa (ā., ((())) ) (o(t).))
(2)
arg max,
(QTee(T))
Q (ï,) + Σ (000) (ar,
(3)
5

---

Published as a conference paper at ICLR 2026
2.1.2 FROM CLASSICAL MESSAGE PASSING TO DISTRIBUTIONAL OPTIMIZATION
In the same way that an EDA transforms arg max, f() into a distributional optimization prob- lem, we can rewrite the equivalent message- message-passing objective in Eq. as DO. Because the original optimization problems over x are equivalent, their DO formulations are equivalent too,
arg max E [po(t)]((z)] = arg max 0
(p (5r))
Σ (p(1).))EE
Bpo(s) Ppott)(Q.t0 (Ti, (Bi #p(o))) (4)
8
where we've ve used linearity of expectations on the right side. However, a generic joint search dis- tribution cannot take advantage of the linear additivity in value functions over the junction tree topology. That is, while the classical traceback equations perform maximization over each meta- variable separately, Eq. 4 usss single, unfactorized search distribution over all variables, po pz(x), to optimize each Q-function. We address this next, by factorizing the search distribution.
Factorized search distribution. Classical message-passingg (Eq. 2 independently maximizes each Q-function conditional on the choice of Fp(s). instead of explicitly maximizing all variables in x jointly (Eq. (3). This is possible because each Q; captures all relevant global information needed to choose T.. It stands to reason then that DO can do something similar. Specifically, instead of training a single joint search distribution, we train smaller search distributions over each TG conditional on Tp(s). to separately maximize each Q. (ā., "p(3)). That is, we factor the search distribution accord- ing to T", resulting in DAG, po == po(ūr) ((pt).))e Po(ū; (p())), ). for root node r, non-root nodes i, and parents p(t), connected by directed edges E', and with parameters, 0. This factorized search distribution can be plugged into Eq. + for equivalent optimization problem. We refer to each element of this product as one of the factors of the search distribution. In our implementation, each factor has completely separate parameters, though write shared @ for conciseness. Notice that each factor of the search distribution interacts with each other factor through the directed edges, 8l That is, the distribution of ã. depends its parent's factor, and through it, all of its parent's ancestors: pe ( = Pa (ã. =pte) Pa (.)) In turn, each of node i's children's factors depends on (o.).. Due to this coupling, we cannot optimize each factor fully independently. But we can still update each factor separately from the others in dlooally-consistee
manner via message-passen passing. In particular, each factor will only be responsible for directly optimizing its own meta-variable, but will need to coordinate with its neighboring factors by getting information from them about how they are optimizing their meta- vetaavariables in a manner analogous to classical message-passingg Our current messages, Quess convey the value of each intermediate objective when all meta-variables are chosen via maximization. For DO, we'll require messages that communicate values when meta-variables are chosen according to the factorized search distribution, po (z).
Distributional value functions. While one certainly could choose to optimize classical value functions using an EDA search distribution (Eq. 4). it doesn't it make sense for two main reasons. As we just mentioned, DO aims train (1))) such that maximizes f(x), or equivalently, the sum of Q-functions, in expectation. Therefore, the DO objective should consider intermediate objective values for meta-variables chosen according to the current search distribution, not those chosen by explicit maximization, in V. (Fp(s)) Additionally, computing classical value functions requires enumerating all assignments of each ã; for the maximum used in V.TTa (ii(t)). When o, contains more than a few design variables, this max operation quickly becomes intractable. One reason for doing distributional optimization is to avoid enumerating massive design spaces. To address both issues, we define corresponding distributional value functions that fulfill both desiderata:
(Q/ V (Tp(t) = Eps (ā. (apcn) 1Q #p())); where
(ū,, #p(t)) =
G. + and (.,) = f(..)) + Σ @Echildren())
V2
Compared to Eq. [ the distributional V -functions compute the value in expectation under pe instead of max. Because they use an expectation instead of max operation, these value functions can be approximated tractably and without bias by drawing Monte-Carlo samples from the search distribu- tion. Moreover, each distributional value function lower bounds each corresponding classical value function because the expectation of function cannot exceed its maximum (details in Sec, A.44. As result, the sum of classical value functions in the objective in Eq. 4 is bounded below by the sum

---

Published as a conference paper at ICLR 2026
of distributional value functions, which we'll optimize instead:
(p (2+))] (2,)) + Σ (p())))e
12 Tp(i))l 2
(ā,, Epo (2, + E (p(0),)ee Pnota)QQ(( Spra)] (5)
Distributional optimization with value functions. All that remains is to derive an update rule that treats each search distribution factor separately. Since each expectand in Eq. doesn't depend on descendant meta-variabel we can replace po(x) with each meta-variable's marginal distribution:
arg max po, (Q.. + Σ (p(t),))ee 18 [Q] (G,, #p()))
We then follow the EDA derivation to arrive at an update rule (details Sec. A.4 in which each term is optimized by only single factor; dependence on po(p((s) is approximated by sampling. We write DADO's 's update rule (sharing a single set of samples; see Fig. [T) as a sum of weighted likelihoods for each search distribution factor,
gm+1_ arg max
Σ
(0. (x,, log (e())) + Σ
of (i,, (n()) log po @p(o)),
xx)pon(2)
(P()),))EE
which is equivalent separate updates because the factors don't share parameters, as desired:
0,+11 arg max
E Ex-pon (w) (itr) log po, (Fr), and
0.+11 arg max Σ pap (a) Of" (T Sp()) log po, (Tx Ep(s)), ViENA {r).
Each factor is weighted by its corresponding value function, enabling it to coordinate with all its descendant factors despite their being updated separately. The whole DO is tied together the top by the root factor. We emphasize that DADO's update is more statistically efficient than a naive EDA's because DADO gets use all K samples for weighted MLE on each lower-dimensional factor distribution (i.e., ratio of number of samples to number of dimensions is larger). Our resulting algorithm (Alg. 1 fits into the three EDA steps: (1) designs are sampled from pe- (z). (2) weights, here each Q-function instead of fI)) are computed, and (3), each search distribution factor receives its own independent weighted maximum likelihood update based on these weights (Fig. [D). These are repeated until convergence, or for some fixed number of iterations. The factor updates are coupled only through the Q-functions, the messages across edges. {Q²" Hew are only valid while 0 is close to 0", meaning one must balance how many gradient steps are taken before drawing new samples. If too many gradient steps are taken without resampling. a factor may be changing its parameters to collaborate with another factor which has already changed its behavior. It's common for EDAs to include an additional hyperparameter W (-).. monotonic shaping function applied to the weights, to alter optimization dynamics (Brookes et al. 2020), included in Alg. [ Choosing W be the identity recovers our derivation. Notice that in place of the classical message-passing traceback equations, DADO simply performs sequential conditional sampling from its search distribution. Interestingly, a very similar algorithm can be derived without using lower bound on the value functions (Fq. 5) if one adds an entropy-maxinn maximizing term to the initial DO objective (Sec. A.5).
RELATED WORK
The closest work to ours in the EDA community is that of the Factorized Distribution Algorithm (FDA), wherein a decomposition is leveraged to factorize the conditional probability table (CPT) that parameterizes the search distribution (Mühlenbein Mahnig, 1999). However, FDA cannot leverage message passing to give different weights to different variables so to more efficiently update the search model. In related work, PPlikann (2005) replaced the CPTs with Bayesian networks, but still had no means to coordinate among variables through message passing the like.
In a complementary line of work on policy optimization in reinforcement learning (analogous to learning the search distribution in EDAs), coordination between variables is enabled by message passing on a factorized search distributio/ppolicy. but this line of work is only suitable for graph topologies that are Markov chains that is, chain graphs. These graphs arise from reward functions,
7

---

Published as a conference paper at ICLR 2026
Algorithm 1 DECOMPOSITION DISTRIBUTIONAL OPTIMIZATION (DADO)
1: Given: junction tree T (V, E), component functions {fr, (8, By hy)oeee and (ff(Ee)
2: Root T such that tree height is minimized, yielding directed junction tree, T' (W.
3: Factorize search distribution according to T' as po( = po(är) [((p)))) Po(ăs I #p(1);
initialize 0
4: Sort nodes topologically: order ON such that 0 is leaf index and O is root index
5: for iteration n ,,2....,A do
6:
{a ~ po(x)
b Sample from the factorized search distribution > Estimate distributional value functions from samples
7:
for E ON do
8:
for k 1,2,.. K do (,4)) ((())) Eeechildeen()) Σ Qe 25)]
9:
10:
i's parent, p(i), exists:
11:
Q, ) t ( p()): 24 + Qi(
12:
end for end for
Update each search distribution factor using value functions as weights i N, in parallel do
14:
for
all
15:
i =:: if
16:
0, f arg maxe, Σ k= E ((Q, log po, (i+)
else
0, f arg maxe, Σxm ((Q J xb())) log (o+ T p(O))
18:
19: end for
20:
end for
fz)), that decompose in time according to a first-order Markov assumption (Peters & Schaal (007))
Peng et al. (2019). Nair et al. (2020)) Such approaches cannot be used on arbitrary junction trees.
Adjacent to our problem of interest, because they focus solely on real-valued design spaces, Grudzien et al. (2024) introduces a new way to discover functional decomposition from data, and then demonstrate that doing so helps improve out-of-distribution generalization, Separately, the Bayesian optimization (BO) community has developed methods for dynamically inferring func- tion decomposition from data during active learning. Although they make use of message passing for coordination, they do not employ distributional optimization, nor can these methods be readily generalized to do so Kandasamy et al. (2015), Han et al. (2021), HHang et al. (2018), Rolland et al. (2018), Bardoo al. (2024)). This community has shown that using approximate decompositions can be helpful ZZooeek & Ammar 2023). DADO could, in principle, be used within the BO inner loop, although such investigation is beyond the scope of the present work.
random
EXPERIMENTAL RESULTS
We perform two sets of experiments with increasing resemblance to real-world scientific design. In each setting we compare standard EDA, that is unaware of the function decomposition, to DADO, which aware of it. In the first setting, we create synthetic functions f(x), while in the second setting, we focus on functions derived from real protein data. In all experiments we designed sequences of fixed length L, where each position is one of D = 20 amino acids. We used N 100 EDA training iterations and G 1 gradient steps for each weighted maximum likelihood update. We use MLP-based search distributions (details in Sec. A.8.1 for DADO and all baselines. We also compare to FDA (Mühlenbein Mahnig 1999) and PPO (Schulman et al., 2017) (details in Sec. A.7). For the shaping function we choose (s)) exp 3 and tune the temperature, B. To pick P and the learning rate, n, we performed hyperparameter sweep separately for each method and each f(x) taking (n") with the largest sample mean at the last iteration. For synthetic experiments,
swept over 100 combinations between n7 E [10 5 5 10-2] and E [0.1, 8). For protein experiments, which take longer, we swept only 54 pairs between 1 e the boundary of the
[5 x 10-5 5 10-2] and 3 E [0.01, 5). If either method's 17" or 3r
8

---

Published as a conference paper at ICLR 2026
26
L 25
L 50
200
150
75
EDA
BA60
10
P 7: 12
P- 7,
20
Trainin
Traning
Traning
Figure 2: Comparison of a naive EDA to DADO on synthetic problems. created three random functions, f(r), each with a randomly chosen junction tree decomposition with maximum node size of one, and randomly chosen parameters. Each experiment used alphabet size D = 20 and sequence length, a, L = 25, b, L = 50, and L 200. Each of the two methods drew K = 100 samples per iteration. For each iteration, we show the mean (solid line) and 95% confidence interval (shaded envelope) of the 100 samples evaluated averaged across results from 20 random sceds. P-values are from two-sided paired t-test that the mean at the final iteration is different between DADO and the most competitive baseline, over the 20 sceds.
[(x),
sweep, we broadened it. (")","" were then used for all replicate runs over random seeds dictating the initial search distribution and sampling. Although not essential our method as described in Alg. [ in our implementation we reduce the variance of the search distribution update by using mean-shifted Q-functions, Qi, Xp) [s ( (a }p]], which are unbiased (Williams 1992). Statistical significance of differences in performance computed by comparing the sample mean fx) at the final iteration, using paired two-sided t-tests over random seeds.
COMPARISON ON FULLY SYNTHETIC FUNCTIONS
Here we used junction trees with meta-variable nodes containing only one original variable, i.e., an exact tree decomposition. We did so because this allowed us to better control the difficulty of the synthetic functions in that we could specify each component function as a dense lookup table of size D and thereby ensure that there would be non-smoothness all throughout the design space. In contrast, larger nodes would have increased lookup table sizes exponentially with the number of variables in each meta-variable, requiring more compact representation. Sparse lookup tables functions. Neural networks
are one such option but generally result in saall-----ahahaa---llllae another, but random initialization tends to yield overly smooth functions that unrcalistically easy to optimize. Nor is it obvious how to produce neural landscapes that are realistically difficult across the entire design space via training. We simulated one fx) for each of three sequence lengths, = {25, 50, 200}. by randomly sampling a junction tree and component functions (details in Sec. A.8.2 Search model weights were initialized from (0, 0.0004) and biases were
All methods drew K = 100 samples from their search distribution at each iteration. For all three sequence lengths, DADO outperforms all three baselines (Fig. 2 When T, 25, the baselines catch up to DADO by the end of 100 iterations, but for larger L, corresponding to larger design spaces, the baseline methods converge to substantially lower values of fz) than DADO. Results from experiments in which we increase D further support DADO's superiority (Fig. A3
COMPARISON ON PROTEIN PROPERTY PREDICTORS
Here we anchor our experiments on f(x) fit to real protein datasets, We use four protein prop- erty datasets, AAV, Amyloid, Gcn4, and TDP-43 (details in App. (A.)). For each dataset, we used AlphaFoldd--rediicte structures (Abramson al. 2024) on the wild-type sequence to obtain 3D structure, from which we constructed a contact graph by thresholding the distance between pairs of residues with threshold t. Following Brookes et al. ( Romero al.) (2013): Voigt et al. (2002), we use a threshold of /. = 4.5A We interpret this contact map as a graph adjacency matrix, from which we algorithmically construct a junction tree (Lauritzen & Spiegelhalter 1988). This defines the topology needed for DADO (Fig. AI). and then we fit the component functions the protein assay-labeled data (details in Sec, A.8.3). We initialized the search distribution by training on the 1,000 sequences with lowest (x) for 1,000 iterations, with learning rate of 2 x 10
All methods drew K = 1000 samples from their search distribution at each iteration. For all four proteins, DADO converges to designs with higher fx) than the baselines both visually and by statistical test (Fig. 3). For three of the four (Amyloid, Gen4, and TDP-43), the sved-averaged 95% confidence intervals of DADO and the best baseline the final iteration do not overlap. We plotted

---

Published as a conference paper at ICLR 2026
AAV2 capsid protein (L=28, D=20)
Amyloid beta protein (L=42, D=21)
9.42e
08
P 1.49e 05
FDA
-log(c-((X))
PPO
-log(c-((x))
0
Training lteration
Training Iteration
Gcn4 protein (L=44, D=20)
TDP-43 protein (L=84, D=21)
d
= 2.53e
P 3.42e 06
8
-log(c-f(0)
-log(c-f(X))
2
20
40
60
80 100
20
40
60
80 100
Training lteration
Training Iteration
Figure Comparison of naive EDA DADO protein property predietive models. For each of four proteins of varying length, a, AAV, b, Amyloid, c, Gcn4, and TDP-43, we neural network property function, [x), adhering to junction tree decomposition derived from the protein's 's 3D structure, and then used naive EDA and DADO to optimize them. Each approach drew K = 1000 samples per EDA iteration. (solid line) and 95% confidence interval (shaded envelope) of the 1000
For each iteration, show the samples evaluated on log(c fz)). averaged across results from random seeds. We plot this quantity make clear the differences between methods when f(x)) is large; c is the largest f(x on given plot, plus a small constant for numerical stability. P-values are from two-sided paired t-test that the mean at the final iteration is different between DADO and the most competitive baseline, over the 20 sceds.
loog(c-f(r)) f(x)) on the y-axis to make this difference clear visually; plots with f(z) on the y-axis included in the appendix (Fig. A5)). Finally, Sec. (e.33. we observe that our choice of AlphaFold3- based topology is robust mild modifications far as holdout performance is concerned.
DISCUSSION
We have proposed new method for distributional optimization that can leverage arbitrary decom- posability the function being optimized. We have shown that it works as expected on synthetic problems-namely. better than naive EDA that is not aware of the decomposition. We have also demonstrated the potential for practical utility on the problem protein design. Importantly, it is necessary that a problem strictly adhere the specified decomposition in order be useful. Specifically, we showed that using the heuristic of thresholding AlphaFold3 estimated contacts, we can obtain a range of decomposed functions, including some that maintain predictive accuracy while providing level of decomposability helpful for DADO. Obtaining similarly useful decomposability other domains will require further investigation; however, as the real world typically is structured, we expect many areas will be amenable to doing so. Estimating decomposability from labeled data an active area of research (Poelwijk et al. (2016); Grudzien al.) (2024): Park et al. (2024)).
It will be insightful to conduct future work interrogating the use of DADO with generalizations of the standard EDA, such as conditioning an unconditional generative model with an inaccurate property predictor for Al-guided design, which should enable us to design while accounting for the fact that f() is not truly known everywhere in the space due to finite data. This particular problem unconditional model into that respects the decomposition, although
may require distilling alternative strategies could prove useful. Finally, use of DADO within Bayesian Optimization for optimization of the acquisition function is exciting direction.
10

---

Published as a conference paper at ICLR 2026
6 REPRODUCIBILITY STATEMENT
We specify exhaustively experimental setup, hyperparaaeterrs model architectures, and algo- rithm details Sec. 4 Datasets used and their sources are detailed in Sec. A.6 Code for our method and experiments is available publicly GitHub: https: /ttp:://githuln wwwpt.ttt.oooooooo...oooooootttt http:///hhha...c////mmeeoorooooo Jamesm/bwweeebddd...... DADO
ACKNOWLEDGMENTS
We thank Kuba Grudzien for helpful discussions regarding function decomposition, Hanlun Jiang and Peter Ma for discussions regarding scientific applications. Toby Kreiman and Bear Xiong for feedback on the manuscript, and Aileen Zhang for feedback on the figures.
JCB was supported in part by an NSF GRFP fellowship and UC Berkeley Chancellor's fellowship. This material is based upon work supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE 2146752. Any opinions, findings, and conclusions or recommendations expressed in this material are those the author(s) and do not necessarily reflect the views of the National Science Foundation. JCB and JL were supported in part by the Office of Naval Research (ONR) under grant N00014-23-1-2587, and the U.S. Department of Energy, Office of Science, Office of Biological and Environmental Research, Lawrence Livermore National Laboratory BioSecure SFA within the Secure Biosystems Design program (SCW1710)
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493-000, 2024.
Anthony Bardou, Patrick Thiran, and Thomas Begin. Relaxing the additivity constraints in decen- tralized no-regret high -igh-dimensional bayesian optimization. In The Twelfth International Confer- ence on Learning Representations, 2024.
Benedetta Bolognesi, Andre J Faure, Mireia Seuma, Jör M Schmiedel, Gian Gaetano Tartaglia, and Ben Lehner. The mutational landscape of prion-like domain. Nature communications, 10(1):
4162, 2019.
David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning. volume 97 of Proceedings of Machine Learn- ing Research, pp. 773-782, PMLR, 09-15 Jun 2019. URL hatps: https::/ppocceedggg.s.. mml.
9a ht.ml
David Brookes, Akosua Busia, Clara Fannjiang. Kevin Murphy, and Jennifer Listgarten. A view of estimation of distribution algorithms through the lens of pnpeettaioo-maxiiizatiin In Proceed- ings of the 2020 Genetic and Evolutionary Computation Conference Companion, pp. 189-190, 2020.
David H Brookes and Jennifer Listgarten. urXiv:1810.033 2018.
Design by adaptive sampling. arXiv preprint
David H. Brookes, Amirali Aghazadeh, and Jennifer Listgarten. On the sparsity of fitness func- tions and implications for learning. Proceedings of the National Academy of Sciences. 119(1); 12109649118, 2022. doi: 10. 210994991 URL htt.h:tp///// http://wwwwwpww....o.c// pnas www.:pp.aa.o.ooo....... org/doi/ abs/10.1073/pmaa 21t.mmass21999911111111
Drew H Bryant, Ali Bashir, Sam Sinai, Nina K Jain, Pierce J Ogden, Patrick F Riley, George M Church, Lucy Colwell, and Eric D Kelsic. Deep diversification of an aav capsid protein by machine learning. Nature Biotechnology, 39(6):691-696, 2021.

---

Published as a conference paper at ICLR 2026
Yongcan Chen, Ruyun Hu, Keyi Li, Yating Zhang, Lihao Fu, Jianzhi Zhang, and Tong Si. Deep mutational scanning of an oxygen-ineee -igee-inddeendeent fluorescent protein creilov for comprehensive pro- filing of mutational and epistatic effects. ACS Synthetic Biology, 1265):1461-147,, 2023.
DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Clau- dio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King. Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http:
7/github. com/ google deee-deeepmnnn
Clara Fannjiang and Jennifer Listgarten. Autofocused oracles for model-based design. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 12945-12956. Curran Associates, Inc.,
2020. URL htttp:://proceeeets wew.eeeedddnn..eeeeeeeeeeeeeeeee neuri cttppppcccpppppppeeeeeeeeeeeeeee http://aaereeeeeeeeeeeeeeeeae wwtps/flllesppape/222. file flc22722caae22222560007777772223
Kuba Grudzien, Masatoshi Uchara, Sergey Levine, and Pieter Abbeel. Functional graphical models: Structure enables offline data-driven optimization. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li (eds.), Proceedings of The 27th International Conference on Artifi- cial Intelligence and Statistics, volume 238 of Proceedings Machine Learning Research, pp- 2449-2457, PMLR, 02-04 May 2024. URL https: htt.:tttss:/ppoooo htttts://prooeediggs.ooo mlr press/v238/
grudzien24a.ht.c html
Eric Han, Ishank Arora, and Jonathan Scarlett. High-dimensional bayesian optimization via tree- structured additive models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 7630-7638, 2021.
Trong Nghia Hoang, Quang Minh Hoang, Ruofei Ouyang, and Kian Hsiang Low. Decentralized high-dimensional bayesian optimization with factor graphs In Proceedings of the AAAI Confer- ence on Artificial Intelligence, volume 32, 2018.
Edwin Jaynes. Information theory and statistical mechanics. Physical review, 106(4):620, 1957.
Kirthevasan Kandasamy, Jeff Schneider, and Barnabas Poczos. High dimensional bayesian op- timisation and bandits via additive models. In Francis Bach and David Blei (eds.), Proceed- ings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 295-304, Lille, France, 07 Jul 2015. PMLR. URL. https: http::///oceeeiieg.e malr.press v37/kandasamyl15 5.htm_
Diederik P Kingma and Max Welling. Auto-encoding variational bayes.
arXiv preprint
arXiv:1312.6114, 2013.
Pedro Larrañaga and Jose A Lozano. Estimation of distribution algorithms: A new tool for evolu- fionary computation, volume 2. Springer Science & Business Media, 2001.
Steffen I. Lauritzen and David Spiegelhalter. Local computations with probabilities on graphical structures and their application to expert systems. Journal of the Royal Statistical Society: Series B (Methodological), 1572):157-194, 1988.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: 08059 2018
Francesca-Zhoufan Li, Jason Yang. Kadina E Johnston, Emre Gürsoy, Yisong Yue, and Frances H -aariing-assitted
Amold. Evaluation of machine learning-ast assisted directed evolution across diverse combinatorial landscapes. bioRxiv, pp. 2024-10, 2024.
Rosalie Lipsh-Sokolis and Sarel J Fleishman. Addressing epistasis in the design of protein function. Proceedings of the National Academy of Sciences, 121 (3933)))333449991 2024.
12

---

Published as a conference paper at ICLR 2026
Ily Loshchilov and Frank Hutter. Decoupled weight decay regularization. a 2017.
arXiv preprint
Heinz Mühhennbei and Thilo Mahnig. Fda scalable evolutionary algorithm for the optimization of additively decomposed functions. Evolutionary Computation, 7(4):353-376 1999. doi: 10.
(166/evco.1999.7.4 4.3.333 httthttt/mmmhhhete
org/ Eda99.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online rein- forcement learning with offline datasets. arXiv preprint arXiv:2006.09359, 2020.
Pascal Notin, Aaron Kollasch, Daniel Ritter, Lood Van Niekerk, Steffanie Paul, Han Spinner, Nathan Rollins, Ada Shaw, Rose Orenbuch, Ruben Weitzman, al. Proteingym: Large-scale benchmarks for protein fitness prediction and design. Advances in Neural Information Processing Systems, 36: 64331-64379, 2023.
Anders Olson, Nicholas Wu, and Ren Sun. A comprehensive biophysical description of pairwise epistasis throughout entire protein domain. Current biology, 24(22):243326651 2014.
Jakub Otwinowski and lya Nemenman. Genotype to phenotype mapping and the fitness landscape of the e. coli lac promoter. PloS one, 8(5):e61570, 2013.
Yeonwoo Park, Brian PH Metzger, and Joseph W Thortoo.. The simplicity of protein sequence- function relationships. Nature Communications, 15(1):7953, 2024.
Judea Pearl. Probabilistic reasoning in intelligent systems: Networks of plausible inference, 1988.
Martin Pelikan. Bayesian optimization algorithm. In Hierarchical Bayesian optimization algorithm: foward a new generation of evolutionary algorithms, pp. 31 31-48. Springer, 2005.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine Advantage-weiigted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv: 191vv01910.00077 2019.
Jan Peters and Stefan Schaal. Reinforcement learning by reward-weighted regression for opera- tional space control. In Proceedings of the 24th International Conference on Machine Learn- ing, ICML '07. pp. 745-750, New York, NY, USA, 2007. Association for Computing Machin- cry. ISBN 9781595937933. doi: (1145/1127346.12755555 URL http::::/dii.ogg// 1145/127345 1273590
Frank J Poelwijk, Vinod Krishna, and Rama Ranganathan. The context-dependence of mutations: linkage of formalisms. PLoS computational biology, 1004771,,, 2016.
Paul Rolland, Jonathan Scarlett, Ilija Bogunovic, and Volkan Cevher. High-dimensional bayesian optimization via additive models with overlapping groups. In Amos Storkey and Fernando Perez-Cruz (eds.), Proceedings of the Twenty-First International Conference on Artificial In- telligence and Statistics, volume
Proceedings of Machine Learning Research, pp. 298-307, PMLR, 09-11 Apr 2018. URL httph::t/sstg htttp:://proceedingg. ml mlt.press/vv87 rol land18a. html
Philip A Romero, Andreas Krause, and Frances H Arnold. Navigating the protein fitness landscape with gaussian processes. Proceedings of the National Academy of Sciences, E033)E113--20,
2013.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889.1897. PMLR, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.063477 2017.
Mireia Seuma, Andre J Faure, Marta Badia, Ben Lehner, and Benedetta Bolognesi. The genetic landscape for amyloid beta fibril nucleation accurately discriminates familial alzheimer's disease mutations. elife, 10:e63364, 2021.
13

---

Published as a conference paper at ICLR 2026
Devavrat Shah.
Algorithms for inference 2014.
(course 6.438. fall 2014). MIT
OpenCourse Ware,
URL
https: https:///ow..o/ hew.mmit.. www.miiteedu. courses/
6-438-aloo aal--aaggrrrhssss htw.hhmsso----teeeeeeeee inffrenneeeee- fall-2014/ 2025-10-30.
Lecture notes. Accessed:
Max V Staller, Alex S Holehouse, Devjanee Swain-Lenz, Rahul K Das, Rohit V Pappu, and Barak A Cohen. A high-throughput mutational scan of an intrinsically disordered acidic transcriptional activation domain. Cell systems, 6(4):444-455, 2018.
Tyler N Starr and Joseph W Thornton. Epistasis in protein evolution. Protein science, 25(7):1204- 1218, 2016.
Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
Ammar Tareen, Mahdi Kooshkbaghi, Anna Posfai, William T. Ireland, David M. McCandlish, and Justin B. Kinney. Mave-nn: learning genotype-phen phenotype maps from multiplex assays of variant effect. Genome Biology, 23:1-27, 12 2022. ISSN 1474760X. doi: 10.1 10.1186/
URL https: http///ttetggooooooooooooo/ http///gnoommooogggeeeoe bi omedcentral..
com/ http///rrclessl11..m 1ttp://.....1111111111 022-0026----2---21-1-
Brandon Trabucco, Aviral Kumar, Xinyang Geng, and Sergey Levine. Conservative objective mod- els for effective offline model- based optimization. In Marina Meila and Tong Zhang (eds.). Pro- ceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. (0358-10366) PMLR, 18-24 Jul 2021.
Kotaro Tsuboyama, Justas Dauparas, Jonathan Chen, Elodie Laine, Yasser Mohseni Behbahani, Jonathan Weinstein, Niall M Mangan, Sergey Ovchinnikov, and Gabriel J Rocklin. Mega-scale experimental analysis of protein folding stability in biology and design. Nature, 620(7973):434- 444, 2023.
Masatoshi Uchara, Yulai Zhao, Ehsan Hajiramezanali, Gabriele Scalia, Gokeen Eraslan, Avantika Lal, Sergey Levine, and Tommaso Biancalani. Bridging model-based optimization and genera- tive modeling via conservative fine-tuning of diffusion models. Advances in Neural Information Processing Systems, 37:127511-12755 2024.
N. Vlassis, R. Elhorst, and J.R. Kok. Anytime algorithms for multiagent decision making using coordination graphs. In 2004 IEEE International Conference on Systems, Man and Cyhernetics (IEEE Cat. No.04CH37583), volume 1, pp. 953-957 vol.1, 2004. doi: 10. 10.1109/05M..22004 1398426.
Christopher A Voigt, Carlos Martinez, Zhen-Gang Wang, Stephen L Mayo, and Frances H Arnold. Protein building blocks preserved by recombination. Nature structural biology, 9(7):553-558 2002.
Ronald J. Williams. Simple statistical gradient-followin algorithms for connectionist reinforce- ment learning- Mach. Learn., 8(3 8(3-4))229 256, May 1992. ISSN 0885-6125. doi: 10.1007/ BF00992696. URL
Fa- Yueh Wu. The potts model. Reviews of modern physics, 54(1):235, 1982.
Danqing Zhu, David H Brookes, Akosua Busia, Ana Carneiro, Clara Fannjiang, Galina Popova, David Shin, Kevin C Donohue, Li Lin, Zachary M Miller, et al. Optimal trade-off control in machine learning-basse based library design, with application to adeno-associaae virus (aav) for gene therapy. Science Advances, 10(4)read33786, 2024.
Brian Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 1433-1438. Chicago, II., USA, 2008
Juliusz Krzysztof Ziomek and Haitham Bou Ammar. Are random decompositions all we need in high dimensional bayesian optimisation? In International Conference Machine Learning, pp- 43347-43368. PMLR, 2023.
14

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 DIFFERENT WAYS OF WRITING FUNCTION DECOMPOSITION
In the introduction (Sec. D. we define a decomposable function as one that can be written like f(x) = C, (.1) - C2(i2)...... C with meta-variables '; that are generic sets of original vari- ables z. This formulation is most helpful for building intuition about delyy-decopposabee functions, in which there are no variables z, that appear in multiple meta-variables (i.e., meta variables don't t overlap). It also encompasses functions for which there are overlapping meta-variables in the de- composition too.
f(x) = C(()) - C2(52), (xxx also provides the helpful intuition that someone looking to use DADO need only specify a decomposition at this level of detail-sets of design variables that directly interact-rather than as a graph. In the protein binding example given in Sec. [ user might specify three component functions: one over sequence positions in the binding interface, To- one over sequence positions in the scaffold, ê1, and one over a subset of positions tying together the binding interface and the scaffold, 22. Specifying decomposition at this level may be easier than in its graph form. In the event that one has an exact decomposition given, this information can be represented either this way, or in terms of the graph introduced below. Situations in which an exact decomposition may be given include hardware design (e.g., due to manufacturing constraints, topology of wires in a circuit is pre-specified) and bi-level optimizations (e.g., to design the best telescope, a scientist optimizes both the arrangement of lenses and their physical properties; an outer loop optimizes topologies and an inner loop using DADO evaluates single topology by optimizing each lens' parameters).
However, f(x)= (i + C2(2)) C.(ft)) written has interpretation in terms of graph, which is necessary for message-passingg and for factorizing the search distribution (Sec. 2. To bridge this gap, we introduce graph in which nodes design variables and edges specify which design variables directly interact to influence f. Users of DADO may also specify the decomposi- tion this level; we expect this to be helpful when domain knowledge is available in the negative form, i.e., when the user can begin with the fully-commn cully-connected graph and remove edges between design variables they think don't directly interact. This graph can then be automatically transformed into a junction tree (Lauritzen & Spiegelhalter 1988), which is what our algorithm actually takes as input. That is, DADO takes an undirected junetion tree, T == (W,E), with nodes, NN, and edges, E. Each node is a set of design variable indices, such that indexing x with this set yields new meta-variable, I; correspondingly exactly to a node in the junction tree, whereas the old meta-variables 2; had no such interpretation. f(x) must be decomposed instead according to T.
We can rewrite the generic function decomposition, f(x) = C(t))) + C2(T2), ...,, Cn (.x)) equiva- lently in terms of N and E: f(x) ExeN fi():) + +uune fij 2,) (Fig. [b.. We' ve replaced component functions (.,)) on generic sets of design variable indices with either an equivalent "node component function", fi (ā:), or "edge component function", (ā,, 22). Correspondingly, i; in the original formulation either corresponds exactly to some node in the junction tree, such that é;; ãz.. or it corresponds to some edge in the junction tree, such that d; ă u k- The fully-decomposed, no-overlap EDA described in Sec. corresponds to an edgeless junction tree- each component function is simply node component function, C, fi;)), and there no edge component functions- which is why a fully- factorifce search distribution suffices. But if, for example, in 21) and d1 = (21, 222 overlap, then they must instead correspond to edge component functions. For a simplified case where the junction tree nodes have cardinality 1 (e.g., á = I.2), might have Co(ā)) fo, (ão, and C₁ (.,) = fi.2 ( .t2). Fig. [ only depicts edge component functions for clarity and simplicity; notice also that node component functions can technically be subsumed into appropriate edge component functions without loss of generality. The presence of edge component functions requires the coupling of the search distribution's 's factors in a manner compatible with E; faly--factorized search distribution with independent distributions for each meta-variable will no longer suffice.
A.1.1 CONNECTION TO EPISTASIS
Function decomposition corresponds intimately to notions of the epistatic landscape for protein property function (e.g.. P.ellli et al. (2016). Wu (1982). Otwinowski & Nemenman (013 Lipsh-
15

---

Published as a conference paper at ICLR 2026
Sokolik & Fleishman (2024) ). An "epistatic expansion" of f,
1.
f(x== - 3+ Σ [ + ΣΣ
B Øy] + .... + Bo.1....-1 T1, TL-I], 1=0 f31
(6)
decomposition of into bias term, first-order terms, second-order terms, and so forth up L-order epistasis. For the first-order terms, each 3 is a vector which is indexed by the particular value of z;. For higher-order terms, is d-dimensional tensor. One can obtain a decomposition graph from an epistatic expansion as follows: for each nonzero term in the expansion, add an edge between each design variable in the term. A function in which only few, lower-order terms were nonzero would be relatively easy to optimize.
An expansion of this sort can be thought of as spectral decomposition of f into components of lower and higher frequencies, and can be computed for small landscapes using a discrete Fourier transform. function with primarily lower- foeuu-fryy components would be smooth and easy to optimize whereas function with higher-frequency components would be rugged and difficult to find the global optimum of. One can imagine a worst-case scenario, that L-order epistasis, of which a needle-in-a-haye a-hast-atacaakk function is example. In general, full epistatic landscapes are intractable to compute for design spaces practical sizes, but estimation of subset of the terms is an active area of research (e.g.. Otwinowski & Nemenman (2013); Park et al. (2024)) due to beliefs that natural proteins tend to exhibit sparse and mostly lower-order epistasis.
In comparison the function decompositions used by DADO, full epistatic landscape specifies both the graph topology and the component functions, everywhere. Completely specifying an order- d interaction if each variable has D states requires choosing Ddd values. The high-dimensional synthetic functions used in our experiments are fully specified everywhere because they 're defined with primarily lower-order epistasis, which can be specified with far fewer values. In contrast, our protein property predictors can have high-orde order component functions represented by neural networks fit from limited data, meaning they' 'r underspecified beyond the training data. some sense, fitting neural network component functions allows to specify higher-order epistasis in locale without paying the price for defining it over the full design space. Real-world design procedures typically only have knowledge of f in small locale, which is why locally-valid decompositions suffice.
A.2 DECOMPOSITION CIMIPOOSSTIIO--NNWAEEEE EDA DERIVATION
Given objective, fx)), to maximize, decomposition- unaware unaaaar EDA (or "naive" EDA) transforms the original problem into a distributional optimization problem as follows (Brookes et al. 2020):
max f(x) = max EE (=))
For the equivalence to hold, the search distribution, po (x), must be capable representing a point mass on x=. Intuitively, one can think of an EDA as having search distribution, po(x),, that acts as a spotlight on the design space, which is iteratively moved toward areas of the space with high fxx
in expectation. Modern EDAs parameterize the search distribution with neural network generative models, such as a Variational Autoencoder (Kingma & Welling (2013); e.g., Brookes et al. (2119).
Procedurally, after having initialized the search distribution, Pao (z), this naive EDA then iterates through these three steps either N times or until some convergence criteria met (Larrañaga & Lozano 2001
Brookes al. 2020): :
1. Draw K samples from the current search distribution, {a*K
Ik=1 p&r (2)
2. Score each sample with the objective function to get its weight, w = f(x*).
3. Update the search distribution parameters with weighted maximum likelihood estimation (MLE), using the weighted samples, g=+1 arg max; E [w* log pe (= In older EDAs, this step was often instead a Iruncated maximum likelihood estimation, in which po(x) modeled only samples with the largest weights.
Often, a predefined monotonic "shaping" function, W(-). is additionally applied to the weights control optimization dynamics. Its monotonicity guarantees that
16

---

Published as a conference paper at ICLR 2026
arg max; f()) = arg max, (x)).)). We've written the EDA above and its derivation be- low without W, but f(x) can be equivalently replaced with (f())) everywhere.
For clarity, sketch out one way of deriving the EDA update (step 3) from Eq. 7 At each EDA iteration, we want to improve the current search distribution with an update, Pan (2) + P.(()) oc f()) pan (z:). To do this, we maximize
1 Z BBon (z) log po + (x))
-DKL (p po(x))
(8)
where the entropy term can be dropped because it has no dependence on 0 and division by Z can be dropped without changing the objective's maximizer. The sample approximation of the remaining terms is exactly step 3. Assuming an exact update (i.e., infinite samples such that the expectation is evaluated exactly, (xx has sufficient capacity to represent pi, (x)) and the KL divergence reaches O),
the resulting search distribution at iteration n. is P. (3) ix [(o (2r). As n - oc, this distribution concentrates all of its mass the global maxima of f and a result, f(x*) = Epp(t)) [f:)]] The derivation in Brookes & Listgarten 2018 is perhaps closest to this one, though their search distribution improvement operator is motivated by Bayes' rule. It can be generalized to the derivation we give by simply writing P. (1:) x (zz) pa (I)) for some C. Bayes' rule yields a special case of
this improvement operator, namely one where the multiplier is C(x) Σn Pur (x)p(S)z) We choose not to interpret fx) as normalized distribution in our derivation, though one can definitely do so.
p(8(x)
In practice, the weighted MLE problem in step 3 is typically solved with gradient descent; we will choose to use only a fixed number of gradient steps, which can be theoretically justified through the equivalence of EDAs to Appectation-Maximiaattee
(Brookes et al. 2020).
A.3 ALGORITHMIC SKETCH: FROM EDA TO DADO
In this section, we describe how infusing a naive EDA (Sec. A.2 with awareness of a known de- composition of fz) stands to make the EDA more statistically efficient.
think of an EDA having search distribution, p=(x), that acts as spotlight on the design space, which is iteratively moved toward areas of the space with high f(x) in expectation. In decomposition-unavee
Intuitively,
EDA, after having initialized the search distribution, pgo (z), with initial
parameters, 80, the EDA proceeds by iterating through these three steps for n = 1...N N or until some convergence criteria has been met (Sec.
1. Draw K samples from the current search distribution, {a^M ~ po-
2. Score each sample with the objective function to get its weight, w² = f(at).
3. Update the search distribution parameters with weighted maximum likelihood estimation (MLE), using the weighted samples, 0=+1 arg maxe E [w> log (=()))]
Now that a standard EDA is fresh our minds, we can consider what might mean to infuse with knowledge of how fIx is decomposed, and how this might prove useful. Before discussing the general case, here we first build intuition by considering a simpler example, one that is "fully". decomposable rather than only soft-decomposable. That is, consider f(x) (i + C2(aa) where the two meta-variables in and d2 do not share any original variables, 2i, corresponding to a junction tree without any edges, only independent nodes. In this setting, we can break the EDA down into two independent EDAs, each with reduced dimensionality compared to the original EDA. Specifically, we can employ fully-factorized search distribution ("fully" again denoting that no variable appears in more than one factor), (o = Pe, (fa ) (f2). We modify the three steps of the EDA as follows:
1. Instead of drawing samples from one search distribution over all design variables, we sample each meta-variable from its own search distribution separately, of ~ pay (ft) and atj ~ pay (ā2).
2. Instead of scoring I² with f(x). we score each variabariable separately with its component function, for i E (1, 2). to obtain separate weights.
3. Instead of one update to joint search distribution, dently of the others with weighted MLE on samples from the relevant search distribution, gx+1 arg maxo; Eegay [uo log po, (++)], for i E (1,2).
update each factor indepen-
17

---

Published as a conference paper at ICLR 2026
What has this bought us? We have split optimization problem combinatorial space into two independent optimization problems, each over a smaller combinatorial space. For example, for binary design variables, where each of the two meta- variables comprises say 15 variables, we have transformed one joint optimization over a space of size 230 2 10 to two independent optimization problems each of size 2 x 10', thereby reducing the effective search space size by four orders of magnitude. For larger alphabets and sequence lengths, these differences will be larger still. Impor- lantly, the weighted MLE update with fixed number of samples, K, is more statistically efficient for each divee-iimeniional search distribution factor than it is for decomposition- unaware EDA using a single joint search distribution over all design variables.
In the general case, that is, beyond full decomposability, the same variable. 24, can participate in multiple component functions C, (z)) In other words, the meta-variables ž may overlap, corre- sponding to there being edges in our junction tree. Because of this "component coupling", the EDA can no longer independently optimize each component function and will require coordination be- tween meta-variables. That is, we can no longer simply divide and conquer as above; rather, we must divide, and then collaborate to conquer. Collaboration among meta-variables will occur by way of a new form of pressage-passing that allows perform distributional optimization while leveraging the decomposability reflected by the junction tree.
Earlier described decomposabilityy f(z) C2 C2 C. (²x) to help provide some intuition. However, to develop DADO it will be more convenient to equivalently describe the decomposability with junction tree defined by nodes that contain multiple design variables and undirected edges that connect the nodes. We define this junction tree decomposition in Sec. [2. S.c.2.1 and and explain its connection to the above decomposability in Sec. A. It suffices say here that design variables x; that occurred uniquely in one meta-variable é now sit in a single node, Ty; variables that previously occurred in multiple verablaas now appear in multiple adjacent nodes. With the idea of a junction tree in hand, we are now set to complete our intuitive overview of DADO, which will use a factorized search distribution matching the topology of the junction tree. Namely, we will convert the undirected junction tree into directed, acyelic graph (DAG) by rooting it and then directing the edges outward from the root. This operation will yield a search distribution that can be written as (o(r) (ā,)) >, po (ā; Tp(o)) ), for root node ã, and non-root nodes, ã; with parents, pp()) Note that each non-root node has exactly one parent in DAG. Each search distribution factor is conditional probability distribution that will be parameterized by neural network such as an autoregressive model.
Consequently, sampling from such a factorized search distribution is both computationally efficient and easy, so step of the EDA remains straightforward. However, steps 2 and 3 of the EDA are not so easily generalized. For these steps to make use of the decomposability, leverage a special kind of message-passing to coordinate updates to different factors of the search distribution as achieve globally consistent update for the entire search distribution. At a high level, our decomposition-awaa -wwre EDA will proceed as follows for the n + 1 iteration, compared to the deloppolcoe EDA presented previously:
will need to
1. Instead of sampling each meta-variable from an independent search distribution factor, we sample meta-variables sequentially from conditional search distribution factors, starting with the root node, 5.. ~ pa- (i and conditionally sampling down the leaves, ã ~ pan (3: 5k =
2. Instead of scoring each meta-variable á separately with its component function, C,att) ) we use message-passing on the junction tree to obtain scores accounting for each meta-variable's coupling to its parents and children, w = (aah)) for the root and w Q, ah ) for non-root nodes.
3. We perform a coordinated update of the search distribution via far-factor weighted MLE on the relevant samples, 6p+1 = arg maxo, EqzØy [w log Po (xt) for the root and Po
0,++1 arg maxe; Eqat 2 w log (a ) for non-root nodes. These updates can be performed separately, as in the fully-decompssed EDA, but are coupled according to the junction tree by the weights.
For formal derivation of DADO, see Sec. 2.1
Sec. A.4
18

---

Published as a conference paper at ICLR 2026
A.4 DADO DERIVATION
Herein, we give an augmented version of the derivation presented in Sec. 2.1 derivation of the naive EDA given in Sec. A.2 The goal of DADO is to obtain a generative model, po()), that maximizes the EDA objective, arg maxo Po)( [=)). while leveraging de- composability in f(x) for optimization efficiency. DADO requires as input a decomposed ver- sion of fx)), which can be described by an undirected junction tree. T= (V, with nodes, N, and edges, E. As noted in the introduction and shown in our experiments on proteins, iden- tifying useful decomposability feasible in practice. Given the junetion tree decomposition, we write ((x) Ji(a) + E(ijje [āg Tz), where we refer to (i: as "node component functions" and ( R5y) as "edge component functions" (Fig. Dp. This functional structure is intimately related to notions of the "epistatic landscape" of a protein's 's property function and spec- tral analysis (see Sec. A.I.I). When the component functions are not known a priori, they can be parameterized and fit to labeled data.
We build on the
We will begin our exposition by recalling how do decomposition aware exact (ton-distributionaa) optimization, that is, to solve arg max, f(x). This problem is efficiently solved with a classical cessaae--pssing algorithm, which coordinates local optimizations across parts of the junction tree to obtain a single global optimum. Its efficiency comes from breaking optimization over all variables jointly into separate optimizations for each (smaller) meta- variable. Having loaded the reader with this intuition, we then adapt these ideas to distributional optimization, yielding DADO.
A.4.1 CLASSICAL MESSAGE- PASSAGE-PASSING FOR NON-DISTRIBUTIONE OPTIMIZATION
Although classical message- passsagggg has been largely used for probabilistic inference on probabilistic graphical models (Pearl 1988, Shah, 2014)), it can also be used for exact optimization of a function, xx = arg max; (() (Vlassis et al. 2004). In particular, message-pann paaggegassing can be used to find global maximum of a function defined on an undirected junction tree, T, by making use of its topology for optimal time-complexity, In particular, one takes the junction tree, which is undirected, and roots it to obtain directed tree, which induces a hierarchy among the meta-variables from root to leaves. Each node is responsible for accumulating information from all nodes its sub-tree and then passing this information on to its parent. Consequently, the root node receives information from the entire tree, which is sufficient to set its variables in a globally optimal manner. Then, starting with the root, each parent communicates its variables' optimal settings to its children, which can in turn set their variables optimally, and so forth.
To obtain the rooted tree, T := UN",E") from the junction tree, one keeps the same nodes, N' N, chooses a root node, r, and directs all edges in & outward from r, yielding directed edges, 8. Although rooting at any node will suffice, we choose r such that T' has the shortest height possible. Message-passing finds a global optimum in two passes through the tree: one round of passing messages up from leaves to root, and then one round passing messages back down to the leaves.
Given T', classical message-passing first uses dynamic programming to accumulate information from the leaves up to the root. Similarly to any dynamic programming procedure, we accumulate solutions to increasingly larger intermediate sub- problems. In this case, each sub- sub-problem is to find the value of f(x) evaluated on only a subset of meta-variables, rather than on the full set of variables in æ. Each sub-problem is tractable owing to the decomposition of the objective function into compo- nent functions for each node and edge, and by respecting the partial order of sub- sub-problems induced by T'. Specifically, one computes a value function. Vmmak (#p()) ). for each node E A/ \ {r]. which tells us for each setting of its parent, Sp(i). the value of the intermediate objective function defined by the edge component function, fp(s), ā.), plus all component functions over the sub-tree rooted at i, given that all nodes maximize their respective intermediate objectives. Computing value functions comprises the first pass through the tree, from leaves to root:
(max((pp))) == maxs, (A((E)) + (Fpl), ū.)) + Σ eEchildren())
/max (i:)
Notably, V.max (1) provides sufficient information about all nodes in the sub-tree rooted at to optimally choose the value of #p(o) with respect to its children. Thus it follows that once all value functions have been computed, the root's assignment can be set in a globally optimal manner from
19

---

Published as a conference paper at ICLR 2026
its children's 's value functions,
+.. = arg max (fr(ā,) +
E eEchildren(r) Ve(at)).
Having chosen the root assignment, we then pass it down the tree as Tp(i) = p()) to its children, which successively pass their chosen assignments to their children, all the way to the leaves,
x = arg maxa, (f((T.) + = 2.) + Σ (Echildren())
Vo(ā;)).
resulting in a global maximizer x² of f(x). This dynamic programming "traceback" of optimal assignments back down the tree constitutes our second and final pass of messages.
Alternative notation. For convenience of our generalization to distributional optimization, we re-write the parent value functions V. (Tp(e) terms of child-parent, QTTH Tp(t)), and single- node, Q,Te. (z:) value functions:
(T
ViTan (pp(t) == maxi, Qut (pt))) where Quia (Tx Tp(o) = ( T.) + Quia (x+, and
(9)
(10)
Q. (i+) == f(((t) + > Eeecchitden(() ViLa (.:))
(11)
In contrast to the original parent value functions, Q Tp)) represents the effect of the choice of both ()()) and i; on their edge component function plus component functions over the sub- tree rooted at i, assuming all descendants of maximize their corresponding value functions. In- tuitively, Qut Tp(o) is the value function on edge (p(i), i) prior to ã; being maximized out, which will become useful if we want to, say, sample I; according to some distribution instead. Quaa (n())) ) is composed of two terms, one of which depends on its parent, and one of which doesn't, Q.aax (Z:). Written using the Q-functions just defined, the equivalent traceback equa- tions for selecting globally optimal assignment, x²,, simply žç == arg max2, E == arg max;, QTaa((
Q.aa Qrax(z,)) and
(ā,, ). In other words,
x² = arg max, f(x) = {arg maxs. QTTIK (ār)) + {arg maxs, Q (T; (
(12)
arg max
(QTEa (T+) + + (p()).))ES QTTE (Te, =p()))
(13)
A.4.2 FROM CLASSICAL MESSAGE PASSING TO DISTRIBUTIONAL OPTIMIZATION
In the same way that an EDA transforms arg max, fx) into a distributional optimization problem (Sec. (.2c.
rewrite the equivalent peessage-passing objective in Eq. 1 as DO. Because the original optimization problems over x are equivalent, their DO formulations are equivalent too,
arg maxo E.po(r) [f(x)] arg max E. (po(z)
Q (mT))()))))))))))))))))())) + E En Tm
(e Q. (T,,
(14)
(E)o QTTaN (2r))+ Σ (P(0,)) Bno(n) 1 (Z pp())))
arg max #
(15)
where we' ve used linearity of expectations. However, a generic joint search distribution cannot take advantage of the linear additivity value functions over the junction tree topology. That is, while the classical traceback equations perform maximization over each meta-variable separately, Eq. [5] uses single, unfactorized search distribution over all variables, px(x), to optimize each Q-function. We address this next, by factorizing the search distribution.
Factorized search distribution. Classical pressage-passingg (Eq, 12) independently maximizes each Q-function conditional on the choice of (p))), instead of explicitly maximizing all variables in jointly (3). This is possible because each Q; captures all relevant global information needed to choose T; It stands to reason then that DO can do something similar. Specifically, instead of train- ing a single joint search distribution, we train smaller search distributions over each ž,, conditional on Tp(a). to separately maximize each Q (ā,, Fp(t)) That is, factor the search distribution according to T resulting in DAG search distribution, p= = po (T.) [[(()))e po (āx =
20

---

Published as a conference paper at ICLR 2026
for root node r, non-root nodes i, and parents p()), connected by directed edges E², and with param- eters, 0. This factorized search distribution can be plugged into Eq. [5r an equivalent optimization problem. We refer to each element of this product one of the factors of the search distribution. In our implementation, each factor has completely separate parameters, though we write a shared 0 for conciseness. Notice that each factor of the search distribution interacts with each other factor through the directed edges, E". That is, the distribution of ã; depends on its parent's factor, and through it, all of its parent's ancestors: (=():) = pe (ā; pp(() (T(()))) In turn, each of node i's children's 's factors depends on po()). Due this coupling. we cannot optimize each factor fully in- dependently. But we can still update each factor separately from the others in a globally-consisten manner via message-passingg mizing its own meta- varaalar, but will need to coordinate with its neighboring factors by getting information from them about how they are optimizing their meta-variables in a manner analogous to classical message- mresaag-passing. Our current messages, Q convey the value of each intermediate objective when all meta-variables are chosen via maximization. For DO, we'll require messages that communicate values when meta-variables are chosen according to the factorized search distribution, po(r).
In particular, each factor will only be responsible for directly opti-
Distributional value functions. While one certainly could choose to optimize classical value functions using an EDA search distribution (Eq. (5). it doesn't make sense for two main reasons. As we just mentioned, DO aims to train p9 such that maximizes f(z), or equivalently, the sum of Q-functions, in expectation. Therefore, the DO objective should consider intermediate objective values for meta-variables chosen according to the current search distribution, not those chosen by explicit maximization, in (Tp())))))) Additionally, computing classical value functions requires enumerating all assignments of each ã; for the maximum used in (īp(e) When ã, contains more than a few design variables, this max operation quickly becomes intractable. One reason for doing distributional optimization is to avoid enumerating massive design spaces. To address both issues, we define corresponding distributional value functions that fulfill both desiderata:
((p(o) Epo(2, 12 (0,, (p())) where (T),,))))) = (Tp), 74) + (ā:) and
(16)
(17)
(Z.) = f((a)) + Σ (Echildren()) Va (ā.).
(18)
Compared to Eq. the distributional fffunctions compute the value in expectation under pa instead of a max. Because they use an expectation instead of a max operation, these value functions can be approximated tractably and without bias by drawing Monte-Carlo samples from the search dis- tribution. The connection to classical value functions can be made clearer by defining a particular distribution which recovers them, the one placing all of its mass on the arg max, and its correspond- ing distributional value functions:
Panax (i:: A (p(o)) = {{1
if A arg arg maxi, Q.IN (Zt Tp(i))
(19)
otherwise
V.max (Tp(t)) = Ep (2 12 (ā,, "ps)) = max fi(ū:) + fp()) 7:)+
V (i.).
2i
(Echildren(s)
(20)
( (p(o) = Spen)a(pp()) .2.) + Q.ax for Q (i:) f(()))+
+
Vinux
(i:)
eEchildren()
(21)
Moreover, each distributional value function lower bounds each corresponding classical value func- tion because the expectation of function cannot exceed its maximum. This relation is evident for the base case of a leaf node,
V.nax (5 (ip()) = max f (at)). 2 Epota, ( (f(
2.)] = V. (Tp(s):
(22)
and by inductive argument, also true for all other nodes' V-functions, and similarly, the Q- functions, which only differ in which V-functions they sum over. This implies that the sum of
21

---

Published as a conference paper at ICLR 2026
classical value functions in the objective in Eq. 15 is bounded below by the sum of distributional value functions. This bound still holds when we take the expectation of each sum under po(x),
Epu(s) (2, (2+)]+
Σ
Eputt)( (øi, R(() 2 Bpo(=) (+.,++ (p(O).))
E
(P(S..)E
(23)
We'll optimize the lower bound with its distributional value functions instead of the sum of classical value functions.
Distributional optimization with value functions. All that remains is to derive an update rule that treats each search distribution factor separately. Since each expectand in Eq. [3] doesn't I depend on descendant meta-variables, we can replace (or) with each meta-variable's 's marginal distribution:
arg max poo (Q.. + E (p(O),)E Bpots, (oel),
(24)
We then follow the EDA derivation (Sec. [A.2] to arrive at aanple-aprroximated update rule
ing the previous iteration's 's distribution, par (x) in which each term is optimized by only a single
factor. We omit the derivation for the root factor because it'ss exactly the same as in Sec. A.2 us-
ing O (æ,) as the weight instead of fx). To update conditional factor from Eq. 24 in an EDA
loop, we define pr, (i. Eppa)) x Qº (āe, Fp(s)) Pa- (To Ep(t)), and minimize its divergence from y (#e (p(or) using the law of iterated expectations:
Epon (Fx(s) [DKL (p% (āe āp(1) ll po =p(o))]
(25)
= - (Fp(s) [e (2.15 16 Ep(e) log (ā [x(()]] (p [)(o))
(26)
Again, when taking the arg max with respect to 0, we can equivalently remove the entropy term which bears no dependence on 0 and division by constant Z. Notice that the dependence of the conditional terms on po (āp(t) in Eq. 24 approximated by sampling, which is why only individ- ual factors of the search distribution, log Po &p(s)), appear in each summand, opposed to log (5 (p(n))p(($p())) Updating pe (áp(t) to maximize (Zc. would be redundant and against the spirit of message-passing since parent Q Q-functions already include that info from chil- dren. The resulting update rule over all search distribution factors (sharing a single set of samples; see Fig. [I]) can be written as sum of weighted likelihoods for each factor,
ga+1__
arg max
Σ
(of" (i.-) log (ā,)) + (p()).)s Q¹" (ă (pt)) log (o I (p(1))
Eaopon (x)
(27)
which can be written equivalently as separate updates because the factors don't 't share parameters, as desired:
0+++1= = arg max
Σ Zxppon (u) Oº" (ā.r) log Po, (īx), and
(28)
0r
0-+1 arg max
pop (a) Q9" (ā., Ep(o) log po. (ā. Zpo)), ViENA (r).
(29)
Each factor is weighted by its corresponding value function, enabling it to coordinate with all its descendant factors despite their being updated separately. The whole DO is tied together at the top by the root factor. Our resulting algorithm (Alg. [ fits into the three EDA steps: (1) designs are sampled from (x), (2) weights, here each Q-function instead of (x). are computed, and (3), each search distribution factor receives its own independent weighted maximum likelihood update based illustrated in Fig. [D). These are repeated until convergence, for some fixed number of iterations. The factor updates are coupled only through the Q-functions, the messages across edges. {Q²" JEN are only valid while 0 is close to 6". meaning one must balance how many gradient steps are taken before drawing new samples. If too many gradient steps are taken without resampling, factor may be changing its parameters to collaborate with another factor which has already changed its behavior. It's common for EDAs to include an additional hyperparameter W(.), a monotonic shaping function applied to the weights, to alter optimization dynamics (Brookes) et al. 2020), so we include it in Alg. 1(] on lines 16 and 18. Choosing W to be the identity recovers our derivation above. Notice that in place of the classical message-passe passing traceback equations, DADO simply performs sequential conditional sampling from its search distribution. Interestingly, a very similar algorithm can be derived without using a lower bound on the value functions (Eq. (3) if one adds an entropy- maximizing term to the initial DO objective (Sec. A.5)
on these weights (steps I and 2
22

---

Published as a conference paper at ICLR 2026
A.5 ALTERNATIVE DADO DERIVATION: MAXIMUM ENTROPY OBJECTIVE
If, by comparison to Sec. A.4 we instead seek to solve closely related objective which addi- tionally maximizes the search distribution's entropy, we can arrive at a similar method by lever- aging the probabilistic graphical modeling (PGM) framework in Levine (2018). This view will provide another intuition for the relationship between classical value functions and distributional value functions approximated with Monte-Carlo samples. It also highlights the core components of our conceptual framework, which are shared even by derivations starting from different objec- tives. In other words, one need not use our specific method as outlined in Alg. to reap the benefits of decomposition-awarc distributional optimization, and variations may have desirable properties and/or simply perform better in certain scientific design settings. This derivation will assume that you'l 'e already read key parts of Sec. 2eSc. and are familiar with the function decomposition, search distribution factorization, and value functions.
Maximum entropy decomposition distributional optimization.
We begin from the max-
imum entropy problem (Jaynes 1957: Ziebart al., 2008 Zhu al. 2024
arg max E (f)) + (H)))) 820. 0.
(30)
This objective is arguably more consistent with the desired end result for many scientific design problems. Often, a distribution over good solutions is preferred to a single solution (i.e., a distribu- tion with no entropy When optimizing a predictive model or more broadly, dealing with uncer- tainty or inaccuracy in the objective function, having a distribution over good solutions can be partic- ularly important. The maximum entropy objective has a closed-fore form solution, p" (z) x exp fx)/3, motivating the solution of an equivalent variational objective, arg ming (PL(Pe ll from which we will obtain decomposed update rule. We first plug in decomposed versions of f and po (1r), as described in Sec. 2.1 Sec. A.4, to get
pe(z; =p(o)) exp EIEN (ip( + f((3s) (EV
arg min DKL (po p²) arg min DKL
po(īr)
II
B
(31)
Notice that we've dropped the normalizing constant for p" because it has no dependence on A. In what follows, for compactness and clarity, we will abuse notation slightly by including the root factor, (T.),, with the other nodes, writing it as conditional on parent even though it has no parent and is an unconditional distribution.
EIEN (i,p Fp(o)) + Si(T:)
DRL
II pa(a; | (p(()) || NEW
exp
(32)
3
EieN Se Sp())) + f((((())) log II
= (x)(
po(as ap())
(33)
B
iEN
(ap Sp(i)) fo fa(pp(t) - log (2; (p(s))
= Eps(a))
(34)
3
LiEN
Fip( + f((āp())
E ((o(
log =p(o))
(35)
3
The final line decomposes over nodes, but still has an explicit dependence on all ancestor node search distribution factors. One way to obtain separate update for each search distribution factor is derive form of p" (z) factorized in accordance with (a()'s factorization, so that we can simply minimize each factor's divergence from the optimal one.
In our experiments entropy bonus), the search distribution avoids collapse because we don't until convergence (fixed number of iterations), because we sweep learning rate (if it's very high, the distri- bution may quickly collapse to a point, limiting further improvement). Similar strategies were used by Brookss et al. (2019) instead of an explicit entropy bonus. EDAs were originally used for discrete optimization problems where a single solution was desired. Using distributional optimization to solve more nuanced problems calls for modifications.
23

---

Published as a conference paper at ICLR 2026
Optimal search distribution factors via wressage-psssing To derive the optimal factors, p" (ā; Ep(o)), ). we'll use pesssag--assingg for which we now introduce some notation. Note that we adopt the same convention
above of including the root with the rest of the value functions, meaning writing it with a parent, although doesn't have a parent and all terms involving its nonexistent parent simply O. We recursively define value functions (consistent with Levine (2018) as:
exp (z, Zp(s)): Vi(āp()) = log Σ
(36)
Bi
Q((2s; = (Tpt), + (z.)., and Q( = f((0,) E (Echilddren())
Vo(T.).
(37)
Though they might seem a bit arbitrary, substituting the value functions into Eq. 35 will help il- luminate p² (ā; Zp(s)). Notice that the Q-functions have the same recursive definition as in DADO, though the V-functions are different. We first consider the base case leaf node, where
Qr Ep(a) fpCo, (ān()): 2.)) + fz())). We substitute, add 0 Vi((p()))/B (i($p())) and use the definition of KL divergence:
logexp ( Po #p(()) + Vi(āpro)
log exp (())))
pep ((cc))
3
B
B
(38)
= (P | DEL pol (it =p(t)) || exp Qi xp(o) - Vi(āp(e))
Vi((p())
(39)
8
B
Because pi (p(r)) only appears in the KL divergence, the overall expectation will be max- imized, as far po po I Xp(s)) is concerned, when the divergence is O. Therefore, for leaf node, p² ("; (ppt)) = exp((Qi(aa Fp(r)) (TTp))))) which is a proper distribution because exp (i())()) is exactly the normalizing constant of exp (i, Tp(e)) Notice that we still have term, (F,o)) (( outside of the divergence (Eq. 39.. Maximizing this term must therefore be the responsibility of po(p(()), not po( Tp(o) In fact, these leftover terms are in- cluded in the parent's Q-function via the sum over all child V V-functions in Q. as we defined above (Eq. 37 It follows that the recursive case has the same form as the leaves, which had (.)).) f (b:) because they've 've no children (more detailed derivation in Levine (018)). That is, for all non-root nodes! i. p" (ā; (p(a)) = exp((Q (T,, Tp(t) - Vi(āpto))/B)) for the root node, (z,,) exp((Q, - V)) where V, = log Σ,, exp (ă,, is a normalizing constant with no dependence on 6. We now have optimal form for each factor and can proceed to minimize our search factors' divergences from them completely separately.
Distributional optimization by matching optimal factors. Given factorization of the optimal search distribution, p" (z)) we can train (o(x''s 's factors to match their corresponding optimal factor distributions, resulting in an algorithm very similar to DADO. The objective for node i's factor minimizes its divergence to the optimal factor p" (#; Zp(s).
(i: xp(e)) || exp Qi (p(r) vi(āpce))
arg min E #: Epone (Ep))) DKI. Po. Po.
(40)
8
arg max E. Eeg (Rxi) [Epa, (Relagca)) ((()) Vi(āp(n) + BHo (Te ()())))
(41)
6,
and an analogous objective can be written for the root node. Notice that V1((p()) can be dropped completely or
any other constant (with respect to X.) without altering the arg max, though certain settings may yield more favorable optimization dynamics We approximate the expectations in the objective using samples from the current search distribution (as in EDAs) to yield a weighted maximum likelihood update rule for each node i,
Σ ((Q(())))))))) - Vi(Ep(o)) lo po. (ā; (p(9) - |3 log Po, (#x (p(())
g arg max e.
Brepan (a)
(42)
(F)tt) is often thought of as "baseline" in the RL literature, the choice of which doesn't 't bias the objec- tive's s gradient, but can reduce its variance substantially and lead to more efficient policy optimization (Sutton)
al. 1998).
24

---

Published as a conference paper at ICLR 2026
and since the weight should always be non-negative, monotonic shaping function, W, outputting non-negative weights, can be introduced too,
0u+ arg max
E (w(e(((s
(pm))))) - Vi(Tp())) - 3) log
logpo, Po. (T. (ppcs)
(43)
O:
Pmpp (x))
There are a few minor differences from DADO. First, there's an entropy-maximizing term. Sec- ond, notice that the value function definitions are a little different from in DADO. Here we use log Σ exp to define Vi(āp())), whereas DADO uses an expectation over the search distribution in- stead. Third, the Q-function has (i(()()) subtracted from it, whereas DA DO's derivation doesn't. t. As mentioned above though, ()(())))) 's role in the objective is that of "baseline", and it can be changed out or dropped. In our implementation, we do subtract ((p the expectation one, not log) exp-as mean baseline function to reduce weight variance and obtain more stable update (Sec. Overall though, the result is quite similar to DADO (Sec. (A.4), supporting the gen- erality of the decomposition-awarr distributional optimization framework, and suggesting that there are a variety of related objectives and algorithms one might use under this umbrella. In particular, one might adapt methods from RL. For example, Peng et al. (2019)'s derivation could be adapted to get a similar weighted maximum likelihood objective with the difference in value functions (also called the advantage) exponentiated instead. TRPO, PPO, or CbAS (Schulman et al.) 2015 2017 Brookes et al. (219)) might be adapted to regularize the search distribution toward some (identically factorized) prior, or the previous iteration's search distribution for stability. We emphasize that the core idea of DADO is to leverage message-passing (i.e., value functions) on arbitrarily decomposed objectives to derive a distributional optimization procedure which updates search distribution factors separately, regardless of the specific objective and algorithm used.
From classical value functions to distributional value functions. Under this definition of value functions using log Σ exp (Eq. 36). we can view (((p)) as a soft maximum, which becomes a hard maximum when Q; are large, such that for āq. VT(Tp()) =s Qi F7())) The hard maximum case resembles classical message-passee passing, from which we began. When Q; are relatively small, information about multiple é; will pass through log exp, thus propagating info for a distribution of descendant states in accordance with y" Tp(t) x exp Qi /3, Another way of seeing the transition from classical message-passing to distributional optimization is that for 33 near 0. p" ( fp(t) will concentrate all its mass on the arg max. But for larger B, p² (ă. Ep()) will concentrate on multiple designs proportional to their Q values. As B - co, p (.: Ep(e)) becomes a uniform distribution.
(,(āt,
A.6 PROTEIN DATASET DETAILS
We investigated seven protein property datasets from Tareen et al. (2022) and Notin et al. (2023).
These proteins span various lengths and all have more than single mutations away from wild-type. These datasets are:
Adeno-associated virus 2 capsid protein (AAV2 capsid or AAV; L 28, D = 20), with data from Bryant et al.) (202I), which assayed 42,329 sequences for virus viability, includ- ing sequences with as many as 27 mutations from the wild-type. The full protein of length 735. The dataset was accessed via Notin et al. (2023).
Amyloid-beta (Amyloid: L. 42, D = 21, where the extra state is stop codon), with data from [Seuma et al.) (202T), which assayed 16,066 sequences for aggregation with a nucleation score, 97% of which are double mutants. The dataset was accessed via Tareen et al.] (2022).
Yeast transcription factor Gen4 (L = 44, D = 20), with data from Staller et al. 2018 which assayed 2,639 sequences for activity, including sequences with as many as 44 mu- tations from the wild-type. The full protein is of length 281. The dataset was accessed via Notin et al. (2023).
TAR DNA-binding protein 43 (TDP-43; L = 84, D 21, where the extra state is stop codon), with data [Bolognesi et al. (2019). which assayed 57,996 sequences for cell toxicity, 98% of which are double mutants. The dataset accessed via Tareen et al. (2022).
25

---

Published as a conference paper at ICLR 2026
Protein G B1 domain (GB1; L = 55, 20), with data from Olson et al. (2014). who scanned all possible single and double mutants and assayed their 536,963 binding affinities to IgGFC, of which 99.8% are double mutants. The dataset was accessed via Tareen et al. (2022).
ynzC, small protein domain (L =39, D = 20), with data from T'uboyymaa al. 2023) assayed for its folding stability on 2,301 variants, 68% of which are double mutants. The dataset was accessed via Notin al. (2023).
Chlamydomonas reinharliiü's light-oxygen-voltage domain (CreiLOV or Phot; L = 118, D 20), with data from Chen et al. (202)). which assayed 167,530 sequences for fluorescence, including sequences with as many as 15 mutations from the wild-type. The dataset was accessed via Notin al. (2023).
Of these, the four shown in the main text were chosen according to two criteria. First, we wanted proteins for which the Alphafold3-derivedd junction tree had relatively small nodes (i.e., more de- composable proteins), discussed more detail below. Second, we also prioritized proteins for which the datasets contained sequences many mutations away from the wild-type, as the resulting predictive models are more likely to be realistic in a larger area of the design space than datasets that only assay concentrated ball of sequences. Based these criteria, we decided to focus on AAV, Amyloid, Gen4, and TDP-43 (Fig.B) These four proteins' decomposition junction trees have nodes of cardinality five fewer, whereas the junction trees for GB1 and ynzC have many nodes with cardinality over 10, and CreiLOV has nodes with cardinality greater than 20. Optimization performance on the remaining datasets is included in the appendix (Ai.
We visualize all seven proteins' AlphaFoldd-prediette structures and corresponding contact graphs and junction trees (Fig. AT). Panels 3-6 (the rightmost four columns) illustrate that AAV, Amyloid, Gen4, and TDP-43 yield junction trees with smaller nodes and fewer positions overlapping in multi- ple junction tree nodes compared to GBI, ynzC, and CreiLOV. It would be interesting to reduce the contact threshold for the latter problems in order to obtain smaller nodes. As the possible efficiency gain from using decomposition-avaal like this graph is, expect DADO would yield more substantial performance improvements over standard EDAs in such a setting. Our studies of predictive model accuracy when varying the contact threshold (Figs. A6 A7 A9 suggest that one may be able to obtain aimost-as--ccuree predictive models under decomposition graphs with lower contact thresholds, but we did not further investigate optimization performance on these models.
method like DADO determined by how sparse and tree-
A.7 BASELINE METHOD DETAILS
We compare DADO with three baseline methods. The first is the naive, decomposition-unavve unaawar EDA, which we refer to as "EDA". Algorithmic details of this method are in Sec. P and Sec. [A.2 also includes derivation of the EDA algorithm. We take this as the starting point for the other baseline methods and now describe how the EDA is modified for each.
The second baseline we consider is the factorized distribution algorithm ("FDA") from Mühlenbein & Mahnig (199)). The only modification to the naive EDA is to replace its search distribution, ordinarily joint distribution over all design variables, with a factorized search distribution. In all of our experiments, we give FDA the same search distribution factorization as DADO for fairness; as such, in our protein experiments, FDA's factorization is based on the AlphaFold3-baar Comparing to FDA sheds light on how much of DADO's 's improved performance can be attributed to it using factorized search distribution versus the contribution of using message-passing value functions for the factorized search distribution update because FDA only uses the former and not the latter, whereas DADO uses both.
contact map.
The third baseline we consider is proximal policy optimization version of the EDA, which we call "PPO". This baseline is inspired by /Schulman al.) (2017). We only modify the EDA's up- date rule to match the PPO update rule. That is, instead of minimizing the naive EDA's loss, ((0) = Eytrpon (z) [f log (o we instead minimize a proximal version of it,
min { f(xx)) Peu (x) I (x+)) clip(
po(a*)
Po (o - 6, 1 +0)
T(()) = Ext (a)
Pau (x)
26

---

Published as a conference paper at ICLR 2026
JT node edge cardinalities
70
Npde Anriuei
Agpiapuco e0pg
Mode Ageuieen
(outline = root)
JI node
postions in
Junction
XXXXXXXXXXXXX  1
JT node hulls (2D PCA)
Theign pesatin mode
PC2 (8%)
poultrr
PC 1
POd (
63X
Membership positions in junction nodes
#JT nodes contenin postion
Dontiet mage Posttin
G3A
Contact graph (designable positions)
edge Atom
Backboon
TEEA
42A
Designable positions context of 3D structure
Bacibona Designable postion
3, AAV2 capsid
b, Amyloid-8
27

---

Published as a conference paper at ICLR 2026
JT node edge cardinalities
322
MONl Me
8 S09g Aapupi
Junction topology (outline = root)
Jf node
postions
$00000
JT node hulls (2D PCA)
Theign pestton mode
(88))
poultre nn.I
15000 10d
P21 (N666)
Membership positions in junction nodes
PJT nodes contening postion
Dontiet mage Posttin
43
Contact graph (designable positions)
Backdooe
Atom
Designable positions context of full 30 structure
Bacibona Designablo postion
d. TDP-43
Gen4
28

---

Published as a conference paper at ICLR 2026
JT node edge cardinalities
Mode Anneaano
Aapupiro
spoN
Junction topology (outline = root)
JI node
postions in.TT
JT node hulls (2D PCA)
Theign pestion mode
(28%)
(2%)
poultrrs n..I
1800 FDd
200A
Membership positions in junction tree nodes
#JT nodes contenin postion
Dontiet mage
204
Contact graph (designable positions)
Backdooe
Atom
Designable positions context of full 30 structure
Bacibona Designablo postion
GB1
ynzC.
29

---

Published as a conference paper at ICLR 2026
Figure A1: structure, contact graph, and junction tree decomposition each protein. For each the seven protein datasets considered. a, AAV; b, Amyloid 8: c. Gen4: 43:-43:e, GBI: f, ynzC: and E, CreiLOV, we show the AlphaFold3 - predicted 3D structure and the contact graph and junction tree (JT) decomposition obtained by imposing a 4.5 A contact threshold. From left to right, each panel shows: (1) the full 3D backbone structure, with the a- cr-caroon atoms of designable positions highlighted in red: (2) the contact graph restricted to designable positions, where edges connect residuc pairs with any atoms within 4.5 A of each other; (3) positions in the contact graph colored by how many JT nodes each participates in (global scale 25);; (4) 2D PCA projection positions overlaid with convex hulls representing cach JT node (hulls may encompass positions not contained as result of dimensionality reduction), colored by cardinality (global scale 1-26) and with the root rootnoo node marked by a thick black outline; (5) the directed topology, with nodes colored by cardinality (global scale 1-26) and the rool node marked by a thick black outline: and (6) each node's 's cardinalities (top) and edge cardinalities that is, of the intersection of the node's 's variables with its parent's 's variables (bottom) with shared global scale (1-52). To examine high-resolution details (e.g. 3D residue orientations) please view on a screen and zoom in.
70
cardinalities
50
20
Mooo do0gpees.000
Junction topology (outline = root)
J mode
porsions
JT hulls (2D PCA)
Deeign pootton mode
(29%))
posiione
Membership positions in Junction tree nodes
nodes contamingggggnoooooooo.... postion
Contact edge
Posilin
Contact (designable positions)
berkboons
Abm
Designable nositions context of full 3D structure
Designabee poslioo
Backbore
g, CreiLov V
30

---

Published as a conference paper at ICLR 2026
where e is hyperparameter that set to 0.2 suggested in the original PPO paper. PPO-style updates can be interpreted as stabilizing the search distribution update of the naive EDA by penal- izing it from deviating too far from the previous iteration's 's search distribution. This approach could easily be added onto DADO and FDA; we did not explore this but expect it would be helpful for these decomposition-awee methods too.
A.8 ADDITIONAL. IMPEEMENTATION DETAILS
We record various implementation details for reproducibility (also code released upon publica- tion). We implemented of our code using Jax, Flax, and Optax (DeepMind et al. 2020).
A.8.1 SEARCH DISTRIBUTION
For the search distribution DADO, we used MLP-based autoregressive neural networks for each position [1..... L] specifically, MLP of size [64, 64] that takes input all conditioning variables (parents in the tree) and outputs logits for each step of autoregressive decoding. Our naive EDA implementation is the same, only assuming all variables are in one meta-variable- a standard autoregressive model over all variables. The FDA baseline uses the same architecture as DADO, whereas our PPO baseline uses the same architecture the naive EDA. For gradient descent to fit the search distribution, we used the AdamW optimizer (Loshchilov & Hutter 2017)
with default momentum parameters 31 0.9, 32 0.999.
A.8.2 FULLY SYNTHETIC FUNCTIONS
We simulated one fz) for each of three sequence lengths, {25, 50, 200), by first randomly sampling a junction tree topology, and then randomly specifying the component functions. Each node function, D-vector, was sampled J, RDD (o, 0.01), where W denotes a Gaussian distribution. Each edge function, D x D matrix, was sampled, fij E DRD ~ ~N,0 0.0025). To ensure reasonable degree of non-smoothness (xx, we further explicitly added what biology is known as reciprocal sign epistasis (Starr & Thornton 2016 Li et al. 2024). Specifically, for each edge function we, twice, randomly assigned one of the 20 alphabet letters to each node, i) sampled an effect size, À NN(0 4). Finally,
Ia = A, == B and ii) H; == C, J.j
D. Next we set fr = A, Tj = B) O and fü, (2, = C, = B) = X, and also half of the time, fi, A, D) 0 and Ji.j (xt A, Tzj D)
A.8.3 DECOMPOSED PROTEIN PROPERTY PREDICTIVE MODELS
For each dataset, used A IphhaFold3-redicted structures (Abramson al. 2024) on the wild-type sequence to obtain a 3D structure, from which we constructed a contact graph by thresholding the distance between pairs of residues with threshold t. Following Brookes al. (2022); Romero al. (2013); Voigt et al. (002)), we use a threshold of t = A.5A We interpret this contact map as graph adjacency matrix, from which we algorithmically construct a junction tree (Lauritzen & Spiegelhal- ter 1988). This defines the topology needed for DADO, but we must also fit the component func- tions on the protein assay-labeled data. To do so, we employ an MLP-based predictive model that strictly enforces the decomposition defined by the junction tree. In particular, we use MI.P that takes as input the sequence, and, critically, also bit-vector specifiying active variables. One func- tion evaluation requires calling this MLP for every node and edge function and summing. We used 5-fold cross validation to sweep through hidden layers of size (116, 16). [128, 16). or [128, 128, 16)), learning rate (0.001 0.0001), and number of training iterations (5,000 or 50,000), to choose hy- perparameters with the lowest cross-vass cols-valldotion mean- error. Finally, we trained the model with those parameerssssssss using all data available.
A.9 ADDITIONAL EXPERIMENTAL RESULTS
Herein, we include additional experimental results synthetic landscapes for completeness primarily studying the relationship between performance and number of design variables (L), and larger alphabet sizes (D). For protein landscapes, we also perform case study one protein dataset in order to study the effect of the distance threshold, t, on decomposability, predictive accuracy, and downstream optimization performance.
31

---

Published as a conference paper at ICLR 2026
A.9.1 SYNTHETIC FUNCTIONS: L=400
We wanted to compare DADO and naive EDA on even larger design spaces, but this can take quite while to run when considering the hyperparameter sweep and replicates (Sec. [.9.7). For expediency, we chose the same hyperparameters (lemperature and learning rate) as were used by each method for L 200 (Fig.) 2) We then ran 20 replicates of each method.
At T. = 400, the trend we observed in Fig. continues- as grows, the gap between DADO and the EDA grows increasingly larger (Fig. A2). expected, because the full design space grows at faster rate than the decomposed design space that DADO operates in.
L= 400
300
DADO
6.59e
200
150
50
0
0 20
40
60
80 100
Training Iteration
Figure A2: Comparison of naive EDA to DADO a synthetic problem with I. 100. We created random function, (x). with randomly chosen junction decomposition with maximum node size of one, and randomly chosen parameters. We used alphahet size = 20 and sequence length two methods drew K 100 samples per iteration. For each iteration, we show the mean (solid line) and 95% confidence interval (shaded envelope) of the 100 samples evaluated on f(x), averaged across results from 20 random seeds. P-value from two-sided paired t-test that the mean at the final iteration is different between methods, over the 20 seeds.
400. Each of the
A.9.2 SYNTHETIC FUNCTIONS: INCREASING D
For all of our synthetic experiments, we used an alphabet size of D = 20, which reflects the typical alphahet for protein design problems. Design problems in other scientific domains might have larger alphahets we also considered D 50 and D 100 here, keeping T, fixed at 100.
We observe that DADO finds designs with higher fx) than the naive EDA across all three alphabet sizes (Fig. A3). As one would expect, as D grows and the design space grows combinatorially larger, it becomes increasingly difficult for both methods to optimize f()).. In general, such problems require using a larger sampling budget (K) and/or more iterations. Although the performance gup between DADO and EDA seems to shrink as D grows, the EDA has converged to suboptimal region of the design space, whereas DADO still has a positive slope and a lot of sampling diversity. This suggests that were we to run more iterations, DADO would continue to improve, but the EDA would not.
A.9.3 COLLECTED PROTHIN EXPERIMENTS
Here, we collect the four proteins shown in the main text (Fig. 3) alongside three other proteins we tested (details of why those four were chosen in Sec. A.6). First, we plot them with log((c f())) on the y-axis, for clarity when f(z) is high (Fig. A4). For comparison, we then show the same experimental results with just (x) on the y-axis (Fig. A5). Whereas AAV, Amyloid, Gen4, and TDP-43 have decomposition junction trees with nodes of cardinality five or less, the junction trees for GBI and ynzC have many nodes with cardinality over 10, and CreiLOV has nodes with cardinality greater than 20. For GB1 and ynzC, DADO still outperforms the decomposition-unawan baselines, but not FDA, which also operates in the decomposed design space. We hypothesize that FDA performs well relative to DADO in these cases because as the junction tree nodes grow larger, the computed value functions become higher variance estimates, such that can be better to weight the search distribution by (x) directly. CreiLOV has even larger nodes, potentially amplifying
32

---

Published as a conference paper at ICLR 2026
L= 100
L= 100, D
L= 100, D 100
4)-
30
2.40a
2 40
g
1.30
2,66
10
100
Trainin
Training
Trairing
Figure A3: Comparison of naive EDA DADO synthetic problems with large alphabets. We created three random functions, f(z), each with the same randomly chosen junction tree decomposition with maximum node size of one, and randomly chosen parameters. Each experiment used sequence length L = 100 and alphabet size D = 20, b, D = 50, and c, D = 100. Each of the two methods drew K = 100 samples show the mean (solid line) and 95% confidence interval (shaded envelope) of the 100 samples evaluated on f(z), averaged across results from 20 random seeds. P-values are from a two-sided paired t-test that the mean the final iteration is different between methods, over the 20 seeds.
per iteration. For each iteration,
this effect. For visualization and comparison of junction trees across proteins, Fig. AT It would be interesting to test this hypothesis by implementing further variance-reduction techniques for DADO's value functions.
a
AAV2 capsel prosen [L-25, DH20)
Artytt oi proten [Lm42 D-231
Gord proteet (Lmid, D=200)
TDP-43 proten (L-84, D=23)
1000 1000
Ierette
Thainig Inndion
Thiiding kindon
Traiiin ta
GB1 protein (L=55, D=20)
ynzC protein (L-39, D=20)
CreiLOV protein (L=118, D=20)
0.000
Pu 8.700--00
Pa P.5000e-05
Rogic
10C
100
20
140
1000
Training Iteration
Training Iteradion
Training Iteraaion
Figure A4: Optimization performance protein problems, plotted on a negative log scale. For each of
seven proteins of varving length, a, AAV (also Fig. 31). b, Amyloid (also Fig. Bb). c, Gcn4 (also Fig. 3
d, TDP-43 (also Fig. (31), e, GB1, f. ynzC, and g, CreiLOV, adhering to a junction tree decomposition derived from the protein's 3D structure, and then used standard EDA
fit a neural network property function, f(x),
and DADO to optimize them. Each approach drew K = 1000 samples per EDA iteration. For each iteration, we show the mean (solid line) and 95% confidence interval (shaded envelope) of the 1000 samples evaluated on - log(e- f())). averaged across results from 20 random seeds. We plot this quantity make clear the differences between methods when (x) is large; e is largest f() on a given plot, plus small constant for numerical stability. P-values are from two-sided paired t-test that the mean at the final iteration is different between methods, over the 20 seeds.
A.9.4 INVESTIGATION OF PROTEIN PROPERTY PREDICTIVE MODEL DECOMPOSABILITY
We also sought investigate how changing the threshold that determines the complexity of the junction tree would affect both predictive performance and optimization performance. Specifically, we varied the distance threshold, t, for which pairs of residues were considered contacting to ex- plore the tradeoff between accuracy of the model and decomposability. The lower the value of t, the stricter the decomposition (the smaller the cardinality of the largest meta-variabley, thus lower t should give DADO a larger advantage over the standard EDA, but may be overly restrictive and yield a worse predictive model. We decided to do a case study of GB for K = 100 samples at each iteration because in this setting, using the default threshold of t = 4.5A leads to comparable per- formance between DADO and the EDA. We wanted to see if decomposing the model further would result in a function that's 's easier for DADO to optimize without sacrificing accuracy. We find that the largest provides the highest holdout predictive accuracy, which is expected because the resulting decomposed model isn't t restricted. However, decreasing this distance down 2.75A allows
33

---

Published as a conference paper at ICLR 2026
A412 cassid (L-28 0-20)
Amylcld beca proccn (L-42, (=22))
protch (L-44, D-20)
TOP proneh 0L-84 0-23)
P- 30
P- P-256--0
7-240
GB1 protein (L:55. D=20)
yrzC protein (L=39, D=20)
CreiLCV protein (L=118, D=20)
1.
05-
0.5
0.1
02
0.4
P-
022
P- 0.70a 03
P- 0-000s--0 05
022
20
80
80
100
20
80
300
20
40
200
Traning Beration
Traiiing Hanation
Trainin Baration
Figure A5: Optimization performance on protein problems, plotted on a standard scale. For each of seven proteins of varying length, a, AAV (also Fig. (). b, Amyloid (also Fig. (bg. Gen4 (also Fig. B.) d. TDP-43 (also Fig. D. e, GBI, f, ynzC, and g, CreiL.V,, we fit a neural network property function, f(z), adhering to junction tree decomposition derived from the protein's 's 3D structure, and then used standard EDA and DADO to optimize them. Each approach drew K = 1000 samples per EDA iteration. For each iteration, we show the mean (solid line) and 95% confidence interval (shaded envelope) of the 1000 samples evaluated on f(x), averaged across results from 20 random seeds. P-values are from two-sided paired I-test that the mean at the final iteration different between methods, over the 20 seeds.
most of the predictive signal to remain while substantially reducing the complexity of the junction tree (Fig. A6a)) enabling DADO to consistently converge on designs with high (xx) within only 10 iterations (Fig. (A6)), whereas for t = 4.5A and & 9A, DADO's mean does not clearly con- verge even after 100 iterations, and its distribution remains dispersed as evidenced by wide shaded envelopes (Fig. [6)..)) Additionally, DADO definitively outperforms the decomposition unnaaae EDA when t = 2.75A (Fig. A66). When the predictive model is less decomposed, DADO is not as distinguishable from the EDA (Fig. A6c,d).
A.9.5 DECOMPOSITION GRAPH ROBUSTNESS
How crucial is it guess the decomposition perfectly a priori? One way to investigate this is to mutate the decomposition graph and observe how the resulting decomposed predictive models' predictive accuracies change. In our first set of experiments, we varied the contact threshold (f) used on the AlphaFold 3D structure to determine connectivity; lowering Ê will gradually remove more distant contacts, whereas increasing t will gradually add more distant contacts. Given different decomposition graphs, we then trained decomposed predictive model for each. For our second set of experiments, we performed random mutations to the t = 4.5A decomposition graph (used in Sec. 4 to study how robust prediction is to missing/ / extra edges. In particular, for each of N e [-50,- 5, 1,5,, 10, 50). we randomly sampled N edges to remove/add. We implemented a check to ensure that the graph doesn't 'tt become disconnected, so for some experiments, we off N at the largest number of edges that could be removed resulting in chain graph. We repeated this procedure 10 times.
In both of our experiments, we considered holdout accuracy as measured by the Pearson correlation coefficient between assay labels and model predictions. All decomposed models for each protein used the same hyperparameters that were chosen via 5-fold cross validation on the full dataset for the base 4.5A decomposed model, for expediency. We also compared to a neural network without any decomposition (i.e., an all-edges model), denoted "Naive NN" in our plots. For this model we performed an additional 5-fold cross validation on the full dataset for each protein.
We observe that the holdout accuracy of the non-decomposed predictive model and the decomposed models of varying all tend to fall within a relatively small range (Fig.A7)) suggesting that using our decomposed predictive model does not constitute substantial sacrifice compared to full (naive) model, and that our decomposition is somewhat robust to changing the contact threshold. When randomly mutate the decomposition graph, we find that holdout accuracy is generally concen- trated within an even smaller range, with some exceptions when many edges are added or removed (Fig. A8 Interestingly, for Gcn4, the decomposed predictive models outperform the naive model,
34

---

Published as a conference paper at ICLR 2026
GB1 predictive model decompositions
Largest meta- variavvrriblle cardinality
b
GB1 protein, 2.75 A contact threshold
Holdout Pearson correlation
1.2
0.97
P 4.099 03
0.98
EDA Sample Mean
0.95
10
Sample
DADO 95% CI
20
40
60
a0
100
Residue contact threshold (A)
Training Iteration
GB1 protein, 4.5 A contact threshold
GB1 protein, A contact threshold
d
1.0
0.9
0.9
0.8
0.8
0.7
0.8
0.5
0.06
0.4
0.3
0.2
20
40
80
80 100
0
20
40
80
80 100
Training lteration
Training Iteration
Figure A6: Investigation of GB1 predictive model decomposability. We varied the distance threshold, for which pairs of residues in protein GB1's 3D structure considered neighbors in the junction tree, using several values in 1 E [0A, 15A], so as to explore the tradeoff between accuracy of model decom- posability. The lower the value of f, the stricter the decomposition (the smaller the cardinality of the largest variabaria, thus lower t should give DADO a larger relative advantage over the standard EDA, but may yield
predictive model from imposing more restrictive functional structure. The Pearson correlation 10% holdout set remains quite good down to and including statistically significant win over standard EDA as seen in panel Blue stars denote values corresponding to the experiments in panels b, f = 2.75A, c, t 4.5A, and d, f = 9A. Fach method drew 100 samples per iteration. For each iteration, show the mean (solid line) and 95% confidence interval (shaded envelope) of the 100 samples evaluated (æ), averaged across results from 20 random seeds that dictated initialization of the search distribution. P-values are from two-sided paired I-lests that AUC of the per-iteration mean different between methods, using the 20 mean curves.
2.75A, a point which DADO provides
though it's worth noting that this prediction task is especially hard due to the relatively uniform dispersion of sequences throughout the design space, and the small size of the dataset. Both models fit the training data well (not shown here).
A.9.6 DECOMPOSITION GRAPH ROBUSTNESS: BOTTOM 50% DATA
Herein, we repeated the same analyses as in Sec. A.9.5 except using the bottom half of the training set (i.e., the datapoints with the lowest assay labels). We used the same exact holdout sets as in Sec. A.9.5 such that decomposed predictive models are tested against sequences with both high and low assay labels. We used the same hyperparameters for all of the predictive models before, which were cross-validated on the full dataset.
Overall, we observe that compared to Fig. A7 and Fig. A8 holdout accuracy measured by the Pearson correlation coefficient between assay labels and model predictions is lower for models and all proteins (Fig. A9 Fig. A10). Interestingly, for several proteins, the more decomposed pre- dictive models (lower t) performed better than the less decomposed models (higher t), as well as the full (naive) model (Fig. A9). This makes sense, more complex models are more prone to overfitting when the training data is shifted from the holdout distribution, and as the training dataset gets smaller. This trend is also reflected in the random graph mutation experiments (Fig, A10). Gen- here, such as predictive accuracy being relatively
erally, the same trends observed in Sec. A.c.5..9.5] robust changing / and to randomly mutating the graph (up point).
35

---

Published as a conference paper at ICLR 2026
A4V2
capsid predictive model decompositions
Amyloid beta predictive model decompositions
Holdout Pearson correlation
Naive NN
Holdout Pearson correlation
Neive MN
12
Largest cardimality
0.98
Largest cardinality
0.950
10
0.923
0.900
0.92
0.875
0.90
0.850
0.825
0.0
2.5
5.0
7.5
10.0
12.5
15.0
0.0
2.5
5.0
7.5
10.0
12.5
15.0
Residue contact threshold (A)
Residue contact threshold (A)
Goné predistive model decomposioos
TDP-43 predictive model decompositinns
GB1 predictive model decompositions
Make
savaas
837
000
240
100
32.5
500
10.0
12.5
15.0
0.0 3.5
13.5
Residue
Figure A7: Investigation of predictive model decomposability by contact threshold. We varied the contact threshold, for which pairs of residues each protein's 's 3D structure are considered neighbors the decom- position graph, using several values in E [OÅ, 15AJ, so as to explore the tradeoff between accuracy of the model and decomposability. We studied a, AAV, b, Amyloid, c, Gcn4, d, TDP-43, and e, GBI.
1.00
capeid robusmess
1:30
Anyyfold beta robustness edge
0.9
on8
089
0.6
0.4
19:03
3-50
6000
090X
10/00
0.00
5-1311
(P18)
3-201
Investigation of predictive model decomposability by random graph mutation. For each of AAV. b, Amyloid, Gcn4, d, TD,-43, and e, GBI, randomly added removed edges from the t 4.5A decomposition graph. In particular, we randomly sampled N E [-50, -10, -5, -1, 5, 10, 50] edges to remove/add. For certain experiments we limit N to avoid creating disconnected graphs. We repeated this procedure 10 times.
Figure A8:
A.9.7 RUNTIME ANALYSIS
We measure the wall-clock time it takes to run single distributional optimization algorithm for both DADO and the naive EDA (Table D] We fix both to run for 100 iterations, drawing 100 samples at each, using the same architectures as described in Sec. I and on a single GPU. We report times just to give a rough sense of runtime and scaling; one could definitely further optimize our code.
few trends stand out. First, the EDA and DADO take roughly the same amount of time to run, with DADO being faster sometimes. This speed-up may come from using a factorized search dis- tribution, which requires only conditioning on parent variables as opposed to all preceding variables (autoregressive). In general, problems with larger L take longer run, and this scaling is worse than linear for our implementation. We also notice that when increasing D and holding L fixed at 100, the runtime is similar. A single run of either DADO or EDA does generally require the full memory of a 16GB GPU, so in practice, we're re able to parallelize runs within GPUs and across multiple GPUs in order to perform our hyperparameter sweep and random seed replicates efficiently.
36

---

Published as a conference paper at ICLR 2026
AAV2 capsid, predictive model decompasitions (bottor 50% data)
Amyloid beto, predcctee model decompositions (bottom 50% data)
Ho corceation
Wave
Largest mode cardralty
Melve NM
Recorpeset
00.00
0555
2.5
15.0
5.0
10.0
C
Besithe contect threchold
Resiue contect theeshold (A)
Gord prediciive model descmppos fomm (bothom 50% dals)
spd decompoaations
pryjictvee decomsovtions
2.
12.5
250
Figure A9: Investigation of predictive model decomposability by contact threshold (bottom 50 % data). We varied the contact threshold, f, for which pairs of residues in each protein's 's 3D structure are considered neighbors in the decomposition graph, using several values in t E [OA, 15AJ. so to explore the tradeoff between accuracy of the model and decomposability. We studied a, AAV. b, Amyloid, c, Gen4, d, TDP-43, and e, GBI. We used the same holdout set as in Fig. [7 but only trained on the hottom 50% of data by assay label.
1.60
AAY2 capsid robusiness to edge mutations (bottom 50% data!
Arylliid beta rebusiness to edge mutations (bottom 50% data)
0.55
0.95
9.50
0.90
0.50
0.75
0.75
1.10
4.24
24
18-15K
3-5081
100-0 neur egai daki Datc 3P% 2011
900 OOAEENE Dotan J0N
Figure A10: Investigation of predictive model decomposability by random graph mutation (bottom 50% data). For each a, AAV, b, Amyloid, c, Gen4, d, TDP-43, and e, GBI, we randomly added or removed edges from the = 1.5A decomposition graph. In particular, we randomly sampled N e [-50,- 10, 55, 1, 1, 10, 50] edges to remove/add. For certain experiments we limit N to avoid creat- ing disconnected graphs. We repeated this procedure 10 times. We used the same holdout set in Fig A8 but only trained on bottom 50% of data by assay label.
Table I: Runtime comparison of EDA and DADO, varying L and D.
Problem
Length (L)
Alphabet size (D)
EDA runtime
DADO runtime
Synthetic tree
25
20
2.8 min.
2.9 min.
Synthetic tree Synthetic tree Synthetic tree Synthetic tree
50
20
6.3 min.
7.0 min.
100
20
16.3 min.
16.1 min.
200
20
42.7 min.
36.2 min.
400
20
139.5 min.
114.2 min.
Synthetic tree
100
50
14.0 min.
16.0 min.
Synthetic tree
100
100
14.5 min.
18.3 min.
AAV2 capsid protein
28
20
2.1 min.
3.4 min.
TDP-43 protein
84
21
6.7 min.
10.5 min.
37


Reviews:
[{"review_id": "XPOITVk9c5", "reviewer": "Reviewer_75LP", "summary": "This paper proposes Decomposition-Aware Distributional Optimization (DADO) that leverages decomposability and employs a generative model to guide the search for the target distribution.\n\nDADO is compared to standard Estimation of Distribution Algorithm (EDA), a form of Expectation-Maximization that is not decomposition-aware.\n\nThe core of DADO is its use of an objective function decomposed into a junction tree, which enables node-level estimation of value (Q- and V-) functions that represent the choice of variables at each edge and node respectively. These value functions are used to update the search distribution via dynamic programming in the form of message-passing.\n\nDADO outperforms standard EDA on synthetic functions as well as a multi-layer perceptron (MLP)-based predictive model of protein property functions.", "questions": "- How long does it take to run DADO and standard EDA?\n- Can DADO still be used for larger alphabet sizes or sequence lengths?\n- How would DADO change if positions did not share the same alphabet?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "- DADO is well-motivated and extensively derived\n- DADO clearly converges faster than standard EDA\n- DADO is a novel, more efficient algorithm than fills a gap in the literature", "weaknesses": "- This paper is hard to read due to large amounts of text, in-line math, and few subsections.\n- The evaluation is limited to three synthetic functions and four learned protein property functions\n- Standard EDA that is unaware of function decomposability can outperform DADO in the absence of ad hoc hyperparameter tuning.\n- The GB1 evaluation appears prematurely ended, as the EDA does not appear to converge by the final training iteration\n\nTypo:263 shaing (shaping)", "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": "8EkTwQeLBh", "reviewer": "Reviewer_59py", "summary": "This paper presents an algorithm (DADO) for learning a distribution over a discrete design space to maximize an expected reward, where that reward function can be expressed as a kind of generalized additive model. The authors show favorable performance relative to classical EDA in both synthetic experiments as well as on a task involving the optimization of the score assigned by a protein property predictor developed with the required structure.", "questions": "How restrictive is the assumed compositionality? From the paper's definition \"consider the form f(x) = C1(˜x1) + C2(˜x2), . . . , Cκ(˜xκ), where Ci denotes an arbitrary function on a set of design variables, x˜i,\" it seems rather general, but the experiments focus on a fairly narrow set of circumstances (i.e., shallow junction trees), so it would be helpful to understand what this definition covers. For example, an arbitrary neural network that takes x as input, transforms it through layers into some d-dimensional embedding, which are then transformed through a linear layer to a scalar output would seem to apply to this definition (in this case, the Ci functions are the individual embedding dimensions which are functions of the x's times the associated weight from the linear layer). Which would indicate to me that this can be applied arbitrarily to neural network predictors, although the experiments suggest this is not in fact the case. Can you clarify?\n\nHow well does the method scale to larger sequence lengths compared to those considered in the paper (L < 100)? What about to larger junction trees? Are there important qualitative differences in predictive performance for the underlying property predictors as sequence length increases?\n\nHow robust is DADO in settings where the assumed decomposition is incorrect or incomplete?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The authors are focused on an important and challenging problem in scientific design, as the design space in such settings tends to be discrete or combinatorial and therefore is challenging to explore or optimize over. Their focus on reducing the combinatorial complexity of this search problem to one that is significantly more manageable is a worthwhile pursuit and has real applications. The authors motivated the problem well and provided strong justification for an approach such as theirs for more efficiently optimizing the design space. Further, the generalized additive model structure induced by a junction tree that defines the relationship between the design variables and the property being optimized is a unique choice and lends itself to the paper's originality. Improved efficiency for exploring large combinatorial spaces is of great significance to the field, and the authors have proposed a method that performs capably on the considered problems.", "weaknesses": "The requirement that the decompositional form be known is a very strong one in practice. Indeed, in the protein property prediction experiments, the authors resorted to a very specifically designed predictor to make use of this decomposition--a design which limits the accuracy of the proposed property predictor. Hence, it is unclear to me whether this method has much utility to the community. It would be helpful if the authors could better profile the impact of the decomposition on the accuracy of the underlying property predictor.\n\nAdditionally, I found the the experiments to be rather limited. There is one set of experiments on small synthetically curated systems with shallow junction trees, and another on a protein property prediction task using the same model architecture for each of the four considered proteins. It is difficult from these limited experiments to build good intuition on the method's ability to scale to problem size and complexity. It would be helpful if the authors could better profile or characterize the behavior of their method on a more representative set of problems involving the optimization over a discrete design space. It would also be helpful to establish performance against other baselines for such problems, not just a vanilla EDA.", "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": "wuFjOghCaY", "reviewer": "Reviewer_6Cz8", "summary": "The paper addresses a critical bottleneck in AI-driven scientific design—scaling distributional optimization to high-dimensional discrete spaces (e.g., proteins, circuits)—by exploiting function decomposability. Its contributions are well-aligned with the needs of both machine learning (ML) for optimization and domain sciences (e.g., protein engineering), making it a valuable addition to the field.", "questions": "1. For domains without 3D structural data (e.g., novel peptides or synthetic materials), what alternative methods would you recommend to derive decomposability for DADO? For example, could unsupervised learning (e.g., clustering design variables by co-occurrence) be used to \n2. If you intentionally introduce errors into the junction tree (e.g., remove 10% of true residue contacts for TDP43), how much does DADO’s performance degrade relative to using the correct tree? Are there any heuristics (e.g., adding \"redundant\" edges to the junction tree) to mitigate this error?\n3. For D=50 (e.g., small molecules with 50 building blocks), would DADO’s current MLP-based search distribution require prohibitive compute? If so, what modifications would you propose to scale DADO to larger alphabets?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Unlike standard Estimation of Distribution Algorithms (EDAs) or reinforcement learning (RL) policy optimization, DADO explicitly leverages the decomposability of objective functions (via junction trees) to avoid optimizing over intractable full combinatorial spaces. This is a departure from \"black-box\" optimization approaches that treat the objective as monolithic.\n 2. It generalizes classical max-plus message passing (used for exact global optimization) to distributional optimization, replacing hard maximization with soft, sample-based expectations. This enables DADO to balance exploration (via a factorized generative model) and exploitation (via coordinated message passing)—a key innovation not seen in prior work like Factorized Distribution Algorithms (FDAs) or chain-structured RL policies.\n3. The paper also explores a practical tradeoff: tuning decomposability (via residue distance thresholds in proteins) to balance predictive model accuracy and optimization efficiency. This empirical insight bridges theoretical algorithm design and real-world scientific constraints.", "weaknesses": "1. The paper relies heavily on structure-based decomposability (AlphaFold3 contact graphs) for proteins but does not explore other practical sources of decomposability: For example, in protein design, decomposability could also be derived from sequence homology (conserved vs. variable regions) or functional annotations (binding sites vs. structural loops). Similarly, in circuit design, decomposability might come from modular components.\n2. The paper empirically shows that \"loose\" decomposability (e.g., t=2.75 Å for GB1) retains predictive accuracy, but lacks theoretical bounds on how much decomposition error DADO can tolerate. \n3. The paper uses D=20 (amino acids) for all experiments but does not address scalability to larger discrete alphabets (e.g., D=100 for small molecules or circuit components)", "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": "Hhw2JRGcwY", "reviewer": "Reviewer_dP4b", "summary": "This paper introduces Decomposition-Aware Distributional Optimization (DADO) as a new algorithm for optimizing discrete functions that exhibit decomposable structures. The method learns a decomposed functional form of an objective and then incorporates a factorized search and message-passing schema to coordinate updates. The method is evaluated against Estimation of Distribution Algorithm (EDA) on both synthetic and real protein design tasks.", "questions": "6. Is it ever possible for the functional form to enforce the wrong prior over the input space (for example, if the training data is messy, noisy, affected by a confounding variable, or otherwise unreliable)? If so, what happens to the performance of DADO in these situations?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- To the best of my knowledge, the core idea of applying message passing for distributional optimization is novel and an interesting approach to the problem of scientific design.\n- I appreciate the statistical tests reported in the manuscript.", "weaknesses": "At a high level, my main concerns/questions stem from the learned decomposition of the objective function. While there is evidence that DADO works if a good function decomposition exists, it is not clear that (1) such a good function decomposition always exists; and (2) that the method outperforms black box optimization methods that do not even take advantage of a function decomposition.\n\n1. In the paragraph starting from line 100, the manuscript describes a method to \"guess\" the functional form of the predictive model and then fitting a model that enforces the guessed decomposition to the available data. However, this might result in a potential distribution shift between the fitted model and the true objective, which may or may not adhere to the \"guessed\" functional form. This is of particular concern in situations where the \"high scoring\" candidate designs are rare in the training data and might obey a different underlying functional form compared to other designs. It would be nice to better understand how much this possible distribution shift is actually a concern.\n2. Somewhat related to the above comment, but it would be good to ablate the size of the training dataset used to learn the MLP-based function decomposition, instead of using all the data available (line 451). In particular, I feel it would be important to investigate the performance of DADO if only the bottom $x$th percentile of designs were used to learn the surrogate function.\n3. The manuscript seems to specifically focus on applications of DADO for protein design (in addition to the synthetic function testing). However, it seems from the title, abstract, introduction, and other parts of the text that DADO is proposed as a general method for many different possible discrete design problems in AI4Science. In this light, it would be good to include additional evaluations across different domains (e.g., circuit, molecule, and material design, for instance) to better illustrate how DADO actually generalizes across different scientific domains. \n4. It seems like a number of baselines are missing from experimental evaluation - in particular, how does DADO compare with methods that do not involve learning a junction tree decomposition at all? This would include any black-box optimization method - PPO, BO, FDA, MCTS (some of which are non-distributional but can still be evaluated using the distributional optimization framework). This would help better clarify the added benefit of even learning a functional decomposition in the first place.\n5. It would be good to include results on ablating $K$ in step 1 (line 154).", "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 (Poster)