redacted dwf<EMAIL_ADDRESS>
# Solving MaxSAT with Matrix Multiplication

David Warde-Farley  Equal contribution  Google DeepMind  Vinod Nair  Equal
contribution  Google DeepMind  Yujia Li  Ivan Lobov  Felix Gimeno  Simon
Osindero

###### Abstract

We propose an incomplete algorithm for Maximum Satisfiability (MaxSAT)
specifically designed to run on neural network accelerators such as GPUs and
TPUs. Given a MaxSAT problem instance in conjunctive normal form, our
procedure constructs a Restricted Boltzmann Machine (RBM) with an equilibrium
distribution wherein the probability of a Boolean assignment is exponential in
the number of clauses it satisfies. Block Gibbs sampling is used to
stochastically search the space of assignments with parallel Markov chains.
Since matrix multiplication is the main computational primitive for block
Gibbs sampling in an RBM, our approach leads to an elegantly simple algorithm
(40 lines of JAX) well-suited for neural network accelerators. Theoretical
results about RBMs guarantee that the required number of visible and hidden
units of the RBM scale only linearly with the number of variables and
constant-sized clauses in the MaxSAT instance, ensuring that the computational
cost of a Gibbs step scales reasonably with the instance size. Search
throughput can be increased by batching parallel chains within a single
accelerator as well as by distributing them across multiple accelerators. As a
further enhancement, a heuristic based on unit propagation running on CPU is
periodically applied to the sampled assignments. Our approach, which we term
_RbmSAT_ , is a new design point in the algorithm-hardware co-design space for
MaxSAT. We present timed results on a subset of problem instances from the
annual MaxSAT Evaluation’s Incomplete Unweighted Track for the years 2018 to
2021. When allotted the same running time and CPU compute budget (but no
TPUs), RbmSAT outperforms other participating solvers on problems drawn from
three out of the four years’ competitions. Given the same running time on a
TPU cluster for which RbmSAT is uniquely designed, it outperforms all solvers
on problems drawn from all four years.

##  1 Introduction

Given a Boolean formula $f(\mathbf{x})$ over a set of $N$ variables
$\mathbf{x}\in\\{\text{True},\text{False}\\}^{N}$, _Boolean Satisfiability_
(SAT) is the problem of deciding if there exists an assignment of $\mathbf{x}$
such that $f(\mathbf{x})=\text{True}$. It was the first problem shown to be
NP-complete [19], and SAT solvers have a wide range of practical applications
[38, 13]. Typically $f(\mathbf{x})$ is in _Conjunctive Normal Form_ (CNF),
which is a conjunction (AND) of _clauses_ , where each clause is a disjunction
(OR) of a subset of variables or their negations. _Maximum Satisfiability_
(MaxSAT) [10] generalizes SAT, requiring determination of the maximum number
of clauses of $f(\mathbf{x})$ that can be satisfied with a single assignment.
A weighted variant assigns weights to clauses, with solvers aiming to maximize
the total weight of satisfied clauses. An _incomplete_ algorithm [32] for
MaxSAT aims to find a good solution quickly without attempting to prove its
optimality.

State-of-the-art solvers for SAT and MaxSAT have been designed with decades of
research to run efficiently on CPUs (see, e.g., annual SAT Competitions [4]
since 2002 and MaxSAT Evaluations [2] since 2006, both of which evaluate CPU-
based solvers). Both algorithmic and hardware improvements have enabled
solvers to scale to larger problems [34, 22] and solve them faster, with one
study estimating that both kinds of improvements have contributed roughly
equally [22]. However, as Moore’s law slows down, CPU improvements can no
longer be relied upon to deliver significant performance gains for solvers.

Special-purpose hardware accelerators are a potential path forward. In
particular, neural network accelerators such as Tensor Processing Units (TPUs)
[31] and Graphics Processing Units (GPUs), provide immense throughput for
specific operations such as matrix multiplication. For example, the TPUv4i
chip has a peak performance of 138 TeraFLOPS, with a cluster of 1024 chips
providing a peak performance of $>100$ PetaFLOPS [30]. Both GPUs and TPUs have
also shown 5-10$\times$ performance improvement over the past several years
[3, 30], offering a tempting alternative to the diminishing returns of Moore’s
law. How to best implement a MaxSAT solver on neural network accelerators is
an open question. Given the CPU-centric design of state-of-the-art solvers and
their use of search algorithms that are not easy to express as matrix
multiplication (e.g., tree search), attempting a direct retargeting of such
solvers is unlikely to be effective [50].

Figure 1: Given a Boolean formula $f(\mathbf{x})$ in conjunctive normal form
with $N$ variables and $C$ clauses ($N=4$, $C=3$ above), RbmSAT constructs a
Restricted Boltzmann Machine (RBM) $p_{\theta}(\mathbf{x})$ with $\mathbf{x}$
as its binary visible units, weights $\theta$, and $O(K)$ binary hidden units
per clause of size $K$ ($K=3$ above). By construction $p_{\theta}(\mathbf{x})$
gives higher probability to an assignment of $\mathbf{x}$ that satisfies more
clauses. RBM’s block Gibbs sampling is used to stochastically search for high
probability assignments, which requires matrix multiplication as the main
operation and can be done scalably on a TPU/GPU cluster. Samples are
periodically updated with a heuristic (running on a CPU) that unassigns values
for a subset of variables in a Gibbs chain and re-assigns them values
logically implied by the assigned variables. The sampling iterations continue
till the time limit, when the best assignment seen so far is output.

We propose an incomplete algorithm for unweighted MaxSAT uniquely designed for
neural network accelerators (Figure 1). The key idea is elegantly simple:
given $f(\mathbf{x})$ in CNF with $N$ variables and $C$ clauses, a
distribution $p_{\theta}(\mathbf{x})$ parameterized by $\theta$ is constructed
as a _Restricted Boltzmann Machine_ (RBM) [49, 23, 26], an energy-based
probabilistic model for high-dimensional distributions [35, 41]. $\theta$ is
chosen such that the more clauses an assignment satisfies, the higher its
probability under $p_{\theta}(\mathbf{x})$. Theoretical results on the
representational capabilities of RBMs [40, 24] guarantee that such an RBM can
be constructed with $N$ visible units and $O(KC)$ hidden units for clause size
$K$ (section 4). Gibbs sampling for $p_{\theta}(\mathbf{x})$ is used to find
highly probable, and thus highly satisfying, variable assignments. The RBM’s
bipartite graph structure between visible and hidden units enable efficient
block Gibbs sampling using matrix multiplication (section 3.1). The result is
a scalable stochastic search algorithm for MaxSAT well-suited for a TPU/GPU
cluster. We further augment the search by periodically improving the sampled
assignments with a simple classical heuristic (section 4.4). We refer to the
algorithm as _RbmSAT_. See section 6.1, figure 5 for a simple 40-line JAX
implementation.

This paper introduces the concepts behind RbmSAT and presents an initial
evaluation on a subset of unweighted problem instances (selected by size to
fit our implementation constraints, see section 5) from the Incomplete
Unweighted Track of four recent annual MaxSAT Evaluations (2018-2021, see
[2]). Following the methodology of the MaxSAT Evaluations, we compare RbmSAT
to the participating solvers by the _average incomplete score_ (section 5) at
a wall clock time limit of 300 seconds. The results (section 5.1) show that in
a controlled comparison using the same _CPU_ , memory, and time budget as the
other solvers, RbmSAT achieves the best average incomplete score for problems
drawn from three out of the four years. In an uncontrolled comparison for the
same 300 second time limit but that provides RbmSAT the advantage of a 64-TPU
cluster, it outperforms all other solvers on problems drawn from all four
years. Since both comparisons are on a subset of instances that omit the
largest ones, and the second comparison is across solvers that run on very
different hardware (CPUs and TPUs) that are inherently difficult to compare,
we do _not_ claim that RbmSAT is a state-of-the-art solver. We nonetheless
believe these results provide evidence for the viability of RbmSAT that
justifies further research.

Our contributions are the following:

  1. 1.

We propose RbmSAT, a MaxSAT algorithm specifically designed for neural network
accelerators.

  2. 2.

We describe a practical implementation of RbmSAT that easily scales parallel
Gibbs sampling with more accelerators.

  3. 3.

