We consider the two related problems of detecting if an example is misclassified or out-of-distribution. We present a simple baseline that utilizes probabilities from softmax distributions. Correctly classified examples tend to have greater maximum softmax probabilities than erroneously classified and out-of-distribution examples, allowing for their detection. We assess performance by defining several tasks in computer vision, natural language processing, and automatic speech recognition, showing the effectiveness of this baseline across all. We then show the baseline can sometimes be surpassed, demonstrating the room for future research on these underexplored detection tasks.
1 INTRODUCTION
When machine learning classifiers are employed in real-world tasks, they tend to fail when the training and test distributions differ. Worse, these classifiers often fail silently by providing highconfidence predictions while being woefully incorrect (Goodfellow et al., 2015; Amodei et al., 2016). Classifiers failing to indicate when they are likely mistaken can limit their adoption or cause serious accidents. For example, a medical diagnosis model may consistently classify with high confidence, even while it should flag difficult examples for human intervention. The resulting unflagged, erroneous diagnoses could blockade future machine learning technologies in medicine. More generally and importantly, estimating when a model is in error is of great concern to AI Safety (Amodei et al., 2016).
These high-confidence predictions are frequently produced by softmaxes because softmax probabilities are computed with the fast-growing exponential function. Thus minor additions to the softmax inputs, i.e. the logits, can lead to substantial changes in the output distribution. Since the softmax function is a smooth approximation of an indicator function, it is uncommon to see a uniform distribution outputted for out-of-distribution examples. Indeed, random Gaussian noise fed into an MNIST image classifier gives a “prediction confidence” or predicted class probability of 91%, as we show later. Throughout our experiments we establish that the prediction probability from a softmax distribution has a poor direct correspondence to confidence. This is consistent with a great deal of anecdotal evidence from researchers (Nguyen & O’Connor, 2015; Yu et al., 2010; Provost et al., 1998; Nguyen et al., 2015).
However, in this work we also show the prediction probability of incorrect and out-of-distribution examples tends to be lower than the prediction probability for correct examples. Therefore, capturing prediction probability statistics about correct or in-sample examples is often sufficient for detecting whether an example is in error or abnormal, even though the prediction probability viewed in isolation can be misleading.
These prediction probabilities form our detection baseline, and we demonstrate its efficacy through various computer vision, natural language processing, and automatic speech recognition tasks. While these prediction probabilities create a consistently useful baseline, at times they are less effective, revealing room for improvement. To give ideas for future detection research, we contribute
∗Work done while the author was at TTIC. Code available at github.com/hendrycks/error-detection
one method which outperforms the baseline on some (but not all) tasks. This new method evaluates the quality of a neural network’s input reconstruction to determine if an example is abnormal.
In addition to the baseline methods, another contribution of this work is the designation of standard tasks and evaluation metrics for assessing the automatic detection of errors and out-of-distribution examples. We use a large number of well-studied tasks across three research areas, using standard neural network architectures that perform well on them. For out-of-distribution detection, we provide ways to supply the out-of-distribution examples at test time like using images from different datasets and realistically distorting inputs. We hope that other researchers will pursue these tasks in future work and surpass the performance of our baselines.
In summary, while softmax classifier probabilities are not directly useful as confidence estimates, estimating model confidence is not as bleak as previously believed. Simple statistics derived from softmax distributions provide a surprisingly effective way to determine whether an example is misclassified or from a different distribution from the training data, as demonstrated by our experimental results spanning computer vision, natural language processing, and speech recognition tasks. This creates a strong baseline for detecting errors and out-of-distribution examples which we hope future research surpasses.
2 PROBLEM FORMULATION AND EVALUATION
In this paper, we are interested in two related problems. The first is error and success prediction: can we predict whether a trained classifier will make an error on a particular held-out test example; can we predict if it will correctly classify said example? The second is in- and out-of-distribution detection: can we predict whether a test example is from a different distribution from the training data; can we predict if it is from within the same distribution?1 Below we present a simple baseline for solving these two problems. To evaluate our solution, we use two evaluation metrics.
Before mentioning the two evaluation metrics, we first note that comparing detectors is not as straightforward as using accuracy. For detection we have two classes, and the detector outputs a score for both the positive and negative class. If the negative class is far more likely than the positive class, a model may always guess the negative class and obtain high accuracy, which can be misleading (Provost et al., 1998). We must then specify a score threshold so that some positive examples are classified correctly, but this depends upon the trade-off between false negatives (fn) and false positives (fp).
Faced with this issue, we employ the Area Under the Receiver Operating Characteristic curve (AUROC) metric, which is a threshold-independent performance evaluation (Davis & Goadrich, 2006). The ROC curve is a graph showing the true positive rate (tpr = tp/(tp + fn)) and the false positive rate (fpr = fp/(fp + tn)) against each other. Moreover, the AUROC can be interpreted as the probability that a positive example has a greater detector score/value than a negative example (Fawcett, 2005). Consequently, a random positive example detector corresponds to a 50% AUROC, and a “perfect” classifier corresponds to 100%.2
The AUROC sidesteps the issue of threshold selection, as does the Area Under the Precision-Recall curve (AUPR) which is sometimes deemed more informative (Manning & Schütze, 1999). This is because the AUROC is not ideal when the positive class and negative class have greatly differing base rates, and the AUPR adjusts for these different positive and negative base rates. For this reason, the AUPR is our second evaluation metric. The PR curve plots the precision (tp/(tp+ fp)) and recall (tp/(tp + fn)) against each other. The baseline detector has an AUPR approximately equal to the precision (Saito & Rehmsmeier, 2015), and a “perfect” classifier has an AUPR of 100%. Consequently, the base rate of the positive class greatly influences the AUPR, so for detection we must specify which class is positive. In view of this, we show the AUPRs when we treat success/normal classes as positive, and then we show the areas when we treat the error/abnormal classes as positive. We can treat the error/abnormal classes as positive by multiplying the scores by −1 and labeling them positive. Note that treating error/abnormal classes as positive classes does not change the AU-
1We consider adversarial example detection techniques in a separate work (Hendrycks & Gimpel, 2016a). 2A debatable, imprecise interpretation of AUROC values may be as follows: 90%—100%: Excellent,
80%—90%: Good, 70%—80%: Fair, 60%—70%: Poor, 50%—60%: Fail.
ROC since if S is a score for a successfully classified value, and E is the score for an erroneously classified value, AUROC = P (S > E) = P (−E > −S). We begin our experiments in Section 3 where we describe a simple baseline which uses the maximum probability from the softmax label distribution in neural network classifiers. Then in Section 4 we describe a method that uses an additional, auxiliary model component trained to reconstruct the input.
3 SOFTMAX PREDICTION PROBABILITY AS A BASELINE
In what follows we retrieve the maximum/predicted class probability from a softmax distribution and thereby detect whether an example is erroneously classified or out-of-distribution. Specifically, we separate correctly and incorrectly classified test set examples and, for each example, compute the softmax probability of the predicted class, i.e., the maximum softmax probability.3 From these two groups we obtain the area under PR and ROC curves. These areas summarize the performance of a binary classifier discriminating with values/scores (in this case, maximum probabilities from the softmaxes) across different thresholds. This description treats correctly classified examples as the positive class, denoted “Success” or “Succ” in our tables. In “Error” or “Err” we treat the the incorrectly classified examples as the positive class; to do this we label incorrectly classified examples as positive and take the negatives of the softmax probabilities of the predicted classes as the scores.
For “In,” we treat the in-distribution, correctly classified test set examples as positive and use the softmax probability for the predicted class as a score, while for “Out” we treat the out-of-distribution examples as positive and use the negative of the aforementioned probability. Since the AUPRs for Success, Error, In, Out classifiers depend on the rate of positive examples, we list what area a random detector would achieve with “Base” values. Also in the upcoming results we list the mean predicted class probability of wrongly classified examples (Pred Prob Wrong (mean)) to demonstrate that the softmax prediction probability is a misleading confidence proxy when viewed in isolation. The “Pred. Prob (mean)” columns show this same shortcoming but for out-of-distribution examples.
Table labels aside, we begin experimentation with datasets from vision then consider tasks in natural language processing and automatic speech recognition. In all of the following experiments, the AUROCs differ from the random baselines with high statistical significance according to the Wilcoxon rank-sum test.
3.1 COMPUTER VISION
In the following computer vision tasks, we use three datasets: MNIST, CIFAR-10, and CIFAR100 (Krizhevsky, 2009). MNIST is a dataset of handwritten digits, consisting of 60000 training and 10000 testing examples. Meanwhile, CIFAR-10 has colored images belonging to 10 different classes, with 50000 training and 10000 testing examples. CIFAR-100 is more difficult, as it has 100 different classes with 50000 training and 10000 testing examples.
In Table 1, we see that correctly classified and incorrectly classified examples are sufficiently distinct and thus allow reliable discrimination. Note that the area under the curves degrade with image recognizer test error.
Next, let us consider using softmax distributions to determine whether an example is in- or outof-distribution. We use all test set examples as the in-distribution (positive) examples. For out-ofdistribution (negative) examples, we use realistic images and noise. For CIFAR-10 and CIFAR-100, we use realistic images from the Scene UNderstanding dataset (SUN), which consists of 397 different scenes (Xiao et al., 2010). For MNIST, we use grayscale realistic images from three sources. Omniglot (Lake et al., 2015) images are handwritten characters rather than the handwritten digits in MNIST. Next, notMNIST (Bulatov, 2011) consists of typeface characters. Last of the realistic images, CIFAR-10bw are black and white rescaled CIFAR-10 images. The synthetic “Gaussian” data
3We also tried using the KL divergence of the softmax distribution from the uniform distribution for detection. With divergence values, detector AUROCs and AUPRs were highly correlated with AUROCs and AUPRs from a detector using the maximum softmax probability. This divergence is similar to entropy (Steinhardt & Liang, 2016; Williams & Renals, 1997).
is random normal noise, and “Uniform” data is random uniform noise. Images are resized when necessary.
The results are shown in Table 2. Notice that the mean predicted/maximum class probabilities (Pred. Prob (mean)) are above 75%, but if the prediction probability alone is translated to confidence, the softmax distribution should be more uniform for CIFAR-100. This again shows softmax probabilities should not be viewed as a direct representation of confidence. Fortunately, out-of-distribution examples sufficiently differ in the prediction probabilities from in-distribution examples, allowing for successful detection and generally high area under PR and ROC curves.
For reproducibility, let us specify the model architectures. The MNIST classifier is a three-layer, 256 neuron-wide, fully-connected network trained for 30 epochs with Adam (Kingma & Ba, 2015). It uses a GELU nonlinearity (Hendrycks & Gimpel, 2016b), xΦ(x), where Φ(x) is the CDF of the standard normal distribution. We initialize our weights according to (Hendrycks & Gimpel, 2016c), as it is suited for arbitrary nonlinearities. For CIFAR-10 and CIFAR-100, we train a 40-4 wide residual network (Zagoruyko & Komodakis, 2016) for 50 epochs with stochastic gradient descent using restarts (Loshchilov & Hutter, 2016), the GELU nonlinearity, and standard mirroring and cropping data augmentation.
3.2 NATURAL LANGUAGE PROCESSING
Let us turn to a variety of tasks and architectures used in natural language processing.
3.2.1 SENTIMENT CLASSIFICATION
The first NLP task is binary sentiment classification using the IMDB dataset (Maas et al., 2011), a dataset of polarized movie reviews with 25000 training and 25000 test reviews. This task allows us to determine if classifiers trained on a relatively small dataset still produce informative softmax
distributions. For this task we use a linear classifier taking as input the average of trainable, randomly initialized word vectors with dimension 50 (Joulin et al., 2016; Iyyer et al., 2015). We train for 15 epochs with Adam and early stopping based upon 5000 held-out training reviews. Again, Table 3 shows that the softmax distributions differ between correctly and incorrectly classified examples, so prediction probabilities allow us to detect reliably which examples are right and wrong.
Now we use the Customer Review (Hu & Liu, 2004) and Movie Review (Pang et al., 2002) datasets as out-of-distribution examples. The Customer Review dataset has reviews of products rather than only movies, and the Movie Review dataset has snippets from professional movie reviewers rather than full-length amateur reviews. We leave all test set examples from IMDB as in-distribution examples, and out-of-distribution examples are the 500 or 1000 test reviews from Customer Review and Movie Review datasets, respectively. Table 4 displays detection results, showing a similar story to Table 2.
3.2.2 TEXT CATEGORIZATION
We turn to text categorization tasks to determine whether softmax distributions are useful for detecting similar but out-of-distribution examples. In the following text categorization tasks, we train classifiers to predict the subject of the text they are processing. In the 20 Newsgroups dataset (Lang, 1995), there are 20 different newsgroup subjects with a total of 20000 documents for the whole dataset. The Reuters 8 (Lewis et al., 2004) dataset has eight different news subjects with nearly 8000 stories in total. The Reuters 52 dataset has 52 news subjects with slightly over 9000 news stories; this dataset can have as few as three stories for a single subject.
For the 20 Newsgroups dataset we train a linear classifier on 30-dimensional word vectors for 20 epochs. Meanwhile, Reuters 8 and Retuers 52 use one-layer neural networks with a bag-of-words input and a GELU nonlinearity, all optimized with Adam for 5 epochs. We train on a subset of subjects, leaving out 5 newsgroup subjects from 20 Newsgroups, 2 news subjects from Reuters 8, and 12 news subjects from Reuters 52, leaving the rest as out-of-distribution examples. Table 5 shows that with these datasets and architectures, we can detect errors dependably, and Table 6 informs us that the softmax prediction probabilities allow for detecting out-of-distribution subjects.
3.2.3 PART-OF-SPEECH TAGGING
Part-of-speech (POS) tagging of newswire and social media text is our next challenge. We use the Wall Street Journal portion of the Penn Treebank (Marcus et al., 1993) which contains 45 distinct POS tags. For social media, we use POS-annotated tweets (Gimpel et al., 2011; Owoputi et al., 2013) which contain 25 tags. For the WSJ tagger, we train a bidirectional long short-term memory recurrent neural network (Hochreiter & Schmidhuber, 1997) with three layers, 128 neurons per layer, with randomly initialized word vectors, and this is trained on 90% of the corpus for 10 epochs with stochastic gradient descent with a batch size of 32. The tweet tagger is simpler, as it is twolayer neural network with a GELU nonlinearity, a weight initialization according to (Hendrycks & Gimpel, 2016c), pretrained word vectors trained on a corpus of 56 million tweets (Owoputi et al., 2013), and a hidden layer size of 256, all while training on 1000 tweets for 30 epochs with Adam and early stopping with 327 validation tweets. Error detection results are in Table 7. For out-ofdistribution detection, we use the WSJ tagger on the tweets as well as weblog data from the English Web Treebank (Bies et al., 2012). The results are shown in Table 8. Since the weblog data is closer in style to newswire than are the tweets, it is harder to detect whether a weblog sentence is outof-distribution than a tweet. Indeed, since POS tagging is done at the word-level, we are detecting whether each word is out-of-distribution given the word and contextual features. With this in mind, we see that it is easier to detect words as out-of-distribution if they are from tweets than from blogs.
3.3 AUTOMATIC SPEECH RECOGNITION
Now we consider a task which uses softmax values to construct entire sequences rather than determine an input’s class. Our sequence prediction system uses a bidirectional LSTM with two-layers and a clipped GELU nonlinearity, optimized for 60 epochs with RMSProp trained on 80% of the TIMIT corpus (Garofolo et al., 1993). The LSTM is trained with connectionist temporal classification (CTC) (Graves et al., 2006) for predicting sequences of phones given MFCCs, energy, and first and second deltas of a 25ms frame. When trained with CTC, the LSTM learns to have its phone label probabilities spike momentarily while mostly predicting blank symbols otherwise. In this way, the softmax is used differently from typical classification problems, providing a unique test for our detection methods.
We do not show how the system performs on correctness/incorrectness detection because errors are not binary and instead lie along a range of edit distances. However, we can perform out-of-
distribution detection. Mixing the TIMIT audio with realistic noises from the Aurora-2 dataset (Hirsch & Pearce, 2000), we keep the TIMIT audio volume at 100% and noise volume at 30%, giving a mean SNR of approximately 5. Speakers are still clearly audible to the human ear but confuse the phone recognizer because the prediction edit distance more than doubles. For more outof-distribution examples, we use the test examples from the THCHS-30 dataset (Wang & Zhang, 2015), a Chinese speech corpus. Table 9 shows the results. Crucially, when performing detection, we compute the softmax probabilities while ignoring the blank symbol’s logit. With the blank symbol’s presence, the softmax distributions at most time steps predict a blank symbol with high confidence, but without the blank symbol we can better differentiate between normal and abnormal distributions. With this modification, the softmax prediction probabilities allow us to detect whether an example is out-of-distribution.
4 ABNORMALITY DETECTION WITH AUXILIARY DECODERS
Having seen that softmax prediction probabilities enable abnormality detection, we now show there is other information sometimes more useful for detection. To demonstrate this, we exploit the learned internal representations of neural networks. We start by training a normal classifier and append an auxiliary decoder which reconstructs the input, shown in Figure 1. Auxiliary decoders are sometimes known to increase classification performance (Zhang et al., 2016). The decoder and scorer are trained jointly on in-distribution examples. Thereafter, the blue layers in Figure 1 are frozen. Then we train red layers on clean and noised training examples, and the sigmoid output of the red layers scores how normal the input is. Consequently, noised examples are in the abnormal class, clean examples are of the normal class, and the sigmoid is trained to output to which class an input belongs. After training we consequently have a normal classifier, an auxiliary decoder, and what we call an abnormality module. The gains from the abnormality module demonstrate there are possible research avenues for outperforming the baseline.
4.1 TIMIT
We test the abnormality module by revisiting the TIMIT task with a different architecture and show how these auxiliary components can greatly improve detection. The system is a three-layer, 1024- neuron wide classifier with an auxiliary decoder and abnormality module. This network takes as input 11 frames and must predict the phone of the center frame, 26 features per frame. Weights are initialized according to (Hendrycks & Gimpel, 2016c). This network trains for 20 epochs, and the abnormality module trains for two. The abnormality module sees clean examples and, as negative examples, TIMIT examples distorted with either white noise, brown noise (noise with its spectral density proportional to 1/f2), or pink noise (noise with its spectral density proportional to 1/f ) at various volumes.
We note that the abnormality module is not trained on the same type of noise added to the test examples. Nonetheless, Table 10 shows that simple noised examples translate to effective detection of realistically distorted audio. We detect abnormal examples by comparing the typical abnormality
module outputs for clean examples with the outputs for the distorted examples. The noises are from Aurora-2 and are added to TIMIT examples with 30% volume. We also use the THCHS-30 dataset for Chinese speech. Unlike before, we use the THCHS-30 training examples rather than test set examples because fully connected networks can evaluate the whole training set sufficiently quickly. It is worth mentioning that fully connected deep neural networks are noise robust (Seltzer et al., 2013), yet the abnormality module can still detect whether an example is out-of-distribution. To see why this is remarkable, note that the network’s frame classification error is 29.69% on the entire test (not core) dataset, and the average classification error for distorted examples is 30.43%—this is unlike the bidirectional LSTM which had a more pronounced performance decline. Because the classification degradation was only slight, the softmax statistics alone did not provide useful outof-distribution detection. In contrast, the abnormality module provided scores which allowed the detection of different-but-similar examples. In practice, it may be important to determine whether an example is out-of-distribution even if it does not greatly confuse the network, and the abnormality module facilitates this.
4.2 MNIST
Finally, much like in a previous experiment, we train an MNIST classifier with three layers of width 256. This time, we also use an auxiliary decoder and abnormality module rather than relying on only softmax statistics. For abnormal examples we blur, rotate, or add Gaussian noise to training images. Gains from the abnormality module are shown in Table 11, and there is a consistent out-of-sample detection improvement compared to softmax prediction probabilities. Even for highly dissimilar examples the abnormality module can further improve detection.
5 DISCUSSION AND FUTURE WORK
The abnormality module demonstrates that in some cases the baseline can be beaten by exploiting the representations of a network, suggesting myriad research directions. Some promising future avenues may utilize the intra-class variance: if the distance from an example to another of the same predicted class is abnormally high, it may be out-of-distribution (Giryes et al., 2015). Another path is to feed in a vector summarizing a layer’s activations into an RNN, one vector for each layer. The RNN may determine that the activation patterns are abnormal for out-of-distribution examples. Others could make the detections fine-grained: is the out-of-distribution example a known-unknown or an unknown-unknown? A different avenue is not just to detect correct classifications but to output the probability of a correct detection. In Appendix B, we show a baseline and evaluation metrics that future research can utilize for estimating the probability of a correct classification. These are but a few ideas for improving error and out-of-distribution detection.
We hope that any new detection methods are tested on a variety of tasks and architectures of the researcher’s choice. A basic demonstration could include the following datasets: MNIST, CIFAR, IMDB, and tweets because vision-only demonstrations may not transfer well to other architectures and datasets. Reporting the AUPR and AUROC values is important, and so is the underlying classifier’s accuracy since an always-wrong classifier gets a maximum AUPR for error detection if error is the positive class. Also, future research need not use the exact values from this paper for comparisons. Machine learning systems evolve, so tethering the evaluations to the exact architectures and datasets in this paper is needless. Instead, one could simply choose a variety of datasets and architectures possibly like those above and compare their detection method with a detector based on the softmax prediction probabilities from their classifiers. These are our basic recommendations for others who try to surpass the baseline on this underexplored challenge.
6 CONCLUSION
We demonstrated a softmax prediction probability baseline for error and out-of-distribution detection across several architectures and numerous datasets. We then presented the abnormality module, which provided superior scores for discriminating between normal and abnormal examples on tested cases. The abnormality module demonstrates that the baseline can be beaten in some cases, and this implies there is room for future research. Our hope is that other researchers investigate architectures which make predictions in view of abnormality estimates, and that others pursue more reliable methods for detecting errors and out-of-distribution inputs because knowing when a machine learning system fails strikes us as highly important.
ACKNOWLEDGMENTS
We would like to thank John Wieting, Hao Tang, Karen Livescu, Greg Shakhnarovich, and our reviewers for their suggestions. We would also like to thank NVIDIA Corporation for donating several TITAN X GPUs used in this research.
A ABNORMALITY MODULE EXAMPLE
B METRICS FOR MEANINGFUL CONFIDENCE VALUES
Throughout this work, we show that the softmax probability of a predicted class is consistently high and thus is a poor proxy for estimating the “confidence” in a classification. In addition to being frequently high, the softmax prediction probability is bounded below by a value greater than zero. To see this, note that a binary sentiment classifier’s minimum confidence in an out-of-distribution example is 50% since the highest softmax probability must be at least 50%. Yet for an out-ofdistribution example we may want less confidence if confidence is to reflect the probability of correct classification.
To evaluate how well a confidence model reflects the probability of correct classification, we present two metrics. First, let us establish notation. Denote the confidence model by c : X → [0, 1], where X is a set of inputs. An output near 1 from c reflects high confidence, or a high probability that the underlying classification model f : X → Y correctly classified its input. Also, let 1(b) be 1 if b is True, and 0 if b is False. A possible confidence model evaluation metric is accuracy. To compute this, we must threshold the confidence at 0.5. However, the model can output probabilities slightly above or below 0.5 and obtain perfect accuracy. Probabilities all near 0.5 are insufficiently informative and lack numerical resolution, and for this reason we avoid thresholding in our metrics. Instead, we use the confidence itself in a metric which we call the probability alignment score,
E(x,y)∼(X ,Y)[21(f(x) = y)− 1][2c(x)− 1].
Note that if f is a highly accurate model, then the confidence model can obtain a high probability alignment by predicting 1 constantly. We meet this challenge by also calculating a “soft” F1 score (Pastor-Pellicer et al., 1998) for a test set {(xi, yi)}mi=1. We define the soft F1 score by first letting tp = ∑m i=1(1 − c(xi))1(f(xi) 6= yi), fp = ∑m i=1(1 − c(xi))1(f(xi) = yi),
fn = ∑m
i=1 c(xi)1(f(xi) 6= yi), pr = tp/(tp + fp), and rc = tp/(tp + fn). Then the soft F1 score is
2 pr · rc
pr + rc ,
and this is our second and final metric for evaluating a confidence model.
Using these metrics, we now establish a baseline using the logits of a softmax. Let li represent the logits vector (the input to the softmax) for example xi. Our baseline confidence model is c(xi) = σ(max li), where σ is the logistic sigmoid. Let us evaluate this baseline by using the same MNIST, Twitter, and IMDB setups from before. We take the median result over five runs. In the case of CIFAR-10, we train using a 40-2 WideResNet. The results are in Table 12, and they indicate that the probabilities are positively aligned and have a nonzero soft F1 score. However, these results also indicate wide room for improvement.
In this appendix, we presented a simple baseline method for estimating classification confidence and metrics for evaluating the confidence model. The baseline draws on the idea of using softmax classification information to create a confidence model. We hope this appendix gives ideas for future research directions, serves as a baseline, or provides metrics for confidence model evaluation.

