Learning Where to Sample in Structured Prediction

Tianlin Shi Jacob Steinhardt Percy LiangTsinghua Universitytianlinshi@gmail.com

Stanford Universityjsteinhardt@cs.stanford.edu

Stanford Universitypliang@cs.stanford.edu

Abstract

In structured prediction, most inference al-gorithms allocate a homogeneous amount ofcomputation to all parts of the output, whichcan be wasteful when different parts varywidely in terms of difficulty. In this paper,we propose a heterogeneous approach thatdynamically allocates computation to the dif-ferent parts. Given a pre-trained model, wetune its inference algorithm (a sampler) toincrease test-time throughput. The inferencealgorithm is parametrized by a meta-modeland trained via reinforcement learning, whereactions correspond to sampling candidateparts of the output, and rewards are log-likelihood improvements. The meta-modelis based on a set of domain-general meta-features capturing the progress of the sam-pler. We test our approach on five datasetsand show that it attains the same accuracyas Gibbs sampling but is 2 to 5 times faster.

1 Introduction

For many structured prediction problems, the outputcontains many interdependent variables, resulting inexponentially large output spaces. These propertiesmake exact inference intractable for models with hightreewidth [Koller et al., 2007], and thus we must relyon approximations such as variational inference andMarkov Chain Monte Carlo (MCMC). A key charac-teristic of many such approximate inference algorithmsis that they iteratively modify only a local part of theoutput structure at a small computational cost. Forexample, Gibbs sampling [Brooks et al., 2011] updatesthe output by sampling one variable conditioned on

Appearing in Proceedings of the 18th International Con-ference on Artificial Intelligence and Statistics (AISTATS)2015, San Diego, CA, USA. JMLR: W&CP volume 38.Copyright 2015 by the authors.

the rest. Often, a large number of local moves is re-quired.

One source of inefficiency in Gibbs sampling is that itdedicates a homogeneous amount of inference to eachpart of the output. However, in practice, the diffi-culty and inferential demands of each part is hetero-geneous. For example, in rendering computer graph-ics, paths of light passing through highly reflective orglossy regions deserve more sampling [Veach, 1997];in named-entity recognition [McCallum and Li, 2003],most tokens clearly do not contain entities and there-fore should be allocated less computation. Attemptshave been made to capture the nature of heterogene-ity in such settings. For example, Elidan et al. [2006]schedule updates in asynchronous belief propagationbased on the information residuals of the messages.Chechetka and Guestrin [2010] focus the computationof belief propagation based on the specific variablesbeing queried. Other work has focused on buildingcascades of coarse-to-fine models, where simple modelsfilter out unnecessary parts of the output and reducethe computational burden for complex models [Violaand Jones, 2001, Weiss and Taskar, 2010].

We propose a framework that constructs heteroge-neous sampling algorithms using reinforcement learn-ing (RL). We start with a collection of transition ker-nels, each of which proposes a modification to part ofthe output (in this paper, we use transition kernelsderived from Gibbs sampling). At each step, our pro-cedure chooses which transition kernel to apply basedon cues from the input and the history of proposedoutputs. By optimizing this procedure, we fine-tuneinference to exploit the specific heterogeneity in thetask of interest, thus saving overall computation attest time.

The main challenge is to find signals that consistentlyprovide useful cues (meta-features) for directing the fo-cus of inference across a wide variety of tasks. More-over, it is important that these signals are cheap tocompute relative to the cost of generating a sample,as otherwise the meta-features themselves become thebottleneck of the algorithm. In this paper, we provide

Learning Where to Sample in Structured Prediction

general principles for constructing such meta-features,based on reasoning about staleness and discord of vari-ables in the output. We cache these ideas out as acollection of five concrete meta-features, which empir-ically yield good predictions across tasks as diverse aspart-of-speech (POS) tagging, named-entity recogni-tion (NER), handwriting recognition, color inpainting,and scene decomposition.

In summary, our contributions are:

• The conceptual idea of learning to sample: wepresent a learning framework based on RL, anddiscuss meta-features that leverage heterogeneity.

• The practical value of the framework: given apre-trained model, we can effectively optimize thetest-time throughput of its Gibbs sampler.

2 Heterogeneous Sampling

Before we formalize our framework for heterogeneoussampling, let’s consider a motivating example.

x I think now is the right timepass 1: y PRP VBP RB VBZ DT NN NNpass 2: y PRP VBP RB VBZ DT JJ NN

Table 1: A POS tagging example. Outputs arerecorded after each sweep of Gibbs sampling. Onlythe ambiguous token “right” (NN: noun, JJ: adjective)needs more inference at the second sweep.

Suppose our task is part-of-speech (POS) tagging,where the input x P X is a sentence and the out-put y P Y is the sequence of POS tags for the words.An example is shown in Table 1. Suppose that thefull model is a chain-structured conditional randomfield (CRF) with unigram potentials on each tag andbigram potentials between adjacent tags. Exact infer-ence algorithms exist for this model, but for illustrativepurposes we use cyclic Gibbs sampling, which samplesfrom the conditional distribution of each tag in cyclicorder from left to right.

