Cliqueformer: Model-Based Optimization With Structured Transformers

## Abstract

Abstract Expressive large-scale neural networks enable training powerful models for prediction tasks. However, in many engineering and science domains, such models are intended to be used not just for prediction, but for design—e.g., creating new proteins that serve as effective therapeutics, or creating new materials or chemicals that maximize a downstream performance measure. Thus, researchers have recently grown an interest in building deep learning methods that solve offline model-based optimization (MBO) problems, in which design candidates are optimized with respect to surrogate models learned from offline data. However, straightforward application of predictive models that are effective at predicting in-distribution properties of a design are not necessarily the best suited for use in creating new designs. Thus, the most successful algorithms that tackle MBO draw inspiration from reinforcement learning and generative modeling to meet the in-distribution constraints. Meanwhile, recent theoretical works have observed that exploiting the structure of the target black-box function is an effective strategy for solving MBO from offline data. Unfortunately, discovering such structure remains an open problem. In this paper, following first principles, we develop a model that learns the structure of an MBO task and empirically leads to improved designs. To this end, we introduce Cliqueformer —a scalable transformer-based architecture that learns the black-box function’s structure in the form of its functional graphical model (FGM), thus bypassing the problem of distribution shift, previously tackled by conservative approaches. We evaluate Cliqueformer on various tasks, ranging from high-dimensional black-box functions from MBO literature to real-world tasks of chemical and genetic design, consistently demonstrating its state-of-the-art performance.

## 1 Introduction

Most of the common use cases of deep learning (DL) so far have taken the form of prediction tasks (Hochreiter & Schmidhuber, ; He et al., ; Krizhevsky et al., ; Vaswani et al., ). However, in many applications, *e.g.* protein synthesis or chip design, we might want to use powerful models to instead solve *optimization* problems. Clearly, accurate predictions of a target score of an object could be used to find a design of the object that maximizes that score. Such a methodology is particularly useful in engineering problems in which evaluating solution candidates comes with big risk. For example, synthesizing a proposed protein requires a series of wet lab experiments and induces extra cost and human effort (Gómez-Bombarelli et al., ; Brookes et al., ). Thus, to enable proposing *de-novo* generation of strong solution candidates, researchers have drawn their attention to offline *black-box optimization* (BBO), often referred to as *model-based optimization* (MBO). In this paradigm, first, a surrogate model of the score is learned from offline data. Next, a collection of designs is trained to maximize the surrogate, and then proposed as candidates for maximizers of the target score (Gómez-Bombarelli et al., ; Kumar et al., ).

Unfortunately, model-based optimization (MBO) introduces unique challenges not encountered in classical prediction tasks. The most significant issue arises from the incomplete coverage of the design space by the data distribution. This limitation leads to a phenomenon known as *distribution shift*, where optimized designs drift away from the original data distribution. Consequently, this results in poor proposals with significantly overestimated scores (Trabucco et al., ; Geng, ). To address it, popular MBO algorithms have been employing techniques from offline reinforcement learning (Kumar et al., ; Trabucco et al., ) and generative modeling (Kumar & Levine, ; Mashkaria et al., ) to enforce the in-distribution constraint. Meanwhile, much of the recent success of DL has been driven by domain-specific neural networks that, when scaled together with the amount of data, lead to increasingly better performance. While researchers have managed to establish such models in several fields, it is not immediately clear how to do it in MBO. Recent theoretical work, however, has shown that MBO methods can benefit from information about the target function’s *structure*, which can be implemented as a decomposition of the surrogate over the target’s functional graphical model (Grudzien et al., , FGM). This insight opens up new possibilities for developing more effective MBO models by injecting such structure into their architecture. However, how to integrate such decompositions into scalable neural networks remains an open question, and addressing this challenge is the focus of this work.

Figure: (a) Contour curves in XY-planes
Refer to caption: x1.png

In contrast to previous works, in our paper, we develop a scalable model that tackles MBO by learning the structure of the black-box function through the formalism of functional graphical models. Our architecture aims to solve MBO by *1)* decomposing the predictions over the *cliques* of the function’s FGM, and *2)* enforcing the cliques’ marginals to have wide coverage with our novel form of the variational bottleneck (Kingma & Welling, ; Alemi et al., ). However, building upon our Theorem 2, we do not follow grudzien2024functional during the FGM discovery step, and instead subsume it in the learning algorithm. To enable scaling to high-dimensional problems and large datasets, we employ the transformer backbone (vaswani2017attention). Empirically, we demonstrate that our model, *Cliqueformer*, inherits the scalability guarantees of MBO with FGM (see Figure 1). We further demonstrate its superiority to baselines in a suite of tasks with latent radial-basis functions (grudzien2024functional) and real-world chemical (hamidieh2018data) and DNA design tasks (trabucco2022design; uehara2024understanding).

## 2 Preliminaries

In this section, we provide the necessary background on model-based optimization. Additionally, we cover the basics of functional graphical models on top of which we build Cliqueformer.

### 2.1 Model-based Optimization

We consider a model-based optimization problem, where we are given a dataset $\mathcal{D}=\{{\mathbf{x}}^{i},{\textnormal{y}}^{i}\}_{i=1}^{N}$ of examples ${\mathbf{x}}\in\mathcal{X}$, following distribution $p({\mathbf{x}})$, and their values ${\textnormal{y}}=f({\mathbf{x}})\in\mathbb{R}$ under an unkown (black-box) function $f:\mathcal{X}\rightarrow\mathbb{R}$.
Our goal is to optimize this function offline—to find its maximizer

$$ $\displaystyle{\mathbf{x}}^{\star}=\operatorname*{arg\,max}_{{\mathbf{x}}\in \mathcal{X}}f({\mathbf{x}})$ (1) $$

by only using information provided in $\mathcal{D}$ (kumar2020model).
Sometimes, a more general objective in terms of a *policy* over $\pi({\mathbf{x}})$ is also used, $\eta({\mathbf{x}})=\mathbb{E}_{{\mathbf{x}}\sim\pi}[f({\mathbf{x}})]$.
In either form, unlike in Bayesian optimization, we cannot make additional queries to the black-box function (brochu2010tutorial; kumar2020model).
This formulation represents settings in which obtaining such queries is prohibitively costly, such as tests of new chemical molecules or of new hardware architectures (kim2016pubchem; kumar2021data).

To solve MBO, it is typical to learn a model $f_{\theta}({\mathbf{x}})$ of $f({\mathbf{x}})$ parameterized by a vector $\theta$ with a regression method and data from $\mathcal{D}$,

$$ $\displaystyle L(\theta)=\mathbb{E}_{({\mathbf{x}},{\textnormal{y}})\sim \mathcal{D}}\big{[}\big{(}f_{\theta}({\mathbf{x}})-{\textnormal{y}}\big{)}^{2} \big{]}+\text{Reg}(\theta)$ (2) $$

where $\text{Reg}(\theta)$ is an optional regularizer.
Classical methods choose $\text{Reg}(\theta)$ to be identically zero, while conservative methos use the regularizer to bring the values of examples out of $\mathcal{D}$ down.
For example, the regularizer of Conservative Objective Model’s (trabucco2021conservative, COMs) is

$$ $\displaystyle\text{Reg}_{\text{com}}(\theta)=\alpha\big{(}\mathbb{E}_{{\mathbf {x}}\sim\mu_{\theta_{\perp}}}[f_{\theta}({\mathbf{x}})]-\mathbb{E}_{{\mathbf{x }}\sim\mathcal{D}}[f_{\theta}({\mathbf{x}})]\big{)},\quad\alpha>0,$ $$

where $(\cdot)_{\perp}$ is the stop-gradient operator
and $\mu_{\theta_{\perp}}({\mathbf{x}})$ is the distribution obtained with a few gradient ascent steps on ${\mathbf{x}}$ initialized from $\mathcal{D}$. This distribution depends on the value of $\theta$ but is not differentiated through while computing the loss’s gradient, and thus we denote it by $\theta_{\perp}$.
Unfortunately, in addition to the extra computational cost that the inner-loop gradient ascent induces, COMs’s regularizer limits the amount of improvement that it allows its designs to make.

