Further, by that $|\widetilde{\theta}^{m}_{j}-\theta_{j}|<G^{m}_{j}$ (see (i))
and

$\displaystyle
2G^{m}_{j}+G^{*}_{j}=\frac{11}{4}G^{m}_{j}+\frac{1}{4}G^{M}_{j}\leq\frac{11}{4}G_{(j)}+\frac{1}{4}G^{M}_{j}\leq\frac{41}{48}\min(\theta_{j+1}-\theta_{j},\theta_{j}-\theta_{j-1}),$

we have

$\displaystyle\\{\widetilde{\theta}^{m}_{j}-G^{m}_{j}-G^{*}_{j}+1,\ldots,\widetilde{\theta}^{m}_{j}-G^{m}_{j}\\}\cap\\{\widetilde{\theta}^{m}_{j}+G^{m}_{j}+1,\ldots,\widetilde{\theta}^{m}_{j}+G^{m}_{j}+G^{*}_{j}\\}\cap\Theta=\emptyset.$
(C.26)

From (C.25) and Assumption 5′ (b), we have

$\displaystyle\delta_{j}^{2}G_{j}^{*}\geq
C_{1}\max\left\\{\omega^{-2}\mathfrak{s}\rho_{n,p}^{2},\left(\omega^{-1}\mathfrak{s}\log(p)\right)^{1/(1-\tau)}\right\\}$

and from (C.26) and Lemma C.2, we have $\bm{\Delta}^{\text{\tiny
L}}_{j}=\widehat{\bm{\beta}}^{\text{\tiny L}}_{j}-\bm{\beta}_{j-1}$ and
$\bm{\Delta}^{\text{\tiny R}}_{j}=\widehat{\bm{\beta}}^{\text{\tiny
R}}_{j}-\bm{\beta}_{j}$ satisfy

$\displaystyle\max\left(\left|\bm{\Delta}^{\text{\tiny
L}}_{j}\right|_{2},\left|\bm{\Delta}^{\text{\tiny
R}}_{j}\right|_{2}\right)\leq\frac{12\sqrt{2\mathfrak{s}}\lambda}{\omega\sqrt{G_{j}^{*}}}\leq\frac{24\sqrt{2\mathfrak{s}}\lambda}{\omega\sqrt{G_{(j)}}},$
$\displaystyle\left|\bm{\Delta}^{\text{\tiny
L}}_{j}(\mathcal{S}_{j-1}^{c})\right|_{1}\leq 3\left|\bm{\Delta}^{\text{\tiny
L}}_{j}(\mathcal{S}_{j-1})\right|_{1}\text{ and
}\left|\bm{\Delta}^{\text{\tiny R}}_{j}(\mathcal{S}_{j}^{c})\right|_{1}\leq
3\left|\bm{\Delta}^{\text{\tiny R}}_{j}(\mathcal{S}_{j})\right|_{1},$

such that the arguments analogous to those employed in the proof of Theorem 1
(ii) are applicable to establish the localisation rate of
$\check{\theta}_{j}$, which completes the proof.

##  Appendix D Algorithms

input : Bandwidth $G$, grid resolution $r$, penalty $\lambda$, threshold $D$,
$\eta\in(0,1]$

initialise : $\widetilde{\Theta}=\emptyset$, $\widehat{\Theta}=\emptyset$

// Stage 1

Compute $T_{k}(G)$ in (6) for all $k\in\mathcal{T}=\mathcal{T}(r,G)$

Add all $\widetilde{\theta}$ satisfying $T_{\widetilde{\theta}}(G)>D$ and
$\widetilde{\theta}={\arg\min}_{k\in\\{\widetilde{\theta}-\lfloor\eta
G\rfloor+1,\ldots,\widetilde{\theta}+\lfloor\eta
G\rfloor\\}\cap\mathcal{T}}T_{k}(G)$ to $\widetilde{\Theta}$, and set
$\widetilde{\Theta}=\\{\widetilde{\theta}_{j},\,1\leq j\leq\widehat{q}\\}$

// Stage 2

for  _$j=1,\ldots,\widehat{q}$_ do

Identify $\widehat{\theta}_{j}={\arg\min}_{\widetilde{\theta}_{j}-G+1\leq
j\leq\widetilde{\theta}_{j}+G}Q(k;\widetilde{\theta}_{j}-G,\widetilde{\theta}_{j}+G,\widehat{\bm{\beta}}^{\text{\tiny
L}}_{j},\widehat{\bm{\beta}}^{\text{\tiny R}}_{j})$ with
$\widehat{\bm{\beta}}^{\text{\tiny L}}_{j}$ and
$\widehat{\bm{\beta}}^{\text{\tiny R}}_{j}$ computed as in (12), and add it to
$\widehat{\Theta}$

end for

return _$\widehat{\Theta}$ _

Algorithm 1 MOSEG: Single-bandwidth two-stage data segmentation methodology
under a regression model.

input : A set of bandwidths $\mathcal{G}$, grid resolution $r$, penalty
$\lambda$, threshold $D$, $\eta\in(0,1]$

initialise : $\widetilde{\Theta}^{A}=\emptyset$, $\check{\Theta}=\emptyset$,
$\mathcal{C}_{j}=\emptyset$ for all $j$

// Pre-estimator generation for  _$h=1,\ldots,H$_ do

Initialise $\widetilde{\Theta}(G_{h})=\emptyset$

Compute $T_{k}(G_{h})$ in (6) for all
$k\in\mathcal{T}_{h}=\mathcal{T}(r,G_{h})$

Add all $\widetilde{\theta}$ satisfying $T_{\widetilde{\theta}}(G_{h})>D$ and
$\widetilde{\theta}={\arg\min}_{k\in\mathcal{I}_{\eta}(\widetilde{\theta})\cap\mathcal{T}_{h}}T_{k}(G_{h})$,
to $\widetilde{\Theta}(G_{h})$

end for

// Anchor change point estimator identification Identify all
$\widetilde{\theta}(G)\in\cup_{h=1}^{H}\widetilde{\Theta}(G_{h})$ satisfying
(18), and add all such estimators to $\widetilde{\Theta}^{A}$, which is
denoted by $\widetilde{\Theta}^{A}=\\{\widetilde{\theta}^{A}_{j},\,1\leq
j\leq\widehat{q}:\,\widetilde{\theta}^{A}_{1}<\ldots<\widetilde{\theta}^{A}_{\widehat{q}}\\}$

for  _$j=1,\ldots,\widehat{q}$_ do

// Pre-estimator clustering

Identify all $\widetilde{\theta}\in\cup_{h=1}^{H}\widetilde{\Theta}(G_{h})$
satisfying (19) and add it to $\mathcal{C}_{j}$

// Location refinement Add $\check{\theta}_{j}$ obtained as in (20) to
$\widehat{\Theta}$

end for

return _$\check{\Theta}$_

Algorithm 2 MOSEG.MS: Multiscale extension of MOSEG.

##  Appendix E Further information on the real datasets

Table E.1 lists the covariates included in the dataset analysed in Section 5.

Table E.1: Covariates contained in the equity premium dataset analysed in
Section 5 (cf. Koo et al., (2020), Table 3)

Name | Description
---|---
d/p | Dividend price ratio: difference between the log of dividends and the log of prices
d/y | Dividend yield: difference between the log of dividends and the log of lagged prices
e/p | Earnings price ratio: difference between the log of earnings and the log of prices
d/e | Dividend payout ratio: difference between the log of dividends and the log of earnings
b/m | Book-to-market ratio: ratio of book value to market value for the Dow Jones Industrial Average
ntis | Net equity expansion: ratio of 12-month moving sums of net issues by NYSE listed stocks over
| the total end-of-year market capitalization of NYSE stocks
tbl | Treasury bill rates: 3-month Treasury bill rates
lty | Long-term yield: long-term government bond yield
tms | Term spread: difference between the long term bond yield and the Treasury bill rate
dfy | Default yield spread: difference between Moody’s BAA and AAA-rated corporate bond yields
dfr | Default return spread: difference between the returns of long-term corporate and government bonds
svar | Log of stock varianceobtained as the sum of squared daily returns on S&P500 index
infl | Inflation: CPI inflation for all urban consumers
ltr | Long-term return: return of long term government bonds


Functional Graphical Models: Structure Enables
Offline Data-Driven Optimization

Anonymous Author

Anonymous Institution

###### Abstract

While machine learning models are typically trained to solve prediction
problems, we might often want to use them for optimization problems. For
example, given a dataset of proteins and their corresponding fluorescence
levels, we might want to optimize for a new protein with the highest possible
fluorescence. This data-driven optimization (DDO) problem presents additional
challenges beyond the standard prediction problem, with the main challenge
being insufficient data coverage in the optimal regions of the input space
that in turn amplifies the regret rates at a fast pace. Indeed, it is not
clear theoretically when existing approaches can even perform better than a
naïve method that simply selects the best input in the dataset. In this paper,
we study how structure can enable sample-efficient data-driven optimization.
To formalize this notion of structure, we introduce functional graphical
models (FGMs) and show theoretically how they can provide for principled data-
driven optimization by decomposing the original high-dimensional optimization
problem into smaller sub-problems. This allows us to reduce regret bounds for
DDO significantly. Moreover, it has been demonstrated that DDO with FGMs can
achieve nearly optimal designs in situations where the naïve approach fails
due to insufficient coverage of the offline data. Subsequently, we present a
data-driven optimization algorithm that leverages FGM factorization while
learning FGMs, either over the original input variables or a latent variable
representation of the inputs.

##  1 Introduction

While machine learning models are typically trained to solve prediction
problems, we might often want to use them for optimization problems. For
instance, in bioengineering, given a dataset containing examples of protein
and their corresponding fluorescence levels, an engineer might aim to design
proteins with higher fluorescence [brookes2019conditioning,
trabucco2022design]. Similarly, when provided with examples of hardware
accelerators and their performance, a designer might want to infer a more
performant hard accelerator design [kao2020confuciux, kumar2021data]. A direct
machine-learning approach to such problems would be to train a surrogate model
that predicts the performance for a given design, and then find its maximizers
with some optimization method [trabucco2021conservative]. Importantly, here,
our focus is on a challenging yet practical scenario where we can use offline
training data to train a surrogate model but cannot collect data online. This
offline data-driven optimization (DDO) approach is often applied in model-
based optimization [hoburg2014geometric, gomez2018automatic,
trabucco2021conservative, MBO] and offline reinforcement learning
[levine2020offline, fujimoto2021minimalist, RL].