On the Incomplete Unweighted Track of MaxSAT Evaluations 2018 to 2021, RbmSAT
outperforms all participating solvers on a subset of problem instances using a
64-TPU cluster in a comparison that controls for wall clock time but not
compute, and remains competitive even when deployed on a traditional CPU.

##  2 Related Work

The literature on SAT and MaxSAT is vast [14]. Many modern MaxSAT solvers rely
on SAT solvers internally by iteratively augmenting the input problem with
extra constraints and applying a SAT solver to improve upon the previous
solution [10]. Popular approaches for SAT broadly belong to two classes: 1)
tree search based on the Conflict Driven Clause Learning (CDCL) algorithm
(see, e.g. [39]), and 2) stochastic local search (see, e.g., [28]). It is not
immediately obvious how such solvers can make efficient use of neural network
accelerators [50]. Approaches that use large-scale parallel CPU compute, such
as Cube and Conquer [25] and entries in the Parallel Track of SAT
competitions, typically retrofit a conventional approach with parallel
compute, e.g. by processing sub-trees of a search tree in parallel, but they
have not been broadly adopted. Hardware acceleration (e.g., FPGA) has been
used to speed up _unit propagation_ which is a computational bottleneck in
CDCL-based solvers [50]. The classical heuristic upon which RbmSAT relies
(section 4.4) is based on unit propagation and can also benefit from such
accelerators.

Ising models have been used to encode combinatorial optimization problems such
that low energy configurations of the Ising model correspond to high quality
solutions, see, e.g., [37]. Early examples are Hopfield networks for the
Travelling Salesman Problem [29] and (general) Boltzmann Machines for SAT and
MaxSAT [20]. Approaches for simulating Ising models using GPUs [18] and TPUs
[53] have been proposed. RBM’s bipartite graph structure differs from the
lattice structure of a 2D Ising model, and admits efficient block Gibbs
sampling using matrix multiplication, making it a natural fit for neural
network accelerators unlike Ising models and Boltzmann Machines (without
connectivity constraints). It is also possible to encode an Ising model’s
energy function as an RBM’s free energy function using the results from [40,
24] and extend our approach to Ising models.

Survey Propagation [16] constructs a factor graph for a Boolean Satisfiability
problem with variables as one set of nodes and clauses as the other. The
defined distribution assigns zero probability to unsatisfying assignments and
uniform probability to satisfying assignments. Message passing is used to
estimate the marginal distributions of the variables. Those variables with
high marginal probability of being 0 or 1 are fixed to their maximum a
posteriori values, and the resulting (typically much) smaller sub-problem is
solved using other approaches. The distribution defined by RbmSAT assigns non-
zero probability to all assignments, but with the probability of an assignment
exponential in the number of satisfied clauses. It matches the factor graph
distribution only at the zero temperature limit.

Several works learn search policies for SAT that generalize across problem
instances (e.g. [47, 46, 6, 54, 36]). In this work we do not learn a search
policy, but such methods are complementary and can be combined with RbmSAT
(e.g. by learning to generate good initial assignments for RbmSAT).

##  3 Background

###  3.1 Restricted Boltzmann Machines

_Boltzmann Machines_ [27] are undirected graphical models defined by the
energy function
$E_{\theta}(\mathbf{x})=-{\mathbf{x}}^{\mathsf{T}}\mathbf{S}\mathbf{x}-{\mathbf{a}}^{\mathsf{T}}\mathbf{x}$,
on variables $\mathbf{x}\in\\{0,1\\}^{D}$, parameterized by
$\theta=\\{\mathbf{S},\mathbf{a}\\}$ where $\mathbf{S}\in\mathbb{R}^{D\times
D}$ is symmetric and $\mathbf{a}\in\mathbb{R}^{D}$. It defines a probability
distribution $p_{\theta}(\mathbf{x})=\exp(-E_{\theta}(\mathbf{x}))/Z(\theta)$,
where
$Z(\theta)=\sum_{\mathbf{x}^{\prime}}\exp(-E_{\theta}(\mathbf{x}^{\prime}))$
is the partition function. For large $D$ computing $Z(\theta)$ requires an
intractable summation, making $p_{\theta}(\mathbf{x})$ and its gradient
(needed to learn $\theta$) intractable as well.

_Restricted Boltzmann Machines_ (RBMs) [49] are a special case in which
$\mathbf{x}$ consists of observed and latent variables and the model imposes a
bipartite connectivity structure on the two groups: no connections are
permitted between two observed variable elements of $\mathbf{x}$ _or_ between
two latent variable elements of $\mathbf{x}$. Denoting the observed variables
as _visible units_ $\mathbf{v}\in\\{0,1\\}^{N}$ and the latent variables as
_hidden units_ $\mathbf{h}\in\\{0,1\\}^{N_{h}}$, the energy function can now
be written as
$E_{\theta}(\mathbf{v},\mathbf{h})=-{\mathbf{v}}^{\mathsf{T}}\mathbf{W}\mathbf{h}-{\mathbf{v}}^{\mathsf{T}}\mathbf{d}-{\mathbf{h}}^{\mathsf{T}}\mathbf{b}$,
where $\theta=\\{\mathbf{W},\mathbf{d},\mathbf{b}\\}$,
$\mathbf{W}\in\mathbb{R}^{N\times N_{h}},\mathbf{d}\in\mathbb{R}^{N}$ and
$\mathbf{b}\in\mathbb{R}^{N_{h}}$. The _free energy_ is the unnormalized
negative log probability of $\mathbf{v}$, and is defined as
$F_{\mathbf{\theta}}(\mathbf{v})=-\log\sum_{\mathbf{h}^{\prime}}\exp\left(-E_{\mathbf{\theta}}(\mathbf{v},\mathbf{h}^{\prime})\right)$.
Given the RBM’s bipartite graph structure this simplifies to the tractable
expression

$\displaystyle
F_{\mathbf{\theta}}(\mathbf{v})=-{\mathbf{d}}^{\mathsf{T}}\mathbf{v}-\sum_{j=1}^{N_{h}}\log\left[1+\exp\left(b_{j}+\sum_{k=1}^{N}w_{kj}v_{k}\right)\right]$
(1)

The RBM’s structure also renders the joint conditional distributions
$p_{\theta}(\mathbf{v}|\mathbf{h})$ and $p_{\theta}(\mathbf{h}|\mathbf{v})$
tractable. This enables parallel _block_ Gibbs sampling. The transition from a
sample $\tilde{\mathbf{v}}_{t}$ at step $t$ to the sample
$\tilde{\mathbf{v}}_{t+1}$ at $t+1$ is given by

$\displaystyle\tilde{\mathbf{h}}_{t}$ $\displaystyle\sim
p_{\theta}(\mathbf{h}\ |\
\mathbf{v}=\tilde{\mathbf{v}_{t}})=\sigma(\mathbf{b}+{\mathbf{W}}^{\mathsf{T}}\tilde{\mathbf{v}_{t}})$
(2) $\displaystyle\tilde{\mathbf{v}}_{t+1}$ $\displaystyle\sim
p_{\theta}(\mathbf{v}\ |\
\mathbf{h}=\tilde{\mathbf{h}_{t}})=\sigma(\mathbf{d}+\mathbf{W}\tilde{\mathbf{h}_{t}})$
(3)

where $\sigma(x)=(1+e^{-x})^{-1}$ is the logistic function, applied element-
wise. RBM sampling is well suited to neural network accelerators as the
majority of operations necessary to compute (2) and (3) are expressible as
matrix multiplications.

RBMs can be viewed as _products of experts_ [26]: each hidden unit (and
visible unit bias $d_{j}$) contributes a multiplicative constraint on the
unnormalized marginal distribution of visible units. The product of several
RBMs is itself, therefore, an RBM. We use this property to construct an RBM
representing a Boolean formula in CNF as a product of RBMs, each representing
one clause in the formula.

##  4 RbmSAT

###  4.1 MaxSAT as Minimizing RBM Free Energy

Theoretical results on the representational efficiency of RBMs [40, 24] form
the basis of our approach. In particular, [40] proves that any distribution
$p(\mathbf{v})$ over $\mathbf{v}\in\\{0,1\\}^{N}$ with a _symmetric_ free
energy function can be represented by an RBM $p_{\theta}(\mathbf{v})$ with $N$
visible units and $O(N^{2})$ hidden units. A function $g(\mathbf{v})$ is
symmetric if and only if $g(\mathbf{v})=q(\mathcal{V})$ where
$\mathcal{V}=\sum_{i=1}^{N}{v_{i}}$, i.e., $g(\cdot)$ depends only on the
number of ones in its input. Subsequently, [24] improves the result to $O(N)$
hidden units. Both papers provide analytic procedures for constructing such an
RBM and its parameters given the desired free energy function.

