Non-Exchangeable Conformal Language Generation with Nearest Neighbors (2024)

Abstract: Quantifying uncertainty in automatically generated text is important for letting humans check potential hallucinations and making systems more reliable. Conformal prediction is an attractive framework to provide predictions imbued with statistical guarantees, however, its application to text generation is challenging since any i.i.d. assumptions are not realistic. In this paper, we bridge this gap by leveraging recent results on *non-exchangeable* conformal prediction, which still ensures bounds on coverage. The result, *non-exchangeable conformal nucleus sampling*, is a novel extension of the conformal prediction framework to generation based on nearest neighbors. Our method can be used post-hoc for an arbitrary model without extra training and supplies token-level, calibrated prediction sets equipped with statistical guarantees. Experiments in machine translation and language modeling show encouraging results in generation quality. By also producing tighter prediction sets with good coverage, we thus give a more theoretically principled way to perform sampling with conformal guarantees.

## introduction natural language generation (nlg) is a multifaceted field spanning applications such as machine translation (mt), language modeling (lm), summarization, question answering and dialogue generation. owing to the recent success of large language models (llms) such as @xcite , bloom @xcite or llama @xcite , natural language modeling with stochastic decoding (sampling) is increasingly used as an interface with end users. while sampling allows for more fluent and varied text, few methods exist to evaluate the reliability of generated text and adequacy of the underlying sampling method. this is particularly relevant for generation scenarios where pre-trained models are applied to new data with potentially different { } a decoder hidden representation z t is used during inference to retrieve the nearest neighbors and their nonconformity scores s k . their relevance is determined by using their distance to compute weights w k , resulting in the quantile q that forms conformal prediction sets. conformal prediction @xcite @xcite has recently gained popularity by providing calibrated prediction sets that are imbued with statistical guarantees about containing the correct solution. nevertheless, applying conformal prediction to nlg is not trivial and comes with a major obstacle: the conditional generation process breaks the independence and identical distribution (i.i.d.) assumption underlying conformal prediction techniques. we tackle this problem by drawing inspiration from recent advances in nearest neighbor language modeling @xcite @xcite and machine translation @xcite @xcite . this way, we are able to dynamically generate calibration sets during inference that are able to maintain statistical guarantees. we schematically illustrate non-exchangeable conformal nucleus sampling in figure 1 : in the first step, we obtain a (sorted) probability distribution over tokens and a latent representation z t for the current generation step from the model. in a second step, we use the latent representation to query a datastore for similar, previously stored representations and their corresponding non-conformity scores, s k . these scores are then used to compute a threshold q based on the theory of non-excheangeable conformal prediction @xcite , which defines a smaller set of tokens that is sampled from. 1 contributions. we present a general-purpose extension of the conformal framework to nlg by tackling the problems above. our contributions are as follows: 1 to the best of our knowledge, we are the first to present a novel technique based on non-exchangeable conformal prediction and to apply it to language generation to produce calibrated prediction sets. 2 we validate the effectiveness of the method in a language modeling and machine translation context, evaluating the coverage of the calibrated prediction sets and showing that our method is on par with or even outperforms other sampling-based techniques in terms of generation quality, all while maintaining tighter prediction sets and better coverage. @xcite we finally demonstrate that these properties are also maintained under distributional shift induced by corrupting the model's latent representations. 4 we publish all the code for this project in an open-source repository. foot_1 2 related work conformal prediction. conformal prediction is a line of work that has recently regained interest in machine learning by producing prediction sets with certain statistical guarantees about containing the correct prediction @xcite @xcite . as the size of prediction sets is calibrated to fulfill these guarantees, one can also see the size of the prediction set itself as a proxy of the uncertainty of a model-the larger the set, the more possible predictions have to be included in order to maintain the coverage guarantee. conformal prediction has already found diverse applications in nlp for classification @xcite @xcite @xcite and sequence labeling problems @xcite , as well as quality estimation @xcite . unfortunately, generation problems are challenging due to their sequential nature and constant breaking of the i.i.d. assumption, so existing works operate on the sequencelevel instead @xcite @xcite . conformal procedures for time-series @xcite @xcite and @xcite @xcite have been proposed in the literature. the most related work to ours is given by @xcite , who apply the standard conformal prediction setup to nlg, arguing that markov chains are a type of βmixing processes, for which @xcite showed coverage to degrade by an only negligible amount. however, ravfogel et al. do not investigate this claim empirically, and furthermore do not find any benefits when generating sequences. in another related work, @xcite propose an approach that is specifically tailored toward language modeling. however, their prediction sets contain entire sequences instead of single tokens. in contrast, our token-level prediction sets are useful for constraining the options during generation and their widths can represent model uncertainty. ## experiments in the following sections, we conduct experiments in both language modeling and machine translation. for machine translation we opt for the 400 million and 1.2 billion parameter versions of the m2m100 model @xcite on the wmt-2022 shared task datasets for german to english and japanese to english @xcite . for language modelling, we use the 350 million and 5 in this phase, we do not let the model generate freely, but feed it the gold prefix during the decoding process to make sure that conformity scores can be computed correctly. 6 using this formulation of the weights w k that depends on the data deviates from the assumptions of original proof, as discussed in @xcite , §4.5. nevertheless, our results in section 4 and those by @xcite show that the obtained bound in equation (4) still remains useful. ## discussion our experiments have shown that despite the absence of i.i.d. data in nlg and the loss in coverage induced by using dynamic calibration sets, the resulting coverage is still close to the pre-specified desired level for both lm and mt. additionally, even though the coverage gap predicted by the method of @xcite is infeasible to compute for us, we did not observe any critical degradation in practice. further, we demonstrated how sampling from these calibrated prediction sets performs similarly or better than other sampling methods. even though our method is still outperformed by beam search in the mt setting, previous work such as minimum bayes risk (mbr) decoding has shown how multiple samples can be re-ranked to produce better outputs @xcite @xcite . additionally, recent dialogue systems based on llms use sampling instead of beam search for generation. since our prediction sets are more flexible and generally tighter, our results serve as a starting point for future work. for instance, our technique could be used with non-conformity scores that do not consider token probabilities alone (e.g. @xcite or using prediction set widths as a proxy for uncertainty @xcite . ## conclusion we successfully demonstrated the application of a non-exchangeable variant of conformal prediction to machine translation and language modeling with the help of k-nn retrieval. we showed our method to be able to maintain the desired coverage best across different dataset strata while keeping prediction sets smaller than other sampling methods, all while providing theoretical coverage guarantees about coverage that other comparable methods lack. we validated our method to produce encouraging results for generation tasks. lastly, we analyzed the behavior under distributional drift, showing how the k-nn retrieval maintains desirable properties for the estimated prediction sets. we see our method as a step to provide a more principled way to perform sampling with conformal guarantees under more realistic assumptions.