The primary challenge in DDO is insufficient data coverage of optimal
examples, which causes the distribution shift in the optimization process:
when we optimize the design (i.e., the input to the model), we produce a
design very different from the training distribution. Predictions for such
designs are likely to be inaccurate, and an optimizer might exploit the model
to generate designs that result in the largest (optimistic) errors.
Consequently, recent works in MBO and offline RL have focused on incorporating
penalties for exploiting data outside of offline data when constructing a
surrogate model [kumar2020conservative, trabucco2021conservative]. However,
even with these conservative approaches, it remains unclear whether we can
learn the optimal design when the offline training data fails to cover it. The
question in this context is: What enables data-driven optimization to be
feasible in this challenging offline scenario?

To make sample-efficient DDO feasible, a specific structure would be required.
Without any structural bias for the surrogate model, even with the
aforementioned conservative ideas, a naïve approach in the worst case reduces
to choosing the best design observed in the offline data, as we will formalize
in Section 3. As the design space becomes large, the probability of an optimal
design being actually present in the offline data generally becomes extremely
low.

In our work, we introduce the framework of _Functional Graphical Models
(FGMs)_ as a structure that facilitates sample-efficient DDO. FGMs are
designed to describe function’s independence properties, allowing us to
decompose every function into subfunctions over partially intersecting
cliques. Importantly, FGMs provide an effective way to introduce structural
bias for DDO by automatically decomposing the whole optimization problem into
smaller subproblems. We demonstrate that DDO with FGMs can mitigate
distribution shift, as it only requires the optimal design to be covered by
the offline data in a more lenient manner, compared to DDO without any
structural bias. Specifically, FGMs enable the learning of the optimal design
as long as the variables in each clique take on optimal values for some data
point – it is not required to see any single point with optimal values for
_all_ cliques. As a result, we can formally demonstrate that DDO with FGMs
surpasses the naïve approach.

Our contributions can be summarized as follows. First, we introduce FGMs and
demonstrate their ability to reduce regret in data-driven optimization
significantly. In particular, our results imply that we can learn a high-
quality design as long as it is covered by the offline data within each
clique, rather than the entire space. Secondly, we propose a practical data-
driven optimization algorithm that can discover FGMs from the offline data,
under Gaussian assumptions, or from a learned latent space in a more general
setting. We validate the effectiveness of our algorithm, both independently
and in combination with DDO, through numerical examples that we have designed
and in the context of the SuperConductor task from Design-Bench
[trabucco2022design].

##  2 Related Work

Data-driven optimization has connections to several fields, including Bayesian
optimization and model-based optimization, as summarized below.

#### Bayesian Optimization.

The idea of using data and machine learning tools for design and optimization
has been one of the primary motivations of Bayesian optimization (BO)
[brochu2010tutorial, shahriari2015taking]. In BO, one establishes a prior
belief of the considered function and updates it based on the given dataset,
as well as queries for function evaluations that balance maximization and
exploration of the function. However, such an online approach is not
applicable to problems where additional queries cannot be made or are very
cost- and time-consuming (e.g., drug design) that would ideally be solved
fully offline. Nevertheless, we consider our work to be of interest to the BO
community which often makes an assumption of a decomposition of the objective
over a dependency graph, which simplifies the problem and allows to solve it
with fewer queries [kandasamy2015high, rolland2018high,
hoang2018decentralized, ziomek2023random]. Meanwhile, in Section 4, we prove
that such a decomposition exists for every function satisfying mild
conditions, easing tractability of problems with finite amounts of static
data.

#### Probabilistic Graphical Models.

Learning FGMs involves uncovering the structure of functions that map inputs
$x$ to outputs $y$. Similarly, the challenge of discovering structural graphs
has been extensively studied in the field of probabilistic graphical models
[studeny2006probabilistic, pearl2009causality, PGMs]. FGMs and PGMs differ
because PGMs are built upon conditional independence, while FGMs are
constructed based on the concept of _functional independence_ , as we will
explain later in Section 4. Additionally, in the literature on PGMs, the main
focus is on uncovering statistical relationships between variables, and PGMs
are technically used by limiting the capacity of the function space. In
contrast, our emphasis is on DDO, and FGMs are technically used to ensure low
coverage terms, as we will formalize in Section 5.

#### Model-Based Optimization.

In recent years, there has been a growing interest in MBO, driven by the
increasing demand for fully offline design of chemical molecules, proteins,
hardware accelerators, and more [gomez2018automatic, brookes2019conditioning,
kumar2020model, kumar2021data]. MBO and DDO are often used interchangeably.
Most of these methods adopt a “conservative” (i.e., pessimistic) approach to
MBO, employing regularizers to prevent the generation of out-of-distribution
designs [trabucco2021conservative, qi2022data]. Combined with appropriate
inductive bias, these methods found application in practical engineering
problems, such as hardware design [kumar2021data]. In our paper, instead of
designing conservative methods, from a different angle, we introduce a
theoretical framework of FGMs and study the tractability of MBO through its
lenses.

##  3 Preliminaries

We are given an offline dataset
$\mathcal{D}=\\{{\mathbf{x}}^{(i)},{\textnormal{y}}^{(i)}\\}_{i=1}^{N}$, where
the input ${\mathbf{x}}\in\mathcal{X}$ is drawn from a data distribution
$p(\cdot)\in\Delta(\mathcal{X})$, and
${\textnormal{y}}=f({\mathbf{x}})\in\mathbb{R}$ is the corresponding
evaluation of an unknown (black-box) function $f$ at ${\mathbf{x}}$. Our goal
is to find an input (i.e., design) ${\mathbf{x}}$ that takes a high value in
terms of $f({\mathbf{x}})$. In the ideal case, we may want to find an exact
${\mathbf{x}}^{\star}\triangleq\operatorname*{arg\,max}_{\bar{\mathbf{x}}\in\mathcal{X}}f(\bar{\mathbf{x}})$.
A more practical goal might be to obtain a potentially stochastic distribution
$\pi\in\Delta(\mathcal{X})$ that maximizes $J(\pi)$ where $J(\pi)$ is the
value defined by $J(\pi)=\mathbb{E}_{{\mathbf{x}}\sim\pi}[f({\mathbf{x}})]$.

In DDO (a.k.a. MBO), conventional methods typically involve the following
steps: (a) obtaining a surrogate model $\hat{f}:\mathcal{X}\to\mathbb{R}$
using function approximation techniques like deep neural networks, and (b)
optimizing the surrogate model with respect to $x\in\mathcal{X}$. This is
formalized in the following procedure:

$\displaystyle\hat{f}$
$\displaystyle=\operatorname*{arg\,min}_{\bar{f}\in\mathcal{F}}\mathbb{E}_{({\mathbf{x}},y)\sim\mathcal{D}}[\\{\bar{f}({\mathbf{x}})-y\\}^{2}],$
(Step (a)) $\displaystyle\hat{\pi}$
$\displaystyle=\operatorname*{arg\,max}_{\pi\in\Pi}\mathbb{E}_{{\mathbf{x}}\sim\pi}[\hat{f}({\mathbf{x}})],$
(Step (b))

where $\Pi$ consists of probability distributions over $\mathcal{X}$, and
$\mathcal{F}$ is a regression class that consists of functions mapping
$\mathcal{X}$ to $\mathbb{R}$. When $\Pi$ is a fully expressive class111When
$\mathcal{X}$ is discrete, we can always use this fully expressive class.
However, $\mathcal{X}$ is continuous, as we will see in Section B.3, the
coverage term could be infinite when $\Pi$ consists of delta distributions. In
practice, for example, we might use
$\Pi=\\{\mathcal{N}({\mathbf{x}},\sigma^{2});{\mathbf{x}}\in\mathcal{X}$} for
small $\sigma^{2}$ in the continuous setting., i.e.,
$\Pi=\Delta(\mathcal{X})$, Step (b) reduces to

$\textstyle\hat{\mathbf{x}}=\operatorname*{arg\,max}_{{\mathbf{x}}\in\mathcal{X}}\hat{f}({\mathbf{x}}).$

The primary practical challenge in the fully offline DDO setting (i.e., MBO)
lies in dealing with distribution shifts. To see this, consider the naïve
approach: we can set $\hat{f}_{\mathrm{sim}}({\mathbf{x}})$ to assign the
observed $f({\mathbf{x}})$ values and assign $-\infty$ outside of the offline
data. This approach amounts to constructing a surrogate model
$\hat{f}_{\mathrm{sim}}$ in the initial Step (a) without relying on any
structural bias. In the discrete case, it is equivalent to setting
222Technically, we further need to add L2 regularization
$\lambda\|\theta\|^{2}$ for objective functions as pessimism.
$\mathcal{F}=\\{{\mathbf{x}}\mapsto\theta^{\top}\phi({\mathbf{x}});\theta\in\mathbb{R}^{|\mathcal{X}|}\\}$
($\phi({\mathbf{x}})$ is a $|\mathcal{X}|$-dimensional one-hot encoding
vector) in Step (a). Then, it leads to selecting the following input:

$\displaystyle\hat{\mathbf{x}}_{\mathrm{sim}}=\operatorname*{arg\,max}_{{\mathbf{x}}\in\mathcal{X}}\hat{f}_{\mathrm{sim}}({\mathbf{x}})=\operatorname*{arg\,max}_{{\mathbf{x}}\in\mathcal{D}}\hat{f}_{\mathrm{sim}}({\mathbf{x}}).$
(1)

While in this scenario, we can only trust estimates of $\theta_{x}$ for
$x\in\mathcal{D}$, this approach simply chooses a design for which
$\theta_{x}$ is high, regardless of how well the whole dataset explains this
value. Hence, its performance is expected to be poor when the data
distribution $p(\cdot)$ does not cover regions with high output values in
terms of $f(\cdot)$. Hence, this method clearly encounters distribution shift,
particularly when the space $\mathcal{X}$ is large and
$p({\mathbf{x}}^{\star})$ could tend to be low.