Figure: Figure 2: An FGM of a 5D function which decomposes as $f({\mathbf{x}})=f_{-5}({\mathbf{x}}_{-5})+f_{-1}({\mathbf{x}}_{-1})$. By Definition 1, nodes ${\textnormal{x}}_{1}$ and ${\textnormal{x}}_{5}$ are not linked.
Refer to caption: x3.png

This is particularly frustrating since recent work on functional graphical models (grudzien2024functional, FGM) delivered a premise of large improvements in the case when the black-box function’s graph can be discovered, as we explain in the next subsection.

### 2.2 Functional Graphical Models

An FGM of a high-dimensional function $f({\mathbf{x}})$ is a graph over individual components of ${\mathbf{x}}$ that separates ${\textnormal{x}}_{i},{\textnormal{x}}_{j}\in{\mathbf{x}}$ if their contributions to $f({\mathbf{x}})$ are independent of each other.
Knowing such a structure of $f$ one can eliminate interactions between independent variables from a model that approximates it, and thus prevent an MBO algorithm from exploiting them.
We summarize basic properties of FGMs below.
In what follows, we denote $\mathcal{X}_{-i}$ as the design (input) space without the $i^{th}$ subspace, and ${\mathbf{x}}_{-i}$ as its element (^1^11For example, if $\mathcal{X}=\mathcal{X}_{1}\times\mathcal{X}_{2}\times\mathcal{X}_{3}$ and ${\mathbf{x}}=({\textnormal{x}}_{1},{\textnormal{x}}_{2},{\textnormal{x}}_{3})$, then $\mathcal{X}_{-2}=\mathcal{X}_{1}\times\mathcal{X}_{3}$, and ${\mathbf{x}}_{-2}=({\textnormal{x}}_{1},{\textnormal{x}}_{3})$.).
We also write $[K]$ to denote the set $\{1,\dots,K\}$.

###### Definition 1 .

Let ${\mathbf{x}}=({\textnormal{x}}_{v}\ |\ v\in\mathcal{V})$ be a joint variable with index set $\mathcal{V}$, and $f({\mathbf{x}})$ be a real-valued function.
An FGM $\mathcal{G}=(\mathcal{V},\mathcal{E})$ of $f({\mathbf{x}})$ is a graph where the edge set $\mathcal{E}\subset\mathcal{X}^{2}$ is such that,

$\exists\ f_{-i}:\mathcal{X}_{-i}\rightarrow\mathbb{R}\ \text{and}\ \ f_{-j}:
\mathcal{X}_{-j}\rightarrow\mathbb{R},\ \text{with}\ f({\mathbf{x}})=f_{-i}({
\mathbf{x}}_{-i})+f_{-j}({\mathbf{x}}_{-j}),\text{implies}\ (i,j)\notin
\mathcal{E}.$

See Figure 2 for illustration.

The basic result about FGMs is that they allow for decomposition of the target function into sub-functions with smaller, partially-overlapping inputs, from the FGM’s set of maximal cliques $\mathcal{C}$,

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

Intuitively, the decomposition enables more efficient learning of the target function since it can be constructed by adding together functions defined on smaller inputs, which are easier to learn.
This, in turn, allows for more efficient MBO since the joint solution ${\mathbf{x}}^{\star}$ can be recovered by *stitching* individual solutions ${\mathbf{x}}^{\star}_{C}$ to smaller problems.
This intuition is formalized by the following theorem.

###### Theorem 1 ( grudzien2024functional ) .

Let $f({\mathbf{x}})$ be a real-valued function, $\mathcal{C}$ be the set of maximal cliques of its FGM, and $\Pi$ be a policy class.
Let $\text{C}_{\text{stat}}$ and $\text{C}_{\text{cpx}}$ be constants that depend on the probability distribution of ${\mathbf{x}}$ and function approximator class’s complexity, respectively, defined in Appendix A.
Then, the regret of MBO with the FGM information is given by,

$\eta(\pi^{\star})-\eta(\hat{\pi}_{\text{FGM}})\leq\text{C}_{\text{stat}}\text{
C}_{\text{cpx}}\max\limits_{\pi\in\Pi,{\mathbf{x}}\in\mathcal{X},C\in\mathcal{
C}}\frac{\pi({\mathbf{x}}_{C})}{p_{C}({\mathbf{x}}_{C})}$.

The implication of this theorem is that the FGM-equipped function approximator does not require the dataset to cover the entire design space.
Rather, it only requires that the individual cliques of the space be covered, which is a much milder requirement, especially when the cliques are small.
In the next section, we show how these results can be combined with a transformer, mitigate the distribution shift problem, and enable efficient MBO.

###### Definition 1 .

###### Theorem 1 ( grudzien2024functional ) .

## 3 Cliqueformer

This section introduces a neural network model designed to solve MBO problems through standard end-to-end training on offline datasets. We present a new theoretical result, outline the key desiderata for such a model, and propose an architecture—*Cliqueformer*—that addresses these requirements.

### 3.1 Non-Uniqueness of Structure Discovery

Figure: Figure 3: Illustration of construction in the proof of Theorem 2 for $d=2$. Red axes represent ${\mathbf{z}}$ and blue axes represent ${\mathbf{v}}$. When considered a function of ${\mathbf{z}}$, the contour curves (straight lines) of $f$ are a function of both ${\textnormal{z}}_{1}$ and ${\textnormal{z}}_{2}$, but as a function of ${\mathbf{v}}$, they only depend on ${\textnormal{v}}_{1}$. The density of the Gaussian distribution (green curves) are identical circles for both variables.
Refer to caption: x4.png

The regret bound from Theorem 1 applies to methods that use the target function’s FGM in their function approximation.
It implies that such methods can solve even very high-dimensional problems if their underlying FGMs have low-dimensional cliques or, simply speaking, are *sparse*.
Since, in general, no assumptions about the input can be made, this motivates learning a representation of the input for which one can make distributional assumptions and infer the FGM with statistical tests.
Following this reasoning, grudzien2024functional offer a heuristic technique for discovering an FGM over learned, latent, normally-distributed variables.
However, as we formalize with the following theorem, even such attempts are futile in dealing with black-box functions.

###### Theorem 2 .

Let $d\geq 2$ be an integer and ${\mathbf{x}}\in\mathbb{R}^{d}$ be a random variable with positive density in $\mathbb{R}^{d}$.
There exists a function $f({\mathbf{x}})$ and two different reparameterizations, ${\mathbf{z}}=z({\mathbf{x}})$ and ${\mathbf{v}}=v({\mathbf{x}})$, of ${\mathbf{x}}$,
that both follow a standard-normal distribution, but the FGM of $f$ with respect to ${\mathbf{z}}$ is a complete graph (has all possible edges), and with respect to ${\mathbf{v}}$ it is an empty graph (has no edges).

###### Proof Sketch .

The proof is a construction.
We first map ${\mathbf{x}}$ to a standard-normal variable ${\mathbf{z}}$ with tools from high-dimensional statistics.
We then introduce scalar variable y that is a function of ${\mathbf{z}}$ and has a complete FGM. We then show that we can rotate ${\mathbf{z}}$ onto another normal variable ${\mathbf{v}}$ with respect to which y is a function with an empty FGM (see Figure 3 for illustration).
We complete the proof by showing how to express y as a function of ${\mathbf{x}}$. The full proof can be found in Appendix A.

The theorem implies that FGM is not a fixed attribute of a function that can be estimated from the data, but instead should be viewed as a property of the input’s reparameterization.
Furthermore, different reparameterizations feature different FGMs with varied levels of decomposability, some of which may not significantly simplify the target function.
This motivates a reverse approach that begins from defining a desired FGM and learning representations of the input that align with the graph.
In the next subsection, we introduce Cliqueformer, where the FGM is specified as a hyperparameter of the model and a representation of the data that follows its structure is learned.

###### Theorem 2 .

###### Proof Sketch .

### 3.2 Architecture

