mbrs: A Library for Minimum Bayes Risk Decoding

Abstract

Minimum Bayes risk (MBR) decoding is a decision rule of text generation tasks that outperforms conventional maximum a posterior (MAP) decoding using beam search by selecting high-quality outputs based on a utility function rather than those with high-probability. Typically, it finds the most suitable hypothesis from the set of hypotheses under the sampled pseudo-references. mbrs is a library of MBR decoding, which can flexibly combine various metrics, alternative expectation estimations, and algorithmic variants. It is designed with a focus on speed measurement and calling count of code blocks, transparency, reproducibility, and extensibility, which are essential for researchers and developers. We published our mbrs as an MIT-licensed open-source project, and the code is available on GitHub. GitHub: https://github.com/naist-nlp/mbrs

Figures (3)

• Figure 1: Workflow overview of mbrs. Decoder class gets a candidate list and outputs high-quality hypotheses. It internally calls the expected_scores() method implemented in Metric classes that calculate the expected scores for each hypothesis. The expected scores are estimated with Monte Carlo (MC) or the model-based (MB) method. In the figure, the colored boxes without border lines denote the functions or methods and the square boxes with border lines denote the abstract classes. Methods that have the same color as a class indicate that they belong to the class.
• Figure 2: The distribution of expected utility scores in the set of hypotheses. The left and right ones show the examples that have the maximum and minimum variance of the expected utility in the test set, respectively.
• Figure 3: The empirical cumulative distribution of ranks selected by MBR decoding from descending rankings of output probabilities.