To tackle this problem 333A discerning reader might consider addressing
distribution shift by applying softer penalties outside of the offline data
through function approximation. While this approach is beneficial and
recognized as a conservative or pessimistic approach, it may not suffice to
mitigate distribution shift, as we will elaborate in Section 5.3., our work
explores reasonable structural bias that can be incorporated into function
classes $\mathcal{F}$ when constructing surrogate models, can mitigate
distribution shift, and can be learned from the offline data.

#### Notation.

We denote by $\mathcal{V}$ the index set of variables that the input
describes, and $\mathbb{P}(\mathcal{V})$ to denote its powerset (i.e., the set
of all subsets of $\mathcal{V}$). If $\mathcal{X}\subseteq\mathbb{R}^{d}$,
then $\mathcal{V}=[d]\triangleq\\{1,\dots,d\\}$. Writing
${\mathbf{x}}_{\mathcal{V}}$ and ${\mathbf{x}}$ is equivalent. For a subset
$\mathcal{S}\subseteq\mathcal{V}$, we denote ${\mathbf{x}}_{\mathcal{S}}$ as
the vector of variables with indexes in $\mathcal{S}$ and
$\mathcal{X}_{\mathcal{S}}$ as its domain. We denote the rest of the vector in
${\mathbf{x}}$ by ${\mathbf{x}}_{-\mathcal{S}}$. Furthermore, $a\lesssim b$
means inequality up to a problem-dependent constant. The notation
$\mathcal{N}(\mu,\sigma^{2})$ means a normal distribution with mean $\mu$ and
variance $\sigma^{2}$.

##  4 Functional Graphical Models

We present the concept of _functional independence_ , which characterizes the
structural relationship between inputs and outputs, along with its equivalent
definition. Following that, we introduce their graphical representation,
referred to as _Functional Graphical Models (FGMs)_. Then, we will show how to
apply these concepts to DDO.

###  4.1 Functional Independence

We define the concept of functional independence, which decomposes the
structures of functions mapping from $\mathcal{X}$ to real numbers. This
concept can be later used as a useful inductive bias when constructing
surrogate functions for DDO.

######  Definition 1 (Functional Independence).

Let $A,B,S\subseteq\mathcal{V}$. We say that ${\mathbf{x}}_{A}$ and
${\mathbf{x}}_{B}$ are functionally independent given ${\mathbf{x}}_{S}$ if,
there exist functions $f_{-(B\setminus S)}$ and $f_{-(A\setminus S)}$, such
that for every ${\mathbf{x}}\in\mathcal{X}$,

$\displaystyle f({\mathbf{x}})=f_{-(A\setminus
S)}({\mathbf{x}}_{\mathcal{V}\setminus(A\setminus S)})+f_{-(B\setminus
S)}({\mathbf{x}}_{\mathcal{V}\setminus(B\setminus S)}).$

If $S=\\{\emptyset\\}$ then the independence is absolute.

Intuitively, functional independence of ${\mathbf{x}}_{A}$ and
${\mathbf{x}}_{B}$ given ${\mathbf{x}}_{S}$ states that, once we fix the value
of ${\mathbf{x}}_{S}$, the two sets of variables can be optimized
independently. This is a desirable property for DDO because each of the lower-
dimensional variables will be covered by the offline data better than the
joint variable, providing such independent optimizers with more information.
We will formalize this notion in Section 5.

Our first result offers a few alternative criteria when $\mathcal{X}$ is
contained in an Euclidian space. These properties are helpful for discovering
structures in unknown functions, as we will discuss in Section 6.

######  Lemma 1.

Suppose $f({\mathbf{x}})$ is twice-continuously differentiable w.r.t.
${\mathbf{x}}$. Let $A,B,S\subseteq\mathcal{V}$. Then, the following
statements are equivalent.

  1. 1.

${\mathbf{x}}_{A}$ and ${\mathbf{x}}_{B}$ are functionally independent given
${\mathbf{x}}_{S}$.

  2. 2.

For every ${\mathbf{x}}\in\mathcal{X}$,

$\displaystyle\frac{\partial f}{\partial{\mathbf{x}}_{A\setminus
S}}({\mathbf{x}})=F_{-(B\setminus
S)}({\mathbf{x}}_{\mathcal{V}\setminus(B\setminus S)}),$

for some function $F_{-(B\setminus
S)}({\mathbf{x}}_{\mathcal{V}\setminus(B\setminus S)})$ that does not take
${\mathbf{x}}_{(B\setminus S)}$ as input.

  3. 3.

For every ${\mathbf{x}}\in\mathcal{X}$,

$\displaystyle\frac{\partial^{2}f}{\partial{\mathbf{x}}_{A\setminus
S}\partial{\mathbf{x}}_{B\setminus S}}({\mathbf{x}})=0.$

Detailed proofs are provided in Appendix B.

In the following subsection, we extend functional independence into a
graphical representation of the input variables, enabling the encoding of
structures within high-dimensional functions.

###  4.2 FGMs

We first define functional graphical models (FGMs), a general framework for
working with functional independence. An example is shown in Figure 1.

######  Definition 2 (Functional Graphical Model (FGM)).

A graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ is an FGM of a function
$f({\mathbf{x}})$ if, for any $i,j\in\mathcal{V}$ such that $i\neq j$,
$(i,j)\notin\mathcal{E}$, we have that ${\textnormal{x}}_{i}$ and
${\textnormal{x}}_{j}$ are functionally independent given
${\mathbf{x}}_{\mathcal{V}\setminus\\{i,j\\}}$. The graph $\mathcal{G}$ with
the smallest set $\mathcal{E}$ is referred to as the minimal FGM of $f$.

Our construction of FGMs is analogous to that of probabilistic graphical
models (PGMs), where the absence of an edge indicates probabilistic
conditional independence between two random variables
[koller2009probabilistic]. As we prove in Appendix B.1, a result similar to
the Hammersley-Clifford theorem in PGMs [clifford1971markov] holds for FGMs,
decomposing arbitrary functions as follows.

######  Theorem 1 (Function Decomposition).

Assumue that $\int_{\mathcal{X}}\exp f({\mathbf{x}})d{\mathbf{x}}$ exists and
let $\mathcal{G}$ be any FGM of $f({\mathbf{x}})$ and $\mathcal{C}$ be its set
of cliques. Then, there exist functions
$\\{f_{C}({\mathbf{x}}_{C}),C\in\mathcal{C}\\}$, such that

$\displaystyle f({\mathbf{x}})=\sum_{C\in\mathcal{C}}f_{C}({\mathbf{x}}_{C}).$

Figure 1:  The clique set of this graph is given by
$\\{\\{0,1,2\\},\\{2,3,4\\},\\{4,5,6\\}\\}$. Hence, if it is a FGM of
$f({\mathbf{x}})$, then
$f({\mathbf{x}})=f_{0,1,2}({\mathbf{x}}_{0,1,2})+f_{2,3,4}({\mathbf{x}}_{2,3,4})+f_{4,5,6}({\mathbf{x}}_{4,5,6})$.

###  4.3 DDO via FGMs

Now, consider a DDO with FGMs. By introducing a structural bias in the form of
$f({\mathbf{x}})=\sum_{C}f_{C}({\mathbf{x}}_{C})$ to build a surrogate model,
this can be instantiated as follows:

$\displaystyle\hat{\pi}_{\mathrm{FGM}}$
$\displaystyle=\operatorname*{arg\,max}_{\pi\in\Pi}\mathbb{E}_{{\mathbf{x}}\sim\pi}[\hat{f}({\mathbf{x}})],\,\hat{f}({\mathbf{x}})=\sum_{C\in\mathcal{C}}\hat{f}_{C}({\mathbf{x}}_{C}),$
$\displaystyle\\{\hat{f}_{C}\\}_{C\in\mathcal{C}}$
$\displaystyle=\operatorname*{arg\,min}_{\\{\bar{f}_{C}\in\mathcal{F}_{C}\\}_{C\in\mathcal{C}}}\mathbb{E}_{({\mathbf{x}},y)\sim\mathcal{D}}[\\{\sum_{C\in\mathcal{C}}\bar{f}_{C}({\mathbf{x}}_{C})-y\\}^{2}],$

where each $\mathcal{F}_{C}$ is a function class that consists of maps from
$\mathcal{X}_{C}$ to $\mathbb{R}$.

Now, consider an illustrative simple case when $\mathcal{X}$ is discrete. In
this case, natural choices of $\Pi$ and $\mathcal{F}_{C}$ would be fully
expressive classes:
$\Pi=\Delta(\Pi),\mathcal{F}_{C}=\\{{\mathbf{x}}_{C}\mapsto\theta^{\top}_{C}\phi_{C}({\mathbf{x}}_{C});\theta^{\top}_{C}\in\mathbb{R}^{|\mathcal{X}_{C}|}\\}$
where $\phi_{C}(\cdot)$ is a one-hot encoding vector over $\mathcal{X}_{C}$.
Then, the abovementioned method is equivalent to choosing

$\displaystyle\operatorname*{arg\,max}_{{\mathbf{x}}^{\prime}\in\mathcal{X}}$
$\displaystyle\\{\tilde{\phi}^{\top}({\mathbf{x}}^{\prime})\mathbb{E}_{{\mathbf{x}}\sim\mathcal{D}}[\tilde{\phi}({\mathbf{x}})\tilde{\phi}({\mathbf{x}})^{\top}]^{-1}\mathbb{E}_{({\mathbf{x}},y)\sim\mathcal{D}}[\tilde{\phi}({\mathbf{x}})y]\\},$
(2) $\displaystyle\tilde{\phi}({\mathbf{x}})$
$\displaystyle=[\phi({\mathbf{x}}_{1})^{\top},\cdots,\phi({\mathbf{x}}_{C})^{\top}]^{\top}.$

Importantly, this method differs from just choosing an optimal ${\mathbf{x}}$
in the dataset $\mathcal{D}$, i.e., $\hat{\mathbf{x}}_{\mathrm{sim}}$ in Eq. 1
which corresponds to a method without any inductive bias.

In practice, we must still estimate FGMs from the data and handle the
continuous space $\mathcal{X}$. We will delve into the practical
implementation of this in Section 6.

##  5 FGMs Enable Sample-Efficient Data-Driven Optimization