The goal of this subsection is to derive an MBO model
that can simultaneously learn the target function as well as its structure, and thus be readily applied to MBO.
First, we would like the model to decompose its prediction into a sum of models defined over small subsets of the input variables in the manner of Equation (3).
As discussed in the previous subsection, efforts to discover such a structure are impractical since there exists a plethora of reparameterizations of the data and their corresponding FGMs of varied decomposabilities.
Thus, instead, we propose that the FGM be defined first, and a representation of the data be learned to align with the chosen structure.

###### Desideratum 1 .

The model should follow a pre-defined FGM decomposition in a learned space.

Accomplishing this desideratum is simple.
A model that we train should, after some transformations, split the input’s representations, denoted as ${\mathbf{z}}$, into partially-overlapping cliques, $({\mathbf{z}}_{C_{1}},\dots,{\mathbf{z}}_{C_{N_{\text{clique}}}})$,
and process them independently, followed by a summation.
To implement this, we specify how many cliques we want to decompose the function into, $N_{\text{clique}}$, as well as their dimensionality, $d_{\text{clique}}$, and the size of their *knots*—dimensions at which two consecutive cliques overlap, $d_{\text{knot}}=|{\mathbf{z}}_{C_{i}}\cap{\mathbf{z}}_{C_{i+1}}|$.
Then, we pass each clique through an MLP network that is also equipped with a trigonometric clique embedding, similar to vaswani2017attention to express a different function for each clique,

$f_{\theta}({\mathbf{z}})=\frac{1}{N_{\text{clique}}}\sum_{i=1}^{N_{\text{
clique}}}f_{\theta}({\mathbf{z}}_{C_{i}},{\bm{c}}_{i})$,  where  $\begin{bmatrix}{\bm{c}}_{i,2j}&{\bm{c}}_{i,2j+1}\end{bmatrix}=\begin{bmatrix}
\sin(i\cdot\omega_{j})&\cos(i\cdot\omega_{j})\end{bmatrix}$,

and $\omega_{j}=10^{-8j/d_{\text{model}}}$.
Here, we use the arithmetic mean over cliques rather than summation because, while being functionally equivalent, it provides more stability when $N_{\text{clique}}\rightarrow\infty$.
Pre-defining the cliques over the representations allows us to avoid the problem of discovering arbitrarily dense graphs, as explained in Subsection 3.1.
This architectural choice implies that the regret from Theorem 1 will depend on the coverage of such representations’ cliques, $\max_{i\in[N_{\text{clique}}]}1/e_{\theta}({\mathbf{z}}_{C_{i}})$, where $e_{\theta}({\mathbf{z}})=\mathbb{E}_{{\mathbf{x}}\sim\mathcal{D}}[e_{\theta}({
\mathbf{z}}|{\mathbf{x}})]$ is the marginal distribution of the representations learned by an encoder $e_{\theta}$.
This term can become dangerously large if individual distributions $e_{\theta}({\mathbf{z}}_{C})$
put disproportionally more density to some regions of the latent space than to others,
Thus, to prevent that, we propose to train the latent space so that the distribution of cliques attain wide coverage.

###### Desideratum 2 .

The model should learn representations whose cliques have dsirtibutions featured by wide coverage.

To meet this requirement, we leverage tools from representation learning, but in a novel way.
Namely, we put a variational bottleneck (kingma2013auto; higgins2016beta; alemi2016deep, VIB) on individual cliques of our representations that brings their distribution closer to a prior with wide coverage, which we choose to be the standard-normal prior.
To implement it, when computing the loss for a single example, we sample a single clique to compute the VIB for at random, as opposed to computing it for the joint latent variable like in the classical VIB,

$$ $\displaystyle\text{VIB}({\mathbf{x}},\theta)=\mathbb{E}_{i\sim U[N_{\text{ clique}}]}\big{[}\text{KL}\big{(}e_{\theta}({\mathbf{z}}_{C_{i}}|{\mathbf{x}}) ,p_{C_{i}}({\mathbf{z}}_{C_{i}})\big{)}\big{]},$ (4) $$

where $e_{\theta}({\mathbf{z}}_{C}|{\mathbf{x}})$ is the density of clique $C$ produced by our learnable encoder $e_{\theta}({\mathbf{z}}|{\mathbf{x}})$, and $p_{C}({\mathbf{z}}_{C})$ is the density of ${\mathbf{z}}_{C}$ under the standard-normal distrituion.
Note that it is not equivalent to the classical, down-weighted VIB in expectation either since our cliques overlap, meaning that knots contribute to the VIB more often than regular dimensions.
Together with the model $f_{\theta}$ and the encoder $e_{\theta}$, we train a decoder $d_{\theta}({\mathbf{x}}|{\mathbf{z}})$ that reconstructs the designs from the latent variables.
Putting it all together, the training objective of our model is a VIB-style likelihood objective with a regression term,

$$ $\displaystyle L_{\text{clique}}(\theta)=\mathbb{E}_{({\mathbf{x}},{\textnormal {y}})\sim\mathcal{D},{\mathbf{z}}\sim e_{\theta}(\cdot|{\mathbf{z}}),i\sim U[N _{\text{clique}}]}\Big{[}\text{VIB}({\mathbf{x}},\theta)-\log d_{\theta}({ \mathbf{x}}|{\mathbf{z}})+\tau\cdot\big{(}{\textnormal{y}}-f_{\theta}({\mathbf {z}})\big{)}^{2}\Big{]},$ (5) $$

where $\tau$ is a positive coefficient that we set to $10$ in our experiments.

Figure: Figure 4: Illustration of information flow in Cliqueformer’s training. Data are shown in navy, learnable variables in blue, neural modules in pink, and loss functions in green. The input ${\mathbf{x}}$ is passed to a transformer encoder to compute representation ${\mathbf{z}}$ which is decomposed into cliques with small overlapping knots (highlighted in colors on the figure). The representation goes to the parallel MLPs whose outputs, added together, predict target y. The representation ${\mathbf{z}}$ is also fed to a transformer decoder that tries to recover the original input ${\mathbf{x}}$. Additionally, the representation goes through an information bottleneck from Equation (4) during training.
Refer to caption: x5.png

Since in our neural network we impose the FGM decomposition of the predictive module in the latent space, we need to endow our model with high expressivity to learn representations that meet such demands.
Thus, we model our encoder $e_{\theta}$ and decoder $d_{\theta}$ with transformer networks (vaswani2017attention).
To leverage such an architecture, the encoder begins by mapping the input vector ${\mathbf{x}}\in\mathbb{R}^{d}$ into $d$ vectors of dimensionality $d_{\text{model}}$, which it then processes as if they were a sequence of token embeddings.
After a series of transformer blocks, the network then maps the sequence into a normal distribution over representations, $e_{\theta}({\mathbf{z}}|{\mathbf{x}})$.
A vector ${\mathbf{z}}$ can then be sampled from that distribution and arranged into $N_{\text{clique}}$ cliques with dimensionality $d_{\text{clique}}$ and knot size of $d_{\text{knot}}$.
The new sequence can then be fed into the predictive model $f_{\theta}$ and to the decoder $d_{\theta}$, which is a transformer too.
For illustration of the information flow in Cliqueformer’s training consult Figure 4.

Figure: Algorithm 1 MBO with Cliquefrormer

###### Desideratum 1 .

###### Desideratum 2 .

### 3.3 Optimizing Designs With Cliqueformer

Once Cliqueformer is trained, we use it to optimize new designs.
Typically, MBO methods initialize this step at a sample of designs $({\mathbf{x}}^{i_{b}})_{b=1}^{B}$ drawn from the dataset (trabucco2021conservative).
Since in our algorithm the optimization takes place in the latent space $\mathcal{Z}$, we perform this step by encoding the sample of designs
with Cliqueformer’s encoder, ${\mathbf{z}}^{i_{b}}\sim e_{\theta}({\mathbf{z}}|{\mathbf{x}}^{i_{b}})$.
We then optimize the representation ${\mathbf{z}}^{i_{b}}$ of design ${\mathbf{x}}^{i_{b}}$ to maximize our model’s value,

$$ $\displaystyle L_{\text{mbo}}\big{(}({\mathbf{z}}^{i_{b}})_{b=1}^{B}\big{)}= \frac{1}{B}\sum_{b=1}^{B}f_{\theta}({\mathbf{z}}^{i_{b}}),$ (6) $$