Figure 2: Accelerator-side computation of a single Gibbs sampling chain
broadly consists of gather/scatter operations (blue arrows), clause-RBM
batched weight matrix multiplications (orange arrows) and element-wise
Bernoulli sampling (purple arrows). An initial assignment vector is gathered
into a matrix of clause-RBM inputs via batched matrix multiplication, while a
second batched matrix multiply, along with the addition of a per-clause bias
vector $\mathbf{b}^{(c)}$ yields logits for the conditional distribution over
clause-RBM hidden units. A configuration of the hidden units is sampled and
the transpose of the two aforementioned matrix multiplies are applied yielding
a set of visible unit logits, from which a new assignment of the variables is
sampled. In practice, multiple chains are run in parallel on a mini-batch of
variable assignments.

#### OR gate as an RBM:

Boolean OR is a symmetric function. To see this let 0 denote $\mathrm{False}$
and 1 denote $\mathrm{True}$, and note that for an input
$\mathbf{v}\in\\{0,1\\}^{N}$, the gate’s output depends only on whether
$\sum_{i=1}^{N}{v_{i}}>0$. Let $p_{\mathrm{OR}}(\mathbf{v})$ be a distribution
with a free energy function $F_{\mathrm{OR}}(\mathbf{v})$ such that
$F_{\mathrm{OR}}(\mathbf{v}=\mathbf{0})=0$, corresponding to the only
unsatisfying input, and $F_{\mathrm{OR}}(\mathbf{v}\neq\mathbf{0})=F_{s}$
where $F_{s}<0$ is the free energy for all other inputs. $p_{\mathrm{OR}}$
thus assigns equal probability to all satisfying assignments and a lower
probability to the single unsatisfying assignment. The corresponding
$p_{\mathrm{OR}}$ We use the results of [40, 24] to construct an RBM
$p_{\theta}(\mathbf{v})=p_{\mathrm{OR}}(\mathbf{v})$. $F_{s}$ acts as a
temperature parameter: $F_{s}=0$ yields a uniform distribution (the
unsatisfying assignment has equal probability to any satisfying assignment),
and lower values reduce $p_{\theta}(\mathbf{v}=\mathbf{0})$.

Given the parameters $\mathbf{W},\mathbf{b}$ (visible biases $\mathbf{d}$ are
not required) of an RBM for a canonical OR gate that does not negate any
inputs, the parameters of an RBM for an OR gate involving negations of one or
more inputs can be obtained by a simple transformation

$\displaystyle\begin{array}[]{c
c}\mathbf{W}^{\prime}=\mathbf{\Lambda}\mathbf{W};&\mathbf{b}^{\prime}=\mathbf{b}+\textstyle{\frac{1}{2}}\left(1-\mathrm{diag}\left(\mathbf{\Lambda}\right)\right)\mathbf{W}\end{array}$
(5)

where $\mathbf{\Lambda}$ is a diagonal matrix with its $k^{th}$ diagonal
element $\lambda_{kk}\in\\{-1,1\\}$ set to $-1$ if the $k^{th}$ input to the
OR gate is negated.

#### CNF Boolean formula as an RBM:

A “formula-RBM” can be constructed for a Boolean formula in CNF by a)
constructing a “clause-RBM” for each of the $C$ clauses (an OR gate) with $L$
hidden units and the $K$ variables in the clause as its visible units, and b)
taking the product of the clause-RBMs by concatenating their hidden units (for
a total of $C\cdot L$ hidden units) and merging the overlapping sets of
variables participating in each clause into a vector of $N$ visible units. The
free energy of the product is simply the sum of the free energies of the
clause-RBMs. Each satisfied clause contributes $F_{s}<0$ to the total free
energy, with any unsatisfied clauses contributing zero. The total free energy
is thus monotonically decreasing, and probability monotonically increasing, in
the number of satisfied clauses of an assignment.

Let $\theta_{f}=\\{\mathbf{W}_{f},\mathbf{d}_{f},\mathbf{b}_{f}\\}$ be the
parameters of the formula-RBM. Let
$\mathbf{T}\in(\mathbb{N}\cup\\{0\\})^{C\times K}$. We have $t_{ck}\leq N$
$\forall c,k$, and if $t_{ck}=0$ this indicates that group $c$ contains fewer
than $K$ variables (i.e. the $0$ values are used as padding). Let
$\mathbf{W}^{({c})}_{f}$ be the (in general, non-contiguous) sub-matrix of
$\mathbf{W}_{f}$ with rows corresponding to the non-zero elements of
$\mathbf{T}_{c,\cdot}$, and columns $\mathbf{u}_{c}\leavevmode\nobreak\
=\leavevmode\nobreak\ {((c\leavevmode\nobreak\ -\leavevmode\nobreak\
1)L,(c\leavevmode\nobreak\ -\leavevmode\nobreak\ 1)L\leavevmode\nobreak\
+\leavevmode\nobreak\ 1,\ldots,cL)}^{\mathsf{T}}$. Similarly, define (for
$j\neq 0$)

$\displaystyle\begin{array}[]{rcl
rcl}\mathbf{v}^{({c})}&=&{({v}_{j})}^{\mathsf{T}}_{{j}\in{\mathbf{T}_{c,\cdot}}}&\mathbf{d}^{({c})}&=&{({d}_{j})}^{\mathsf{T}}_{{j}\in{\mathbf{T}_{c,\cdot}}}\\\
\mathbf{b}^{({c})}&=&{({b}_{j})}^{\mathsf{T}}_{{j}\in{\mathbf{u}_{c}}}&\mathbf{h}^{({c})}&=&{({h}_{j})}^{\mathsf{T}}_{{j}\in{\mathbf{u}_{c}}}\end{array}$

Denoting
$\theta_{f}^{(c)}=\\{\mathbf{W}^{({c})}_{f},\mathbf{d}^{({c})}_{f},\mathbf{b}^{({c})}_{f}\\}$,
the free energy of the formula-RBM is given by:

$\displaystyle F_{\theta_{f}}(\mathbf{v})=$
$\displaystyle\sum_{c=1}^{C}F_{\theta^{(c)}_{f}}(\mathbf{v}^{({c})})$ (6)

#### Gibbs sampling in a formula-RBM:

Given the groups
$\mathbf{v}^{({1})},\mathbf{v}^{({2})},\ldots,\mathbf{v}^{({C})}$, sampling
values for $\mathbf{h}$ amounts to sampling each disjoint group of units
$\mathbf{h}^{({1})},\mathbf{h}^{({2})},\ldots,\mathbf{h}^{({C})}$ from its
respective clause-RBM:

$\displaystyle p_{\theta_{f}}(\mathbf{h}\ |\ \mathbf{v})=$
$\displaystyle\prod_{c=1}^{C}p_{\theta^{(c)}_{f}}\left(\mathbf{h}^{({c})}\ |\
\mathbf{v}^{({c})}\right)$ (7)

Sampling $\mathbf{v}$ given $\mathbf{h}$ poses a greater challenge, as every
visible unit $v_{i}$ may participate in many clause-RBMs. Let
$\ell^{({c})}=\mathbf{W}_{f}^{(c)}\mathbf{h}^{({c})}$. The logits of the
conditional distribution decompose as a sum of contributions from each
$\ell^{({c})}$ that contains $v_{i}$, added to a global bias $d_{i}$:

$\displaystyle p_{\theta}(v_{i}\ |\
\mathbf{h})=\sigma\left(d_{i}+\sum_{c=1}^{C}\sum_{k=1}^{K}\delta_{i,t_{ck}}\ell_{k}^{({c})}\right)$
(8)

with $\delta_{i,j}=1$ for $i=j$ and 0 for $i\neq j$. A practical
implementation of RBM operations with these _a priori_ connectivity
constraints need never instantiate the $N\times CL$ weight matrix
$\mathbf{W}_{f}$, which will be very large for large $N$ and $C$, and sparse
for $K\ll N$. Instead, each sub-matrix $\mathbf{W}_{f}^{(c)}$ can be stored
contiguously and zero entries of $\mathbf{W}_{f}$ represented only implicitly.