In this section, we present our main result, demonstrating that functions
represented by FGMs can be optimized efficiently with much less stringent data
coverage assumptions than those needed for simple data-driven optimization.
Note our theory in this section still holds when an observation $y$ has a
bounded measurement error $\epsilon$, i.e., $y=f({\mathbf{x}})+\epsilon$.

###  5.1 Regret Guarantees

First, we review the regret guarantee of the naïve approach in Step (a) and
Step (b). A standard regret guarantee in this setting [chen2019information]
states that, with probability $1-\delta$, the following holds:

$\displaystyle\underbrace{J(\pi^{\star})-J(\hat{\pi})}_{\mathrm{Regret\,against}\,\pi^{\star}}\lesssim\underbrace{\max_{\bar{\pi}\in\Pi}\max_{{\mathbf{x}}\in\mathcal{X}}\left|\frac{\bar{\pi}({\mathbf{x}})}{p({\mathbf{x}})}\right|}_{\mathrm{Coverage(a)}}\times\underbrace{\sqrt{\frac{\log(|\mathcal{F}|/\delta)}{n}}}_{\mathrm{Complexity(a^{\prime})}},$
(3)

where $\pi^{\star}$ is an optimal design among $\Pi$, i.e.,
$\pi^{\star}=\operatorname*{arg\,max}_{\pi\in\Pi}J(\pi)$.

The coverage term (a) measures the discrepancy between a distribution in $\Pi$
and a data distribution $p(\cdot)$. The term (a’) signifies the size of the
function class $\mathcal{F}$. Note this $|\mathcal{F}|$ could be easily
replaced with the covering number when $\mathcal{F}$ is infinite
[wainright2019high].

Our main result is the following regret guarantee for DDO with FGMs introduced
in Section 4.3:

######  Theorem 2 (Regret of DDO with FGMs).

For any $C\in\mathcal{C}$, suppose (1) $\mathbb{E}_{{\mathbf{x}}\sim
p}[\bar{f}_{C}({\mathbf{x}}_{C})]=0$ for any $\bar{f}_{C}\in\mathcal{F}_{C}$,
(2) the correlation in each clique is well-controlled, such that

$\displaystyle\max_{\\{\bar{f}_{C}\in\mathcal{F}_{C}\\}}\mathrm{Corr}[\bar{f}_{C_{1}}({\mathbf{x}}_{C_{1}}),\bar{f}_{C_{2}}({\mathbf{x}}_{C_{2}})]\leq\sigma,$
(4)

(3) the models are well-specified, such that $f_{C}\in\mathcal{F}_{C}$, and
(4) $|\bar{f}_{C}({\mathbf{x}}_{C})|\leq 1$ for any
${\mathbf{x}}_{C}\in\mathcal{X}_{C},\bar{f}_{C}\in\mathcal{F}_{C}$. Then, with
probability $1-\delta$, we have

