Lazy-k Decoding: Constrained Decoding for Information Extraction (2023)

Abstract: We explore the possibility of improving probabilistic models in structured prediction. Specifically, we combine the models with constrained decoding approaches in the context of token classification for information extraction. The decoding methods search for constraint-satisfying label-assignments while maximizing the total probability. To do this, we evaluate several existing approaches, as well as propose a novel decoding method called Lazy- k . Our findings demonstrate that constrained decoding approaches can significantly improve the models’ performances, especially when using smaller models. The Lazy- k approach allows for more flexibility between decoding time and accuracy. The code for using Lazy- k decoding can be found at https://github.com/ArthurDevNL/lazyk.

## introduction much of today's information extraction (ie) is done using probability-based token-classification models such as bert @xcite , roberta @xcite , layoutlm @xcite or lilt @xcite . these models aim for the best results by increasingly stacking large amounts of parameters, which comes at the cost of increased computational requirements and training complexity. typically, only the top-1 prediction is used, despite the fact that models produce probabilities for all tokenlabel combinations. ideally, alternative, high-likelihood predictions are explored to improve predictions from existing models. this is especially interesting in structuredprediction tasks, where the model's predictions are parsed into predefined structures. these structures allow for defining constraints that evaluate whether a produced prediction adheres to the expected structure, which can then be used to iterate over multiple high probability predictions until a satisfying solution is found. a concrete example of such a structure is in the case of invoice information extraction. in this task, the model is given the outputs of an optical character recognition (ocr) system and needs to predict which parts of the text correspond to the various elements in an invoice. for example, in fig. 1 , the model is expected to predict the total, cash and change amounts. however, the occlusion of the "cash" text introduces noise into the model's predictions, causing it to incorrectly label the cash amount as another total amount. using the arithmetic semantics of invoices, we know that the total amount should equal the cash amount paid minus the change amount. as such, we know that the model's best prediction is probably incorrect. alternative, high-probability labelassignments can be explored to find a constraintsatisfying solution instead. industrial document processing systems usually have programmatic post-processing logic that detects and sometimes corrects aforementioned semantic constraints. these systems, however, rarely exploit the remaining information "hidden" in the produced probability distributions, and the custom correction code is often complex and hard to maintain. furthermore there is the possibility of ocrinduced errors, which go beyond the scope of the present work but remain an important source of errors in document ie @xcite . in short, we exploit task-specific structures to explore alternative high-likelihood predictions. specifically, we ## background to search over high-probability predictions, we require a probabilistic model that outputs independent probabilities for a given sequence of tokens. given an input sequence @xmath0, x 2 , . . . , x n }, x @xmath1, the goal is to estimate the probability of the output sequence @xmath2, y 2 , . . . , y n }, y @xmath3. as this probability quickly becomes intractable, it is usually estimated by factoring it as: the decoding process refers to the way we obtain an estimate ŷ for y from such a model. the simplest approach consists of taking the arg max as ŷ = arg max in addition, we introduce a global, binary constraint c : x × @xmath4, 1} and formalize our problem of interest as note that for the method proposed in sec. 4, we make no further assumption about the constraint. this is important because many existing constrained decoding approaches require the constraints to be expressed in linear form @xcite . some problems may consist of both linear and non-linear constraints. token-classification models often use the bio labeling scheme @xcite , where the labels are prefixed with b(eginning), i(nside) and o(utside) to be able to classify spans of multiple tokens. in this formulation, an i label must always be preceded by a b or another i label of the same class. this labeling constraint can be expressed using linear constraints. however, the solution to the linear constraints is not guaranteed to also be a solution to the non-linear constraints. the semantic constraint @xmath5, the text corresponding to the labelization must be parsed from text to a float, which is a non-linear operation. an example where the optimal solution satisfying the linear (bio) constraints does not satisfy the non-linear (semantic) constraints is shown in tab. 1(b) one line 4. ## related work several decoding methods for the setting from eq. ( 3 ) have been proposed. an excellent benchmark for learning and decoding under constraints is provided in gluecons @xcite . for decoding, the work mostly explores the usage of integer linear programming (ilp) for finding a constraint-satisfying solution given the model's probabilities. ilp problems can be solved using the branch and bound algorithm @xcite , which is a type of informed search algorithm; it uses a linear formulation of the constraints to guide it more effectively through the search space. another example of a method that uses knowledge about its constraints is viterbi @xcite , which is a dynamic programming approach that can also take into account specific constraints, although more restrictive than ilp as it only works in a markovian setting. the advantage of these informed search methods is that they will always find the optimal answer within a reasonable amount of time, should it exist. however, they also have a non-negligible minimum running time and impose aforementioned requirements on the constraints. informed search methods have previously been applied to the task of information extraction. decoding under constraints using ilp was inspired by the work from roth and yih who explored the application to entity and relation extraction @xcite . they formulate the decoding problem as a linear program with the objective to maximize the overall probability of a sequence given a set of constraints. in these programs, the decision variables are indicator variables 1 j i , indicating the assignment of label j to token i. using this, one can express the constraint "1 label per token" as for more complex constraints we can use uninformed search methods as they do not make any assumptions on the implementation of the constraints and simply iterate over the search space in a greedy manner. the most widely known method for this is beam search (bs) @xcite . while it does not impose any restriction on the type of constraints, it is not ideal to our global decoding setting as it works in a "left-to-right" manner. to illustrate, bs takes as input a parameter k and outputs the top-k sequences by computing the top-k beams at every token, based on the previous top-k beams. in order to evaluate global constraints, beam search first needs to compute all top-k sequences after which the constraint can be evaluated. unfortunately, this means that if the constraintvalidating prediction ends up being the most likely (arg max) sequence, beam search will have computed k -1 too many sequences. in addition, if the constraint-validating prediction is not in the topk beams, a new search with an unknown, higher k ′ needs to be run, which also includes recomputing the initial previous k predictions. several adaptations have been suggested in the context of natural language generation @xcite @xcite , but none of which solve aforementioned problems for global constraints. others propose extending beam search with learnable heuristics that try to predict whether a given label "56" "." "000" b total 0.3 --i total -0.4 0.4 b cash 0.5 --i cash -0.3 0.3 (a) p "56" "." "000" bio sem. 8.0% b cash i total i total no -6.0% b cash i total i cash no -6.0% b cash i cash i total no -4.8% b total i total i total yes no 4.5% b cash i cash i cash yes yes 3.6% b total i total i cash no -. . . . . . . . . . . . . . . . . . (b) our method follows a similar approach to a × with partial expansion @xcite which has previously been applied to the multiple sequence alignment problem. to our knowledge, we are the first to apply a* with partial expansion that allows for more general constraints than ilp for the global constraint decoding setting. ## lazy-k decoding as the name suggests, the lazy-k decoder allows for decoding the k most probable sequences in a lazy manner. this means that it only iterates over the necessary number of sequences and stops once a satisfying solution is found. the hypothesis that this decoder explores is that the constraintsatisfying sequence is somewhere among the other high probability sequences. to do this efficiently, we exploit the fact that the k-th most probable sequence is always within "editdistance" 1 from one of the k -1 more probable sequences. this follows from the independence of each label as shown in eq. 1. we put "editdistance" in quotes here because we use a slightly more strict definition of edit-distance that also takes into account the order between the various label probabilities. more details about this can be found in app. a. at its core, it is a variant of best-first search (russel et al., 1994), where the model's predictions are used to determine the order in which the possible label assignments are explored. each state represents a full label assignment @xmath8, y 2 , . . . , y n } for all n tokens @xmath9, x 2 , . . . , x n }. the cost g(y) of a state is defined as follows: we use y k to denote the k-th lowest cost label assignment, and define the starting point y 1 as: the algorithm for lazy-k decoding is given in alg. 1. it works by maintaining a heap of the k best states, prioritized by the score of the next best unexplored state within 1 edit distance. the heap is initialized with the starting state y 1 . upon exploring a state, it is tested against the constraint and returns directly if it is satisfied. if the constraint is not satisfied, the heap is extended with the newly explored state and the priority score of the originating state y i is updated to reflect the score of the next best unexplored state. different from best-first search, upon exploring a state, we do not add all the children to the heap. instead, we only add the next best state y k and update the priority key for y i to be the score of the next best state within edit distance 1. this significantly reduces the size of the heap, as a classical search implementation adds n possible children at every iteration, whereas in this case, the number of states in the heap is at most equal to the number of iterations. this heap-size reduction in turn translates in better run time complexity as all following heap operations become cheaper. the nextbest function takes as input a state y and the frontier. the frontier is a dictionary that holds the explored states and next best states for all explored states. the values are integers that keep track of the i-th best change for a given state. if @xmath10. as the next best state may already exist in the frontier, the nextbest function is wrapped in addnextbest to make sure algorithm 1 lazy-k decoding require: input sequence: x require: cost function g from eq. 4 require: binary constraint c, max iterations k 1: function lazy-k(x, g, c, k) given that the algorithm iterates over the possible sequences in decreasing order of probability, it is trivial to prove that it will always find the optimal solution should it exist. in practice however, the combinatorial growth in the number of states quickly renders exhaustive search infeasible. to prevent this, an additional stopping condition is used where the iteration stops if no satisfying solution has been found after a fixed number of k iterations. one could also set the stopping condition according to a cumulative probability mass p or some other measure; we leave this exploration for future work. ## experiments to evaluate the relevance of the lazy-k decoder, we perform invoice information extraction on several datasets. the aim of the task is to extract various amounts from invoices such that they satisfy their expected arithmetic structure. for each dataset, we train a token-classification model and generate predictions for the test set. the predictions are then fed into the different decoding algorithms along with the constraints, and return the highest-probability sequence satisfying the constraints. data we evaluate the decoders on a total of three datasets shown in tab. 2: cord @xcite , wildreceipt @xcite and docile @xcite . while the constraints are semantically ## conclusion in summary, we have introduced a novel and efficient decoding method called lazy-k that allows for decoding under global, hard constraints. when applied in the context of invoice information extraction, lazy-k is faster than existing, greedy search methods and allows for more flexibility in trading off computing time and extraction performance compared to ilp. in addition, the possibility of using programmatic constraints directly makes lazyk an easy to use off-the-shelf solution for applying corrections to probabilistic models in the context of structured predictions. future work could explore the application to other structured-prediction problems with non1 scores for constrained decoding on smaller models. lazy-ilp is limited to 8 iterations and lazy-k to 2 14 . linear constraints besides information extraction. additionally, the improvements in extraction performance using the decoding methods are promising, which could also be explored in semisupervised learning settings. another interesting direction to explore would be the combination of lazy-k decoding with confidence calibration methods such as temperature scaling.