SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zifan Wang*†
Center for AI Safety zifan@safe.ai
Saranya Vijayakumar*
Carnegie Mellon University saranyav@andrew.cmu.edu
Kaiji Lu†
Pinterest Inc.
Caleblu95@gmail.com
Vijay Ganesh
Georgia Institute of Technology  
vghanesh@gatech.edu
Somesh Jha
University of Wisconsin-Madison jha@cs.wisc.edu
Matt Fredriskon†
Carnegie Mellon University mfredrik@cmu.edu
Recent techniques that integrate solver layers into Deep Neural Networks (DNNs) have shown promise in bridging a long-standing gap between inductive learning and symbolic reasoning techniques. In this paper we present a set of techniques for integrating Satisfiability Modulo Theories (SMT) solvers into the forward and backward passes of a deep network layer, called SMTLayer. Using this approach, one can encode rich domain knowledge into the network in the form of mathematical formulas. In the forward pass, the solver uses symbols produced by prior layers, along with these formulas, to construct inferences; in the backward pass, the solver informs updates to the network, driving it towards representations that are compatible with the solver's theory. Notably, the solver need not be differentiable. We implement SMTLayer as a Pytorch module, and our empirical results show that it leads to models that 1) require fewer training samples than conventional models, 2) that have consistent performance against certain types of covariate shift, and 3) that ultimately learn representations that are consistent with symbolic knowledge, and thus naturally interpretable. Our code is available at https://github.com/cmu-transparency/smt-layer
Deep neural networks (DNNs) have recently made significant strides, achieving surprising levels of performance on tasks like question-answering [40], text summarization [27], and code generation [4, 21]. However, the ability of models that perform well on these benchmarks to consistently demonstrate sound logical reasoning, even on tasks that may appear to be more simple and self-contained, remains in question [6, 31, 39]. For example, a prediction market on whether GPT-4 will be able to consistently solve "easy" Sikuu puzzles from the LA Times has remained open for several months at the time of this writing, despite the prompt-tuning efforts of traders on the market [19].
One way to address this issue is to encode logical constraints that are essential for certain inference tasks symbolically, making them available to the model either during training or inference. A promising approach to this type of "neuro-symbolic" learning is to integrate solver layers into the model [35, 34, 13, 18], both during training and inference. During training, the model must learn a representation of the data that is compatible with the symbolic knowledge, and during inference this
$^{\dagger}$ Work done as a PhD student in Carnegie Mellon University.
grounded representation is used to provide information from which a symbolic reasoning engine (i.e., a "solver") can extract accurate results.
The most straightforward way to incorporate a solver layer into a deep model is to learn representations that are compatible with symbols used by the solver. For example, if one wanted to leverage symbolic domain knowledge to classify images of birds, or diagnose ailments from CT scans, then one could train a model as in "concept bottlenecking" [14]. This requires detailed supervised labels, which may be prohibitively expensive to obtain and keep consistent with a potentially evolving domain theory.
We present a set of techniques for incorporating a Satisfiability Modulo Theories (SMT) solver into a DNN layer so that symbolic knowledge can be leveraged to learn such a compatible representation, without requiring label supervision. Our method, unlike prior methods, does not require fine-grained labels to learn representations [14]. For example, for learning Visual Sodium, regular supervised learning requires labels of individual digits, rather than just the solution to a set of Sodium puzzles, to learn a dedicated digit classifier first. Our work sidesteps the need to manually break a problem apart and obtain intermediate labels to supervise on, and instead allows for learning a predictor end-to-end from just labels on the targeted task. Our approach is general, and can handle a broad range of domain knowledge encoded as SMT constraints, provided that they interface with the surrounding neural network layers over propositional variables. Unlike related prior work [35], our approach does not approximate the solver's behavior by formulating a differentiable relaxation. Rather, as the solver works on a set of constraints, we extract information that is geared towards checking the correctness of preceding layers, and use that information to construct training updates (Section 4.2).
We present two different approaches for this, one based on unsatisfiable cores, and another based on weighted MaxSMT (Section 4.2). There are several advantages to this approach. Aside from the mild interface constraints mentioned earlier (i.e., solver and neural layers interface with each other via boolean variables), our approach does not place any restrictions on the theory solver embedded in the layer, such as linearity [34] or even decidability—if the solver is capable of efficiently discharging the relevant constraints, then the layer can operate as intended. Because there is no need to provide a differentiable relaxation for each theory or solver technique that one may want to incorporate, we can leverage the continuous and unabated progress being made in solver technology.
We implement our approach as a PyTorch [22] layer, using the Z3 [7] SMT solver as the solver layer to solve SMT and MaxSMT constraints. On three applications involving vision and natural language: visual arithmetic, algebraic equation solving, and a so-called natural language "liar's puzzle," we demonstrate that our implementation can be incorporated into DNN architectures to solve problems more effectively than conventional DNNs (Section 5). In particular, our results show that the data needed to train a DNN with symbolic knowledge may be much simpler than may be necessary otherwise, and that while doing so is more expensive computationally, often times the more efficient (i.e., not involving MaxSMT) algorithms perform well in practice.
1. We present SMTLayer, a framework for incorporating an SMT solver into a DNN, as a layer that leverages symbolic knowledge during training and inference.
2. Our empirical evaluation, over four types of tasks requiring logical reasoning, e.g. visual Sodomoku, shows that models using SMTLayer require significantly less training data, can be trained more efficiently, and generalize a lot better on the much larger test sets compared to those based on closely-related prior work [35, 13].
Section 3 provides background, Section 4 describes SMTLayer, Section 5 gives our empirical evaluation, and Appendix A.1 (supplementary material) discusses the broader impacts of our work.
Solver Layers. Vlastelica et al. [34] integrate a black-box and non-differentiable combinatorial solver on top of a deep network. To propagate the gradient through the solver on the backward pass, they linearly interpolate the loss w.r.t the solver's input and define the gradient of the solver as the slopes of the line segments. CSL solves a set of problems where the solver's objective must be linear w.r.t its input, e.g. finding the shortest path and travel salesman problem (TSP). Further, the authors assume that the only labels available are the outputs of solvers, e.g. the minimum cost in TSP, and
hence their tool has to discover the label for the output of the network itself. These requirements limit the choices one has for the solver layer. Wang et al. [35] present SATNet, a network architecture with a differentiable approximate MAXSAT solver layer. Their approximation is based on a coordinate descent approach to solving the semidefinite program (SDP) relaxation of the MAXSAT problem. SATNet does not assume that the logical structure of the problem is given, and instead attempts to learn it.
Neural Logic Programming. While SATNet integrates a logic-based solver on top of a network, DeepProbLog takes the opposite approach, extending the capability of a probabilistic logic programming language with neural predicates [18]. In the context of our work, the logic program can be viewed as a "solver layer" that explicitly encodes symbolic knowledge. Scallop [13] extends DeepProbLog to scale without sacrificing accuracy compared to DeepProbLog. Similarly to DeepProbLog, each possible result of the sum of two digits in MNIST is given a probability, in the form of a weighted Boolean formula. They prune unlikely clauses of the formula, represented by proofs, only keeping the top- $k$  most likely. Likelihood is computed using weighted model counting [13, 3]. These techniques are well-suited to problems that benefit from probabilistic Datalog, but have inherent limitations: they cannot handle quantifiers, general negation, and the range of supported first-order theories is more restrictive.
More recent work explores different directions in training models to perform neural-symbolic reasoning. For example, SMT solvers and MCMC sampling are utilized to support network training, side-stepping shortcut satisfaction [15, 16]. Hoernle et al. [12] introduce MultiplexNet, a deep model that ensures the satisfaction of symbolic constraints during the inference stage and adds an additional dimension to the neuro-symbolic learning landscape. Marconato et al. [20] has investigated how to mitigate the challenges and limitations associated with neuro-symbolic reasoning shortcuts and Yang et al. [37] use straight-through-estimators and logical constraints for neural network learning.
Differentiable Logic & Semantic Losses. Another recent direction has explored differentiable logics [9, 33, 32], loss functions that incorporate logical semantics [1, 36], or training methods guided by symbolic solvers [10, 28]. These approaches provide ways of integrating symbolic knowledge into training, by making logical formulas differentiable, and therefore amenable to optimization when included in a loss function. This line of work does not explicitly aim to make use of symbolic information during inference. In contrast, the information that our approach extracts from the solver during training is used to condition the model towards a representation that will allow it to communicate effectively with the solver during inference. Additionally, we do not require the logical formulas, or the solver, to be differentiable.
Let  $\mathbf{X}$  denote a domain of features,  $\mathbf{Y}$  a domain of labels, and  $\mathcal{D}$  a distribution over  $\mathbf{X} \times \mathbf{Y}$ . Formally,  $\mathcal{D}$  is a probability measure on a space given by a  $\sigma$ -algebra over subsets of  $\wp(\mathbf{X} \times \mathbf{Y})$ . The goal of a learning algorithm  $A$  is to find a function  $h: \mathbf{X} \to \mathbf{Y}$  that, for  $(x, y) \sim \mathcal{D}$ , can be used to predict  $y$  when given  $x$ . To do this,  $A$  is given a set of training examples  $S = (x_1, y_1), \ldots, (x_m, y_m)$  sampled i.i.d. from  $\mathcal{D}$ , and uses some criterion to select  $h$  from a hypothesis class  $\mathcal{H}$  of functions. We refer to  $h$  as the model learned by  $A$  on  $S$ . When the learning algorithm  $A$  is clear from the context, we will write  $h_S$  to denote the model produced from the given sample. Throughout this paper, we will generally assume that the loss is either the 0-1 loss  $\ell^{01}$  or binary cross-entropy  $\ell^{\mathrm{bce}}$ .
A theory  $T$  consists of a signature  $\Sigma$  of constant, predicate, and function symbols, as well as a set of axioms over  $\Sigma$ . Formulas in a theory are composed of elements of  $\Sigma$ , variables, and logical symbols such as quantifiers and Boolean operations. We use the term decision procedure to refer to an algorithm that is given an open  $T$ -formula, and returns true if it is satisfiable, and false otherwise. Additionally, it may return an assignment to all of the variables that demonstrates satisfiability, or if the formula is not satisfiable, then it may return an unsatisfiable core, which is a subset of clauses taken from the formula's representation in conjunctive normal form that remains unsatisfiable. Loosely, we also refer to such an algorithm as a "solver", but this term is more general, and could also refer to an algorithm that identifies the maximal set of clauses, possibly weighted by some user-defined values, that are satisfiable when conjoined.
Figure 1: Illustrative MNIST addition example (left), and SMT-based forward pass (right, Alg. 1). We use a binary encoding of digits, because SMT solvers support constraints involving integers that are encoded this way. While only four bits per digit are necessary to represent the inputs, five are needed for the output, and we represent all digits as 5-bit vectors for uniformity.
4 Constructing SMTLayer
In this section, we present SMTLayer, a set of algorithms for computing the forward and backward passes of a layer whose behavior is defined by a set of user-defined SMT constraints. SMTLayer does not have trainable parameters, and its functionality is wholly defined by a set of SMT constraints  $\phi$  that are provided by the model designer. SMTLayer can be used in modern deep-learning frameworks as a drop-in replacement for more conventional neural network layers, e.g., dense, convolutional, and LSTM [11] are prominent examples of widely-used layers.
Section 4.1 provides a high-level overview of our approach, Section 4.2 describes them in detail, and Section 4.3 begins an analysis of this setting that we hope future work will continue developing.
We envision SMTLayer being used primarily at the top of a DNN taking inputs from a stack of conventional DNN layers that convert raw input features into ground terms for the constraints  $\phi (z_0,\ldots ,z_{p - 1},y_0,\ldots ,y_{q - 1})$  embedded in SMTLayer, and producing outputs that are consistent with  $\phi$  and the given ground terms. Figure[1]shows an illustrative example, with the previously-studied problem of MNIST addition [18, 13].
During the forward pass the outputs of the previous layer are mapped to designated free variables  $z_0, \ldots, z_{p-1}$ . The layer then checks the satisfiability of  $\phi$ , a formula in an appropriate combination of first-order theories, after substituting these ground terms for the  $z_i$ , and the output of the layer consists of the solver's model for  $y_0, \ldots, y_{q-1}$ . These outputs are converted from Boolean to floating-point values by mapping false to -1 and true to 1. At the moment, the only restriction on  $\phi$  that our layer requires is that  $z$  and  $y$  be vectors of Booleans, so that they can be appropriately mapped to continuous values; any other symbols appearing in  $\phi$  can come from arbitrary domains (e.g. strings) supported by the underlying SMT solver.
In the backward pass, the layer receives the gradient of its output with respect to the function whose derivative is being computed, which we will assume is the binary cross-entropy loss  $\ell(y, y^{\star})$ . Unless stated otherwise, we will assume this loss for the remainder of the section. This gradient is used, along with the inputs and outputs of the corresponding forward pass, to first compute an amended output  $\hat{y}$  which corresponds to an output that would have yielded a smaller loss. Because the outputs are Boolean, it is always possible to determine the ground truth output  $y^{\star}$  from this information. Using  $\hat{y}$ , the layer determines which of components of its inputs are inconsistent with  $\phi$  and  $\hat{y}$ , and provides the corresponding gradients to the previous layer. Section 4.2 details the manner in which these gradients are computed.
4.2 SMTLayer, forward and backward
We now present the details of the forward and backward passes of SMTLayer. There are two algorithms for each pass,  $\mathsf{F}_{smt}^{\phi}$  and  $\mathsf{F}_{max}^{\phi}$  are forward passes, and  $\mathsf{B}_{core}^{\phi}$ ,  $\mathsf{B}_{max}^{\phi}$  are backward passes.  $\mathsf{F}_{max}^{\phi}$  and  $\mathsf{B}_{max}^{\phi}$  both make use of MaxSMT solvers, whereas  $\mathsf{F}_{smt}^{\phi}$  and  $\mathsf{B}_{core}^{\phi}$  rely on satisfiability solvers (SMT). They are all compatible with each other. That is,  $\mathsf{F}_{smt}^{\phi}$  (and  $\mathsf{F}_{max}^{\phi}$ ) can be used with either  $\mathsf{B}_{max}^{\phi}$  or  $\mathsf{B}_{core}^{\phi}$ . Algorithms for  $\mathsf{F}_{smt}^{\phi}$  and  $\mathsf{B}_{core}^{\phi}$  are included in this section, and the MaxSMT-based variants are detailed in Appendix A.2
Forward pass. Algorithm 1 illustrates  $F_{smt}^{\phi}$ , our SMT-based forward pass, and Algorithm 3 for computing  $F_{max}^{\phi}$  is given in Appendix A.2. Both of the algorithms are parameterized by a user-provided first-order formula  $\phi$ , and take a single vector-valued input consisting of unscaled floating-point values (logits). These values are cast to Booleans by taking their sign on line 2 of both algorithms, and they are equated with the corresponding free variables  $z_0, \ldots, z_{p-1}$ .
The key difference between  $\mathsf{F}_{\text{max}}^{\phi}$  and  $\mathsf{F}_{\text{smt}}^{\phi}$  is the way in which they handle inputs that are inconsistent with  $\phi$  when interpreted as Booleans.  $\mathsf{F}_{\text{smt}}^{\phi}$  addresses this by providing an output that is also inconsistent with  $\phi$ , i.e., a vector of zeroes, effectively signaling that the network below it did not provide consistent inputs. Alternatively, we can interpret the values provided by the network as Booleans enriched with "confidence" values. To obtain MaxSMT weights, Algorithm 3 scales the input logits to a formal probability distribution via the softmax function (line 3). With this approach, SMTLayer will always provide a valid, although not necessarily correct, output that is consistent wrt  $\phi$  with the inputs of which the network below is most "confident." (line 4).
Backward pass. The backward pass is responsible for computing the gradient of the loss with respect to the layer inputs. It is given the gradient of the loss with respect to the layer outputs, and is assumed to have memoized the inputs and outputs from the forward pass. The gradients returned by this pass are then used by the backward pass of previous layers, and ultimately to derive updates to trainable parameters.
The key issue in designing a backward pass for SMTLayer is the geometry of the functions computed by either forward pass. For any vector  $v \in \{-1, 0, 1\}^p$  and  $x, x'$  with  $\operatorname{sign}(x) = \operatorname{sign}(x') = v$ , then  $\mathsf{B}^\phi(x) = \mathsf{B}^\phi(x')$ , so these functions are piece-wise constant step functions ranging over the corners of the  $\mathbb{R}^q$  unit hypercube. Thus, while they are differentiable almost everywhere, the gradient is not helpful for training because it is always zero. Prior work
on integrating such functions into deep networks [35, 34] addresses this problem by relaxing the function computed by the forward pass, so that its gradients are hopefully more informative.
In contrast,  $\mathsf{B}_{core}^{\phi}$  (Algorithm 2) and  $\mathsf{B}_{max}^{\phi}$  (Algorithm 4 in Appendix A.2) do not attempt to provide gradients for a relaxation of the forward pass. Instead, they use information provided by the solver in its computation of the forward pass to identify which components of the input may have contributed to higher loss. The gradient is then computed by constructing a variant of the input provided to the forward pass, which differs on the identified components, and returning the gradient of the BCE loss of the original input on this variant. The two algorithms differ in the information that they extract from the solver, i.e., either solutions to a MaxSMT instance or an unsatisfiable core.
Both algorithms begin by initializing the gradient to be the loss between the logit inputs, and their hard labels (line 2). Recall that we assume the loss  $\ell$  is binary cross-entropy, so the result will not
be zero. The purpose of this initialization is to emulate the dynamics of training with cross-entropy loss with a conventional layer; when the rounded output matches the target, the loss is not zero, and training will continue to move the parameters in a direction that makes them agree "more" with the hard target.
One line 3, they then use the provided gradient from the subsequent layer together with the memoized output from the forward pass to construct  $\hat{y}$ , a "corrected" output that satisfies  $\ell(\hat{y}, y^{\star}) \leq \ell(y, y^{\star})$ . To understand why, observe that the sign of  $\hat{y}$  computed on line 3 of both algorithms must be equal to that of  $y^{\star}$ . This follows from two facts: (1) at any coordinate  $i$  where  $y[i] \neq y^{\star}[i]$ ,  $\operatorname{sign}(\partial_y \ell(y, y^{\star}))[i] = \operatorname{sign}(y)[i]$ ; (2) at any coordinate  $i$  where  $y[i] = y^{\star}[i]$ ,  $\operatorname{sign}(\partial_y \ell(y, y^{\star}))[i] = -1 \cdot \operatorname{sign}(y)[i]$ . If the sign of  $\hat{y}$  is the same as that of  $y$ , then both algorithms return the initialized gradient. Otherwise, they extract information from the solver using  $z$  and  $\hat{y}$ .
Algorithm 2 that runs  $\mathsf{B}_{core}^{\phi}$  identifies a set of constraints  $z_{i} = z_{b}[i]$  that are inconsistent with  $\phi \wedge \phi_{y}$ . Note that line 7 specifies a minimal unsatisfiable core, but this is not necessary. All that is needed is that none of the clauses in the core be superfluous, i.e., deleting any singleton clause from  $I$  will cause it to be satisfiable. If a superfluous clause remains in the core, then the gradient returned for the corresponding input will have the incorrect sign, which may lead to issues with training.  $\mathsf{B}_{core}^{\phi}$  then updates the gradient at each input identified in the core using the loss of  $z[i]$  with respect to  $1 - \mathrm{sign}(z[i])$ , which will lead to updates in a direction that would have modified the input such that  $i$  was not in the unsat core. The indices not in the unsat core have their gradients set to zero, as their absence in the core is not evidence that these inputs were correct or incorrect.
Algorithm 4 (in Appendix A.2) that runs  $\mathsf{B}_{\text{max}}^{\phi}$  instead constructs a set of clauses  $\phi_y$  that constrain the free  $y_0, \ldots, y_{q-1}$  to take the values of  $\hat{y}_b$ , the Boolean conversion of  $\hat{y}$ . It then computes the softmax values of the absolute incoming logits  $|z|$ , and uses them to find the maximally-weighted set of clauses ( $\text{softmax}(|z|)[i], z_i = z_b[i]$ ) that are consistent with  $\phi \wedge \phi_y$ . Intuitively, these are the inputs that the previous layer is most confident in that can be made consistent with the corrected label  $\hat{y}$  by changing some of the less confident inputs.  $\mathsf{B}_{\text{max}}^{\phi}$  then updates the initialized gradient at each index for which the solution to this instance does not match the sign of the original input.
To understand the settings where SMTLayer will provide optimal results, we introduce a class of decomposable learning problems (Definition I).
Definition 1 (Decomposable problem). Let  $T$  be a first-order theory with constants in  $\mathbf{Z}$ . An ERM problem  $\mathcal{D}$ ,  $\mathcal{H}$  is decomposable by  $T$  if there exists a function  $f: \mathbf{X} \to \mathbf{Z}$ , companion hypothesis class  $\mathcal{H}_f \subseteq \mathbf{X} \to \mathbf{Z}$ , and  $T$ -formula  $\phi$  such that:
1. For any  $h \in \mathcal{H}$ , there exists  $h_f \in \mathcal{H}_f$  and  $h'$  such that  $h = h' \circ h_f$ .
2. There exists a random function  $g: \wp(\mathbf{Y}) \to \mathbf{Y}$  such that for any  $n > 0$  and  $\forall S$  in the support of  $\mathcal{D}^n$ ,
$$
\Pr_ {(x, y) \sim \mathcal {D}} [ (x, y) \in S ] = \Pr_ {(x, \cdot) \sim \mathcal {D}} [ (x, g (\langle x \rangle_ {f, \phi})) \in S ]
$$
where  $\langle x\rangle_{f,\phi} = \{y:\phi (f(x),y)$  is satisfied}.