$\displaystyle J(\pi^{\star})-J(\hat{\pi}_{\mathrm{FGM}})\lesssim$
$\displaystyle\underbrace{\sqrt{\frac{1}{1-\sigma}}}_{\mathrm{Corr}(b^{\prime\prime})}\underbrace{\max_{C\in\mathcal{C}}\max_{\bar{\pi}\in\Pi}\max_{{\mathbf{x}}_{C}\in\mathcal{X_{C}}}\left|\frac{\bar{\pi}({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\right|}_{\mathrm{Coverage(b)}}\underbrace{\sqrt{\frac{|\mathcal{C}|\sum_{C}\log(|\mathcal{F}_{C}|/\delta)}{n}}}_{\mathrm{Complexity(b^{\prime})}},$

where $p({\mathbf{x}}_{C}),\bar{\pi}({\mathbf{x}}_{c})$ are marginal
distributions of $p({\mathbf{x}}),\bar{\pi}({\mathbf{x}})$ with respect to
${\mathbf{x}}_{C}$, respectively, with slight abuse of notation.

While the bound of the naïve approach in Eq. 3 implies that MBO without FGM
structure necessitates the data distribution $p(\cdot)$ to cover the optimal
design itself, the FGM factorization only necessitates the data distribution
to cover components within each “clique” in the optimal design 444Technically,
Theorem 2 states that not only to the best distribution $\pi^{\star}$, but all
distributions in the class $\Pi$ need to be covered. As we will mention in
Section 5.3, $\max_{\bar{\pi}\in\Pi}$ can be replaced with $\pi^{\star}$ by
incorporating conservative approaches.. This is a significantly less stringent
requirement, as the best design in the offline data might be much less
effective than the best design achievable by combining favorable settings of
the values in each of the cliques, as described in Figure 2. We will delve
further into the implications of this statement below.

Figure 2: Consider a function $f(x_{1},x_{2})=-(x_{1}-1)^{2}-(x_{2}-2)^{2}$.
Clearly, the singleton cliques $\\{x_{1}\\}$ and $\\{x_{2}\\}$ are
functionally independent, while the data coming from a correlated normal
distribution are not statistically independent. While the dataset does not
jointly cover the optimal solution $x^{\star}=(1,2)$, it does cover individual
components $x^{\star}_{1}=1,x^{\star}_{2}=2$, and does a method that can learn
the component functions can compose them into $x^{\star}$.

#### Coverage term.

The term (b) corresponds to a coverage term that takes into account the FGM
factorization. This term quantifies how a distribution $\pi\in\Pi$ differs
from a data distribution $p$ in terms of _marginal_ distributions with respect
to ${\mathbf{x}}_{C}$. The following lemma demonstrates that term (b) is
smaller than the coverage term (a) across the entire space $\mathcal{X}$:

######  Lemma 2 (Improvment of Coverage terms).

For any $\pi\in\Delta(\mathcal{X})$, we have

$\displaystyle\max_{C\in\mathcal{C}}\max_{{\mathbf{x}}_{C}\in\mathcal{X}_{C}}\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\leq\max_{{\mathbf{x}}\in\mathcal{X}}\frac{\pi({\mathbf{x}})}{p({\mathbf{x}})}.$

Consider the following illustrative example: suppose the space $\mathcal{X}$
is discrete, denoted by
$\mathcal{X}=\mathcal{X}_{1}\otimes\mathcal{X}_{2}\otimes\cdots\otimes\mathcal{X}_{d}$
where $\mathcal{X}_{1}=\mathcal{X}_{2}=\cdots=\mathcal{X}_{d}$. Now, let’s
assume that $f({\mathbf{x}})=\sum_{C}f_{C}({\mathbf{x}}_{C})$, and
$p(x_{1},x_{2},\cdots,x_{d})$ follows a jointly independent uniform
distribution. Term (a) exhibits exponential growth with the dimension $d$:

$\displaystyle\max_{\bar{\pi}\in\Delta(\mathcal{X})}\max_{{\mathbf{x}}\in\mathcal{X}}\frac{\bar{\pi}(x_{1},x_{2},\cdots,x_{d})}{p(x_{1},x_{2},\cdots,x_{d})}=|\mathcal{X}_{1}|^{d}.$
(5)

Therefore, the naïve data-driven optimization procedure suffers from a curse
of dimensionality, requiring an exponentially large number of samples to
obtain a near-optimal design as the dimension increases. On the contrary, the
term (b) remains reasonably small:

$\displaystyle\max_{\bar{\pi}\in\Delta(\mathcal{X})}\max_{C\in\mathcal{C}}\max_{{\mathbf{x}}_{C}\in\mathcal{X}_{C}}\frac{\bar{\pi}({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\leq|\mathcal{X}_{1}|^{\max_{C}d_{C}}$
(6)

where $d_{C}$ represents the dimension of ${\mathbf{x}}_{C}$. Thus, the DDO
with FGMs can effectively overcome the curse of dimensionality when
$\max_{C}d_{C}$ is reasonably small.

Furthermore, in a more extreme example, illustrated in Figure 2, the coverage
term (b) is bounded, but the coverage term (a) is infinite. This illustrates
that DDO with FGMs can identify the optimal design, when DDO without FGMs
cannot do so even with infinite samples. It’s worth noting that in this
scenario, a discerning reader might express concerns about the potential
explosion of the additional term $(b^{\prime\prime})$, which could adversely
affect the regret of DDO with FGMs. However, as we will discuss in Section
5.2, this term $(b^{\prime\prime})$ is moderately constrained.

#### Statistical complexity term.

We can anticipate that $(b^{\prime})$ is generally much smaller than
$(a^{\prime})$ because the function class to obtain a surrogate model is more
restricted. Now, let’s examine a specific scenario where
$\mathcal{X}_{1}=\cdots=\mathcal{X}_{d}$ and each $\mathcal{X}_{i}$ is
discrete. Suppose that $\mathcal{F}$ is a linear combination of all polynomial
basis functions over $[x_{1},\cdots,x_{d}]$, i.e., the fully expressive class.
In this case, a standard argument based on covering [wainright2019high] shows
that term (a’) would be on the order of $O(\sqrt{|\mathcal{X}_{1}|^{d}/n})$
because the number of parameters in the model is
$|\mathcal{X}|=|\mathcal{X}_{1}|^{d}$. Now, let’s consider the term (b’). When
we set each $\mathcal{F}_{C}$ to be a linear combination of all polynomial
basis functions over the entire $\mathcal{X}_{C}$, its scaling becomes
$O(\sqrt{|\mathcal{C}||\mathcal{X}_{1}|^{\max_{C}d_{C}}/n})$ because the
number of parameters in the model is bounded from above by
$|\mathcal{C}||\mathcal{X}_{1}|^{\max_{C}d_{C}}$. Again, this implies that the
DDO with FGMs can effectively overcome the curse of dimensionality when
$\max_{C}d_{C}$ is reasonably small.

#### Summary.

Theorem 2 reveals that the original data-driven optimization problem naturally
breaks down into subproblems for each clique $C\in\mathcal{C}$. More
specifically, this result shows that, while the naïve approach requires the
data distribution to cover the optimal design, with FGMs, it is only necessary
for the data distribution to cover the optimal values of each clique
separately. With non-trivial FGMs, as we see in Eq. 5 and Eq. 6, the
difference in the number of samples needed to obtain a near-optimal design
between DDO with FGMs and naïve DDO (i.e., DDO without FGMs) can be
exponential 555Note the term (b”) is still $1$ in this case. Hence, we can
ignore the term (b”) for the comparison.. Furthermore, although there could be
a potential concern related to the additional term (b”), in certain instances,
DDO with FGMs has the capability to learn the optimal design, whereas the
naive DDO approach cannot achieve this even with an infinite amount of data.
The intuition for this is that naïve DDO requires the samples from the data
distribution to randomly set _all_ of the variables to near-optimal values in
at least some of the training points in the offline data, whereas DDO with
FGMs only requires some training points for each clique to have near-optimal
values, and does not require any single training point to have near-optimal
values for _all_ cliques. In the naïve case, this can lead to catastrophic
sample complexity, essentially reducing the method to selecting the best point
in the offline data (i.e., the naïve approach in Eq. 1), whereas with FGMs and
relatively small cliques, it is possible to recombine the best values for each
clique and find good designs potentially with exponentially fewer samples.

Note that the literature on statistics, in the context of additive models
[hastie1987generalized], has previously explored the reduction of statistical
complexity. However, our innovation lies in characterizing the coverage term
within the context of DDO. It’s important to highlight that simply reducing
statistical complexity in a polynomial manner (i.e., improving term (a’)) is
insufficient for overcoming the curse of dimensionality (i.e., requiring an
exponentially large number of samples in DDO), as the coverage term (a) might
still grow exponentially.

Figure 3:  In this scenario, there are $d$ cliques, and each variable
$\mathcal{X}_{1}$ to $\mathcal{X}_{d}$ is binary, i.e., ${0,1}$. We define
$f(x)$ as $x_{1}x_{2}+x_{2}x_{3}+x_{3}x_{4}+\cdots+x_{d}x_{1}$ and set
$n=1000$. The optimal design is $x_{1}=1,\cdots,x_{d}=1$. We run $50$
simulations and show average performance with $95\%$ confidence intervals. The
regret of $\hat{\pi}_{\mathrm{FGM}}$ (i.e., Eq. 2) is almost always $0$.
However, the regret of $\hat{\pi}$ without FGM (i.e.,
$\hat{\mathbf{x}}_{\mathrm{sim}}$ in Eq. 1) suffers from the curse of
dimensionality.

###  5.2 Assumptions on DDO with FGMs

We make four assumptions. The first and last assumptions are primarily
technical. The first one can always be satisfied with negligible error by
substituting $f_{C}(x^{(i)})$ and $y^{(i)}$ with
$f_{C}({\mathbf{x}}^{(i)})-\mathbb{E}_{\mathcal{D}}[f({\mathbf{x}}_{C})]$ and
$y^{(i)}-\mathbb{E}_{\mathcal{D}}[y]$. The third assumption is standard, and
we can easily account for potential misspecification errors in the theorem.
The second assumption is substantial. When $\sigma=1$, our guarantee is
essentially void. In essence, this assumption implies that each clique
exhibits mild independence, automatically breaking down the entire problem
into subproblems. Technically, this enables us to translate the error of
$\hat{f}=\sum_{C}\hat{f}_{C}$ for the entire clique into the error of
$\hat{f}_{C}$ for each individual clique.

Readers might have concerns about whether the second assumption could
potentially lead to issues, particularly in cases where cliques overlap. In
this regard, we first note that this assumption can be significantly relaxed.
In (4), while the maximum is taken over the entire function class
$\mathcal{F}_{C}$, we can replace it with a much smaller set consistent with
$f(x)$, as we explain in Appendix B.3. Secondly, in many cases, we can still
observe huge improvements in terms of regret. For instance, we conducted a
numerical experiment comparing naïve DDO in Eq. 1 with DDO using FGMs in Eq. 2
when $f(x)=x_{1}x_{2}+x_{2}x_{3}+\cdots+x_{d-1}x_{d}+x_{d}x_{1}$, where
$\forall i;\mathcal{X}_{i}=\\{0,1\\}$ and $p(x_{1},\cdots,x_{d})$ is a jointly
uniform distribution. Figure 3 shows significant improvements, even when every
node is shared in the FGM.

###  5.3 Relation with Pessimism

In our work, to facilitate the comparison between the naïve procedure without
FGMs and the one with FGMs, we did not include penalties for exploiting
outside of offline data (i.e., pessimism) in the algorithms. When we
incorporate pessimism into the algorithms, we can formally show that it
alleviates distribution shift in that we can obtain a regret bound where
$\max_{\bar{\pi}\in\Pi}$ in the coverage term is replaced with just an optimal
design $\pi^{\star}$ [rashidinejad2021bridging, xie2021bellman]. In this
sense, pessimism alleviates the distribution shift. However, in the instance
we show, even if we replace $\max_{\bar{\pi}\in\Pi}$ with an optimal design
$\pi^{\star}$, the coverage term Eq. 5 for the procedure with FGMs is still
$O(|\mathcal{X}_{1}|^{d})$, and the coverage term Eq. 6 for the procedure with
FGMs is still $O(|\mathcal{X}_{1}|^{\max_{C}d_{C}})$. Consequently, pessimism
itself cannot resolve distribution shift when the optimal design is not well-
covered by the data distribution $p(\cdot)$. FGMs relax the notion of the
coverage in terms of each clique but not the whole space $\mathcal{X}$. We can
observe that FGMs and pessimism work together to alleviate distribution shift
in a distinct manner.

###  5.4 Relation with Offline RL.

Readers might wonder about the connection to offline RL. In non-contextual
offline RL, a typical objective is to optimize
$\mathbb{E}_{\pi}[\sum_{t=1}^{H}f_{t}(a_{t})]$, where $H$ is the horizon and
each $a_{t}$ is an action at time $t$. Consequently, this setting can be
represented as an FGM with $H$ cliques $[a_{1},\cdots,a_{H}]$. However, there
are two significant distinctions. First, our decomposition is much more
versatile as it allows for potential overlaps between each clique. Second, in
offline RL, we obtain a signal for each clique, namely $f_{t}(a_{t})$, given
$a_{t}$. On the contrary, in our case, we can only observe the sum
$\sum_{t}f_{t}(a_{t})$.

##  6 Graph Discovery

So far, we’ve shown how functions with known FGMs can be optimized effectively
and under much milder coverage conditions than naïve DDO. However, in many
cases, we do not know the FGM for a given function. Next, we’ll discuss how to
construct FGMs only from offline data practically, provided that the data
satisfies a Gaussian distribution assumption:

######  Assumption 1.

The inputs in the offline data follow the standard Gaussian distribution,
${\mathbf{x}}\sim\mathcal{N}(0,I)$.

This assumption may appear strong. However, whitening to remove correlations
is a common procedure in machine learning, and in cases where the data
distribution has higher order moments, we can use more advanced representation
learning methods, such as variational auto-encoders [kingma2013auto,
rezende2015variational, dai2019diagnosing, VAE]. Indeed, these methods are
designed specifically to acquire latent variable representations of the data
that follow the standard Gaussian distribution.

To derive our graph discovery method, we start with the following lemma, known
as the second-order Stein’s identity [stein2004use, erdogdu2015newton]:

######  Lemma 3.

Let $\bar{{\mathbf{x}}}\in\mathbb{R}^{d_{{\mathbf{x}}}}$ and $i\neq
j\in[d_{{\mathbf{x}}}]$. Then

$\displaystyle\frac{\partial^{2}}{\partial\bar{{\textnormal{x}}}_{j}\bar{{\textnormal{x}}}_{i}}\mathbb{E}_{{\mathbf{x}}\sim
N}\big{[}f(\bar{{\mathbf{x}}}+{\mathbf{x}})\big{]}=\mathbb{E}_{{\mathbf{x}}\sim
N}[{\textnormal{x}}_{i}{\textnormal{x}}_{j}f(\bar{{\mathbf{x}}}+{\mathbf{x}})].$

Below, we demonstrate how this lemma can be used to discover FGMs from data in
combination with the third criterion (in Lemma 1) for functional independence.
Suppose that ${\textnormal{x}}_{i}$ and ${\textnormal{x}}_{j}$ are not linked
in a FGM of $f$, such that $(i,j)\notin\mathcal{E}$. By that criterion,

$\displaystyle\frac{\partial^{2}}{\partial\bar{{\textnormal{x}}}_{j}\partial\bar{{\textnormal{x}}}_{i}}\mathbb{E}_{{\mathbf{x}}\sim
N}\big{[}f(\bar{{\mathbf{x}}}+{\mathbf{x}})\big{]}$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}\sim
N}\Big{[}\frac{\partial^{2}f}{\partial\bar{{\textnormal{x}}}_{j}\partial\bar{{\textnormal{x}}}_{i}}(\bar{{\mathbf{x}}}+{\mathbf{x}})\Big{]}=\mathbb{E}_{{\mathbf{x}}\sim
N}[0]=0.$

Letting $\bar{{\mathbf{x}}}=0_{d_{{\mathbf{x}}}}$, for
$(i,j)\notin\mathcal{E}$, we have $\mathbb{E}_{{\mathbf{x}}\sim
N}\big{[}{\textnormal{x}}_{i}{\textnormal{x}}_{j}f({\mathbf{x}})\big{]}=0.$
666Note that while it’s uncertain whether the reverse claims hold in general,
this graph discovery method remains practical for providing a suitable graph
for our downstream tasks, i.e., DDO.  Hence, to reconstruct edges in the
unknown FGM of $f$, one can take $M$ samples
$\\{{\mathbf{x}}^{(k)},{\textnormal{y}}^{(k)}\\}_{k=1}^{M}$ and estimate

$\displaystyle\hat{H}_{ij}=\frac{1}{M}\sum_{k=1}^{M}{\textnormal{x}}^{(k)}_{i}{\textnormal{x}}^{(k)}_{j}{\textnormal{y}}^{(k)},$
(7)

and infer that an edge is absent if $\hat{H}_{ij}$ is close to zero.

Algorithm 1 Gaussian Graph Discovery & Optimization (GGDO)

1: Input: dataset
$\mathcal{D}=\\{{\mathbf{x}}^{(i)},{\textnormal{y}}^{(i)}\\}_{i=1}^{N}$

2: Approximate the cliques of the FGM of $f({\mathbf{x}})$ with Equation (7)

3: Learn the functions $\hat{f}_{C}({\mathbf{x}}_{C})$ by minimizing the mean
squared error  $\mathbb{E}_{({\mathbf{x}},{\textnormal{y}})\sim
p}\big{[}\big{(}\sum_{C\in\hat{\mathcal{C}}}\hat{f}_{C}({\mathbf{x}}_{C})-{\textnormal{y}}\big{)}^{2}\big{]}$

4:
$\hat{\pi}=\operatorname*{arg\,max}_{\pi}\mathbb{E}_{{\mathbf{x}}\sim\pi}\big{[}\sum_{C\in\hat{\mathcal{C}}}\hat{f}_{C}({\mathbf{x}}_{C})\big{]}$.

##  7 Experiments

In this section, we present numerical experiments to examine how readily
applicable our theoretical findings are in practice. Specifically, once the
FGM clique set $\hat{\mathcal{C}}$ is approximated with Equation (7), we can
learn functions $\\{\hat{f}_{C}({\mathbf{x}}_{C}),C\in\mathcal{C}\\}$ via
neural network regression with standard MSE losses. Here, note that by
estimating standard deviation of
$\\{{\textnormal{x}}_{i}^{(k)}{\textnormal{x}}_{j}^{(k)}{\textnormal{y}}^{(k)}\\}_{k=1}^{K}$,
$\forall i,j\in[d_{{\mathbf{x}}}]$, that we denote as $\hat{\sigma}_{ij}$, we
performed the Gaussian hypothesis test $|\hat{H}_{ij}|\geq
c_{\alpha/2}\hat{\sigma}_{ij}/\sqrt{M}\implies(i,j)\in\mathcal{E}$, where
$c_{\alpha/2}$ is the top-$\alpha/2$ quantile of standard Gaussian
distribution. Then, we find new candidate values $\hat{{\mathbf{x}}}$ for
${\mathbf{x}}^{\star}$ by performing gradient ascent on ${\mathbf{x}}$. For
example, we can introduce a probability distribution $\pi({\mathbf{x}})$,
sample ${\mathbf{x}}\sim\pi$, evaluate
$\hat{f}({\mathbf{x}})=\sum_{C\in\mathcal{C}}\hat{f}_{C}({\mathbf{x}}_{C})$,
and take a gradient step on the parameters of $\pi$. We summarize the entire
DDO algorithm in Algorithm 1.

###  7.1 Syenthetic Experiments

We begin by studying the effectiveness of our algorithms in a simple hand-
crafted example, where we can gauge whether the method recovers the _true_
underlying structure. The data are generated as follows. First, we generate
inputs ${\mathbf{x}}\sim N(0_{d},I_{d})$. For this input, we specify the set
of cliques $\mathcal{C}$. We then construct the function
$f({\mathbf{x}})=\sum_{C\in\mathcal{C}}{\textnormal{w}}_{C}\cdot\exp(-||{\mathbf{x}}_{C}-\mu_{C}||^{2}),$
where the centers $\\{\mu_{C},C\in\hat{\mathcal{C}}\\}$ are sampled from a
Gaussian distribution and the positive weights
$\\{{\textnormal{w}}_{C},C\in\hat{\mathcal{C}}\\}$ are also random. First, we
examine independencies discovered by Equation (7). We apply the algorithm to
problems with $d=4$ and $d=7$ variables, plot the reconstructed graphs, and
measure their graph-editing distance (GED), normalized to the range $[0,1]$.
We then learn the function $\hat{f}$ with the discovered decomposition and
maximize a design distribution $\pi({\mathbf{x}})$, which is initialized at
the maximum of a random batch, against it. We found that our method discovers
the majority of functional independencies, bringing GED low to 0.17 and 0.14
for 4- and 7-dimensional problems, respectively (see Appendix C for
visualization). Consequently, Figure 4 show that Algorithm 1 is able to
largely improve over the maximal value in the batch, in which it was
initialized.

Figure 4: Values of $\hat{f}(\hat{{\mathbf{x}}})$, where
$\hat{{\mathbf{x}}}\sim\pi$, along the course of gradient ascent on $\pi$, for
4- and 7-dimensional problems (x-axis: iterations, y-axis: outputs). The
algorithm improves from the initial value of $2$ to $3.3$, and from $2$ to
$6.5$, in 4D and 7D cases, respectively. Results are averaged over 5 random
seeds.

###  7.2 Semi-Syenthetic Experiments

Unfortunately, in many real-world applications, the data distribution will not
satisfy the Gaussian assumptions. Hence, despite its lack of theoretical
guarantees, we propose a simple approach to tackle such problems with VAEs
[kingma2013auto]. Namely, we first reparameterize the data so that they
satisfy Assumption 1 by fitting a VAE with Gaussian prior. Then, we apply
Algorithm 1 to these learned representations. Note that we do not calculate
$\hat{{\bm{H}}}_{ij}$ from $M$ samples at once—we keep track of exponential
moving averages of such estimators throughout training to decrease the
variance.

We apply this method to the previous $7$-dimensional problem, which we modify
by nonlinearly transforming the ground-truth arguments to the function
$f({\mathbf{x}})$. We also apply it to the Superconductor task for Design-
Bench [trabucco2022design]—where the goal is to optimize the design of a
superconductor, given a dataset of designs. For comparison, we run an off-the-
shelf MBO method, COMs [trabucco2021conservative], in each of these tasks.
Excitingly, Figure 5 shows that our method was capable of handling these two
tasks, while COMs failed to solve the former problem.

Figure 5:  The score of the policy optimized with our FGM- and VAE-based
method (blue) vs COMs (pink) in the 7-D task (left) and Superconductor
(right).

##  8 Conclusion

We showed how DDO can be made feasible and tractable for functions that
exhibit factorization structure. We capture this structure with functional
graphical models (FGMs), which describe how a function can factorize over
cliques analogous to how PGMs factorize probability distributions. We then
showed that functions with FGMs admit sample-efficient optimization without
requiring that the data distribution place high probability directly on
optimal designs, only requiring coverage of optimal values for each clique in
the FGM individually.

##  Appendix A Auxiliary Facts

######  Theorem 3 (Hammersley-Clifford [clifford1971markov]).

Let $\mathcal{G}=(\mathcal{V},\mathcal{E})$ be a simple graph with clique set
$\mathcal{C}$, and let ${\mathbf{x}}={\mathbf{x}}_{\mathcal{V}}\sim p$, where
$p({\mathbf{x}})$ is a strictly positive probability distribution. If $p$ is
Markov with respect to $\mathcal{G}$, then,
$\forall{\mathbf{x}}\in\mathcal{X}$,

$\displaystyle
p({\mathbf{x}})=\prod_{C\in\mathcal{C}}\phi_{C}({\mathbf{x}}_{C}),$

for some functions $\phi_{C}({\mathbf{x}}_{C})>0$.

See 3

###### Proof.

We start from proving a known result. Namely,

$\displaystyle\frac{\partial}{\partial\bar{{\textnormal{x}}}_{i}}\mathbb{E}_{{\mathbf{x}}\sim
N}\big{[}f(\bar{{\mathbf{x}}}+{\mathbf{x}})]=\mathbb{E}_{{\mathbf{x}}\sim
N}[{\textnormal{x}}_{i}f(\bar{{\mathbf{x}}}+{\mathbf{x}})].$ (8)

To do it, we use integration by parts,

$\displaystyle\frac{\partial}{\partial\bar{{\textnormal{x}}}_{i}}\mathbb{E}_{{\mathbf{x}}\sim
N}[f(\bar{{\mathbf{x}}}+{\mathbf{x}})]$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}\sim
N}\Big{[}\frac{\partial}{\partial\bar{{\textnormal{x}}}_{i}}f(\bar{{\mathbf{x}}}+{\mathbf{x}})\Big{]}$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}\sim
N}\Big{[}\frac{\partial}{\partial{\textnormal{x}}_{i}}f(\bar{{\mathbf{x}}}+{\mathbf{x}})\Big{]}$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}_{-i}\sim
N}\Big{[}\int_{{\textnormal{x}}_{i}}d{\textnormal{x}}_{i}\frac{1}{\sqrt{2\pi}}\exp\Big{(}\frac{-1}{2}{\textnormal{x}}^{2}_{i}\Big{)}\frac{\partial}{\partial{\textnormal{x}}_{i}}f(\bar{{\mathbf{x}}}+{\mathbf{x}})\Big{]}$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}_{-i}\sim
N}\Big{[}\Big{[}\frac{1}{\sqrt{2\pi}}\exp\Big{(}\frac{-1}{2}{\textnormal{x}}^{2}_{i}\Big{)}f(\bar{{\mathbf{x}}}+{\mathbf{x}})\Big{]}^{\infty}_{-\infty}+\int_{{\textnormal{x}}_{i}}d{\textnormal{x}}_{i}\frac{{\textnormal{x}}_{i}}{\sqrt{2\pi}}\exp\Big{(}\frac{-1}{2}{\textnormal{x}}^{2}_{i}\Big{)}f(\bar{{\mathbf{x}}}+{\mathbf{x}})\Big{]}$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}_{-i}\sim
N}\Big{[}0+\int_{{\textnormal{x}}_{i}}d{\textnormal{x}}_{i}\frac{{\textnormal{x}}_{i}}{\sqrt{2\pi}}\exp\big{(}\frac{-1}{2}{\textnormal{x}}^{2}_{i}\big{)}f(\bar{{\mathbf{x}}}+{\mathbf{x}})\Big{]}$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}\sim
N}[{\textnormal{x}}_{i}f(\bar{{\mathbf{x}}}+{\mathbf{x}})].$