The example in Table 1 shows at least two sweeps ofcyclic Gibbs sampling are required, because it is hardto know whether “right” is an adjective or a noun un-til the tag for the following word “time” is sampled.However, the second pass wastes computation by sam-pling other tags that are mostly determined at the firstpass. This inspires the following inference strategy:

pass 1 sample the tags for each word.pass 2 sample the tag for “right”.

In general, it is desirable to have the inference algo-rithm itself figure out which locations to sample, and

Algorithm 1 Template for a heterogeneous sampler

1: Initialize yr0s „ P0pyq for some initializing P0p¨q.2: for t “ 1 to TMpxq do3: select transition kernel Aj for some 1 ďjď m4: sample yrts „ Ajp¨ | yrt´ 1sq5: end for6: output y “ yrT s

at a higher level, which test instances to sample.

2.1 Framework

We now formalize the intuition from the previous ex-ample. Assume our pre-trained model specifies a dis-tribution ppy | xq. On a set of test inputs X “

pxp1q, . . . ,xpnqq, we would like to infer the outputsY “ pyp1q, . . . ,ypnqq using some inference algorithmM. To simplify notation, we will focus on a single in-stance px,yq, though our final algorithm considers alltest instances jointly. Notice that we can reduce fromthe multiple-instance case to the single-instance caseby just concatenating all the instances into a singleinstance.

We further assume that a single output y “

py1, . . . , ymq is represented by m variables. For in-stance, in POS tagging, yj (j “ 1, . . . ,m) is the part-of-speech of the j-th word in the sentence x. We aregiven a collection of transition kernels which target thedistribution ppy | xq. For Gibbs sampling, we have thekernels tAjpy

1 | yq : j “ 1, . . . ,mu, where the transi-tion Aj samples y1j conditioned on all other variablesy j , and leaves y1 j equal to y j .

Algorithm 1 describes the form of samplers we con-sider. A sampler generates a sequence of outputsyr1s,yr2s, . . . by iteratively selecting a variable indexjrts and sampling yrt ` 1s „ Ajrtsp¨ | yrtsq. Ratherthan applying the transition kernels in a fixed order,our samplers select the transition Ajrts to apply basedon the input x together with the sampling history.The total number of Markov transitions TMpxq madebyM characterizes its computational cost on input x.

How do we choose which transition kernel to apply?A natural objective is to maximize the expected log-likelihood under the model ppy | xq of the sampleroutput Mpxq:

maxM

EqMpy|xqrlog ppy | xqs (1)

s.t.: TMpxq ď T˚,

where qMpy | xq is the probability that M outputs yand T˚ is the computation budget. Equation (1) saysthat we want qM to place as much probability massas possible on values of y that have high probability

Tianlin Shi, Jacob Steinhardt, Percy Liang

under p, subject to a constraint T˚ on the amountof computation. Note that if T˚ “ 8, the optimalsolution would be the posterior mode.

Solving this optimization problem at test time is in-feasible, so we will instead optimize M on a trainingset, and then deploy the resulting M at test time.

3 Reinforcement Learning ofHeterogeneous Samplers

We would like to optimize the objective in (1), butsearching over all samplers M and evaluating the ex-pectation in (1) are both difficult. We use reinforce-ment learning (RL) to find an approximate solution.

Reduction to RL. Recall yr0s,yr1s, . . . is the se-quence of outputs generated by our sampler, whereyrt ` 1s „ Ajrtsp¨ | yrtsq. To cast our prob-lem into the RL framework, let the state st “

pyr0s . . . ,yrts, jr0s, . . . , jrt ´ 1sq be the entire historyof samples, and the action at “ Ajrts refers to thetransition kernel that produces yrt` 1s from yrts. Welet the reward be the improvement in log-probability:

Rpst, at, st`1q “ log ppyrt 1̀s | xq´log ppyrts | xq. (2)

Let S be the space of states and A be the space ofactions as defined above. The goal of RL is to find apolicy π : S Ñ A to maximize the expected cumulativereward ErRT˚s, where

RT “

T´1ÿ

t“0

Rpst, at, st`1q. (3)

Clearly, the total reward RT is equal to log ppyrT s |xq ´ log ppyr0s | xq; since yr0s is independent of theparticular policy, maximizing cumulative reward isequivalent to maximizing the original objective in (1).Although reward only depends on the last output yrtsof a state st, we will allow the policy to depend on thefull history encapsulated in st.

Learning algorithm. Our algorithm learns anaction-value function Qpst, Ajq that predicts the valueof using Aj in state st. Standard reinforcement learn-ing methods, such as Q-learning [Watkins and Dayan,1992] and SARSA [Rummery and Niranjan, 1994], donot work well for our purpose. The issue is that theyattempt to learn a function Qpst, Ajq that models theexpected cumulative future reward if we take actionAj in state st. In our setting this cumulative rewardis hard to predict for the following two reasons:

• It is difficult to estimate how far the current sam-ple is from the global optima. As an approxima-

tion, we use Qpst, Ajq to predict the cumulativereward over a shorter time horizon H ! T˚.

• The reward over time H also depends on the con-text of yj . By subtracting the contextual part ofthe reward, we can hope to isolate the contribu-tion of action Aj . Thus, we use Qpst, Ajq to modelthe difference in reward from taking a transitionAj , relative to the baseline of making no transi-tion.

Formally, we learn Q using sample backup [Sutton andBarto, 1998, Chapter 9.5]. We start with some state-action sequence ps0, a0, s1, a1, . . . , sT˚q sampled froma fixed exploration policy π1. Then, for each indext pt “ 0, . . . , T˚ ´ 1q in the sequence, we generate arollout starting from initial state sc1 “ st`1: for i “1, 2, . . . ,H, we generate action aci “ arg maxaQps

