1 Introduction↩︎

Natural language generation (NLG) is a multi-faceted 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 GPT-4 [1], BLOOM [2] or LLaMA [3], 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 distribution to the training data, increasing the risk of generating erroneous, misleading, and potentially harmful text [4]–[8].

Conformal prediction [9]–[11] 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 [12]–[14] and machine translation [15]–[18]. 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 1: In the first step, we obtain a (sorted) probability distribution over tokens and a latent representation \(\mathop{\mathrm{\mathbf{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 \(\hat{q}\) based on the theory of non-exchangeable conformal prediction [19], which defines a smaller set of tokens that is sampled from.¹

1.0.0.1 Contributions.

We present a general-purpose extension of the conformal framework to NLG by tackling the problems above. Our contributions are as follows: 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. 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. We finally demonstrate that these properties are also maintained under distributional shift induced by corrupting the model’s latent representations. We publish all the code for this project in an open-source repository.²

2 Related Work↩︎

2.0.0.1 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 [9]–[11]. 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 [20]–[25] and sequence labeling problems [26], as well as quality estimation [27], [28]. 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 sequence-level instead [29]–[31]. Conformal procedures for time-series [32]–[35] and general non-i.i.d. data [19], [36]–[38] have been proposed in the literature. The most related work to ours is given by [39], who apply the standard conformal prediction setup to NLG, arguing that Markov chains are a type of \(\beta\)-mixing processes, for which [34] showed coverage to degrade by an only negligible amount. However, [39] do not investigate this claim empirically, and furthermore do not find any benefits when generating sequences. In another related work, [29] 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.

2.0.0.2 Uncertainty in NLP.

Modeling uncertainty in NLP has already been studied in classification [40]–[42] and regression settings [43]–[45]. However, NLG proves more challenging due to it non-i.i.d. and combinatorial nature. Some works have proposed Bayesian Deep Learning methods for NLG: [46] use Monte Carlo Dropout [47] to produce multiple generations for the same input and measure their pair-wise BLEU scores. [48] define extensions of mutual information for structured prediction. Other existing approaches try to account for the paraphrastic nature of language by modeling the entropy over meaning classes [49], investigate the use of linguistic markers to indicate uncertainty [50] or ask the model directly for its confidence [51], [52]. [53] provide an extensive overview of the theory and current state of the field.

3 Background↩︎

3.0.0.1 Conformal Prediction.

Conformal prediction is an attractive method for uncertainty quantification due to its statistical coverage guarantees [9]–[11]. Given some predictor, a held-out calibration set \(\{(\mathop{\mathrm{\mathbf{x}}}_i, y_i)\}_{i=1}^N\), and a pre-defined miscoverage level \(\alpha\) (e.g., 0.1), the calibration set is used to obtain prediction sets \(\mathcal{C}(\mathop{\mathrm{\mathbf{x}}}^*)\) for a new test point \(\mathop{\mathrm{\mathbf{x}}}^*\) satisfying \[\label{eq:conformal-prediction-guarantees} p\Big(y^* \in \mathcal{C}(\mathop{\mathrm{\mathbf{x}}}^*)\Big) \ge 1 - \alpha,\tag{1}\] that is, the probability of the prediction set \(\mathcal{C}(\mathop{\mathrm{\mathbf{x}}}^*)\) containing the correct label \(y^*\) is at least \(1-\alpha\). This is achieved by the following recipe: Firstly, one has to define a non-conformity score, that provides an estimate of the distance of the test point to the rest of the data, i.e., a proxy for the uncertainty over the test point predictions. In this context, the score can be as simple as \(s_i = 1 - p_{\mathop{\mathrm{\boldsymbol{\theta}}}}(y|\mathop{\mathrm{\mathbf{x}}})\), i.e. one minus the softmax probability of the true class, which will be higher when the model is wrong or less confident. Next, we define \(\hat{q}\) as the \(\big\lceil(N+1)(1 - \alpha)/N\big\rceil\)-th quantile of the non-conformity scores. Then, when we make a new prediction for a test point \(\mathop{\mathrm{\mathbf{x}}}^*\), we can create prediction sets defined as

\[\label{eq:conformal-prediction-sets} \mathcal{C}(\mathop{\mathrm{\mathbf{x}}}^*) = \Big\{y\;\Big|\;p_{\mathop{\mathrm{\boldsymbol{\theta}}}}(y|\mathop{\mathrm{\mathbf{x}}}^*) \ge 1 - \hat{q}\Big\},\tag{2}\]

which is guaranteed to fulfil the coverage requirement in 1 for i.i.d. data [9], [11].

3.0.0.2 Non-exchangeable Conformal Prediction.

[19] address a major shortcoming in the method above: When a test point and the calibration data are not i.i.d.,³ the distributional drift causes any previously found \(\hat{q}\) to be miscalibrated, and thus the intended coverage can no longer be guaranteed. However, we can still perform conformal prediction by assigning a weight \(w_i \in [0, 1]\) to every calibration data point, reflecting its relevance—i.e. assigning lower weights to points far away from the test distribution. Then, by normalizing the weights with \(\tilde{w}_i = w_i / (1 + \sum_{i=1}^N w_i)\), we define the quantile as

\[\label{eq:non-exchangeable-quantile} \hat{q} = \inf \Big\{q\;\Big| \sum_{i=1}^N \tilde{w}_i \mathbf{1}\big\{s_i\le q \big\} \ge 1 - \alpha \Big\},\tag{3}\]

with \(\mathbf{1}\{\cdot\}\) denoting the indicator function. The construction of the prediction sets then follows the same steps as before. Most notably, the coverage guarantee in 1 now changes to

\[\label{eq:non-exchangeable-conformal-prediction-guarantees} p\Big(y^* \in \mathcal{C}(\mathop{\mathrm{\mathbf{x}}}^*)\Big) \ge 1 - \alpha - \sum_{i=1}^N \tilde{w}_i\varepsilon_i,\tag{4}\]

with an extra term including the total variation distance between the distribution of a calibration and a test point, \(\varepsilon_i = d_\text{TV}\big((\mathop{\mathrm{\mathbf{x}}}_i, y_i), (\mathop{\mathrm{\mathbf{x}}}^*, y^*)\big)\).⁴ Unfortunately, this term is hard to estimate or bound, nevertheless, the selection of appropriate weights that can capture the relevance of calibration points to the test set should moderate both the impact of the distant data points on the estimation of the prediction set and the impact of \(d_{\mathrm{TV}}\) on the coverage bound. In other words, for large \(d_{\mathrm{TV}}\) values we expect to have smaller weights, that allow us to achieve coverage close to the desired values. We show in our experiments that the loss of coverage when using nearest neighbor weights is limited and revisit the practical implications in 5.

3.1 Method: Non-exchangeable Conformal Language Generation through Nearest Neighbors↩︎

We now present a novel method to apply conformal prediction in NLG by synthesizing the non-exchangeable approach of [19] with \(k\)-NN search-augmented neural models [12], [15]. The related approach by [39] calibrates prediction sets within bins of similar entropies using the non-exchangeable procedure described in 3. However, this implies that we would use semantically unrelated (sub-)sequences to calibrate the model—in fact, we show experimentally that this approach obtains generally trivial coverage by producing extremely wide prediction sets. Instead, we propose to perform a dynamic calibration step during model inference, only considering the most relevant data points from the calibration set. We do this in the following way: Given a dataset \(\{(\mathop{\mathrm{\mathbf{x}}}^{(i)}, y^{(i)})\}\) of sequences \(\mathop{\mathrm{\mathbf{x}}}^{(i)} = (\mathop{\mathrm{\mathbf{x}}}_1^{(i)}, \ldots, \mathop{\mathrm{\mathbf{x}}}_S^{(i)})\) and corresponding references consisting of gold tokens \(y^{(i)} = (y_1^{(i)}, \ldots, y_T^{(i)})\), we extract the model’s decoder activations \(\mathop{\mathrm{\mathbf{z}}}_t^{(i)} \in \mathbb{R}^d\) and conformity scores \(s_t^{(i)}\).⁵ We save those in a datastore allowing for fast and efficient nearest neighbor search using FAISS [54]. In the inference phase, during every decoding step, we then use the decoder hidden state \(\mathop{\mathrm{\mathbf{z}}}_t^*\) to query the data store for the \(K\) nearest neighbors and their conformity scores and record their distances. We use the squared \(l_2\) distance to compute the weight \(w_k\) as

\[\label{eq:weight-equation} w_k = \exp\Big(-\big|\big|\mathop{\mathrm{\mathbf{z}}}_t^* - {\mathop{\mathrm{\mathbf{z}}}_k}\big|\big|^2_2\;/\;\tau\Big),\tag{5}\]

where \(\tau\) corresponds to a temperature hyperparameter.⁶ This formulation is equivalent to a RBF kernel with scale parameter \(\tau\). Finally, we use the weights to compute the quantile \(\hat{q}\) as in 3 . The entire algorithm is given in 7.5.

3.1.0.1 Adaptive Prediction Sets.

The efficacy of conformal prediction hinges on the choice of non-conformity score, with the simple non-conformity score \(s_i = 1 - p_{\mathop{\mathrm{\boldsymbol{\theta}}}}(y_t|\mathop{\mathrm{\mathbf{x}}}, y_{<t})\) known to undercover hard and overcover easy subpopulations of the data. Due to the diverse nature of language, we therefore opt for adaptive prediction sets [55], [56]. Adaptive prediction sets redefine the non-conformity score as the cumulative probability over classes (after sorting descendingly) necessary to reach the correct class. Intuitively, this means that we included all classes whose cumulative probability does not surpass \(\hat{q}\). Compared to the simple conformity score, this produces wider predictions sets for hard inputs, encompassing more potentially plausible continuations in a language context. A more formal definition is given in 7.1.

4 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 [57] on the WMT-2022 shared task datasets for German to English and Japanese to English [58]. For Language Modelling, we use the 350 million and 1.3 billion parameter versions of the OPT model [59] and replicate the setup by [39]: We calibrate our model on \(10000\) sentences from a 2022 English Wikipedia dump [60] and test coverage and generation on \(1000\) sentences from OpenWebText [61].⁷ All models are used in a zero-shot setup without extra training or finetuning. For the datastore, we use the implementation by FAISS library [54], computing \(2048\) clusters in total and probing \(32\) clusters per query. We also summarize the environmental impact of our experiments in 7.6.

4.1 Evaluating Coverage↩︎

First of all, we demonstrate that the retrieved information from the data store enables us to successfully apply the proposed method. Coverage is an important notion in conformal prediction, referring to the correct label being covered by a prediction set or intervals. Since we can always achieve trivial coverage by choosing the largest possible prediction set, an ideal method would strike a balance between high coverage and small prediction sets. While it is not possible to measure coverage in a free generation setting (see next section), we can assess whether the correct class is contained in the prediction set if we feed the actual reference tokens into the decoder and check whether we include the true continuation.⁸ For our MT task, this is reminiscent of an interactive translation prediction setup [63]–[65], where we would like to suggest possible continuations to a translator, suggesting the next word from a set of words that (a) contains plausible options and (b) is limited in size, in order to restrict the complexity for the end user. Before we run our experiments, we need to determine \(\tau\), which we tune on the calibration set using a stochastic hill-climbing procedure described in 7.2. We compare our non-exchangeable conformal nucleus sampling (Non-Ex. CS) with nucleus sampling [66] and conformal nucleus sampling (Conf. Sampl.; [39]). The latter bin predictions on a calibration set by the entropy of the output distribution, and compute one \(\hat{q}\) per such entropy bin using the standard conformal procedure given in the beginning of 3.

4.1.0.1 Evaluation.

We measure the total coverage using different distance metrics, namely, squared \(l_2\) distance, normalized inner product, and cosine similarity (see 1 2),⁹ as well as binning predictions by set size and then measuring the per-bin coverage in 2 (more results given in 7.3). We also summarize the plots in 2 via the Expected Coverage Gap (ECG)¹⁰ that we define as

\[\label{eq:ecg} \text{ECG} = \sum_{b=1}^B \frac{|\mathcal{B}_b|}{N} \max \Big( 1 - \alpha - \text{Coverage}\big(\mathcal{B}_b\big), 0 \Big),\tag{6}\] where \(\mathcal{B}_b\) denotes a single bin and \(N\) the total number of considered predictions in the dataset.¹¹ The ECG thus captures the average weighted amount of undercoverage across bins. In our experiments, we use \(75\) bins in total. The same bins are used to also evaluate the Size-Stratified Coverage metric (SSC) proposed by [68], with a well-calibrated method resulting in a SCC close to the desired coverage \(1-\alpha\):

\[\text{SCC} = \min_{b \in \{1,\ldots, B\}} \text{Coverage}\big(\mathcal{B}_b\big).\]

We can therefore understand the SCC as the worst-case coverage across all considered bins. We present some additional experiments where we assess the impact of key hyperparameters in 7.4.

4.1.0.2 Results.

We found our method to miss the desired coverage of \(90 \%\) for MT by \(8 \%\) or less. Beyond the reported values, we were not able to further increase coverage by varying the temperature parameter without avoiding trivial coverage (i.e., defaulting to very large set sizes), which is likely due to the impossible-to-estimate coverage in 4 . Most notably, our method was able to achieve better SCC scores while maintaining considerably smaller prediction sets than the baselines on average. The reason for this is illustrated in 2: while standard nucleus sampling produces some prediction sets that are small, the total coverage seems to mostly be achieved by creating prediction sets between 60k–80k tokens. The behavior of conformal nucleus sampling by [39] is even more extreme in this regard, while our method focuses on producing smaller prediction sets, with the frequency of larger set sizes decreasing gracefully. In 2 (d), we can see that the larger M2M100 models also tend to produce larger prediction sets, but still noticeably smaller than the baselines. Importantly, for both M2M100 models, even very small prediction sets (size \(\leq 1000\)) achieve non-trivial coverage, unlike the baseline methods. For LM, we always found the model to slightly overcover. This does not contradict the desired lower bound on the coverage in 4 and suggests a more negligible distributional drift. While nucleus sampling produces the smallest average prediction sets, we can see that based on the SCC values some strata remain undercovered. Instead, our method is able to strike a balance between stratified coverage and prediction set size. With respect to distance measures, we find that the difference between them is minimal, indicating that the quality largely depends on the retrieved local neighborhood of the decoder encoding and that finding the right temperature can help to tune the models to approximate the desired coverage. We would now like to find out whether this neighborhood retrieval mechanism can prove to be robust under distributional shift as well. Since we did not observe notable differences between the distance metrics, we continue with the \(l_2\) distance.

4.2 Coverage Under Shift↩︎

To demonstrate how the retrieval of nearest neighbors can help to maintain coverage under distributional shift, we add Gaussian noise of increasing variance—and therefore intensity—to the last decoder hidden embeddings (for MT) and the input embeddings (LM).¹² This way, we are able to simulate distributional drift while still keeping the original sequence of input tokens intact, allowing us to measure the actual coverage. We show the achieved coverage along with the average set size (as a percentage of the total vocabulary) and the average quantile \(\hat{q}\) in 3. We can see that the conformal sampling method deteriorates into returning the full vocabulary as a prediction set. Thus it behaves similarly to simple sampling as indicated by the \(\hat{q}\) values being close to \(1\). Nucleus sampling provides smaller prediction sets compared to conformal sampling, but they seem invariant to noise. As such, the method is not robust to noise injection in the open text generation task, and the obtained coverage deteriorates with noise variance \(\geq 0.025\). Instead, the use of nearest neighbors allows for the estimation of prediction sets that are small but amenable to increase, such that the obtained coverage remains close to the desired one. We can specifically observe that the prediction set size increases considerably to mitigate the injected noise in the open-text generation case.

4.2.0.1 Neighbor Retrieval.

We further analyze how the retrieval enables this flexibility by relating it to the entropy of the output distribution of the 400M parameters M2M100 on German to English. Intuitively, the baseline methods, faced by high-entropy output distributions, need to produce wide prediction sets in order to maintain coverage. In fact, we report such results by correlating entropy levels and prediction set sizes using Spearman’s \(\rho\) in 3, showing strong positive correlations. Our method in contrast shows consistently an anticorrelation between these two quantities, enabled by decoupling the creation of prediction sets from statistics of the output distribution to instead considering the non-conformity scores of similar subsequences. The fact that the prediction set size is not just dependent on the entropy of the predictions while maintaining coverage demonstrates the value of the nearest neighbors: In this way, model uncertainty becomes more flexible and is corroborated by evidence gained from similar inputs.

4.3 Generation Quality↩︎

Crucially, our method should not degrade and potentially even improve generation quality. Thus, we evaluate generation quality for the same tasks without supplying the gold prefix. For language modeling, we follow [39] and use the first \(35\) tokens from the original sentence as input. We compare against a set of generation strategies including top-\(k\) sampling [76]–[78], nucleus sampling and conformal nucleus sampling. We also test a variant of our method using constant weights \(w_k = 1\) for retrieved neighbors (Const. Weight CS) to assess the impact of the weighted neighbor retrieval procedure. We further compare with beam search [79], [80] with a softmax temperature of \(0.1\), and greedy decoding. Evaluation is performed using BLEU [81], COMET-22 [82], [83] and chrF [84] for MT as well as MAUVE [85] and BERTscore [86] for text generation.¹³

4.3.0.1 Results.

We show the results for the different methods in ¿tbl:tab:translations-results?. We see that beam search outperforms all sampling methods for MT. This corroborates previous work by [87] who argue that (nucleus) sampling methods, by pruning only the bottom percentile of the token distribution, introduce some degree of randomness that is beneficial for open text generation but may be less optimal for conditional language generation, where the desired output is constrained and exact matching generations are preferred (which is the case for MT). Among sampling methods, we find nucleus sampling and conformal sampling to perform similarly (being in agreement with the findings of [39]) but are sometimes on par or even outperformed by our non-exchangeable conformal sampling for MT. For text generation, our method performs best for the smaller OPT model but is slightly beaten by conformal nucleus sampling in terms of MAUVE. When using constant weights, performance deteriorates to the conformal sampling setup, emphasizing the importance of not considering all conformity scores equally when computing \(\hat{q}\), even though the effect seems to be less pronounced for larger models. This illustrates the benefit of creating flexible prediction sets that are adapted on token-basis, suggesting that both the latent space neighborhoods as well as the conformity scores are informative. We discuss examples of generated text in 7.7.

5 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 [19] 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 [88]–[91]. 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. [92]) or using prediction set widths as a proxy for uncertainty [55].

6 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.

Limitations↩︎

We highlight two main limitations of our work here: Potential issues arising from different kinds of dataset shift as well as efficiency concerns.

6.0.0.1 Distributional Drifts.

Even though any loss of coverage due to the term quantifying distributional drift in 4 was limited in our experiments (see [sec:retrieval-quality] [sec:shift-robustness]), this might not hold across all possible setups. As long as we cannot feasibly approximate the shift penalty, it is impossible to determine a priori whether the loss of coverage might prove to be detrimental, and would have to be checked in a similar way as in our experiments. Furthermore, we only consider shifts between the models’ training distributions and test data distributions here, while many other, unconsidered kinds of shifts exist [93], [94].

6.0.0.2 Computational Efficiency.

Even using optimized tools such as FAISS [54], moving the conformal prediction calibration step to inference incurs additional computational cost during generation. Nevertheless, works such as [18], [95] show that there are several ways to improve the efficiency of \(k\)-NN approaches, and we leave such explorations to future work.

Ethical Considerations↩︎

The main promise of conformal prediction lies in its correctness—i.e. producing prediction sets that contain the correct prediction and are thus reliable. In an application, this could potentially create a false sense of security. On the one hand, the conformal guarantee holds in expectation, and not necessarily on a per-sample basis. On the other hand, our experiments have demonstrated that coverage might also not hold when distributional shifts are at work or when looking at specific subpopulations. Therefore, any application should certify that coverage is maintained for potentially sensitive inputs.

Acknowledgements↩︎

We thank the anonymous reviewers for the constructive feedback and useful discussions. This work was supported by EU’s Horizon Europe Research and Innovation Actions (UTTER, contract 101070631), by the project DECOLLAGE (ERC-2022-CoG 101088763), by the Portuguese Recovery and Resilience Plan through project C645008882-00000055 (Center for Responsible AI), and by Fundacão para a Ciência e Tecnologia through contract UIDB/50008/2020.

7 Appendix↩︎

Aside from 7.1 giving more detail on the construction of adaptive prediction sets, we use this appendix to bundle more details about experiments and their results. 7.2 details the procedure to determine the temperature in 5 . We present more results from the experiments in 4.1 in 7.3.
We illustrate the overall algorithm in 7.5 and estimate environmental impact of our work in 7.6.

7.1 Adaptive Prediction Sets↩︎

Here we provide a more formal definition of the adaptive prediction sets. Let \(\pi\) be a permutation function mapping all possible output tokens \(\{1, \ldots, C\}\) to the indices of a permuted version of the set, for which tokens are sorted by their probability under the model, descendingly. We define the non-conformity score as \[\label{eq:adaptive-non-conformity-score} s_i = \sum_{j=1}^{\pi(y_t)} p_{\mathop{\mathrm{\boldsymbol{\theta}}}}\big(\pi^{-1}(j)\big|\mathop{\mathrm{\mathbf{x}}}, y_{<t}\big).\tag{7}\] Since we only include the cumulative mass up until the gold label, the summation stops at \(\pi(y)\). The prediction sets are then defined as \[\begin{align} \label{eq:adaptive-prediction-set} \mathcal{C}(\mathop{\mathrm{\mathbf{x}}}^*, y^*_{<t}) & = \Big\{\pi^{-1}(1), \ldots, \pi^{-1}(\hat{c})\Big\}, \end{align}\tag{8}\] with \(\hat{c} = \sup \{c^\prime\;|\;\sum_{j=1}^{c^\prime} p_{\mathop{\mathrm{\boldsymbol{\theta}}}}(\pi^{-1}(j)\,|\,\mathop{\mathrm{\mathbf{x}}}^*, y^*_{<t}) < \hat{q} \} + 1\), where we add one extra class to avoid empty sets.

7.2 Temperature Search↩︎

In order to determine the temperature used in 5 for the different distance metrics in 1, we adopt a variation of a simple hill-climbing procedure. Given user-defined bounds for the temperature search \(\tau_\text{min}\) and \(\tau_\text{max}\), we sample an initial candidate \(\tau_0 \sim \mathcal{U}[\tau_\text{min}, \tau_\text{max}]\), and then evaluate the coverage of the method given the candidate on the first \(100\) batches of the calibration dataset. The next candidate then is obtained via

\[\begin{align} \tau_{t+1} & = \tau_{t} + \eta \cdot \varepsilon \cdot \text{sgn}\big(1 - \alpha - \text{Coverage}(\tau_t)\big); \nonumber \\ \varepsilon & \sim \mathcal{N}(0, \tau_\text{max} - \tau_\text{min}), \end{align}\]

where \(\eta\) is a predefined step size (in our case \(0.1\)) and \(\text{Coverage}(\tau_t)\) the achieved coverage given a candidate \(\tau_t\). The final temperature is picked after a fixed number of steps (\(t=20\) in our work) based on the smallest difference between achieved and desired coverage.
Overall, we found useful search ranges to differ greatly between datasets, models, and distance metrics, as illustrated by the reported values in 1 and 2. In general, the stochastic hill-climbing could also be replaced by a grid search, even though we sometimes found the best temperature to be “hidden” in a very specific value range. It also has to be noted that temperature for the \(l_2\) distance is the highest by far since FAISS returns squared \(l_2\) distances by default.

7.3 Additional Coverage Results↩︎

We show additional plots illustrating the coverage per set size-bins in 4. We can see the counterparts for 2 using the larger M2M100\(_\text{(1.2B)}\) model in 4 (a) 4 (b): Instead of leveling off like for the smaller model, most prediction set sizes are either in a very small range or in a size of a few ten thousand. In 4 (c) 4 (d), we show similar plots for the two different OPT model sizes. Since in both cases, most prediction set sizes are rather small, we zoom in on the the sizes from \(1\) to \(100\). Here, we can observe a similar behavior to the smaller M2M100\(_\text{(400m)}\), gradually leveling off. We do not show similar plots for other distance metrics as they show similar trends.

7.4 Impact of Coverage Threshold and Neighborhood Size Choice↩︎

In this section, we present experiments surrounding the two most pivotal parameters of our method: The desired confidence level \(\alpha\), as well as the number of neighbors.

7.4.0.1 Coverage Threshold.

In 4, we investigate the impact of different values on \(\alpha\) on our evaluation metrics. We show that the increase in \(\alpha\) does indeed produce the expected decrease in coverage, however with a certain degree of overcoverage for the de \(\rightarrow\) en MT and the LM task. The loss in coverage always goes hand in hand with a decrease in the average prediction set width as well, as the model can allow itself to produce tighter prediction sets at the cost of higher miscoverage. As this also produces bin in which all contained instances are uncovered, this produces zero values for the SCC, while we cannot discern clear trends for the ECG.

7.4.0.2 Neighborhood Size.

In 5, we vary the effect of the chosen neighborhood size (with \(100\) being the value we use in our main experiments). We make the following, interesting observations: Coverage on the MT task seems to decrease with an increase in the neighborhood size as prediction set widths get smaller on average, with a neighborhood size around \(100\) striking a balance between coverage, width, computational cost and SCC / ECG. For LM, coverage seems to be mostly constant, with prediction set width hitting an inflection point for \(100\) neighbors. We speculate that initially there might be a benefit to considering more neighbors to calibrate \(\hat{q}\), but that considering too large neighborhoods might introduce extra noise. While we found \(100\) to be a solid choice for the purpose of our experiments, we leave more principled ways to determine the neighborhood size to future work.

7.5 Algorithm↩︎

We show the algorithm that was schematically depicted in 1 in pseudo-code in 5. It mostly requires that we have pre-generated a datastore of latent representations of the model on a held-out set along with their non-conformity scores (in our case, using the score defined in 7 and the FAISS [54] as the datastore architecture). Furthermore, we need to have determined an appropriate value for the temperature \(\tau\) in advance (see 7.2). Then, the algorithm involves the following steps:

1. Extract the latent encoding for the current time step \(\mathop{\mathrm{\mathbf{z}}}_t\) from the model. Even though different options are imaginable, we utilize the activations of the uppermost layer.

2. Retrieve \(K\) neighbors and their corresponding non-conformity scores from the datastore.

3. Compute the weights \(w_k\) based on the squared \(l_2\) distance between \(\mathop{\mathrm{\mathbf{z}}}_t\) and its neighbors in the datastore and normalize the weights to obtain \(\tilde{w}_k\).

4. Use 3 to find the quantile \(\hat{q}\).

5. Use \(\hat{q}\) to create prediction sets, for instance the adaptive prediction sets defined in 8 .

6. Finally, generate the new token \(y_t\) by sampling from the prediction set.

The main computational bottleneck of this algorithm is the retrieval process that fetches the closest neighbors from the datastore during every generation step. However, while not explored further in this work, there are some potential avenues to reduce this load: On the one hand, works such as [18], [95] have demonstrated ways to reduce the computational load of \(k\)-NN based approaches. On other hand, we treat the number of neighbors \(K\) fixed during every generation step. However, it seems intuitive that the number of neighbors necessary to create good prediction sets would not be the same for all tokens. Future research could explore setting \(K\) dynamically during every time step, thus reducing the overall slowdown.

7.6 Environmental Impact↩︎

We track the carbon emissions produced by this work using the codecarbon tracking tool [96]–[98]. The carbon efficiency was estimated to be \(0.12\) kgCO\(_2\)eq / kWh. 159.5 hours of computation were performed on a NVIDIA RTX A6000. Total emissions are estimated to be \(6.99\) kgCo2eq. All of these values are upper bound including debugging as well as failed or redundant runs, and thus any replication of results will likely be shorter and incur fewer carbon emissions.

7.7 Qualitative Analysis↩︎

In 6 7 8 9, we show a few samples from the different methods on our used datasets. We could observe some general patterns from the given and further samples: For the translation task, our method and other sampling techniques seem to sometimes prefer literal but unfitting translations. As expected for text generation, greedy sampling produces the worst translations, often getting stuck in loops, while beam search produces incoherent outputs. Interestingly, non-exchangeable conformal sampling and conformal sampling seem to often produce similar sampling before diverging.

References↩︎