at the same time minding the enumerator of the regret bound from Theorem 1.
That is, we don’t want the optimizer to explore regions under which the marginal densities $e_{\theta}({\mathbf{z}}_{C})=\mathbb{E}_{{\mathbf{x}}\sim\mathcal{D}}[e_{
\theta}({\mathbf{z}}_{C}|{\mathbf{x}})]$ are small.
Fortunately, since the encoder was trained with standard-normal prior on the cliques, $p({\mathbf{z}}_{C})=N(0_{d_{\text{clique}}},I_{d_{\text{clique}}})$, we know that values of ${\mathbf{z}}$ closer to the origin have unilaterally higher marginals.
This simple property of standard-normal distribution allows us to confine the optimizer’s exploration to designs with in-distribution cliques by exponentially decaying the design at every optimization step.
Thus, we use AdamW as our optimizer (loshchilov2017fixing).
We provide the pseudocode of the whole procedure of designign with Cliqueformer in Algorithm 1.

## 4 Related Work

The idea of using machine learning models in optimization problems has existed for a long time, and has been mainly cultivated in the literature on Bayesian optimization (williams2006gaussian; brochu2010tutorial; snoek2012practical, BO).
The BO paradigm relies on two core assumptions: availability of data of examples paired with their target function values, as well as access to an oracle that allows a learning algorithm to query values of proposed examples.
Thus, similarly to reinforcement learning, the challenge of BO is to balance exploitation and exploration of the black-box function modeled by a Gaussian process.

Recently, to help BO tackle very high-dimensional problems, techniques of decomposing the target function have become more popular (kandasamy2015high; rolland2018high).
Most commonly, these methods decompose the target function into functions defined on the input’s partitions.
While such models are likely to deviate far from the functions’ ground-truth structure, one can derive theoretical guarantees for a range of such decompositions under the BO’s query budget assumptions (ziomek2023random).
However, these results do not apply to our setting of offline MBO, where no additional queries are available and the immediate reliability on the model is essential.
Furthermore, instead of partitioning the input, we decompose our prediction over a latent variable that is learned by a transformer, enabling the model to acquire an expressive structure over which the decomposition is valid.

Offline model-based optimization (MBO) has been recently attracting attention of researchers and practitioners from domains where BO assumptions cannot be easily met,
offering an attractive premise of producing solutions directly after training on static datasets, without the need for additional queries.
One of the first tasks tackled by MBO was molecule design (gomez2018automatic), where a variational auto-encoder (kingma2013auto, VAE) was used to learn continuous representations of molecular data.
In contrast to our work, however, this work does not study learning structural properties of the target function, but rather is a proof of concept of applying deep learning to molecular design.
A data type-agnostic algorithm was introduced by brookes2019conditioning, dubbed Conditioning by Adaptive Sampling (CBaS), that iteratively refines its design proposals in response to predictions of a non-differentiable oracle.
While one can use this refinement procedure in combination with trainable models by means of *auto-focusing* (fannjiang2020autofocused), this setting is different than ours since we assume the ability to model the black-box function with a neural network.
Since such a model is differentiable, we can simply rely on automatic differentiation (paszke2019pytorch) to refine our designs.
trabucco2021conservative introduced a neural network-based method, exactly for our setting, dubbed Conservative Objective Models (COMs), where a surrogate model is trained to both predict values of examples that can be found in the dataset, and penalize those that are not.
COMs differs from our work fundamentally, since its contribution lies in the formulation of the conservative regularizer applied to arbitrary neural networks, while we focus on scalable model architectures that facilitate computational design.
Another recent line of work proposes to tackle the design problem through means of generative modeling.
BONET (mashkaria2023generative) and DDOM (pmlr-v202-krishnamoorthy23a) are examples of works that bring the most recent novelties of the field to address design tasks.
BONET does so by training a transformer to generate sequences of designs that monotonically improve in their value,
and DDOM by training a value-conditioned diffusion model.
That is, these methods attempt to generate high-value designs through novel conditional generation mechanisms.
Instead, we model MBO as a maximization problem, and propose a scalable model that acquires the structure of the black-box function through standard gradient-based learning.
To this end, we bring powerful techniques from deep learning and generative modeling, like transformers (vaswani2017attention) and variational-information bottlenecks (kingma2013auto; alemi2016deep).

The work of grudzien2024functional introduced the theoretical foundations of functional graphical models (FGMs), including Theorem 1.
However, as we have shown in Theorem 2, their graph discovery heuristic for neural networks renders learning the black-box function’s structure an open problem.
Our work addresses this issue by subsuming the graph discovery step in the architecture of our *Cliqueformer* that learns to abide by a pre-defined FGM.
As such, the model learns the structure of the target function, as well as learns to predict its value, in synergy within an end-to-end training.

## 5 Experiments

In this section, we provide the empirical evaluation of Cliqueformer.
We begin by benchmarking Cliqueformer against prior methods on tasks from the MBO literature.
Then, we finish by evaluating the benefit of the novel FGM decomposition layer in Cliqueformer through an ablation study.

### 5.1 Benchmarking

We compare our model to three classes of algorithms, each represented by a proven prior method.
As a *naïve* baseline, we employ gradient ascent on a learned model (*Grad. Asc.*).
To compare to *exploratory* methods, we use Reward-Weighted Regression (peters2007reinforcement, *RWR*) which learns by regressing the policy against its most promising perturbations.
To represent the recently proposed *conservative* algorithms, we compare to state-of-the-art Conservative Objective Models (trabucco2021conservative; kumar2021data, *COMs*).
Additionally, to evaluate the importance of the novel elements of our architecture and training, we evaluate gradient ascent with the transformer backbone (vaswani2017attention, *Transformer*).
For COMs and other methods, we use the recommended hyperparameter setting from (trabucco2021conservative).
While Cliqueformer can be tuned for each task, we keep most of the hyperparameters the same: refer to Appendix B for more details.

For every method, we report an empirical estimate of its $100^{th}$ percentile, similarly to trabucco2021conservative.
We estimate it by averaging the values of the top 10 designs out of 1000 candidates, averaged across 5 seeds.
In the following paragraphs, we introduce benchmark tasks, from MBO literature, that we use in our experiments, and analyze their results.

#### Latent Radial-Basis Functions (Lat. RBF).

This is a suite of tasks designed to expose vulnerability of MBO models (grudzien2024functional).
The data pairs $({\mathbf{x}},{\textnormal{y}})$ are generated by first drawing a standard normal vector ${\mathbf{z}}\sim N(0_{d_{z}},I_{d_{z}})$, where $d=\{11,31,41,61\}$, then computing y as a sum of radial-basis functions of $d_{C}$-dimensional cliques of a pre-defined FGM.
The observed vector ${\mathbf{x}}$ is a non-linear transformation of ${\mathbf{z}}$, *i.e.*, ${\mathbf{x}}=T({\mathbf{z}})\in\mathcal{T}\subset\mathbb{R}^{d}$, where $d>d_{z}$, while ${\mathbf{z}}$ itself is hidden from the data.
Such tasks allow us to study whether an MBO method learns to produce *valid* designs by verifying that *ground-truth* inputs ${\mathbf{z}}$ can be recovered from them.
That is, a design $\hat{{\mathbf{x}}}$ for which the map $T^{-1}(\hat{{\mathbf{x}}})$ is ill-defined is considered invalid.
In our experiments, an invalid design receives a value of $-\infty$.
We report the average validity of designs produced by the methods in blue.
The results in Table 1 show that the most able method at keeping its proposals at the manifold of valid designs is Cliqueformer.
Importantly, this ability does not diminish even in higher-dimensional tasks, while the second-best such method, COMs, gradually loses this ability.
We also use these tasks to study if a model is capable of exploiting the RBF’s structure in the optimization step by varying the effective dimensionality $d_{z}$ of the data while keeping $d_{C}$ fixed.
While the naive, the exploratory, and Transformer baselines perform poorly on these tasks overall,
COMs’s performance clearly drops as the task dimension increases.
Meanwhile, as predicted by Theorem 1, Cliqueformer attains similar, strong performance across all tasks.