The authors present results on a number of different tasks where the goal is to determine whether a given test example is out-of-domain or likely to be mis-classified. This is accomplished by examining statistics for the softmax probability for the most likely class; although the score by itself is not a particularly good measure of confidence, the statistics for out-of-domain examples are different enough from in-domain examples to allow these to be identified with some certainty.
My comments appear below:
1. As the authors point out, the AUROC/AUPR criterion is threshold independent. As a result, it is not obvious whether the thresholds that would correspond to a certain operating point (say a true positive rate of 10%) would be similar across different data sets. In other words, it would be interesting to know how sensitive the thresholds are to different test sets (or different splits of the test set). This is important if we want to use the thresholds determined on a given held-out set during evaluation on unseen data (where we would need to select a threshold).
2. Performance is reported in terms of AUROC/AUPR and models are compared against a random baseline. I think it’s a little hard to look at the differences in AUC/AUPR to get a sense for how much better the proposed classifier is than the random baseline. It would be useful, for example, if the authors could also report how strongly statistically significant some of these differences are (although admittedly they look to be pretty large in most cases).
3. In the experiments on speech recognition presented in Section 3.3, I was not entirely clear on how the model was evaluated. In Table 9, for example, is an “example” the entire utterance or just a single (stacked?) speech frame. Assuming that each “example” is an utterance, are the softmax probabilities the probability of the entire phone sequence (obtained by multiplying the local probability estimates from a Viterbi decoding?)
4. I’m curious about the decision to ignore the blank symbol’s logit in Section 3.3. Why is this required?
5. As I mentioned in the pre-review question, at least in the speech recognition case, it would have been interesting to compare performance obtained using a simple generative baseline (e.g., GMM-HMM). I think that would serve as a good indication of the ability of the proposed model to detect out-of-domain examples over the baseline.


