Differentiable Weighted Finite-State Transducers

1 INTRODUCTION . Weighted finite-state transducers ( WFSTs ) are a commonly used tool in speech and language processing ( Knight & May , 2009 ; Mohri et al. , 2002 ) . They are most frequently used to combine predictions from multiple already trained models . In speech recognition , for example , WFSTs are used to combine constraints from an acoustic-to-phoneme model , a lexicon mapping words to pronunciations , and a word-level language model . However , combining separately learned models using WFSTs only at inference time has several drawbacks , including the well-known problems of exposure bias ( Ranzato et al. , 2015 ) and label bias ( Bottou , 1991 ; Lafferty et al. , 2001 ) . Given that gradients may be computed for most WFST operations , using them only at the inference stage of a learning system is not a hard limitation . We speculate that this limitation is primarily due to practical considerations . Historically , hardware has not been sufficiently performant to make training with WFSTs tractable . Also , no implementation exists with the required operations which supports automatic differentiation in a high-level yet efficient manner . We develop a framework for automatic differentiation through operations on WFSTs . We show the utility of this framework by leveraging it to design and experiment with existing and novel learning algorithms . Automata are a more convenient structure than tensors to encode prior knowledge into a learning algorithm . However , not training with them limits the extent to which this prior knowledge can be incorporated in a useful manner . A framework for differentiable WFSTs allows the model to jointly learn from training data as well as prior knowledge encoded in WFSTs . This enables the learning algorithm to incorporate such knowledge in the best possible way . Use of WFSTs conveniently decomposes operations from data ( i.e . graphs ) . For example , rather than hand-coding sequence-level loss functions such as Connectionist Temporal Classification ( CTC ) ( Graves et al. , 2006 ) or the Automatic Segmentation Criterion ( ASG ) ( Collobert et al. , 2016 ) , we may specify the core assumptions of the criteria in graphs and compute the resulting loss with graph operations . This facilitates exploration in the space of such structured loss functions . We show the utility of the differentiable WFST framework by designing and testing several algorithms . For example , bi-gram transitions may be added to CTC with a transition WFST . We scale transitions to large token set sizes by encoding pruning and back-off in the transition graph . Word pieces are commonly used as the output of speech recognition and machine translation models ( Chiu et al. , 2018 ; Sennrich et al. , 2016 ) . The word piece decomposition for a word is learned with a task-independent model . Instead , we use WFSTs to marginalize over the latent word piece decomposition at training time . This lets the model learn decompositions salient to the task at hand . Finally , we show that WFSTs may be used as layers themselves intermixed with tensor-based layers . We propose a convolutional WFST layer which maps lower-level representations to higher-level representations . The WFST convolution can be trained with the rest of the model and results in improved accuracy with fewer parameters and operations as compared to a traditional convolution . In summary , our contributions are : • A framework for automatic differentiation with WFSTs . The framework supports both C++ and Python front-ends and is available at https : //www.anonymized.com . • We show that the framework may be used to express both existing sequence-level loss functions and to design novel sequence-level loss functions . • We propose a convolutional WFST layer which can be used in the interior of a deep neural network to map lower-level representations to higher-level representations . • We demonstrate the effectiveness of using WFSTs in the manners described above with experiments in automatic speech and handwriting recognition . 2 RELATED WORK . A wealth of prior work exists using weighted finite-state automata in speech recognition , natural language processing , optical character recognition , and other applications ( Breuel , 2008 ; Knight & May , 2009 ; Mohri , 1997 ; Mohri et al. , 2008 ; Pereira et al. , 1994 ) . However , the use of WFSTs is limited mostly to the inference stage of a predictive system . For example , Kaldi , a commonly used toolkit for automatic speech recognition , uses WFSTs extensively , but in most cases for inference or to estimate the parameters of shallow models ( Povey et al. , 2011 ) . In some cases , WFSTs are used statically to incorporate fixed lattices in discriminative sequence criteria ( Kingsbury , 2009 ; Kingsbury et al. , 2012 ; Su et al. , 2013 ; Veselỳ et al. , 2013 ) . Implementations of sequence criteria in end-to-end style training are typically hand-crafted with careful consideration for speed ( Amodei et al. , 2016 ; Collobert et al. , 2019 ; Povey et al. , 2016 ) . The use of hand-crafted implementations reduces flexibility which limits research . In some cases , such as the fully differentiable beam search of Collobert et al . ( 2019 ) , achieving the necessary computational efficiency with a WFST-based implementation may not yet be tractable . However , as a first step , we show that in many common cases we can have the expressiveness afforded by the differentiable WFST framework without paying an unacceptable penalty in execution time . The ideas of learning with WFSTs ( Eisner , 2002 ) and automatic differentiation through operations on graphs ( Bottou et al. , 1997 ) are not new . However , no simple and efficient frameworks exist . Especially related to this work , and inspiring the name of our framework , are the graph transformer networks of Bottou et al . ( 1997 ) . Generalized graph transducers ( Bottou et al. , 1996 ) are more expressive than WFSTs , allowing arbitrary data as edge labels . When composing these graphs , one defines an edge matching function and a “ transformer ” to construct the resulting structure . While not this general , differentiable WFSTs nevertheless allow for a vast design space of interesting algorithms , and perhaps make a more pragmatic trade-off between flexibility and efficiency . Highly efficient libraries for operations on WFSTs exist , notably OpenFST and its predecessor FSM ( Allauzen et al. , 2007 ; Mohri et al. , 2000 ) . We take inspiration from OpenFst in the interface and implementation of many of our functions . However , the design implications of operating on WFSTs with automatic differentiation are quite different than those of the use cases OpenFST has been optimized for . We also draw inspiration from libraries for automatic differentiation and deep learning ( Collobert et al. , 2011 ; Paszke et al. , 2019 ; Pratap et al. , 2019 ; Tokui et al. , 2015 ) . Some of the algorithms we propose , with the goal of demonstrating the utility of the differentiable WFST library , are inspired by prior work . Prior work has explored pruning the set of allowed alignments with CTC , and in particular limiting the spacing between output tokens ( Liu et al. , 2018 ) . Learning n-gram word decompositions with both differentiable ( Liu et al. , 2017 ) and nondifferentiable ( Chan et al. , 2017 ) loss functions has also been explored . 3 DIFFERENTIABLE WEIGHTED FINITE-STATE TRANSDUCERS . A weighted finite-state acceptor A is a 6-tuple consisting of an alphabet Σ , a set of states Q , start states Qs , accepting states Qa , a transition function π ( q , p ) which maps elements of Q × Σ to elements of Q , and a weight function ω ( q , p ) which maps elements of Q × Σ to R. A weighted finite-state transducer T is a 7-tuple which augments an acceptor with an output alphabet ∆ . The transition function π ( q , p , r ) and weight function ω ( q , p , r ) map elements of Q×Σ×∆ to elements of Q and R respectively . In other words , each edge of a transducer connects two states and has an input label p , an output label r and a weight w ∈ R. We denote an input by p = [ p1 , . . . , pT ] where each pi ∈ Σ . An acceptor A accepts the input p if there exists a sequence of states qi+1 = π ( qi , pi ) ∈ Q such that q1 ∈ Qs and qT+1 ∈ Qa . With a slight abuse of notation , we let p ∈ A denote that A accepts p. The score of p is given by s ( p ) =∑T i=1 ω ( qi , pi ) . Let r = [ r1 , . . . , rT ] be a path with ri ∈ ∆ . A transducer T transduces the input p to the output r ( i.e . ( p , r ) ∈ T ) if there exists a sequence of states qi+1 = π ( qi , pi , ri ) ∈ Q such that q1 ∈ Qs and qT+1 ∈ Qa . The score of the pair ( p , r ) is given by s ( p , r ) = ∑T i=1 ω ( qi , pi , ri ) . We restrict operations to the log and tropical semirings . In both cases the accumulation of weights along a path is with addition . In the log semiring the accumulation over path scores is with logsum-exp which we denote by logaddi si = log ∑ i e si . In the tropical semiring , the accumulation over path scores is with the max , maxi si . We allow transitions in both acceptors and transducers . An input ( or output ) on an edge means the edge can be traversed without consuming an input ( or emitting an output ) . This allows inputs to map to outputs of differing length . 3.1 OPERATIONS . We briefly describe a subset of the most important operations implemented in our differentiable WFST framework . For a more detailed discussion of operations on WFSTs see e.g . Mohri ( 2009 ) . Intersection We denote by B = A1 ◦ A2 the intersection of two acceptors A1 and A2 . The intersected graph B contains all paths which are accepted by both inputs . The score of any path in B is the sum of the scores of the two corresponding paths in A1 and A2 . Composition The same symbol denotes the composition of two transducers U = T1 ◦ T2 . If T1 transduces p to u and T2 transduces u to r then U transduces p to r. As in intersection , the score of the path in the composed graph is the sum of the path scores from the input graphs . Forward and Viterbi Score The forward score of T is logadd ( p , r ) ∈T s ( p , r ) . Similarly , the Viterbi score of T is max ( p , r ) ∈T s ( p , r ) . 0 1a : a/0 a : ε/0 ( a ) Token graph Ta 0 1a : ε/0 2b : ab/0 ( b ) Label graph Y 0 1a : ε/0 a : ε/0 2b : ab/0 b : ε/0 ( c ) Alignment graph T ◦ Y Viterbi Path The Viterbi path of T is given by arg max ( p , r ) ∈T s ( p , r ) . The forward score , Viterbi score , and the Viterbi path for an acceptorA are defined in the same way using just p. In a directed acyclic graph these operations can be computed in time linear in the size of the graph with the well-known forward and Viterbi algorithms ( Jurafsky , 2000 ; Rabiner , 1989 ) . We support standard rational operations including the union ( A + B ) , concatenation ( AB ) , and the Kleene closure ( A∗ ) of WFSTs . To enable automatic differentiation through complete computations , we support slightly non-standard operations . We allow for the negation of all the arcs weights in a graph and the addition or subtraction of the arc weights of two identically structured graphs .

Summary: This paper presents how weighted finite-state transducers (WFST) and a few common operations performed on them can be integrated in a differentiable model, and therefore contribute to the training of complete systems. The authors propose a few case studies, mainly in language applications, where the WFSTs are used to compute a sequence-level loss function, to keep ambiguity and let the model decide word-peices decomposition, or gracefully replace convolution layers. The code associated with the presented methods will be available.