# Aligning Language Models to Explicitly Handle Ambiguity

Hyuhng Joon Kim1, Youna Kim1, Cheonbok Park2, Junyeob Kim1,
Choonghyun Park1, Kang Min Yoo1 2 3, Sang-goo Lee1, Taeuk Kim4 *
1Seoul National University, 2NAVER Cloud
3NAVER AI LAB, 4Hanyang University
{heyjoonkim, anna9812, juny116, pch330<EMAIL_ADDRESS>
{cbok.park<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

In spoken languages, utterances are often shaped to be incomplete or vague for
efficiency. This can lead to varying interpretations of the same input, based
on different assumptions about the context. To ensure reliable user-model
interactions in such scenarios, it is crucial for models to adeptly handle the
inherent ambiguity in user queries. However, conversational agents built upon
even the most recent large language models (LLMs) face challenges in
processing ambiguous inputs, primarily due to the following two hurdles: (1)
LLMs are not directly trained to handle inputs that are too ambiguous to be
properly managed; (2) the degree of ambiguity in an input can vary according
to the intrinsic knowledge of the LLMs, which is difficult to investigate. To
address these issues, this paper proposes a method to align LLMs to explicitly
handle ambiguous inputs. Specifically, we introduce a proxy task that guides
LLMs to utilize their intrinsic knowledge to self-disambiguate a given input.
We quantify the information gain from the disambiguation procedure as a
measure of the extent to which the models perceive their inputs as ambiguous.
This measure serves as a cue for selecting samples deemed ambiguous from the
models’ perspectives, which are then utilized for alignment. Experimental
results from several question-answering datasets demonstrate that the LLMs
fine-tuned with our approach are capable of handling ambiguous inputs while
still performing competitively on clear questions within the task.

Aligning Language Models to Explicitly Handle Ambiguity

Hyuhng Joon Kim1, Youna Kim1, Cheonbok Park2, Junyeob Kim1, Choonghyun Park1,
Kang Min Yoo1 2 3, Sang-goo Lee1, Taeuk Kim4 * 1Seoul National University,
2NAVER Cloud 3NAVER AI LAB, 4Hanyang University {heyjoonkim, anna9812,
juny116, pch330<EMAIL_ADDRESS>{cbok.park,
<EMAIL_ADDRESS><EMAIL_ADDRESS>

**footnotetext: Corresponding author.

##  1 Introduction

Large Language Models (LLMs) (Ouyang et al., 2022; Team et al., 2023; Achiam
et al., 2023) have demonstrated remarkable capabilities in text generation,
proving particularly effective for question-answering (QA) tasks (Zhang et
al., 2023; Etezadi and Shamsfard, 2023). QA systems in the wild are frequently
confronted with unexpected inputs from users, such as unanswerable (Kim et
al., 2023b; Yin et al., 2023) or ambiguous questions (Cole et al., 2023; Lee
et al., 2023; Kim et al., 2023a). To build a reliable user-friendly model, it
is essential for the model to robustly handle such inputs. In this work, we
seek to extend the scope of research to effectively handle invalid inputs.
Specifically, we focus on managing "ambiguity" (Gleason, 1963; Mackay and
Bever, 1967), which poses a significant challenge in Natural Language
Processing (NLP) (Jurafsky, 1996).

Figure 1:  An example of an ambiguous query from AmbigQA (Min et al., 2020).
The phrase "national championship" poses diverse denotations, causing
ambiguity. Model possessing various related knowledge of the query could
perceive it as ambiguous (left). On the other hand, if the model does not have
sufficient related knowledge (right), the query can be perceived as
unambiguous. Thus, the degree of ambiguity perceived by the model may vary
even with identical inputs.

Ambiguity refers to cases where an expression conveys multiple denotations
(Wasow et al., 2005). Users may pose queries with clear intentions that,
possibly due to insufficient domain knowledge, result in ambiguous requests.
If the model arbitrarily responds to such ambiguity, there is a risk of
misinterpreting the user’s original intent, potentially harming the model’s
reliability. This is especially pronounced in sensitive domains such as legal
(Schane, 2002; Choi, 2024) or medical (Stevenson and Guo, 2010; Gyori et al.,
2022) domains, where misinterpretations can lead to serious drawbacks. Despite
the importance, approaches to robustly manage ambiguity are still
significantly unexplored. In this paper, we endeavor to utilize the model’s
intrinsic knowledge to align the model in a manner that effectively handles
ambiguity.

Properly processing ambiguous inputs is challenging primarily due to the
following two hurdles. Firstly, models are not directly trained to explicitly
express ambiguity. Even if a model perceives ambiguity, it is challenging to
verify the recognition without explicit feedback. The second challenge is that
the degree of ambiguity for the query can vary depending on the intrinsic
knowledge of the model. Consider the scenario depicted in Figure 1. The
initial query is ambiguous as the phrase "national championship" poses various
denotations, such as "national tennis championship" or "national golf
championship". If a model possesses comprehensive knowledge across the
possible denotations, it is plausible for the model to recognize the ambiguity
(left). However, if the model’s knowledge is limited to "national tennis
championship", it would perceive the query as unambiguous (right). Therefore,
it is essential to verify whether the input is deemed ambiguous from the
model’s point of view.

To overcome these issues, this paper proposes a method to align models to
explicitly handle ambiguous queries. Specifically, we design a proxy task that
guides the model to self-disambiguate a given query by utilizing its intrinsic
knowledge. Then, we quantify the information gain from the disambiguation as
an implicit measure of the extent to which the models perceive their inputs as
ambiguous. This measure serves as a cue for selecting samples deemed ambiguous
from the model’s perspective, which are then utilized for alignment.
Experimental results from several QA datasets demonstrate that the alignment
process enables the model to properly clarify ambiguous inputs while
maintaining its inherent capabilities. The findings underscore the value of
assessing the perceived ambiguity, rather than relying solely on the ground-
truth ambiguity. Furthermore, to provide a comprehensive framework for
assessing ambiguity, we construct a new dataset dubbed AmbigTriviaQA. The
dataset facilitates a more extensive evaluation of models’ robustness in
addressing ambiguity, thus contributing to the further expansion of related
research.

##  2 Related Work

#### Ambiguity in NLP

An expression is defined as ambiguous if it has two or more distinct
denotations (Wasow et al., 2005). Ambiguity challenges NLP applications by
obscuring the intended meaning of expressions, leading to difficulties in
accurately performing specific tasks. Efforts addressing this issue span
across various domains, including machine translation (Pilault et al., 2023),
coreference resolution (Poesio and Artstein, 2005; Yuan et al., 2023), and
natural language inference (Liu et al., 2023).

The challenge intensifies in the scope of QA as ambiguous questions may yield
various answers, potentially not aligning with the user’s initial intent. Min
et al. (2020) introduce the AmbigQA dataset to tackle ambiguity in open-domain
QA and Stelmakh et al. (2022) expands it to long-form generation. Furthermore,
Cole et al. (2023) discovered that quantifying sampling repetition presents a
reliable uncertainty measure for ambiguity, while Kim et al. (2023a) generates
tree-of-clarification (ToC) that refines ambiguity within the inputs. As we
share the goal of handling ambiguity, we adopt a novel approach of directly
aligning the model to address ambiguity.

#### Alignment of LLMs

LLMs are fundamentally trained through causal language modeling, a process
essential for understanding and generating text of high fluency and
consistency. To better harness these models, approaches have been developed to
align them with human preferences (Leike et al., 2018; Ji et al., 2023b). This
has taken various forms, notably through Reinforcement Learning from Human
Feedback (RLHF) (Ouyang et al., 2022; Bai et al., 2022a; Chakraborty et al.,
2024), as well as Supervised Fine-tuning (SFT) (Dong et al., 2023; Yang et
al., 2023; Zhou et al., 2024).

Previous works focused on preferences such as helpfulness (Ding et al., 2023;
Köpf et al., 2023; Xu et al., 2024) and safety (Bai et al., 2022b; Ji et al.,
2023a; Liu et al., 2024b). Recent studies have concentrated on the factuality
(Yang et al., 2023; Tian et al., 2024), avoiding hallucinations. Building on
this foundation, our research extends the scope and aims to align models to
effectively understand and handle ambiguities, which is a relatively
unexplored area within the field of model alignment. This stands in contrast
to previous methods which typically bypass the nuanced interpretation of
ambiguous contexts inherent in language.

#### Data Quality Control for Alignment

Data-centric AI (Chu et al., 2016; Majeed and Hwang, 2023; Kumar et al., 2024)
emphasizes the importance of data quality in training models. In the context
of the instruction following, LIMA (Zhou et al., 2024) demonstrates that
models can be effectively aligned even with 1,000 human-curated samples.
Similarly, AlpaGasus (Chen et al., 2024) shows effective alignment can be
achieved by utilizing only a small subset of the Alpaca dataset (Taori et al.,
2023) selected by ChatGPT. Various approaches for data selection have been
explored, including those based on pre-defined quality factors such as length
and complexity (Liu et al., 2024a), and utilizing gradient similarity from
validation sets as a selection criterion (Xia et al., 2024). In this paper, we
distinctly define data quality to effectively align models to handle
ambiguity. To do so, we make use of the model’s perceived ambiguity as an
implicit cue for data quality.

##  3 Methodology

Figure 2:  The overall process of our four-stage alignment pipeline. We
initially filter samples that the model cannot explicitly handle (Stage 1),
and self-disambiguate them to measure the information gain (Stage 2). Samples
with high information gain are deemed ambiguous and utilized for supervised
fine-tuning (Stage 3 $\&$ 4).

The objective of our approach is to align models to explicitly handle
potentially ambiguous inputs leveraging intrinsic model knowledge. To this
end, we propose a four-stage alignment pipeline, depicted in Figure 2. In this
section, we first formulate the problem and describe each stage in detail.

#### Problem Formulation

The goal of a QA task is to generate a factually correct answer $y$, given an
unambiguous input $x_{\text{unambig}}$, a pre-defined inference template
$t(\cdot)$, and a language model $M$. As we expand our input scope to
ambiguous queries $x_{\text{ambig}}$, the model is expected to generate a
clarification request $y_{\text{clarify}}$222 We have considered various
approaches to handle ambiguity but were concluded to be impractical.
Arbitrarily offering one of the valid answers may fail to reflect the user’s
intent, and presenting all possible answers is often impractical due to the
potentially vast number of valid answers.  for $x_{\text{ambig}}$ to resolve
the ambiguity, where the user is best positioned to clarify their intent.

###  3.1 Explicit Prediction Stage

This initial stage involves assessing whether the model can appropriately
handle each sample and identifying samples that the model currently fails to
explicitly manage. By comparing the model’s prediction with the ground-truth
label, samples are categorized based on their response accuracy. We collect
$n$ correct samples, which the model can properly handle, as
$D_{\text{correct}}=\\{(x^{i}_{\text{correct}},y^{i}_{\text{correct}})\\}^{n}_{i=1}$
and incorrect samples are classified as $D_{\text{incorrect}}$.

###  3.2 Implicit Ambiguity Detection Stage

The objective of this stage is to identify samples that the model perceives as
ambiguous from $D_{\text{incorrect}}$. Given that it is challenging for the
model to explicitly express ambiguity, we construct a proxy task to estimate
the ambiguity from the model’s point of view.

The proxy task is designed to self-disambiguate $x$ and implicitly measure the
perceived ambiguity. Specifically, the model is first prompted to generate a
disambiguation $\hat{x}_{\text{disambig}}$ for the input $x$. In this process,
the model leverages its intrinsic knowledge related to $x$ and generates
further details. If $x$ lacks specifications and the model possesses related
knowledge necessary to compensate, then $\hat{x}_{\text{disambig}}$ would
yield a higher certainty (lower entropy) for the model. On the other hand, if
$x$ requires no specification or the model lacks the necessary knowledge,
$\hat{x}_{\text{disambig}}$ would exhibit a similar level of uncertainty to
$x$. To quantify the uncertainty associated with $x$ and
$\hat{x}_{\text{disambig}}$, we employ the model’s average entropy (Malinin
and Gales, 2021; Abdar et al., 2021). Formally, the entropy of an output
distribution is defined as follows:

$\mathcal{H}_{x,i}=-\sum_{v\in\mathcal{V}}p_{x,i}(v)\log p_{x,i}(v)$ (1)

where $p_{x,i}(v)$ is the probability of the $i^{\text{th}}$ token $v$ of a
sentence $x$ from the full vocabulary set $\mathcal{V}$. The average entropy
for $x$ can be defined as:

$\mathcal{H}_{x}=\frac{1}{I}\sum_{i}\mathcal{H}_{x,i}$ (2)

where $x$ is composed of $I$-tokens. We quantify the changes in input
uncertainty by the difference in average entropy, which we define as
information gain (Infogain). The Infogain from the disambiguation can be
defined as the following:

$\text{{Infogain}}_{x,\hat{x}_{\text{disambig}}}=\mathcal{H}_{x}-\mathcal{H}_{\hat{x}_{\text{disambig}}}$
(3)

If the disambiguation results in a meaningful specification by utilizing
intrinsic knowledge, a substantial Infogain would be measured, suggesting that
the model considers $x$ as ambiguous. Conversely, a negligible Infogain
indicates that the model does not perceive $x$ as ambiguous. Samples with
Infogain greater than the threshold $\epsilon$ are classified as ambiguous,
denoted as $x_{\text{ambig}}$.

###  3.3 Data Construction Stage

In this stage, we construct datasets for the alignment process. This involves
labeling $m$ samples identified as ambiguous and constructing an ambiguous
dataset
$D_{\text{ambig}}=\\{(x^{j}_{\text{ambig}},y_{\text{clarify}})\\}^{m}_{j=1}$.
$y_{\text{clarify}}$ serves as the ground-truth label for ambiguous samples,
which are randomly selected from pre-defined clarification requests stipulated
in Appendix D. To prevent the potential loss of the model’s existing
knowledge, we also incorporate $D_{\text{correct}}$ for training. We balance
the number of samples from both datasets so that $n=m$. The final training
dataset is thus established as $D=D_{\text{correct}}+D_{\text{ambig}}$.

###  3.4 Supervised Fine-tuning (SFT) Stage

Utilizing the dataset $D=\\{(x^{k},y^{k})\\}^{n+m}_{k=1}$, the model is
trained to generate ground-truth label $y$ for input $x$, employing the
identical inference template $t(\cdot)$. The model $M$ with parameter $\theta$
is trained as follows:

$\min_{\theta}\sum_{(x,y)\in D}\sum_{i=1}^{|y|}-\log
M_{\theta}(y_{i}|y_{<i},t(x))$ (4)

##  4 Experimental Setting

###  4.1 Datasets

The capability of the model to perform within the trained domain is pivotal.
However, its ability to generalize to out-of-distribution (OOD) is essential
for real-world applicability, as queries that deviate from the training data
are frequently confronted in the wild. To this end, we employ one training
dataset and three OOD test sets to evaluate in diverse domains. All the
datasets include both ambiguous and unambiguous queries.

#### AmbigQA

Introduced by Min et al. (2020), AmbigQA is a derivative of the Natural
Questions dataset (Kwiatkowski et al., 2019), designed to verify data points
deemed ambiguous. The dataset covers diverse sources of ambiguity such as
event and entity references. We set AmbigQA as the in-domain dataset and
utilize it for training.

#### SituatedQA

SituatedQA (Zhang and Choi, 2021) specifically focuses on temporal and
geographic ambiguity from the input query. As the cause of ambiguity and its
construction process are distinct, we assess performance on the temporal split
and the geographic split separately, denoted as Temp and Geo, respectively.

#### AmbigTriviaQA

Since there are limited datasets for evaluating ambiguity in open-domain QA,
we construct a new dataset, namely AmbigTriviaQA. By taking questions from the
widely-used TriviaQA dataset (Joshi et al., 2017), we prompt gpt-3.5-turbo to
ambiguate the initial query and verify the results. More details on dataset
construction are described in Appendix B.

###  4.2 Baselines

To assess the effectiveness of our approach, we establish two sets of
baselines: inference-only methods and trained methods. Further implementation
details are described in Appendix A.

#### Inference-Only Methods

Inference-only methods address ambiguity by directly prompting the model. We
employ naïve prompting (Naïve) as a fundamental baseline, applying a simple QA
prompt. Furthermore, we explore ambiguity-aware prompting (Ambiguity-aware),
which additionally provides instructions on handling ambiguity. We also
examine Sample Repetition (Cole et al., 2023) by measuring the consistency of
the sampled generations. Finally, we evaluate Self-Ask (Amayuelas et al.,
2023), where the model initially generates an answer and subsequently
determines the ambiguity based on the generation.

#### Trained Methods

Given the lack of directly comparable prior work, we compare a fine-tuned
baseline wherein the model is trained with the in-domain training set. Full-
set applies the full in-domain training dataset and Subset is trained on a
randomly selected subset of size $n+m$, which is equivalent in size to our
approach. Additionally, we compare Honesty-tuned (Yang et al., 2023), which
takes a similar approach utilizing an implicit measure named "expected
accuracy" to estimate the model’s factuality. The expected accuracy is
measured as the average accuracy of sampled prediction for a single input
query. We have revised the method so that incorrect samples based on "expected
accuracy" are selected from the ground-truth ambiguous samples. Although
Honesty-tuned shares similarities with our approach in the way of estimating
the model’s knowledge with an implicit measure, a notable distinction lies in
our main focus on handling ambiguity beyond factuality.

###  4.3 Evaluation Metrics

As we expand the input scope to possibly ambiguous questions, the model should
be capable of handling both unambiguous and ambiguous queries simultaneously.
Therefore, we employ two widely used evaluation metrics to assess performance
on both types of inputs. A successful alignment should preserve the model’s
capability to handle unambiguous inputs while successfully managing ambiguous
queries. All evaluations are conducted by comparing the greedy generation to
the ground truth.

#### Unambiguous Accuracy (Unambig. Acc.)

While expanding the task scope, it remains crucial for the model to preserve
the ability to handle unambiguous inputs. Thus, our analysis persists in
exclusively evaluating the model’s accuracy in processing unambiguous queries.
We measure the quality of the generation by employing
RougeL333https://huggingface.co/spaces/evaluate-metric/rouge (Lin and Och,
2004) with all the possible answers, where the prediction is regarded as
correct if the score is above 0.3.

#### Ambiguity Detection F1-score (Ambig. F1)

The model should be capable of detecting ambiguity and generating
clarification requests for ambiguous inputs. However, especially for trained
methods, models may exhibit biased predictions toward clarification requests.
Taking these aspects into account, we evaluate the model’s ambiguity detection
capability with F1-score, which captures both the precision and recall of
prediction, offering a balanced view of the model’s ambiguity detection
performance. Further details on the detection process are described in
Appendix C.

###  4.4 Implementation Details

For our experiments, we utilize Llama2 7B $\&$ 13B (Touvron et al., 2023) and
Mistral 7B (Jiang et al., 2023). We utilized QLoRA (Dettmers et al., 2023) to
facilitate efficient training. Implementation details are stipulated in
Appendix D.

Method |  | # Train
---
Samples
AmbigQA |  | SituatedQA
---
(Geo)
| SituatedQA
---
(Temp)
| Ambig-
---
TriviaQA
|  |  | Unambig.
---
Acc.
| Ambig.
---
F1
| Unambig.
---
Acc.
| Ambig.
---
F1
| Unambig.
---
Acc.
| Ambig.
---
F1
| Unambig.
---
Acc.
| Ambig.
---
F1
Llama2 7B
Naïve | 0 | 28.43 | 0.00 | 22.53 | 0.00 | 21.72 | 0.00 | 62.46 | 0.00
Ambiguity-aware | 0 | 4.94 | 68.95 | 3.95 | 32.44 | 1.72 | 35.53 | 22.63 | 61.03
Sample Repetition | 0 | 5.42 | 74.96 | 3.56 | 34.43 | 4.40 | 38.43 | 29.58 | 65.88
Self-Ask | 0 | 26.75 | 13.41 | 21.54 | 8.18 | 19.68 | 18.48 | 61.29 | 3.84
Honesty-tuned | 3,088 | 18.19 | 68.00 | 9.49 | 36.56 | 7.98 | 37.57 | 52.42 | 54.60
Subset | 3,088 | 29.16 | 64.23 | 23.72 | 36.85 | 13.95 | 36.74 | 51.71 | 55.70
Full-set | 10,036 | 37.11 | 59.98 | 26.09 | 41.15 | 18.93 | 35.38 | 58.25 | 49.96
Ours | 3,088 | 27.23 | 63.69 | 24.51 | 42.05 | 21.90 | 40.77 | 53.41 | 61.34
Mistral 7B
Naïve | 0 | 13.01 | 0.00 | 7.51 | 0.00 | 10.91 | 0.00 | 37.52 | 0.00
Ambiguity-aware | 0 | 9.76 | 54.74 | 2.17 | 26.01 | 5.33 | 22.48 | 27.87 | 44.94
Sample Repetition | 0 | 4.70 | 51.21 | 2.57 | 29.25 | 2.15 | 29.64 | 25.54 | 32.54
Self-Ask | 0 | 13.01 | 0.00 | 7.51 | 0.00 | 10.91 | 0.00 | 37.52 | 0.00
Honesty-tuned | 1,382 | 16.51 | 69.28 | 11.66 | 33.84 | 8.09 | 39.16 | 41.70 | 65.19
Subset | 1,382 | 33.49 | 60.91 | 29.05 | 35.64 | 19.11 | 37.38 | 58.87 | 54.61
Full-set | 10,036 | 43.73 | 62.85 | 24.11 | 40.81 | 26.76 | 24.14 | 65.04 | 49.63
Ours | 1,382 | 37.23 | 50.31 | 32.21 | 42.18 | 35.74 | 40.17 | 58.14 | 58.93
Llama2 13B
Naïve | 0 | 31.33 | 0.00 | 22.53 | 0.00 | 22.90 | 0.00 | 60.49 | 0.14
Ambiguity-aware | 0 | 3.37 | 70.44 | 3.16 | 33.10 | 2.22 | 36.66 | 16.96 | 64.17
Sample Repetition | 0 | 10.00 | 71.10 | 6.32 | 32.85 | 9.45 | 37.87 | 39.41 | 63.40
Self-Ask | 0 | 31.33 | 0.00 | 22.53 | 0.00 | 22.90 | 0.00 | 60.49 | 0.14
Honesty-tuned | 3,216 | 17.83 | 68.57 | 3.16 | 33.29 | 3.58 | 38.03 | 48.13 | 60.39
Subset | 3,216 | 35.18 | 62.89 | 23.12 | 36.19 | 20.50 | 39.00 | 59.72 | 57.26
Full-set | 10,036 | 43.61 | 62.87 | 23.12 | 38.37 | 19.68 | 24.07 | 66.80 | 48.81
Ours | 3,216 | 37.83 | 58.15 | 24.51 | 41.59 | 24.36 | 41.09 | 63.74 | 55.23

Table 1:  Experimental results of in-domain dataset and three OOD datasets. We
report unambiguous accuracy (Unambig. Acc.) and ambiguity detection F1-score
(Ambig. F1). For each dataset, the best method is highlighted in bold and the
second-best method is underlined. Our method demonstrates comparable
performance in-domain and outperforms all the baselines in the OOD setting.

##  5 Experimental Results

The main results of our experiments are presented in Table 1. Inference-only
methods exhibit a pronounced deficiency in handling ambiguous queries.
Specifically, Naïve establish poor performance in responding to ambiguous
queries, resulting in a notably low F1-score. Ambiguity-aware demonstrates a
strong bias towards clarification requests, as it achieves a relatively high
F1-score at the expense of accuracy. Similarly, Sample Repetition exhibits a
substantial trade-off between F1-score and accuracy. Self-Ask displays a
subpar F1-score, indicating that it is challenging to resolve ambiguity by
explicitly "self-asking" the model.

Trained methods exhibit enhanced performance overall compared to inference-
only approaches. Honesty-tuned struggles to handle ambiguity, as it also
demonstrates biased detection performance. This is likely because the implicit
measure from Honesty-tuned can be influenced by various factors but not
specifically ambiguity. The results underscore the necessity of distinct
methods for perceiving ambiguity. Compared to Honesty-tuned, Subset exhibits
relatively balanced performance across both metrics. Full-set demonstrates the
most superior performance among the baselines, particularly in the in-domain
setting, as it has access to the ground-truth ambiguity.

Our approach yields comparable results in in-domain and demonstrates superior
OOD performances. Despite employing identical inference templates as Naïve,
our method demonstrates equal or improved unambiguous accuracy. This indicates
the effectiveness of our alignment in managing ambiguity while preserving the
inherent capabilities of the model. We can also observe an improvement in
unambiguous accuracy, particularly in SituatedQA splits. It is especially
surprising given that our method was trained on $D_{\text{correct}}$, which
the model is already capable of handling. Compared to Full-set, we note a
slight decline in the in-domain performance, an expected result given that
Full-set is optimized with ground-truth ambiguity of the in-domain data.
However, our method outperforms Full-set across OOD datasets in F1-score up to
17 points. This discrepancy underscores the effectiveness of utilizing
perceived ambiguity for alignment, facilitating superior generalization and
robustness. The efficacy of leveraging only the data perceived ambiguous
(about 32% in the Llama2 family and 13% in Mistral) emphasizes the importance
of data quality over quantity (Zhou et al., 2024; Chen et al., 2024).

Method | AmbigQA |  | SituatedQA
---
(Geo)
|  | Unambig.
---
Acc.
| Ambig.
---
F1
| Unambig.
---
Acc.
| Ambig.
---
F1
Random | 29.40 | 64.74 | 18.38 | 36.23
Implicit | 29.52 | 62.23 | 24.90 | 41.19
Ours | 27.23 | 63.69 | 24.51 | 42.05
Method |  | SituatedQA
---
(Temp)
| Ambig-
---
TriviaQA
|  | Unambig.
---
Acc.
| Ambig.
---
F1
| Unambig.
---
Acc.
| Ambig.
---
F1
Random | 16.99 | 38.76 | 56.13 | 50.63
Implicit | 21.36 | 40.11 | 55.65 | 59.16
Ours | 21.90 | 40.77 | 53.41 | 61.34

Table 2:  Ablation results of ambiguous data selection. In-domain dataset is
highlighted in gray. The best method is in bold, and the second best method is
underlined. Results show that perceived ambiguity measured by Infogain is an
effective cue for data selection.

Method | VAR ($\uparrow$) | MCR ($\downarrow$) | OAP ($\uparrow$)
---|---|---|---
AmbigQA
Honesty-tuned | 73.81 | 44.49 | 20.48
Subset | 62.12 | 33.05 | 20.79
Full-set | 52.82 | 18.64 | 21.48
Ours | 65.19 | 30.51 | 22.65
SituatedQA (Geo)
Honesty-tuned | 93.80 | 58.77 | 19.34
Subset | 63.57 | 35.96 | 20.35
Full-set | 72.09 | 35.96 | 23.08
Ours | 89.15 | 28.95 | 31.67
SituatedQA (Temp)
Honesty-tuned | 85.73 | 56.34 | 18.71
Subset | 68.61 | 54.86 | 15.48
Full-set | 58.56 | 46.95 | 15.53
Ours | 72.95 | 36.08 | 23.31
AmbigTriviaQA
Honesty-tuned | 42.39 | 11.64 | 18.73
Subset | 48.10 | 17.57 | 19.83
Full-set | 38.93 | 9.81 | 17.56
Ours | 57.45 | 16.91 | 23.87

Table 3:  VAR, MCR, and OAP of trained methods. A high OAP is preferred, with a high VAR and a low MCR. The best performance is in bold, and the second best performance is underlined. Our approach demonstrates the best OAP across all datasets, with the least trade-off between VAR and MCR.  | Model
---
Prediction
| Ground
---
Truth
Type | Generated Text
Unambig. | Ambig. | $x$ |  Who sings don’t mess around with jim?
$\hat{x}_{\text{disambig}}$ |  Who sings don’t mess around with jim, in the 1960s?
Unambig. | Unambig. | $x$ |  Who is winner in bigg boss season 5 kannada?
$\hat{x}_{\text{disambig}}$ |  Who is the winner of the fifth season of the kannada version of the indian reality television series bigg boss?
Ambig. | Ambig. | $x$ |  How many jury members in a criminal trial?
$\hat{x}_{\text{disambig}}$ |  How many jury members are required in a criminal trial in the united states?
Ambig. | Unambig. | $x$ |  How many pages in a brave new world?
$\hat{x}_{\text{disambig}}$ |  How many pages are in the 1932 novel brave new world by aldous huxley?
Table 4:  Example of initial query $x$ and its disambiguation
$\hat{x}_{\text{disambig}}$ generated by Llama2 7B. Additional specification
from the model is in bold. Unambig. and Ambig. refers to Unambiguous and
Ambiguous, respectively.

##  6 Ablation Study

###  6.1 Impact of Infogain for Data Selection

For a deeper analysis of the influence of Infogain for data selection within
our pipeline, we conduct an ablation study by varying the criteria for
selecting ambiguous data. While maintaining the same $D_{\text{correct}}$ for
unambiguous samples, we alter the selection of $m$ samples labeled as
ambiguous. We compare the following data selection strategies:

  * •

Random Selection (Random) We randomly select $m$ ground-truth ambiguous
samples, without any consideration of Infogain.

  * •

Implicit Measure-based Selection (Implicit) We select top-$m$ samples with the
largest Infogain among those that are ground-truth ambiguous. It differs from
our approach as our method utilizes samples perceived as ambiguous, allowing
the potential inclusion of unambiguous samples.

Table 2 is the ablation results on Llama2 7B. For AmbigQA, baselines
leveraging ground-truth ambiguity slightly outperform our method, which is a
similar tendency from Section 5 where Full-set exhibit better in-domain
performance. However, across OOD datasets, our approach demonstrates
significantly superior performance. Specifically, Random demonstrates a
notable drop in F1-score by up to 10 points, illustrating the limitations of
simply utilizing the ground-truth ambiguity, which might not align with the
model’s perceived ambiguity. Furthermore, Implicit surpasses Random by up to 5
points in F1-score, validating the effectiveness of Infogain as a cue for data
selection. Finally, our approach outperforms Implicit across the majority of
metrics, even with unambiguous samples selected for training, again
highlighting the effectiveness of Infogain as the data selection measure.

###  6.2 Analysis on Sample-level Prediction Change

Our method is designed to align the model to generate clarification requests
for ambiguous queries. However, the process may lead to a potential trade-off,
where the model erroneously generates clarification requests for unambiguous
inputs that were previously well-handled. To assess this balance, we introduce
three metrics: Valid Alignment Rate (VAR) measures the proportion of ambiguous
samples incorrectly handled before alignment that are correctly addressed
post-alignment and Misaligned Clarification Rate (MCR) measures the rate of
correct unambiguous samples before training that erroneously generates
clarification requests after alignment. A high VAR is desirable, whereas a low
MCR is preferred simultaneously. Inspired by Yang et al. (2023), we
additionally define Overall Alignment Performance (OAP) that measures the
balance between VAR and MCR.

$\texttt{OAP}=\frac{\texttt{VAR}+(1-\texttt{MCR})}{2}$ (5)

Table 3 compares the results on Llama2 7B. Honesty-tuned exhibits high VAR but
poor MCR, implying a tendency to misinterpret known knowledge as ambiguous.
This aligns with the previous results where Honesty-tuned displays biased
generation towards clarification requests. On the other hand, Full-set and
Subset demonstrate a good MCR and a relatively low VAR. Our method performs
superior OAP, successfully addressing ambiguities (high VAR) while preserving
existing capabilities (low MCR).

##  7 Self-disambiguation Case Study

Table 4 demonstrates examples of initial query $x$ and its disambiguation
$\hat{x}_{\text{disambig}}$ generated by Llama2 7B. The first example is when
$x$ is inherently ambiguous, yet the model perceives it as unambiguous.
Specifically, the model generates hallucination ("in the 1960s") where the
song "don’t mess around with jim" was originally released in 1972. This non-
factual generation would not provide any information gain to the model,
classifying $x$ as ambiguous. In such a case, $x$ should be considered
"unknown" with no related knowledge within the model. The second and third
examples are correctly classified, as the model properly applies its intrinsic
knowledge to perceive ambiguity. Regardless of the quantity of additional
context generated, the model is capable of verifying its ambiguity. The last
example is a misclassification as ambiguous. Despite disambiguation provides
factually correct information ("1932 novel" and "by Aldous Huxley") for "brave
new world", we speculate that the misclassification may arise from the
existence of various media, such as movies and songs, sharing the title "brave
new world", leading to an erroneous integration of knowledge.

##  8 Conclusion

In this paper, we present a novel alignment pipeline designed to enhance the
ability of LLMs to address ambiguities within queries, leveraging the model’s
intrinsic knowledge. Our method employs an implicit measure, dubbed Infogain,
to quantify ambiguity as perceived by the model. Through alignment based on
the measure, the model learns to explicitly handle ambiguous as well as
unambiguous queries. Experimental results demonstrate the effectiveness of our
alignment, particularly pronounced in out-of-distribution scenarios. Results
indicate the importance of alignment based on the model’s perceived ambiguity.
Future work may explore the extension of this methodology to broader domains
and more complex types of ambiguities, further solidifying the role of LLMs in
managing the inherent uncertainty present in NLP tasks.

## Limitations

For the experiments, we explore the most widely used models for evaluation,
specifically Llama2 and Mistral. Despite this, a more comprehensive evaluation
encompassing a broader consideration of LLMs could have enriched our findings,
providing insights across different architectures and capabilities. Larger
models in scale could demonstrate different tendencies and should be explored
for future work. Furthermore, our work mainly focuses on supervised fine-
tuning (SFT) as the alignment method. However, alternative methods, such as
Reinforcement Learning from Human Preference (RLHF) (Ouyang et al., 2022) or
Direct Preference Optimization (DPO) (Rafailov et al., 2023) could offer
distinct advantages towards our objective. Finally, the experiments are mainly
focused on short-form QA tasks. The research scope could be expanded to long-
form generation tasks such as detailed reasoning.

## References

  * Abdar et al. (2021) Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U. Rajendra Acharya, Vladimir Makarenkov, and Saeid Nahavandi. 2021.  A review of uncertainty quantification in deep learning: Techniques, applications and challenges.  _Information Fusion_ , 76:243–297.
  * Achiam et al. (2023) Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. 2023.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_.
  * Amayuelas et al. (2023) Alfonso Amayuelas, Liangming Pan, Wenhu Chen, and William Wang. 2023.  Knowledge of knowledge: Exploring known-unknowns uncertainty with large language models.  _Preprint_ , arXiv:2305.13712.
  * Bai et al. (2022a) Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, Nicholas Joseph, Saurav Kadavath, Jackson Kernion, Tom Conerly, Sheer El-Showk, Nelson Elhage, Zac Hatfield-Dodds, Danny Hernandez, Tristan Hume, Scott Johnston, Shauna Kravec, Liane Lovitt, Neel Nanda, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, Ben Mann, and Jared Kaplan. 2022a.  Training a helpful and harmless assistant with reinforcement learning from human feedback.  _Preprint_ , arXiv:2204.05862.
  * Bai et al. (2022b) Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, Carol Chen, Catherine Olsson, Christopher Olah, Danny Hernandez, Dawn Drain, Deep Ganguli, Dustin Li, Eli Tran-Johnson, Ethan Perez, Jamie Kerr, Jared Mueller, Jeffrey Ladish, Joshua Landau, Kamal Ndousse, Kamile Lukosuite, Liane Lovitt, Michael Sellitto, Nelson Elhage, Nicholas Schiefer, Noemi Mercado, Nova DasSarma, Robert Lasenby, Robin Larson, Sam Ringer, Scott Johnston, Shauna Kravec, Sheer El Showk, Stanislav Fort, Tamera Lanham, Timothy Telleen-Lawton, Tom Conerly, Tom Henighan, Tristan Hume, Samuel R. Bowman, Zac Hatfield-Dodds, Ben Mann, Dario Amodei, Nicholas Joseph, Sam McCandlish, Tom Brown, and Jared Kaplan. 2022b.  Constitutional ai: Harmlessness from ai feedback.  _Preprint_ , arXiv:2212.08073.
  * Chakraborty et al. (2024) Souradip Chakraborty, Jiahao Qiu, Hui Yuan, Alec Koppel, Furong Huang, Dinesh Manocha, Amrit Singh Bedi, and Mengdi Wang. 2024.  Maxmin-rlhf: Towards equitable alignment of large language models with diverse human preferences.  _Preprint_ , arXiv:2402.08925.
  * Chen et al. (2024) Lichang Chen, Shiyang Li, Jun Yan, Hai Wang, Kalpa Gunaratna, Vikas Yadav, Zheng Tang, Vijay Srinivasan, Tianyi Zhou, Heng Huang, and Hongxia Jin. 2024.  Alpagasus: Training a better alpaca model with fewer data.  In _The Twelfth International Conference on Learning Representations_.
  * Choi (2024) Jonathan H Choi. 2024.  Measuring clarity in legal text.  _U. Chi. L. Rev._ , 91:1.
  * Chu et al. (2016) Xu Chu, Ihab F. Ilyas, Sanjay Krishnan, and Jiannan Wang. 2016.  Data cleaning: Overview and emerging challenges.  In _Proceedings of the 2016 International Conference on Management of Data_ , SIGMOD ’16, page 2201–2206, New York, NY, USA. Association for Computing Machinery.
  * Cole et al. (2023) Jeremy Cole, Michael Zhang, Daniel Gillick, Julian Eisenschlos, Bhuwan Dhingra, and Jacob Eisenstein. 2023.  Selectively answering ambiguous questions.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 530–543, Singapore. Association for Computational Linguistics.
  * Dettmers et al. (2023) Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. 2023.  Qlora: Efficient finetuning of quantized llms.  _Preprint_ , arXiv:2305.14314.
  * Ding et al. (2023) Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. 2023.  Enhancing chat language models by scaling high-quality instructional conversations.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 3029–3051, Singapore. Association for Computational Linguistics.
  * Dong et al. (2023) Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, KaShun SHUM, and Tong Zhang. 2023.  RAFT: Reward ranked finetuning for generative foundation model alignment.  _Transactions on Machine Learning Research_.
  * Etezadi and Shamsfard (2023) Romina Etezadi and Mehrnoush Shamsfard. 2023.  The state of the art in open domain complex question answering: a survey.  _Applied Intelligence_ , 53(4):4124–4144.
  * Gleason (1963) H.A. Gleason. 1963.  _Linguistics and English Grammar_.  H.A. Gleason jr.
  * Gyori et al. (2022) Benjamin M Gyori, Charles Tapley Hoyt, and Albert Steppi. 2022.  Gilda: biomedical entity text normalization with machine-learned disambiguation as a service.  _Bioinformatics Advances_ , 2(1):vbac034.
  * Ji et al. (2023a) Jiaming Ji, Mickel Liu, Juntao Dai, Xuehai Pan, Chi Zhang, Ce Bian, Boyuan Chen, Ruiyang Sun, Yizhou Wang, and Yaodong Yang. 2023a.  Beavertails: Towards improved safety alignment of LLM via a human-preference dataset.  In _Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track_.
  * Ji et al. (2023b) Jiaming Ji, Tianyi Qiu, Boyuan Chen, Borong Zhang, Hantao Lou, Kaile Wang, Yawen Duan, Zhonghao He, Jiayi Zhou, Zhaowei Zhang, Fanzhi Zeng, Kwan Yee Ng, Juntao Dai, Xuehai Pan, Aidan O’Gara, Yingshan Lei, Hua Xu, Brian Tse, Jie Fu, Stephen McAleer, Yaodong Yang, Yizhou Wang, Song-Chun Zhu, Yike Guo, and Wen Gao. 2023b.  Ai alignment: A comprehensive survey.  _Preprint_ , arXiv:2310.19852.
  * Jiang et al. (2023) Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2023.  Mistral 7b.  _Preprint_ , arXiv:2310.06825.
  * Joshi et al. (2017) Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. 2017.  TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension.  In _Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1601–1611, Vancouver, Canada. Association for Computational Linguistics.
  * Jurafsky (1996) Daniel Jurafsky. 1996.  A probabilistic model of lexical and syntactic access and disambiguation.  _Cognitive science_ , 20(2):137–194.
  * Kim et al. (2023a) Gangwoo Kim, Sungdong Kim, Byeongguk Jeon, Joonsuk Park, and Jaewoo Kang. 2023a.  Tree of clarifications: Answering ambiguous questions with retrieval-augmented large language models.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 996–1009, Singapore. Association for Computational Linguistics.
  * Kim et al. (2023b) Najoung Kim, Phu Mon Htut, Samuel R. Bowman, and Jackson Petty. 2023b.  (QA)2: Question answering with questionable assumptions.  In _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 8466–8487, Toronto, Canada. Association for Computational Linguistics.
  * Köpf et al. (2023) Andreas Köpf, Yannic Kilcher, Dimitri von Rütte, Sotiris Anagnostidis, Zhi Rui Tam, Keith Stevens, Abdullah Barhoum, Duc Nguyen, Oliver Stanley, Richárd Nagyfi, Shahul ES, Sameer Suri, David Glushkov, Arnav Dantuluri, Andrew Maguire, Christoph Schuhmann, Huu Nguyen, and Alexander Mattick. 2023.  Openassistant conversations - democratizing large language model alignment.  In _Advances in Neural Information Processing Systems_ , volume 36, pages 47669–47681. Curran Associates, Inc.
  * Kumar et al. (2024) Sushant Kumar, Sumit Datta, Vishakha Singh, Sanjay Kumar Singh, and Ritesh Sharma. 2024.  Opportunities and challenges in data-centric ai.  _IEEE Access_.
  * Kwiatkowski et al. (2019) Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. 2019.  Natural questions: A benchmark for question answering research.  _Transactions of the Association for Computational Linguistics_ , 7:452–466.
  * Lee et al. (2023) Dongryeol Lee, Segwang Kim, Minwoo Lee, Hwanhee Lee, Joonsuk Park, Sang-Woo Lee, and Kyomin Jung. 2023.  Asking clarification questions to handle ambiguity in open-domain QA.  In _Findings of the Association for Computational Linguistics: EMNLP 2023_ , pages 11526–11544, Singapore. Association for Computational Linguistics.
  * Leike et al. (2018) Jan Leike, David Krueger, Tom Everitt, Miljan Martic, Vishal Maini, and Shane Legg. 2018.  Scalable agent alignment via reward modeling: a research direction.  _CoRR_ , abs/1811.07871.
  * Lin and Och (2004) Chin-Yew Lin and Franz Josef Och. 2004.  Automatic evaluation of machine translation quality using longest common subsequence and skip-bigram statistics.  In _Proceedings of the 42nd Annual Meeting of the Association for Computational Linguistics (ACL-04)_ , pages 605–612, Barcelona, Spain.
  * Liu et al. (2023) Alisa Liu, Zhaofeng Wu, Julian Michael, Alane Suhr, Peter West, Alexander Koller, Swabha Swayamdipta, Noah Smith, and Yejin Choi. 2023.  We’re afraid language models aren’t modeling ambiguity.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 790–807, Singapore. Association for Computational Linguistics.
  * Liu et al. (2024a) Wei Liu, Weihao Zeng, Keqing He, Yong Jiang, and Junxian He. 2024a.  What makes good data for alignment? a comprehensive study of automatic data selection in instruction tuning.  In _The Twelfth International Conference on Learning Representations_.
  * Liu et al. (2024b) Zixuan Liu, Xiaolin Sun, and Zizhan Zheng. 2024b.  Enhancing llm safety via constrained direct preference optimization.  _Preprint_ , arXiv:2403.02475.
  * Loshchilov and Hutter (2019) Ilya Loshchilov and Frank Hutter. 2019.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_.
  * Mackay and Bever (1967) Donald G Mackay and Thomas G Bever. 1967.  In search of ambiguity.  _Perception & Psychophysics_, 2:193–200.
  * Majeed and Hwang (2023) A. Majeed and S. Hwang. 2023.  Data-centric artificial intelligence, preprocessing, and the quest for transformative artificial intelligence systems development.  _Computer_ , 56(05):109–115.
  * Malinin and Gales (2021) Andrey Malinin and Mark Gales. 2021.  Uncertainty estimation in autoregressive structured prediction.  In _International Conference on Learning Representations_.
  * Mangrulkar et al. (2022) Sourab Mangrulkar, Sylvain Gugger, Lysandre Debut, Younes Belkada, Sayak Paul, and Benjamin Bossan. 2022.  Peft: State-of-the-art parameter-efficient fine-tuning methods.  https://github.com/huggingface/peft.
  * Min et al. (2020) Sewon Min, Julian Michael, Hannaneh Hajishirzi, and Luke Zettlemoyer. 2020.  AmbigQA: Answering ambiguous open-domain questions.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 5783–5797, Online. Association for Computational Linguistics.
  * Ouyang et al. (2022) Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. 2022.  Training language models to follow instructions with human feedback.  _Advances in neural information processing systems_ , 35:27730–27744.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. 2019.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32.
  * Pilault et al. (2023) Jonathan Pilault, Xavier Garcia, Arthur Bražinskas, and Orhan Firat. 2023.  Interactive-chain-prompting: Ambiguity resolution for crosslingual conditional generation with interaction.  In _Proceedings of the 13th International Joint Conference on Natural Language Processing and the 3rd Conference of the Asia-Pacific Chapter of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 455–483, Nusa Dua, Bali. Association for Computational Linguistics.
  * Poesio and Artstein (2005) Massimo Poesio and Ron Artstein. 2005.  The reliability of anaphoric annotation, reconsidered: Taking ambiguity into account.  In _Proceedings of the Workshop on Frontiers in Corpus Annotations II: Pie in the Sky_ , pages 76–83, Ann Arbor, Michigan. Association for Computational Linguistics.
  * Rafailov et al. (2023) Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. 2023.  Direct preference optimization: Your language model is secretly a reward model.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Schane (2002) Sanford Schane. 2002.  Ambiguity and misunderstanding in the law.  _T. Jefferson L. Rev._ , 25:167.
  * Stelmakh et al. (2022) Ivan Stelmakh, Yi Luan, Bhuwan Dhingra, and Ming-Wei Chang. 2022.  ASQA: Factoid questions meet long-form answers.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pages 8273–8288, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
  * Stevenson and Guo (2010) Mark Stevenson and Yikun Guo. 2010.  Disambiguation in the biomedical domain: the role of ambiguity type.  _Journal of biomedical informatics_ , 43(6):972–981.
  * Taori et al. (2023) Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. 2023.  Stanford alpaca: An instruction-following llama model.  https://github.com/tatsu-lab/stanford_alpaca.
  * Team et al. (2023) Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. 2023.  Gemini: a family of highly capable multimodal models.  _arXiv preprint arXiv:2312.11805_.
  * Tian et al. (2024) Katherine Tian, Eric Mitchell, Huaxiu Yao, Christopher D Manning, and Chelsea Finn. 2024.  Fine-tuning language models for factuality.  In _The Twelfth International Conference on Learning Representations_.
  * Touvron et al. (2023) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. 2023.  Llama 2: Open foundation and fine-tuned chat models.  _Preprint_ , arXiv:2307.09288.
  * Wasow et al. (2005) Thomas Wasow, Amy Perfors, and David Beaver. 2005.  The puzzle of ambiguity.  _Morphology and the web of grammar: Essays in memory of Steven G. Lapointe_ , pages 265–282.
  * Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. 2020.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pages 38–45, Online. Association for Computational Linguistics.
  * Xia et al. (2024) Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. 2024.  Less: Selecting influential data for targeted instruction tuning.  _Preprint_ , arXiv:2402.04333.
  * Xu et al. (2024) Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, Qingwei Lin, and Daxin Jiang. 2024.  WizardLM: Empowering large pre-trained language models to follow complex instructions.  In _The Twelfth International Conference on Learning Representations_.
  * Yang et al. (2023) Yuqing Yang, Ethan Chern, Xipeng Qiu, Graham Neubig, and Pengfei Liu. 2023.  Alignment for honesty.  _arXiv preprint arXiv:2312.07000_.
  * Yin et al. (2023) Zhangyue Yin, Qiushi Sun, Qipeng Guo, Jiawen Wu, Xipeng Qiu, and Xuanjing Huang. 2023.  Do large language models know what they don’t know? In _Findings of the Association for Computational Linguistics: ACL 2023_ , pages 8653–8665, Toronto, Canada. Association for Computational Linguistics.
  * Yuan et al. (2023) Yuewei Yuan, Chaitanya Malaviya, and Mark Yatskar. 2023.  AmbiCoref: Evaluating human and model sensitivity to ambiguous coreference.  In _Findings of the Association for Computational Linguistics: EACL 2023_ , pages 1023–1030, Dubrovnik, Croatia. Association for Computational Linguistics.
  * Zhang et al. (2023) Lingxi Zhang, Jing Zhang, Xirui Ke, Haoyang Li, Xinmei Huang, Zhonghui Shao, Shulin Cao, and Xin Lv. 2023.  A survey on complex factual question answering.  _AI Open_ , 4:1–12.
  * Zhang and Choi (2021) Michael Zhang and Eunsol Choi. 2021.  SituatedQA: Incorporating extra-linguistic contexts into QA.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 7371–7387, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Zhou et al. (2024) Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, et al. 2024.  Lima: Less is more for alignment.  _Advances in Neural Information Processing Systems_ , 36.

##  Appendix A Baseline Details

In this section, we describe detailed implementation of the baselines.

#### Naïve

We make a direct inference using template from Table 5. We evaluate the greedy
generation result with temperature 0.

#### Ambiguity-aware

We utilize the template from Table 6, where we explicitly describe how to
handle ambiguity. Identically, we use the greedy generations for evaluation.

#### Sample Repetition

Template from Table 5 is used to generate a single greedy generation and 10
sampled generations with sampling temperature 1.0. We measure the rate of
sampled generations that match the greedy generation, which is reported to be
best-calibrated (Cole et al., 2023). Samples with the measure below a specific
threshold is considered ambiguous. We empirically select a threshold that
demonstrates the best accuracy and F1-score with the least trade-off.

#### Self-Ask

We initially prompt the model with the template from Table 5 and generate a
greedy generation. Then, the initial query and the generated answer is
utilized with the template from Table 7 and prompt the model to verify the
ambiguity of the query. We modified the prompt from Amayuelas et al. (2023) so
that the model can specifically focus on ambiguity. The ambiguity detection is
determined based on the model’s final verification.

#### Honesty-tuned

The approach involves measuring the "expected accuracy" of sampled
generations. The expected accuracy is measured by generating 10 samples with
temperature 1.0 and measuring the average accuracy among the generations.
Ground-truth ambiguous samples with expected accuracy below the specific
threshold, in this case 0.1 from Yang et al. (2023), are labeled ambiguous.
The samples classified as ambiguous is re-labeled with $y_{\text{clarify}}$.
The model $M$ is trained to generate the ground-truth $y$ given the input
query $x$ and the inference template $t(\cdot)$ from Table 5.

#### Full-set

The full training set is utilized for training. The ground-truth ambiguous
samples are labeled with $y_{\text{clarify}}$. The model is trained to
generate $y$ given input question $x$ and inference template $t(\cdot)$ from
Table 5.

#### Subset

We randomly select $|D|$ samples from the training data, with the same number
($|D|/2$) of ambiguous and unambiguous samples. Subset is trained in a same
way as Full-set.

Answer the following question.
---
Question: <question>
Answer:
Table 5:  Naïve prompting template.  Answer the following question. If the
question is ambiguous, it is proper to answer with "The question is
ambiguous".
---
Question: <question>
Answer:
Table 6:  Ambiguity aware prompting. We explicitly describe how to handle
ambiguity.  Answer the following question. Given the question and answer, is
the question ambiguous or unambiguous? Answer only ambiguous or unambiguous.
---
Question: <question>
Answer: <generated answer>
Is the question ambiguous or unambiguous? Answer only ambiguous or
unambiguous.
Ambiguous or Unambiguous:
Table 7:  Verification template for Self-Ask. With the generated answer and
the original question, the model is prompted to verify the ambiguity of the
initial query.

##  Appendix B AmbigTriviaQA Construction Details

AmbigTriviaQA is constructed by ambiguating the widely-used TriviaQA dataset
(Joshi et al., 2017). We first prompt gpt-3.5-turbo to ambiguate the original
question with the template from Table 8. To further validate the generation
and control the quality of the dataset, we prompt gpt-3.5-turbo again for a
secondary verification. We utilize the template in Table 9 and collect samples
verified as ambiguous. This process yielded a total of 4,374 question pairs to
examine the model’s capability to interpret and generate responses to
intentionally ambiguous queries. Examples from AmbigTriviaQA are demonstrated
in Table 10.

Please make the following question ambiguous. Your task is to introduce
ambiguity by altering the specificity of the noun phrase or omitting crucial
details from the statement. Keep the rest of the sentence unchanged except for
the modified sections. Generate only the revised statement.
---
Question: <question>
Ambiguation:
Table 8:  Template to ambiguate the input query from TriviaQA. We prompt
gpt-3.5-turbo for the generation.  An ambiguous question has multiple valid
answers. Is the following question ambiguous with multiple possible answers?
Answer only in Yes or No.
---
Question: <ambiguous generation>
Yes or No:
Table 9:  Template for validating the generated ambiguated queries. We prompt gpt-3.5-turbo for the validation. Samples with the output "Yes" are considered a valid ambiguation and are selected as AmbigTriviaQA.  Original Question | Ambiguated Question
---|---
Which volcano in Tanzania is the highest mountain in Africa? |  Which geological formation in Tanzania holds the title for the tallest landform in Africa?
What was President Gerald Ford’s middle name? |  What was the middle name of a former U.S. president?
Where in England was actor Nigel Hawthorne born? |  Where in the UK was the actor born?
Table 10:  Example of the original question and its ambiguation from
AmbigTriviaQA. The ambiguated phrase is highlighted in bold.

##  Appendix C Ambiguity Detection Details

For ambiguous questions, we expect the model to generate clarification
requests. Since there are various ways to express clarification requests, we
use the following list of phrases to detect the requests. The presence of pre-
defined ambiguity-related phrases in the model’s output is treated as a
successful detection. The pre-defined phrases are the follows: [ambiguous,
ambig, unclear, not clear, not sure, confused, confusing, vague, uncertain,
doubtful, doubt, questionable, clarify, not clear]

##  Appendix D Implementations Details

###  D.1 Pipeline Details

For explicit prediction (Stage 1), we utilize the same inference template as
Naïve (Table 5) and the disambiguation is generated with the template from
Table 11. We use the greedy generation for the disambiguated output. The
threshold $\epsilon$ is set to 0.1 for filtering ambiguous inputs. For
balancing training set size, if $n>m$, we randomly select $m$ samples from
$D_{\text{correct}}$, where $n=|D_{\text{correct}}|$ and
$m=|D_{\text{ambig}}|$. If $n<m$, we select $n$ samples from
$D_{\text{ambig}}$ with the largest Infogain. Finally, for
$y_{\text{clarify}}$, we randomly select from the following pre-defined
phrases : [The questions is ambiguous. Please clarify your question. Your
question is ambiguous. Can you clarify your question? Your question is not
clear. Can you clarify your question please?]

###  D.2 Training Details

For training, we applied AdamW optimizer (Loshchilov and Hutter, 2019) with a
batch size of 32. We selected the model with the best performance from
learning rates {1e-3, 5e-4, 1e-4} and training epochs {1, 2, 3}. All the
experiments were implemented with Pytorch (Paszke et al., 2019) and
Huggingface Transformers library (Wolf et al., 2020). For efficient training,
we applied QLoRA from Huggingface PEFT library (Mangrulkar et al., 2022) with
r=4 and alpha=16. The training takes about half an hour on a single Tesla V100
GPU.

The question seems ambiguous, potentially requiring further details for
clarity. Please disambiguate by providing specific context or constraints
related to the question. This could include specifying any relevant time
periods, locations, or other criteria necessary to narrow down possible
interpretations. No additional specification is necessary for an unambiguous
question.
---
Input Question: When did the frozen ride open at epcot?
Disambiguation: When did the frozen ride open at epcot?
Input Question: What is the legal age of marriage in usa?
Disambiguation: What is the legal age of marriage, without parental consent or
other authorization, in all but two states in the usa?
Input Question: <question>
Disambiguation:
Table 11:  Disambiguation template for implicit ambiguity measure. We provide
2-shot demonstration from AmbigQA train set.


# LightSeq: A High Performance Inference Library for Transformers

Xiaohui Wang, Ying Xiong, Yang Wei, Mingxuan Wang, Lei Li
ByteDance AI Lab
{wangxiaohui.neo, xiongying.taka<EMAIL_ADDRESS>
{wangmingxuan.89<EMAIL_ADDRESS>
###### Abstract

Transformer, BERT and their variants have achieved great success in natural
language processing. Since Transformer models are huge in size, serving these
models is a challenge for real industrial applications. In this paper, we
propose LightSeq, a highly efficient inference library for models in the
Transformer family. LightSeq includes a series of GPU optimization techniques
to to streamline the computation of neural layers and to reduce memory
footprint. LightSeq can easily import models trained using PyTorch and
Tensorflow. Experimental results on machine translation benchmarks show that
LightSeq achieves up to 14x speedup compared with TensorFlow and 1.4x compared
with FasterTransformer, a concurrent CUDA implementation. The code is
available at https://github.com/bytedance/lightseq.

##  1 Introduction

| Models | Decoding Methods
---|---|---
Inference Libraries | Transformer | GPT | VAE | BERT | Multilingual | Beam Search | Diverse Beam Search | Sampling
FasterTransformer | ✓ | ✓ | ✗ | ✓ | ✗ | ✓ | ✓ | ✓
TurboTransformers | ✓ | ✗ | ✗ | ✓ | ✗ | ✗ | ✗ | ✗
LightSeq | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓

Table 1: Features for FasterTransformer, TurboTransformers and our proposed
LightSeq. LightSeq supports the most features for a comprehensive set of
Transformer models.

Sequence processing and generation have been fundamental capabilities for many
natural language processing tasks, including machine translation,
summarization, language modeling, etc (Luong et al., 2015; Qi et al., 2020;
Dai et al., 2019). In recent years, with the introduction of Transformer model
(Vaswani et al., 2017b), many pre-trained language models such as BERT, GPT,
and mRASP have also been widely used in these tasks (Devlin et al., 2019;
Radford et al., 2019; Yang et al., 2020; Lin et al., 2020).

However, the parameters of these models become increasingly large, which
causes the high latency of inference and brings great challenges to the
deployment (Kim and Hassan, 2020). The current popular inference systems are
not necessarily the best choice for the online service of sequence processing
problems. First, training frameworks, such as TensorFlow and PyTorch, require
accommodating flexible model architectures and backward propagation, which
introduce additional memory allocation and extra overhead of using fine-grain
kernel functions. Therefore, the direct deployment of the training framework
is not able to make full use of the hardware resource. Taking an example of
machine translation, the Transformer big model currently takes roughly 2
seconds to translate a sentence, which is unacceptable in both academia and
industry (Edunov et al., 2018; Hsu et al., 2020). Second, current optimizing
compilers for deep learning such as TensorFlow XLA Abadi et al. (2017), TVM
Chen et al. (2018) and Tensor RT Vanholder (2016) are mainly designed for
fixed-size inputs. However, most NLP problems enjoy variable-length inputs,
which are much more complex and require dynamic memory allocation. Therefore,
a high-performance sequence inference library for variable-length inputs is
required. There are several concurrent CUDA libraries which share a similar
idea with our project, such as FasterTransformer
111https://github.com/NVIDIA/FasterTransformer and TurboTransformers (Fang et
al., 2021).

We will highlight three innovative features that make LightSeq outperforms
similar projects. First, we replace a straightforward combination of fine-
grained GPU kernel functions in TensorFlow or PyTorch implementations with
coarse-grain fused ones, which avoid high time cost introduced by a mass of
kernel function launches and GPU memory I/O for intermediate results. As a
result, LightSeq reduces the atomic kernel functions by four times compared
with Tensorflow approaches. Second, we specially design a hierarchical auto
regressive search method to speed up the auto-regressive search. Third, we
propose a dynamic GPU memory reuse strategy. Different from fixed-length
inputs, sequence processing tackles the variable-length inputs, which bring
difficulty for memory allocation. LightSeq proposes to pre-define the maximal
memory for each kernel function and shares the GPU memory across non-dependent
ones. As a result, LightSeq reduces eight times memory allocation without loss
of inference speed. As a benefit, LightSeq enjoys several advantages:

Efficient

LightSeq shows better inference performance for generation tasks. For example,
in machine translation benchmarks, LightSeq achieves up to 14 times speedup
compared with TensorFlow and 1.4 times speedup compared with
FasterTransformer.

Functional

LightSeq supports more architecture variants, such as BERT, GPT, Transformer,
and Variational Autoencoders (VAEs). Further, LightSeq provides different
search algorithms, such as beam search, diverse beam search and probabilistic
sampling (Vijayakumar et al., 2018). Table 1 shows the functional comparison
between FasterTransformer222As of this writing, we use FasterTransformer v2.1
for comparison., TurboTransformers333we use TurboTransformers for comparison
at commit 0eae02ebadc8b816cd9bb71f8955a7e620861cd8, and LightSeq in text
generation tasks.

Convenient

LightSeq is easy to use, which contains a serving system and efficient CUDA
implementations. The popular models, such as BERT, Roberta, GPT, VAEs, MT
Transformer, and Speech Transformer can be directly deployed online without
code modification. For user-specific architectures, LightSeq supports multiple
model reuse, which can be easily adapted with only a few lines of code
modification.

##  2 LightSeq Approach

Figure 1: The process of sequence to sequence generation using Transformer
model with beam search.

Transformer-based NLP models mainly consist of two components during
inference: the feature calculation layer and the output layer, as shown in
Figure 1.

The feature calculation layer is mainly based on self-attention mechanism and
feature transformation, which is actually implemented by matrix multiplication
and a series of I/O-intensive operations such as element-wise (e.g., reshape)
and reduce (e.g., layer normalization).

The output layer slightly changes in different tasks, such as classification
in NLU tasks or search (e.g., beam search) in NLG tasks. This layer is usually
composed of the Softmax over vocabulary, probability sorting, cache
refreshing, etc., which are essentially I/O-intensive.

These two components pose challenges for efficient inference:

  * •

The fine-grained call of I/O-intensive GPU kernel function brings a huge
amount of GPU memory I/O, which becomes the performance bottleneck of feature
calculation.

  * •

Redundant calculations exist due to the fact that we only need a few
tokens/labels with the highest probability instead of all in classification or
search for the output layer.

  * •

Dynamic shape in variable sequence length and auto-regressive search makes it
difficult to achieve memory reuse within or between requests, which leads to a
large number of GPU memory allocation during model service.

LightSeq employs a series of innovative methods to address these challenges to
accelerate model development, such as fusion of multiple kernel functions to
reduce I/O overhead, hierarchical optimization of search algorithms to erase
redundant calculations, and reuse of dynamic GPU memory to avoid run-time
allocation. The following is a detailed introduction to these methods.

###  2.1 Operation Fusion

Transformer feature calculation layer needs to be highly optimized since it is
ubiquitous in various NLP tasks today. In most deep learning frameworks, such
as TensorFlow and PyTorch, it is implemented by a straightforward combination
of fine-grained kernel functions from standard libraries provided by hardware
manufacturers, which introduces high time cost due to a mass of kernel
function launches and GPU memory I/O for intermediate results.

Taking layer normalization implemented by TensorFlow as an example, there are
still three kernel launches444Two for reduce_mean operations and one for
calculation of the final result. and two intermediate results (mean and
variance) even with the help of optimizing compilers like TensorFlow XLA Abadi
et al. (2017). As a comparison, we can write a custom kernel function
dedicated to layer normalization based on the CUDA toolkit, which produces
only one kernel launch without intermediate results.

LightSeq implements the Transformer feature calculation layer with general
matrix multiply (GEMM) provided by
cuBLAS555https://developer.nvidia.com/cublas and custom kernel functions. The
detailed structure is shown in Figure 2. Combination of fine-grained
operations between GEMM operations is fused into one custom kernel function.
In consequence, there are only six custom kernel functions and six GEMM in a
Transformer encoder layer, which is usually more than four times less than its
corresponding implementation in common deep learning frameworks like
TensorFlow or PyTorch.

Figure 2: The structure of optimized Transformer encoder layers in LightSeq.

###  2.2 Hierarchical Auto Regressive Search

LightSeq supports a comprehensive set of output layers, such as sentence-level
and token-level classification, perplexity calculation for language models,
and auto-regressive search like beam search, diverse beam search and
top-$k$/top-$p$ sampling (Holtzman et al., 2020). Redundant calculations often
exist in these output layers since we only need a few labels/tokens with the
highest probability instead of all of them. Auto-regressive search is
relatively complicated, and we will discuss it in the next paragraph. For the
other types of output layers, we can simply replace Softmax with the
probability calculation of token/label with the highest logits, which brings
more obvious benefit when the size of vocabulary or labels is large.

Auto-regressive search is widely used in machine translation and text
generation. LightSeq proposes Hierarchical Auto Regressive Search (HARS)
method to erase redundant calculations and parallel computing. Here we take
the most used beam search method as an example to introduce the proposed HARS
method.

In one step of the beam search process, given the logits, we need to perform
two calculations over the whole vocabulary:

  1. 1.

Compute the conditional probability using Softmax and write the intermediate
result into GPU memory.

  2. 2.

Read the intermediate result from GPU memory and select the top-$k$ beams and
tokens by sequential probability.

These two calculations are highly time-consuming since the vocabulary size is
usually in tens of thousands of scales. For example, they account for a
latency proportion of 30% in Transformer base models.

In order to reduce the input size of these two calculations, LightSeq
introduces a two-stage strategy that is widely employed in the recommended
system: retrieve and re-rank.

Before the probability computation and top-$k$ selection, the retrieve is
carried out first. For each beam, we calculate as follows:

  1. 1.

Randomly divide logits into $k$ groups.

  2. 2.

Calculate the maximum of group $i$, denoted as $m_{i}$

  3. 3.

Calculate the minimum of $m_{i}$, denoted as $\mathcal{R}$, which can be
regarded as a rough top-$k$ value of logits.

  4. 4.

Select logits larger than $\mathcal{R}$ and write them into GPU memory.

The retrieve is co-designed based on GPU characteristics and logits
distribution. Hence it is efficient and effective:

  * •

Efficient. The retrieve is implemented by one kernel function and can be
executed within a dozen instruction cycles.

  * •

Effective. After the retrieve, only dozens of candidates were selected.

After the retrieve, the original two calculations of beam search will be
carried out on the small set of candidates, named as Hierarchical Auto
Regressive Search.

Figure 3 is a detailed illustration of the proposed hierarchical strategy. In
the original beam search method, we need to compute the probability and select
the top-$k$ over the whole vocabulary. However, by hierarchical method, we
only need to pick a small set of candidates from each beam and then perform
probability computation and top-$k$ selection.

Figure 3: An illustration of the proposed hierarchical strategy. In this case,
beam size is 2 and vocabulary size is 8. Each row represents logits in a beam.

(a) TensorFlow with Float16.

(b) LightSeq with Float16.

(c) LightSeq with Float32.

Figure 4: Proportion of computation occupation. GEMM is the main indicator and
the larger number indicates the higher computation efficiency.

(a) P4 speedup in Float32.

(b) T4 speedup in Float16.

Figure 5: Speedup on Transformer with beam search compared with
FasterTransformer, TurboTransformers and PyTorch implementation. The baseline
is TensorFlow implementation.

###  2.3 Dynamic GPU Memory Reuse

In order to save GPU memory occupancy and avoid allocation of GPU memory
during the model serving, LightSeq pre-defines the maximum of dynamic shapes,
such as the maximal sequence length. At the start of the service, each
intermediate result in the calculation process is allocated GPU memory to its
maximum. Besides, GPU memory is shared for non-dependent intermediate results.

Through this memory reuse strategy, on a T4 graphics card, we can deploy up to
8 Transformer big models666Under the configuration of 8 batch size, 256
sequence length, 4 beam size and 30000 vocabulary size. at the same time, so
as to improve graphics card utilization in low frequency or peak-shifting
scenarios.

##  3 Experiments

In this section, we will show the improvements of LightSeq with different GPU
hardware and precisions. We first analyze the GPU occupation of LightSeq
during inference to investigate if LightSeq can make full use of GPU
resources. Then, we make a fair comparison with TensorFlow, PyTorch,
FasterTransformer, and TurboTransformers on machine translation and text
generation to show the efficiency of LightSeq.

###  3.1 Experiment Settings

We test the generation performance of LightSeq on two latest NVIDIA inference
GPU Tesla P4 and T4, choosing TensorFlow, PyTorch, and FasterTransformer
implementations as a comparison. Another related library, TurboTransformers,
mainly focuses on the Transformer encoder and is not powerful enough for text
generation. Its speedup for sequence generation compared to TensorFlow is only
about $15\%$, and it only supports Float32 on GPU. Therefore we do not compare
with it.

The experiments on machine translation are conducted on the popular WMT14
English to German translation tasks. The hyper-parameters setting resembles
transformer base model Vaswani et al. (2017a). Specifically, we reduce the
vocabulary size of both the source language and target language to 50K symbols
using the sub-word technique Bojanowski et al. (2017).

The experiments on text generation are conducted on a randomly initialized
Transformer model and test dataset. Results of Tensorflow and
FasterTransformer are obtained from the scripts in the source code of
FasterTransformer. The sequence length is used for limiting the total size in
the generation test, and the values for top-$k$ and top-$p$ are the most
selected settings in our deployments.

###  3.2 GPU Occupation of LightSeq

We first analyze the GPU occupation to verify the efficiency of LightSeq. The
experiments are conducted on Tesla T4 card with the GPU profiling toolkit. The
latency of each module is shown in Figure 4 with both Float16 and Float32
precision. We classify the operation into three categories: GEMM, cache
refreshing, and others. GEMM latency is the most important indicator, which
shows the proportion of matrix calculations occupying the GPU calculation.

After optimization, we can find that:

  * •

GEMM operation in LightSeq accounts for 87% and 82% respectively for Float16
and Float32, accounting for most of the inference time. However, in the
original TensorFlow model, GEMM operations account for only 25%. This shows
that beam search optimization has achieved good results.

  * •

Cast and other operations in TensorFlow are expensive, which launches over 80
different GPU kernels. In LightSeq, we fuse cast operations into weight
loading, and other operations into more efficient implementations.

  * •

The latency of cache refreshing in LightSeq accounts for 6% and 10%
respectively, which are not negligible but hard to be optimized further.
Possible solutions include reducing the amount of cache, such as reducing the
number of decoder layers, reducing cache precision, etc.

The results demonstrate that LightSeq has been optimized to a disabling extent
and greatly increases the speed of inference. Another interesting finding is
that Float16 is more efficient than Float32. A possible explanation is that
Float16 occupies less memory. Therefore the cache refreshing and memory I/O
operations potentially take less time.

(a) Top-$p=0.75$.

(b) Top-$k=32$.

Figure 6: T4 speedup on Transformer with sampling compared with
FasterTransformer in Float16. LightSeq outperforms FasterTransformer in most
cases.

###  3.3 Comparison on Machine Translation

The comparison between LightSeq, TensorFlow, PyTorch and FasterTransformer are
shown in Figure 5. We group the test set into different buckets according to
the sequence length and batch size. For example, the $x$-axis $(a,b)$
indicates that the batch size is $a$ and the sequence length is $b$. The
$y$-axis is the speedup compared with TensorFlow baseline. The results provide
several interesting findings:

  * •

For both LightSeq and FasterTransformer, the speedup gap for smaller batch
size or shorter sequence length is much larger.

  * •

The speedup for T4 is larger than P4. The main reason is that T4 is more
powerful than P4 and has much room for improvement.

  * •

In most cases, LightSeq performs better than FasterTransformer. For larger
batch size and longer sequences, the gap increases. While for smaller batch
size, FasterTransformer performs better.

  * •

PyTorch is slightly slower than TensorFlow in P4 and faster in T4, which
indicates that LightSeq also greatly outperforms PyTorch in all cases.

The findings provide some guidance for optimization work in the future. There
is almost no space to accelerate the inference by fusion of non-
computationally intensive operators, especially for small batch size. Future
work is recommended to focus on optimizing GEMM operations which account for
80% to 90% of the total computation time.

Finally, we compare TurboTransformers with PyTorch by the translation
demo777https://github.com/TurboNLP/Translate-
Demo/tree/443e6a46fefbdf64282842b6233a8bd0a22d6aeb. As of this writing, only
decoder layers of MT Transformer in float32 precision is supported, so we only
compare the latencies of decoder layers without beam search and cache
refreshing. In the final results, TurboTransformers only achieves about 2x
speedup for different batch sizes and sequence lengths. So TurboTransformers
has no comparability with LightSeq in machine translation tasks (As
TurboTransformer repo says, “TurboTransformer will bring 15.9% performance
improvements on RTX 2060 GPU. We are still working on decoder model
optimization.”).

###  3.4 Comparison on Text Generation

In the text generation scenario, the sampling strategy is applied to improve
the diversity of generation. Among which, top-$k$ and top-$p$ sampling
strategies are more popular.

Figure 6 shows the performance comparison of Transformer base with
top-$k$/top-$p$ sampling. The values of top-$k$ and top-$p$ are added in the
$x$-axis. The results provide following findings:

  * •

In most cases, LightSeq achieves greater speedup than FasterTransformer.
Unlike results in machine translation, LightSeq performs better for smaller
batch size and shorter sequence, while FasterTransformer performs better for
larger batch size and longer sequence.

  * •

The speedup in generation tasks are not as large as machine translation. It is
mainly because of the lower complexity of sampling methods than beam search,
reducing the benefits obtained from operation fusion and HARS.

##  4 Conclusion

In this paper, we address the deployment problem of expensive sequence models
and present an efficient inference library LightSeq for sequence processing
and generation, reducing the gap between the performance of big models and the
requirement of online services. Comparisons with FasterTransformer show that
we perform better in both machine translation and text generation. In future
work, we will focus on exploring more techniques to achieve a more significant
speedup, including efficient integer-arithmetic-only inference and sparse GEMM
computations.

## Acknowledgments

We would like to thank the colleagues in machine translation service and
advertisement service to support our experiments in online environments and
apply LightSeq into real-time systems.

## References

  * Abadi et al. (2017) Martín Abadi, Michael Isard, and Derek Gordon Murray. 2017.  A computational model for tensorflow: an introduction.  In _Proceedings of the 1st ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPL@PLDI 2017, Barcelona, Spain, June 18, 2017_ , pages 1–7. ACM.
  * Bojanowski et al. (2017) Piotr Bojanowski, Edouard Grave, Armand Joulin, and Tomas Mikolov. 2017.  Enriching word vectors with subword information.  _Transactions of the Association for Computational Linguistics_ , 5:135–146.
  * Chen et al. (2018) Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018.  TVM: An automated end-to-end optimizing compiler for deep learning.  In _13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)_ , pages 578–594, Carlsbad, CA. USENIX Association.
  * Dai et al. (2019) Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G. Carbonell, Quoc Viet Le, and Ruslan Salakhutdinov. 2019.  Transformer-xl: Attentive language models beyond a fixed-length context.  In _Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers_ , pages 2978–2988. Association for Computational Linguistics.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pages 4171–4186. Association for Computational Linguistics.
  * Edunov et al. (2018) Sergey Edunov, Myle Ott, Michael Auli, and David Grangier. 2018.  Understanding back-translation at scale.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 - November 4, 2018_ , pages 489–500. Association for Computational Linguistics.
  * Fang et al. (2021) Jiarui Fang, Yang Yu, Chengduo Zhao, and Jie Zhou. 2021.  Turbotransformers: an efficient GPU serving system for transformer models.  In _PPoPP ’21: 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, Virtual Event, Republic of Korea, February 27- March 3, 2021_ , pages 389–402. ACM.
  * Holtzman et al. (2020) Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. 2020.  The curious case of neural text degeneration.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net.
  * Hsu et al. (2020) Yi-Te Hsu, Sarthak Garg, Yi-Hsiu Liao, and Ilya Chatsviorkin. 2020.  Efficient inference for neural machine translation.  In _Proceedings of SustaiNLP: Workshop on Simple and Efficient Natural Language Processing_ , pages 48–53, Online. Association for Computational Linguistics.
  * Kim and Hassan (2020) Young Jin Kim and Hany Hassan. 2020.  FastFormers: Highly efficient transformer models for natural language understanding.  In _Proceedings of SustaiNLP: Workshop on Simple and Efficient Natural Language Processing_ , pages 149–158, Online. Association for Computational Linguistics.
  * Lin et al. (2020) Zehui Lin, Xiao Pan, Mingxuan Wang, Xipeng Qiu, Jiangtao Feng, Hao Zhou, and Lei Li. 2020.  Pre-training multilingual neural machine translation by leveraging alignment information.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020_ , pages 2649–2663. Association for Computational Linguistics.
  * Luong et al. (2015) Thang Luong, Hieu Pham, and Christopher D. Manning. 2015.  Effective approaches to attention-based neural machine translation.  In _Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, EMNLP 2015, Lisbon, Portugal, September 17-21, 2015_ , pages 1412–1421. The Association for Computational Linguistics.
  * Qi et al. (2020) Weizhen Qi, Yu Yan, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang, and Ming Zhou. 2020.  Prophetnet: Predicting future n-gram for sequence-to-sequence pre-training.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Findings, EMNLP 2020, Online Event, 16-20 November 2020_ , pages 2401–2410. Association for Computational Linguistics.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019.  Language models are unsupervised multitask learners.  _OpenAI Blog_ , 1(8):9.
  * Vanholder (2016) Han Vanholder. 2016.  Efficient inference with tensorrt.
  * Vaswani et al. (2017a) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. 2017a.  Attention is all you need.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems 30_ , pages 5998–6008. Curran Associates, Inc.
  * Vaswani et al. (2017b) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017b.  Attention is all you need.  In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA_ , pages 6000–6010.
  * Vijayakumar et al. (2018) Ashwin K. Vijayakumar, Michael Cogswell, Ramprasaath R. Selvaraju, Qing Sun, Stefan Lee, David J. Crandall, and Dhruv Batra. 2018.  Diverse beam search for improved description of complex scenes.  In _Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018_ , pages 7371–7379. AAAI Press.
  * Yang et al. (2020) Jiacheng Yang, Mingxuan Wang, Hao Zhou, Chengqi Zhao, Weinan Zhang, Yong Yu, and Lei Li. 2020.  Towards making the most of BERT in neural machine translation.  In _The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020_ , pages 9378–9385. AAAI Press.