Conformal Autoregressive Generation: Beam Search with Coverage Guarantees

TL;DR

This paper presents two conformal prediction-based extensions to beam search that provide theoretical coverage guarantees for sequence generation, with one offering simple dynamic subsets and the other adaptive beam widths.

Contribution

It introduces novel conformal prediction methods integrated into beam search to ensure coverage guarantees in sequence generation tasks.

Findings

Both methods achieve theoretical coverage bounds.

Empirical evaluation on NLP and chemistry tasks demonstrates effectiveness.

Adaptive method provides coverage guarantees with a variable beam width.

Abstract

We introduce two new extensions to the beam search algorithm based on conformal predictions (CP) to produce sets of sequences with theoretical coverage guarantees. The first method is very simple and proposes dynamically-sized subsets of beam search results but, unlike typical CP procedures, has an upper bound on the achievable guarantee depending on a post-hoc calibration measure. Our second algorithm introduces the conformal set prediction procedure as part of the decoding process, producing a variable beam width which adapts to the current uncertainty. While more complex, this procedure can achieve coverage guarantees selected a priori. We provide marginal coverage bounds for each method, and evaluate them empirically on a selection of tasks drawing from natural language processing and chemistry.