ci , aq

and state sci`1 using aci , and define the utility due totaking at:

U c “ Rpst, at, st`1q `

Hÿ

i“1

Rpsci , aci , sci`1q. (4)

Next, consider starting from state sb1 “ st and nottaking at, and letting abi “ arg maxaQps

bi , aq. The

resulting states and actions sbi , abi define the following

utility:

U b “

Hÿ

i“1

Rpsbi , abi , sbi`1q. (5)

To model the Q-function, we use a single-layer neuralnetwork with one hidden node [Tesauro, 1995]:

Qps, aq “ w σpα ¨ φps, aqq ` b, (6)

where σp¨q is the logistic function, φps, aq P RL aremeta-features and α P RL, w P R, b P R are the meta-parameters; we write θ “ pw, b, αq.

We update θ with a temporal difference update [Sut-ton and Barto, 1998] based on our rollout:

θ Ð θ ` ηt ˚ dt, (7)

where ηt P RL`2 are step sizes, “ ˚ ” is element-wisemultiplication and the vector dt is

dt “´

U c ´ U b ´Qpst, atq¯

∇θQpst, atq. (8)

For choosing ηt, we use AdaGrad from Duchi et al.[2010]:

ηt “η

b

δ `řt

i“0 di ˚ di

, (9)

where η is the meta learning rate and δ is a smoothingparameter (we use η “ 1 and δ “ 10´4 in experi-ments).

Learning Where to Sample in Structured Prediction

We use the cyclic Gibbs sampler as the fixed explo-ration policy π1 to generate the initial states for therollouts. The entire training procedure is shown inAlgorithm 2:

Algorithm 2 Learning a heterogeneous sampler.

1: Input Dataset X, transition kernels Aj , numberof epochs E , cyclic Gibbs policy π1, time horizonH.

2: Initialize y „ P0pyq. Set s0 “ pyr0sq.3: for epoch = 1, . . . , E do4: for t “ 0, . . . , T˚ ´ 1 do5: Get action at “ π1pstq.6: Sample st`1 from at.7: Extract φpst, atq.8: Estimate gradient using (8)9: Update meta-parameters θ via (7).

10: end for11: end for

Test time. At test time, we apply Algorithm 1. Inparticular, we maintain a priority queue over locationsj, where priorities are the Q-values.

• During the select step, computearg maxAj

Qpst, Ajq for state st, popping offthe maximum element from the priority queue.

• After the sample step, re-compute the meta-features of the actions that depend on jrts andupdate the corresponding Q-values in the priorityqueue.

In terms of computational complexity, select takesOplogmq, and re-scoring takes OpLnj logmq, wherem is the number of variables, and nj is the numberof neighbors of variable yj . The complexity of meta-feature computation will be discussed in Section 4.

4 Meta-Features

The effectiveness of our framework hinges on the exis-tence of good meta-features for the Q-function in (6).Our goal is to develop general meta-features that ex-hibit strong predictive power across many datasets. Inthis section, we offer a few guiding principles, whichculminates in a set of five meta-feature templates.

Principles. One necessary criterion is that comput-ing the meta-features should be computationally cheaprelative to the rest of the inference algorithm. With-out loss of generality, we assume each output variableyj takes one of K values and has nj neighbors. Thensuch meta-features would take, for example, OpKq or

Opnjq to compute, and they should not be computedmore often than being sampled.

In order to satisfy this criterion, an idea that we havefound consistently useful is stale values. For example,suppose we would like to use the entropy of ppyj |y jq as a meta-feature. Computing it would be justas expensive as sampling from Aj . Instead, we keeptrack of a stale version of conditional entropy. Everytime we sample from Aj , we compute the entropy ofyj according to the sampling distribution and store itin memory. Then the stale conditional entropy of yjis defined as the current entropy in memory. As thestale conditional entropy is a meta-feature, we leave itup to the learning algorithm to determine how muchto trust it.

Reasoning about staleness is valuable in another wayas well: we want to know how different the Markovblanket of yj is relative to the last time it was sampled.If it is very different, this tells us two things: first,any stored quantities such as entropy that we havecomputed are probably out of date. More importantly,the conditional distribution ppyj | y jq is probablyvery different than when we last sampled from Aj , soit is probably a good idea to sample from Aj again.