#### Learning OR gate RBM parameters:

The RBM parameters for an OR gate can be computed analytically using the
procedure in [40] or [24]. However, we have observed empirically that learning
them by directly fitting the free energy function for an OR gate yields RBMs
for which Gibbs sampling mixes faster, which is crucial for efficient search.
The analytically computed weights have high magnitude, which slows mixing.
Typically SAT and MaxSAT problem instances have a small enough number of
variables $K$ in a single clause (e.g., we assume $K<=7$ in this work) that
all possible $2^{K}$ clause inputs can be exhaustively enumerated, in which
case the RBM parameters $\theta$ can be learned by regressing the free energy
to the target values as
$\theta_{F_{s}}=\operatorname*{arg\,min}_{\mathbf{\theta}}\|F_{\mathbf{\theta}}(\mathbf{0})\|^{2}+\sum_{\mathbf{v}\neq\mathbf{0}}\|F_{\mathbf{\theta}}(\mathbf{v})-F_{s}\|^{2}$.

An equivalent concatenative approach to constructing RBMs for representing
Boolean circuits from smaller pre-trained RBMs was proposed by [44], who
perform learning of the constituent RBMs using the contrastive divergence
approximation to the maximum likelihood gradient. This unfortunately
introduces unnecessary noise and does not optimize the RBM log likelihood (or,
indeed, any objective function [51]). If $K$ is too large to allow for
enumeration, it is possible to relax the fully deterministic learning
described above while still avoiding MCMC gradient approximations by sampling
satisfying assignments uniformly at random on each learning step, but we have
not yet empirically examined this possibility.

###  4.2 Stochastic Search with Gibbs Sampling

Performing stochastic search consists of instantiating parallel Markov chains,
initialized from a uniform distribution over assignments, and performing block
Gibbs sampling as explained above. We use an ensemble of values for the
temperature $F_{s}$, train an OR gate RBM for each temperature, and
instantiate and simulate all corresponding formula-RBMs in parallel. The
clause satisfaction count of an assignment can be computed at each step with
very little overhead, as explained below. We retain the best assignment seen
by any Markov chain, and output the best identified assignment once we have
exhausted the time limit.

###  4.3 Scalable TPU Implementation

We implement our approach in JAX [15] targeting TPUs [1], proprietary ASICs
developed by Google to accelerate neural network workloads, which rely heavily
on matrix multiplication. JAX programs are just-in-time compiled using the XLA
(Accelerated Linear Algebra) [5] domain-specific compiler. A simplified
version of the RbmSAT code is provided in Figure 5 in the Appendix.

We perform block Gibbs sampling in multiple RBMs trained with different OR
gate satisfying free energy targets synchronously on slices of 64 TPUv4 chips
containing 128 TPU cores. At a high level, performing a single full round of
Gibbs sampling involves gathering variables for each clause, sampling hidden
states associated with each clause, computing logit contributions for each
participating variable of each clause, summing the logit contributions across
clauses with a scatter-add operation, and sampling new visible states.

Weights of the RBMs are assembled from pretrained weights for a canonical OR
gate with a suitable maximum clause size $K$ and the desired free energy
target, and transformed as needed according to (5). A clause with fewer than
$K$ participating variables can be padded to length $K$ by repeating one of
its variables with the same polarity, which leaves the set of satisfying
assignments invariant.

For the gather and scatter operations, our implementation stores a table
$\mathcal{T}\in\mathbb{Z}^{K\times C}$ of signed integers111Typically $K<<C$,
storing this table variable-major, clause-minor incurs less losses to padding
given the TPU’s tiled memory model, which is allocated in blocks of $8\times
128$., with
$\mathbf{T}={(\mathrm{sign}(\mathcal{T})\odot\mathcal{T})}^{\mathsf{T}}$, and
signs representing the polarity (negation) of each position in the clause,
i.e. a column ${(1,2,-3,4)}^{\mathsf{T}}$ represents the clause $v_{1}\vee
v_{2}\vee\neg v_{3}\vee v_{4}$. Both the gather and scatter-add operations can
be implemented via batched matrix multiplication (via XLA’s DotGeneral
primitive, exposed by JAX’s einsum operator) on the TPU’s matrix
multiplication unit (MXU) by transforming $\mathcal{T}$ into a 3-dimensional
one-hot representation $\mathbf{Q}\in\\{0,1\\}^{C\times N\times K}$ where
$q_{cik}=\delta_{i,t_{ck}}$. Conveniently, the XLA compiler fuses the creation
of this intermediate array (implemented as a broadcasted integer equality
test) with the DotGeneral operation, eliminating the memory cost of
instantiating $\mathbf{Q}$ explicitly. To perform a gather, a batch of
assignments of size $B\times N$ is right-multiplied by each slice
$\mathbf{Q}_{c,\cdot,\cdot}$, yielding a result of size $B\times C\times K$. A
scatter-add is performed, conceptually, by right-multiplying each set of logit
contributions $\ell^{({c})}$, for each batch item and each clause $c$, by
${\mathbf{Q}}^{\mathsf{T}}_{c,\cdot,\cdot}$. Storing the polarities in
$\mathcal{T}$ is not necessary for sampling, but storing them on device in
this manner allows for efficient evaluation of the number of satisfying
clauses by performing a parallel reduction following the gather step by
comparing thresholded signs of $\mathcal{T}$ to the Boolean values in each
gathered clause, performing the OR operation for each clause, treating the
Boolean result as integers in $\\{0,1\\}$, and summing the result for each
chain state.

###  4.4 Unit Propagation-based Heuristic

We augment the Gibbs sampling-based stochastic search with a heuristic
proposed in [48]. It is based on _unit propagation_ , which attempts to extend
a partial assignment of the variables for a formula by computing the partial
assignment’s logical implications for the unassigned variables. This is a
sequential procedure: an initial partial assignment may get extended by
implied values for some unassigned variables, which may further imply values
for more variables, and so on. For example, in the expression $(a\vee b\vee
c)\wedge(\neg c\vee\neg d)$, a partial assignment $a=b=\mathrm{False}$ implies
$c$ must be $\mathrm{True}$ for the first clause to be true. However, $c$ is
negated in the second clause, which implies a value of $\mathrm{False}$ for
$d$.

We rank each variable $v_{i}$ by a moving average of the conditional variance
$\nu_{i}=\rho_{i}(1-\rho_{i})$ of the Bernoulli distribution with parameter
$\rho_{i}=p_{\mathbf{\theta}_{f}}(v_{i}=1\ |\ \mathbf{h})$. The variables with
the lowest variance are selected to be unassigned and recomputed using unit
propagation. The intuition is that variables with the lowest average
conditional variance for a given Markov chain flip state infrequently, which
could indicate that the chain is “stuck” due to such variables. Starting from
a partial assignment that excludes such variables and applying unit
propagation can result in a change in value for these “stuck” variables, and
allow the assignment to make highly non-local moves. The
$\lfloor\frac{1}{2}N\rfloor$ variables with the _highest_ value of $\nu_{i}$
are held fixed and the rest are updated. See [48] for more details. Since unit
propagation does not use matrix multiplication, it is not well-suited to run
on a TPU. The heuristic is applied periodically, with the current samples sent
to the host for CPU computation, and the updated samples returned to resume
sampling. Results show that the heuristic significantly improves the
performance of RbmSAT.

##  5 Evaluation

#### Datasets:

We use publicly available instances from the annual MaxSAT Evaluations for the
years 2018 to 2021. Since RbmSAT is an incomplete solver and currently does
not handle weighted clauses, we use instances from the Incomplete Unweighted
Track of each Evaluation222https://maxsat-
evaluations.github.io/2021/tracks.html. As a first step towards demonstrating
the usefulness of our approach, we select a subset of instances with maximum
clause size $\leq 7$, number of variables $\leq 10000$, and number of clauses
$\leq 100000$. We chose this subset such that the parameters of the formula-
RBM and a batch of assignments would fit in the memory of a single TPU core
for all assignments, in order to ease experimentation and resource allocation
decisions. Larger instances can be tackled by sharding clauses across multiple
accelerators at the cost of a cross-device reduction at every sampling step,
but optimizing per-instance Gibbs sampling throughput becomes more challenging
in this regime. With regard to clause size in particular, instances with large
maximum clause size tend to have few large and many small clauses, and thus a
fully homogeneous formula-RBM implementation of the form described here would
incur significant memory overhead due to padding smaller clauses to the size
of the largest clause through variable repetition. This could be addressed by
clustering clauses by size and essentially running several independent
formula-RBMs which are coupled by summing their visible logits, at the cost of
higher implementational complexity.

