L EARNING TO M AKE D ECISIONS VIA S UBMODULAR R EGULARIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Many sequential decision making tasks can be viewed as combinatorial optimization problems over a large number of actions.
When the cost of evaluating an action is high, even a greedy algorithm, which iteratively picks the best action given the history, is prohibitive to run.
In this paper, we aim to learn a greedy heuristic for sequentially selecting actions as a surrogate for invoking the expensive oracle when evaluating an action.
In particular, we focus on a class of combinatorial problems that can be solved via submodular maximization (either directly on the objective function or via submodular surrogates).
We introduce a data-driven optimization framework based on the submodular-norm loss, a novel loss function that encourages the resulting objective to exhibit diminishing returns .
Our framework outputs a surrogate objective that is efﬁcient to train, approximately submodular, and can be made permutation-invariant.
The latter two properties allow us to prove strong approximation guarantees for the learned greedy heuristic.
Furthermore, our model is easily integrated with modern deep imitation learning pipelines for sequential prediction tasks.
We demonstrate the performance of our algorithm on a variety of batched and sequential optimization tasks, including set cover, active learning, and Bayesian optimization for protein engineering.
I NTRODUCTION
In real-world automated decision making tasks, we seek the optimal set of actions that jointly achieve the maximal utility.
Many of such tasks—either deterministic/non-adaptive or stochastic/adaptive— can be viewed as combinatorial optimization problems over a large number of actions.
As an example, consider the active learning problem, where a learner seeks the maximally-informative set of training examples for learning a classiﬁer.
The utility of a training set could be measured by the mutual information (Lindley, 1956) between the training set and the remaining (unlabeled) data points, or by the expected reduction in generation error, should one train the model on the candidate training set.
Similar problem arises in a number of other domains, such as experimental design (Chaloner and Verdinelli, 1995), document summarization (Lin and Bilmes, 2012), recommender system (Javdani et al., 2014), and policy making (Runge et al., 2011).
Identifying the optimal set of actions (e.g., optimal training sets, most informative experiments) amounts to evaluating the expected utility over a combinatorial number of candidate sets; when the underlying model class is complex and the evaluation of the utility function is expensive, these tasks are notoriously difﬁcult to optimize (Krause and Guestrin, 2009).
It has been shown that for a broad class of decision making problem—including the aforementioned applications in active learning and experimental design—where the optimization criterion is to maximize the decision-theoretic value of information , it is possible to design surrogate objective function that are (approximately) submodular while being aligned with the original objective at the optimal solutions (Javdani et al., 2014; Chen et al., 2015b; Choudhury et al., 2017).
Here, the information gathering policies no longer aim to directly optimize the target objective value, but rather choose to follow a greedy trajectory governed by the underlying surrogate function.
These insights have led to principled algorithms allowing a signiﬁcant gain in efﬁciency of the decision making process with a greedy policy, while enjoying strong performance guarantees that are competitive with the optimal policy.
Despite the promising performance, a caveat for these “submodular surrogate”-based approaches is that it is often challenging to engineer such a surrogate objective without an ad-hoc design and analysis.
Furthermore, it has been established that for certain classes of surrogate functions, it is
NP-hard to compute/evaluate the function value (Javdani et al., 2014).
In such cases, even a greedy policy, which iteratively picks the best action given the (observed) history, can be prohibitively costly to design or run.
Therefore it remains a critical challenges for solving more general classes of (sequential) decision making tasks, especially when involving expensive function evaluations.
Overview of main results Inspired by contemporary work in data driven decision making, we aim to learn a greedy heuristic for sequentially selecting actions, as a surrogate for invoking the expensive oracle when evaluating an action.
Our key insight is that many practical algorithms can be interpreted as greedy approaches that follow an (approximate) submodular surrogate objective.
In particular, we focus on the class of combinatorial problems that can be solved via submodular maximization (either directly on the objective function or via a submodular surrogates).
We highlight some of the key results below:
• Focusing on utility-based greedy policies, we introduce a data-driven optimization framework based on the “ submodular-norm ” loss, which is a novel loss function that encourages learning a utility function that exhibits “diminishing returns”.
Our framework, called L EAS U R E (Learning with Submodular Regularization), outputs a surrogate objective that is permutation-invariant, efﬁcient to train, and approximately submodular.
The latter property allows us to prove strong approximation guarantees for the resulting greedy heuristic.
• We show that our approach can be easily integrated with modern imitation learning pipelines for sequential prediction tasks.
We provide rigorous analysis of the proposed algorithm and prove strong performance guarantees for the learned objective.
• We demonstrate the performance of our approach on a variety of decision making tasks, including set cover, active learning for classiﬁcation, and data-driven protein design.
Our results suggest that L EA S U R E not only requires much fewer oracle calls at training phase to learn the target objective (i.e., to minimize the approximation error against the oracle objective), but also outperforms standard learning-based
baselines for solving the corresponding optimization task (i.e., to minimize the regret for the original combinatorial optimization task).
In particular, L EA S U R E has shown promising performance in the protein design task, and will be incorporated into the real-world protein design workﬂow.
R ELATED W ORK
Near-optimal decision making via submodular optimization.
Submodularity is a property of a set function that resembles diminishing returns and it has wide applications from information gathering to document summarization (Krause and Golovin, 2014).
The maximization of a submodular function has been an active area of study in various settings such as centralized (Nemhauser et al., 1978;
Buchbinder et al., 2014; Mitrovic et al., 2017), streaming (Badanidiyuru et al., 2014; Kazemi et al., 2019; Feldman et al., 2020), continuous (Bian et al., 2017b; Bach, 2019) and approximate (Horel and Singer, 2016; Bian et al., 2017a).
The greedy algorithm, which iteratively selects an element that maximizes the marginal gain, and its variants feature prominently in the algorithm design.
For example, in the case of maximizing a monotone submodular function subject to a cardinality constraint, it is shown that the greedy algorithm achieves an approximation ratio of (1 − 1 /e ) of the optimal solution (Nemhauser et al., 1978).
In applications where we need to make a sequence of decisions, such as information gathering, we usually need to adapt our future decisions based on past outcomes.
Adaptive submodularity is the corresponding property where an adaptive greedy algorithm enjoys a similar guarantee for maximizing an adaptive submodular function (Golovin and Krause, 2011).
Recent works have explored optimizing the value of information (Chen et al., 2015b) and Bayesian active learning (Javdani et al., 2014; Chen et al., 2017a) with this property.
Learning submodular functions.
Given the wide applications of submodular functions, the problem of learning submodular functions is the focus of intense study.
Lin and Bilmes (2012) study a mixture of ”submodular shells”, which is an abstract representation of submodular functions, and apply it to the task of document summarization.
Deep submodular functions (Dolhansky and Bilmes, 2016) mimic the structure of deep neural networks and represent a more expressive class of submod- ular functions than sums of concave composed with modular functions.
The theoretical question of the learnability of general submodular functions is analyzed in (Balcan and Harvey, 2018).
Learning to optimize.
Learning from demonstration (LfD), or imitation learning, explores techniques for learning a policy (i.e., a mapping from states to actions) directly from examples provided by an expert (e.g., an expensive computational oracle, or a human instructor)
(Chernova and
Thomaz, 2014).
Classical work on imitation learning (e.g., the Dataset Aggregation (DAgger) algorithm (Ross et al., 2011)) takes reduction-based approaches that reduce the policy learning problem to the supervised learning setting.
Recent works have explored using imitation learning to improve combinatorial optimization solvers.
In (He et al., 2014; Song et al., 2018), the authors learn a node selection policy to replace the default selection rule in a branch-and-bound integer program solver.
Similar ideas have also been applied to learn branching policies (Khalil et al., 2016; Gasse et al., 2019) and decomposition-based approaches (Song et al., 2020).
Learning active learning.
Recently, Konyushkova et al.
(2017) proposed a learning-based framework for active learning (LAL) with manually-engineered feature representation of historical active learning trajectories, and employed standard supervised learning approaches for learning an active learning policy, with the goal of generalizing the learned active learning policy to an unseen dataset.
In (Liu et al., 2018), the authors adapt the DAgger algorithm to provide feedback on which new data points should be added in order to reduce the training loss.
Then an active learning policy is trained to imitate the feedback on selecting new data points.
The batch version of the problem is also studied.
Wei et al.
(2015) focuses on selecting a diverse batch by exploring the submodularity of the problem.
Ash et al.
(2020) uses the gradient signal from the classiﬁer to assess both the “usefulnes” and diversity of the selected samples.
B ACKGROUND AND P ROBLEM S TATEMENT
In this section, we formalize the optimal decision making problem addressed in this paper, and then introduce our learning-based decision making protocol.
3. D ECISION M AKING VIA S UBMODULAR S URROGATE
Given a ground set of items V to pick from.
Let u : 2 V → R be a set function that measures the value of any given subset 1 A ⊆ V .
For example, for experimental design, u ( A ) captures the utility of the output of the best experiment; for active learning u ( A ) captures the generalization error after training with set A .
We denote a policy π : 2 V → V to be a partial mapping from the set/sequence of items already selected, to the next item to be picked.
We use Π to denote our policy class.
Each time a policy picks an item e ∈ V , it incurs a unit cost.
Given the ground set V , the utility function u , and a budget k for selecting items, we seek the optimal policy π that achieves the maximal utility:
[Equation]
Here S π,k denotes the sequence of items picked by π , with S π,i =
S π,i − 1 ∪ { π ( S π,i − 1 ) } for i >and S π, 0 = ∅ .
As we have discussed in the previous sections, many sequential decision making problems can be characterized as constrained monotone submodular maximization problem, i.e., u is
Monotone :
For any A ⊆ V and e ∈ V \ A , f ( A ) ≤ f
( A ∪ { e } ) .
Submodular :
For any A ⊆ B ⊆ V and e ∈ V \ B , f ( A ∪ { e } ) − f ( A ) ≥ f
( B ∪ { e } ) − f ( B ) .
In such cases, a mypopic algorithm following the greedy trajectory of u admits a near-optimal policy.
Note that when u is not monotone submodular, one strategy is to design a surrogate function f : 2 V → R which is
1. globally aligning with u : For instance, f lies within a factor of u : f ( A ) ∈
[ c 1 · u ( A ) , c 2 · u ( A ))] for some constants c 1 , c 2 and any set A ⊆ V ; or within a small margin with u : f ( A ) ∈
[ u ( A ) − (cid:15), u ( A )
+ (cid:15) ] for a ﬁxed (cid:15) > 0 and any set A ⊆ V ;
monotone submodular: Intuitively, a submodular surrogate function encourages selecting items that are beneﬁcial in the long run, while ensuring that the decision maker does not miss out any actions that are “surprisingly good” by following a myopic policy (i.e., future gains for any item are diminishing).
Examples that fall into this category include machine teaching (Singla et al., 2014), active learning (Chen et al., 2015a), etc.
We argue that in real-world decision making scenarios—as validated later in section 6—the decision maker is following a surrogate objective that aligns with the above characterization.
In the following context, we will assume that such surrogate function exists.
Our goal is thus to learn from an expert policy that behaves greedily according to such surrogate functions.
3. L EARNING TO M AKE D ECISIONS
We focus on the regime where the expert policy is expensive to evaluate.
Let g : 2 V × V
→ R be the score function that quantiﬁes the beneﬁt of adding a new item to an existing subset of V .
For the expert policy and submodular surrogate f discussed in section 3.1,
∀ A ⊆ V and e ∈ V
[Equation]
For example, in the active learning case, g exp ( A, e ) could be the expert acquisition function that ranks the importance of labelling each unlabelled point, given the currently labelled subset.
In the set cover case, g exp ( A, e ) could be the function that gives the score to each vertex and determines the next best vertex to add to the cover set.
Given a loss function (cid:96) , our goal is to learn a score function ˆ g that incurs the minimal expected loss when evaluated against the expert policy: ˆ g = arg min g E A,e [ (cid:96) ( g ( A, e ) , g exp ( A, e ))] .
Subsequently, the utility by the learned policy is u ( S ˆ π,k ) , where for any given history A ⊆ V , ˆ π ( A ) ∈ arg max e ∈V ˆ g ( A, e ) .
L EARNING WITH S UBMODULAR R EGULARIZATION
To capture our intuition that a greedy expert policy tends to choose the most useful items, we introduce L EA S U R E , a novel regularizer that encourages the learned score function (and hence surrogate objective) to be submodular.
We describe the algorithm below.
Given the groundset V , let f : 2 V → R be any approximately submodular surrogate such that f ( A ) captures the “usefulness” of the set A .
The goal of a trained policy is to learn a score function g : 2 V × V → R that mimics g exp ( A, x ) =
f ( A ∪ { x } ) − f ( A ) , which is often prohibitively expensive to evaluate exactly.
Then, given any such g , we can deﬁne a greedy policy π ( A ) =
argmax x ∈V g ( A, x ) .
With L EA S U R E , we aim to learn such function g that approximates g exp well while being inexpensive to evaluate at test time.
Let D real
= { ( (cid:104)
A i , x i (cid:105) , y exp i =
g exp ( A i , x i )) } m be the gathered tuple of expert scores for each set - element pair.
If the set 2 V × V was not too large, the L EA S U R E could be trained on the randomly collected tuples D real .
However, 2 V tends to be too large to explore, and generating ground truth labels could be very expensive.
To leverage that, we also collect an unsupervised synthetic dataset of tuples D synth = {
( (cid:104) A, x (cid:105) , (cid:104) A (cid:48) , x (cid:105)| A (cid:22) A (cid:48) } n where (cid:104) A, x (cid:105) denote
a labelled tuple (cid:104) A, x
(cid:105) ∈
D real , and A (cid:48) denote a randomly selected superset of
A .
Deﬁne:
Loss ( g, g exp ) =
(cid:88)
[Equation]
where λ > 0 is the regularization parameter.
Intuitively [Equation] such regularization term will force the learned function g to be close to submodular, as it will lead to larger losses every time g ( A (cid:48) , x ) >
g ( A, x ) .
If we expect f to be monotonic, we also introduce a second regularizer ReLu ( − g ( A (cid:48) , x )) which pushes the learned function to be positive.
Combined, the loss function becomes (used in line 11 in Algorithm 1):
[Equation]
where γ is another regularization strength parameter.
Such loss should push g to explore a set of approximately submodular, approximately monotonic functions.
Thus, if f exhibits the submodular and monotonic behavior, g trained on this loss function should achieve a good local minima.
Finally, before we provide the algorithm, note that since 2 V is too large to explore, instead of sampling random tuples for D real , we use modiﬁed DAgger.
Then g can learn not only from the expert selections of (cid:104) A, x (cid:105) , but it can also see the labels of the tuples the expert would not have chosen.
The algorithm is provided below:
Here
(cid:105) ..., (cid:104){ x 1 , ..., x T − 1 } , x T (cid:105) ) , collected using a mixed policy π i .
On line 8, expert feedback of selected actions is collected to form D i .
We also want to note that in some settings, even collecting exact expert labels g exp at train time could be too expensive.
In that case, g exp can be replaced with a less expensive, noisy approximate expert g exp (cid:15)
≈ g exp .
In fact, all three of our experiments use noisy experts in one form or another.
A NALYSIS
Estimating the expert’s policy We ﬁrst consider to bound the loss of the learned policy measured against the expert’s policy.
Since L EA S U R E could be viewed as a specialization of DA GGER (Ross et al., 2011) for learning a submodular function, it naturally inherits the performance guarantees from DA GGER , which show that the learned policy efﬁciently converges to the expert’s policy.
Concretely, the following result, which is proved in appendix, shows that the learned policy is consistent with the expert policy and thus is a no-regret algorithm:
Theorem 1 (Theorem 3.3, Ross et al.
Denote the loss of ˆ π at history state H as l ( H, ˆ π ) :
= (cid:96) ( g ( H, ˆ π ( H )) , g exp ( H, π exp ( H ))) .
Let
d
ˆ π be the average distribution of states if we follow ˆ π for a ﬁnite number of steps.
Furthermore, let D i be a set of m random trajectories sampled with π i at round i ∈ { 1 , . . .
, N } ,
and ˆ (cid:15) N = min π 1 N (cid:80) Ni =1 E H i ∼ D
i [ l ( H i , ˆ π )] be the training loss of the best policy on the sampled trajectories.
If N is O (cid:0) T 2 log(1 /δ ) (cid:1) and m is O (1) then with probability at least 1 − δ there exists a ˆ π among the N policies, with E H ∼ d ˆ π [ l ( H, ˆ π )] ≤ ˆ (cid:15) N + O (cid:0) 1 T (cid:1) .
Approximating the optimal policy Note that the previous notion of regret corresponds to the average difference in score function between the learned policy and the expert policy.
While this result shows that L EA S U R E is consistent with the expert, it does not directly address how well the learned policy performs in terms of the gained utility.
We then provide a bound on the expected value of the learned policy, measured against the value of the optimal policy.
For speciﬁc decision making tasks where the oracle follows an approximately submodular objective, our next result, which is proved in the appendix, shows that the learned policy behaves near-optimally.
Theorem 2.
Assume that the utility function u is monotone submodular.
Furthermore, assume the expert policy π exp follows a surrogate objective f such that for all A ⊆ V , | f ( A ) − u
( A ) | < (cid:15) E where (cid:15) E > 0 .
Let ˆ (cid:15) N = min
π 1 N (cid:80) Ni =1 l ( H i , ˆ π ) be the training loss of the best policy on the sampled trajectories.
If N is O (cid:0) T 2 log(1 /δ ) (cid:1) then with probability at least 1 − δ , the expected utility achieved by running ˆ π for k steps is
[Equation]
E XPERIMENTS
In this section, we demostrate the performance of L EA S U R E on three diverse sequential decision making tasks, namely set cover, learning active learning for classiﬁcation and protein engineering.
Baselines We compare our approach to the Deep Submodular Function (Dolhansky and Bilmes (2016)) and Deep Batch Active Learning by Diverse, Uncertain Gradient Lower Bounds (BADGE,
Ash et al. In their paper, Dolhansky and Bilmes (2016) learn a submodular surrogate function f : 2 V → R that produces a score for each set A ⊂ V .
The architecture of the Deep Submodular Function (DSF) forces the function f to be exactly submodular, as opposed to L EA S U R E , which is only encouraged to be submodular through a regularizer.
However, the architecture and the training procedure of the DSF are quite restrictive, which does not allow the DFS to explore a large domain during training and restricts how expressive it can be compared to a standard neural network.
Moreover, DSF are restricted to small V , and the number of parameters increases with the size of V .
That is not true for L EA S U R E , which number of parameters grows with the dimensionality of elements in V .
This makes DSF useful for small datasets, but makes it prohibitively expensive to use on larger problems.
In fact, we could not compare L EA S U R E to DSF on Learning Active Learning or Protein
Engineering Tasks, as it was not feasible to train DSF on these sets.
For Learning Active Learning experiment, we also compare with a recent deep active learning approach (Ash et al., 2020).
Finally, we want to add that L EA S U R E can be seamlessly integrated with any standard Machine Learning library, and since the architecture of the learned policy in L EA S U R E is not restrictive, any available optimization trick can be used to achieve better performance.
On the other hand, DSF cannot be as easily implemented, and the standard libraries are not optimized for the DSF architecture.
6. S ET C OVER
Before testing our approach on a real-world scenario, we showcase its performance on a simple submodular and monotonic maximization problem.
Set cover is a classical example: given a set of elements U = { 1 , 2 , ..., n } (called the universe) and a collection of m sets S = { s 1 , .., s m } whose union equals the universe, the set cover problem is to identify the smallest sub-collection of S whose union equals the universe.
Formulated as a policy learning problem, the goal is to learn the score function g : 2 S × S → R such that for any S l ⊂ S, x ∈ S ,
[Equation]
Given such a function g , we can then deﬁne a policy π : 2 s → S as π ( S l ) = argmax x ∈ S g ( S l , x ) .
During training, tuples { ( S l , x ) , g exp } are collected, and then g is trained on this set.
We trained four different policies: a function g parametrized by a neural network with MSE ( g, g exp ) as the loss, a function g with the same loss and just a monotonicity regularizer, a function g trained using both monotonicity and submodular regularizers, as well as the Deep Submodular Function baseline (Dolhansky and Bilmes (2016)).
Our dataset is the subset of the Mushroom dataset Lim (2015), consisting of 1000 sets.
Each set contains 23 mushroom species, and there are a total of 119 species.
The goal is then to train a policy to select the largest superset of these sets.
We train and test learned policies in two settings: Exact Set Cover, where we collect tuples { ( S l , x ) , g exp } for training, and
Noisy Set Cover, where we have access only to { ( S l , x ) , g exp (cid:15) } , where g exp (cid:15) is a noisy score.
The networks are trained on rollouts of length 20 (i.e. on sets { S l : | S l | ≤ 20 } ), and tested on rollout of length up to 100 .
The results of training policies on these two frameworks are in Figure 1.
In 1a and 1b, we plot the value of set cover as a function of the size of the superset.
L EA S U R E signiﬁcantly outperforms other learned policies, although Deep Submodular Function generalized better to larger rollout lengths - L EA S U R E gets most of its set cover gains in the ﬁrst 10 − 20 selected points, while
Deep Submodular Function continues to noticeably improve past the training rollout length.
Note that in 1a and 1b, the competing baselines all exhibit an “diminishing returns” effect, resulting in a concave-shaped value function.
With a submodular-norm regularizer, L EA S U R E quickly idenﬁed the sets with large marginal gains.
This observation aligns with our analysis in section 5. [Figure]
6. L EARNING A CTIVE L EARNING ON F ASHION MNIST
In this section we demonstrate the performance of L EA S U R E on a real-world task that is not submodular or monotonic, but that usually exhibits submodular and monotonic behaviour.
In particular, in an active learning framework, there is a partially labelled dataset S =
{ S l , S u } and a policy π : 2 S → S .
The labelled subset S l can be used to infer from data (learn the image classiﬁer, predict unlabelled protein ﬁtness, etc).
The goal of the policy is to select the smallest subset S π ⊂ S u to label such that the accuracy of supervised learning from S π ∪ S l is maximized.
Since selecting a subset is a prohibitively expensive combinatorial task, the policy is usually sequential.
In particular, it selects points to add to S π one by one (or in batches) using some score function g ( S π ∪ S l , · ) :
S u
→ R to score each point x ∈ S u and then the policy labels the point with the largest score.
If g were to be the ﬁrst order difference of a submodular function f , i.e. g ( A, e ) = f
( A ∪{ e } ) − f ( A ) , then the policy would be near-optimal.
Moreover, as discussed above, intuitively we expect g to have this property in most cases, since adding an extra point to a larger set usually has less effect than adding the same point to a smaller subset of the set.
The above motivates the use of L EA S U R E in active learning (Figure 2).
In this experiment, the set
S is a fashion MNIST dataset of 28 × 28 pixels, greyscale images of clothes that come from one of the 10 classes.
The goal was to learn a policy that greedily selects a point x ∈ S u to label, such that a neural network classiﬁer trained on the labelled set S l ∪ { x } produces the most accurate classiﬁcation of the unlabelled images.
In particular, we trained the above function g to predict the accuracy gain g exp from labelling a point.
The accuracy gain g exp was measured by training the neural network classiﬁer on both S l and S l ∪ { x } and then recording the difference in validation set classiﬁcation accuracy.
Since obtaining exact g exp for each datapoint is very expensive, we instead collected noisy labels g exp (cid:15)
≈ g exp , obtained by training the classiﬁer for only 10 epochs.
The tuples { ( S l , x ) , g exp (cid:15) ) } were collected using DAgger with rollouts of length 30 (starting from a random batch of 20 images).
For training, we used an initially unlabelled dataset with 60000 images,of which were set aside to use for evaluating validation accuracy.
We trained two neural networks to approximate g - an unregularized one, and one with a monotonicity and a submodularity regularizer (i.e. L EA S U R E ). [Figure]
The trained policies were tested on a set of 8000 images, with additional 2000 set aside for validation.
At test time, we again started with a random batch of size 20 and then used each policy to sequentially select additional 200 images to label (Figure 2).
The recorded test error rate was collected using real g exp , i.e. a classiﬁer trained until training loss reaches a certain threshold.
Even though L EA S U R E was trained on much shorter rollouts using very noisy labels, it still outperformed all other baselines.
The submodular regularizer allows the learned score function g to ﬁnd a local minima that generalizes well to out of sample.
6. P ROTEIN E NGINEERING
By employing a large protein engineering database containing mutation-function data (Wang et al., 2019), we demonstrate that L EA S U R E enables the learning of an optimal policy for imitating expert design of protein sequences (see Appendix for detailed discussion of datasets).
As in (Liu et al., 2018) we construct a fully data-driven expert which evaluates via one step roll-out the effect of labeling each candidate data (in our case a protein mutant) with the objective of minimizing loss on a downstream regression task (predicting protein ﬁtness).
When training the policy to emulate the algorithmic expert via imitation learning, we represent each state as two merged representations: (1) a ﬁxed dimensional representation of the protein being considered (as the last dense layer of the network described in Appendix C), and (2) a similar ﬁxed dimensional representation of the data already included in the training set (as a sum of their embeddings), including their average label value.
At each step, the expert policy greedily selects a protein to label, which minimizes the expected regression loss on the downstream regression task (prediction of protein ﬁtness). Once the expert policy selects a protein sequence, we then add it into the training set, retrain the model, and use the updated model to predict a protein with the maximum ﬁtness. Once the complete pool of data has been evaluated, the states are stored along with their associated preference score, taken as their ability to reduce the loss in the one-step roll out. This paired state action data is used to train the policy model at the end of each episode, as described in (Liu et al., 2018).
As we observe in Figure 5a, this method trains a policy which performs nearly identically to this limited rollout, greedy oracle.
The use of submodular regularization enables the learning of a policy which generalizes to a fundamentally different protein engineering task.
In our experiments, L EA S U R E is trained to emulate a greedy oracle for maximizing the stability of protein G, a small bacterial protein used across a range of biotechnology applications (Sjbring et al., 1991).
We evaluate our results by applying the trained policy to select data for the task of predicting antibody binding of a small molecule.
As is the case with all protein ﬁtness landscapes, the evaluation dataset is highly imbalanced, with the vast majority of mutants conferring no improvement at all.
Because data is expensive to label in biological settings (proteins must be synthesized, puriﬁed and tested), we are often limited in how many labels can feasibly be generated, and the discriminative power among the best results is often more important than among the worst.
To construct a metric with real-world applicability we assess each model by systemically examining the median Kd of the next ten data points selected at each budget, from 10 to 110 total labels.
This method is utilized in recognisance of the extreme ruggedness of protein engineering landscapes, wherein the vast majority of labels are of null ﬁtness, and the ability to select rare useful labels for the next experimental cycle is of key importance.
We observe that L EA S U R E outperforms all evaluated baselines, and that the inclusion of submodular optimization is mandatory to its success (ﬁgure 7b).
A greedy active learner which labels the antibody mutation with the best predicted Kd (the smallest) preforms approximately equivalently with selecting random labels.
Use of dropout as an approximation of model uncertainty as in (Gal and
Ghahramani, 2015) improves upon these baselines, although signiﬁcant betterment is not achieved until approximately 35 labels are added.
In comparison, the results from L EA S U R E diverge from all others nearly immediately, and the best model, which uses a lambda of 0.1, achieves a notable improvement in Kd, 5 . 81 µ M, vs 7 .
27 µ M achieved by entropy sampling.
In support of methods success, we note that the learned policy preforms approximately as well as the greedy oracle which it emulates (Appendix ﬁgure 7b).
We observe that the results are robust within a range of possible lambda values (Fig 2b), and that without the use of submodular regularization the trained policy fails to learn a policy better than the selection of random labels. This is an important ﬁnding, as the method proposed by (Liu et al., 2018) without L EA S U R E , has been shown to be a state-of-the-art method for imitation learning. Based on these empirical results, L EA S U R E demonstrates signiﬁcant potential as computational tool for real-world automated experimental design tasks : In particular, in the protein engineering task, L EA S U R E achieves the SOTA on the benchmark data-sets considered in this work.
While
L EA S U R E does involve repeated retraining of the protein engineering network, we observe that it returns strong results even with a single step of training.
Additionally, the networks that are employed are very simple (Appendix C). This allows for reasonable training time (36 hours) and nearly instantaneous inference. Given the considerable time and cost of protein engineering, these computational budgets are quite modest. Protein engineering is a time consuming (months to years) and expensive undertaking (10’s of thousands to millions of dollars).
These projects usually strive to achieve the best possible results given a ﬁxed budget.
We have demonstrated in our work the ability deliver signiﬁcant improvements in protein potency for the modest ﬁxed budgets.
Although the cost savings of engineering and testing an individual protein (or label) vary signiﬁcantly based on the system, ranging tens to hundreds of dollars, we observe that to achieve a Kd of 8e-6 M L EA S U R E delivers an approximate cost savings of 65%, or 40 fewer labels than the next best method.
The sequential synthesis and evaluation of each of these labels would likely span several months and additionally incur several thousands of dollars of materials costs.
In this paper, we introduce L EA S U R E , a data-driven decision making framework based on a novel submodular-regularized loss function.
The algorithm was inspired by the recent developments of submodular-surrogate-based near-optimal algorithms for sequential decision making.
We have demonstrated L EA S U R E on several diverse set of decision making tasks.
Our results suggest that L EA S U R E can be easily integrated with modern deep imitation learning pipelines, and that it is efﬁcient to run, while still reaching the best performance among the competing baselines.
In addition to demonstrating the strong empirical performance on several use cases, we believe our work also provides useful insights in the design and analysis of novel information acquisition heuristics where traditional ad-hoc approaches are not feasible.
A P ROOF FOR SECTIONA. P ROOF OF T HEOREM
Proof.
The high-level idea is to ﬁrst connect the total expected utility of the learned policy ˆ π with the expected utility of the expert policy π exp , following the analysis in DAgger (Ross et al., 2011).
Then, we will use the fact that π exp is greedy with respect to f , an approximation to the submodular utility function u , to bound the one step gain of the π exp against the k step gain of running the optimal policy, and subsequently bound the total utility of the expert policy against the optimal policy.
We would eventually obtain a similar result as Theorem 2, detailed as follows.
More concretely, following Theorem 3.4 in DAgger, we obtain that
[Equation]
Here ∆ max is the largest one-step deviation from π exp that ˆ π can suffer.
It is equivalent to the term u in the DAgger paper.
Since f is (cid:15) -close to a monotone submodular function u , we know that ∆ max ≤ max A ⊂V , | A | = k f ( A ) ≤ max A ⊂V , | A | = k u ( A ) + (cid:15) E , which is a constant once u is given.
Next, since π exp is greedily optimizing an (cid:15) E -approximation to a monotone submodular function u , we know that
[Equation]
following the proof from Theorem 5 in (Chen et al., 2017b).
Combining both steps, we [-get ∆ b ≤ (cid:0) 1 − 1 k (cid:1) b ∆ 
 + (cid:80) bi =0 (cid:0) 1 − 1 k (cid:1) i . Hence-]have that
[Equation]
which completes the proof.
B S UPPELEMENTAL D ETAILS FOR THE S ET C OVER AND MNIST A CTIVE L EARNING E XPERIMENTS
We provide additional results for the set cover experiments, under the same experimental setup as
Figure 1a and 1b.
The subplots 4a and 4b show the mean square error of learned policy g as a function of the size of S l .
We provide a zoomed-in version of 4b in Figure 4c.
In Figure 4c, we show it is clear that training the neural network on the monotonicity regularizer only does not help it learn out of sample - the error rapidly increases as soon as the test rollout length becomes larger than the training rollout length. In Noisy Set Cover experiment (Figure 4a), each label of the element added to the superset was perturbed with N (0 , 1) noise. As a result, the variance of the total noise is linear in the number of sets.
So, it is reasonable that the MSE error grows with number of sets - the policies cannot learn to predict random noise. While stochastic MSE of L EA S U R E and the no-regularizer policy are similar, L EA S U R E outperforms in the number of elements added, which is what matters in practice (Figure 1).
These two Figures conﬁrm our intuition that when the problem is not exactly submodular, Leasure will still generalize better than no regularizer by learning to ignore small deviations from submodularity. Finally, it is also expected that DSF has a lower MSE than Leasure when the label noise is too large - Deep Submodular Functions are required to be submodular. When the stochasticity in the MSE becomes overwhelmingly large, that restrictive requirement becomes an advantage.
However, when the MSE variance is not too large, the lack of expressiveness and the difﬁculty of optimization of DSF makes it lose its advantage compared to Leasure.
[Figure]
For completion, we also provide our architecture and parameter choices for both set cover and
Learning Active Learning on MNIST experiments.
For set cover, the problem is too simple to require DA GGER (Ross et al., 2011).
Instead, the tuples are generated randomly.
For active learning on MNIST, the tuples are indeed generated using Algorithm 1.
For MNIST, we ﬁrst preprocessed our dataset with PCA, leaving the number of vectors necessary to achieve 80% covariance on the training set (24 vectors).
For set cover, each element was a set v containing 23 elements v 1 , v 2 , .., v 23 , where v i was an integer corresponding to the label of the species.
As a neural network input, v was simply represented as a vector of [ v 1 , ..., v 23 ] .
Both set cover and MNIST used a modiﬁed Deepset architecture (Zaheer et al.
for score networks as follows:
Given a set A = { v 0 , ..., v k } ⊂ V and a datapoint v ∈ V , the score network g ﬁrst preprocesses all inputs v 0 , ..., v k , v to obtain learned embeddings ¯ v 0 , .., ¯ v k , ¯ v .
(See Figure 5)
Then, the elements in A are combined using Deepsets architecture to produce a learned set embedding ¯ A .
Finally, ¯ A and ¯ v are concatenated and then a learned linear layer and a Leaky ReLu nonlinearity are applied to produce g ( A, v ) .
(See Figure 6).
All dense layers have 64 neurons and a bias term.
Using this Deepsets-like framework, we achieve permutation invariance of elements in set A while also keeping the network expressive enough to learn a wide range of functions.
[Figure]
Learning element representation
[Figure]
Combining element representation using DeepSets
For both tasks, the score networks are trained using ADAM with a learning rate of 1e-3. Beta parameter from Line 2 in L EA S U R E was picked randomly to be 45 . From experiments, the exact value of the parameter did not matter as long as it starts with at least 12 and degrades towards almost 0 after N iterations. The λ and γ parameters were picked using a hyperparameter sweep in log space.
As per our intuition, we have found that the strength of the parameters should reﬂect your certainty that the task is submodular and/or monotone.
For set cover, λ = 0 . 1 , γ = 0 . 5 , while for active learning λ = 0 . 001 and γ = 0 .
Notice that the values are not comparable between different experiments: for MNIST Learning Active Learning (LAL), g exp ( A, v ) ∈ [0 , 1) outputs the accuracy gain of adding v to A and training a supervised model on it; for set cover, g exp ( A, v ) ∈ { 0 , 1 } outputs the number of new elements added to the set by adding x to A .
For LAL, the values of g exp are usually much smaller than 1 , particularly for larger sets.
Thus, the values for the two regularizers had to be smaller so that the model learns not just the regularizer.
Finally, we wanted to discuss our baselines in Fashion MNIST experiments.
In Figure 2, we have four baselines: random, uncertainty, BADGE (Ash et al., 2020), and no regularizer.
The no regularizer baseline was trained identically to L EA S U R E , except for the absense of submodularity and monotonicity regularizers.
The no regularizer baseline performed well on the sets with up toadditional points - corresponding exactly to the length of the training rollouts.
However, it failed to generalize.
On the other hand, the submodular regularizer allowed the learned score function to ﬁnd a local minima that generalized well to out of sample.
Finally, BADGE did not seem to perform well when the number of datapoints in the set was large, likely because the gradient signal from adding any one additional datapoint was too weak and thus the selection of the next best datapoint was too noisy.
Some more details regarding BADGE (Ash et al., 2020).
The authors do not learn a policy, instead, they use gradients of the classiﬁer (gradient embedding) to select a useful, diverse batch.
Although
BADGE was originally made for a batch setting, the authors’ main idea is still applicable to our case: they argued that the next datapoint(s) can be selected by looking at which ﬁctitious labels would produce the largest gradients in the classiﬁer network.
Therefore, we replaced the kmeans++ algorithm the authors suggested with simply selecting the datapoint that corresponds to the largest gradient norm.
This algorithm has an advantage that it does not require a trained policy network.
However, it provides no guarantees about submodularity of the resulting policy, and, in our experiments, the performance degrades with the size of the set - likely because the gradient signal from adding any one additional datapoint was too weak and thus the selection of the next best datapoint was too noisy.
Since BADGE requires a neural network classiﬁer/regressor, we could not use it as a baseline for Set Cover (Set Cover regression function is simply adding all elements in the superset).
The no-regularizer baseline is similar to that of Konyushkova et al.
However, the problem considered in Konyushkova et al.
(2017) is not compatible with most of the tasks we considered here (for MNIST, yes if we use random forest classiﬁers; but for others not).
Furthermore, Konyushkova et al.
(2017) treated the problem under a classical supervised learning setting this is often not desirable, given that we are learning a policy from non i.i.d.
data samples. C S UPPELEMENTAL D ETAILS FOR THE P ROTEIN E NGINEERING E XPERIMENTS Dataset Our datasets were identiﬁed in Protabank (Wang et al., 2019) for training of active learning policies and benchmarking of performance.
In selecting datasets upon which to train our active learning models several factors were considered.
As the state space of possible protein variants for typical engineering application is very large, size is our foremost criteria.
Additionally it will be advantageous to use datasets which characterize mutations to all amino acids (as opposed to Alanine scans), and those which include epistatic interactions.
We also desire to identify datasets which have a high quality, quantitative readout, such as calorimetry, ﬂuorescence, or SPR data.
Protein Engineering Methods Embeddings of protein sequences were created using the TAPE repository (Rao et al., 2019) according to the UniRep system as ﬁrst proposed in Alley et al.
UniRep produces protein embeddings as a matrix of shape (length protein sequence, 1900), although we average together the embeddings only of positions being engineered to produce a consistent embedding of shape (1900,).
We have implemented the active learning imitation learning algorithm proposed in Liu et al.
(2018) to work with the protein embedding representations described above.
Pseudocode for this method is presented in Algorithms 1 and 2 from the original work. As in
Liu’s work, our policy network consists of a single dense unit which acts sequentially on the pool of samples being considered to produce a preference score.
Our downstream protein engineering network (which was used to compute the preference score of the expert policy) acts on the protein embeddings prepared using TAPE.
The network consists of an attention layer, followed by a 1dimensional convolution layer (128 ﬁlters, kernel size 3), before being ﬂattened and applying two fully connected layers of 128 units each.
When predicting protein ﬁtness, dropout is applied with a probability of 0.5 and an additional dense layer is applied with one unit and linear activation. Both networks are trained using ADAM with a learning rate of 1e-3. The implementation of this part of the project is nearly identical to Liu, only changing the data representation, protein ﬁtness network structure, and values of K (30), B (100) and T (20) as listed in the appendix of our work. Beta is ﬁxed at 0.5, although the method was shown to be robust to a range of values.
At training time,labels are randomly selected for evaluating the effect of the greedy oracle, and 10 data are randomly selected to form the initial data set for learning.
The superset is appended at each step of training the policy to maintain a size of 2x the labeled dataset.
The training of a policy using these settings takes 36 hours on a modern multiprocessor computer equipped with an NVIDIA Titan V GPU.
[Figure]