SKIP-GRAPH: LEARNING GRAPH EMBEDDINGS WITH AN ENCODER-DECODER MODEL
In this work, we study the problem of feature representation learning for graphstructured data. Many of the existing work in the area are task-specific and based on supervised techniques. We study a method for obtaining a generic feature representation for a graph using an unsupervised approach. The neural encoderdecoder model is a method that has been used in the natural language processing domain to learn feature representations of sentences. In our proposed approach, we train the encoder-decoder model to predict the random walk sequence of neighboring regions in a graph given a random walk along a particular region. The goal is to map subgraphs — as represented by their random walks — that are structurally and functionally similar to nearby locations in feature space. We evaluate the learned graph vectors using several real-world datasets on the graph classification task. The proposed model is able to achieve good results against state-ofthe-art techniques.
1 INTRODUCTION
The skip-gram model (Mikolov et al., 2013) was originally introduced in the natural language processing (NLP) domain as a model for learning vector representations of words. Recently, it has been adapted successfully to solve the problem of learning node representations for graph-structured data (Grover & Leskovec, 2016; Perozzi et al., 2014). The learned vectors can then be used directly in problems such as link prediction (Miller et al., 2009), or clustering of nodes on a graph (Vinayak et al., 2014). However, in many real-world applications we need to learn a feature representation for the entire graph instead of representations for just the nodes in the graph. In this paper, we study the graph representation learning problem, where the task is to learn a feature representation for any graph object. We propose a novel solution based upon the encoder-decoder model.
Graph-structured data can be found in many different domains including biology, chemistry, and the study of social networks. For instance, in chemistry, chemical compounds can be represented as molecular graphs (Duvenaud et al., 2015). In social network analysis, the interaction among different entities of a community can be captured using a social graph (Yanardag & Vishwanathan, 2015). A natural question that arises in these scenarios is what the structure of a graph tells us about the properties of the graph (e.g., what does the molecular graph tell us about the compound’s aqueous solubility, or its anti-cancer activity?). In other words, we are often interested in performing machine learning tasks on graph-structured data. Many techniques have been proposed to solve this problem, these include learning graph kernels (Vishwanathan et al., 2010), identifying discriminative subgraphs (Kong et al., 2011), using specially designed neural network models such as the graph neural network (Scarselli et al., 2009), and learning the graph fingerprint (Duvenaud et al., 2015). Most of the approaches for learning graph features are supervised and task-specific. Our approach, on the other hand, is unsupervised and general-purpose. The learned features can be used directly with off-the-shelf machine learning methods on different tasks, such as classification or clustering. Perhaps the work that resembles this work the most is the one in (Yanardag & Vishwanathan, 2015). We argue, however, that our approach is different and this is good motivation to pursue the study as there has not been many work published in the area. For one, we use the skip-thought model (Kiros
et al., 2015) and we are not just interested in structurally similar subgraphs but also functionally similar ones.
Our approach is based on the encoder-decoder model (Kalchbrenner & Blunsom, 2013; Cho et al., 2014); in particular, we are interested in the skip-thought model. In (Kiros et al., 2015), tuples composed of three consecutive sentences from word documents are fed into an RNN model and the model attempts to reconstruct the previous and next statements given the middle sentence. After training on a large text corpus, the hidden vector values for an input sentence can be used as that input sequence’s feature representation. It has been shown that the model learns a function that maps semantically and syntactically similar sentences close to one another in feature space. In this work, the idea is to take instead a sequence generated by a random walk along a labeled graph and to divide it into three parts, feeding these into the encoder-decoder model. Since the structure of the graph determines the random walk sequences that can be generated, we can treat each sub-sequence as a representation of a particular subgraph in the graph. We argue that by training an encoderdecoder model on a large number of random walk sequences, we can learn a feature representation that groups structurally and functionally similar subgraphs together. Figure 1 shows an example of how we can train the model using a random walk over a graph. A simple example that illustrates how the model may learn to identify functionally similar subgraphs is shown in Figure 2.
After the model is trained on a large sample of random walks generated from a dataset of labeled graphs, we can then freeze the model and use the encoder as a feature extractor. In particular, we obtain a feature representation of a graph by sampling multiple short random walks and aggregating the information encoded in the feature representations of these short walks. We borrow an analogy from the NLP domain to highlight the idea. In order to obtain a good feature representation for a text document, short of sampling all the words in the document one may sample a set of sentences from the document and use these to construct the features for the document. Similarly, to obtain a feature representation for a graph, we sample a set of subgraphs (as represented by the short walks) and use the aggregate subgraph features to construct the final graph feature vector. Since we use the trained encoder as our feature extractor, graphs that share structural and functional properties will tend to have more similar feature vectors.
2 PROPOSED METHOD
2.1 SKIP-THOUGHT
Since our proposed approach is based on the encoder-decoder model of (Kiros et al., 2015), we begin by briefly introducing the model. The encoder-decoder model uses an RNN with GRU (Chung et al., 2014) activation as the encoder and an RNN with a conditional GRU as the decoder. The model is trained using the Adam stochastic optimization algorithm (Kingma & Ba, 2015).
The input to the model is a tuple of sentences (si−1, si, si+1), with xti being the word embedding for the t-th word, wti , of sentence si. The word embeddings for the middle sentence, si, are fed sequentially as input to the the encoder. The encoder generates a hidden vector hti at each time step t, this is the information the model retained after processing sequence x1i , ... , x t i and can be thought of as the sequence representation. The hidden state hNi can thus be considered the sentence representation, given si is of length N . Given a sequence to encode, the encoder iterates through the following equations, as given in (Kiros et al., 2015). Here the subscripts i are dropped for simplicity.
rt = σ(Wrx t + Urh t−1) (1)
zt = σ(Wzx t + Uzh t−1) (2)
h̄t = tanh(Wxt + U(rt ht−1)) (3) ht = (1− zt) ht−1 + zt h̄t (4)
where rt is the forget gate, zt is the update gate, h̄t is the proposed hidden state, and is the component-wise product. Here rt decides what information to discard from the previous state, zt decides what new information to encode, and the new hidden vector ht is calculated accordingly. Values in rt and zt are in the range [0, 1].
Two decoders with separate parameters are used to reconstruct the previous statement si−1 and the next statement si+1. The computation for the decoder is similar to that of the encoder, except this time the models are also conditioned on the encoder output hi. Decoding involves iterating through the following statements. Again the subscript i+ 1 (similarly, i− 1) is dropped.
rt = σ(Wdrx t−1 + Udrh t−1 + Crhi) (5) zt = σ(Wdzx t−1 + Udzh t−1 + Czhi) (6)
h̄t = tanh(Wdxt−1 + Ud(rt ht−1) + Chi) (7) hti+1 = (1− zt) ht−1 + zt h̄t (8)
here the C matrices are used to bias the computation by the sentence vector produced by the encoder. Also, note that the word embeddings are from the previous and next statements since these are what is given to the decoders. The probability of word wti+1 can be calculated by
P (wti+1|w<ti+1,hi) ∝ exp(vwti+1h t i+1) (9)
where vwti+1 is the row vector in the vocabulary vector V corresponding to the word w t i+1. The vocabulary matrix, V, is a weight matrix shared by both decoders connecting the decoder’s hidden states for computing a distribution over words. Finally, given a sentence tuple, the training objective is given by∑ t logP (wti+1|w<ti+1,hi) + ∑ t logP (wti−1|w<ti−1,hi) (10)
which is the sum of log-probabilities for the words in the previous and next statements, si−1 and si+1, conditioned on the sentence representation for si. The total objective would then be the above summed for all tuples in the training data.
2.2 SKIP-GRAPH
In this work, we are interested in graph-structured data in particular. In our setting, we are given a set of labeled graphs D = {G1,G2, ... ,Gn} with each graph associated with a class label. A graph G = (V, E , `v) is comprised of a vertex set V , an edge set E ⊆ V × V , and a node labeling function `v : V → LV which assigns each node to a label in LV . Additionally, the edges may also be labeled in which case we also have an edge labeling function `e : E → LE . Nodes and edges can also have associated feature vectors, these are fv ∈ RDv , and fe ∈ RDe , respectively.
2.2.1 UNLABALED GRAPHS
Although we will be working primarily with labeled graphs, our method can be easily extended to support unlabeled graphs by including an additional pre-processing step. Algorithms like the Weisfeiler-Lehman algorithm (Weisfeiler & Lehman, 1968; Shervashidze et al., 2011) or the Morgan algorithm (Rogers & Hahn, 2010) for calculating molecular fingerprints are iterative algorithms that work by repeatedly calculating the attribute for a node via hashing of the attributes of its neighboring nodes. The final node attributes capture the local structure or topology of the graph. For unlabeled graphs, all node attributes can be initialized to a constant value and after the algorithm is run, we can treat the node attributes as the labels for the nodes in the graph.
2.2.2 TRAINING SET GENERATION
Given a set of graphs D, a sample size K, a minimum random walk length lmin, and a maximum random walk length lmax, we take each graph G ∈ D and generate K random walk sequences. Specifically, for a graph G, K sequences of the form
`v(v1), ..., `v(vk), `v(vk+1), ..., `v(vk+k′), `v(vk+k′+1), ..., `v(vk+k′+k′′)
are generated. Here, v1 ∈ V is a randomly selected start node, (vi, vi+1) ∈ E for i from 1 ... k + k′+k′′−1, and lmin ≥ k, k′, k′′ ≥ lmax. We can split each sequence into three sub-sequences with s1 = `v(v1), ..., `v(vk), s2 = `v(vk+1), ..., `v(vk+k′), and s3 = `v(vk+k′+1), ..., `v(vk+k′+k′′). For each sequence, k, k′, and k′′ are randomly drawn to be between the constraints. Since the length of the sub-sequences do not need to have fixed lengths and can instead be between lmin and lmax, regions of varying sizes can easily be considered.
In the above formulation, we assume that only the vertices in the graph are labeled and node and edge features are not given. When nodes, or edges, are labeled and feature vectors are provided we can use a one-hot embedding to represent each unique combination of labels and features. This treats each distinct combination as a unique “word” and does not capture the relationship between nodes or edges that share labels or certain features. A better approach is to simply use a one-of-|L| vector to encode the label and concatenate this with the feature vector, this allows the node or edge embedding to capture shared features and labels.
Once all the tuples of random walk sequences have been generated, they can be used to train the encoder-decoder1 in an unsupervised fashion.
2.2.3 OBTAINING FINAL GRAPH REPRESENTATION
After the encoder-decoder has been trained, we can freeze the model and use the encoder to generate representations, hi, for any arbitrary random walk sequence. Ultimately, however, we are interested in obtaining a representation for entire graphs so we try several strategies for aggregating the encoder representations obtained from a set of independent random walks sampled from a given graph.
1. Single walk: In this approach we do not use several encoder representations. Instead, we train the model on relatively long (relative to the size of the graphs in the dataset) random walk sequences and use a single long walk over the graph to obtain its representation.
2. Average: We compute the component-wise average of the encoder representations of the sampled random walk sequences. This is then used as the graph representation.
1We use the implementation in https://github.com/ryankiros/skip-thoughts.
3. Max: As in (Kiela & Bottou, 2014), we take the component-wise absolute maximum of all encoder representations.
4. Cluster: The encoder representations are first fed into a clustering technique like Kmeans (Hamerly & Elkan, 2003) and we use the cluster information to create a bag-ofcluster vector that serves as the graph’s representation.
The procedure for obtaining the graph embeddings is summarized in Algorithm 1. The calculated graph embeddings can now be used with any off-the-shelf machine learning method.
Algorithm 1: Calculate graph embedding Input : Training set D, sample size K, walk lengths lmin and lmax, aggregate sample size K ′, and aggregate method agg Output: Graph embeddings
1 Generate set of K × |D| random walk tuples, S; 2 Train encoder-decoder model using S; 3 for each G in D do 4 Randomly select K ′ random walks; 5 Obtain encoder representations h1, ...,hK′ from the random walks; 6 Compute graph embedding with agg(h1, ...,hK′); 7 end 8 Return final graph embeddings;
3 EXPERIMENTS
3.1 DATASET
We evaluate our proposed method on the binary classification task using four chemical compound datasets (Kong et al., 2011). The datasets contain chemical compounds encoded in the SMILES format (Weininger, 1988); class labels indicate the anti-cancer properties (active or inactive) of each compound. We use the RDKit2 package to obtain the molecular graphs from the SMILES data. We also use RDKit to obtain the labels for the nodes (atom type) and edges (bond type). Additionally, we used the number of attached hydrogens as a node feature and bond conjugation as an edge feature. Since the edges in the datasets we evaluate on are also labeled, the generated random walk sequences include edges. The datasets are all highly skewed with far more negative samples than positive ones, we tested the methods on balanced datasets by selecting a random set of negative samples equal to the positive ones. Table 1 shows a summary of the datasets used. The average size of the molecular graphs in each of the four datasets is around 30.
3.2 COMPARED METHODS
We compared our proposed approach with several state-of-the-art techniques. Since the method is a task-irrelevant way to obtain graph representations, the goal of the paper isn’t necessarily to come up with a method that achieves absolute best performance on the tested datasets so we do not test against an exhaustive list of methods. Our primary objective is to see whether the method can
2http://www.rdkit.org/
potentially be used to learn useful graph embeddings as a starting point for future investigation in the area. Since we are testing the method using molecular graph datasets, we chose to compare against techniques that have achieved state-of-the-art performance on these type of graphs. We also compare against a method that learns node embeddings instead of an entire graph embedding. The tested methods are:
• ECFP (Rogers & Hahn, 2010): Extended-connectivity circular fingerprints, which are a refinement of the Morgan algorithm (Morgan, 1965), use an iterative approach to encode information about substructures in a molecular graph in a fingerprint vector. In this method a hash function is used to map the concatenated features from a neighborhood to an index in the fingerprint vector.
• NeuralFPS (Duvenaud et al., 2015): Neural fingerprints replace the function that is used to compute a fingerprint vector with a differentiable neural network. This allows the method to learn from the data, prioritizing useful or discriminative features.
• DeepWalk (Perozzi et al., 2014): The DeepWalk model learns representations for nodes in a single graph. However, we can also train the model using random walks from multiple graphs if the various graphs share the same kind of nodes. The model will then learn to generate similar representations for nodes that co-occur frequently across all the graphs. To generate the final embedding for a graph, we can simply apply average pooling to the vectors of all the nodes in the graph – which is a reasonable strategy to capture the overall profile of the graph.
• Skip-graph: Our proposed method. We train an encoder-decoder model using random walks generated from the graphs and use the encoder’s random walk representation to calculate the graph embedding.
To test ECFP and NeuralFPS, we used the library3 provided by (Duvenaud et al., 2015). The size of the graph embedding was restricted to 164 for all methods and a grid-search was done to optimize the parameters of the various methods. For ECFP and NeuralFPS, we tested different values for the following parameters: fingerprint radius, `2 regularization penalty, step size for the optimization, hidden layer dimension, and convolution layer dimension (only for NeuralFPS). All results reported are the average over 5-fold cross validation. Since a neural network, with a single hidden layer, was used as the classifier in Duvenaud et al. (2015), we chose to use the same classifier for our model and the grid-search was performed over the same set of values for classifier-related parameters. In particular, for the neural network, we tested various settings with hidden layer size selected from {70, 100, 140}, and `2 regularization chosen from {0.0001, 0.001, 0.01, 0.1}.
3.3 CLASSIFICATION RESULTS
We show the classification accuracy of the different methods in Table 2. The proposed method achieves top performance in three of the four datasets we tested. It is a little surprising, however, to find that NeuralFPS performs slightly worse than ECFP. This seems to suggest that it is overfitting the data as NeuralFPS is a generalization of ECFP and should, in theory, be at least as good as ECFP. Also, we find that averaging the DeepWalk embeddings trained from random walks generated from the entire training set can be a simple yet effective way to generate a graph representation.
3.4 PARAMETER STUDY
We tested the performance of the method using the various aggregation methods. The performance was extremely poor when we trained the encoder-decoder model on long random walks and used a single long walk to generate the graph representation. The other three aggregation strategies yielded better results. Figure 3(a) shows the performance of these methods. Averaging the hidden vector representations seems to yield the best performance, calculating the component-wise maximum yielded the second best results while the method that had the additional cluster pre-processing step performed slightly worse.
We plot the accuracy of the method over the number of training epochs in Figure 3(b). With the exception of the HIV dataset, which has a relatively few number of samples, the results show a gradual increase in the classification accuracy as the number of training epochs is increased. This is consistent with results in other work that show that given a large number of training data, recurrent neural models generally achieve better results when trained longer.
Figure 3(c) shows the accuracy in the classification task over different sample sizes K ′, or the number of samples aggregated to obtain the final graph representation. It is clear from the results that a better graph representation is obtained if we use more samples to calculate the final graph representation. This is quite intuitive as a limited sample may not be representative and may fail to capture the properties of the graph well enough.
We tested several different values for lmin and lmax and the one that seemed to perform best in our case was lmin = 7 and lmax = 12. This is a reasonable constraint on the random walk length given that the average size of the molecular graphs was around 30. We used K = 100 when generating a set of random walks to train the encoder-decoder.
3.5 VISUALIZATION OF GRAPH EMBEDDINGS
We show a scatterplot of the HIV graph embeddings learned by our model in Figure 4. In particular, we highlight two pairs of graphs that had very similar embeddings. We note that the first pair of graphs (the one on the right) are structurally similar, that is they have a large sub-structure in common. The graphs in the second pair each contain two similar substructures that are joined by segments that appear to be “functionally” similar.
3.6 USING AN ENSEMBLE OF CLASSIFIERS
Since it is possible to generate many different sets of random walks to train the encoder-decoder model, we tried training five encoders on five separate sets of random walks. An ensemble (Opitz & Maclin, 1999) of five classifiers is then created with each classifier trained on the graph representations obtained from one of the five encoders. We compare the predictive accuracy of the ensemble versus the single classifier when all other settings are fixed. We observed a slight improvement (around 1 − 3%) in the accuracy of the model. All the results reported above are for the single classifier case.
4 CONCLUSION
We introduced an unsupervised method, based on the encoder-decoder model, for generating feature representations for graph-structured data. The model was evaluated on the binary classification task on several real-world datasets. The method outperformed several state-of-the-art algorithms on the tested datasets.
There are several interesting directions for future work. For instance, we can try training multiple encoders on random walks generated using very different neighborhood selection strategies. This may allow the different encoders to capture different properties in the graphs. We would also like to test the approach using different neural network architectures. Finally, it would be interesting to test the method on other types of heterogeneous information networks.

The paper presents a method to learn graph embeddings in a unsupervised way using random walks. It is well written and the execution appears quite accurate. The area of learning whole graph representations does not seem to be very well explored in general, and the proposed approach enjoys having very few competitors.
In a nutshell, the idea is to linearize the graph using random walks and to compute the embedding of the central segment of each walk using the skip-thought criterion. Being not an expert in biology, I can not comment whether or not this makes sense, but the gains reported in Table 2 are quite significant.
An anonymous public comment compared this work to a number of others in which the problem of learning representations of nodes is considered. While this is arguably a different goal, one natural baseline would be to pool these representations using mean- or max- pooling. It would very interesting to do such a comparison, especially given that the considered approach heavily relies on pooling (see Figure 3(c))
To sum up, I think it is a nice paper, and with more baselines I would be ready to further increase the numerical score.


Code super-optimization is the task of transforming any given program to a more efficient version while preserving its input-output behaviour. In some sense, it is similar to the paraphrase problem from natural language processing where the intention is to change the syntax of an utterance without changing its semantics. Code-optimization has been the subject of years of research that has resulted in the development of rule-based transformation strategies that are used by compilers. More recently, however, a class of stochastic search based methods have been shown to outperform these strategies. This approach involves repeated sampling of modifications to the program from a proposal distribution, which are accepted or rejected based on whether they preserve correctness and the improvement they achieve. These methods, however, neither learn from past behaviour nor do they try to leverage the semantics of the program under consideration. Motivated by this observation, we present a novel learning based approach for code super-optimization. Intuitively, our method works by learning the proposal distribution using unbiased estimators of the gradient of the expected improvement. Experiments on benchmarks comprising of automatically generated as well as existing (“Hacker’s Delight”) programs show that the proposed method is able to significantly outperform state of the art approaches for code super-optimization.
1 INTRODUCTION
Considering the importance of computing to human society, it is not surprising that a very large body of research has gone into the study of the syntax and semantics of programs and programming languages. Code super-optimization is an extremely important problem in this context. Given a program or a snippet of source-code, super-optimization is the task of transforming it to a version that has the same input-output behaviour but can be executed on a target compute architecture more efficiently. Superoptimization provides a natural benchmark for evaluating representations of programs. As a task, it requires the decoupling of the semantics of the program from its superfluous properties, the exact implementation. In some sense, it is the natural analogue of the paraphrase problem in natural language processing where we want to change syntax without changing semantics.
Decades of research has been done on the problem of code optimization resulting in the development of sophisticated rule-based transformation strategies that are used in compilers to allow them to perform code optimization. While modern compilers implement a large set of rewrite rules and are able to achieve impressive speed-ups, they fail to offer any guarantee of optimality, thus leaving room for further improvement. An alternative approach is to search over the space of all possible programs that are equivalent to the compiler output, and select the one that is the most efficient. If the search is carried out in a brute-force manner, we are guaranteed to achieve super-optimization. However, this approach quickly becomes computationally infeasible as the number of instructions and the length of the program grows.
In order to efficiently perform super-optimization, recent approaches have started to use a stochastic search procedure, inspired by Markov Chain Monte Carlo (MCMC) sampling (Schkufza et al., 2013). Briefly, the search starts at an initial program, such as the compiler output. It iteratively suggests modifications to the program, where the probability of a modification is encoded in a proposal distribution. The modification is either accepted or rejected with a probability that is dependent on the improvement achieved. Under certain conditions on the proposal distribution, the above procedure can be shown, in the limit, to sample from a distribution over programs, where the probability of a program is related to its quality. In other words, the more efficient a program, the more times it is encountered, thereby enabling super-optimization. Using this approach, high-quality implementations of real programs such as the Montgomery multiplication kernel from the OpenSSL library were discovered. These implementations outperformed the output of the gcc compiler and even expert-handwritten assembly code.
One of the main factors that governs the efficiency of the above stochastic search is the choice of the proposal distribution. Surprisingly, the state of the art method, Stoke (Schkufza et al., 2013), employs a proposal distribution that is neither learnt from past behaviour nor does it depend on the syntax or semantics of the program under consideration. We argue that this choice fails to fully exploit the power of stochastic search. For example, consider the case where we are interested in performing bitwise operations, as indicated by the compiler output. In this case, it is more likely that the optimal program will contain bitshifts than floating point opcodes. Yet, Stoke will assign an equal probability of use to both types of opcodes.
In order to alleviate the aforementioned deficiency of Stoke, we build a reinforcement learning framework to estimate the proposal distribution for optimizing the source code under consideration. The score of the distribution is measured as the expected quality of the program obtained via stochastic search. Using training data, which consists of a set of input programs, the parameters are learnt via the REINFORCE algorithm (Williams, 1992). We demonstrate the efficacy of our approach on two datasets. The first is composed of programs from “Hacker’s Delight” (Warren, 2002). Due to the limited diversity of the training samples, we show that it is possible to learn a prior distribution (unconditioned on the input program) that outperforms the state of the art. The second dataset contains automatically generated programs that introduce diversity in the training samples. We show that, in this more challenging setting, we can learn a conditional distribution given the initial program that significantly outperforms Stoke.
2 RELATED WORKS
Super-optimization The earliest approaches for super-optimization relied on brute-force search. By sequentially enumerating all programs in increasing length orders (Granlund & Kenner, 1992; Massalin, 1987), the shortest program meeting the specification is guaranteed to be found. As expected, this approach scales poorly to longer programs or to large instruction sets. The longest reported synthesized program was 12 instructions long, on a restricted instruction set (Massalin, 1987).
Trading off completeness for efficiency, stochastic methods (Schkufza et al., 2013) reduced the number of programs to test by guiding the exploration of the space, using the observed quality of programs encountered as hints. In order to improve the size of solvable instances, Phothilimthana et al. (2016) combined stochastic optimizers with smart enumerative solvers. However, the reliance of stochastic methods on a generic unspecific exploratory policy made the optimization blind to the problem at hand. We propose to tackle this problem by learning the proposal distribution.
Neural Computing Similar work was done in the restricted case of finding efficient implementation of computation of value of degree k polynomials (Zaremba et al., 2014). Programs were generated from a grammar, using a learnt policy to prioritise exploration. This particular approach of guided search looks promising to us, and is in spirit similar to our proposal, although applied on a very restricted case.
Another approach to guide the exploration of the space of programs was to make use of the gradients of differentiable relaxation of programs. Bunel et al. (2016) attempted this by simulating program execution using Recurrent Neural Networks. However, this provided no guarantee that the network parameters were going to correspond to real programs. Additionally, this method only had the
possibility of performing local, greedy moves, limiting the scope of possible transformations. On the contrary, our proposed approach operates directly on actual programs and is capable of accepting short-term detrimental moves.
Learning to optimize Outside of program optimization, applying learning algorithms to improve optimization procedures, either in terms of results achieved or runtime, is a well studied subject. Doppa et al. (2014) proposed imitation learning based methods to deal with structured output spaces, in a “Learning to search” framework. While this is similar in spirit to stochastic search, our setting differs in the crucial aspect of having a valid cost function instead of searching for one.
More relevant is the recent literature on learning to optimize. Li & Malik (2016) and Andrychowicz et al. (2016) learn how to improve on first-order gradient descent algorithms, making use of neural networks. Our work is similar, as we aim to improve the optimization process. However, as opposed to the gradient descent that they learn on a continuous unconstrained space, our initial algorithm is an MCMC sampler on a discrete domain.
Similarly, training a proposal distribution parameterized by a Neural Network was also proposed by Paige & Wood (2016) to accelerate inference in graphical models. Similar approaches were successfully employed in computer vision problems where data driven proposals allowed to make inference feasible (Jampani et al., 2015; Kulkarni et al., 2015; Zhu et al., 2000). Other approaches to speeding up MCMC inference include the work of Salimans et al. (2015), combining it with Variational inference.
3 LEARNING STOCHASTIC SUPER-OPTIMIZATION
3.1 STOCHASTIC SEARCH AS A PROGRAM OPTIMIZATION PROCEDURE
Stoke (Schkufza et al., 2013) performs black-box optimization of a cost function on the space of programs, represented as a series of instructions. Each instruction is composed of an opcode, specifying what to execute, and some operands, specifying the corresponding registers. Each given input program T defines a cost function. For a candidate programR called rewrite, the goal is to optimize the following cost function:
cost (R, T ) = ωe × eq(R, T ) + ωp × perf(R) (1) The term eq(R; T ) measures how well the outputs of the rewrite match the outputs of the reference program. This can be obtained either exactly by running a symbolic validator or approximately by running test cases. The term perf(R) is a measure of the efficiency of the program. In this paper, we consider runtime to be the measure of this efficiency. It can be approximated by the sum of the latency of all the instructions in the program. Alternatively, runtime of the program on some test cases can be used.
To find the optimum of this cost function, Stoke runs an MCMC sampler using the Metropolis (Metropolis et al., 1953) algorithm. This allows us to sample from the probability distribution induced by the cost function:
p(R; T ) = 1 Z exp(−cost (R, T ))). (2)
The sampling is done by proposing random moves from a different proposal distribution:
R′ ∼ q( · |R). (3) The cost of the new modified program is evaluated and an acceptance criterion is computed. This acceptance criterion
α(R, T ) = min ( 1, p(R′; T ) p(R; T ) ) , (4)
is then used as the parameter of a Bernoulli distribution from which an accept/reject decision is sampled. If the move is accepted, the state of the optimizer is updated to R′. Otherwise, it remains inR. While the above procedure is only guaranteed to sample from the distribution p( · ; T ) in the limit if the proposal distribution q is symmetric (q(R′|R) = q(R|R′) for all R,R′), it still allows us
to perform efficient hill-climbing for non-symmetric proposal distributions. Moves leading to an improvement are always going to be accepted, while detrimental moves can still be accepted in order to avoid getting stuck in local minima.
3.2 LEARNING TO SEARCH
We now describe our approach to improve stochastic search by learning the proposal distribution. We begin our description by defining the learning objective (section 3.2.1), followed by a parameterization of the proposal distribution (section 3.2.2), and finally the reinforcement learning framework to estimate the parameters of the proposal distribution (section 3.2.3).
3.2.1 OBJECTIVE FUNCTION
Our goal is to optimize the cost function defined in equation (1). Given a fixed computational budget of T iterations to perform program super-optimization, we want to make moves that lead us to the lowest possible cost. As different programs have different runtimes and therefore different associated costs, we need to perform normalization. As normalized loss function, we use the ratio between the best rewrite found and the cost of the initial unoptimized program R0. Formally, the loss for a set of rewrites {Rt}t=0..T is defined as follows:
r({Rt}t=0..T ) = ( mint=0..T cost (Rt, T )
cost (R0, T )
) . (5)
Recall that our goal is to learn a proposal distribution. Given that our optimization procedure is stochastic, we will need to consider the expected cost as our loss. This expected loss is a function of the parameters θ of our parametric proposal distribution qθ:
L(θ) = E{Rt}∼qθ [r({Rt}t=0..T )] . (6)
3.2.2 PARAMETERIZATION OF THE MOVE PROPOSAL DISTRIBUTION
The proposal distribution (3) originally used in Stoke (Schkufza et al., 2013) takes the form of a hierarchical model. The type of the move is initially sampled from a probability distribution. Additional samples are drawn to specify, for example, the affected location in the programs ,the new operands or opcode to use. Which of these probability distribution get sampled depends on the type of move that was first sampled. The detailed structure of the proposal probability distribution can be found in Appendix B.
Stoke uses uniform distributions for each of the elementary probability distributions the model samples from. This corresponds to a specific instantiation of the general stochastic search paradigm. In this work, we propose to learn those probability distributions so as to maximize the probability of reaching the best programs. The rest of the optimization scheme remains similar to the one of Schkufza et al. (2013).
Our chosen parameterization of q is to keep the hierarchical structure of the original work of Schkufza et al. (2013), as detailed in Appendix B, and parameterize all the elementary probability distributions (over the positions in the programs, the instructions to propose or the arguments) independently. The set θ of parameters for qθ will thus contain a set of parameters for each elementary probability distributions. A fixed proposal distribution is kept through the optimization of a given program, so the proposal distribution needs to be evaluated only once, at the beginning of the optimization and not at every iteration of MCMC.
The stochastic computation graph corresponding to a run of the Metropolis algorithm is given in Figure 1. We have assumed the operation of evaluating the cost of a program to be a deterministic function, as we will not model the randomness of measuring performance.
3.2.3 LEARNING THE PROPOSAL DISTRIBUTION
In order to learn the proposal distribution, we will use stochastic gradient descent on our loss function (6). We obtain the first order derivatives with regards to our proposal distribution parameters using the REINFORCE (Williams, 1992) estimator, also known as the likelihood ratio estimator (Glynn, 1990) or the score function estimator (Fu, 2006). This estimator relies on a rewriting of
the gradient of the expectation. For an expectation with regards to a probability distribution x ∼ fθ, the REINFORCE estimator is:
∇θ ∑ x f(x; θ)r(x) = ∑ x r(x)∇θf(x; θ) = ∑ x f(x; θ)r(x)∇θ log(f(x; θ)), (7)
and provides an unbiased estimate of the gradient.
A helpful way to derive the gradients is to consider the execution traces of the search procedure under the formalism of stochastic computation graphs (Schulman et al., 2015). We introduce one “cost node” in the computation graphs at the end of each iteration of the sampler. The associated cost corresponds to the normalized difference between the best rewrite so far and the current rewrite after this step:
ct = min
( 0, ( cost (Rt, T )−mini=0..t−1 cost (Ri, T )
cost (R0, T )
)) . (8)
The sum of all the cost nodes corresponds to the sum of all the improvements made when a new lowest cost was achieved. It can be shown that up to a constant term, this is equivalent to our objective function (5). As opposed to considering only a final cost node at the end of the T iterations, this has the advantage that moves which were not responsible for the improvements would not get assigned any credit.
For each round of MCMC, the gradient with regards to the proposal distribution is computed using the REINFORCE estimator which is equal to
∇̂θ,iL(θ) = (∇θ log qθ(Ri|Ri−1)) ∑ t>i ct. (9)
As our proposal distribution remains fixed for the duration of a program optimization, these gradients needs to be summed over all the iterations to obtain the total contribution to the proposal distribution. Once this gradient is estimated, it becomes possible to run standard back-propagation with regards to the features on which the proposal distribution is based on, so as to learn the appropriate feature representation.
4 EXPERIMENTS
4.1 SETUP
Implementation Our system is built on top of the Stoke super-optimizer from Schkufza et al. (2013). We instrumented the implementation of the Metropolis algorithm to allow sampling from parameterized proposal distributions instead of the uniform distributions previously used. Because the proposal distribution is only evaluated once per program optimisation, the impact on the optimization throughput is low, as indicated in Table 3.
Our implementation also keeps track of the traces through the stochastic graph. Using the traces generated during the optimization, we can compute the estimator of our gradients, implemented using the Torch framework (Collobert et al., 2011).
Datasets We validate the feasibility of our learning approach on two experiments. The first is based on the Hacker’s delight (Warren, 2002) corpus, a collection of twenty five bit-manipulation programs, used as benchmark in program synthesis (Gulwani et al., 2011; Jha et al., 2010; Schkufza et al., 2013). Those are short programs, all performing similar types of tasks. Some examples include identifying whether an integer is a power of two from its binary representation, counting the number of bits turned on in a register or computing the maximum of two integers. An exhaustive description of the tasks is given in Appendix C. Our second corpus of programs is automatically generated and is more diverse.
Models The models we are learning are a set of simple elementary probabilities for the categorical distribution over the instructions and over the type of moves to perform. We learn the parameters of each separate distribution jointly, using a Softmax transformation to enforce that they are proper probability distributions. For the types of move where opcodes are chosen from a specific subset, the probabilities of each instruction are appropriately renormalized. We learn two different type of models and compare them with the baseline of uniform proposal distributions equivalent to Stoke.
Our first model, henceforth denoted the bias, is not conditioned on any property of the programs to optimize. By learning this simple proposal distribution, it is only possible to capture a bias in the dataset. This can be understood as an optimal proposal distribution that Stoke should default to.
The second model is a Multi Layer Perceptron (MLP), conditioned on the input program to optimize. For each input program, we generate a Bag-of-Words representation based on the opcodes of the program. This is embedded through a three hidden layer MLP with ReLU activation unit. The proposal distribution over the instructions and over the type of moves are each the result of passing the outputs of this embedding through a linear transformation, followed by a SoftMax.
The optimization is performed by stochastic gradient descent, using the Adam (Kingma & Ba, 2015) optimizer. For each estimate of the gradient, we draw 100 samples for our estimator. The values of the hyperparameters used are given in Appendix A. The number of parameters of each model is given in Table 1.
Model Training Test
Uniform 57.01% 53.71% Bias 36.45 % 31.82 % MLP 35.96 % 31.51 %
Table 2: Final average relative score on the Hacker’s Delight benchmark. While all models improve with regards to the initial proposal distribution based on uniform sampling, the model conditioning on program features reach better performances.
4.2 EXISTING PROGRAMS
In order to have a larger corpus than the twenty-five programs initially present in “Hacker’s Delight”, we generate various starting points for each optimization. This is accomplished by running Stoke with a cost function where ωp = 0 in (1), and keeping only the correct programs. Duplicate programs are filtered out. This allows us to create a larger dataset from which to learn. Examples of these programs at different level of optimization can be found in Appendix D.
We divide this augmented Hacker’s Delight dataset into two sets. All the programs corresponding to even-numbered tasks are assigned to the first set, which we use for training. The programs corresponding to odd-numbered tasks are kept for separate evaluation, so as to evaluate the generalisation of our learnt proposal distribution.
The optimization process is visible in Figure 2, which shows a clear decrease of the training loss and testing loss for both models. While simply using stochastic super-optimization allows to discover programs 40% more efficient on average, using a tuned proposal distribution yield even larger improvements, bringing the improvements up to 60%, as can be seen in Table2. Due to the similarity between the different tasks, conditioning on the program features does not bring any significant improvements.
In addition, to clearly demonstrate the practical consequences of our learning, we present in Figure 3 a superposition of score traces, sampled from the optimization of a program of the test set. Figure 3a corresponds to our initialisation, an uniform distribution as was used in the work of Schkufza et al. (2013). Figure 3d corresponds to our optimized version. It can be observed that, while the uniform proposal distribution was successfully decreasing the cost of the program, our learnt proposal distribution manages to achieve lower scores in a more robust manner and in less iterations. Even using only 100 iterations (Figure 3e), the learned model outperforms the uniform proposal distribution with 400 iterations (Figure 3c).
4.3 AUTOMATICALLY GENERATED PROGRAMS
While the previous experiments shows promising results on a set of programs of interest, the limited diversity of programs might have made the task too simple, as evidenced by the good performance of a blind model. Indeed, despite the data augmentation, only 25 different tasks were present, all variations of the same programs task having the same optimum.
To evaluate our performance on a more challenging problem, we automatically synthesize a larger dataset of programs. Our methods to do so consists in running Stoke repeatedly with a constant cost function, for a large number of iterations. This leads to a fully random walk as every proposed programs will have the same cost, leading to a 50% chance of acceptance. We generate 600 of these programs, 300 that we use as a training set for the optimizer to learn over and 300 that we keep as a test set.
The performance achieved on this more complex dataset is shown in Figure 4 and Table 4.
Model Training Test
Uniform 76.63% 78.15 % Bias 61.81% 63.56% MLP 60.13% 62.27%
Table 4: Final average relative score. The MLP conditioning on the features of the program perform better than the simple bias. Even the unconditioned bias performs significantly better than the Uniform proposal distribution.
5 CONCLUSION
Within this paper, we have formulated the problem of optimizing the performance of a stochastic super-optimizer as a Machine Learning problem. We demonstrated that learning the proposal distribution of a MCMC sampler was feasible and lead to faster and higher quality improvements. Our approach is not limited to stochastic superoptimization and could be applied to other stochastic search problems.
It is interesting to compare our method to the synthesis-style approaches that have been appearing recently in the Deep Learning community (Graves et al., 2014) that aim at learning algorithms directly using differentiable representations of programs. We find that the stochastic search-based approach yields a significant advantage compared to those types of approaches, as the resulting program can be run independently from the Neural Network that was used to discover them.
Several improvements are possible to the presented methods. In mature domains such as Computer Vision, the representations of objects of interests have been widely studied and as a result are successful at capturing the information of each sample. In the domains of programs, obtaining informative representations remains a challenge. Our proposed approach ignores part of the structure of the program, notably temporal, due to the limited amount of existing data. The synthetic data having no structure, it wouldn’t be suitable to learn those representations from it. Gathering a larger dataset of frequently used programs so as to measure more accurately the practical performance of those methods seems the evident next step for the task of program synthesis.
A HYPERPARAMETERS
A.1 ARCHITECTURES
The output size of 9 corresponds to the types of move. The output size of 2903 correspond to the number of possible instructions that Stoke can use during a rewrite. This is smaller that the 3874 that are possible to find in an original program.
A.2 TRAINING PARAMETERS
All of our models are trained using the Adam (Kingma & Ba, 2015) optimizer, with its default hyper-parameters β1 = 0.9, β2 = 0.999, = 10−8. We use minibatches of size 32.
The learning rate were tuned by observing the evolution of the loss on the training datasets for the first iterations. The picked values are given in Table 7. Those learning rates are divided by the size of the minibatches.
B STRUCTURE OF THE PROPOSAL DISTRIBUTION
The sampling process of a move is a hierarchy of sampling step. The easiest way to represent it is as a generative model for the program transformations. Depending on what type of move is sampled, different series of sampling steps have to be performed. For a given move, all the probabilities are sampled independently so the probability of proposing the move is the product of the probability of picking each of the sampling steps. The generative model is defined in Figure 5. It is going to be parameterized by the the parameters of each specific probability distribution it samples from. The default Stoke version uses uniform probabilities over all of those elementary distributions.
C HACKER’S DELIGHT TASKS
The 25 tasks of the Hacker’s delight Warren (2002) datasets are the following:
1. Turn off the right-most one bit
2. Test whether an unsigned integer is of the form 2(n− 1) 3. Isolate the right-most one bit 4. Form a mask that identifies right-most one bit and trailing zeros 5. Right propagate right-most one bit 6. Turn on the right-most zero bit in a word 7. Isolate the right-most zero bit 8. Form a mask that identifies trailing zeros 9. Absolute value function
10. Test if the number of leading zeros of two words are the same 11. Test if the number of leading zeros of a word is strictly less than of another work 12. Test if the number of leading zeros of a word is less than of another work 13. Sign Function 14. Floor of average of two integers without overflowing 15. Ceil of average of two integers without overflowing 16. Compute max of two integers 17. Turn off the right-most contiguous string of one bits 18. Determine if an integer is a power of two 19. Exchanging two fields of the same integer according to some input 20. Next higher unsigned number with same number of one bits 21. Cycling through 3 values 22. Compute parity 23. Counting number of bits 24. Round up to next highest power of two 25. Compute higher order half of product of x and y
Reference implementation of those programs were obtained from the examples directory of the stoke repository (Churchill et al., 2016).
D EXAMPLES OF HACKER’S DELIGHT OPTIMISATION
The first task of the Hacker’s Delight corpus consists in turning off the right-most one bit of a register.
When compiling the code in Listing 6a, llvm generates the code shown in Listing 6b. A typical example of an equivalent version of the same program obtained by the data-augmentation procedure is shown in Listing 6c. Listing 6d contains the optimal version of this program.
Note that such optimization are already feasible using the stoke system of Schkufza et al. (2013).

This work builds on top of STOKE (Schkufza et al., 2013), which is a superoptimization engine for program binaries. It works by starting with an existing program, and proposing modifications to it according to a proposal distribution. Proposals are accepted according to the Metropolis-Hastings criteria. The acceptance criteria takes into account the correctness of the program, and performance of the new program. Thus, the MCMC process is likely to converge to correct programs with high performance. Typically, the proposal distribution is fixed. The contribution of this work is to learn the proposal distribution as a function of the features of the program (bag of words of all the opcodes in the program). The experiments compare with the baselines of uniform proposal distribution, and a baseline where one just learns the weights of the proposal distribution but without conditioning on the features of the program. The evaluation shows that the proposed method has slightly better performance than the compared baselines.
The significance of this work at ICLR seems to be quite low., both because this is not a progress in learning representations, but a straightforward application of neural networks and REINFORCE to yet another task which has non-differentiable components. The task itself (superoptimization) is not of significant interest to ICLR readers/attendees. A conference like AAAI/UAI seem a better fit for this work.
The proposed method is seemingly novel. Typical MCMC-based synthesis methods are lacking due to their being no learning components in them. However, to make this work compelling, the authors should consider demonstrating the proposed method in other synthesis tasks, or even more generally, other tasks where MH-MCMC is used, and a learnt proposal distribution can be beneficial. Superoptimization alone (esp with small improvements over baselines) is not compelling enough.
It is also not clear if there is any significant representation learning is going on. Since a BoW feature is used to represent the programs, the neural network cannot possibly learn anything more than just correlations between presence of opcodes and good moves. Such a model cannot possibly understand the program semantics in any way. It would have been a more interesting contribution if the authors had used a model (such as Tree-LSTM) which attempts to learn the semantics the program. The quite naive method of learning makes this paper not a favorable candidate for acceptance.