**Table 1: Experimental results of Cliqueformer and the baselines. Each score is the mean of values of $TopK$=10 of $B$=1000 designs, averaged over 5 runs. The values were normalized with the min-max scheme, where the minimum and the maximum are taken from the dataset. We also provide the standard deviation estimates as the standard deviation for each of the $TopK$ samples, averaged over the runs. Additionally, for Lat. RBF tasks, we provide the average validity in blue, calculated as the percentage of valid designs from the $B$ produced candidates, averaged across the runs. We provide average score (the higher the better) and the average rank for each method. For the average score $-\infty$ was taken into calculation as zero.**
| Task\Method | Grad.Asc. | RWR | COMs | Transformer | Cliqueformer |
| --- | --- | --- | --- | --- | --- |
| Lat. RBF 11 | $-\infty$ $\pm$ 0.00 (0%) | 0.08 $\pm$ 0.08 (1%) | 0.66 $\pm$ 0.04 (72%) | 0.47 $\pm$ 0.05 (60%) | 0.65 $\pm$ 0.07 (74%) |
| Lat. RBF 31 | $-\infty$ $\pm$ 0.00 (0%) | 0.31 $\pm$ 0.10 (3%) | 0.50 $\pm$ 0.05 (32%) | $-\infty$ $\pm$ 0.00 (0%) | 0.64 $\pm$ 0.05 (76%) |
| Lat. RBF 41 | $-\infty$ $\pm$ 0.00 (0%) | 0.35 $\pm$ 0.08 (3%) | 0.45 $\pm$ 0.06 (16%) | 0.20 $\pm$ 0.01 (75%) | 0.66 $\pm$ 0.05 (72%) |
| Lat. RBF 61 | $-\infty$ $\pm$ 0.00 (0%) | 0.29 $\pm$ 0.10 (4%) | 0.25 $\pm$ 0.04 (7%) | 0.16 $\pm$ 0.03 (64%) | 0.66 $\pm$ 0.05 (68%) |
| Superconductor | 1.13 $\pm$ 0.08 | 1.03 $\pm$ 0.07 | 0.97 $\pm$ 0.08 | 0.96 $\pm$ 0.05 | 1.43 $\pm$ 0.04 |
| TF-Bind-8 | 0.99 $\pm$ 0.00 | 1.58 $\pm$ 0.03 | 1.57 $\pm$ 0.02 | 1.48 $\pm$ 0.03 | 1.58 $\pm$ 0.01 |
| DNA HEPG2 | 2.16 $\pm$ 0.07 | 1.91 $\pm$ 0.12 | 1.20 $\pm$ 0.09 | 2.13 $\pm$ 0.06 | 2.10 $\pm$ 0.07 |
| DNA k562 | 2.11 $\pm$ 0.06 | 1.91 $\pm$ 0.11 | 1.80 $\pm$ 0.12 | 2.60 $\pm$ 0.19 | 3.15 $\pm$ 0.07 |
| Ave.score($\uparrow$) | 0.80 | 0.93 | 0.93 | 0.80 | 1.46 |
| Ave.rank($\downarrow$) | 3.75 | 3.00 | 3.13 | 3.50 | 1.38 |

#### Superconductor.

This task poses a challenge of designing a superconducting material, represented by an 81-dimensional vector, with as high critical temperature as possible (hamidieh2018data).
It tests abilities of MBO models in real-world continuous problems.
As the very high score of gradient ascent shows, in contrast to Lat. RBF, this task rewards greedy optimizer updates.
While we developed Cliqueformer paying attention to distribution shift,
its ability to compose individual in-distribution cliques into a joint solution (*stitching*) allows for large improvements that outperform all baselines.
In particular, it greatly outperforms COMs whose conservative objective caps the extent of improvement in the design scores, and Transformer which does not attain Cliqueformer’s decomposition.

#### TFBind-8 & DNA Enhancers.

In these discrete tasks we optimize DNA sequences of length 8 and 200, respectively.
In TFBind-8 the target is the sequence’s binding affinity with a particular transcription factor, while in DNA Enhancers we maximize HEPG2 and k562 activity levels.
Being low-dimensional, TFBind-8 (trabucco2022design) is a testbed that allows us to verify MBO models’ ability to solve discrete tasks.
Cliqueformer solves this problem very efficiently, largely improving upon the dataset, but it is worth noting that most baselines, with the exception of gradient ascent with and without Transformer backbone, performed similarly.
Thus, while the TFBind-8 task does not favor greedy design optimization, Cliqueformer is still able to leverage its other strengths to achieve great performance.
Then, we use the DNA Enhancers tasks (uehara2024understanding)
to study the scalability of our method to very high-dimensional problems and large datasets (approximately $2\times 10^{5}$ examples in this case).
The results in Table 1 and the high score of gradient ascent show that these tasks favor direct, greedy optimization of designs more than they benefit from conservatism of COMs.
Nevertheless, Cliqueformer is able to greatly exceed the quality of observed designs, and performs on par with gradient ascent in HEPG2, and greatly outperforms all baselines in k562, confirming its ability to learn structure within discrete and very high-dimensional data.
Overall, averaging across all tested tasks, Cliqueformer achieves the best overall performance.

### 5.2 Ablation

While the decomposing mechanism of Cliqueformer is entirely novel, other components, such as transformer blocks (vaswani2017attention) and variational-information bottlenecks (kingma2013auto; alemi2016deep) are known and powerful deep learning tools.
In this subsection, we verify the utility of the decomposing component with an ablation study, in which we sweep over the number of cliques of Cliqueformer for a few representative tasks.
In each of the tested tasks we fix the size of the latent variable ${\mathbf{z}}$ and sweep over the number of cliques $N_{\text{clique}}$ into which it can be decomposed.
We cover the case $N_{\text{clique}}=1$ to compare Cliqueformer to an FGM-oblivious VAE with transformer backbone and AdamW design optimizer.

Figure: (a) TFBind-8
Refer to caption: x6.png

Results in Figure 5 show that Cliqueformer does benefit from the FGM decomposition in all tested tasks.
However, the results suggest that the optimal number of cliques varies between tasks, *e.g.* 4 for TFBind-8 and Superconductor, but 2 for Lat. RBF 41.
In our experiments, we consistently obtained good performance by setting the clique size to $d_{\text{clique}}=3$ and choosing the number of cliques so that the total latent dimension approximately matches that of the design.
In DNA Enhancers, we doubled the clique size and halved the number of cliques to decrease the computational cost of attention.
More details of hyperparameters can be found in Appendix B.
An in-depth analysis of the relation between hyperparameters and the performance is an exciting avenue of future work.

## 6 Conclusion

In this work, we proposed Cliqueformer—a scalable neural model for model-based optimization.
We derived its building blocks following recent advances in MBO theory centered in functional graphical models, equipping it with an ability to acquire structure of the black-box target function.
This ability sets the model free from requiring explicit conservative regularization or iterative retraining to propose in-distribution designs.
Empirically, Cliqueformer achieves state-of-the-art performance, which it sustains across all tested tasks.
Cliqueformer opens an exciting avenue of research in MBO focused on scaling design datasets and large neural networks.

## Acknowledgements

Pieter Abbeel holds concurrent appointments as a Professor at UC Berkeley and as an Amazon Scholar. This paper describes work performed at UC Berkeley and is not associated with Amazon.
We thank Google TPU Research Cloud (TRC) for granting us access to TPUs for research.
We also thank Olivia Liu for helping us design the illustration of Cliqueformer.

## Appendix

## Appendix A Theoretical details

The full statement of the following theorem considers an MBO algorithm with function clas $\mathcal{F}=\mathcal{F}_{C_{1}}\oplus...\oplus\mathcal{F}_{C_{N_{\text{clique}
}}}$, so that every of its element has form

$f({\mathbf{x}})=\sum_{i=1}^{N_{\text{clique}}}f_{C_{i}}({\mathbf{x}}_{C_{i}})$.

As described in Section 3, Cliqueformer’s architecture forms such a function class on top of the learned latent space.
We define the statistical constant as

