Efficient Computation of Expectations under Spanning Tree Distributions (2021)

Abstract: We give a general framework for inference in spanning tree models. We propose unified algorithms for the important cases of first-order expectations and second-order expectations in edge-factored, non-projective spanning-tree models. Our algorithms exploit a fundamental connection between gradients and expectations, which allows us to derive efficient algorithms. These algorithms are easy to implement with or without automatic differentiation software. We motivate the development of our framework with several cautionary tales of previous research, which has developed numerous inefficient algorithms for computing expectations and their gradients. We demonstrate how our framework efficiently computes several quantities with known algorithms, including the expected attachment score, entropy, and generalized expectation criteria. As a bonus, we give algorithms for quantities that are missing in the literature, including the KL divergence. In all cases, our approach matches the efficiency of existing algorithms and, in several cases, reduces the runtime complexity by a factor of the sentence length. We validate the implementation of our framework through runtime experiments. We find our algorithms are up to 15 and 9 times faster than previous algorithms for computing the Shannon entropy and the gradient of the generalized expectation objective, respectively.

## introduction dependency trees are a fundamental combinatorial structure in natural language processing. it follows that probability models over dependency trees are an important object of study. in terms equal contribution. the matrix-tree theorem (kirchhoff, 1847)more specifically, its counterpart for directed graphs @xcite -appeared before the nlp community in an onslaught of contemporaneous papers @xcite @xcite ) that leverage the classic result to efficiently compute the normalization constant of a distribution over trees. the result is still used in more recent work @xcite . we build upon this tradition through a framework for computing expectations of a rich family of functions under a distribution over trees. expectations appear in all aspects of the probabilistic modeling process: training, model validation, and prediction. therefore, developing such a framework is key to accelerating progress in probabilistic modeling of trees. our framework is motivated by the lack of a unified approach for computing expectations over spanning trees in the literature. we believe this gap has resulted in the publication of numerous inefficient algorithms. we motivate the importance of developing such a framework by highlighting the following cautionary tales. our framework recovers @xcite algorithm, and additionally provides the gradient of entropy in o n 3 . we summarize our main results below: we demonstrate how our framework calculates several new quantities, such as the kullback-leibler divergence, which (to our knowledge) had no prior algorithm in the literature. ## distributions over trees we consider the distribution over trees in weighted directed graphs with a designated root node. a (rooted, weighted, and directed) graph is given by @xmath0, e, ρ). @xmath1, . . . , n} ∪ {ρ} is a set of n +1 nodes where ρ is a designated root node. e is a set of weighted edges where each edge (i in natural language processing applications, these weights are typically parametric functions, such as log-linear models @xcite or neural networks @xcite , which are learned from data. a tree foot_0 d of a graph g is a set of n edges such that all non-root nodes j have exactly one incoming edge and the root node ρ has at least one outgoing edge. furthermore, a tree does not contain any cycles. we denote the set of all trees in a graph by d and assume that |d| > 0 (this is not necessarily true for all graphs). the weight of a tree @xmath2: normalizing the weight of each tree yields a probability distribution: of course, for (2) to be a proper distribution, we require w @xmath3, and @xmath4. ## expectations in this section, we characterize the family of expectations that our framework supports. our framework is an extension of @xcite to distributions over spanning trees. in contrast, their framework considers expectations over distributions that can be factored as b-hypergraphs @xcite . our distributions over trees cannot be cast as polynomial-size b-hypergraphs. another important distinction between our framework and that of @xcite is that we do not use the semiring abstraction as it is algebraically too weak to compute the determinant efficiently. 7 without any assumptions on f , computing (9) is intractable. 8 in the remainder of this section, we will characterize a class of functions f whose expectations can be efficiently computed. the first type of functions we consider are functions that are additively decomposable along the edges of the tree. formally, a function r : @xmath5 the second type of functions we consider are functions that are second-order additively decomposable along the edges of the tree. formally, a function r: @xmath6, r : @xmath7: thus, t(d) ∈ r r×s is generally a matrix. an example of such a function is the gradient of entropy (see §6.2) or the ge objective @xcite (see §6.4 with respect to the edge weights. another example of a secondorder additively decomposable function is the single-root and multi-root settings @xcite . 8 of course, one could use sampling methods, such as monte carlo, to approximate (9). sampling methods may be efficient if the variance of f under p is not too large. thus, it is second-order additively decomposable function as long as r(d) and s(d) are additively decomposable. one family of functions which can be computed efficiently but we will not explore here are those who are multiplicatively decomposable over the edges. a function q : @xmath8 ## connecting gradients and expectations in this section, we build upon a fundamental connection between gradients and expectations @xcite . this connection allows us to build on work in automatic differentiation to obtain efficient gradient algorithms. while the propositions in this section are inspired from past work, we believe that the presentation and proofs of these propositions have previously not been clearly presented. 11 we find it convenient to work with unnormalized expectations, or totals (for short). we denote the total of a function f as f @xmath9. we recover the expectation with e p [f ] = f/z. we note that totals (on their own) may be of 10 the ℓ k norm of the distribution p often denoted as it is computable from a zero th -order expectation because it can be written as @xcite .1) provide a similar derivation to proposition 3 and proposition 4 for hypergraphs. the first-order case. specifically, the partial derivative ∂z ∂w ij is useful for determining the total weight of trees which include the edge (@xmath10, proposition'4 will establish a connection between the unnormalized expectation r and ∇z. proposition 4. for any additively decomposable function r: @xmath11, the total r can be computed using a gradient-vector product @xmath12. we can similarly use proposition 5. for any pair of edges @xmath13, proposition 6 will relate ∇ 2 z to ∇r. this will be used in proposition 7 to establish a connection between the total t and ∇ 2 z, and additionally establishes a connection between t and ∇r. 13 as each edge can only appear once in a tree, w ij, @xmath14. proposition 6. for any additively decomposable function r : @xmath15, 14 and edge @xmath16, proposition 7. for any second-order additively decomposable function t: @xmath17, which is expressed as the outer product of additively decomposable functions, r : @xmath18: @xmath19, t(d) = r(d)s(d) , where r does not depend on w, the total t can be computed using a jacobian-matrix product proof. we first prove (19 then (20) immediately follows by substituting (18) into ( 19 ) and expanding the summation. remark. there is a simple recipe to compute ∇r n for each @xmath20, . . . , r. first, some notation; let -→ 1 ij be a vector over e with a 1 in dimension (@xmath21, and zeros elsewhere. by plugging [r ij ] n and s @xmath22, we can compute t @xmath23. 15 however, if r depends on w, we must add the following first-order term, which is due to the product rule we provide the details for computing the gradients of two first-order quantities, shannon entropy and the kl divergence, using this recipe in §6.2 and §6.3, respectively. ## algorithms having reduced the computation of r and t to finding derivatives of z in §4, we now describe efficient algorithms that exploit this connection. the main algorithmic ideas used in this section are based on automatic differentiation (ad) techniques @xcite . these are general-purpose techniques for efficiently evaluating gradients given algorithms that evaluate the functions. in our setting, the algorithm in question is an efficient procedure for evaluating z, such as the procedure we described in §2.1. while we provide derivatives §5.1 in our algorithms, these can also be evaluated using any ad library, such as jax @xcite , pytorch (our choice) (paszke et al., 2019), or tensorflow @xcite . proposition 4 is realized as t 1 in fig. 2 and ( 19 ) and ( 20 ) are realized as t v 2 and t h 2 in fig. 3 , respectively. we provide the runtime complexity of each step in the algorithms. these will be discussed in more detail in §5.2. ## applications and prior work in this section, we apply our framework to compute a number of important quantities that are used when working with probabilistic models. we relate our approach to existing algorithms in the literature (where applicable), and mention existing and potential applications. many of our quantities were covered in @xcite for b-hypergraphs; we extend their results to spanning trees. in most applications that involve training a probabilistic model, the edge weights in the model will be parameterized in some fashion. traditional approaches @xcite @xcite use log-linear parameterizations, whereas more recent work @xcite @xcite use neural-network parameterizations. our algorithms are agnostic as to how edges are parameterized. ## conclusion we presented a general framework for computing first-and second-order expectations for additively decomposable functions. we did this by exploiting a key connection between gradients and expectations that allows us to solve our problems using automatic differentiation. the algorithms we provide are simple, efficient, and extendable to many expectations. the automatic differentiation principle has been applied in other settings, such as weighted context-free grammars @xcite and chain-structured models @xcite . we hope that this paper will also serve as a tutorial on how to compute expectations over trees so that the list of cautionary tales does not grow further. particularly, we hope that this will allow for the kl divergence to be used in semi-supervised training of dependency parsers. our aim is for our approach for computing expectations to be extended to other structured prediction models.