Document-Level Text Generation with Minimum Bayes Risk Decoding using Optimal Transport

2025-05-29

AI Summary

Document-level text generation faces two key challenges: (1) difficulty in modeling long-range contextual dependencies, and (2) the absence of document-granularity utility functions in existing Minimum Bayes Risk (MBR) decoding frameworks. This work pioneers the integration of optimal transport theory into the MBR paradigm, proposing a document-level utility estimation method grounded in the Wasserstein distance. By leveraging this geometrically principled metric, our approach overcomes the limitation that sentence-level utility functions cannot generalize to document-level evaluation. Crucially, it requires no model retraining and operates as a plug-and-play post-hoc refinement of decoder outputs. Empirical evaluation across document-level machine translation, text simplification, and dense image captioning demonstrates consistent and significant improvements over standard MBR decoding—evidenced by gains in BLEU, SARI, and CIDEr—thereby validating both the effectiveness and cross-task generalizability of document-level utility modeling.

Technology Category

• Natural Language Processing: Generation
• Reasoning under Uncertainty: Decision/Utility Theory
• Machine Learning: Calibration & Uncertainty Quantification

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Document-level text generation tasks are known to be more difficult than sentence-level text generation tasks as they require the understanding of longer context to generate high-quality texts. In this paper, we investigate the adaption of Minimum Bayes Risk (MBR) decoding for document-level text generation tasks. MBR decoding makes use of a utility function to estimate the output with the highest expected utility from a set of candidate outputs. Although MBR decoding is shown to be effective in a wide range of sentence-level text generation tasks, its performance on document-level text generation tasks is limited as many of the utility functions are designed for evaluating the utility of sentences. To this end, we propose MBR-OT, a variant of MBR decoding using Wasserstein distance to compute the utility of a document using a sentence-level utility function. The experimental result shows that the performance of MBR-OT outperforms that of the standard MBR in document-level machine translation, text simplification, and dense image captioning tasks. Our code is available at https://github.com/jinnaiyuu/mbr-optimal-transport

Problem

• Adapting MBR decoding for document-level text generation
• Improving utility estimation for long documents with Wasserstein distance
• Enhancing performance in document-level machine translation and captioning

Innovation

• MBR-OT uses Wasserstein distance for documents
• Adapts sentence-level utility to document-level
• Outperforms standard MBR in multiple tasks