Now, applying this result twice,

$\displaystyle\frac{\partial^{2}}{\partial\bar{{\textnormal{x}}}_{j}\bar{{\textnormal{x}}}_{i}}\mathbb{E}_{{\mathbf{x}}\sim
N}\big{[}f(\bar{{\mathbf{x}}}+{\mathbf{x}})]$
$\displaystyle=\frac{\partial}{\partial\bar{{\textnormal{x}}}_{j}}\frac{\partial}{\partial\bar{{\textnormal{x}}}_{i}}\mathbb{E}_{{\mathbf{x}}\sim
N}[f(\bar{{\mathbf{x}}}+{\mathbf{x}})]$
$\displaystyle=\frac{\partial}{\partial\bar{{\textnormal{x}}}_{j}}\mathbb{E}_{{\mathbf{x}}\sim
N}[{\textnormal{x}}_{i}f(\bar{{\mathbf{x}}}+{\mathbf{x}})]=\mathbb{E}_{{\mathbf{x}}\sim
N}[{\textnormal{x}}_{i}{\textnormal{x}}_{j}f(\bar{{\mathbf{x}}}+{\mathbf{x}})],$

as required. ∎

##  Appendix B Omitted Proofs Of Our Results

###  B.1 Proof of Lemma 1

See 1

