GPU-Accelerated Forward-Backward algorithm with Application to Lattice-Free MMI

TL;DR

This paper introduces a GPU-optimized forward-backward algorithm expressed through sparse matrix operations in a semiring, enabling faster training of TDNNs with LF-MMI without approximations.

Contribution

It presents a novel GPU-friendly formulation of the forward-backward algorithm using semiring algebra, improving training speed for LF-MMI-based models.

Findings

Implementation is about twice as fast as existing C++/CUDA systems.

No need for approximations like leaky-HMM in training.

Easily implemented in Julia or similar languages.

Abstract

We propose to express the forward-backward algorithm in terms of operations between sparse matrices in a specific semiring. This new perspective naturally leads to a GPU-friendly algorithm which is easy to implement in Julia or any programming languages with native support of semiring algebra. We use this new implementation to train a TDNN with the LF-MMI objective function and we compare the training time of our system with PyChain - a recently introduced C++/CUDA implementation of the LF-MMI loss. Our implementation is about two times faster while not having to use any approximation such as the "leaky-HMM".