After selecting problem instances according to the above size criteria and
removing any duplicates, we are left with a total of 92 unique instances
across the four Evaluations. We present results on these instances.

#### Incomplete score metric:

We use the same evaluation metric as the MaxSAT Evaluations, the _incomplete
score_. Let $c_{i}^{j}$ be the number of unsatisfied clauses achieved by the
$i^{th}$ solver on the $j^{th}$ problem instance, and $c_{best}^{j}\leq
c_{i}^{j}$ be the best known number of unsatisfied clauses for that instance.
Then the incomplete score for solver $i$ on problem $j$ is defined as
$s_{i}^{j}=(c_{best}^{j}+1)/(c_{i}^{j}+1)$. Note that $s_{i}^{j}\in(0,1]$,
with higher scores indicating better performance. As in the Evaluations, we
rank solvers in decreasing order of their average incomplete score achieved on
a set of instances at a wall clock time limit of 300 seconds.

#### Baselines:

We compare RbmSAT to the solvers that participated in the MaxSAT Evaluations
2018-2021. The number of unsatisfied clauses achieved by a solver on each
problem instance from a given Evaluation year is publicly available. We use
this data to compute the average incomplete score of the participant solvers
and RbmSAT on the selected instances. Since the set of solvers changes from
year to year, the scores and the ranking of the solvers are computed
separately for each year. The participant solvers use state-of-the art SAT
solvers and stochastic local search algorithms and represent the current
state-of-the-art for MaxSAT. Detailed descriptions can be found in [11, 12, 8,
9].

#### Hardware configurations:

Our best results showcase RbmSAT on each problem instance using a cluster of
64 TPUv4 chips. Solvers participating in Evaluations use conventional CPUs
(Intel Xeon 2.4GHz CPU with 32 GB of RAM for the Incomplete Unweighted Track).
These two hardware configurations are not easily comparable, so we run RbmSAT
on CPU as well using the Evaluation machine configuration, which we denote as
“RbmSAT (CPU)” in the results333Due to the complexity of porting our TPU-
focused implementation to the competition execution environment, these results
were obtained on a closely matched machine configuration, with a same-
generation Intel Xeon CPU available in our own environment, and the same RAM
limit..

#### Hyperparameter tuning:

RbmSAT has several hyperparameters, such as the number of parallel chains to
use444Note that more Markov chains is not always obviously better: running
fewer chains may be preferable if it results in those chains making greater
progress due to a lower time requirement per-step. (i.e. batch size),
frequency of applying the unit propagation-based heuristic to the samples, how
to order variables in the heuristic, etc. We set the values of these
hyperparameters using a grid search for each Evaluation year separately. We
use problem instances from all Evaluations other than the one for which we are
selecting the hyperparameters to evaluate the goodness of a hyperparameter
choice.

###  5.1 Results

Table 1 summarizes the main results of the paper. RbmSAT on 64 TPUs achieves a
better average incomplete score at the 300 second time limit than Evaluation
participants on instances from each of the four Evaluation years. RbmSAT (64
TPUs) achieves substantial improvements over the next best solver for the 2018
and 2019 Evaluations (24.9% and 23.23% better scores, respectively), with
smaller improvements for the 2020 and 2021 Evaluations. It achieves the lowest
number of unsatisfied clauses on 35 out of the 92 instances across all four
Evaluations. RbmSAT on CPU is also competitive, outperforming the other
Evaluation participants on three out of the four years. This demonstrates that
our stochastic search approach itself is powerful even without the compute
throughput of a TPU cluster.

MaxSAT 2018
---
Solver | Average
| incomplete
| score
RbmSAT (64 TPUs) | 0.947
RbmSAT (CPU) | 0.942
SATLike | 0.758
maxroster | 0.706
SATLike-c | 0.699
LinSBPS | 0.624
Open-WBO-Inc-OBV | 0.620
Open-WBO-Riss | 0.599
Open-WBO-Gluc | 0.579
Open-WBO-Inc-MCS | 0.559
MaxSAT 2019
RbmSAT (64 TPUs) | 0.968
RbmSAT (CPU) | 0.908
SATLike | 0.785
sls-mcs | 0.783
sls-mcs-lsu | 0.782
Loandra | 0.762
LinSBPS2018 | 0.745
Open-WBO-g | 0.712
Open-WBO-ms | 0.709
MaxSAT 2020
RbmSAT (64 TPUs) | 0.930
RbmSAT (CPU) | 0.884
sls-mcs | 0.866
SATLike-c | 0.865
sls-lsu | 0.865
StableResolver | 0.863
Loandra | 0.839
TT-Open-WBO | 0.800
MaxSAT 2021
RbmSAT (64 TPUs) | 0.977
StableResolve | 0.975
SATlike-c | 0.973
Satlike-ck | 0.972
TT-Open-WBO-Inc-21 | 0.973
Loandra | 0.940
RbmSAT (CPU) | 0.908
Exact | 0.722
Table 1: Incomplete score at a time limit of 300s for RbmSAT and solvers that
participated in the annual MaxSAT Evaluations 2018-2021, averaged over a
subset of instances from the unweighted, incomplete solver track of the
Evaluations that are selected to have maximum clause size $\leq 7$, number of
variables $\leq 10000$, and number of clauses $\leq 100000$ to fit our
implementation constraints (see section 5). Higher incomplete score $\in(0,1]$
is better. “RbmSAT (64 TPUs)” runs on a cluster of 64 TPUs, while “RbmSAT
(CPU)” runs on a CPU. Incomplete scores are computed using the number of
unsatisfied clauses achieved by the participating solvers on the problem
instances as reported at https://maxsat-evaluations.github.io/ and by RbmSAT.

Table 2 presents an ablation study that evaluates the effect of both the core
RBM sampling component and the prioritized unit propagation heuristic. Where
the RBM is ablated and therefore conditional distribution variances are
unavailable, a uniform random ranking prioritizes variables for unit
propagation instead. As these results demonstrate, both the RBM component and
unit propagation contribute significantly to RbmSAT’s performance. For the
2018 and 2021 Evaluations, ablating the RBM seems to worsen results more than
removing unit propagation, while for the other two Evaluations, the opposite
is the case. These results also demonstrate that while unit propagation is a
powerful improvement operator for candidate assignments, repeated application
of unit propagation to randomly chosen subsets is insufficient: a diverse set
of strong candidates for improvement generated by stochastic search is
critical, with the RBM performing better in this regard than candidates chosen
uniformly at random.

| Average incomplete score
---|---
Solver | MaxSAT 2018 | MaxSAT 2019 | MaxSAT 2020 | MaxSAT 2021
RbmSAT | 0.947 | 0.968 | 0.930 | 0.977
RbmSAT - UP | 0.938 | 0.842 | 0.782 | 0.840
RbmSAT - Gibbs | 0.855 | 0.899 | 0.842 | 0.785
UP only | 0.357 | 0.290 | 0.222 | 0.082
Table 2: Ablation results for RbmSAT on 64 TPUs. “RbmSAT - UP” removes unit
propagation, “RbmSAT - Gibbs” replaces Gibbs sampling with uniform random
sampling, and “UP only” removes sampling and uses only unit propagation with a
prioritization order chosen uniformly randomly.

Figure 3: Average incomplete score achieved by RbmSAT for various batch sizes
(number of parallel Markov chains) per TPU device (left) and number of TPUs
(right).

Figure 4: Number of unsatisfied clauses after 300 seconds on factoring
instances as a function of the number of clauses (left) and number of
variables.

Figure 3 shows the effect of increasing the number of parallel Markov chains
(i.e., the batch size used during sampling) and the number of TPUs on the
average incomplete score. As expected, the results improve with more chains
and more TPUs. Unsurprisingly, the graphs show diminishing returns for both,
particularly for the number of TPUs (which corresponds to larger jumps in the
number of chains).

