S yn J ax: Structured Probability Distributions for JAX (2023)

Abstract: The development of deep learning software libraries enabled significant progress in the field by allowing users to focus on modeling, while letting the library to take care of the tedious and time-consuming task of optimizing execution for modern hardware accelerators. However, this has benefited only particular types of deep learning models, such as Transformers, whose primitives map easily to the vectorized computation. The models that explicitly account for structured objects, such as trees and segmentations, did not benefit equally because they require custom algorithms that are difficult to implement in a vectorized form. SynJax directly addresses this problem by providing an efficient vectorized implementation of inference algorithms for structured distributions covering alignment, tagging, segmentation, constituency trees and spanning trees. This is done by exploiting the connection between algorithms for automatic differentiation and probabilistic inference. With SynJax we can build large-scale differentiable models that explicitly model structure in the data. The code is available at https://github.com/google-deepmind/synjax

## introduction in many domains, data can be seen as having some structure explaining how its parts fit into a larger whole. this structure is often latent, and it varies depending on the task. for examples of discrete structures in natural language consider figure 1 . the words together form a sequence. each word in a sequence is assigned a part-of-speech tag. these tags are dependent on each other, forming a linearchain marked in red. the words in the sentence can be grouped together into small disjoint contiguous groups by sentence segmentation, shown with bubbles. a deeper analysis of language would show that the groupings can be done recursively and thereby produce a syntactic tree structure. structures can also relate two languages. for instance, in the same figure, a japanese translation can be mapped to an english source by an alignment. these structures are not specific to language. similar structures appear in biology as well. nucleotides of any two rna sequences are matched with monotone alignment @xcite , genomic data is segmented into contiguous groups @xcite and tree-based models of rna capture the hierarchical nature of the protein folding process @xcite @xcite . most contemporary deep learning models attempt to predict output variables directly from the input without any explicit modeling of the intermediate structure. modeling structure explicitly could improve these models in multiple ways. first, it could allow for better generalization trough the right inductive biases @xcite . this would improve not only sample efficiency but also downstream performance @xcite @xcite . explicit modeling of structure can also enable incorporation of problem specific algorithms (e.g. finding shortest paths; @xcite or constraints (e.g. enforcing alignment @xcite or enforcing compositional calculation @xcite . discrete structure also allows for better interpretability of the model's decisions @xcite . finally, sometimes structure is the end goal of learning itself -for example we may know that there is a hidden structure of a particular form explaining the data, but its specifics are not known and need to be discovered @xcite . auto-regressive models are the main approach used for modeling sequences. non-sequential structures are sometimes linearized and approximated with a sequential structure @xcite . these models are powerful as they do not make any independence assumptions and can be trained on large amounts of data. while sampling from auto-regressive models is typically tractable, other common inference problems like finding the optimal structure or marginalizing over hidden variables are not tractable. approximately solving these tasks with auto-regressive models requires using biased or high-variance approximations that are often computationally expensive, making them difficult to deploy in large-scale models. alternative to auto-regressive models are models over factor graphs that factorize in the same way as the target structure. these models can efficiently compute all inference problems of interest exactly by using specialized algorithms. despite the fact that each structure needs a different algorithm, we do not need a specialized algorithm for each inference task (argmax, sampling, marginals, entropy etc.). as we will show later, synjax uses automatic differentiation to derive many quantities from just a single function per structure type. large-scale deep learning has been enabled by easy to use libraries that run on hardware accelerators. research into structured distributions for deep learning has been held back by the lack of ergonomic libraries that would provide acceleratorfriendly implementations of structure componentsespecially since these components depend on algorithms that often do not map directly onto available deep learning primitives, unlike transformer models. this is the problem that synjax addresses by providing easy to use structure primitives that compose within jax machine learning framework. to see how easy it is to use synjax consider example in figure 2 . this code implements a policy gradient loss that requires computing multiple quantities -sampling, argmax, entropy, logprobability -each requiring a different algorithm. in this concrete code snippet, the structure is a nonprojective directed spanning tree with a single root if the user wants only to change slightly the the tree requirements to follow the projectivity constraint they only need to change one flag and synjax will in the background use completely different algorithms that are appropriate for that structure: it will use @xcite for argmax and variations of @xcite algorithm for other quantities. the user does not need to implement any of those algorithms or even be aware of their specifics, and can focus on the modeling side of the problem. ## structured distributions distributions over most structures can be expressed with factor graphs -bipartite graphs that have random variables and factors between them. we associate to each factor a non-negative scalar, called potential, for each possible assignment of the random variables that are in its neighbourhood. the potential of the structure is a product of its factors: ## computing probability of a structure and partition function equation 2 shows the definition of the probability of a structure in a factor graph. computing the numerator is often trivial. however, computing the denominator, the partition function, is the complicated and computationally demanding part because the set of valid structures t is usually exponentially large and require specialized algorithms for each type of structure. as we will see later, the algorithm for implementing the partition function accounts for the majority of the code needed to add support for a structured distribution, as most of the other properties can be derived from it. here we document the algorithms for each structure. ## computing marginals in many cases we would like to know the probability of a particular part of structure appearing, regardless of the structure that contains it. in other words, we want to marginalize (i.e. sum) the probability of all the structures that contain that part: computing these factors was usually done using specialized algorithms such as inside-outside or forward-backward. however, those solutions do not work on distributions that cannot use belief propagation like non-projective spanning trees. a more general solution is to use an identity that relates gradients of factor's potentials with respect to the log-partition function: this means that we can use any differentiable implementation of log-partition function as a forward pass and apply backpropagation to compute the marginal probability @xcite . eisner (2016) has made an explicit connection that "inside-outside and forward-backward algorithms are just backprop". this approach also works for non-projective spanning trees that do not fit belief propagation framework @xcite . for template models like pcfg, we use again the sticky log-potentials because usually we are not interested in marginal probability of the rules but in the marginal probability of the instantiated constituents. the derivative of log-partition with respect to the constituent's sticky log-potential will give us marginal probability of that constituent. ## computing most probable structure for finding the score of the highest scoring structure we can just run the same belief propagation algorithm for log-partition, but with the max-plus semiring instead of the log-plus semiring (goodman, 1999). to get the most probable structure, and not just its potential, we can compute the gradient of part potentials with respect to the viterbi structure potential @xcite . the only exceptions to this process are nonmonotone alignments and spanning trees because they do fit easily in belief propagation framework. for the highest scoring non-monotone alignment, we use the jonker-volgenant algorithm as implemented in scipy @xcite . maximal projective spanning tree can be found by combining eisner's algorithm with maxplus semiring, but we have found kuhlmann's tabulated arc-hybrid algorithm to be much faster @xcite (see figure 4 in the appendix). this algorithm cannot be used for any inference task other than argmax because it allows for spurious derivations. to enforce singleroot constraint with kuhlmann's algorithm we use the reweighting trick from @xcite . for non-projective spanning trees synjax uses a combination of reweighting+tarjan algorithm as proposed in @xcite . ## sampling a structure strictly speaking, there is no proper sampling semiring because semirings cannot have nondeterministic output. however, we can still use the semiring framework and make some aspect of them non-deterministic. @xcite and rush (2020) use a semiring that in the forward pass behaves like a log-semiring, but in the backward pass instead of computing the gradient it does sampling. this is in line of how forward-filtering backward-sampling algorithm works @xcite . non-projective spanning trees do not support the semiring framework so we use custom algorithms for them described in @xcite . it contains colbourn's algorithm that has a fixed runtime of o(n 3 ) but is prone to numerical issues be-cause it requires matrix-inversion @xcite , and wilson's algorithm that is more numerically stable but has a runtime that depends on concrete values of log-potentials @xcite . synjax also supports vectorized sampling without replacement (swor) from @xcite . ## differentiable sampling the mentioned sampling algorithms provide unbiased samples of structures useful for many inference tasks, but they are not differentiable because the gradient of sampling from discrete distributions is zero almost everywhere. this problem can be addressed with log-derivative trick from rein-force algorithm @xcite , but that provides high variance estimates of gradients. to address this problem there have been different proposals for differentiable sampling algorithms that are biased but can provide low-variance estimates of gradients. synjax implements majority of the main approaches in the literature including structured attention @xcite , relaxed dynamic programming @xcite , perturband-map @xcite , gumbel-crf @xcite , stochastic softmax-tricks @xcite , and implicit maximum-likelihood estimation @xcite . it also include different noise distributions for perturbations models, including sum-of-gamma noise @xcite that is particularly suited for structured distributions. ## entropy and kl divergence to compute the cross-entropy and kl divergence, we will assume that the two distributions factorize in exactly the same way. like some other properties, cross-entropy can also be computed with the appropriate semirings @xcite @xcite , but those approaches would not work on non-projective spanning tree distributions. there is a surprisingly simple solution that works across all distributions that factorize in the same way and has appeared in a couple of works in the past @xcite @xcite . here we give a full derivation for cross-entropy: this reduces the computation of cross-entropy to finding marginal probabilities of one distribution, and finding log-partition of the other -both of which can be computed efficiently for all distributions in synjax. given the method for computing cross-entropy, finding entropy is trivial: kl divergence is easy to compute too: most inference algorithms apply a large number of elementwise and reshaping operations that are in general fast but create a large number of intermediate tensors that occupy memory. to speed this up we use checkpointing @xcite to avoid memorization of tensors that can be recomputed quickly. that has improved memory usage and speed, especially on tpus. all functions that could be vectorized are written in pure jax. those that cannot, like wilson sampling (1996) and @xcite , are implemented with numba @xcite all synjax distributions inherit from equinox modules @xcite which makes them simultaneously pytrees and dataclasses. thereby all synjax distributions can be transformed with jax.vmap and are compatible with any jax neural framework, e.g. haiku and flax. pgmax @xcite is a jax library that supports inference over arbitrary factor graphs by using loopy belief propagation. after the user builds the desired factor graph, pgmax can do automatic inference over it. for many structured distributions building a factor graph is the difficult part of implementation because it may require a custom algorithm (e.g. cky or needleman-wunsch). synjax implements those custom algorithms for each of the supported structures. with synjax the user only needs to provide the parameters of the distribution and synjax will handle both building of the factor graph and inference over it. for all the included distributions, synjax also provides some features not covered by pgmax, such as unbiased sampling, computation of entropy, cross-entropy and kl divergence. optax @xcite provides ctc loss implementation for jax but without support for computation of optimal alignment, marginals over alignment links, sampling alignments etc. all the mentioned jax libraries focus on continuous or categorical distributions and, with the exception of hmms and ctc loss, do not contain distributions provided by synjax. synjax fills this gap in the jax ecosystem and enables easier construction of structured probability models. the most comparable library in terms of features is torch-struct (rush, 2020) that targets pytorch as its underlying framework. torch-struct, just like synjax, uses automatic differentiation for efficient inference. we will point out here only the main differences that would be of relevance to users. synjax supports larger number of distributions and inference algorithms and provides a unified interface to all of them. it also provides reproducable sampling trough controlled randomness seeds. synjax has a more general approach to computation of entropy that does not depend on semirings and therefore applies to all distributions. synjax is fully implemented in python and compiled with jax.jit and numba.jit while torch-struct does not use any compiler optimizations except a custom cuda kernel for semiring matrix multiplication. if we compare lines of code and speed ( synjax also provides the fastest and most feature rich implementation of spanning tree algorithms. so far the most competitive libraries for spanning trees were by zmigrod et al. and stanojević and cohen. synjax builds on stanojević and cohen code and annotates it with numba instructions which makes it many times faster than any other alternative (see figure 3 in the appendix). ## conclusion one of the main challenges in creating deep neural models over structured distributions is the difficulty of their implementation on modern hardware accelerators. synjax addresses this problem and makes large scale training of structured models feasible and easy in jax. we hope that this will encourage research into finding alternatives to auto-regressive modeling of structured data.