In addition to staleness, another important notion isdiscord. If at least one neighbor of yj is inconsistentwith the current value of yj , then it is probably worthsampling from Aj .

Templates. Based on the ideas of staleness and dis-cord, we introduce the following meta-feature tem-plates:

vary: the number of variables in the Markov blanketof yj that have changed since the last time we sam-pled from Aj. The more neighbors of yj that havechanged, the more stale yj could be. This meta-feature is computed as follows: upon sampling yj , weset varypyjq “ 0, and increment vary of its neighborsby 1. So the computational complexity of vary is onlyOpnjq.

nb: discord with neighbors. We define meta-featuresnb-y-y1 (y, y1 P t1, . . . ,Ku) which are binary indicatorsto track the occurrences of each of the possible valuepairs py, y1q for yj and one of its neighbors. The intu-ition is that certain pairs of values are very unlikely tooccur as part of a legitimate output, and thus repre-sent a discordance that requires taking more samples;the nb meta-features allow us to learn these pairs fromdata. Although the total number of nb meta-featuresis K2, the computational complexity is still Opnjq dueto sparsity.

cond-ent: the conditional entropy of yj at the last

Tianlin Shi, Jacob Steinhardt, Percy Liang

time it was sampled. Variables with high conditionalentropy have a high degree of uncertainty and can ben-efit from being sampled further.

However, as noted earlier, keeping track of an up-to-date conditional entropy would be too computationallyexpensive. We therefore use a stale version based onthe last time that the variable yj in question was sam-pled by Aj . The computational overhead of cond-entis OpKq, since we can just compute the entropy basedon ppyj | y jq, which already needs to be computed inorder to sample from Aj .

unigram-ent: Sometimes, we can fit a simpler un-igram model to the training dataset, where all thevariables yj are independent (given the input). Insuch cases, unigram entropy of ppyjq can be used asan over-estimate of the degree of ambiguity for a givenvariable [Bishop, 2006]. If the unigram entropy is verylow, the variable is probably not worth sampling. Tocapture this, we use the indicator function for the uni-gram entropy being below some threshold (10´4 in ourexperiments).

sp: number of times yj has been sampled thus far. Apotential problem with all of the above meta-featuresis that they might overly explore possibilities for thesame variable. So we need some way to reason aboutthe fact that at some point sampling the same variablemore is unlikely to lead to improvements. We do thisby keeping track of the number of times a variablehas already been sampled: once a variable has beensampled too many times, it is unlikely that samplingit further will be fruitful.

5 Experiments

In this section, we provide empirical evaluation of ourmethod, which we call HeteroSampler.

5.1 Datasets

To evaluated our method on five tasks: part-of-speech tagging (POS), named-entity recognition(NER), handwriting recognition, color inpainting, andscene decomposition.

The general setup is as follows. First, we use RL tolearn the parameters of the meta-model on the trainingdataset. We then run Algorithm 1 on the test set.Unless otherwise stated, in all experiments we use E “3 training epochs, step size η “ 1, and time horizonH “ 1. In addition to using cyclic Gibbs to generatea base policy for training, we also compare to cyclicGibbs at test time.

We evaluated on the following datasets:

Words   Japan coach Shu Kamo said : ' ' The Syrian own goal proved lucky for usTruth   B-LOC O B-PER I-PER O O O O O B-MISC O O O O O O

1 I-ORG O B-PER I-PER O O O O O B-LOC O O O O O O

2 B-LOC O B-PER I-PER O O O O O B-MISC O O O O O O

3 B-LOC O B-PER I-PER O O O O O B-MISC O O O O O O

(a) Cyclic Gibbs sampler on NER-f4

Words   Japan coach Shu Kamo said : ' ' The Syrian own goal proved lucky for usTruth   B-LOC O B-PER I-PER O O O O O B-MISC O O O O O O

1 I-ORG O B-PER I-PER O O O O O B-LOC O O O O O O

2 B-LOC O B-PER I-PER O O O O O B-MISC O O O O O O

3 B-LOC O B-PER I-PER O O O O O B-MISC O O O O O O

(b) HeteroSampler on NER-f4

Figure 1: Visualization of computational resource allo-cation on a test example from NER-f4. Each row isa snapshot of the sample after k “ 1, 2, 3 sweeps. Adarker color means that more cumulative samples havebeen taken. For HeteroSampler, one sweep correspondsto making m transitions, where m is the number of vari-ables. HeteroSampler has learned to sample harder partsof the instance, such as ambiguous tokens.

POS/NER Tagging. The POS tagging datasetcomes from the standard Wall Street Journal (WSJ)section of the Penn Treebank, and the NER tag-ging dataset is taken from the 2003 CoNLL SharedTask. We trained an CRF model with fea-tures between each token and its corresponding tag(i.e. features rgpxiq, yis with feature extractor gp¨q),and higher-order features between/among tags (i.e.ryi, yi`1, ..., yi`qs) [Liang et al., 2008]. We refer to q asthe factor size, and call the corresponding tasks POS-fq and NER-fq. The feature extraction functions gp¨qinclude prefixes and suffixes of the word (up to length4), lowercase word, word signature (e.g. McDiarmidinto AA and AaAaaaaaa, banana into a and aaaaaa)and an indicator of the word’s being capitalized. TheCRF is trained using AdaGrad [Duchi et al., 2010] with5 passes over the training set, using Gibbs sampling forinference with 8 total sweeps over each instance and 5sweeps as burn-in. The performance is evaluated viatag accuracy for POS and F1 score for NER.