Figure 4 illustrates preliminary results on the scaling behaviour of RbmSAT as
a function of the problem size in terms of the number of variables and clauses
in the problem. The curve illustrates performance of 64 TPUv2 chips with a 300
second wall clock time budget on an arbitrarily chosen subset of the “long”
factoring instances encoded as SAT problems, as proposed in [42]555These
instances can be downloaded at https://github.com/sebastianv89/factoring-
sat/tree/master/instances where the product to be factored ranges in size from
6 to 100 bits. The number of unsatisfied clauses has a roughly linear
relationship with the number of variables or clauses. However, the problems
used here are relatively small; further study on larger instances is needed to
better understand RbmSAT’s scaling behavior.

##  6 Conclusions and Future Work

We have derived a correspondence between RBMs and MaxSAT problems and proposed
RbmSAT, a highly parallelizable algorithm that can exploit the compute power
of neural network accelerators for MaxSAT. We have demonstrated its efficacy
on instances drawn from four MaxSAT Evaluations, and showcased the scaling
benefits of our approach. Our results, including a baseline which demonstrates
lower but nonetheless surprisingly strong performance simply by filtering
high-throughput randomly generated assignments and improving them with unit
propagation, suggest that our contribution will be far from the last word on
hardware-accelerated stochastic search for MaxSAT and related problems. Just
as deep learning workloads benefited from advances in GPU technology driven by
commodity video gaming hardware, our RBM construction of SAT is well
positioned to exploit the wide availability of accelerated matrix
multiplication hardware primarily built for training and inference in
deterministic deep neural networks.

One future direction is to apply RbmSAT to SAT and #SAT. Though RbmSAT is most
naturally thought of as a MaxSAT search algorithm given its construction of a
distribution with monotonically increasing probability in the number of
satisfied clauses, RbmSAT can identify fully satisfying assignments.
Improvements which encourage assignment diversity may adapt RbmSAT more
readily to SAT. #SAT is the problem of finding the total number of fully
satisfying assignments, which can be posed as the problem of computing the
partition function of the formula-RBM at the zero temperature limit when all
fully satisfying assignments have equal probability while the rest have zero
probability. Techniques such as Annealed Importance Sampling (AIS) [43], which
has already been explored for the estimation of RBM partition functions [45],
can potentially be applied to finding approximate solutions for #SAT.

Our method involves running an ensemble of RBMs trained at different
temperatures, and could in principle be extended to use parallel tempering
[52, 21, 17], which simulates an ensemble of systems at different temperatures
and periodically proposes swaps between particles in adjacent replicas. This
procedure allows particles to mix between disparate modes at higher
temperatures, while particles at lower temperatures explore the fine structure
of the energy landscape within modes, with exchanges allowing a single
particle to mix more freely on finite time horizons than Gibbs dynamics would
ordinarily permit.

## Acknowledgements

We would like to thank Alex Davies for detailed technical feedback on the
manuscript as well as Oriol Vinyals, Guillaume Desjardins and Geoffrey Hinton
for helpful research discussions, Aidan Clark for early engineering
discussions, Jonas Adler, Andy Brock, John Jumper and Sander Dieleman for
valuable advice regarding TPUs, and Richard Tanburn for debugging support. We
would also like to acknowledge the foundational work of the JAX and XLA teams,
in particular Lena Martens, Blake Hechtman and Peter Hawkins, as well as all
contributors to the DeepMind JAX ecosystem [7].