###### Proof.

1$\iff$2: Let us, for clarity, write $\underline{A}=A\setminus S$ and
$\underline{B}=B\setminus S$. We suppose that 1 holds as in Definition 1,

$\displaystyle\frac{\partial
f}{\partial{\mathbf{x}}_{\underline{A}}}({\mathbf{x}})=F_{-\underline{B}}({\mathbf{x}}_{\mathcal{V}\setminus\underline{B}}).$
(9)

Integrating it with respect to ${\mathbf{x}}_{\underline{A}}$ gives

$\displaystyle f({\mathbf{x}})$
$\displaystyle=\int_{{\mathbf{x}}_{\underline{A}}}F_{-\underline{B}}({\mathbf{x}}_{\mathcal{V}\setminus\underline{B}})d{\mathbf{x}}_{\underline{A}}+f_{-\underline{A}}({\mathbf{x}}_{\mathcal{V}\setminus\underline{A}})$
$\displaystyle=f_{-\underline{B}}({\mathbf{x}}_{\mathcal{V}\setminus\underline{B}})+f_{-\underline{A}}({\mathbf{x}}_{\mathcal{V}\setminus\underline{A}}),$

for some $f_{-\underline{B}}$ and $f_{-\underline{A}}$. This proves statement
2 and allows us to recover Equation (9) for $\underline{A}$ ($\underline{B}$)
by differentiating with respect to ${\mathbf{x}}_{\underline{A}}$
(${\mathbf{x}}_{\underline{B}}$).

1$\iff$3: Suppose that 1 holds as in Definition 1. Since the right-hand side
of the equation is not a function of ${\mathbf{x}}_{\underline{B}}$,
differentiating the equation with respect to ${\mathbf{x}}_{\underline{B}}$
gives $0$, which proves statement 3. We recover statement 1 from statement 3
by integrating with respect to ${\mathbf{x}}_{\underline{B}}$ (or
${\mathbf{x}}_{\underline{A}}$). ∎

###  B.2 Proof of Theorem 1

See 1

###### Proof.

Consider a probability distribution defined by

$\displaystyle e_{f}({\mathbf{x}})=\frac{\exp f({\mathbf{x}})}{Z},\ \
\text{where}\ \ Z\triangleq\int_{{\mathbf{x}}}\exp
f({\mathbf{x}})d{\mathbf{x}}.$

Then, let $i$ and $j$ be nodes in $\mathcal{G}$ that are not linked:
$(i,j)\notin\mathcal{E}$. By Definition 2, it follows that

$\displaystyle e_{f}({\mathbf{x}})$
$\displaystyle=\frac{\exp\big{(}f_{-i}({\mathbf{x}}_{-i})+f_{-j}({\mathbf{x}}_{-j})\big{)}}{Z}$
$\displaystyle=\frac{\exp f_{-i}({\mathbf{x}}_{-i})}{\sqrt{Z}}\frac{\exp
f_{-j}({\mathbf{x}}_{-j})}{\sqrt{Z}}.$

That is, measure $e_{f}({\mathbf{x}})$ factorizes over
$\\{\mathcal{V}\setminus\\{i\\},\mathcal{V}\setminus\\{j\\}\\}$. Hence, random
variables ${\textnormal{x}}_{i}$ and ${\textnormal{x}}_{j}$ are
(probabilistically) independent given
${\mathbf{x}}_{\mathcal{V}\setminus\\{i,j\\}}$ [studeny2006probabilistic]. As
the pair $(i,j)$ was arbitrary, this implies that $e_{f}({\mathbf{x}})$
satisfies the pairwise Markov property with respect to $\mathcal{G}$
(equivalently, $\mathcal{G}$ is an I-map of $e_{f}({\mathbf{x}})$). As the
measure induced by $e_{f}({\mathbf{x}})$ is strictly positive, Hammersley-
Clifford Theorem [clifford1971markov] implies that $e_{f}({\mathbf{x}})$
factorizes according to $\mathcal{G}$. That is, for the clique set
$\mathcal{C}$ of $\mathcal{G}$, there exist positive functions
$\\{u_{C}({\mathbf{x}}_{C})\ |\ C\in\mathcal{C}\\}$, such that

$\displaystyle
e_{f}({\mathbf{x}})=\prod_{C\in\mathcal{C}}u_{C}({\mathbf{x}}_{C}).$

Taking the logarithm on both sides, we get

$\displaystyle f({\mathbf{x}})=\sum_{C\in\mathcal{C}}\log
u_{C}({\mathbf{x}}_{C})+\log
Z\triangleq\sum_{C\in\mathcal{C}}f_{C}({\mathbf{x}}_{C})+\log Z.$

That is, the original function $f({\mathbf{x}})$ can be expressed as a sum of
functions of cliques of graph $\mathcal{G}$ (up to an input-independent
constant). Of course, the constant can be subsumed into the decomposing
functions. This lets us conclude the proof of the theorem.

∎

###  B.3 Proof of Theorem 2

*

###### Proof.

Following the standard literature of regression [agarwal2019reinforcement] and
using Assumption (3), by leveraging $f_{C}\in\mathcal{F}_{C}$, we can obtain
with probability $1-\delta$,

$\displaystyle\mathbb{E}_{{\mathbf{x}}\sim
p}\Big{[}\big{(}\hat{f}({\mathbf{x}})-f({\mathbf{x}})\big{)}^{2}\Big{]}\leq\mathrm{Stat},\quad\mathrm{Stat}=D\times\frac{\sum_{C\in\mathcal{C}}\log(|\mathcal{F}_{C}|/\delta)}{n}.$
(10)

where
$\hat{f}({\mathbf{x}})=\sum_{C}\hat{f}_{C}({\mathbf{x}}_{C}),f({\mathbf{x}})=\sum_{C}f_{C}({\mathbf{x}}_{C})$,
and $D$ is universal constant, noting
$\log(|\mathcal{F}|)=\sum_{C}\log(|\mathcal{F}_{C}|).$ Note this statement
states that the MSE guarantee is ensured for the whole
$\hat{f}({\mathbf{x}})$.

Now, we have

$\displaystyle\mathbb{E}_{{\mathbf{x}}\sim
p}\big{[}\big{(}\hat{f}({\mathbf{x}})-f({\mathbf{x}})\big{)}^{2}\big{]}$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}\sim
p}\Bigg{[}\Big{(}\sum_{C}\\{\hat{f}_{C}({\mathbf{x}}_{C})-f_{C}({\mathbf{x}}_{C})\\}\Big{)}^{2}\Bigg{]}$
$\displaystyle=\sum_{C}\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C}({\mathbf{x}}_{C})-f_{C}({\mathbf{x}}_{C})\\}^{2}\right]+\sum_{C_{1}\neq
C_{2}}\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C_{1}}({\mathbf{x}}_{C_{1}})-f_{C_{1}}({\mathbf{x}}_{C_{1}})\\}\\{\hat{f}_{C_{2}}({\mathbf{x}}_{C_{2}})-f_{C_{2}}({\mathbf{x}}_{C_{2}})\\}\right].$

Here, using Assumption (1) and (2), note

$\displaystyle\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C_{1}}({\mathbf{x}}_{C_{1}})-f_{C_{1}}({\mathbf{x}}_{C_{1}})\\}\\{\hat{f}_{C_{2}}({\mathbf{x}}_{C_{2}})-f_{C_{2}}({\mathbf{x}}_{C_{2}})\\}\right]$
$\displaystyle\leq\sigma\times\left\\{\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C_{1}}({\mathbf{x}}_{C_{1}})-f_{C_{1}}({\mathbf{x}}_{C_{1}})\\}^{2}\right]^{1/2}\right\\}\times\left\\{\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\left\\{\hat{f}_{C_{2}}({\mathbf{x}}_{C_{2}})-f_{C_{2}}({\mathbf{x}}_{C_{2}})\right\\}^{2}\right]^{1/2}\right\\}.$

Hence, we get

$\displaystyle\mathbb{E}_{{\mathbf{x}}\sim
p}[\\{\hat{f}({\mathbf{x}})-f({\mathbf{x}})\\}^{2}]$ (11)
$\displaystyle\geq\sum_{C}\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C}({\mathbf{x}}_{C})-f_{C}({\mathbf{x}}_{C})\\}^{2}\right]+\sigma\times\left\\{\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C_{1}}({\mathbf{x}}_{C_{1}})-f_{C_{1}}({\mathbf{x}}_{C_{1}})\\}^{2}\right]^{1/2}\right\\}\times\left\\{\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\left\\{\hat{f}_{C_{2}}({\mathbf{x}}_{C_{2}})-f_{C_{2}}({\mathbf{x}}_{C_{2}})\right\\}^{2}\right]^{1/2}\right\\}$
(12) $\displaystyle\geq\sigma\left\\{\sum_{C}\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C}({\mathbf{x}}_{C})-f_{C}({\mathbf{x}}_{C})\\}^{2}\right]^{1/2}\right\\}^{2}+\\{1-\sigma\\}\sum_{C\in\mathcal{C}}\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C}({\mathbf{x}}_{C})-f_{C}({\mathbf{x}}_{C})\\}^{2}\right].$
(13)
$\displaystyle\geq\\{1-\sigma\\}\sum_{C\in\mathcal{C}}\mathbb{E}_{{\mathbf{x}}\sim
p}\left[\\{\hat{f}_{C}({\mathbf{x}}_{C})-f_{C}({\mathbf{x}}_{C})\\}^{2}\right].$
(14)

Therefore, we have

$\displaystyle\sum_{C\in\mathcal{C}}\mathbb{E}_{{\mathbf{x}}\sim
p}[\\{\hat{f}_{C}({\mathbf{x}}_{C})-f_{C}({\mathbf{x}}_{C})\\}^{2}]\leq\frac{1}{1-\sigma}\times\mathrm{Stat}.$
(15)