Handwriting Recognition. We use the handwrit-ing recognition dataset from Weiss and Taskar [2010].The data were originally collected by Kassel [1995],and contain 6877 handwritten words from 150 subjectswith 55 distinct words. Each instance of the dataset isa word, which is a sequence of characters. Associatedwith each output character is a corresponding 16 ˆ 8input binary optical image. The dataset is split intotraining and test set, where the training set had 6251words and the test set had 626 words.

Similar to POS/NER tagging, the baseline algo-rithm is an CRF. We have a feature for each (pixelvalue, location, character) triple, as well as higher-order n-gram potentials between consecutive charac-

Learning Where to Sample in Structured Prediction

(a) NER (factor size 2) (b) NER (factor size 4) (c) POS (factor size 4)

0 20 40 60 80 100 120Average Number of Transitions

0.70

0.72

0.74

0.76

0.78

0.80F1

sco

re

HeteroSampler

cyclic Gibbs

0 10 20 30 40 50 60 70 80Average Number of Transitions

0.70

0.72

0.74

0.76

0.78

0.80

F1

sco

re

HeteroSampler

cyclic Gibbs

0 20 40 60 80 100 120Average Number of Transitions

0.92

0.93

0.94

0.95

0.96

0.97

Acc

ura

cy

HeteroSampler

cyclic Gibbs

(d) OCR (factor size 4) (e) Color Inpainting (f) Scene Decomposition

0 20 40 60 80 100 120 140Average Number of Transitions

0.70

0.75

0.80

0.85

0.90

0.95

Acc

ura

cy

HeteroSampler

cyclic Gibbs

0 2 4 6 8 10 12 14 16

Average Number of Transitions (x 105 )

−20

−18

−16

−14

−12

−10

−8

−6

Log P

robabili

ty (

x 102

)

HeteroSampler

cyclic Gibbs

0 200 400 600 800 1000

Average Number of Transitions

830

840

850

860

870

880

890

Log P

robabili

ty

HeteroSampler

cyclic Gibbs

Figure 2: Accuracy vs.number of transitionsacross several tasks.HeteroSampler con-verges much faster thancyclic Gibbs sampling.On the color inpaint-ing task, dynamicallychoosing the order ofsampling also allowsthe algorithm to find abetter local optimum.

ters. The training scheme of the full model is similarto POS/NER, except that we use 16 Gibbs samplingsweeps in total with 5 burn-in sweeps. The results areevaluated via character-wise accuracy.

Color Inpainting. The three-class color inpaintingtask is borrowed from ?. The input is a corruptedcolor image in a circular domain, and the target im-age is an equipartition of the circle using three colors.We use a pre-trained model from the OpenGM bench-mark [Kappes et al., 2013]. The baseline is Gibbs sam-pling with 100 sweeps over the instance. There are twoinstances in this dataset and we use one to train theHeteroSampler and the other to test it. Performanceis evaluated based on the log-probability of the output.

Scene Decomposition. The scene decompositiondataset is obtained from the source in Gould et al.[2009]. The goal is to segment a natural image intoeight semantic categories, such as grass and sky. Weuse the subset of 715 images included in the OpenGMtoolkit [Kappes et al., 2013], for which an existinggraphical model is publicly available. The graphicalmodel is a superpixel factor graph, and each super-pixel has 773 feature dimensions. Among the 715 im-ages, we randomly pick 358 instances for training, andthe rest are used for testing. The baseline Gibbs sam-pling uses 16 sweeps over each image, and we train theHeteroSampler in the same way as Color Inpainting.Performance is again evaluated via log-probability.

5.2 Visualization of Resource Allocation

Figure 1 visualizes the allocation of sampling op-erations on an NER instance. While cyclic Gibbssampling uniformly distributes its computational re-sources, HeteroSampler is able to focus on harder

parts of the task such as names.

5.3 Performance under Different Budgets

To measure the performance under different budgets,we gradually increase the total number of transitionsat test time for the HeteroSampler and the overallnumber of sweeps for the Gibbs baseline. The numberof transitions for training are held fixed.

Figure 2 plots the performance versus average num-ber of transitions per instance. As we see, giventhe same budget, HeteroSampler achieves equalor better performance for all tasks. In addition,HeteroSampler reaches the ceiling accuracy 2 to 5times faster regardless of the problem domain. For theColor Inpainting problem, which is the most challeng-ing of the tasks, HeteroSampler also achieves betterend performance when it converges. This is due tothe fact that, by optimizing the order of sampling, themeta-algorithm is able to find a better local optimum.