## References

  * [1] Cloud TPU System Architecture.  https://cloud.google.com/tpu/docs/system-architecture-tpu-vm.  Accessed: 2022-05-19.
  * [2] MaxSAT Evaluations.  https://maxsat-evaluations.github.io/.  Accessed: 2022-05-19.
  * [3] Nvidia h100 tensor core gpu architecture.  https://resources.nvidia.com/en-us-tensor-core.  Accessed: 2022-10-13.
  * [4] SAT Competitions.  http://www.satcompetition.org/.  Accessed: 2022-05-19.
  * [5] Xla: Optimizing compiler for machine learning, 2017.
  * [6] Saeed Amizadeh, Sergiy Matusevych, and Markus Weimer.  Learning to solve circuit-sat: An unsupervised differentiable approach.  In ICLR, 2019.
  * [7] Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola.  The DeepMind JAX Ecosystem, 2020.
  * [8] Fahiem Bacchus, Jeremias Berg, Matti Järvisalo, and Ruben Martins.  Maxsat evaluation 2020: Solver and benchmark descriptions.  University of Helsinki Department of Computer Science Report Series B, B-2020-2, 2020.
  * [9] Fahiem Bacchus, Jeremias Berg, Matti Järvisalo, and Ruben Martins.  Maxsat evaluation 2021: Solver and benchmark descriptions.  University of Helsinki Department of Computer Science Report Series B, B-2021-2, 2021.
  * [10] Fahiem Bacchus, Matti Järvisalo, and Ruben Martins.  Maximum Satisfiability, pages 929 – 991.  Frontiers in Artificial Intelligence and Applications. IOS PRESS, Netherlands, 2 edition, 2021.
  * [11] Fahiem Bacchus, Matti Järvisalo, and Ruben Martins.  Maxsat evaluation 2018: Solver and benchmark descriptions.  University of Helsinki Department of Computer Science Series of Publications B, B-2018-2, 2018.
  * [12] Fahiem Bacchus, Matti Järvisalo, and Ruben Martins.  Maxsat evaluation 2019: Solver and benchmark descriptions.  University of Helsinki Department of Computer Science Report Series B, B-2019-2, 2019.
  * [13] A. Biere, A. Biere, M. Heule, H. van Maaren, and T. Walsh.  Handbook of Satisfiability: Volume 185 Frontiers in Artificial Intelligence and Applications.  IOS Press, NLD, 2009.
  * [14] A. Biere, M. Heule, and H. van Maaren.  Handbook of Satisfiability: Second Edition.  Frontiers in Artificial Intelligence and Applications. IOS Press, 2021\.
  * [15] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.
  * [16] Alfredo Braunstein, Marc Mezard, and Riccardo Zecchina.  Survey propagation: An algorithm for satisfiability.  Random Struct. Algorithms, 27:201–226, 09 2005.
  * [17] KyungHyun Cho, Tapani Raiko, and Alexander Ilin.  Parallel tempering is efficient for learning restricted boltzmann machines.  In The 2010 international joint conference on neural networks (ijcnn), pages 1–8. IEEE, 2010.
  * [18] Chase Cook, Hengyang Zhao, Takashi Sato, Masayuki Hiromoto, and Sheldon X. D. Tan.  Gpu based parallel ising computing for combinatorial optimization problems in vlsi physical design, 2019.
  * [19] Stephen A. Cook.  The complexity of theorem-proving procedures.  In Proceedings of the Third Annual ACM Symposium on Theory of Computing, STOC ’71, pages 151–158, New York, NY, USA, 1971. ACM.
  * [20] A. d’Anjou, M. Grana, F.J. Torrealdea, and M.C. Hernandez.  Solving satisfiability via boltzmann machines.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 15(5):514–521, 1993.
  * [21] Guillaume Desjardins, Aaron Courville, Yoshua Bengio, Pascal Vincent, Olivier Delalleau, et al.  Parallel tempering for training of restricted boltzmann machines.  In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 145–152. MIT Press Cambridge, MA, 2010.
  * [22] Johannes K. Fichte, Markus Hecher, and Stefan Szeider.  A time leap challenge for sat-solving.  In Principles and Practice of Constraint Programming: 26th International Conference, CP 2020, Louvain-La-Neuve, Belgium, September 7–11, 2020, Proceedings, page 267–285, Berlin, Heidelberg, 2020. Springer-Verlag.
  * [23] Yoav Freund and David Haussler.  Unsupervised learning of distributions on binary vectors using two layer networks.  In J. Moody, S. Hanson, and R. P. Lippmann, editors, Advances in Neural Information Processing Systems, volume 4. Morgan-Kaufmann, 1992.
  * [24] Linyan Gu, Jianfeng Huang, and Lihua Yang.  On the representational power of restricted boltzmann machines for symmetric functions and boolean functions.  IEEE Transactions on Neural Networks and Learning Systems, 30:1335–1347, 2019.
  * [25] Marijn J. H. Heule, Oliver Kullmann, Siert Wieringa, and Armin Biere.  Cube and conquer: Guiding cdcl sat solvers by lookaheads.  In Kerstin Eder, João Lourenço, and Onn Shehory, editors, Hardware and Software: Verification and Testing, pages 50–65, Berlin, Heidelberg, 2012. Springer Berlin Heidelberg.
  * [26] Geoffrey E Hinton.  Training products of experts by minimizing contrastive divergence.  Neural Computation, 14(8):1771–1800, 2002.
  * [27] Geoffrey E Hinton, Terrence J Sejnowski, et al.  Learning and relearning in boltzmann machines.  In Parallel distributed processing: Explorations in the microstructure of cognition, volume 1, pages 282–317. MIT Press, 1986.
  * [28] Holger Hoos and Thomas Stützle.  Local search algorithms for sat: An empirical evaluation.  Journal of Automated Reasoning, 24:421–481, 01 2000.
  * [29] John Hopfield and D Tank.  Neural computation of decisions in optimization problems.  Biological cybernetics, 52:141–52, 02 1985.
  * [30] Norman P. Jouppi, Doe Hyun Yoon, Matthew Ashcraft, Mark Gottscho, Thomas B. Jablin, George Kurian, James Laudon, Sheng Li, Peter Ma, Xiaoyu Ma, Thomas Norrie, Nishant Patil, Sushma Prasad, Cliff Young, Zongwei Zhou, and David Patterson.  Ten Lessons from Three Generations Shaped Google’s TPUv4i, page 1–14.  IEEE Press, 2021.
  * [31] Norman P. Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le, Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon.  In-datacenter performance analysis of a tensor processing unit.  In Proceedings of the 44th Annual International Symposium on Computer Architecture, ISCA ’17, page 1–12, New York, NY, USA, 2017. Association for Computing Machinery.
  * [32] Henry A. Kautz, Ashish Sabharwal, and Bart Selman.  Incomplete algorithms.  In Armin Biere, Marijn Heule, Hans van Maaren, and Toby Walsh, editors, Handbook of Satisfiability - Second Edition, volume 336 of Frontiers in Artificial Intelligence and Applications, pages 213–232. IOS Press, 2021.
  * [33] Diederik Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [34] Lars Kotthoff, Alexandre Fréchette, Tomasz Michalak, Talal Rahwan, Holger H. Hoos, and Kevin Leyton-Brown.  Quantifying algorithmic improvements over time.  In Proceedings of the 27th International Joint Conference on Artificial Intelligence, IJCAI’18, page 5165–5171. AAAI Press, 2018.
  * [35] Nicolas Le Roux and Yoshua Bengio.  Representational power of restricted boltzmann machines and deep belief networks.  Neural Computation, 20(6):1631–1649, 2008.
  * [36] Zhuwen Li, Qifeng Chen, and Vladlen Koltun.  Combinatorial optimization with graph convolutional networks and guided tree search.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * [37] Andrew Lucas.  Ising formulations of many NP problems.  Frontiers in Physics, 2:5, February 2014.
  * [38] Joao Marques-Silva.  Practical applications of boolean satisfiability.  In 2008 9th International Workshop on Discrete Event Systems, pages 74–80, 2008.
  * [39] Joao Marques-Silva, Inês Lynce, and Sharad Malik.  Conflict-driven clause learning sat solvers.  In Handbook of Satisfiability, 2009.
  * [40] James Martens, Arkadev Chattopadhya, Toni Pitassi, and Richard Zemel.  On the representational efficiency of restricted boltzmann machines.  In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
  * [41] Guido Montufar and Nihat Ay.  Refinements of Universal Approximation Results for Deep Belief Networks and Restricted Boltzmann Machines.  Neural Computation, 23(5):1306–1319, 05 2011.
  * [42] Michele Mosca and Sebastian R Verschoor.  Factoring semi-primes with (quantum) sat-solvers.  arXiv preprint arXiv:1902.01448, 2019.
  * [43] Radford M. Neal.  Annealed importance sampling.  Statistics and Computing, 11, 2001.
  * [44] Saavan Patel and Sayeef Salahuddin.  Combining learned representations for combinatorial optimization, 2019\.
  * [45] Ruslan Salakhutdinov and Iain Murray.  On the quantitative analysis of deep belief networks.  In Proceedings of the 25th International Conference on Machine Learning, ICML ’08, page 872–879, New York, NY, USA, 2008. Association for Computing Machinery.
  * [46] Daniel Selsam and Nikolaj Bjørner.  Guiding high-performance sat solvers with unsat-core predictions.  In Mikoláš Janota and Inês Lynce, editors, Theory and Applications of Satisfiability Testing – SAT 2019, pages 336–353, Cham, 2019. Springer International Publishing.
  * [47] Daniel Selsam, Matthew Lamm, Benedikt Bünz, Percy Liang, Leonardo de Moura, and David L. Dill.  Learning a SAT solver from single-bit supervision.  In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
  * [48] Xujie Si, Yujia Li, Vinod Nair, and Felix Gimeno.  Prioritized unit propagation with periodic resetting is (almost) all you need for random SAT solving.  CoRR, abs/1912.05906, 2019.
  * [49] P Smolensky.  Information processing in dynamical systems: foundations of harmony theory.  In Parallel distributed processing: explorations in the microstructure of cognition, volume 1, pages 194–281. MIT Press, 1986.
  * [50] Ali Asgar Sohanghpurwala, Mohamed W. Hassan, and Peter Athanas.  Hardware accelerated sat solvers—a survey.  Journal of Parallel and Distributed Computing, 106:170–184, 2017\.
  * [51] Ilya Sutskever and Tijmen Tieleman.  On the convergence properties of contrastive divergence.  In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pages 789–795, Chia Laguna Resort, Sardinia, Italy, 13–15 May 2010. PMLR.
  * [52] Robert H Swendsen and Jian-Sheng Wang.  Replica monte carlo simulation of spin-glasses.  Physical review letters, 57(21):2607, 1986.
  * [53] Kun Yang, Yi-Fan Chen, Georgios Roumpos, Chris Colby, and John Anderson.  High performance monte carlo simulation of ising model on tpu clusters.  In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’19, New York, NY, USA, 2019. Association for Computing Machinery.
  * [54] Emre Yolcu and Barnabas Poczos.  Learning local search heuristics for boolean satisfiability.  In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.

## Appendix

###  6.1 JAX implementation of RbmSAT

1# Copyright 2023 DeepMind Technologies Limited.2# SPDX-License-Identifier:
Apache-2.034from jax import nn5import jax.numpy as jnp6from jax.random import
PRNGKey, bernoulli, fold_in, split789def gather_and_count(v, Q, polarity):10 c
= jnp.einsum(’bv,kcv->bkc’, v, Q)11 return c, ((polarity > 0) *
c).any(axis=-2).sum(axis=-1)121314def rbmsat(W_c, b_c, T, B, N, upp, upw,
seed, $\alpha$):15 s_max = jnp.zeros((B,))16 init_key, v_key, h_key =
split(PRNGKey(seed), 3)17 v = bernoulli(init_key, shape=(B, N))18 $\nu$ = 0.25
* jnp.ones_like(v)19 polarity = jnp.sign(T) # strictly {1, -1}, not 020 Q =
nn.one_hot(abs(T) - 1, N) # indices in T are 1-based21 W = polarity[:,
jnp.newaxis, :] * W_c[:, :, jnp.newaxis]22 b = b_c[:, jnp.newaxis] + W_c.T @
(1 - polarity) / 223 t, d = 1, -124 step = 025 while time_remaining():26 step
+= 127 if upp > 0 and d == 0:28 v_u = jnp.vstack([v,
fetch_unit_prop_result()])29 _, s_u = gather_and_count(v_u, Q, polarity)30
ranks = jnp.argsort(s_u)31 v = v_u[ranks[-B:]]32 $\nu$ = jnp.vstack([$\nu$,
$\nu$])[ranks[-B:]]33 if upp > 0 and t % upp == 0:34 d = upw35
send_for_unit_propagation(v, $\nu$)36 c, s = gather_and_count(v, Q,
polarity)37 s_max = jnp.maximum(s_max, s)38 h_logits = b +
jnp.einsum(’bkc,khc->bhc’, c, W)39 h = bernoulli(fold_in(h_key, step),
nn.sigmoid(h_logits))40 $\rho$ = nn.sigmoid(jnp.einsum(’bhc,khc,kcv->bv’, h,
W, Q))41 $\nu$ = (1 - $\alpha$) * $\nu$ + $\alpha$ * $\rho$ * (1 - $\rho$)42 v
= bernoulli(fold_in(v_key, step), $\rho$)43 t, d = t + 1, max(d - 1, -1)44
return s_max.max()

