Torch-Struct: Deep Structured Prediction Library

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 this https URL.

Structured prediction is an area of machine learning focusing on representations of spaces with combinatorial structure, and 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 BIBREF0, BIBREF4, discriminative dependency and constituency parsing BIBREF10, BIBREF8, unsupervised learning for labeling and alignment BIBREF11, BIBREF12, approximate translation decoding with beam search BIBREF9, 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 final layer for deep models BIBREF13, BIBREF14. Another has incorporated structured prediction within deep learning models, exploring novel models for latent-structure learning, unsupervised learning, or model control BIBREF15, BIBREF16, BIBREF17. We aspire to make both of these use-cases as easy to use as standard neural networks.
The practical challenge of employing structured prediction is that many required algorithms are difficult to implement efficiently and correctly. Most projects reimplement custom versions of standard algorithms or focus particularly on a single well-defined model class. This research style makes it difficult to combine and try out new approaches, a problem that has compounded with the complexity of research in deep structured prediction.
With this challenge in mind, we introduce Torch-Struct with three specific contributions:
Modularity: models are represented as distributions with a standard flexible API integrated into a deep learning framework.
Completeness: a broad array of classical algorithms are implemented and new models can easily be added in Python.
Efficiency: implementations target computational/memory efficiency for GPUs and the backend includes extensions for optimization.
In this system description, we first motivate the approach taken by the library, then present a technical description of the methods used, and finally present several example use cases.
Several software libraries target structured prediction. Optimization tools, such as SVM-struct BIBREF18, focus on parameter estimation. Model libraries, such as CRFSuite BIBREF19 or CRF++ BIBREF20, implement inference for a fixed set of popular models, such as linear-chain CRFs. General-purpose inference libraries, such as PyStruct BIBREF21 or TurboParser BIBREF22, 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 BIBREF23, allow for specification and inference over some discrete domains. Most ambitiously, inference libraries such as Dyna BIBREF24 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.
While structured prediction is traditionally presented at the output layer, recent applications have deployed structured models broadly within neural networks BIBREF15, BIBREF25, BIBREF16. Torch-Struct aims to encourage this general use case.
To illustrate, we consider a latent tree model. ListOps BIBREF26 is a dataset of mathematical functions. Each data point 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)$ BIBREF16, BIBREF27.
A popular approach is a latent-tree RL model which we briefly summarize. The objective is to maximize the probability of the correct prediction under the expectation of a prior tree model, $p(z|x ;\phi )$,
Computing the expectation is intractable so policy gradient is used. First a tree is sampled $\tilde{z} \sim p(z | x;\phi )$, then the gradient with respect to $\phi $ is approximated as,
where $b$ is a variance reduction baseline. A common choice is the self-critical baseline BIBREF28,
Finally an entropy regularization term is added to the objective encourage exploration of different trees, $ O + \lambda \mathbb {H}(p(z\ |\ x;\phi ))$.
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 tree model $p(z\ | x; \phi )$:
[description]font=
[itemsep=-2pt]
Policy gradient, $\tilde{z} \sim p(z \ |\ x ; \phi )$
Score policy samples, $p(z \ | \ x; \phi )$
Backpropagation, $\frac{\partial }{\partial \phi } p(z\ |\ x; \phi )$
Self-critical, $\arg \max _z p(z \ |\ x;\phi )$
Objective regularizer, $\mathbb {H}(p(z\ |\ x;\phi ))$
For structured models, each of these terms is non-trivial to compute. A goal of Torch-Struct is to make it seamless to deploy structured models for these complex settings. To demonstrate this, Torch-Struct includes an implementation of this latent-tree approach. With a minimal amount of user code, the implementation achieves near perfect accuracy on the ListOps dataset.
The library design of Torch-Struct follows the distributions API used by both TensorFlow and PyTorch BIBREF29. 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 (scores) output from a deep network $\ell $, the user can request samples $z \sim \textsc {CRF}(\ell )$, probabilities $\textsc {CRF}(z;\ell )$, modes $\arg \max _z \textsc {CRF}(\ell )$, or other distributional properties such as $\mathbb {H}(\textsc {CRF}(\ell ))$. 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 FIGREF11 demonstrates this API for a binary tree CRF over an ordered sequence, such as $p(z \ | \ y ;\phi )$ from the previous section. The distribution takes in log-potentials $\ell $ 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.
Table TABREF2 shows all of the structures and distributions implemented in Torch-Struct. While each is internally implemented using different specialized algorithms and optimizations, from the user's perspective they all utilize the same external distributional API, and pass a generic set of distributional tests. This approach hides the internal complexity of the inference procedure, while giving the user full access to the model.
We now describe the technical approach underlying the library. To establish notation first consider the implementation of a categorical distribution, Cat($\ell $), with one-hot categories $z$ with $z_i = 1$ from a set $\cal Z$ and probabilities given by the softmax,
Define the log-partition as $A(\ell ) = \mathrm {LSE}(\ell )$, i.e. log of the denominator, where $\mathrm {LSE}$ is the log-sum-exp operator. Computing probabilities or sampling from this distribution, requires enumerating $\cal Z$ to compute the log-partition $A$. A useful identity is that derivatives of $A$ yield category probabilities,
Other distributional properties can be similarly extracted from variants of the log-partition. For instance, define $A^*(\ell ) = \log \max _{j=1}^K \exp \ell _j$ then: $\mathbb {I}(z^*_i = 1) = \frac{\partial }{\partial \ell _i} A^*(\ell ) $.
Conditional random fields, CRF($\ell $), extend the softmax to combinatorial spaces where ${\cal Z}$ is exponentially sized. Each $z$, is now represented as a binary vector over polynomial-sized set of parts, $\cal P$, i.e. ${\cal Z} \subset \lbrace 0, 1\rbrace ^{|\cal P|}$. Similarly log-potentials are now defined over parts $\ell \in \mathbb {R}^{|\cal P|}$. For instance, in Figure FIGREF11 each span is a part and the $\ell $ vector is shown in the top-left figure. Define the probability of a structure $z$ as,
Computing probabilities or sampling from this distribution, requires computing the log-partition term $A$. In general computing this term is now intractable, however for many core algorithms in NLP there are exist efficient combinatorial algorithms for this term (as enumerated in Table TABREF2).
Derivatives of the log-partition again provide distributional properties. For instance, the marginal probabilities of parts are given by,
Similarly derivatives of $A^*$ correspond to whether a part appears in the argmax structure. $\mathbb {I}(z^*_p = 1) = \frac{\partial }{\partial \ell _p} A^*(\ell ) $.
While these gradient identities are well-known BIBREF30, they are not commonly deployed. Computing CRF properties is typically done through two-step specialized algorithms, such as forward-backward, inside-outside, or similar variants such as viterbi-backpointers BIBREF31. In our experiments, we found that using these identities with auto-differentiation on GPU was often faster, and much simpler, than custom two-pass approaches. Torch-Struct is thus designed around using gradients for distributional computations.
Torch-Struct is a collection of generic algorithms for CRF inference. Each CRF distribution object, $\textsc {CRF}(\ell )$, is constructed by providing $\ell \in \mathbb {R}^{|{\cal P}|}$ where the parts $\cal P$ are specific to the type of distribution. Internally, each distribution is implemented through a single Python function for computing the log-partition function $A(\ell )$. From this function, the library uses auto-differentiation and the identities from the previous section, to define a complete distribution object. The core models implemented by the library are shown in Table TABREF2.
To make the approach concrete, we consider the example of a linear-chain CRF.
latent](a)$z_1$; latent, right = of a](b)$z_2$; latent, right = of b](c)$z_3$; (a) – (b) – (c);
The model has $C$ labels per node with a length $T=2$ edges utilizing a first-order linear-chain (Markov) model. This model has $2\times C \times C$ parts corresponding to edges in the chain, and thus requires $\ell \in \mathbb {R}^{2\times C \times C}$. The log-partition function $A(\ell )$ factors into two reduce computations,
Computing this function left-to-right using dynamic programming yield the standard forward algorithm for sequence models. As we have seen, the gradient with respect to $\ell $ produces marginals for each part, i.e. the probability of a specific labeled edge.
We can further extend the same function to support generic semiring dynamic programming BIBREF34. A semiring is defined by a pair $(\oplus , \otimes )$ with commutative $\oplus $, distribution, and appropriate identities. The log-partition utilizes $\oplus , \otimes = \mathrm {LSE}, +$, but we can substitute alternatives.
For instance, utilizing the log-max semiring $(\max , +)$ in the forward algorithm yields the max score. As we have seen, its gradient with respect to $\ell $ is the argmax sequence, negating the need for a separate argmax (Viterbi) algorithm. Some distributional properties cannot be computed directly through gradient identities but still use a forward-backward style compute structure. For instance, sampling requires first computing the log-partition term and then sampling each part, (forward filtering / backward sampling). We can compute this value by overriding each backpropagation operation for the $\bigoplus $ to instead compute a sample.
Table TABREF16 shows the set of semirings and backpropagation steps for computing different terms of interest. We note that many of the terms necessary in the case-study can be computed with variant semirings, negating the need for specialized algorithms.
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. Figure FIGREF17 shows the impact of these optimizations on the core algorithms.
The commutative properties of semiring algorithms allow flexibility in the order in which we compute $A(\ell )$. 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 BIBREF35, typically used for computing prefix sums. To compute, $A(\ell )$ in this manner we first pad the sequence length $T$ out to the nearest power of two, and then compute a balanced parallel tree over the parts, shown in Figure FIGREF21. Concretely each node layer would compute a semiring matrix multiplication, e.g. $ \bigoplus _c \ell _{t, \cdot , c} \otimes \ell _{t^{\prime }, c, \cdot }$. Under this approach, we only need $O(\log N)$ steps in Python and can use parallel GPU operations for the rest. Similar parallel approach can also be used for computing sequence alignment and semi-Markov models.
Computational complexity is even more of an issue for parsing algorithms, which cannot be as easily parallelized. The log-partition for parsing is computed with the Inside algorithm. This algorithm must compute each width from 1 through T in serial; however it is important to parallelize each inner step. 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 reindex the chart. Instead of using a single chart $C$, we split it into two parts: one right-facing $C_r[i, d] = C[i, i+d]$ and one left facing, $C_l[i+d, T-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[\cdot , d]$ in one operation. Variants of this same approach can be used for all the parsing models employed.
The two previous optimizations reduce most of the cost to semiring matrix multiplication. In the specific case of the $(\sum , \times )$ semiring these can be computed very efficiently using matrix multiplication, which is highly-tuned on GPU hardware. Unfortunately for other semirings, such as log and max, these operations are either slow or very memory inefficient. For instance, for matrices $T$ and $U$ of sized $N \times M$ and $M \times O$, we can broadcast with $\otimes $ to a tensor of size $N \times M \times O$ and then reduce dim $M$ by $\bigoplus $ at a huge memory cost. To avoid this issue, we implement custom CUDA kernels targeting fast and memory efficient tensor operations. For log, this corresponds to computing,
where $q = \max _n T_{m,n} + U_{n, o}$. To optimize this operation on GPU we utilize the TVM language BIBREF36 to layout the CUDA loops and tune it to hardware.
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. 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.
We thank Yoon Kim, Xiang Lisa Li, Sebastian Gehrmann, Yuntian Deng, and Justin Chiu for discussion and feedback on the project. The project was supported by NSF CAREER 1845664, NSF 1901030, and research awards by Sony and AWS.

Figure and table captions:
Figure 1: Distribution of binary trees over an 1000- token sequence. Coloring shows the marginal probabilities of every span. Torch-Struct is an optimized collection of common CRF distributions used in NLP designed to integrate with deep learning frameworks.
Table 1: Models and algorithms implemented in Torch-Struct. Notation is developed in Section 5. Parts are described in terms of sequence lengths N,M , label size C, segment length K, and layers / grammar size L,G. Lines of code (LoC) is from the log-partition (A(`)) implementation. T/S is the tokens per second of a batched computation, computed with batch 32, N = 25, C = 20,K = 5, L = 3 (K80 GPU run on Google Colab).
Figure 2: Latent Tree CRF example. (a) Logpotentials ` for each part/span. (b) Marginals for CRF(`) computed by backpropagation. (c) Mode tree argmaxz CRF(z; `). (d) Sampled tree z ∼ CRF(`).
Table 2: (Top) Semirings implemented in Torch-Struct. Backprop/Gradients gives overridden backpropagation computation and value computed by this combination. (Bot) Example of gradients from different semirings on sequence alignment with dynamic time warping.
Figure 3: Speed impact of optimizations. Time is given in seconds for 10 runs with batch 16 executed on Google Colab. (a) Speed of a linear-chain forward with 20 classes for lengths up to 500. Compares left-to-right ordering to parallel scan. (b) Speed of CKY inside with lengths up to 80. Compares inner loop versus vectorization. (c) Speed of linear-chain forward of length 20 with up to 100 classes. Compares broadcast-reduction versus CUDA semiring kernel. (Baseline memory is exhausted after 100 classes.)
Figure 4: Parallel scan implementation of the linearchain CRF inference algorithm. Here ⊕ ⊗ represents a semiring matrix operation and I is padding.

Questions about this paper:
Is this library implemented into Torch or is framework agnostic?

Answers:
It uses deep learning framework (pytorch)