Next, we justify measuring computational cost interms of number of samples. To do this, we mea-sured the overhead of computing the policy relativeto the cost of sampling. As long as this overheadis low, computing samples is the bottleneck in thebase algorithm, and so number of samples is a reason-able measure of computational cost. Figure 3 showshow many wall-clock seconds were spent computingthe HeteroSampler policy for each dataset. For mosttasks, sampling involves computing all of the featuresof the full model, while computing the policy uses onlya few meta-features and therefore has negligible cost.The exception is color inpainting, where the full modelhas only a few features.

Tianlin Shi, Jacob Steinhardt, Percy Liang

(a) NER (factor size 2) (b) NER (factor size 4)

40 60 80 100 120 140 160 180 200Average Number of Transitions

0

20

40

60

80

100

120

140

160

180W

all-

clock

Seco

nds

Policy

Overall

40 60 80 100 120 140 160 180 200Average Number of Transitions

0

20

40

60

80

100

120

140

160

180

Wall-

clock

Seco

nds

Policy

Overall

(c) POS (factor size 4) (d) OCR (factor size 4)

100 200 300 400 500 600 700 800 900Average Number of Transitions

0

100

200

300

400

500

600

700

800

900

Wall-

clock

Seco

nds

Policy

Overall

0 50 100 150 200 250Average Number of Transitions

0

50

100

150

200

250

Wall-

clock

Seco

nds

Policy

Overall

(e) Color Inpainting (f) Scene Decomposition

0 50 100 150 200 250 300 350Average Number of Transitions

0

10

20

30

40

50

60

70

80

Wall-

clock

Seco

nds

Policy

Overall

100 150 200 250 300Average Number of Transitions

0

50

100

150

200

250

Wall-

clock

Seco

nds

Policy

Overall

Figure 3: Wall-clock time vs. average number of transi-tions across different datasets. This measures the overheadof policy evaluation. The red line “sampling” shows thetime spent without policy evaluation, while the blue line“policy” shows the time spent on policy evaluation.

5.4 Cumulative Rewards

We would like to verify two facts: first, training withcumulative rewards is useful relative to just using im-mediate reward; second, our meta-features can predictthe cumulative rewards. First, Figure 4(a) shows ac-curacy vs. number of transitions with H “ 0 (im-mediate rewards) and H “ 1 (cumulative rewards)on NER-f4. As we can see, with cumulative rewards,HeteroSampler performs better. Table 2 provides anintuitive explanation.

To see which meta-features are contributing to the pre-

10 20 30 40 50 60 70Average Number of Transations

0.70

0.72

0.74

0.76

0.78

0.80

F1

H=0

H=1

oracle bias cond-ent sp0.0

0.2

0.4

0.6

0.8

1.0

oracle bias cond-ent sp−0.2

0.0

0.2

0.4

0.6

0.8

1.0

(a) (b) (c)

Figure 4: Effect of training with cumulative rewards. (a)Convergence curve of the inference algorithm trained withone-step look-ahead (H “ 1) vs. immediate rewards (H “

0). (b) Weights of meta-features when trained with H “

0 plus the oracle meta-feature. (c) Weights of the meta-features when trained with H “ 1 plus the oracle meta-feature.

input KANSAS CITY AT OAKLANDimmediate B-LOC I-LOC O B-LOCcumulative B-ORG I-ORG O B-LOC

Table 2: Cumulative rewards are often helpful whenthere is high correlation between variables. In this ex-ample, “KANSAS CITY ” is initially labeled as an lo-cation. Two coordinated actions are needed to changeit to an organization and improve log-likelihood. Ameta-model trained with immediate rewards would notrecognize the value of sampling “KANSAS” alone.

diction of cumulative rewards, we intentionally add anoracle meta-feature, which is the immediate rewardof sampling. Figures 4(b) and (c) visualize the weightsof some meta-features for H “ 0 and H “ 1 respec-tively. As expected, when learned with immediate re-wards, all weights concentrates on the oracle meta-feature. The learned meta-model does not encourageexploration and therefore may omit positions that havecumulative reward. When trained with cumulative re-wards, the meta-model also distributes some weight tocond-ent, which leads to better exploration and bet-ter performance.

5.5 Meta-feature Ablation Analysis

To evaluate the contribute of individual meta-featuresand to understand their role in predicting reward, wedid a meta-feature ablation analysis. With one meta-feature removed at a time, we run the meta-algorithmand produce a convergence curve, shown in Figure 5for NER-f2 and POS-f4. All meta-features play an im-portant role. sp is the most important; without it,we would repeatedly sample variables with high un-certainty.

6 Related Work and Discussion

At a high-level, our approach is about fine-tuning theinference algorithm of a pre-trained model to makemore effective use of computational resources at testtime. Specifically, we use reinforcement learning totrain a sampler that operates on variables heteroge-neously based on the promise of likelihood improve-ments. We demonstrated substantial speed improve-ments on several structured prediction tasks.

The idea of treating structured prediction as a se-quential decision-making problem has been exploredby SEARN [Daume et al., 2009] and DAGGER [Rosset al., 2011a]. Both train a multiclass classifier to buildup a structured output incrementally in a fixed order.Similar ideas have been applied in dependency parsing[Goldberg and Nivre, 2013].