Figure 5: A simplified implementation of RbmSAT; unit propagation and wall
clock timing have been stubbed out for clarity. rbmsat takes as arguments
canonical OR gate RBM parameters, the table $\mathbf{T}$ (section 4.3), a
number of chains $B$ and total number of variables $N$, the unit propagation
period upp, the number of steps before fetching unit propagation results upw,
a random seed and a parameter $\alpha$ for the moving average of the variance.
Returns the maximum number of clauses satisfied when the time budget is
exhausted.

###  6.2 List of problem instances

Table 3 lists the names of the MaxSAT Evaluation instances that are used in
this work, split by year. The instances are available for download from the
Evaluation websites linked from the top row of the table.

MaxSAT 2018 | MaxSAT 2019 | MaxSAT 2020 | MaxSAT 2021
---|---|---|---
(17 instances) | (29 instances) | (23 instances) | (23 instances)
maxcut-hamming6-4.clq.cnf | scpcyc06_maxsat.cnf | ram_k4_n20.ra0.cnf | maxcut-140-630-0.7-6.cnf
maxcut-brock400_4.clq.cnf | scpcyc10_maxsat.cnf | normalized-par32-1.cnf | uaq-min-nr-nr50-np400-rpp5-nc0-rs0-t0-plb100-n9.cnf
bcp-msp-normalized-par32-1-c.cnf | MANN_a27.clq.cnf | ram_k4_n19.ra0.cnf | ram_k3_n18.ra0.cnf
set-covering-scpcyc07_maxsat.cnf | ram_k4_n18.ra0.cnf | MANN_a9.clq.cnf | maxcut-140-630-0.7-4.cnf
bcp-msp-normalized-f600.cnf | p_hat700-3.clq.cnf | sanr400_0.7.clq.cnf | maxcut-140-630-0.8-26.cnf
bcp-msp-normalized-f2000.cnf | normalized-par32-5.cnf | MANN_a81.clq.cnf | scpcyc09_maxsat.cnf
maxcut-brock400_2.clq.cnf | normalized-par32-4-c.cnf | hamming6-2.clq.cnf | data.243.cnf
maxcut-t6pm3-8888.spn.cnf | brock200_1.clq.cnf | normalized-g100x100.opb.msat.cnf | hamming10-4.clq.cnf
maxcut-brock200_1.clq.cnf | sanr200_0.9.clq.cnf | ram_k3_n19.ra0.cnf | brock200_3.clq.cnf
bcp-msp-normalized-par32-1.cnf | san200_0.7_1.clq.cnf | scpcyc08_maxsat.cnf | ram_k3_n14.ra0.cnf
maxcut-p_hat500-3.clq.cnf | p_hat500-2.clq.cnf | brock400_4.clq.cnf | maxcut-140-630-0.8-18.cnf
set-covering-scpcyc08_maxsat.cnf | ram_k3_n16.ra0.cnf | ram_k3_n13.ra0.cnf | ram_k3_n15.ra0.cnf
set-covering-scpcyc09_maxsat.cnf | ram_k3_n12.ra0.cnf | normalized-par32-2.opb.msat.cnf | maxcut-140-630-0.7-45.cnf
set-covering-scpcyc06_maxsat.cnf | san400_0.7_2.clq.cnf | ram_k3_n20.ra0.cnf | ram_k3_n11.ra0.cnf
set-covering-scpcyc10_maxsat.cnf | maxcut-140-630-0.8-7.cnf | c-fat500-10.clq.cnf | MinFill_R0_queen5_5.cnf
bcp-msp-normalized-f1000.cnf | ram_k3_n10.ra0.cnf | normalized-par32-3-c.cnf | maxcut-140-630-0.7-12.cnf
maxcut-MANN_a9.clq.cnf | normalized-par32-1.opb.msat.cnf | ram_k3_n17.ra0.cnf | data.729.cnf
| normalized-par32-5.opb.msat.cnf | scpcyc07_maxsat.cnf | data.135.cnf
| brock400_3.clq.cnf | MANN_a45.clq.cnf | maxcut-140-630-0.8-39.cnf
| normalized-par32-4.opb.msat.cnf | p_hat500-3.clq.cnf | maxcut-140-630-0.8-1.cnf
| keller5.clq.cnf | johnson8-4-4.clq.cnf | maxcut-140-630-0.7-20.cnf
| brock400_2.clq.cnf | hamming8-2.clq.cnf | maxcut-140-630-0.7-28.cnf
| normalized-par32-3.opb.msat.cnf | brock400_1.clq.cnf | data.405.cnf
| brock800_1.clq.cnf |  |
| hamming10-2.clq.cnf |  |
| maxcut-140-630-0.7-31.cnf |  |
| p_hat300-3.clq.cnf |  |
| sanr400_0.5.clq.cnf |  |
| maxcut-140-630-0.8-17.cnf |  |
Table 3: Names of MaxSAT Evaluation instances used for evaluating RbmSAT. Evaluation | Prioritized UP Interval (Gibbs steps) | Markov chains per device
---|---|---
MaxSAT 2018 | 5000 | 128
MaxSAT 2019 | 5000 | 128
MaxSAT 2020 | 5000 | 128
MaxSAT 2021 | 4000 | 64
Table 4: RbmSAT hyperparameter values used for MaxSAT Evaluations.

###  6.3 Additional implementation details

Pretraining: We pre-train OR-gate RBMs with different satisfying assignment
free energy targets for maximum clause sizes of 3 through 7 as described in
section 4. We use Adam [33] with a learning rate of $10^{-3}$. We use $3$
hidden units for clause size 3 and $K+1$ hidden units for clause sizes
$K\in[4,7]$.

Our experiments use pre-trained RBMs at either 2, 8 or 16 distinct free energy
targets. In the 1 TPU condition in figure 3, the two logical cores use OR-gate
RBMs trained with $F^{\star}\in\\{0.068,0.528\\}$. In the 4 TPU condition,
targets were 0.068, 0.128, 0.188, 0.248, 0.308, 0.368, 0.428, and 0.488. In
the remaining conditions additional intermediate targets were added for a
total of 16 targets: 0.068, 0.098, 0.128, 0.158, 0.188, 0.218, 0.248, 0.278,
0.308, 0.338, 0.368, 0.398, 0.428, 0.458, 0.488, and 0.518. For 16 TPUs with
32 logical cores we repeat each distinct set of RBM parameters twice, with
distinct random initializations. For 64 TPUs with 128 logical cores, we repeat
each distinct set of RBM parameters 4 times.

Precision: The TPU matrix multiplication unit (MXU) natively performs
computations in 16-bit _bfloat16_ floating point arithmetic with the same
approximate dynamic range as IEEE 754 single precision, employing the same
exponent bit width but a reduced resolution significand (8-bit rather than
24-bit). Single precision floating point can be emulated with multiple passes
but by default, 32-bit operands are truncated. SIMD operations on the Vector
Processing Unit (VPU) (such as the addition of $\mathbf{b}$ in (7)) are
performed in 32-bit arithmetic for 32-bit operands. We store RBM parameters in
32-bit floating point but empirically the speed benefits of 16-bit MXU
operation outweigh any envisioned gains from increased precision.

Gather and scatter-add operations: We describe a procedure for indexing on the
MXU using a one-hot expansion of the index table, $\mathbf{Q}$. Depending on
the instance size, we found performing these computations on the VPU (using
JAX “fancy indexing” for gathers and jax.lax.segment_sum for scatter-adds)
could result in higher throughput. Note also that while binary visible unit
states are exactly representable in _bfloat16_ , the scatter-add computation
summing logit contributions will have different precision characteristics when
performed on the MXU and VPU, as described in section 6.3. We use this
alternative VPU-based implementation in our experimental results.

Batching operations: We are able to significantly increase sampling throughput
by compiling multiple sampling steps at a time using XLA control flow
(jax.lax.scan). For every dispatch to the TPU, we perform 500 rounds of block
Gibbs sampling, where each round consists of sampling hidden unit values and
then sampling new visible unit values.

###  6.4 Hyperparameters

We tune RbmSAT’s hyperparameters for each MaxSAT Evaluation year separately
using a grid search. For a given Evaluation year, each grid search point is
evaluated on the instances from all the other Evaluations and the best
hyperparameter values are selected. The selected values are shown in Table 4.