Torch-Struct: Deep Structured Prediction Library (2020)

Abstract: The literature on structured prediction for NLP describes a rich collection of distributions and algorithms over sequences, segmentations, alignments, and trees; however, these algorithms are difficult to utilize in deep learning frameworks. We introduce Torch-Struct, a library for structured prediction designed to take advantage of and integrate with vectorized, auto-differentiation based frameworks. Torch-Struct includes a broad collection of probabilistic structures accessed through a simple and flexible distribution-based API that connects to any deep learning model. The library utilizes batched, vectorized operations and exploits auto-differentiation to produce readable, fast, and testable code. Internally, we also include a number of general-purpose optimizations to provide cross-algorithm efficiency. Experiments show significant performance gains over fast baselines and case-studies demonstrate the benefits of the library. Torch-Struct is available at https://github.com/harvardnlp/pytorch-struct .

## introduction structured prediction is an area of machine learning focusing on representations of spaces with combinatorial structure, as well as algorithms for inference and parameter estimation over these structures. core methods include both tractable exact approaches like dynamic programming and spanning tree algorithms as well as heuristic techniques such linear programming relaxations and greedy search. structured prediction has played a key role in the history of natural language processing. example methods include techniques for sequence labeling and segmentation @xcite , discriminative dependency and constituency parsing @xcite , unsupervised learning for labeling and alignment @xcite , approximate translation decoding with beam search @xcite , among many others. in recent years, research into deep structured prediction has studied how these approaches can be integrated with neural networks and pretrained models. one line of work has utilized structured prediction as the final layer for deep models @xcite . another has incorporated structured prediction within deep learning models, exploring novel models for latentstructure learning, unsupervised learning, or model control @xcite @xcite . we aspire to make both of these use-cases as easy to use as standard neural networks. name structure (z) parts (p) algorithm (a(ℓ)) loc t/s sample reference linear-chain, hmm labeled chain edges (n c 2 ) forward-backward 20 390k (lafferty et al., 2001) factorial-hmm labeled chains trans. (lc 2 ) obs. (n c l ) factorial f-b 20 25k (ghahramani and jordan, 1996) alignment alignment match (n m ) skips (2n m ) dtw, ctc 50 13k (needleman and wunsch, 1970) semi-markov seg. labels edge(n kc 2 ) segmental f-b 30 87k (baum and petrie, 1966) (sarawagi and cohen, 2005) context-free labeled tree cf rules (g) term. (cn ) i-o cky 70 37k @xcite simple cky labeled tree splits (cn 2 ) 0-th order cky 30 118k @xcite dependency proj. tree arcs (n 2 ) eisner alg 40 28k @xcite dep (np) non-proj. tree arcs (n 2 ) matrix-tree chiu-liu (map) 40 1.1m @xcite ) auto-regressive sequence prefix (c n ) greedy search, beam search 60 - @xcite ## related work several software libraries target structured prediction. optimization tools, such as svm-struct @xcite , focus on parameter estimation. model libraries, such as crfsuite @xcite , crf++ @xcite , or ncrf++ @xcite , implement inference for a fixed set of popular models, usually linear-chain crfs. general-purpose inference libraries, such as pystruct (müller and behnke, 2014) or tur-boparser @xcite , utilize external solvers for (primarily map) inference such as integer linear programming solvers and admm. probabilistic programming languages, for example languages that integrate with deep learning such as pyro @xcite , allow for specification and inference over some discrete domains. most ambitiously, inference libraries such as dyna @xcite allow for declarative specifications of dynamic programming algorithms to support inference for generic algorithms. torch-struct takes a different approach and integrates a library of optimized structured distributions into a vectorized deep learning system. we begin by motivating this approach with a case study. ## motivating case study while structured prediction is traditionally presented at the output layer, recent applications have deployed structured models broadly within neural networks @xcite @xcite . torch-struct aims to encourage this general use case. to illustrate, we consider a latent tree model. listops @xcite is a dataset of mathematical functions. each input/output pair consists of a prefix expression x and its result y, e.g. models such as a flat rnn will fail to capture the hierarchical structure of this task. however, if a model can induce an explicit latent z, the parse tree of the expression, then the task is easy to learn by a tree-rnn model p(y|x, z) @xcite . let us briefly summarize a latent-tree rl model for this task. the objective is to maximize the probability of the correct prediction under the expectation of a prior tree model, p(z|x; φ), computing the expectation is intractable so policy gradient is used. first a tree is sampled z ∼ p(z|x; φ), then the gradient with respect to φ is approximated as, finally an entropy regularization term is added to the objective encourage exploration of different trees, obj + λh(p(z | x; φ)). even in this brief overview, we can see how complex a latent structured learning problem can be. to compute these terms, we need 5 different properties of the structured prior model p(z |x; φ): for structured models, each of these terms is nontrivial to compute. a goal of torch-struct is to make it seamless to deploy structured models for these complex settings. ## library design the library design of torch-struct follows the distributions api used by both tensorflow and py-torch @xcite . for each structured model in the library, we define a conditional random field (crf) distribution object. from a user's standpoint, this object provides all necessary distributional properties. given log-potentials ℓ output from a deep network, the user can request samples z ∼ crf(ℓ), probabilities crf(z; ℓ), modes arg max z crf(ℓ), or other distributional properties such as h(crf(ℓ)). the library is agnostic to how these are utilized, and when possible, they allow for backpropagation to update the input network. the same distributional object can be used for standard output prediction as for more complex operations like attention or reinforcement learning. figure 2 demonstrates this api for a binary tree crf over an ordered sequence, such as p(z | x; φ) from the previous section. the distribution takes in log-potentials ℓ which score each possible span in the input. the distribution converts these to probabilities of a specific tree. this distribution can be queried for predicting over the set of trees, sampling a tree for model structure, or even computing entropy over all trees. ## optimizations torch-struct aims for computational and memory efficiency. implemented naively, dynamic programming algorithms in python are prohibitively slow. as such torch-struct provides key primitives to help batch and vectorize these algorithms to take advantage of gpu computation and to minimize the overhead of backpropagating through chart-based dynamic programmming. we discuss three optimizations: a) parallel scan, b) vectorization, and c) semiring matrix multiplications. figure 3 shows the impact of these optimizations on the core algorithms. parallel scan inference the commutative properties of semiring algorithms allow flexibility in the order in which we compute a(ℓ). typical implementations of dynamic programming algorithms are serial in the length of the sequence. on parallel hardware, an appealing approach is a parallel scan ordering @xcite , typically used for computing prefix sums. to compute, a(ℓ) in this manner we first pad the sequence length n out to the nearest power of two, and then compute a balanced parallel tree over the parts, shown in figure 4 . concretely each node layer would compute a semiring matrix multiplication, e.g. c ℓ n,•,c ⊗ ℓ n+1,c,• . under this approach, assuming enough parallel cores, we only need o(log n ) steps in python and can use parallel operations for the rest. similar parallel approach can also be used for computing sequence alignment and semi-markov models. vectorization computational complexity is even more of an issue for algorithms that cannot easily be parallelized. for example, parsing algorithms the generalize cky are common in nlp. the cky algorithm has a bottleneck that it must compute each width from 1 through n in serial; however internally each one of these steps can be vectorized. assuming we have computed all inside spans of width less than d, computing the inside span of width d requires computing for all i, in order to vectorize this loop over i, j, we need to reindex the chart. instead of using a single chart c, we split it into two parts: one rightfacing c r [i, d] = c[i, i + d] and one left facing, c l [i+d, n -d] = c[i, i+d]. after this reindexing, the update can be written. unlike the original, this formula can easily be computed as a vectorized semiring dot product. this allows use to compute c r [•, d] in one operation. variants of this same approach can be used for many more complex dynamic programs. ## conclusion and future work we present torch-struct, a library for deep structured prediction. the library achieves modularity through its adoption of a generic distributional api, completeness by utilizing crfs and semirings to make it easy to add new algorithms, and efficiency through core optimizations to vectorize important dynamic programming steps. in addition to the problems discussed so far, torch-struct also includes several other example implementations including supervised dependency parsing with bert, unsupervised tagging, structured attention, and connectionist temporal classification (ctc) for speech. code demonstrates that the model is able to replicate standard deep learning results, although we focus here on the fidelity and implementation approach of the core library. the full library is available at https: //github.com/harvardnlp/pytorch-struct. in the future, we hope to support research and production applications employing structured models. we also believe the library provides a strong foundation for building generic tools for interpretablity, control, and visualization through its probabilistic api. finally, we hope to explore further optimizations to make core algorithms competitive with highly-optimized neural network components. these approaches provide a benchmark for improving autodifferentiation systems and extending their functionality to higher-order properties.