Last: Scalable Lattice-Based Speech Modelling in Jax

Ke Wu; Ehsan Variani; Tom Bagby; Michael Riley

ICASSP 2023accepted0 citations

Last: Scalable Lattice-Based Speech Modelling in Jax

Ke Wu, Ehsan Variani, Tom Bagby, Michael Riley

Abstract

We introduce LAST, a LAttice-based Speech Transducer library in JAX. With an emphasis on flexibility, ease-of-use, and scalability, LAST implements differentiable weighted finite state automaton (WFSA) algorithms needed for training & inference that scale to a large WFSA such as a recognition lattice over the entire utterance. Despite these WFSA algorithms being well-known in the literature, new challenges arise from performance characteristics of modern architectures, and from nuances in automatic differentiation. We describe a suite of generally applicable techniques employed in LAST to address these challenges, and demonstrate their effectiveness with benchmarks on TPUv3 and V100 GPU.

BibTeX

@inproceedings{icassp2023_lastscalablelatt,
  title = {Last: Scalable Lattice-Based Speech Modelling in Jax},
  author = {Ke Wu and Ehsan Variani and Tom Bagby and Michael Riley},
  booktitle = {ICASSP 2023},
  year = {2023}
}