Incorporating domain knowledge into neural-guided search via in situ priors and constraints
Abstract
Many AutoML problems involve optimizing discrete objects under a black-box reward.
Neural-guided search provides a ﬂexible means of searching these combinatorial spaces using an autoregressive recurrent neural network.
A major beneﬁt of this approach is that builds up objects sequentially —this provides an opportunity to incorporate domain knowledge into the search by directly modifying the logits emitted during sampling.
In this work, we formalize a framework for incorporating such in situ priors and constraints into neuralguided search, and provide suﬃcient conditions for enforcing constraints.
We integrate several priors and constraints from existing works into this framework, propose several new ones and demonstrate their eﬃcacy in informing the task of symbolic regression
1. Introduction
“Any practical algorithm must avoid exploring all but a tiny fraction of the state space.”
Russell and Norvig (2002)
Many problems in automated machine learning (AutoML) fall into the category of symbolic optimization (SO).
We consider the following discrete optimization problem: arg max n ∈ N ,τ 1 ,...,τ n [ R ( τ 1 , . . .
, τ n )] with τ i ∈ L = { α, β, . . .
, ζ }
RNN
[Figure]
pp l y c on s t
r
a
i n t s softmax
[Figure]
A pp l y p r i o r s ⌧ i ⌧ i
[Figure]
In this formulation, τ = [ τ 1 , . . . , τ n ] is a discrete object represented by a variable-length sequence of discrete symbols or “tokens” τ i selected from a library L , and R is a black-box (i.e. non-diﬀerentiable) reward function.
A popular SO problem is neural architecture search (NAS), in which tokens represent architectural hyperparameters, the sequence represents a speciﬁcation of a neural network architecture, and the reward is the validation accuracy when instantiating the speciﬁed network and training it on some downstream task (Zoph and Le, 2016).
Other examples include program synthesis (Abolaﬁa et al., 2018), symbolic regression (Petersen et al., 2021), de novo molecular design (Popova et al., 2019), automated theorem proving (Bibel, 2013), and many traditional combinatorial optimization problems (e.g. traveling salesman problem) (Bello et al., 2016).
For many SO problems, the object is not traditionally represented as a sequence; however, these can be reduced to SO by establishing a sequential representation of the object.
For example, small molecules can be represented by simpliﬁed molecular-input line-entry system (SMILES) strings (Weininger, 1988), and mathematical expressions can be represented by their pre-order traversals.
SO is challenging due to its combinatorial nature, as the size of the search space grows exponentially with the length of the sequence.
Because of this, it is commonly solved using evolutionary approaches (e.g. genetic programming), which are easily scalable and broadly applicable (Banzhaf et al., 1998).
More recently, neural-guided search has proven to be a successful alternative (Zoph and Le, 2016; Bello et al., 2016).
In this framework, an autoregressive recurrent neural network (RNN) generates objects sequentially.
Speciﬁcally, the RNN emits a vector of logits (cid:96) that deﬁnes a categorical distribution over tokens in the library L , conditioned on the previously sampled tokens via the RNN state.
A token is then sampled: τ i ∼ p ( τ i | τ 1:( i − 1) ; θ ) = Categorical(Softmax
( (cid:96) )), where θ are the RNN parameters.
The RNN is trained using reinforcement learning or other heuristic methods (Abolaﬁa et al., 2018) aimed at increasing the likelihood of high-reward objects.
A key methodological aspect of autoregressive approaches is that they generate objects sequentially , i.e. one token at a time.
In contrast, evolutionary approaches generate new objects by making edits (via mutations and crossovers) to existing objects.
A beneﬁt of sequential object generation is that it aﬀords an opportunity to incorporate prior knowledge into the search phase during each step of the generative process.
More speciﬁcally, one can incorporate knowledge in situ (i.e. during the autoregressive generative process) by directly modifying the emitted logits (cid:96) before sampling each token.
In contrast, non-sequential approaches resort to post hoc methods such as rejection sampling, which is ineﬃcient and can lead to intractable likelihoods.
As a simple example, consider expert knowledge that the current token should not be the same as the previous token, i.e. we want to impose the constraint
τ i (cid:54) =
τ i − 1 .
This particular constraint can be incorporated by simply adding negative inﬁnity to the logit corresponding to τ i − 1 each step of autoregressive sampling.
While several individual works incorporate priors and/or constraints into their search, to our knowledge there is no existing work that formalizes this approach.
Thus, we make the following contributions: (1) deﬁning a simple yet ﬂexible framework for imposing in situ priors and constraints for autoregressive neural-guided search, (2) formalizing suﬃcient conditions for imposing in situ constraints, (3) providing a centralized list of previously proposed priors and constraints, (4) proposing several new priors and constraints, and (5) empirically demonstrate their eﬃcacy in the task of symbolic regression.
Our overarching goal is to provide a set of priors and constraints that is suﬃciently broad to spark new ideas and encourage readers to design their own novel priors and constraints suited to their particular AutoML tasks.
2. Related Work
In many areas of AI, the idea of pruning , or eliminating possibilities from consideration without having to examine them, is critically important (Russell and Norvig, 2002).
In adversarial search for games, where branching factors are high and exploring all possible moves is infeasible,
alpha-beta pruning (Edwards and Hart, 1961) can be used to eliminate branches of the game tree that have no eﬀect on the ﬁnal evaluation and would be otherwise futilely evaluated by the standard minimax algorithm.
Other approaches, like forward pruning (Greenblatt et al., 1967) or futility pruning (Heinz, 1998), reduce the space by eliminating moves that appear to be poor moves based on heuristic evaluation functions.
The recent achievements in the game of Go (Silver et al., 2016) are also largely attributed to using policy and value networks (trained over millions of games) to guide Monte Carlo Tree
Search (MCTS) rollouts, thereby avoiding exploring paths that would result in sub-optimal moves while focusing computational resources in exploring the most promising paths.
In the ﬁeld of SO, Popova et al.
(2019) use valency constraints to follow rules of organic chemistry to optimize molecular structures.
Li et al.
(2019) use a neural-guided approach to identify asymptotic constraints of leading polynomial powers and use those constraints to guide MCTS for the problem of symbolic regression.
Kim et al.
(2021b) employ several expression-based constraints to inform the search for symbolic control policies in reinforcement learning.
3. A Framework for in situ Priors and Constraints
In this work, we deﬁne a prior as a logit adjustment vector , denoted (cid:96) ◦ , that is added to the emitted RNN logits (cid:96) during autoregressive sampling:
τ i ∼ Categorical(Softmax
( (cid:96) + (cid:96) ◦ )).
We deﬁne constraints as a special case of “hard” priors in which logit adjustments are either zero (no eﬀect) or negative inﬁnity (the token cannot be selected), denoted (cid:96) (cid:20) .
Thus, priors bias the search but do not reduce the size of the search space, whereas constraints prune (i.e. eliminate possible sequences of) the search space.
Typically, multiple priors and constraints can be composed simply by summing their logit adjustments, yielding the ﬁnal distribution for the i th token:
τ i ∼ Categorical(Softmax
( (cid:96) + (cid:80) j
(cid:96) ( j ) ◦ + (cid:80) k
(cid:96) ( k ) (cid:20) )), where (cid:96) ( j ) ◦ is the j th prior and (cid:96) ( k ) (cid:20) is the k th constraint.
Tables 1 and 2 provide succinct summaries of several broadly applicable classes of constraints and priors, respectively.
In the sections below, we provide details for three novel constraints, two additional constraints which appear in existing works but we generalize in this work, and two novel priors.
Detailed descriptions of other existing constraints and priors are provided in the Appendix.
3.1 Example classes of in situ constraints
Lexicographical constraint.
Many SO problem exhibit large semantic equivalence classes , or sets of objects whose semantics are identical, e.g. x + y and y + x in mathematics,
CHand H 4 C in chemistry, or A ∧ B and B ∧
A in logical reasoning.
To reduce the number of semantically equivalent sequences generated for tree-based objects, we can assign a predeﬁned, arbitrary order for the operands of commutative operators (e.g. + or × ).
To do so, we posit a lexicographical ordering of the tokens, deﬁned by an injective function l ( τ i ), where l : L → N .
Given a commutative operator token of arity n , the lexicographical constraint is enforced by requiring the lexicographical value of a child token to be no less than that of the previous child token, i.e. children are lexicographically sorted.
Given n diﬀerent child tokens, there are n ! possible ways to place them as children of an n -ary commutative operator.
This constraint reduces the number of possible orderings to one.
Further, since this reduction propagates through the subtrees, it reduces the search space exponentially with the number of commutative tokens.
Note that this constraint reduces the size of the search space without reducing the number of semantically unique objects in the search space.
Example.
Under the lexicographical constraint applied to the commutative operator +, if l (cos) > l (sin), then the expression cos( (cid:3) )
+ sin( (cid:3) )
(in this order) is constrained but sin( (cid:3) )
+ cos( (cid:3) ) is allowed.
Subtree length constraint.
Similar to the lexicographical constraint, the subtree length constraint reduces the number of semantically repeated sequences generated.
Here it does so by limiting the length of a subtree by the length of the preceding sibling subtree.
When sampling a tree-based object by its pre-order traversal, a given subtree completes before the next sibling subtree begins.
Thus, we can implement constrain each subtree to have a maximum length given by preceding sibling subtree length.
Example.
Under the subtree length constraint applied to the commutative binary operator +, consider the partial sequence [+ , sin , x ] (corresponding to the expression sin( x ) + (cid:3) ).
The ﬁrst subtree (sin( x )) has length 2, so the next subtree ( (cid:3) ) length is constrained to be at most 2.
Type and unit constraints.
In NAS, diﬀerent positions along the sequence have diﬀerent types, (e.g. activation function, number of neurons), each of which has its own set of allowable tokens.
A type constraint ensures that all positions satisfy the speciﬁed types.
More generally, as in the case of mathematical expressions, tokens may have speciﬁc input/output types or units.
For example, the × token has no requirements on input units, but the output units must be equal to the product of the input units.
Example.
Given the partial sequence [ × , x ] (corresponding to the expression x · (cid:3) ), where x has units kg and the ﬁnal output has units kg 2 , tokens whose output types cannot be kg are constrained (e.g. trigonometric functions, input variables with units other than kg).
Relational constraint.
A broad class of constraints for tree-based objects involves preventing one arbitrary set of tokens (called “targets”) from having a particular structural relationship with another arbitrary set of tokens (called “eﬀectors”).
Common relationships include descendants, children, and siblings.
This type of constraint is useful in expression search spaces for preventing nested trigonometric functions, inverse unary operators cancelling out (e.g. log(exp( x ))), or redundantly adding together two constants.
A few instances of relational constraints are used in Petersen et al.
; here, we generalize them into a single constraint class.
Example.
Under the relational constraint “[trigonometric functions] cannot be the [descendant] of [trigonometric functions],” given the partial sequence [sin , + , x, × , x ] (corresponding to the expression sin( x + x · (cid:3) )), all trigonometric tokens are constrained because they are a descendant of sin.
Blacklist constraint.
For all object types, one can constrain a speciﬁed set of “blacklisted” sequences (or partial sequences) from being sampled.
This may be useful to prune the search space of previously known sub-optimal solutions.
Example.
Liang et al.
(2018) implement a special case of this constraint called “systematic exploration” in which the set
[Table]
[Table]
of blacklisted sequences grows over time
[Table]
and is deﬁned by the history of sequences sampled so far.
Under this constraint, the same sequence will never be sampled more than once.
3.2 Suﬃcient conditions for imposing in situ constraints
Here we describe suﬃcient conditions for the ability to enforce a particular in situ constraint.
Let C i be the desired set of constrained/disallowed tokens before sampling τ i .
Autoregressive sampling can enforce any constraint that can be expressed as a function f ( · ) of the partial sequence τ 1: i − 1 , the library of tokens L
[Table] and any constraint parameters Ω (e.g. [min] and [max] for the length constraint), provided that C i does not include all tokens in L .
That is, if there exists an f ( · ) such that C i = f
( τ 1: i − 1 , L , Ω) and C i (cid:40)
L , then autoregressive sampling can enforce the constraint induced by f ( · ).
Further, k independent constraints
C (1) i , . .
, C ( k ) i can be combined if and only if C (1) i ∪ · · · ∪ C ( k ) i (cid:40) L .
Notably, these conditions preclude constraints based on the semantics or knowledge of the complete sequence τ .
For example, one cannot constrain the search to all expressions with a particular range, or to molecules with a speciﬁc melting point.
We describe several examples under this formalism:
Lexicographical : C i
= { v : l ( v ) <
l ( τ left ) , v ∈ L} , where τ left is the left sibling of τ i .
The constraint parameters are the lexicographical values: Ω = { l ( v ) : v ∈ L} .
Subtree length :
Set N ← i + L ( τ 1: i − 1 ) and apply the length constraint with [min] =and [max] = N , where L ( τ 1: i − 1 ) is the size of the subtree rooted at the left sibling of τ i .
Blacklist : C i = { v : ( τ 1: i − 1 (cid:107) v ) ∈
T , v ∈ L} , where T is the set of blacklisted sequences and (cid:107) represents concatenation.
In this case, Ω = T .
3.3 Example classes of in situ priors
Token-speciﬁc priors.
Token-speciﬁc priors allow users to diﬀerentially control the relative prior probability of particular tokens.
Speciﬁcally, the user selects a vector λ of length |L| that speciﬁes the desired relative prior probability of each token.
The token-speciﬁc prior adds logits (cid:96) ◦ such that the resulting probability vector is multiplied element-wise by λ and renormalized to sum to unity.
To compute this prior, we seek (cid:96) ◦ such that
Softmax( (cid:96) + (cid:96) ◦ ) =
( λ (cid:12) Softmax( (cid:96) )) /
( λ · Softmax( (cid:96) )), where (cid:12) represents element-wise mul- [Table] tiplication and · represents the dot product.
The solution is simply (cid:96) ◦
= log λ (deﬁned up to an additive constant); interestingly, this does not depend on (cid:96) [Table]
Positional priors.
For sequence-based objects [Table] it may be desirable to apply diﬀerent token-speciﬁc priors at diﬀerent positions along the sequence.
In this case, separate tokenspeciﬁc priors can be applied at particular positions along the sequence. For example, in NAS
[Table]
given a known, high-performing reference architecture [Table] one can guide the search toward similar architectures by applying positional priors that bias each position toward the corresponding value in the reference architecture
[Table]
4. Experiments and Discussion
[Table]
For empirical analysis, we consider symbolic regression , the task of discovering tractable mathematical expressions to ﬁt a dataset ( X, y ), where X ∈
R n and y ∈ R .
In this SO problem, tokens represent mathematical operators (e.g. sin , × ).
A sequence τ is instantiated as a function τ = f ( X ), where f : R n → R , which is used to predict values ˆ y = f
( X ).
Reward is based on mean-square error between ˆ y and y .
Table 3 shows performance on the Nguyen symbolic regression benchmarks (Uy et al., 2011), using either deep symbolic regression (DSR) (Petersen et al., 2021) or random search, i.e. DSR with learning rate set to 0.
We observe that adding individual priors or constraints (rows 2 − 7) generally improves performance over no priors or constraints (row 1), for both
DSR and random search.
Combining all priors and constraints (rows 8 and 9; note that the lexicographical constraint (L) and subtree length constraint (S) are mutually incompatible) greatly improves performance.
Lastly, it is interesting to note that random search with all priors and constraints (55.4%) achieves nearly the same performance as DSR with no priors or constraints (57.4%), demonstrating the ability of priors and constraints to eﬀectively bias and prune the search.
5. Conclusion
We consider a generic framework for incorporating in situ priors and constraints into neuralguided search, which we believe is well-suited for integrating expert knowledge into symbolic optimization tasks.
By contextualizing many existing priors and constraints within this framework and proposing several new ones, we hope to encourage researchers to design their own priors and constraints suited to their particular AutoML tasks. Acknowledgments This work was performed under the auspices of the U S. Department of Energy by Lawrence
Livermore National Laboratory under contract DE-AC52-07NA27344.
Lawrence Livermore National Security, LLC. LLNL-CONF-822798. Appendix A. A.1 Descriptions of existing constraints
To describe existing constraints, we deﬁne δ as the delta function, d i as the number of unselected or “dangling” nodes at position i (for tree-based objects), 1 condition as the indicator function (which returns 1 if condition is true and 0 otherwise), and A ( n ) as the set of tokens in L with arity n (for tree-based objects).
Length constraint.
Many SO problems allow variable-length objects, e.g. de novo molecular design, program synthesis, symbolic regression, and some formulations of NAS.
One can constrain the length to prevent objects that are either too complicated or too simple.
An interesting special case of this constraint is when the minimum and maximum length are equal.
For tree- or graph-based objects, this constrains the search to diﬀerent structures of the same length.
This can be used for the task of searching over the space of molecular isomers, e.g. straight-chain alkanes.
Given constraint parameters Ω = { [min] , [max] } , the length constraint is given by:
[Equation]
Example 1.
Given the partial sequence [sin , + , x ] (corresponding to the expression sin( x + (cid:3) )) with a maximum length of 5, all tokens with arity greater than 1 are constrained because the sequence would no longer be able to complete before exceeding the maximum length.
Example 2.
Given the partial sequence [+ , x ] (corresponding to the expression x + (cid:3) ) with a minimum length of 4, terminal tokens are constrained because they would complete the sequence prematurely.
Repeat constraint.
Given a set T of target tokens and a [min] and [max] value, the repeat constraint aims to limit the number of times each target token appears to fall between [min] and [max], inclusive.
Set N v ← (cid:80) i − 1 k =
δ τ k ,v , ∀ v ∈ T .
Then, for each target token v ∈ T , set
[Equation]
The ﬁnal constraint is given by C i =
(cid:83) v ∈ T C ( v ) i .
In this case, Ω = {T , [min] , [max] } .
Example.
Given T = { x } , [min] = 1, and [max] = 2, and the partial sequence [+ , + , x, x ] (corresponding to the expression x + x + (cid:3) ), x cannot be selected because it would exceed [max].
Valency constraint.
This constraint can be enforced using the length constraint , since valency corresponds to arity.
Thus, given the current value of d i and the valency value n , we can apply the length constraint with [min] =
d i + n .
A.2 Descriptions of existing priors
Soft length prior.
The length constraint is known to result in highly skewed distribution over lengths (Kim et al., 2021b;
Landajuela et al., 2021).
To alleviate this phenomenon,
Landajuela et al.
propose combining the length constraint with a “soft” version that reduces the probability of terminal tokens early on (to discourage short sequences) and of non-terminal tokens later on (to discourage long sequences).
Speciﬁcally, for libraries with arities in { 0 , 1 , 2 } , they deﬁne the soft length prior at position i as:
[Equation]
where ( · ) n denotes that element ( · ) is repeated n times, λ and σ are hyperparameters, and the tokens corresponding to logits (cid:96) are sorted by decreasing arity.
Uniform arity prior.
With zero-initialized RNN weights, the initial distribution is uniform over all tokens in L .
One can convert this to instead to be uniform over all arities in L by applying a uniform arity prior (Landajuela et al., 2021) given by:
[Equation]
where k is the maximum arity of tokens in L and the tokens corresponding to logits (cid:96) are sorted by decreasing arity.
This is typically combined with the soft length prior.
Language model prior.
Kim et al.
(2021a) introduce a mathematical language model trained on a “corpus” of mathematical expressions derived from Wikipedia, then use this language model as a prior to demonstrate an increase in search eﬃciency for symbolic regression.
Given a language model M that produces logits (cid:96) M from a given partial sequence τ 1: i − 1 , the language model prior is given by:
[Equation]
where λ is a hyperparameter controlling the strength of the language model prior, acting like an inverse temperature on the contribution of (cid:96) M to the softmax computation.