Learning Where to Sample in Structured Prediction

14 16 18 20 22 24Average Transations Per Example

0.68

0.70

0.72

0.74

0.76

0.78

0.80F1

all

all\unigram-ent

all\nb

all\vary

all\cond-ent

all\sp

(a) Meta-feature ablation on NER (factor size is 2)

25 30 35 40 45 50 55 60Average Transations Per Example

0.92

0.93

0.94

0.95

0.96

Acc

ura

cy all

all\unigram-ent

all\nb

all\vary

all\cond-ent

all\sp

(b) Meta-feature ablation on POS (factor size is 4)

Figure 5: Meta-feature ablation study on various datasets. (a) shows the convergence curves of HeteroSampler withone meta-feature removed on (a) NER-f2 and (b) POS-f4. F is the entire meta-feature set, and “z” denotes excluding ameta-feature. We see that each meta-feature matters for at least some of the tasks, and sp is the most important.

To obtain speedups, it is beneficial to learn the or-der in which the structured output is constructed;this flexibility is the cornerstone of our work. Forexample, Goldberg and Elhadad [2010] proposed anapproach that learns to construct a dependency treeby adding “easy” arcs first. More generally, Jianget al. [2012] maintains a priority queue over partiallyconstructed hypotheses for constituency parsing andlearns to choose which one to process first. While theaforementioned work builds up outputs incrementally,our heterogeneous sampler makes modifications to fulloutputs, which can be more flexible.

Other work also operate in the space of full outputs.For example, Doppa et al. [2014b,a] perform severalsteps of local search around a baseline prediction.Zhang et al. [2014] performed greedy hill-climbingfrom multiple random starting points for dependencyparsing. Ross et al. [2011b] used DAGGER to learnmessage-passing inference algorithms. However, un-like our method, none of these papers deal with theissue of determining which locations are useful to op-erate on without explicitly evaluating the model scorefor each candidate modification. We use lightweightmeta-features for this purpose.

Some methods use a fixed strategy to prioritize infer-ence in a fixed model. For example, residual beliefpropagation [Elidan et al., 2006] selects the messagebetween two variables that has changed the most fromthe previous iteration. In cases where we are inter-ested in a particular query variable, Chechetka andGuestrin [2010] prioritizes messages based on impor-tance to the query. Wick and McCallum [2011] imple-ments the same intuition in the context of MCMC.

SampleRank [Wick et al., 2011] also performs learningin the context of sampling, but is complementary toour work: SampleRank fixes a sampling strategy andtrains the underlying model, whereas we fix the un-derlying model and train the sampling strategy usingdomain-general meta-features. Wick et al. [2009] tack-

les the local optima problem in structured predictionby using RL to train policies that could select fruitfuldownward jumps, which is the same issue that our useof cumulative rewards attempts to address.

More generally, the goal of speeding up inference attest time is quite established by now. Viola and Jones[2001] used a sequence of models from simple to com-plex for face detection, at each successive stage prun-ing out unlikely locations in the image. Weiss andTaskar [2010] trained a sequence of Markov models ofincreasing order, each successive stage pruning out un-likely local configurations.

As feature extraction is often the performance bottle-neck, it is a promising place to look for speed improve-ments. Weiss and Taskar [2013] used RL to train poli-cies that adaptively determine the value of informationof each feature at test time. For dependency parsing,He et al. [2013] considers a sequence of increasinglycomplex features, and uses DAGGER to learn whicharcs to commit to before adding more features.

Our work is superficially related to the work on adap-tive MCMC [Andrieu and Thoms, 2008], but the goalsare quite different. Adaptive MCMC samplers attemptto preserve the stationary distribution, while our ap-proach seeks to directly maximize log-likelihood withina fixed number of time steps.

As a final remark, in this paper, we only focused on theissue of “where to sample”, but the general framework,which merely learns which transition kernels to apply,could also be applied to determine “how to sample” tooby supplying a richer family of transition kernels—forexample, ones based on models with different featuresets or blocked samplers. This opens up a vast set ofpossibilities for finer-grained adaptivity.

Acknowledgements. This work was supported bythe Fannie & John Hertz Foundation for the secondauthor and the Microsoft Research Faculty Fellowshipfor the third author.

Tianlin Shi, Jacob Steinhardt, Percy Liang

References

C. Andrieu and J. Thoms. A tutorial on adaptiveMCMC. Statistics and Computing, 18(4):343–373,2008.

C. M. Bishop. Pattern recognition and machine learn-ing. Springer New York, 2006.

S. Brooks, A. Gelman, G. Jones, and X. Meng. Hand-book of Markov Chain Monte Carlo. CRC Press,2011.

Antonin Chambolle, Daniel Cremers, and ThomasPock. A convex approach to minimal partitions.SIAM Journal on Imaging Sciences, 5(4):1113–1158, 2012.

A. Chechetka and C. Guestrin. Focused belief propa-gation for query-specific inference. In Artificial In-telligence and Statistics (AISTATS), 2010.