Then, for any policy $\pi\in\Pi$, we get

$\displaystyle J(\pi)-J(\hat{\pi})$ $\displaystyle\leq
J(\pi)-\mathbb{E}_{{\mathbf{x}}\sim\pi}[\hat{f}({\mathbf{x}})]+\mathbb{E}_{{\mathbf{x}}\sim\pi}[\hat{f}({\mathbf{x}})]-\mathbb{E}_{{\mathbf{x}}\sim\hat{\pi}}[\hat{f}({\mathbf{x}})]+\mathbb{E}_{{\mathbf{x}}\sim\hat{\pi}}[\hat{f}({\mathbf{x}})]-J(\hat{\pi})$
$\displaystyle=J(\pi)-\mathbb{E}_{{\mathbf{x}}\sim\pi}[\hat{f}({\mathbf{x}})]+\mathbb{E}_{{\mathbf{x}}\sim\hat{\pi}}[\hat{f}({\mathbf{x}})]-J(\hat{\pi})$
$\displaystyle=\mathbb{E}_{{\mathbf{x}}\sim\pi}[f({\mathbf{x}})]-\mathbb{E}_{{\mathbf{x}}\sim\pi}[\hat{f}({\mathbf{x}})]+\mathbb{E}_{{\mathbf{x}}\sim\hat{\pi}}[\hat{f}({\mathbf{x}})]-\mathbb{E}_{{\mathbf{x}}\sim\hat{\pi}}[f({\mathbf{x}})].$

We use the following to upper-bound last terms:

$\displaystyle|\mathbb{E}_{{\mathbf{x}}\sim\pi}[f({\mathbf{x}})-\hat{f}({\mathbf{x}})]|\leq\sum_{C\in\mathcal{C}}\mathbb{E}_{{\mathbf{x}}\sim\pi}[|f_{C}({\mathbf{x}}_{C})-\hat{f}_{C}({\mathbf{x}}_{C})|]$
(Triangle inequality)
$\displaystyle\leq\left\\{\max_{C\in\mathcal{C}}\max_{{\mathbf{x}}_{C}\in\mathcal{X}_{C}}\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\right\\}\sum_{C\in\mathcal{C}}\mathbb{E}_{{\mathbf{x}}\sim
p}[|f_{C}({\mathbf{x}})-\hat{f}_{C}({\mathbf{x}})|]$ (Importance sampling)
$\displaystyle\leq\left\\{\max_{C\in\mathcal{C}}\max_{{\mathbf{x}}_{C}\in\mathcal{X}_{C}}\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\right\\}\sum_{C\in\mathcal{C}}\sqrt{\mathbb{E}_{{\mathbf{x}}\sim
p}[\\{f_{C}({\mathbf{x}}_{C})-\hat{f}_{C}({\mathbf{x}}_{C})\\}^{2}]}$
(Jensen’s inequality)
$\displaystyle\leq\left\\{\max_{C\in\mathcal{C}}\max_{{\mathbf{x}}_{C}\in\mathcal{X}_{C}}\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\right\\}\sqrt{|\mathcal{C}|}\sqrt{\sum_{C\in\mathcal{C}}\mathbb{E}_{{\mathbf{x}}\sim
p}[\\{f_{C}({\mathbf{x}}_{C})-\hat{f}_{C}({\mathbf{x}}_{C})\\}^{2}]}$ (CS
inequality)
$\displaystyle=\max_{C\in\mathcal{C}}\max_{{\mathbf{x}}_{C}\in\mathcal{X}_{C}}\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\times\sqrt{|\mathcal{C}|}\times\sqrt{\frac{\mathrm{Stat}}{1-\sigma}}.$
(Use Ineq. 15)

Finally, combining everything, we have

$\displaystyle J(\pi)-J(\hat{\pi})$ $\displaystyle\leq
2\max_{\pi\in\Pi}\max_{C\in\mathcal{C}}\max_{{\mathbf{x}}_{C}\in\mathcal{X}_{C}}\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\times\sqrt{|\mathcal{C}|}\mathrm{Stat}$
$\displaystyle\lesssim\max_{\pi\in\Pi}\max_{C\in\mathcal{C}}\max_{{\mathbf{x}}_{C}\in\mathcal{X}_{C}}\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\times\sqrt{\frac{|\mathcal{C}|\sum_{C\in\mathcal{C}}\log(|\mathcal{F}_{C}|/\delta)}{n}}\times\sqrt{\frac{1}{1-\sigma}}.$

∎

###  B.4 Proof of Lemma 2

*

###### Proof.

Let
${\mathbf{x}}=[{\mathbf{x}}^{\top}_{C},\\{{\mathbf{x}}^{\prime}\\}^{\top}_{C}]^{\top}$.
Then,

$\displaystyle\max_{{\mathbf{x}}^{\prime}_{C}}\frac{\pi({\mathbf{x}})}{p({\mathbf{x}})}=\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}\times\max_{{\mathbf{x}}^{\prime}_{C}}\frac{\pi({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})}{p({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})}\geq\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}.$
(16)

Here, we use
$\max_{{\mathbf{x}}^{\prime}_{C}}\frac{\pi({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})}{p({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})}\geq
1$. This is because if
$\max_{{\mathbf{x}}^{\prime}_{C}}\frac{\pi({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})}{p({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})}<1$,
we would get contradiction:

$\displaystyle
1=\sum_{{\mathbf{x}}^{\prime}_{C}}\pi({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})\leq\sum_{{\mathbf{x}}^{\prime}_{C}}p({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})\max_{{\mathbf{x}}^{\prime}_{C}}\frac{\pi({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})}{p({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})}<\sum_{{\mathbf{x}}^{\prime}_{C}}p({\mathbf{x}}^{\prime}_{C}|{\mathbf{x}}_{C})=1.$

Finally, using Eq. 16, we obtain

$\displaystyle\max_{{\mathbf{x}}}\frac{\pi({\mathbf{x}})}{p({\mathbf{x}})}=\max_{{\mathbf{x}}_{C}}\max_{{\mathbf{x}}^{\prime}_{C}}\frac{\pi({\mathbf{x}})}{p({\mathbf{x}})}\geq\max_{{\mathbf{x}}_{C}}\frac{\pi({\mathbf{x}}_{C})}{p({\mathbf{x}}_{C})}.$

∎

##  Appendix C Visualizations

Figure 6: Ground-truth graphs (left column) for 4- and 7-dimensional problems
and graphs reconstructed by Equation (7). These achieve the normalized graph-
editing distances from the ground truth of 0.17 and 0.14, respectively.

##  Appendix D Hyperparameters For Experiments

We conducted our experiments with on a Nvidia GP102 GPU, with CUDA 11.6, and
used PyTorch as our programming package.

In all runs, for both Algorithm 1 and COMs [trabucco2021conservative], all
neural networks were single-hidden layer MLPs, with 128 hidden units. Also,
both methods were trained on batches of size 500. The rest of hyper-
parameters: learning rates, gradient ascent steps, COMs gap (for COMs), for
both COMs and our algorithm, were adopted from the COMs paper
[trabucco2021conservative].

In the toy example, both methods were trained for $2e4$ steps. In
SuperConductor, they were trained for $50$ epochs.

We estimate the running average of the pseudo-Hessian from Equation (7) with
momentum parameter $0.99$. The edge test is done at $1-\alpha=0.95$ confidence
level. In the toy example, we update the model’s belief of the graph every
$5e3$ steps, while in SuperConductor every $100$ steps.

## Checklist

  1. 1.

For all models and algorithms presented, check if you include:

    1. (a)

A clear description of the mathematical setting, assumptions, algorithm,
and/or model. [Yes/No/Not Applicable] Yes

    2. (b)

An analysis of the properties and complexity (time, space, sample size) of any
algorithm. [Yes/No/Not Applicable] Not applicable: some of our algorithms use
deep neural networks and generative-modeling algorithms which are not
analytically tractable.

    3. (c)

(Optional) Anonymized source code, with specification of all dependencies,
including external libraries. [Yes/No/Not Applicable] Yes

  2. 2.

For any theoretical claim, check if you include:

    1. (a)

Statements of the full set of assumptions of all theoretical results.
[Yes/No/Not Applicable] Yes

    2. (b)

Complete proofs of all theoretical results. [Yes/No/Not Applicable] Yes

    3. (c)

Clear explanations of any assumptions. [Yes/No/Not Applicable] Yes

  3. 3.

For all figures and tables that present empirical results, check if you
include:

    1. (a)

The code, data, and instructions needed to reproduce the main experimental
results (either in the supplemental material or as a URL). [Yes/No/Not
Applicable] Yes

    2. (b)

All the training details (e.g., data splits, hyperparameters, how they were
chosen). [Yes/No/Not Applicable] Yes

    3. (c)

A clear definition of the specific measure or statistics and error bars (e.g.,
with respect to the random seed after running experiments multiple times).
[Yes/No/Not Applicable] Yes

    4. (d)

A description of the computing infrastructure used. (e.g., type of GPUs,
internal cluster, or cloud provider). [Yes/No/Not Applicable] Yes

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets, check if you include:

    1. (a)

Citations of the creator If your work uses existing assets. [Yes/No/Not
Applicable] Yes

    2. (b)

The license information of the assets, if applicable. [Yes/No/Not Applicable]
Not applicable

    3. (c)

New assets either in the supplemental material or as a URL, if applicable.
[Yes/No/Not Applicable] Not applicable

    4. (d)

Information about consent from data providers/curators. [Yes/No/Not
Applicable] Not applicanble

    5. (e)

Discussion of sensible content if applicable, e.g., personally identifiable
information or offensive content. [Yes/No/Not Applicable] Not applicable

  5. 5.

If you used crowdsourcing or conducted research with human subjects, check if
you include:

    1. (a)

The full text of instructions given to participants and screenshots.
[Yes/No/Not Applicable] Not applicable

    2. (b)

Descriptions of potential participant risks, with links to Institutional
Review Board (IRB) approvals if applicable. [Yes/No/Not Applicable] Not
applicable

    3. (c)

The estimated hourly wage paid to participants and the total amount spent on
participant compensation. [Yes/No/Not Applicable] Not applicable