$C_{\text{stat}}=\sqrt{\frac{1}{1-\sigma}}$,  where
$\sigma=\max\limits_{C_{i}\neq C_{j},\hat{f}_{C_{i}},\hat{f}_{C_{j}}}\mathbb{C}
orr_{{\mathbf{x}}\sim p}[\hat{f}_{C_{i}}({\mathbf{x}}_{C_{i}}),\hat{f}_{C_{j}}]$

and the function approximation complexity constant as

$C_{\text{cpx}}=\sqrt{\frac{N_{\text{clique}}\sum_{i=1}^{N_{\text{clique}}}\log
(|\mathcal{F}_{C_{i}}|/\delta)}{N}}$,

where $\delta$ is the PAC error probability (shalev2014understanding).

See 1

See 2

###### Proof.

Since the density of ${\mathbf{x}}$ is positive and continuous, we can form a bijection that maps ${\mathbf{x}}$ to another random variable ${\mathbf{z}}\in\mathbb{R}^{l}$, where $l\leq d$, that follows the standard-normal distribution (dai2019diagnosing, Appendix E).
We denote this bijection as $Z({\mathbf{x}})$.
Let us define

${\textnormal{y}}=f^{z}({\mathbf{z}})=\exp\big{(}\frac{1}{\sqrt{l}}\sum_{i=1}^{
l}{\textnormal{z}}_{i}\big{)}$.

Then, the FGM of $f^{z}$ has an edge between every two variables since each variable’s partial derivative

$\frac{\partial f^{z}}{\partial{\textnormal{z}}^{i}}=\frac{1}{\sqrt{l}}\exp\big
{(}\frac{1}{\sqrt{l}}\sum_{i=1}^{d}{\textnormal{z}}_{i}\big{)}$

is also a function of all others (grudzien2024functional, Lemma 1).
Consider now a rotation $\rho:{\mathbf{z}}\mapsto{\mathbf{v}}=({\textnormal{v}}_{1},\dots,{\textnormal{
v}}_{l})$ such that ${\textnormal{v}}_{1}=\frac{1}{\sqrt{l}}\sum_{i=1}^{l}{\textnormal{z}}_{i}$.
Then, ${\mathbf{v}}\sim N(0_{l},I_{l})$, and y can be expressed in terms of ${\mathbf{v}}$ as ${\textnormal{y}}=f^{v}({\mathbf{v}})=\exp({\textnormal{v}}_{1})$.
Then, the FGM of $f^{v}$ has no edges, since it depends on only one variable, inducing no interactions between any two variables.
Recall that ${\mathbf{x}}=Z^{-1}({\mathbf{z}})$.
Then, ${\mathbf{x}}$ be represented by standard-normal ${\mathbf{z}}$ and ${\mathbf{v}}$, obtainable by

${\mathbf{z}}=Z({\mathbf{x}})$ and ${\mathbf{v}}=\rho({\mathbf{z}})=\rho\big{(}Z({\mathbf{x}})\big{)}$.

Furthermore, we can define

$f({\mathbf{x}})=f^{z}\big{(}Z({\mathbf{x}})\big{)}$

which is identically equal to $f^{z}({\mathbf{z}})$ and $f^{v}({\mathbf{v}})$, which have a complete and an empty FGM, respectively, thus fulfilling the theorem’s claim.
∎

###### Proof.

## Appendix B Experimental details

#### Datasets.

We use the implementation of grudzien2024functional to generate data with latent radial-basis functions.
Also, we initially wanted to use Design Bench (trabucco2022design) for experiments with practical tasks.
However, at the time of this writing, the benchmark suite was suffering a data loss and was not readily available.
To overcome it, we manually found the data and implemented dataset classes.
TFBind-8 (trabucco2022design) could be fully downloaded since the number of possible pairs $({\mathbf{x}},{\textnormal{y}})$ is quite small.
Hence, a design can be evaluated by looking up its score in the dataset.
For Superconductor (hamidieh2018data), we pre-trained an XGBoost oracle on the full dataset, and trained our model and the baselines to predict the labels produced by the oracle.
The proposed designs of the tested models are evaluated by calling the oracle as well.
We obtained DNA Enhancers dataset from the code of uehara2024understanding, available at