H. Daume, J. Langford, and D. Marcu. Search-basedstructured prediction. Machine Learning, 75:297–325, 2009.

J.R. Doppa, A. Fern, and P. Tadepalli. Hc-search: Alearning framework for search-based structured pre-diction. Journal of Artificial Intelligence Research,50:403–439, 2014a.

J.R. Doppa, A. Fern, and P. Tadepalli. Structured pre-diction via output space search. Journal of MachineLearning Research, 15:1317–1350, 2014b.

J. Duchi, E. Hazan, and Y. Singer. Adaptive sub-gradient methods for online learning and stochas-tic optimization. In Conference on Learning Theory(COLT), 2010.

G. Elidan, I. McGraw, and D. Koller. Residual beliefpropagation: Informed scheduling for asynchronousmessage passing. In Uncertainty in Artificial Intel-ligence (UAI), 2006.

Y. Goldberg and M. Elhadad. An efficient algorithmfor easy-first non-directional dependency parsing. InAssociation for Computational Linguistics (ACL),pages 742–750, 2010.

Y. Goldberg and J. Nivre. Training deterministicparsers with non-deterministic oracles. Transac-tions of the Association for Computational Linguis-tics (TACL), 1, 2013.

S. Gould, R. Fulton, and D. Koller. Decomposing ascene into geometric and semantically consistent re-gions. In ICCV, 2009.

H. He, H. Daume, and J. Eisner. Dynamic fea-ture selection for dependency parsing. In EmpiricalMethods in Natural Language Processing (EMNLP),2013.

J. Jiang, A. Teichert, J. Eisner, and H. Daume.Learned prioritization for trading off accuracy andspeed. In Advances in Neural Information Process-ing Systems (NIPS), 2012.

J. H. Kappes, B. Andres, F. A. Hamprecht, C. Schnorr,S. Nowozin, D. Batra, S. Kim, B. X. Kausler, J. Lell-mann, N. Komodakis, and C. Rother. A compara-tive study of modern inference techniques for dis-crete energy minimization problems. In ComputerVision and Pattern Recognition (CVPR), 2013.

R.H. Kassel. A comparison of approaches to on-linehandwritten character recognition. PhD thesis, Mas-sachusetts Institute of Technology, 1995.

D. Koller, N. Friedman, L. Getoor, and B. Taskar.Graphical models in a nutshell. Statistical RelationalLearning, page 13, 2007.

P. Liang, H. Daume, and D. Klein. Structure compila-tion: Trading structure for features. In InternationalConference on Machine Learning (ICML), 2008.

A. McCallum and W. Li. for named entity recogni-tion with conditional random fields, feature induc-tion and web-enhanced lexicons. In Proceedings ofthe seventh conference on Natural language learningat HLT-NAACL 2003-Volume 4, pages 188–191. As-sociation for Computational Linguistics, 2003.

S. Ross, G. Gordon, and A. Bagnell. A reduction ofimitation learning and structured prediction to no-regret online learning. In Artificial Intelligence andStatistics (AISTATS), 2011a.

S. Ross, D. Munoz, M. Hebert, and J. A. Bagnell.Learning message-passing inference machines forstructured prediction. In Computer Vision and Pat-tern Recognition (CVPR), pages 2737–2744, 2011b.

G.A. Rummery and M. Niranjan. Online Q-learningusing connectionist systems. University of Cam-bridge, Department of Engineering, 1994.

R.S. Sutton and A.G. Barto. Introduction to reinforce-ment learning. MIT Press, 1998.

G. Tesauro. Temporal difference learning and td-gammon. Communications of the ACM, 38(3):58–68, 1995.

E. Veach. Robust Monte Carlo methods for light trans-port simulation. PhD thesis, Stanford University,1997.

Learning Where to Sample in Structured Prediction

P. Viola and M. Jones. Rapid object detection usinga boosted cascade of simple features. In ComputerVision and Pattern Recognition (CVPR), 2001.

C. Watkins and P. Dayan. Q-learning. Machine learn-ing, 8(3-4):279–292, 1992.

D. Weiss and B. Taskar. Structured prediction cas-cades. In Artificial Intelligence and Statistics (AIS-TATS), 2010.

D. Weiss and B. Taskar. Learning adaptive value ofinformation for structured prediction. In Advancesin Neural Information Processing Systems (NIPS),2013.

M. Wick, K. Rohanimanesh, S. Singh, and a. A. Mc-Callum. Training factor graphs with reinforcementlearning for efficient map inference. In Advancesin Neural Information Processing Systems (NIPS),2009.

M. Wick, K. Rohanimanesh, and K. Bellare. Sampler-ank: Training factor graphs with atomic gradients.In International Conference on Machine Learning(ICML), 2011.

M. L. Wick and A. McCallum. Query-aware MCMC.In Advances in Neural Information Processing Sys-tems (NIPS), pages 2564–2572, 2011.

Y. Zhang, T. Lei, R. Barzilay, and T. Jaakkola. Greedis good if randomized: New inference for depen-dency parsing. In Empirical Methods in NaturalLanguage Processing (EMNLP), 2014.