[https://github.com/masa-ue/RLfinetuning_Diffusion_Bioseq/tree/master](https://github.com/masa-ue/RLfinetuning_Diffusion_Bioseq/tree/master).

Following the procedure in

[https://github.com/masa-ue/RLfinetuning_Diffusion_Bioseq/blob/master/tutorials/Human-enhancer/1-Enhancer_data.ipynb](https://github.com/masa-ue/RLfinetuning_Diffusion_Bioseq/blob/master/tutorials/Human-enhancer/1-Enhancer_data.ipynb).

we additionally filter the dataset to keep only sequences featured by chromosomes from 1 to 4.
We use their pre-trained oracle for generating labels and evaluation of proposed designs.
Following fannjiang2020autofocused and trabucco2022design, we train our models on the portions of the datasets with values below their corresponding $80^{th}$.
Upon evaluation, we obtain the ground-truth/oracle value of the proposed design y, and normalize it as

$$ $\displaystyle\bar{{\textnormal{y}}}=\frac{{\textnormal{y}}-{\textnormal{y}}_{ \min}}{{\textnormal{y}}_{\max}-{\textnormal{y}}_{\min}},$ $$

and report $\bar{{\textnormal{y}}}$. ${\textnormal{y}}_{\min}$ and ${\textnormal{y}}_{\max}$ are the minimum and the maximum of the training data.
This normalization scheme is different than, for example, the one in the work by trabucco2022design.
We choose this scheme due to its easy interpretability—a score of $\bar{{\textnormal{y}}}>1$ implies improvement over the given dataset, which is the ultimate objective of MBO methods.
However, we note that a score of less than $1$ does not imply failute of the algorithm, since we initialize our designs at a random sample from the dataset, which can be arbitrarily low-value or far from the optimum.
For some functions, like in latent RBFs, the optima are very narrow spikes in a very high-dimensional space, being nearly impossible to find (see Figure 1(a)).
We choose such an evaluation scheme due to its robustness that allows us to see how good ut improving any design our algorithms are overall.

#### Hyper-parameters.

For baselines, we use hyper-parameters suggested by trabucco2021conservative.
We decreased the hidden layer sizes (at no harm to performance) for Lat. RBF 31 and DNA Enhancers tasks where the performance was unstable with larger sizes.
Also, we haven’t tuned most of the Cliqueformer’s hyper-parameters *per-task*.
We found, however, as set of hyper-parameters that works reasonably well on all tasks.

On all tasks, we use 2 transformer blocks in both the encoder and the decoder, with transformer dimension of 64, and 2-head attention.
The predictive model $f_{\theta}({\mathbf{z}})$ is a multi-layer perceptron with 2 hidden layers of dimension 256.
We change it to 512 only for DNA Enhancers.
The best activation function we tested was GELU (hendrycks2016gaussian), and LeakyReLU(0.3) gives similar results.
We use dropout of rate 0.5 (srivastava2014dropout).
In all tasks, weight of the MSE term to $\tau=$10 (recall Equation (5)).
Additionally, we warm up our VIB term linearly for $1000$ steps (with maximal coefficient of 1).
We train the model with AdamW (loshchilov2017fixing) with the default weight decay of Pytorch (paszke2019pytorch).
We set the model learning rate to 1e-4 and the design learning rate to 3e-4 in all tasks.
We train the design with AdamW with high rates of weight decay (ranging from 0.1 to 0.5).

In all tasks, we wanted to keep the dimension of the latent variable ${\mathbf{z}}$ more-less similar to the dimension of the input variable ${\mathbf{x}}$, and would decrease it, if possible without harming performance, to limit the computational cost of the experiments.
The dimension of ${\mathbf{z}}$ can be calculated from the clique and knot sizes as

$dim({\mathbf{z}})=d_{\text{knot}}+N_{\text{clique}}\cdot(d_{\text{clique}}-d_{
\text{knot}})$.

In most tasks, we used the clique dimension $d_{\text{clique}}=3$ with knot size of $d_{\text{knot}}=1$.
We made an exception for Superconductor, where we found a great improvement by setting $d_{\text{clique}}=21$ and $N_{\text{clique}}=4$ (setting $d_{\text{clique}}=3$ and $N_{\text{clique}}=40$ gives score of 0.99); and DNA Enhancers, where we doubled the clique size (to 6) and halved the number of cliques to (40), to lower the computational cost of attention.
In DNA Enhancers tasks, we additionally increased the MLP hidden dimension to 512 due to greater difficulty of modeling high-dimensional tasks.
We summarize the task-specific hyper-parameters in Table 2.

We want to note that these hyper-parameters are not optimal per-task.
Rather, we chose schemes that work uniformly *well enough* on all tasks.
However, each task can benefit from further alteration of hyper-parameters.
For example, we observed that Lat. RBF tasks benefit from different numbers of design steps; for Lat. RBF 41, we found the optimal number to be 400; for Superconductor, it seems to be 200.
Due to time constraints, we have not exploited scalability of Cliqueformer in DNA Enhancers tasks, but observed pre-training losses to decrease more with increased parameter count and training duration.

**Table 2: Hyper-parameter configuration for different benchmark tasks.**
| Task | N_clique | d_clique | MLP dim | design steps | Weight decay |
| --- | --- | --- | --- | --- | --- |
| Lat. RBF 11 | 10 | 3 | 256 | 50 | 0.5 |
| Lat. RBF 31 | 18 | 3 | 256 | 50 | 0.5 |
| Lat. RBF 41 | 20 | 3 | 256 | 50 | 0.5 |
| Lat. RBF 61 | 28 | 3 | 256 | 50 | 0.5 |
| TFBind-8 | 4 | 3 | 256 | 1000 | 0.5 |
| Superconductor | 4 | 21 | 256 | 1000 | 0.5 |
| Dna Enhancers | 40 | 6 | 512 | 1000 | 0.1 |

## Links

We provide an overview of our work on the project website at

[https://sites.google.com/berkeley.edu/cliqueformer/home](https://sites.google.com/berkeley.edu/cliqueformer/home).

In particular, we release our code at

[https://github.com/znowu/cliqueformer-code](https://github.com/znowu/cliqueformer-code).

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - Is there any performance degradation from not using the graph discovery heuristic? It would be helpful to compare the performance, at least on LRBF.
- Do you have further intuition on how to set $N_{\text{clique}} $ or $ d_{\text{clique}} $?
soundness: 2
strengths: - Cliqueformer effectively addresses the limitations of FGM by eliminating the need for explicit graph discovery steps, while still maintaining the scalability advantages of MBO with FGM.
- The proposed method is carefully designed based on theoretical observations
- It demonstrates significant performance improvements over COMs across a wide variety of tasks.
- well-written with clear motivation
summary: This work proposes a new architecture, Cliqueformer, for solving model-based optimization (MBO) using the functional graphical model (FGM) framework. Cliqueformer subsumes the graph discovery step within its architecture, rather than using the graph discovery heuristic from the previous work (Grudzien et al., 2024). The VIB-style objective further ensures broad coverage of each clique’s distribution, enhancing stability in optimization. Experimental results demonstrate that Cliqueformer outperforms previous methods across a range of tasks.
weaknesses: While the paper makes significant contributions, some claims might be somewhat overstated:

- The authors state that "we consistently obtained good performance by setting the clique size to $ d_{\text{clique}} = 3 $." However, the experimental results suggest that the choice of $ d_{\text{clique}} $ or $ N_{\text{clique}} $ significantly affects performance. For instance, Figure 5(b) shows that only $ N=4 $ achieves better performance than $ N=1 $ (assuming $ N=1 $ represents the whole graph), and depending on how $ N_{\text{clique}} $ is set, Cliqueformer sometimes fails to achieve scores higher than 1.
- The claim that "Cliqueformer achieves state-of-the-art performance" may be premature. Since some recent studies are not included as baselines (e.g., BDI [1], BootGen [2], or the baselines in Design-bench [3] for TFBind and superconductor), and because the scores are normalized differently, it's challenging to assert that Cliqueformer achieves state-of-the-art results conclusively.

[1] Chen, Can, et al. "Bidirectional learning for offline infinite-width model-based optimization." Advances in Neural Information Processing Systems 35 (2022): 29454-29467.

[2] Kim, Minsu, et al. "Bootstrapped training of score-conditioned generator for offline design of biological sequences." Advances in Neural Information Processing Systems 36 (2024).

[3] Trabucco, Brandon, et al. "Design-bench: Benchmarks for data-driven offline model-based optimization." International Conference on Machine Learning. PMLR, 2022.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - Why does Cliqueformer suffer less from distribution shift problems? Can adding an additional conservative regularization benefit the model?
- Why does the decoder have to be a transformer model? What if we use a linear or shallow MLP prediction head instead? This is equivalent to using the last layer of the transformer backbone for the latent representation which is the common practice in many domains.
soundness: 3
strengths: The paper presents a novel method for offline model-based optimization. The idea of learning a latent space and imposing FGM on it is interesting. I also like the simplicity of the model and the clear writing of the paper. The empirical performance of the proposed method is strong against different baselines.
summary: The paper proposes Cliqueformer, a new method for offline model-based optimization. The method is built on the idea of functional graph model (FGM) which decomposes the input space into cliques that contribute independently to the function. The paper proposes to learn a latent representation of the input space, imposes FGM cliques on this space, and simultaneously trains regression models on the latent cliques. To perform optimization, the model takes gradient steps of the regression model against some initial latent representations and decodes back to the original input space.
weaknesses: - Looking at Figure 5, the performance varies significantly with different values of cliques, and there is not a universal value across different datasets. How did the authors pick the value in an offline setting where you're not allowed to query the oracle?
- The baselines seem quite outdated. Have the authors tried comparing with more recent baselines such as generative methods using transformers or diffusion models [1, 2, 3, 4]? Some of these papers should also be cited and discussed in the paper.
- The Superconductor task uses a surrogate function and not the true oracle, which was found to not be reliable in [4].
- Minor comment: Algorithm 1 should remove the training process of Cliqueformer for better clarity.

[1] Krishnamoorthy, Siddarth, Satvik Mehul Mashkaria, and Aditya Grover. "Generative pretraining for black-box optimization." arXiv preprint arXiv:2206.10786 (2022).

[2] Krishnamoorthy, Siddarth, Satvik Mehul Mashkaria, and Aditya Grover. "Diffusion models for black-box optimization." International Conference on Machine Learning. PMLR, 2023.

[3] Chen, Can, et al. "Bidirectional learning for offline infinite-width model-based optimization." Advances in Neural Information Processing Systems 35 (2022): 29454-29467.

[4] Nguyen, Tung, Sudhanshu Agrawal, and Aditya Grover. "Expt: Synthetic pretraining for few-shot experimental design." Advances in Neural Information Processing Systems 36 (2023): 45856-45869.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see the points raised in the weaknesses section. 

## Minor Comments

1. Line 249: "distrituion" should be "distribution".

2. Line 731: "failute" on should be "failure".

3. Line 80: "More results in Section 5" is not a complete sentence.

4. Line 745: "We use dropout of rate 0.5" should be replaced with "We use a dropout rate of 0.5"​.

5. The abstract does not seem fitting. The content prior to the "In this paper" sentence, i.e., Lines 11-25, do not strike me as abstract material, and are best suited elsewhere in the paper. For instance, I don't consider explaining "why RL strategies generally work better on MBO problems" a priority for the abstract. Only Lines 26-33 seem specifically relevant to the paper's contributions.
soundness: 3
strengths: 1. The notations and model definitions were well-described. I had an easy time understanding the approach, and the explanations were easy on the mind.

2. The intuitions in the paper are stated clearly, and the desideratums were easy to follow.

3. The figures seemed effective at explaining the related concepts. Figure 4 presents a clear summary of the model, Figure 3 clearly demonstrates the problem with FGM choices, and Figures 1 and 2 are also helpful.

4. A limited set of ablation studies show the effect of the number of cliques on the downstream score.
summary: This paper aims at proposing a model capable of learning the structure of a model-based optimization task under a functional graphical model's framework. The evaluations include a black box optimization task in the 10-s of dimensions range, and some chemical and genetic design tasks. Unlike previous work, the FGM discovery is merged into the model's learning, and the marginal distributions of the cliques were enforced to mimic a standard normal by minimizing a KL divergence.
weaknesses: 1. As is, I fear the introduction and model description sections are quite disconnected from the experiments. The authors spend a substantial amount of time in Sections 1-4 building intution on how the proposed model should work, what features are desired, and so on. The experiments don't do much more than presenting some benchmarks. This makes it difficult for me to verify that the gained improvements are actually the result of the claims and contributions, rathen than being side effects of auxillary design choices. For instance,

  * Was the transformer architecture adaptation the main source of improvement?

  * Was the VIB loss effective at all? What happens if it is removed? How much of the improvements could be attributed to this structuring of the latent space? Is there data demonstrating the inherent risk by enforcing this loss too weakly or too forcefully?

  * Were the FGM components of the model effective at "exploiting the structure of the target black-box function"? Were there any practical results, *in this paper*, answering this specific claim?

  * Did the method actually result in better coverage of the cliques data, and contributed to the distribution-shift issue? Figure 1 alone seems rather insufficient to address this broad question.

2. The paper is based more around "proposing a solution", than the proper scientific steps of (1) identifying a narrow, previously unknown, problem in the prior methods or datasets, (2) showing evidence about the existence of the problem, and then (3) proposing a solution and demonstrating its efficacy at addressing this issue. 

3. The experimental breadth and the considered methods are awfully inadequate. At this point, there are tens, if not hundreds, of other successful MBO methods, and reward weighted regression and grad asc, for instance, are irrelevant and not remotely representatives of the state of the literature. Since there is a strong emphasis on (1) RL strategies for MBO, and (2) the fact that many of the baselines are sequence-based, the authors should add relevant baselines such as the following:

  * DynaPPO (Model-based reinforcement learning for biological sequence design, ICLR 2020)

  * Ada-Lead (https://arxiv.org/pdf/2010.02141) in FLEXS (https://flexs.readthedocs.io/en/latest/index.html) has a standard implementation and is another contender. 

  * PEX (Proximal Exploration for Model-guided Protein Sequence Design, ICML 2022)

  * FBGAN (Feedback gan (fbgan) for dna: A novel feedback-loop architecture for optimizing protein functions)

  * Other rather weak and old baselines are Cbas, Dbas, CEM-PI, etc.

4. Theorem 1 is a restatement of a basic result in batch reinforcement learning, and offers a rather weak upper bound. While such results are helpful in deriving RL theory, they're nowhere close to being practical; the $C_{stat}$ and $C_{cpx}$ can be arbitrarily large, and the maximum term is almost impossible to estimate. Since this bound is practically almost never tight, it makes the claims and contributions in this paper less verifiable.

**Recommendation**

I feel ambivalent about this paper. On one hand, I see that the authors clearly spent reasonable time and efforts developing an idea, implementing the model, and writing the paper. There is clearly some quality effort spent here. On the other hand, the scientific skeleton of the paper and the supporting evidence is significantly lacking, as I described earlier, and just giving a pass here seems like quite the negligence on my part.

After much delibration, I kindly decided to give the work a weak rejection for now. That being said, I'm keeping an open mind, and I would be happy to raise my score if the authors deliver a strong rebuttal, i.e., a rebuttal that adds significant amounts of supporting experiments and data to the paper and fills up the glaring void. I hope the authors understand that they need much more experimental results in this paper for it to be published at a top-tier venue.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: * How sensitive is the method to the choice of predefined FGM structure? What guidelines would you suggest for selecting appropriate clique sizes and overlaps for new problems?
* How does the computational cost of Cliqueformer compare to baselines? 
* It is not clear to me how this approach would be applicable to problems where the design space is a bit more unstructured (e.g. graphs). Do you have any thoughts on that?
soundness: 2
strengths: * Cliqueformer is a novel approach to handling distribution shift in MBO by learning structured representations rather than using conservative regularization. 
* The training objective is quite interesting, combining useful ideas from VIB with FGMs.
* The paper is generally well written and easy to follow. 
* The submission is accompanied by code to reproduce the results with adequate instructions to reproduce the results.
summary: The paper studies the problem of offline model based optimization. The author first consider outline shortcomings of naively using the perspective of recently proposed functional graphical models FGM) perspective for model based optimization. Specifically, the authors show that the FGM is not an attribute of a function to be inferrred but depends on the parameterization of the input. The paper thus considers fixing the structure of the FGM (as a hyperparameters). The paper introduces a transformer-based arhcitecture called Cliqueformer and proposes a training scheme. Unlike previous approaches that try to discover FGMs directly, Cliqueformer learns latent representations that align with a predefined FGM structure. The model combines transformer blocks with a decomposed predictive module and a novel form of variational bottleneck applied to individual cliques. The approach achieves strong performance across various tasks, including synthetic functions, chemical design, and genetic sequence optimization.
weaknesses: * The authors acknowledge that optimal clique configurations vary between tasks, but use a general set of values for their experiments but is very much unclear to me how the parameters would be selected for a new problem? To be clear, this is a challenge in all offline MBO algorithms where selecting the hyperparamters is not trivial (since that requires access to online evaluations) but with the proposed approach there are a lot more hyperparameters (e.g. knowing the optimal clique configurations, the transformer parameters, etc) which makes the results less convincing.  More systematic study of hyperparameter sensitivity would be valuable (especially related to the transformer configuration. 
* Another significant flaw is the evaluation of the empirical results and discussion of related works. First the paper does a poor job of covering the offline MBO literature. Several recent works are not discussed in the related work (some cited below). The authors also choose just two baselines, which are quite inadequate considering the highly active nature of the field. To be clear, the issue is that the paper claims state-of-the-art performance but fails to compare to the state-of-the-art. More
* Another important aspect missing from the empirical evaluation is a discussion of the computational complexity of the approach. Specifically how does the runtime of using Cliqueformer compare with the runtime of the baselines. Especially since the Cliqueformer uses a whole transformer, whereas (if I understood correctly) the COM baseline is simply a MLP. 
* Finally, there are no formal guarantees about learned representations matching desired FGM structure. 

Minor: 

* This may be considered a minor nitpick but result of Theorem doesn't seem substantial enough to be a theorem. I would suggest calling it a Proposition. 
* The approach of learning a structured representation of a function is very common in the causality and graphical models literature. Even though the paper is specifically about MBO, I encourage the authors to add some discussion about this in the related work section. 
* L118: methos -> methods
* L120: com -> COM 
* L126-7: “This is particularly frustrating since recent work on” -> why frustrating? 
* Figure 3: Axis labels and ticks are too small. 
* L316: designign -> designing 


[1] Kim, M., Berto, F., Ahn, S., & Park, J. (2023). Bootstrapped training of score-conditioned generator for offline design of biological sequences. Advances in Neural Information Processing Systems, 36.

[2] Chen, C. S., Beckham, C., Liu, Z., Liu, X. S., & Pal, C. (2023). Parallel-mentoring for offline model-based optimization. Advances in Neural Information Processing Systems, 36.

[3] Chen, C., Zhang, Y., Liu, X., & Coates, M. (2023, July). Bidirectional learning for offline model-based biological sequence design. In International Conference on Machine Learning (pp. 5351-5366). PMLR.

[4] Yuan, Y., Chen, C. S., Liu, Z., Neiswanger, W., & Liu, X. S. (2023). Importance-aware co-teaching for offline model-based optimization. Advances in Neural Information Processing Systems, 36.

Meta-review: This paper introduces Cliqueformer, a new architecture for model-based optimization (MBO) problems such as chemical and genetic design tasks. The proposed model aims to improve out-of-distribution generalization by learning and predicting a functional graphical model of the underlying black-box function. The field of model-based optimization is relatively new, so there is potential for this method. However, this paper would greatly benefit from another peer review before it is accepted for publication. The authors performed limited experiments, few ablation studies focused on hyperparameter sensitivity and did not describe computational complexity. Most of the paper (the first six pages) focuses on intuitions and desiderata of MBO methods. I suggest that the authors use the space in their paper more strategically, dedicating more space to results, hyperparameter studies, and *experimental evidence* as to why their approach is significant and of interest to the community.

Decision: Reject