Efficient and Asymptotically Unbiased Constrained Decoding for Large Language Models
Extracted source text · arXiv 2504.09135
\documentclass[twoside]{article}
% \usepackage{aistats2025}
% If your paper is accepted, change the options for the package
% aistats2025 as follows:
%
\usepackage[accepted]{aistats2025}
%
% This option will print headings for the title of your paper and
% headings for the authors names, plus a copyright note at the end of
% the first column of the first page.
% If you set papersize explicitly, activate the following three lines:
%\special{papersize = 8.5in, 11in}
%\setlength{\pdfpageheight}{11in}
%\setlength{\pdfpagewidth}{8.5in}
% If you use natbib package, activate the following three lines:
\usepackage[round]{natbib}
\renewcommand{\bibname}{References}
\renewcommand{\bibsection}{\subsubsection*{\bibname}}
% If you use BibTeX in apalike style, activate the following line:
\bibliographystyle{apalike}
\input{math_commands}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage[colorlinks,
linkcolor=red,
anchorcolor=blue,
citecolor=blue,
]{hyperref}
\usepackage{hyperref}
\usepackage{url}
\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage{algorithmicx}
\usepackage{cleveref}
\usepackage{xcolor}
\usepackage{amssymb}
\usepackage{xspace}
\usepackage{graphicx}
\usepackage{pifont}
\newcommand{\alg}{{\texttt{DISC}}\xspace}
\newcommand{\ppv}{{\texttt{PPV}}\xspace}
\newcommand{\ac}{{(\textcolor{red}{addcite})}\xspace}
\newcommand{\haotian}[1]{{\color{teal}{\small\sf [HT: #1]}}}
\newcommand{\theertha}[1]{{\color{blue}{\small\sf [theertha: #1]}}}
\newcommand{\felix}[1]{{\color{red}{\small\sf [felix: #1]}}}
\begin{document}
% If your paper is accepted and the title of your paper is very long,
% the style will print as headings an error message. Use the following
% command to supply a shorter title of your paper so that it can be
% used as headings.
%
\runningtitle{Efficient and Asymptotically Unbiased Constrained Decoding for Large Language Models}
% If your paper is accepted and the number of authors is large, the
% style will print as headings an error message. Use the following
% command to supply a shorter version of the authors names so that
% they can be used as headings (for example, use only the surnames)
%
\runningauthor{Haotian Ye, Himanshu Jain, Chong You, Ananda Theertha Suresh, Haowei Lin, James Zou, Felix Yu}
\twocolumn[
% \aistatstitle{Efficient and Asymptotically Unbiased\\Constrained Autoregressive Generation}
\aistatstitle{Efficient and Asymptotically Unbiased\\Constrained Decoding for Large Language Models}
\aistatsauthor{ Haotian Ye$^{1*}$ \And Himanshu Jain$^2$ \And Chong You$^2$ \And Ananda Theertha Suresh$^2$}
\aistatsauthor{Haowei Lin$^3$ \And James Zou$^1$ \And Felix Yu$^2$}
\aistatsaddress{ $^1$Stanford University \And $^2$Google \And $^3$Peking University}
% \centering{$^*$Work done during internship at Google Research. Correspondence to \url{haotianye@stanford.edu}}
% \aistatsauthor{Haowei Lin \And James Zou \And Felix Yu}
% \aistatsaddress{Peking University \And Stanford University \And Google}
]
% \footnote{Work done during internship at Google Research. Correspondence to \url{haotianye@stanford.edu}}
% \haotian{We can remove asymptotically here because by setting $K$ to infinity, it's unbiased, and the sample step is still bounded in expectation.}
\begin{abstract}
\input{abstract}
\end{abstract}
% Logical flow of the paper:
% \begin{enumerate}
% \item Retrieval is a widely used and essential technique for language model applications such as retrieval augmented generation, entity linking, and so on.
% \item From retrieval to generative retrieval.
% \item To achieve this, constrained decoding is necessary. However, current constrained decoding is myopic and inefficient when a number of valid items is huge. \textcolor{red}{check how they do, whether they are myopic, and if there are inefficient.}
% \item
% \end{enumerate}
% \section{Writing Panels}
% Experiments on extreme classification (XC)\begin{enumerate}
% \item : Will include two datasets (AmazonCAT-13k \& Wiki-500K) that are multi-labels. See \Cref{fig:illustration_XC}.
% \item This will be finished by the weekend.
% \item Existing methods are NOT with LLM, so there is a chance that we can change the game here.
% \end{enumerate}
% Experiments on generative retrieval (GR):\begin{enumerate}
% \item Will include all tasks in the autoregressive entity retrieval paper, as well as NQ320k and MSMARCO (two extremely large scale datasets). Labels are partially seen during training.
% \item This will be finished by next Tue.
% \item There are lots of existing methods, but lots of them are dirty. While we believe that a standard and purely LLM-based method should be the method, we nevertheless will choose 1-2 additional methods and show that adding ours to theirs helps.
% \item We will try to focus more on the OOD setting where even a fine-tuned model does not perform near-optimally (not sure how much dataset is ood);
% \end{enumerate}
% Writting:\begin{enumerate}
% \item Introduction (Wed)
% \item Related works (Himanshu)
% \item Preliminary (Done, please check)
% \item Method (based on alg and theorems, extend to a section) (WED)
% \item Experiments (Haowei) (THU)
% \item Proof (can be put in the appendix) (Ananda) (Wed)
% \end{enumerate}
% \begin{figure}[h]
% \centering
% \includegraphics[width=\linewidth]{Figures/XC-illustration.png}
% \caption{Illustration of extreme classification.}
% \label{fig:illustration_XC}
% \end{figure}
% Question:\begin{enumerate}
% \item Think about whether GR or CD should be the central topic of the work.
% \item Whether XC is multi-tokens
% \item Whether we have enough methods for GR to add ours on.
% \end{enumerate}
\input{Sections/introduction}
\input{Sections/preliminary}
\input{Sections/method}
\input{Sections/experiments}
\input{Sections/related_works}
\section{DISCUSSION \& CONCLUSION}
This paper focuses on comparing trie-based methods and does not consider combining \alg with countless existing techniques. It should be pointed out that \alg can be seamlessly integrated into existing language model frameworks without requiring any modifications to the underlying architecture. This minimally invasive property means it is compatible with most generation strategies such as temperature scaling or top-$k$ sampling, as well as task-specific techniques such as special tokenizations or retrieval tricks.
% This allows practitioners to effortlessly incorporate \alg{} into their current systems, enhancing constrained decoding processes to achieve unbiased and efficient generation across various applications.
In addition, notice that \alg and \ppv can be applied to other types of constraints (beyond set constraints $\mathcal S$ we consider). For instance, constraints defined by context-free grammars, specific formatting requirements like generating valid JSON structures, or other complex criteria can also be incorporated, and the theoretical analysis extends. Future work could explore empirical evaluations of \alg{} and \ppv{} across a broader spectrum of constrained decoding scenarios.
In conclusion, we address the challenges of bias and inefficiency in constrained decoding by introducing the \alg{} and the \ppv{}. We believe that our methods effectively mitigate unintended biases and enhance computational efficiency, broadening the applicability of constrained decoding in real-world scenarios.
\bibliography{reference}
\clearpage
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section*{Checklist}
% %%% BEGIN INSTRUCTIONS %%%
% The checklist follows the references. For each question, choose your answer from the three possible options: Yes, No, Not Applicable. You are encouraged to include a justification to your answer, either by referencing the appropriate section of your paper or providing a brief inline description (1-2 sentences).
% Please do not modify the questions. Note that the Checklist section does not count towards the page limit. Not including the checklist in the first submission won't result in desk rejection, although in such case we will ask you to upload it during the author response period and include it in camera ready (if accepted).
% \textbf{In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.}
% %%% END INSTRUCTIONS %%%
\begin{enumerate}
\item For all models and algorithms presented, check if you include:
\begin{enumerate}
\item A clear description of the mathematical setting, assumptions, algorithm, and/or model. [Yes]
\item An analysis of the properties and complexity (time, space, sample size) of any algorithm. [Yes]
\item (Optional) Anonymized source code, with specification of all dependencies, including external libraries. [Yes]
\end{enumerate}
\item For any theoretical claim, check if you include:
\begin{enumerate}
\item Statements of the full set of assumptions of all theoretical results. [Yes]
\item Complete proofs of all theoretical results. [Yes]
\item Clear explanations of any assumptions. [Yes]
\end{enumerate}
\item For all figures and tables that present empirical results, check if you include:
\begin{enumerate}
\item The code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL). [Yes]
\item All the training details (e.g., data splits, hyperparameters, how they were chosen). [Not Applicable]
\item A clear definition of the specific measure or statistics and error bars (e.g., with respect to the random seed after running experiments multiple times). [Not Applicable]
\item A description of the computing infrastructure used. (e.g., type of GPUs, internal cluster, or cloud provider). [Yes]
\end{enumerate}
\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets, check if you include:
\begin{enumerate}
\item Citations of the creator If your work uses existing assets. [Yes]
\item The license information of the assets, if applicable. [Not Applicable]
\item New assets either in the supplemental material or as a URL, if applicable. [Not Applicable]
\item Information about consent from data providers/curators. [Not Applicable]
\item Discussion of sensible content if applicable, e.g., personally identifiable information or offensive content. [Not Applicable]
\end{enumerate}
\item If you used crowdsourcing or conducted research with human subjects, check if you include:
\begin{enumerate}
\item The full text of instructions given to participants and screenshots. [Not Applicable]
\item Descriptions of potential participant risks, with links to Institutional Review Board (IRB) approvals if applicable. [Not Applicable]
\item The estimated hourly wage paid to participants and the total amount spent on participant compensation. [Not Applicable]
\end{enumerate}
\end{enumerate}
\clearpage
\input{supplement}
\end{document}
%%%%% NEW MATH DEFINITIONS %%%%%
\usepackage{amsmath,amsfonts,bm}
\usepackage{amsthm}
% Mark sections of captions for referring to divisions of figures
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}
% Highlight a newly defined term
\newcommand{\newterm}[1]{{\bf #1}}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% THEOREMS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% \theoremstyle{plain}
\newtheorem{theorem}{Theorem}[section]
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{corollary}[theorem]{Corollary}
% \theoremstyle{definition}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{assumption}[theorem]{Assumption}
% \theoremstyle{remark}
\newtheorem{remark}[theorem]{Remark}
\usepackage{hhline}
\usepackage{diagbox}
\usepackage{arydshln}
% Figure reference, lower-case.
\def\figref#1{figure~\ref{#1}}
% Figure reference, capital. For start of sentence
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
% Section reference, lower-case.
\def\secref#1{section~\ref{#1}}
% Section reference, capital.
\def\Secref#1{Section~\ref{#1}}
% Reference to two sections.
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
% Reference to three sections.
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
% Reference to an equation, lower-case.
\def\eqref#1{equation~\ref{#1}}
% Reference to an equation, upper case
\def\Eqref#1{Equation~\ref{#1}}
% A raw reference to an equation---avoid using if possible
\def\plaineqref#1{\ref{#1}}
% Reference to a chapter, lower-case.
\def\chapref#1{chapter~\ref{#1}}
% Reference to an equation, upper case.
\def\Chapref#1{Chapter~\ref{#1}}
% Reference to a range of chapters
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
% Reference to an algorithm, lower-case.
\def\algref#1{algorithm~\ref{#1}}
% Reference to an algorithm, upper case.
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
% Reference to a part, lower case
\def\partref#1{part~\ref{#1}}
% Reference to a part, upper case
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}
\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}
\def\eps{{\epsilon}}
% Random variables
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
% rm is already a command, just don't name any random variables m
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}
% Random vectors
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}
% Elements of random vectors
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}
% Random matrices
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}
% Elements of random matrices
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}
% Vectors
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}
% Elements of vectors
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}
% Matrix
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}
% Tensor
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}
% Graph
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}
% Sets
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
% Don't use a set called E, because this would be the same as our symbol
% for expectation.
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}
% Entries of a matrix
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}
% entries of a tensor
% Same font as tensor, without \bm wrapper
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}
% The true underlying data generating distribution
\newcommand{\pdata}{p_{\rm{data}}}
% The empirical distribution defined by the training set
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
% The model distribution
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
% Stochastic autoencoder distributions
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}
\newcommand{\laplace}{\mathrm{Laplace}} % Laplace distribution
\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
% Wolfram Mathworld says $L^2$ is for function spaces and $\ell^2$ is for vectors
% But then they seem to use $L^2$ for vectors throughout the site, and so does
% wikipedia.
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}
\newcommand{\parents}{Pa} % See usage in notation.tex. Chosen to match Daphne's book.
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak
\usepackage{color}
\usepackage{multirow}
\usepackage{booktabs}
\usepackage{listings}
\usepackage{color}
\definecolor{codegray}{rgb}{0.5,0.5,0.5}
\definecolor{codepurple}{rgb}{0.58,0,0.82}
\lstdefinestyle{mypython}{
language=Python,
backgroundcolor=\color{white},
commentstyle=\color{codegray},
keywordstyle=\color{blue},
stringstyle=\color{codepurple},
basicstyle=\ttfamily\footnotesize,
breaklines=true,
frame=single,
showstringspaces=false,
captionpos=b,
}
\iffalse
The objective of a retrieval system is to find all the documents relevant to a given query from a large set of documents. Traditionally this is done by representing queries and documents using embeddings and then doing a nearest neighbor search over the entire set of document embeddings.
With the advent of LLMs new approaches for retrieval have emerged whereby LLMs directly generate relevant documents for a given query without explicitly computing embeddings or doing nearest neighbor search. There are primarily two ways of doing this - 1) the entire set of documents is put in context of LLM and then LLM generates relevant documents from this context, and 2) documents are not put in the context but to make sure that the generated documents are from the retrieval set, decoding process is constrained. In large scale applications where the size of retrieval set is $O(B)$ the first approach would require the LLM to have $O(B)$ context length which as of today is infeasible. In this paper we focus on the second set of approaches and propose an efficient constrained decoding algorithm.
In particular, we propose a novel algorithm that leverages sorted index table and dynamic importance sampling to accomplish constrained decoding. Unlike conventional prefix-tree-based methods, our algorithm is optimized to run purely on GPUs, making it significantly efficient. In addition, with the help of the importance sampling, sampling distribution is theoretically guaranteed to be asymptotically unbiased.
\fi
% While contemporary large language models have revolutionized countless applications, some of these applications, such as recommendations, safety checks, and extreme classification, impose additional constraints on outputs. To ensure constraints are met, constrained decoding, a technique used to control which tokens can be generated, has been widely adopted. Unfortunately, existing constrained decoding methods are inefficient and inherently bias token probabilities, resulting in a misaligned output distribution. This paper introduces \textbf{Dynamic Importance Sampling for Constrained decoding} (\alg), a novel algorithm that leverages dynamic importance sampling to ensure asymptotic unbiasedness. \alg is designed to run purely on GPUs, making it highly efficient. It is also minimally invasive, making it compatible with existing task-specific generation techniques. Extensive experiments on generative retrieval, recommendation systems, and extreme classification demonstrate the superiority of \alg over existing methods. These results highlight the potential of \alg to improve constrained generation tasks in real-world applications.
%While contemporary large language models have revolutionized numerous applications, many real-world scenarios impose additional constraints on model outputs.
% \felix{Many real-world applications of LLMs requires additional constraints on the outputs. For example, when using LLMs to solve a retrieval task, the outputs can only be from a fixed vocabulary of items. Similar applications include recommendation, classification, labeling and generation with trust and safety constraints.}
In real-world applications of large language models, outputs are often required to be confined: selecting items from predefined product or document sets, generating phrases that comply with safety standards, or conforming to specialized formatting styles.
To control the generation, constrained decoding has been widely adopted.
However, existing prefix-tree-based constrained decoding is inefficient under GPU-based model inference paradigms, and it introduces unintended biases into the output distribution.
This paper introduces Dynamic Importance Sampling for Constrained Decoding (\alg{}) with GPU-based Parallel Prefix-Verification (\ppv), a novel algorithm that leverages dynamic importance sampling to achieve theoretically guaranteed asymptotic unbiasedness and overcomes the inefficiency of prefix-tree. Extensive experiments demonstrate the superiority of our method over existing methods in both efficiency and output quality. These results highlight the potential of our methods to improve constrained generation in applications where adherence to specific constraints is essential.
% \footnote{Work done during internship at Google Research.
% Correspondence to \url{haotianye@stanford.edu}}
\section{INTRODUCTION}
% \,
% \haotian{mention somewhere that the CD we consider in this paper is mainly huge set constraints.}
% \haotian{Chong's comment:
% I did a bit of rewriting of the first four paragraphs of the intro, please see if it looks good or not.
% Things that may require some more major changes:
% - There is "for example" at the end of both paragraph 1 and 2, to justify the need of constrained decoding. I understand they are not entirely the same, but the difference is nuanced. Maybe combine them somehow.
% - Consequently, the model may start with gen-
% erating high-probability tokens that ultimately lead to
% invalid sequences under the constraints. Forced to ad-
% just, the model then selects lower-probability tokens to
% remain within the feasible set, resulting in outputs that
% do not truly represent the model’s learned language
% patterns.
% This does not seem clear what it means, but I don't know how to improve it. Perhaps you have a figure in mind (currently in TODO) that can help to clarify this.
% }
% Contemporary large language models (LLMs) \ac, due to their striking capacities in a variety of tasks, have been widely used in different applications, including question answering \ac, coding \ac, vision understanding \ac, mathematics \ac, and more. When users have a priori of model outputs that are hard to describe in prompts, e.g., outputs must be from (or not from) a huge given set \ac, follow special formats \ac, or be described by particular grammars like context-free grammars \ac, \textit{constrained encoding}, an on-the-fly technique controlling what tokens are allowed to generate, is often imposed to ensure that these requirements are satisfied \ac.
Large language models (LLMs) \citep{lee2023benefits,gilbert2023large,hwang2023review, skjuve2021my} exhibit remarkable capabilities in a variety of tasks and have been widely used in different applications, including question answering \citep{brown2020language,black2022gpt}, coding \citep{achiam2023gpt}, vision understanding \citep{liu2023visual,zhu2023minigpt4,ye2023mplugowl}, mathematics \citep{trinh2024solving}, and more.
Conventional autoregressive decoding does not restrict the tokens generated in each decoding step, which is reasonable in general use cases.
%
However, many real-world scenarios demand particular control over the output of these models.
For instance, the user may require the outputs to belong to (or be excluded from) a huge predetermined set \citep{Cao2021,rajput2024recommender}, adhere to specific formats \citep{Geng2023}, or conform to particular grammars such as context-free grammars \citep{Kellner2024}.
% Generally speaking, constrained decoding is essential in large-scale applications where user requirements are critical. For instance, in recommendation systems where models are required to select products from a vast product set, we do not want to recommend products that do not exist \ac; for generative retrieval tasks where the models are asked to ``generate'' a document from a document set, the decoding process has to be constrained such that the generated document is indeed from the set \ac. Constrained decoding can also be used to filter unsafe phrases from the outputs by restricting the generation to acceptable tokens \ac.
% By autoregressively constraining the next tokens to only those that lead to valid sequences according to the specified constraints, constrained decoding ensures that the final output of LLMs complies with the desired criteria, making it essential in numerous applications where adherence to user-defined requirements is critical \ac.
% For example, in recommendation systems where the model must generate outputs that are valid identifiers from a large database, such as product codes, URLs, or legal terms, it is imperative that the generated text corresponds exactly to one of the entries in the predefined set \ac. Similarly, when filtering out undesirable or unsafe content, constrained decoding can be used to prevent the generation of prohibited words or phrases by restricting the available tokens at each step \ac.
% By autoregressively constraining the next tokens to only those that lead to valid sequences according to the specified constraints, constrained decoding ensures that the final output of LLMs complies with the desired criteria, making it essential in numerous applications where adherence to user-defined requirements is critical \ac.
Recent studies have demonstrated \textit{constrained decoding} as a helpful technique for aligning the output of LLM with requirements \citep{sun2024learning,Cao2021}.
Given a predefined set of allowed outputs, it dynamically restricts the permissible tokens in each decoding step such that the final output aligns with the predefined constraints, enabling numerous applications where adherence to user-defined requirements is critical.
For example, in recommendation systems where the model is asked to generate an identifier for a product code, URL, or legal term from a large database, constrained decoding ensures that the generated text corresponds to exactly one of the valid identifiers \citep{bevilacqua2022autoregressive}. Similarly, in content filtering, constrained decoding can prevent the generation of prohibited words or phrases by restricting to a set of valid tokens \citep{zhang2023safetybench}.
% This paragraph describes the bias that constrained decoding will lead to. Give an illustrative figure to help clarify. Since CD is widely used, this will potentially bias the distribution and results in sub-optimal generation.
% While constrained decoding ensures that the generated outputs adhere to specific requirements, it can introduce unintended biases into the output distribution. As depicted in Figure \ac, constrained decoding can skew the model's natural probability distribution by overly restricting the available tokens at each step, leading to sub-optimal generation results. This bias arises because the constraints can force the model to generate less likely tokens that satisfy the constraints, rather than the most probable tokens according to its learned distribution. Since constrained decoding is widely used, this bias can have significant implications on the quality and fairness of the generated content.
\begin{figure}[h]
\centering
\includegraphics[width=\linewidth]
{Figures/bias_illustration.pdf}
\vspace{-10pt}
\caption{
% An illustration why the generation is biased with constrained decoding imposed. Assume a user request the model to recommend a product related to ``soccer shoes''. The probability that the model chooses each token is specified in the connection lines, where dash lines with \ding{53} denote invalid tokens (i.e., ``soccer shoes'' and ``expensive soccer-used gloves'' is not in the recommendation list for sell). As a result, despite that the probability that the model output ``soccer gloves'' (0.06) is significantly lower than ``expensive soccer-used shoes'' (0.324), the probability that the model generates the former (0.6) is much higher than the latter (0.324). This is because the model myopically choose ``soccer'' without be aware that ``soccer shoes'' is invalid.
An illustration of how biased generation occurs when constrained decoding is applied. A user requests the model to recommend a product related to ``soccer shoes''. The probability of selecting each token is represented by the connection lines, where dashed lines with \ding{53} indicate invalid tokens (e.g., ``soccer shoes'' and ``used soccer gloves'' are not available in the recommendation list). Even though the probability of ``soccer gloves'' ($0.06$) is much lower than ``used soccer shoes'' ($0.324$), the model generates the former with a higher final probability ($0.6 \times 1$) than the latter ($0.4 \times 0.9 \times 1$), since the model shortsightedly selects ``soccer'', unaware of the invalidity of ``soccer shoes''.
}
\label{fig:bias_illustration}
% \vspace{-5pt}
\end{figure}
% While constrained decoding ensures compliance with the requirements, it inevitably introduces unintended biases into the output distribution—a problem that existing works have not addressed \ac. This bias arises because the model, unaware of the constraints during training, assigns probabilities based on the entire vocabulary. Consequently, it may start with high-probability tokens that lead to invalid sequences under the constraints. Forced to adjust, the model then selects lower-probability tokens to remain within the allowed set, resulting in outputs that do not truly represent the model's learned language patterns. We illustrate this in \textcolor{red}{TODO}. Our theoretical analysis, formalized in \Cref{thm:lower_bound}, quantifies the extent of this bias and demonstrates its universal impact across different models, tasks, and datasets. This bias can significantly degrade the quality and reliability of the generated outputs, posing serious challenges for applications that depend on accurate and unbiased language generation.
While ensuring compliance with user requirements, we demonstrate in this paper that constrained decoding introduces unintended biases into the output distribution, a critical issue that existing works have not addressed \citep{sun2024learning}.
As illustrated in \Cref{fig:bias_illustration}, this bias stems from the unawareness of future constraints during autoregressive decoding.
% This bias stems from the descrepency between model training, where the probabilities are assigned to the entire vocabulary, and the constrained probabilities during decoding.
The model may start with generating high-probability tokens that could ultimately lead to sequences that violate the constraints.
Forced by constrained decoding to adjust, the model has to select subsequent tokens with lower probability, resulting in outputs that do not truly represent the model's learned language patterns.
In \Cref{thm:lower_bound}, we quantify the universal impact of bias across diverse models, tasks, and datasets.
This significantly degrades the quality and reliability of the generated outputs, posing severe challenges for applications that rely on accurate and unbiased generation.
% While this method ensures compliance with the requirements, it inevitably introduces unintended biases into the output distribution, a problem that existing works have not considered \ac.
% Specifically, as the model is forced to generate tokens from a restricted set, the model may select less probable tokens simply because the more probable ones are disallowed by the constraints, leading to outputs that are not truly representative of the model's learned language patterns.
% This bias arises due to the mismatch between the model's learned distribution and the constrained token set. The model is unaware of the constraints during training and thus assigns probabilities based on the entire vocabulary. When constraints are imposed during decoding, the conditional probabilities of the allowed tokens are not properly adjusted to reflect the reduced token set, resulting in skewed output distributions. Our theoretical analysis, as formalized in Theorem \ac, quantifies the extent of this bias and demonstrates its universal impact across different models, tasks, and datasets. As illustrated in Figure \ac, even when the model assigns high probabilities to certain tokens, these tokens may be invalid under the constraints, forcing the model to select lower-probability tokens and thereby altering the overall distribution.
% Moreover, existing methods for constrained decoding often rely on prefix tree (trie), a data structure commonly used in file systems \ac, to efficiently represent the set of valid sequences. However, since tries are typically operated on CPUs, integrating them with language models running on GPUs or TPUs \ac leads to heavy data communication overhead. Each time a token is generated, the GPU must communicate with the CPU to update and query the trie. This frequent GPU-CPU data transfer is incompatible with the high-throughput GPU inference, resulting in increased latency and resource consumption. On the other hand, implementing tries directly on GPUs is challenging due to their fragmented structure and varying node sizes. It remains unclear whether parallel processing and optimization on GPUs is possible.
% These inefficiencies hinder the practical deployment of constrained decoding in real-world applications where speed and scalability are critical.
Moreover, existing constrained decoding methods use prefix trees (a.k.a., \emph{tries}), a data structure used in file systems \citep{bayer1977prefix}, to store all valid outputs.
However, these tries are typically implemented on CPUs, leading to a heavy data communication overhead when integrating them with language models running on GPUs \citep{owens2008gpu} or TPUs \citep{jouppi2017datacenter}.
Each time a token is generated, it must be transferred to CPUs to query the trie, incurring substantial communication overhead.
While implementing tries on GPUs can reduce the communication cost, their fragmented structures and varying node sizes pose significant challenges for efficient parallel processing.
These inefficiencies hinder the practical deployment of constrained decoding in real-world applications where speed and scalability are critical.
% This paragraph describes our method and mentions its advantages.
% To address this issue, we propose a novel approach that mitigates the bias introduced by constrained decoding while maintaining computational efficiency. Our method leverages a dynamic adjustment of token probabilities within the decoding process, ensuring that the model's inherent distribution is more accurately represented even under constraints. By integrating this technique, we can generate outputs that are both compliant with user requirements and statistically consistent with the language model's learned patterns. Our approach does not require additional training or fine-tuning of the model, making it minimally invasive and easily applicable to existing systems.
To address these challenges, we propose a novel approach that is asymptotically unbiased while enhancing computational efficiency for set-based constrained decoding. Our method comprises two key contributions: a \textbf{Dynamic Importance Sampling algorithm for Constrained decoding} (\alg), equipped with a \textbf{Parallel Prefix-Verification} (\ppv) subroutine optimized for high-performance computing environments. \alg dynamically extends existing importance sampling techniques to LLMs decoding and provides a theoretically guaranteed sampling method that adjusts the token probabilities during decoding to reflect the true distributions under the constraints more accurately. As proved in \Cref{thm:alg}, \alg is unbiased with a fixed expected sampling size, ensuring that the generated outputs are both compliant with the constraints and statistically consistent with the language model's learned distribution.
% We prove in \Cref{thm:alg} that \alg is unbiased with an fixed expected sampling steps.
Regarding the efficiency issue, \ppv significantly accelerates the decoding by alternatively ``verifying'' rather than ``returning'' all valid candidate tokens. By leveraging the fact that the language model already computes the scores for all tokens at each decoding step, it verifies only tokens with top scores to determine whether they are valid under the given constraints.
This removes the requirements on prefix-tree and can be realized with a ordered candidate array with logarithmic time complexity $\mathcal O(\log N)$, where $N$ is the total number of candidates in the constraint set.
Importantly, the entire verification process is therefore highly parallelizable across all batches and top tokens. This bypass the inefficient trie loading, frequent data transfers between devices, and fragmented data access, significantly reducing computation time.
% \textcolor{red}{This paragraph discusses all experimental results and highlights its applicability to the existing framework (minimally invasive).}
We run on 20 datasets across four tasks to evaluate our proposed method compared to existing baselines. This includes existing benchmarks such as entity disambiguation and document retrieval, as well as novel use cases such as generative retrieval under in-context learning.
Not surprisingly, we consistently outperform trie-based methods across all datasets while being at most $8.5$ times faster. Remarkably, our method can be understood as a meta-algorithm that can be applied to different techniques for certain applications, as it only improves the constrained decoding process.
We summarize our contributions as follows:
(1) We uncover and theoretically analyze the bias introduced by constrained decoding in autoregressive models, highlighting its universal impact on output distributions.
(2) We propose \alg equipped with \ppv, a novel algorithm that provides asymptotically unbiased sampling for constrained decoding with impressive efficiency. Remarkably, it can be seamlessly integrated with other generation strategies and adapted to various constraints.
(3) We provide comprehensive experimental results across various datasets and tasks, validating the effectiveness and efficiency of our approach.
In summary, we aim to enhance the reliability and fairness of LLM-generated content, thereby broadening the applicability of these powerful models. \footnote{Code is available at \url{https://github.com/YWolfeee/Large-Scale-Selection-for-LLMs}.}
% (4)
% \end{enumerate}
% The last paragraph talks about the organization of the paper and one sentence for elevation.
% The remainder of the paper is organized as follows: In Section 2, we formalize the problem and discuss the limitations of existing constrained decoding methods. Section 3 introduces our proposed approach and provides theoretical analysis. Section 4 presents experimental results and discussions. We discussed related works in Section 5 and conclude the paper in Section 6.
\iffalse
\section{INTRODUCTION (old draft)}
Constrained decoding, where generated sequences must satisfy specific constraints, is essential in numerous applications of large language models (LLMs). These constraints can range from selecting outputs from a predefined list to adhering to context-free grammars (CFGs) or regular expressions (REGEX) \citep{hokamp2017lexically,hu2017controlled}. Enforcing such constraints ensures that the generated outputs are valid, coherent, and appropriate for the task at hand. In this paper, we focus on the challenge of sampling from an extremely large list of valid candidates—a scenario where enumerating all possibilities is computationally infeasible. Generating valid outputs on-the-fly is critical, especially in applications where subsequent generation steps depend on previous retrieval results.
In contemporary applications like generative retrieval and recommendation systems, constrained decoding plays a pivotal role. For instance, in generative retrieval, models generate document identifiers that must correspond to actual documents within a massive and dynamically changing corpus \citep{lewis2020retrieval,guu2020realm}. Similarly, recommendation systems require models to recommend products that are valid and available to users \citep{sun2019bert4rec}. As candidate sets become increasingly large and dynamic, constrained decoding becomes an unavoidable component to ensure the validity of generated outputs.
However, imposing constraints during the auto-regressive decoding process of LLMs introduces significant challenges. A key issue is the introduction of bias in the generated distributions. To illustrate, consider a scenario where the valid candidates are "badminton shoes" and "Adidas badminton paddle," and the user's query is "badminton paddle." An unconstrained model might naturally generate "badminton paddle," aligning well with the user's intent. But when constrained to the valid candidates, the model might generate "badminton shoes" instead. This occurs because the model assigns higher probability to generating "badminton" as the initial token, leading to a sequence that is less preferred or even irrelevant \citep{welleck2020neural}. This bias arises because the model's probability distribution over the constrained set does not align with the true distribution over valid outputs, leading to suboptimal results. The problem is exacerbated when the probability mass assigned by the LLM to the valid candidate set is limited, which is often the case in large-scale applications.
To address these challenges, we propose a novel algorithm called Dynamic Importance Sampling for Constrained Decoding (DIS-CD). Our method customizes importance sampling for auto-regressive generation, ensuring that the sampling distribution is theoretically guaranteed to be asymptotically unbiased. Unlike conventional methods that rely on prefix trees or finite automata \citep{stahlberg2019neural,zhang2021neural}, which are difficult to implement efficiently on GPUs, our algorithm is optimized for GPU computation, leading to significant improvements in computational efficiency. DIS-CD serves as a meta-algorithm for constrained decoding, capable of integrating with various existing systems designed for generative retrieval, recommendation systems, retrieval-augmented generation (RAG), and more. By incorporating DIS-CD, these systems can achieve better performance compared to not using the algorithm, both in terms of accuracy and computational efficiency.
\fi
% \haotian{To Himanshu: I currently put related works to the second last section for better understanding the length of each section. Feel free to adjust the order.}
\section{PROBLEM DEFINITION}
\input{Sections/alg}
This section introduces notations for language modeling and background on trie-based constrained decoding.
We also explain the two major challenges of existing approaches, namely the biasedness of constrained sampling and the inefficiency of trie implementation.
\subsection{Notations and Background}
Throughout the paper, we use $\vx = [x_1, x_2, ..., x_T]$ to represent a sequence where $|\vx| = T$ represents the sequence length.
We use $\vx_{<t} \triangleq [x_1, \cdots, x_{t-1}]$ to represent the sub-sequence of $\vx$ containing the first $t-1$ entries.
%$\vx_{<t} \triangleq [x_1, \cdots, x_{t-1}]$ to represent the sub-sequence of $\vx$.
We say that a sequence $\vy$ is a prefix of the sequence $\vx$, denoted as $\vy \preceq \vx$, if there exists an integer $t \le T$ such that $\vy = \vx_{<t}$.
Given two sequences $\vx$ and $\vy$, we use $\vx + \vy$ to represent their concatenation.
% Finally, we use $|\cdot|$ to represent the length of a sequence or the cardinality of a set.
% $[N]$ to represent $[1, ..., N]$
% \paragraph{Auto-regressive language model} serves as the major backbone of contemporary LLMs by generating texts \textit{token-by-token}, closely mirroring the construction of natural language.
% Given a pre-defined vocabulary set (token set) $\mathcal V$, a sentence (token sequence) $\vx = [x_1,\cdots,x_T]$ where $x_i \in \mathcal V$ are individual tokens, auto-regressive model factorizes the probability of $\vx$ as
% \begin{align*}
% P(\vx) = \prod_{t=1}^T P(x_t|\vx_{<t}).
% \end{align*}
% In practice, we use a language model $L$ to model the conditional distribution with $P_L(x_t| \vx_{<t})$. We omit the subscript $L$ for simplicity if there is no ambiguity.
\textbf{Autoregressive language modeling} is the backbone of contemporary LLMs that outputs text in a \textit{token-by-token} manner, closely mimicking the generation process of natural languages.
Mathematically, let $\mathcal V$ be a pre-defined vocabulary (i.e., token) set\footnote{Mathematically, we assume that $\mathcal V$ is equipped with an order. This holds trivially in modern computers when tokens are represented as integers.}, a sentence is represented as a sequence $\vx = [x_1,\cdots,x_T]$ where $x_i \in \mathcal V, i = 1, \ldots, T$ are individual tokens.
The autoregressive language modeling factorizes the probability of generating $\vx$ as
\begin{align*}
P(\vx) = \prod_{t=1}^T P(x_t|\vx_{<t}).
\end{align*}
Then, each conditional probability $P(x_t|\vx_{<t})$ may be estimated from the output of a language model $L$, which we represent as $P_L(x_t| \vx_{<t})$.
Notice that here we use $L$ to denote the language model with input prompts equipped, i.e., $P_L(x)$ is the output distribution conditioned on inputs that are part of $L$, since our study is regardless of the inputs that users provide.
We will omit the subscript $L$ for simplicity if there is no ambiguity.
% In practice, we use a language model $L$ to model the conditional distribution with $P_L(x_t| \vx_{<t})$. We omit the subscript $L$ for simplicity if there is no ambiguity.
% \textbf{Constrained decoding} aims to apply constraints to LLMs generation \textit{on-the-fly} to ensure that outputs follow user requirements. This paper mainly focuses on \textit{set constraints}, i.e., the generated outputs are from a given set $\mathcal S$ such as a set of product names, phrases, sentences or paragraphs.
% % i.e., by having LLMs generate entities from $\mathcal S$ directly.
% Specifically, given a sequence $\vx$, we only allow tokens in
% \begin{align*}
% \mathcal A_{\vx, \mathcal S} = \{v \in \mathcal V : \exists \vs \in \mathcal S, \vx + [v] \preceq \vs\}
% \end{align*}
% to be generated, unless the model generates the \textit{end-of-sentence} (EOS) token when $\vx \in \mathcal S$. For simplicity, we call $\mathcal A$ a \textit{valid set} in general.
% Conventionally, constrained decoding is implemented via a prefix-tree (aka, a trie) \ac, a standard data structure used in modern file systems.
\textbf{Constrained decoding} is a technique to align the output of a language model with particular constraints.
For this paper, we focus on \textit{set constraints}, i.e., the generated sequence is required to be from a pre-defined set $\mathcal S$ of sequences, which in practice could be a particular set of products, paragraphs, documents, etc.
In autoregressive language generating, this requires that $P(x_t|\vx_{<t})$ is zero for all $x_t \in \mathcal V \setminus \mathcal A_{\vx_{<t}, \mathcal S}$ where $\mathcal A_{\vx, \mathcal S}$ for general $\vx, \mathcal S$ is a \emph{valid set} defined as
\begin{align*}
\mathcal A_{\vx, \mathcal S} \triangleq \{v \in \mathcal V : \exists \vs \in \mathcal S ~~\text{s.t.}~ \vx + [v] \preceq \vs\},
\end{align*}
or unless the model generates the \textit{end-of-keyword} (\texttt{EOK}) token when $\vx_{<t} \in \mathcal S$.
A typical way of checking if a token is in the valid set is through encoding $\mathcal S$ with a prefix-tree, which we explain next.
% \paragraph{Prefix-tree}, also known as trie, is a type of search tree where each node is an individual token in $\mathcal V$ \ac. Given a sequence set $\mathcal S$, a node $x_i$ is the child of node $x_{i-1}$ if and only if $[x_1,\cdots, x_{i-1}, x_i]$ is the prefix of a member in $\mathcal S$, where $x_{j+1}$ is the child of $x_{j}$ for all $j$. When searching a candidate sequence, the algorithm starts from the root node and searches whether the first token is one of its children, and if so, repeats the above process from the child node with the remaining sequence until the entire sequence is found or a token miss happens. \textcolor{red}{TODO} gives a simple illustration. Notice that the time complexity \ac is $\mathcal O(|\vx||\mathcal V|)$ for sequence $\vx$, and for large $\mathcal S$ with large size, $|\vx|$ is typically in the order of $\log(|\mathcal S|)$.
\textbf{Prefix-tree}, also known as trie, is a type of search tree for checking if a particular string is from a given set of strings.
Each node in the tree represents a token in $\mathcal V$ with the root node representing the \textit{begining-of-keyword} (\texttt{BOK}) token.
Given a path from the root to node $\vx$, a node $V \in \mathcal V$ is the child of the node if and only if $V \in \mathcal A_{\vx, \mathcal S}$.
% Each path starts from the root to a particular node corresponds to a prefix of a sequence in $\mathcal S$, and the set of child of this node gives the valid set of the particular prefix.
When generating a sequence using a language model, the algorithm starts from the root node and keeps generating a token that is the child of the last generated token until the entire sequence is returned (reached a leaf node). As an illustration, in \Cref{fig:bias_illustration}, only ``soccer gloves'', ``used shirts'', and ``used soccer shoes'' can be generated.
We denote the output distribution of constrained decoding as $P^\text{CD}_\mathcal S$.
\subsection{Bias Induced by Constrained Decoding}
Ideally, given a language model $L$, we want to sample from the distribution $ P _\mathcal S (\vx) = P_L(\vx) / P_L(\mathcal S)$. Naively, one can score all sequences in $\mathcal S$ and sample in hindsight, but this is impractical when $\mathcal S$ is large, a typical setting in constrained decoding.
Although prefix-tree-based constrained decoding offers an alternative to sample, it inevitably induces bias in output distribution, a problem that existing works fail to consider.
Intuitively, the bias can be attributed to the myopic autoregressive decoding when additional constraints are imposed: the model can start generating a prefix $\vx \preceq \vy$ when $P_L(\vy)$ is large, but ends up generating a low-probability sequence $\vy'$ where $\vx \preceq \vy'$, because $\vy \notin \mathcal S$ and $\vy' \in \mathcal S$. Remarkably, this is an issue only when constrained decoding is applied, because in conventional model generation, suffices with high probability will not be masked out. Below, we present in theory how biased the output is depends on the ``bad'' probability $p_b$ that the model assigns on sequences outside $\mathcal S$, i.e., $1 - P_L(\mathcal S)$.
% \theertha{What does $b$ in $P_b$ denote? Might be good to introduce it before the theorem. Also define $P^{\text{CD}}_\mathcal S$ formally in the constrained decoding section above.}
\begin{theorem}
\label{thm:lower_bound}
% For any distribution $\mathbb P_{all}$ that any autoregressive generative model $L$ and any text prefix induces, denote $\mathbb P_\mathcal S$ the marginal distribution in support set $\mathcal S$ over $\mathbb P_{all}$, and $\mathbb P^{\text{CD}}_\mathcal S$ the distribution that standard constrained decoding on $\mathcal S$ results in. Then, for arbitrarily $P_b $, there exists a candidate set $\mathcal S$ with $P_b = \sum_{w \notin \mathcal S} p_{all}(w)$, such that the KL divergence between $\mathbb P_{\mathcal S}$ and $\mathbb P^{\text{CD}}_\mathcal S$ is lower bounded by
For arbitrary probability value $p_b$, there exists an autoregressive model $L$ and a candidate set $\mathcal S$ with $p_b = 1 - P_L(\mathcal S) $, such that the KL divergence\footnote{For two discrete distribution $P, Q$ over set $\mathcal X$, the KL divergence is defined as $KL(P\|Q) = \sum_{x\in \mathcal X}P(x)\ln \frac{P(x)}{Q(x)}$.} between $ P_{\mathcal S}$ and $ P^{\text{CD}}_\mathcal S$ is lower bounded by
\begin{align}
KL( P_{\mathcal S} \| P^{\text{CD}}_\mathcal S) \geq \Omega \Big( \ln\frac{1}{1 - p_b} \Big).
\end{align}
\end{theorem}
% \theertha{Add a definition of KL diverngence to be safe.}
The proof of \Cref{thm:lower_bound} is deferred to the appendix. It implies that the bias is universal, regardless of language models and tasks. So long as the probability outside $\mathcal S$ is large, the divergence can be significant.
% Notice that a more severe bias occurs when models are not explicitly fine-tuned on $\mathcal S$ and the probability of all sequences outside $\mathcal S$ is substantially large.
% \begin{theorem}
% something here.
% \end{theorem}
% In addition to constrained sampling, constrained beam search for greedy decoding is also proposed \ac. Unfortunately, the greedy nature of beam search could further exacerbate the issue we encounter.
% In short, existing constrained decoding methods are suboptimal for modeling this distribution, and more accurate estimation still needs to be discovered.
\subsection{Inefficiency of Constrained Decoding}
In addition to the estimation error, trie-based constrained decoding methods also suffer from efficiency issues.
Theoretically, the time complexity for constrained decoding is $\mathcal O(|\vx||\mathcal V|)$, and for large $\mathcal S$ with large size, $|\vx|$ is typically in the order of $\log(|\mathcal S|)$.
In comparison to the LLM inference time, this might seems negligible.
However, since LLMs are run on advanced computation devices such as GPUs and TPUs, having a trie run on CPUs means \textit{at each decoding step}, current prefixes have to be sent to CPUs for constrained decoding, and all valid subsequent tokens have to be sent back. Even constrained decoding can be run in parallel in CPUs, data communication itself is highly inefficient, incurring substantial communication overhead in integration with high-throughput GPU-based LLM inference.
A seemingly simple solution is to have a trie ran purely on GPUs for constrained decoding. However, to the best of our knowledge, whether an efficient GPU trie algorithm exists remains to be studied in the first place.
Two fundamental challengs have to be addressed.
First, as the number of children can vary substantially from node to node, maintaining a tire that consists of a huge number of fragments with different sizes on GPUs, even if not impossible, is difficult and unfriendly to parallel computation.
Second, it is unclear how to search multiple candidates on tries in a vectorized manner, such that the resources of GPUs can be fully leveraged. Given that the stopping time and size of children node are intrinsically different from candidates to candidates, organizing them into vectors for parallel computation remains difficult.
Overall, conventional trie-based constrained decoding is unsuitable with LLMs in terms of estimation accuracy and efficiency. Remarkably, this misalignment is irrespective of how the constraints are represented (e.g., whether keywords are sentences, phrases, or ID lists) and how LLMs are capable. It is rather a fundamental issue impeding their downstream applications.
\begin{algorithm*}[t]
\caption{\textbf{D}ynamic \textbf{I}mportance \textbf{S}ampling for \textbf{C}onstrained Decoding (\alg{})
% \haotian{To Ananda: Change the inner sampling process ( for one candidate into a built-in function, and call them $K$ times in Line 19.}
}
\label{alg:disc}
\begin{algorithmic}[1]
\State \textbf{Input:} Language model $L$, constraint set $\mathcal S$, maximum allowed sampling steps $K$, temperature $T$
% \State \textbf{def} \texttt{sample_each}():
% \IndState something
\Function{SampleCandidate()}{}:
\State Initialize candidate sequence $\va = []$, importance score $x = 1$.
\While{$\va \notin \mathcal S$}
\State Get next token distribution $P = L(\va, T)$.
\State Get mask of valid tokens $\texttt{mask} = \texttt{PPV}_{\mathcal{S}}(P, \va)$. \label{code:line_ppv}\Comment{Call \Cref{alg:ppv}}
\State Sample a token $t$ from the distribution $\frac{P \,\odot\, \texttt{mask}}{\left| P \,\odot\, \texttt{mask} \right|_1}$.
\State Update $\va = \va + [t]$, $x = x \times \left| P \odot \texttt{mask} \right|_1$.
% \State Update importance score .
\EndWhile
\State \Return $\va$, $x$
\EndFunction
% \State Initialize candidate buffers $A = []$, score buffer $X = []$
\State Initialize sampling step $k = 0$
\While{$k < K$} \Comment{Sample at most $K$ candidates}
\State $(\va, x) = \texttt{SampleCandidate}()$
\If{$\va \in \mathcal{S}$}
\State Update $k = k + 1$.
\If{$x > \epsilon \sim \mathcal{U}[0,1]$} \Comment{Accept with probability $x$}
\State \textbf{Return:} Sampled candidate sequence $\va$.
% \Else
% \State $A = A \cup \{a\}$, $X = X \cup \{x\}$
% \State Reset $a = []$, $x = 1$
\EndIf
\EndIf
\EndWhile
\State Sample $(\va_k, x_k) = \texttt{SampleCandidate}()$ for $K$ times, then sample $\va$ from $\{\va_k\}_{k=1}^K$ with probability $\frac{x_k}{\sum x_k}$. \label{code:IS}
\State \textbf{Return:} Sampled candidate sequence $\va$
\end{algorithmic}
\end{algorithm*}
\section{METHOD}
% \theertha{In line 10 of the algorithm, we use $A = A + [a]$ and $S = S + [s]$, which according to earlier definition, means concatenation, but I believe we mean adding to the set $A$, please change notation to reflect the same.}
% \theertha{Clarify line 17 to say corresponding $s$ in $S$.}
% \theertha{explicitly state $P^K_{IS}$ is Algorithm 1}
This section presents our method to mitigate the bias and improve computational efficiency.
% We start with a dynamic importance sampling algorithm for constrained decoding (\alg), which reduces bias by adopting an asymptotically unbiased sampling method.
% We then present a GPU-optimized prefix verification subroutine (\ppv) to improve efficiency by effectively handling parallel constrained decoding on GPUs.
% Our approach consists of two key components: a dynamic importance sampling algorithm for constrained decoding (\alg) and a GPU-optimized prefix verification subroutine (\ppv). \alg addresses the bias by providing an asymptotically unbiased sampling method, and its subroutine \ppv improves efficiency by effectively handling parallel constrained decoding on GPUs.
\subsection{Dynamic Importance Sampling for Constrained Decoding (\alg{})}
\alg incorporates importance sampling into LLM decoding.
Unlike traditional fixed-size importance sampling, \alg dynamically adjusts the sampling process, resulting in better efficiency and convergence rate.
% The core idea of \alg{} is to dynamically sample candidate sequences from the constraint set $\mathcal S$ and decide whether to accept or reject the candidates on-the-fly.
% Once a candidate is accepted, it is guaranteed that the distribution is unbiased and \alg returns the candidate. Otherwise, it keeps samples until a pre-defined buffer size $K$ is met and runs importance sampling over all samples.
% The pseudo-code of \alg{} is presented in Algorithm \ref{alg:disc}. Notice that in the algorithm, it calls a core function $\texttt{PPV}(P, a)$ that returns a mask of valid tokens based on the current partial sequence $a$ and the token probability distribution $P$. This sub-routine is GPU-optimized, and we leave the discussion to the next subsection.
The pseudo-code of \alg{} is in Algorithm \ref{alg:disc}.
It keeps sampling candidate sequences using standard constrained decoding and accepts each sample with a probability computed from how well the unconstrained output aligns with the constraints.
This process is continued until a candidate sequence is accepted, upon which the distribution is guaranteed to be unbiased.
If all $K$ samples are rejected, we further run an importance sampling on $K$ new sample sequences to minimize the bias.
% The biasedness of the sampling decreases with an increasing $K$, as we will see below.
% Hence, $K$ controls a trade-off between sampling quality and computation cost.
Notice that the algorithm calls a core function $\texttt{PPV}(P, a)$ that returns a mask of valid tokens based on the current partial sequence $a$ and the token probability distribution $P$. This sub-routine is GPU-optimized and is presented in the following subsection.
% In Algorithm \ref{alg:disc}, we initialize an empty candidate sequence $a$ and an importance score $s$. At each step, we obtain the next token distribution $P$ from the language model $L$ given the current sequence. The constrainer $C$, implemented via the PPV subroutine, provides a mask over valid tokens. We then sample a token $t$ from the constrained distribution. If $t$ is the end-of-sequence token (\texttt{EOK}), we decide whether to accept the candidate sequence $a$ based on the importance score $s$ and a random threshold $\epsilon$. This procedure ensures unbiased sampling by correcting for the probability mass of the constrained tokens.
\paragraph{Theoretical Analysis.}
The following theorem states that \alg{} achieves an upper bound on the Kullback-Leibler (KL) divergence between the distribution induced by \alg and the true constrained distribution $P_{\mathcal{S}}$.
\begin{theorem}
\label{thm:alg}
For any autoregressive model $L$ (equipped with any text prompt), denote $P^{K}_{\text{IS}}$ the distribution induced by \Cref{alg:disc} with a maximum of $K$ sampling steps under constrained decoding. Then, the KL divergence between $P^{K}_{\text{IS}}$ and the true distribution $P_{\mathcal{S}}$, $\text{KL}\left(P_{\mathcal{S}} \| P^{K}_{\text{IS}}\right)$ is bounded by
\begin{align}
\mathcal{O}\left(
\frac{p^K_b}{1 - p_b} \cdot \Big(
% \frac{P_b^K}{1 - P_b} \left( \frac{P_b}{K} + \sqrt{\frac{2\ln K}{K}} \right)
\sqrt{\frac{p_b(1-p_b)}{K}}+ \frac{p_b}{K}
\Big)
\right),
\end{align}
where $p_b = 1 - P_L(\mathcal{S})$ is the total probability mass of sequences outside $\mathcal S$, and $\mathcal{O}$ hides universal constants. In addition, the expected number of sampling steps is
\begin{align*}
\mathbb E [k] = \frac{1 - p_b^K}{1 - p_b} + K p_b^K \leq \frac {1+e}{1 - p_b}.
\end{align*}
Furthermore, the algorithm is unbiased, i.e.,
\begin{align*}
KL(P_\mathcal S \| P_{IS}^K) \to 0,
\end{align*}
with finite expectation number $\mathbb E[k] \to \frac{1+e}{1 - p_b}$.
Here the expectation is over all LLM sampling, acceptance/rejection procedure, and importance sampling.
\end{theorem}
The proof of \Cref{thm:alg} is provided in the appendix.
% This theorem offers an upper bound on the KL divergence between two distributions.
It shows that the estimator is asymptotically unbiased: when $K \to \infty$, the bound shrinks to zero, i.e., the estimated distribution converges to the ground-truth, and the convergence is \textit{exponentially} fast.
Even if we set $K = \infty$ so that the divergence is zero, the expected sample step is still bounded, implying the unbiased constrained decoding is realizable.
Otherwise, the expected number of sampling steps is still generally much smaller than $K$ when $p_b$ is not close to $1$.
Finally, if the invalid probability $p_b$ goes to zero, the unconstrained decoding almost generates valid sequences, and the bound shrinks to zero.
This is expected because the constrained decoding does not (need to) take any effect.
One may fine-tune the language model on $\mathcal S$ to reduce the bound in practice.
% We demonstrate that this leads to an improved performance in the experiments.
% This implies that fine-tuning the model on $\mathcal S$ can substantially improve the performance as well, as observed in experiments.
% leading to an asymptotically unbiased estimator since it shrinks to
% \paragraph{Comparison with Beam Search.}
% Compared to constrained beam search, which is inherently greedy and may exacerbate bias by focusing on high-probability tokens, \alg{} provides a more balanced exploration of the constrained space. By dynamically adjusting the importance scores and allowing for multiple candidate sequences, \alg{} reduces the likelihood of being trapped in suboptimal paths and offers better coverage of the valid sequence space.
% \paragraph{Bound on KL Divergence}
% The bound on the KL divergence provided by Theorem \ref{thm
% } demonstrates that \alg{} effectively reduces the bias introduced by constrained decoding. As $K$ increases, the bound decreases, indicating that the distribution induced by \alg{} converges to the true constrained distribution $P_{\mathcal{S}}$.
\subsection{Parallel Prefix-Verification (\ppv)}
% \haotian{Improve clarification.}
% We have demonstrated the superiority of our method in distribution estimation, and
We next turn to the efficiency issue.
The core operation in constrained decoding is to mask out invalid tokens given already generated prefixes, as specified in \Cref{code:line_ppv} of \Cref{alg:disc}.
% To enhance the computational efficiency of the constrained decoding process, we develop a GPU-optimized parallel prefix verification (PPV) subroutine. Instead of traversing on a trie to obtain all valid candidate tokens, it verifies whether tokens with top scores are valid according to the given constraints at each decoding step. We provide the pseudo-code of \ppv in \Cref{alg:ppv}.
% A standard approach is to traverse the trie and obtain all valid candidate tokens, which is computationally inefficient on GPUs.
We develop the Parallel Prefix-Verification (\ppv) subroutine that can benefit from GPU parallel processing by verifying whether the tokens with the highest probabilities from the language model are valid in each decoding step, avoiding complicated data transfers and trie traverse.
We provide the pseudo-code in \Cref{alg:ppv} with details below, and present an illustration in \Cref{fig:ppv_illustration}.
\begin{figure}
\centering
\includegraphics[width=0.98\linewidth]{Figures/ppv.pdf}
\caption{An illustration of \ppv. Assume the input is $[0,31,555]$, and the top-three token candidates given by the LLM are $\{35, 145,111\}$. \ppv concatenates them to the input and forms a three-row matrix. Then, it compares each row with the prepared array $\mathcal X$ where each row is a keyword in $\mathcal S$, in parallel. Since $\mathcal X$ is ordered alphabetically, the comparison can be performed via a binary search. After the comparison, we can determine whether each of the three partial outputs is the prefix of a row in $\mathcal X$. \ppv finishes by returning a vector mask indicating the validity of each candidate.
}
\label{fig:ppv_illustration}
\end{figure}
\begin{algorithm}[ht]
\caption{\textbf{P}arallel \textbf{P}refix-\textbf{V}erification (PPV)} \label{alg:ppv}
\begin{algorithmic}[1]
% \State \textbf{Prepare:} Given a constraint set $\mathcal{S}$, organize the tokenized $\mathcal{S}$ into a matrix $\mathcal X \in \mathcal V^{|\mathcal S| \times l_\text{max}}$ sorted lexicographically based on the vocabulary order $\mathcal V$. Pad shorter sequences to $l_\text{max}$, i.e., the maximum length of sequences in $\mathcal S$.
\State \textbf{Prepare:} Represent the constraint set $\mathcal{S}$ by a matrix $\mathcal X \in \mathcal V^{|\mathcal S| \times l_\text{max}}$ with rows sorted lexicographically based on the vocabulary order on $\mathcal V$. Pad shorter sequences to $l_\text{max}$, i.e., the maximum length of sequences in $\mathcal S$.\label{code:prepare}
\State \textbf{Input:} Token distribution $P$, partial sequence $\va$
\State Compute the top $M$ entries of $P$ and record their indices $\vt = \{t_1, t_2, \dots, t_M\}$. \label{code:topk}
% \State Compute the top $M$ token indices $\vt = (t_1, t_2, \dots, t_M)$ with the highest probabilities according to $P$.
% \State Construct a candidate matrix $V \in \mathcal{V}^{M \times (|a| + 1)}$ where each row $V_i$ is the concatenation of the partial sequence $a$ with the token $t_i$:
% \[
% V_i = (a_1, a_2, \dots, a_n, t_i) \quad \text{for } i = 1, 2, \dots, M
% \]
\State Construct a candidate matrix $V \in \mathcal{V}^{M \times (|\va| + 1)}$ where each row $V_i = \va + [t_i]$. \label{code:vector}
% \State Binary search whether $\vv$ is a prefix of rows in $A_\mathcal S$ for all $M$ rows in parallel, obtain boolean $\texttt{mask}$ with shape $(M,)$ for all rows.
\State Simultaneously perform a binary search for each candidate sequence $V_i$ in the rows of $\mathcal X$ to determine if $V_i$ is a prefix of any sequence in $\mathcal{S}$. Record the results in a boolean array $\texttt{mask}_V \in \{0,1\}^M$:
\[
\texttt{mask}_V[i] =
\begin{cases}
1 & \text{if } \exists j ~\text{s.t.}~ V_i \preceq \mathcal X _j \\
0 & \text{otherwise}
\end{cases}
\]
\State Extend $\texttt{mask}_V$ to $\texttt{mask} \in \{0,1\}^{|\mathcal V|} $ with $\texttt{mask} [t_i] =1$ if $\texttt{mask}_V[i] =1$ and the remaining to $0$.
% \in \vt$ and from shape $(M, )$ to shape $(|\mathcal V|, )$ by setting values on unslected rows to $\texttt{False}$.
\State \textbf{Return:} Mask of valid tokens $\texttt{mask} $
\end{algorithmic}
\end{algorithm}
\begin{table*}[t]
\centering
\caption{\textbf{Task Statistics}. ``\#Corpus'' refers to the number of valid entities, while ``Avg. corpus token length’’ represents the average number of tokens per entity. ``\#Query'' indicates the number of test queries for each task. ``Vocab size’’ denotes the vocabulary size of the generator LLM, and ``Metrics’’ specifies the evaluation metrics.}
\resizebox{\textwidth}{!}{\begin{tabular}{cccccc}
\toprule
\textbf{Task} & \textbf{\#Corpus} & \textbf{Avg. corpus token length} & \textbf{\#Query} & \textbf{Vocab size} & \textbf{Metrics} \\
\midrule
Entity Disambiguation & 5,903,530 & 7.701 & 24,100 & 50,264 & Micro F1 \\
Document Retrieval & 5,903,530 & 7.701 & 51,304 & 50,264 & R-Precision \\
Entity Retrieval & 4,635,922 & 9.485 & 400 & 128,256 & Relevance score \\
Fact Checking & 5,416,593 & 9.436 & 1,535 & 128,256 & Relevance score \\
\bottomrule
\end{tabular}}
% \vspace{-10pt}
\end{table*}
% Specifically, instead of returning all possible candidate tokens from the constrained set $\mathcal{S}$, we leverage the fact that the language model already computes the scores of all tokens. We select the top $M$ tokens based on their probabilities and verify their validity using the PPV subroutine. This approach resolves the issue of varying sizes and fragmentation inherent in data structures like prefix trees.
% By organizing the candidate tokens and the constrained set $\mathcal{S}$ appropriately, we can perform prefix verification efficiently using parallel search algorithms on GPUs. The verification process can be conducted in $\mathcal{O}(\log N)$ time, where $N$ is the total number of candidates in $\mathcal{S}$. Moreover, since the size $M$ is fixed and typically much smaller than $N$, the computational overhead is minimal.
% An outline of the PPV subroutine is as follows:
\iffalse
\begin{algorithm*}[ht]
\caption{\textbf{P}arallel \textbf{P}refix \textbf{V}erification (PPV)}
\label{alg:ppv+}
\begin{algorithmic}[1]
\Require
\begin{itemize}
\item Constrained set of sequences $\mathcal{S}$
\item Token distribution $P \in \mathbb{R}^{|\mathcal{V}|}$
\item Partial sequence $a = (a_1, a_2, \dots, a_n)$
\item Maximum sequence length $l_{\max}$
\item Number of top candidates $K$
\end{itemize}
\Ensure
\begin{itemize}
\item Mask of valid tokens $\texttt{mask} \in \{0,1\}^{|\mathcal{V}|}$
\end{itemize}
\State \textbf{Prepare Matrix:}
Organize the tokenized set $\mathcal{S}$ into a matrix $M_{\mathcal{S}} \in \mathbb{N}^{|\mathcal{S}| \times l_{\max}}$ sorted lexicographically based on the vocabulary order $\mathcal{V}$. Pad shorter sequences with a designated padding token to ensure uniform length.
\State \textbf{Select Top Candidates:}
Identify the top $K$ token indices $\mathbf{t} = (t_1, t_2, \dots, t_K)$ with the highest probabilities in $P$.
\State \textbf{Generate Candidate Sequences:}
Construct a candidate matrix $V \in \mathcal{V}^{K \times (|a| + 1)}$ where each row $V_i$ is the concatenation of the partial sequence $a$ with the token $t_i$:
\[
V_i = (a_1, a_2, \dots, a_n, t_i) \quad \text{for } i = 1, 2, \dots, K
\]
\State \textbf{Parallel Verification:}
Perform a binary search for each candidate sequence $V_i$ within $M_{\mathcal{S}}$ to determine if $V_i$ is a prefix of any sequence in $\mathcal{S}$. Record the results in a boolean array $\texttt{mask}_V \in \{0,1\}^K$:
\[
\texttt{mask}_V[i] =
\begin{cases}
1 & \text{if } V_i \text{ is a prefix of some } S \in \mathcal{S} \\
0 & \text{otherwise}
\end{cases}
\]
\State \textbf{Expand Mask to Vocabulary Size:}
Initialize $\texttt{mask} \in \{0,1\}^{|\mathcal{V}|}$ to all zeros. For each valid top token $t_i$ where $\texttt{mask}_V[i] = 1$, set $\texttt{mask}[t_i] = 1$.
\State \Return $\texttt{mask}$
\end{algorithmic}
\end{algorithm*}
\fi
% \paragraph{Array preparation.}
% Given a candidate set $\mathcal S$, the sub-routine first conducts a one-time preparation for all subsequent constrained decoding. Specifically, keywords are ordered by initially comparing their first tokens based on their sequence within the vocabulary set $\mathcal{V}$. If the first tokens are identical, the comparison proceeds to the second tokens, continuing sequentially from left to right. A padding token is appended at the end when necessary to ensure uniform length. Once ordered, all keywords are stacked into an sorted array $M_\mathcal S$ on GPUs that allows for efficient binary search for verification.
\paragraph{Array Preparation.}
Given the constraint set $\mathcal{S}$, \ppv first performs a preprocessing in \Cref{code:prepare} so that it is amenable to binary search during the subsequent verification step. % to facilitate all subsequent constrained decoding operations.
Specifically, the sequences in $\mathcal{S}$ are ordered by first comparing their initial tokens according to the order specified by the vocabulary set $\mathcal{V}$. If the initial tokens are identical, the comparison proceeds to the subsequent tokens in a left-to-right manner. A padding token is appended to the end of shorter keywords as necessary to ensure uniform sequence length.
All ordered sequences are aggregated into an array $\mathcal X$ on GPUs, which can be efficiently binary searched during decoding.
In practice, we also split $\mathcal S$ into a few pieces according to the sequence lengths to avoid too much padding that depletes disk storage space.
% \paragraph{Parallel verification.} During generation, given a batch of partially generated keywords and their corresponding distributions of subsequent tokens, \ppv select $M$ candidate tokens with the highest probabilities for each keyword, and concatenate them to original keywords to obtain a $M$ times larger array to be verified. Then, simultaneously for each candidate, \ppv search the smallest keywords in $\mathcal S$ that is larger than the candidate \textit{in a binary manner} on $M_{\mathcal S}$ and obtain a batch of indices. Since $M_\mathcal S$ is ordered left-to-right, \ppv only needs to verify whether the candidate is the prefix of the keyword corresponding to the indices. This generated a mask, which will be further extended to the size of $\mathcal V$ and return to \alg.
% \paragraph{Parallel Verification.} During the generation process, given a batch of partially generated sequences and their corresponding distributions over the next token, \ppv selects the top $M$ candidate tokens with the highest probabilities for each keyword.
% These selected tokens are concatenated to the original keywords, resulting in an array that is $M$ times larger for verification. Then, simultaneously for each candidate in this expanded array, \ppv performs a binary search on the sorted array $A$ to identify the smallest keyword in $\mathcal{S}$ that is lexicographically larger than the candidate. Since $A$ is ordered from left to right, \ppv can find the keyword in $\log(|\mathcal S|)$ time, and check whether each candidate is a prefix of the corresponding keyword identified by the indices. This verification process generates a mask, which is subsequently expanded to match the size of the vocabulary set $\mathcal{V}$ and returned to \alg for further processing.
\paragraph{Parallel Verification.} During the generation process, \ppv selects the top $M$ candidate tokens $\vt = \{t_1,\cdots,t_M\} $ with the highest probabilities for each partially generated sequence $\va$.
Each $t_m$ is concatenated to $\va$, producing an array $\mathcal V$ for verification (\Cref{code:vector}).
Then, simultaneously for each row $V_m$, \ppv performs a binary search on the sorted array $\mathcal X$ to identify the index $i_m$, where the row $\mathcal X_{i_m}$ is the smallest row that is lexicographically larger than the $V_m$. Since $A$ is ordered from left to right, \ppv can find the target in $\log(|\mathcal S|)$ time.
Importantly, we can derive that if $t_m \in \mathcal A_{\va, \mathcal S}$ if and only if $V_m \preceq \mathcal X_{i_m}$.
Therefore, a simple check tells whether $t_m$ is a valid subsequent token.
Notice that all computations are performed on GPUs, and additional batching across multiple queries can further extend parallel processing.
% This verification process returns a mask specifying cal{V}$ and returned to \alg for further processing.
\paragraph{Selection of top candidate tokens.}
In \Cref{code:topk}, \ppv performs a sorting and selects only the top $M$ candidate tokens.
In theory, we should set $M = |\mathcal V|$ to avoid missing valid tokens with low probability.
However, we argue that this does not lead to a practical issue.
First, most modern LLMs, when doing inference, will limit the choices to the tokens with top $50$ probabilities in the first place \citep{dubey2024llama}.
This implies that verification essentially does not miss any valid candidates. In addition, our efficient parallel verification makes the verification cost almost identical when setting $M$ to $50$ or even $500$, as verified in the appendix. Additionally, the top $50$ probabilities already sum close to $1$, so there is no actual difference in setting $M$ to different values.
% The choice of $M$ balances the trade-off between computational efficiency and the coverage of valid tokens. A larger $M$ increases the likelihood of including valid tokens but also increases computational overhead. In practice, $M$ can be set based on the computational resources available and the specific requirements of the application.
% \paragraph{Complexity Analysis on GPUs}
% We now turn to theoretical analysis on computation complexity. For each call to \ppv, we need to verify $M$ candidates, each using a binary search. Therefore, the total complexity is $\mathcal O(M \log |\mathcal S|)$. In comparison, the complexity for trie is
% The PPV subroutine is highly parallelizable and leverages the parallel processing capabilities of GPUs. By performing the verification process for all tokens in the top $M$ candidates simultaneously across all batches, we achieve significant speedups compared to CPU-based implementations. This approach avoids the frequent data transfers between CPU and GPU inherent in traditional methods like prefix trees, enhancing both scalability and efficiency.
\paragraph{Advantages over prefix trees.}
\ppv eliminates the need for CPU-based data structures like prefix trees as well as frequent data transfer between devices. In addition, its inference can be further parallelized across samples in a batch, leading to additional acceleration. Lastly, we notice that trie (often stored as JSON \citep{pezoa2016foundations}) loading is extremely slower compared with array loading. The evaluation time for a dataset can be $8.5$ times different at most.
% By focusing on verifying a fixed number of top candidates and utilizing efficient search algorithms, we overcome the challenges associated with implementing tries on GPUs.
% \paragraph{Integration with \alg{}}
% In the \alg{} algorithm, we utilize the PPV subroutine in the line $mask = C(P, a)$ to obtain the mask of valid tokens. This seamless integration ensures that the constrained decoding process is both unbiased and computationally efficient, making it practical for real-world applications with large candidate sets.
% \subsection{Theorem (vanilla version)}
% \begin{theorem}
% For any distribution $\mathbb P_{all}$ that any auto-regressive generative model and any text prefix induces, denote $\mathbb P_{keywords}$ the distribution of all valid keywords in set $\mathcal W$ over $\mathbb P_{all}$, and $\mathbb P^K_{IS}$ the sequential importance sampling distribution with $K$ samples under constrained decoding. Then the KL divergence between the two distributions can be bounded by
% \begin{align}
% \text{KL}(\mathbb P_{keywords} \| \mathbb P_{IS}^K) \leq \mathcal O \Big(\frac{1- P}{P^2 K^{1-P}} + \sqrt{\frac{(1-P)\ln K}{K}}\Big),
% \end{align}
% where $P = \sum_{w \in \mathcal W}p_{all}(w)$ is the probability on all keywords, and $\mathcal O$ hides universal constants.
% \end{theorem}
% \section{METHOD}
% In regard to the myopic sampling of LLMs due to unawareness of constrained decoding, we propose an improved sampling algorithm named "dynamic importance sampling". Compared with fixed-size importance sampling, this method is much more efficient, and results in a significantly better convergence rate. The pseudo-code of DIS is presented in \Cref{alg:DIS}. \theertha{In line 10 of the algorithm, we use $A = A + [a]$ and $S = S + [s]$, which according to earlier definition, means concatenation, but I believe we mean adding to the set $A$, please change notation to reflect the same.}
% \begin{algorithm}[ht]
% % \small
% \caption{\textbf{D}ynamic \textbf{I}mportance \textbf{S}ampling for \textbf{C}onstrained Decoding}
% \label{alg:DIS}
% \begin{algorithmic}[1]
% \State \textbf{Input:} LLM $L$, Constrainer $C$, Maximum allowed sampling step $K$, prefix token list $q$, temperature $T$.
% \State Initialize candidate $a = []$, importance score $ s = 1, k = 1$
% \State Initialize candidate buffer $A = []$, score buffer $S = []$
% \While {$a[-1]$ is not EOK and $k \leq K$}
% \State Get next token distribution $P= L(q + a, T)$
% \State Get mask of valid tokens $mask = C(P, a)$
% \State Sample $t$ over distribution $\frac{P \odot mask}{\| P \odot mask\|_1}$
% \State Update importance score $s = s \times \| P \odot mask\|_1$ if $a \neq []$
% \If {$t$ is EOK and $s \leq \epsilon \sim \mathcal U[0,1]$}
% % \State Sample $\epsilon \sim \mathcal U[0,1]$
% % \If {$\epsilon > s$}
% \State $A = A + [a], S = S + [s], k = k + 1$
% \State Reset $a = [], s = 1$
% % \EndIf
% \Else
% \State $a = a + [t]$
% \EndIf
% \EndWhile
% \If {$k = K + 1$}
% \State Sample $a$ from $A$ with probability proportional to $S$.
% \EndIf
% \State \textbf{Output:} sampled candidate $a$
% \end{algorithmic}
% \end{algorithm}
% \theertha{Clarify line 17 to say corresponding $s$ in $S$.}
% \theertha{explicitly state $P^K_{IS}$ is Algorithm 1}
% \begin{theorem}
% \label{thm:alg}
% For any distribution autoregressive model $L$ (equipped with any text prefix), denote $ P^K_{IS}$ the dynamic importance sampling distribution with $K$ samples in maximum under constrained decoding. Then, the KL divergence between $ P^K_{IS}$ and $P_\mathcal S$ can be bounded by
% \begin{align}
% \text{KL}(P_\mathcal S \| P_{IS}^K) \leq \mathcal O \Big(\frac{P_{b}^K}{1 - P_b} \Big( \frac{P_{b}}{K} + \sqrt{\frac{2\ln K}{K}}\Big)\Big),
% \end{align}
% and the expected sampling step is $\frac{1 - P_b^K}{1- P_b}$.
% % Here $P_b = \sum_{w \in_{p} \mathcal W, w \notin \mathcal W}p_{all}(w)$ is the probability on all keywords,
% Here $P_b = 1 - P_L(\mathcal S)$ is the probability on all invalid sequences,
% and $\mathcal O$ hides universal constants.
% \end{theorem}
% \begin{proof}
% Just a scratch here.
% % We use $p(x)$ throughout the proof to represent the probability that $L$ generates $x$ (and then $EOK$) conditioned on $q$.
% % We first compute the probability that all $K$ generated candidates do not meet the requirement ($a$ is reset to $[]$). Notice that for a
% \end{proof}
% \paragraph{Comparison with beam search.}
% \paragraph{Expected sampling steps.}
% \paragraph{Bound on KL divergence}
% \subsection{GPU-based implementation}
% An pesudo-code for the sub-routine.
% \paragraph{Selection of top candidate $M$.}
% \paragraph{Complexity analysis in GPUs.}
\input{Sections/large_table}
\section{EXPERIMENTS}
This section conducts comprehensive experiments to demonstrate the efficacy of our methods. Without further explanation, \alg is always equipped with \ppv.
\subsection{Experimental Settings}
%Additional settings can be found in the appendix.
\paragraph{Tasks \& datasets.} We evaluate on four tasks. The first two tasks are \emph{Entity Disambiguation}, which tries to find the most relevant entities from a list given an article, and \emph{Document Retrieval}, which tries to retrieve the most relevant documents.
Each task contains six and eleven datasets, respectively.
\cite{Cao2021} fine-tunes BART models to generate entities using trie-based constrained decoding.
We follow their settings and generate entities from the set $\mathcal S$ that contains approximately six million Wikipedia pages. These datasets help compare our methods with the conventional trie-based approach using models fine-tuned on relevant tasks. Notice that \cite{Cao2021} provides results with additional tricks, such as manually setting $\mathcal S$ to a minimal candidate set for each query to achieve the best performance. Our paper seeks to provide a better constrained decoding algorithm, and it suffices to compare algorithms without using these tricks. That said, our method can be seamlessly combined with these add-ons. For instance, we demonstrate in the appendix that our method is still superior when combined with beam search.
Apart from existing benchmarks, we innovatively study the in-context learning setting for generative retrieval tasks, including \emph{Entity Retrieval} from DBPedia~\citep{hasibi2017dbpedia} and \emph{Fact Checking} with Wikipedia. We use the newly released LLAMA3.2-3B-INSTRUCT\footnote{\url{https://www.llama.com/}} as a generator. For each task, we select three query-answer pairs from the validation set and provide them in-context to the model to clarify the intentions and output formats.
% We use the newly released \texttt{LLaMA-3.2-Instruct-3B}\ac as backbone LLM.
All datasets can be found from the information retrieval benchmark BEIR~\citep{thakur2021beir}.
% We use the datasets from zero-shot information retrieval benchmark BEIR\ac, i.e., DBPedia-Entity\ac, Climate-FEVER~\citep{diggelmann2020climate}, and SCIDOCS\ac for Entity Retrieval, Fact Checking, and Citation Prediction, respectively.
\paragraph{Methods.} We compare vanilla generation (without constraints), trie-based constrained decoding, and \alg equipped with \ppv. For our method, the number of maximum sampling step $K$ is set to $\{1,2,3,4\}$, and the top candidate's size $M$ is set to $50$.
Notice that when $K=1$, the performance of \alg should be the same as trie-based method (but faster).
For sampling, we set model temperature $T$ to $1$ and do not include length penalty (rescale probabilities by length). We use A100 GPUs and set batch size to $128$.
\paragraph{Evaluation metrics.}
For Entity Disambiguation and Document Retrieval, we adopt the \emph{InKB} micro-F1 following~\citet{le2018improving}, and R-precision~\citep{fischer2017privacy} following~\citet{petroni2020kilt} respectively. For in-context learning tasks, we utilize relevance score, which is calculated as the average query-document relevance using TREC evaluation tool~\citep{van2018pytrec_eval}.
\subsection{Main Results}
\begin{figure*}[t]
\centering
\includegraphics[width=1.0\textwidth]{Figures/curve.pdf}
\vspace{-15pt}
\caption{Time vs. Performance Comparison on Document Retrieval and Entity Disambiguation Tasks for all methods. The y-axis presents performance (R-Precision for Document Retrieval and Micro F1 scores for Entity Disambiguation), and the x-axis presents inference time that is average across all tasks belonging to that task.
The improvement on time (x-axis) is resulted from \ppv, and the improvement on performance (y-axis) is resulted from \alg, respectively.
}
\label{fig:time-performance-curve}
\vspace{-5pt}
\end{figure*}
% We evaluate the performance of our proposed methods, \alg equipped with \ppv, on the four tasks outlined in the experimental settings. The results demonstrate that our methods not only significantly improve generation quality over baseline approaches but also enhance efficiency.
\paragraph{Entity Disambiguation and Document Retrieval.} \Cref{tab:entity,tab:document} present the performance of unconstrained decoding (Vanilla), trie-based constrained decoding, and our method \alg with varying values of maximum sampling steps $K$. On both tasks, \alg consistently outperforms the trie-based constrained decoding across all datasets and models. On Entity Disambiguation with $K=2$, the average Micro F1 gain compared to Vanilla is increased from $0.1\%$ ($2.5\%$) to $3.3\%$ ($3.4\%$) for the BLINK (AY2) model.
% With $K=4$, \alg achieves the highest average accuracy of 0.8004 (BLINK) and 0.7889 (AY2), representing an improvement of \textcolor{red}{+4.79\%} and \textcolor{red}{+5.00\%} over the trie-based method, respectively. Notably, on the ACE2004 dataset, \alg attains an accuracy of 0.8444 (BLINK), significantly surpassing the trie-based method's 0.8093.
For Document Retrieval, \alg achieves an average R-Precision gain of $5.2\%$, nearly double the $2.8\%$ improvement achieved by the trie-based method.
% of 0.6399, outperforming the trie-based method by \textcolor{red}{+8.75\%}.
Remarkably, the effect of improving $K$ quickly plateaus when $K$ increases above $2$, while the difference between $K=1,2$ is non-trivial.
The improvements are particularly pronounced on datasets like HoPo and NQ, where the R-Precision improvements are $21.58\%$ and $16.89\%$ compared to the Vanilla model, respectively.
% The inference time for each algorithm and model is presented in \Cref{fig:time-performance-curve}.
% Results are averaged across all datasets in the task.
% The efficiency advantages of \alg are obvious, with $K=1$ being significantly faster, and $K=2$ being much better performance and still more efficient. The fact that the sweet trade-off holds across all tasks clearly implies the advantage of \ppv,
The inference time for each algorithm and model is presented in \Cref{fig:time-performance-curve}. Results are averaged across all datasets. The advantages of \alg{} are evident, with $K = 1$ being significantly faster and $K=2$ offering much better performance while maintaining higher efficiency. The consistency of this optimal trade-off across all tasks clearly highlights the benefits of \ppv{}.
% These findings demonstrate that \alg{} effectively balances speed and performance, making it a highly effective solution for constrained decoding in diverse real-world applications.
\begin{table}[h]
\centering
% \vspace{-1em}
\caption{The relevance score and inference time of all methods on Entity Retrieval and Fack Checking. Models are not tuned and only prompted in context.}
\label{tab:icl}
\resizebox{0.9\linewidth}{!}{\begin{tabular}{ccccc}
\toprule
& \multicolumn{2}{c}{\textbf{DBPedia-Entity}} & \multicolumn{2}{c}{\textbf{Climate-Fever}} \\
\cmidrule{2-5}
& Score & Time (s) & Score & Time (s) \\
\midrule
Vanilla & 0.340 & 20.20 & 0.160 & 73.4 \\
\midrule
Trie-based & 0.600 & 171.73 & 0.271 & 234.8 \\
\midrule
\alg ($K=1$) & 0.640 & 20.41 & 0.270 & 74.0 \\
\alg ($K=2$) & 0.683 & 37.92 & 0.294 & 149.9 \\
\alg ($K=3$) & 0.728 & 58.04 & 0.290 & 225.1 \\
\alg ($K=4$) & \textbf{0.735} & 76.88 & \textbf{0.295} & 304.7 \\
\midrule
Delta (\%) & \textcolor{red}{+116.18} & - & \textcolor{red}{+84.44} & - \\
\bottomrule
\end{tabular}}
\vspace{-1em}
\end{table}
\paragraph{Entity Retrieval and Fact Checking.} \Cref{tab:icl} presents the performance and efficiency comparison among different methods on the Entity Retrieval (DBPedia-Entity) and Fact Checking (Climate-FEVER) tasks.
On the first task, \alg with $K=4$ achieves a score of $0.735$, more than doubling the performance of Vanilla ($0.340$) and $22\%$ higher than the trie-based method.
Surprisingly, this is realized with nearly half of the inference time (since $|\mathcal S|$ is huge is this dataset).
Similarly, on the Climate-FEVER dataset, \alg with $K=4$ attains a score of $0.295$, marking an $84.4\%$ improvement over the Vanilla model and a $2.4\%$ increase over the trie-based method.
As this in-context learning pipeline is novel and purely zero-shot, it clearly demonstrates the potential of \alg's capability to enhance performance without task-specific fine-tuning.
% These results confirm that \alg effectively enhances the generation quality for tasks requiring constrained outputs, without the need for additional tricks or task-specific fine-tuning.
% A $2\times 2$ figure specifying the time-performance curve.
% Two huge table corresponding to two tasks
% \subsection{Analysis of \alg}
% Problems discussed in the section includes:
% \paragraph{Time analysis.} Loading time / inference time under different set size $\mathcal S$, different $K$ and different $M$.
% \paragraph{Model quality.} Fine-tuning vs no fine-tuning (ICL); different models;
% \paragraph{Applying beam searchs} compare the setting where we apply beam search to both cases.
% \begin{table*}[t]
% \centering
% \caption{something}
% \label{tab:entity}
% \resizebox{\textwidth}{!}{\begin{tabular}{ccccccccccccccc}
% \toprule
% \textbf{Dataset} & \multicolumn{2}{c}{\textbf{ACE2004}} & \multicolumn{2}{c}{\textbf{AIDA}} & \multicolumn{2}{c}{\textbf{AQUAINT}} & \multicolumn{2}{c}{\textbf{CWEB}} & \multicolumn{2}{c}{\textbf{MSNBC}} & \multicolumn{2}{c}{\textbf{WIKI}} & \multicolumn{2}{c}{\textbf{Avg.}} \\% & \textbf{Time} \\
% \midrule
% \textbf{Models} & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK \\% & \\
% \midrule
% Vanilla & 0.8016 & 0.7938 & 0.7697 & 0.8297 & 0.7923 & 0.7909 & 0.6150 & 0.6299 & 0.7424 & 0.7607 & 0.7868 & 0.7779 & 0.7513 & 0.7638 \\% & 39.4 \\
% \midrule
% Trie & 0.8483 & 0.8093 & 0.7962 & 0.8303 & 0.8157 & 0.7923 & 0.6346 & 0.6287 & 0.7607 & 0.7698 & 0.8030 & 0.7788 & 0.7764 & 0.7682 \\% & 99.5 \\
% \midrule
% \alg ($K=1$) & 0.8405 & 0.8093 & 0.7971 & 0.8303 & 0.8143 & 0.7923 & 0.6347 & 0.6287 & 0.7652 & 0.7698 & 0.8031 & 0.7788 & 0.7758 & 0.7682 \\% & 55.4 \\
% \alg ($K=2$) & 0.8483 & 0.8444 & 0.8025 & 0.8517 & 0.8212 & 0.8129 & 0.6469 & 0.6637 & 0.7805 & 0.8034 & 0.8110 & 0.8049 & 0.7851 & 0.7968 \\% & 89.2 \\
% \alg ($K=3$) & 0.8483 & 0.8444 & 0.8078 & 0.8524 & \textbf{0.8308} & \textbf{0.8129} & \textbf{0.6510} & 0.6648 & \textbf{0.7820} & 0.7957 & 0.8096 & 0.8082 & 0.7882 & 0.7964 \\% & 118.5 \\
% \alg ($K=4$) & \textbf{0.8488} & \textbf{0.8444} & \textbf{0.8103} & \textbf{0.8598} & 0.8267 & 0.8116 & 0.6507 & \textbf{0.6674} & 0.7805 & \textbf{0.8095} & \textbf{0.8162} & 0.8100 & \textbf{0.7889} & \textbf{0.8004} \\% & 146.6 \\
% \midrule
% Delta & \textcolor{red}{+5.89\%} & \textcolor{red}{+6.37\%} & \textcolor{red}{+5.27\%} & \textcolor{red}{+3.63\%} & \textcolor{red}{+4.34\%} & \textcolor{red}{+2.61\%} & \textcolor{red}{+5.80\%} & \textcolor{red}{+5.95\%} & \textcolor{red}{+5.13\%} & \textcolor{red}{+6.41\%} & \textcolor{red}{+3.73\%} & \textcolor{red}{+4.13\%} & \textcolor{red}{+5.00\%} & \textcolor{red}{+4.79\%} \\%& \\
% \bottomrule
% \end{tabular}}
% \end{table*}
% \begin{table*}[t]
% \centering
% \caption{something}
% \label{tab:entity}
% \resizebox{\textwidth}{!}{\begin{tabular}{ccccccccccccccc}
% \toprule
% \textbf{Dataset} & \multicolumn{2}{c}{\textbf{ACE2004}} & \multicolumn{2}{c}{\textbf{AIDA}} & \multicolumn{2}{c}{\textbf{AQUAINT}} & \multicolumn{2}{c}{\textbf{CWEB}} & \multicolumn{2}{c}{\textbf{MSNBC}} & \multicolumn{2}{c}{\textbf{WIKI}} & \multicolumn{2}{c}{\textbf{Avg.}} \\% & \textbf{Time} \\
% \midrule
% \textbf{Models} & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 \\% & \\
% \midrule
% Vanilla & 0.7938 & 0.8016 & 0.8297 & 0.7697 & 0.7909 & 0.7923 & 0.6299 & 0.6150 & 0.7607 & 0.7424 & 0.7779 & 0.7868 & 0.7638 & 0.7513 \\% & 39.4 \\
% \midrule
% Trie & 0.8093 & 0.8483 & 0.8303 & 0.7962 & 0.7923 & 0.8157 & 0.6287 & 0.6346 & 0.7698 & 0.7607 & 0.7788 & 0.8030 & 0.7682 & 0.7764 \\% & 99.5 \\
% \midrule
% \alg ($K=1$) & 0.8093 & 0.8405 & 0.8303 & 0.7971 & 0.7923 & 0.8143 & 0.6287 & 0.6347 & 0.7698 & 0.7652 & 0.7788 & 0.8031 & 0.7682 & 0.7758 \\% & 55.4 \\
% \alg ($K=2$) & 0.8444 & 0.8483 & 0.8517 & 0.8025 & 0.8129 & 0.8212 & 0.6637 & 0.6469 & 0.8034 & 0.7805 & 0.8049 & 0.8110 & 0.7968 & 0.7851 \\% & 89.2 \\
% \alg ($K=3$) & 0.8444 & 0.8483 & 0.8524 & 0.8078 & \textbf{0.8129} & \textbf{0.8308} & 0.6648 & \textbf{0.6510} & 0.7957 & \textbf{0.7820} & 0.8082 & 0.8096 & 0.7964 & 0.7882 \\% & 118.5 \\
% \alg ($K=4$) & \textbf{0.8444} & \textbf{0.8488} & \textbf{0.8598} & \textbf{0.8103} & 0.8116 & 0.8267 & \textbf{0.6674} & 0.6507 & \textbf{0.8095} & 0.7805 & \textbf{0.8100} & \textbf{0.8162} & \textbf{0.8004} & \textbf{0.7889} \\% & 146.6 \\
% \midrule
% $\Delta$ & \textcolor{red}{+6.37\%} & \textcolor{red}{+5.89\%} & \textcolor{red}{+3.63\%} & \textcolor{red}{+5.27\%} & \textcolor{red}{+2.61\%} & \textcolor{red}{+4.34\%} & \textcolor{red}{+5.95\%} & \textcolor{red}{+5.80\%} & \textcolor{red}{+6.41\%} & \textcolor{red}{+5.13\%} & \textcolor{red}{+4.13\%} & \textcolor{red}{+3.73\%} & \textcolor{red}{+4.79\%} & \textcolor{red}{+5.00\%} \\%& \\
% \bottomrule
% \end{tabular}}
% \end{table*}
% \begin{table*}[t]
% \centering
% % \small
% \caption{something}
% \label{tab:document}
% \resizebox{\textwidth}{!}{
% \begin{tabular}{ccccccccccccc}
% \toprule
% \textbf{Dataset} & \textbf{AY2} & \textbf{CWEB} & \textbf{ELI5} & \textbf{FEV} & \textbf{HoPo} & \textbf{NQ} & \textbf{zsRE} & \textbf{TREx} & \textbf{TQA} & \textbf{WNED} & \textbf{WoW} & \textbf{Avg.} \\% & \textbf{Time} \\
% \midrule
% Vanilla & 0.8842 & 0.6465 & 0.0982 & 0.7568 & 0.2371 & 0.4843 & 0.8230 & 0.6746 & 0.5438 & 0.8177 & 0.5062 & 0.5884 \\% & 54.2 \\
% \midrule
% Trie-based & 0.8942 & 0.6696 & 0.1042 & 0.7863 & 0.2646 & 0.5375 & 0.8601 & 0.7540 & 0.5555 & 0.8522 & 0.5023 & 0.6164 \\% & 116.3 \\
% \midrule
% \alg ($K=1$) & 0.8999 & 0.6707 & 0.1042 & 0.7902 & 0.2655 & 0.5217 & 0.8617 & 0.7494 & 0.5613 & 0.8469 & 0.5082 & 0.6163 \\% & 61.4 \\
% \alg ($K=2$) & 0.9051 & 0.6828 & 0.1155 & 0.8026 & 0.2838 & 0.5629 & 0.8953 & 0.7732 & 0.5906 & \textbf{0.8628} & 0.5275 & 0.6365 \\% & 104.3 \\
% \alg ($K=3$) & 0.9072 & 0.6832 & 0.1128 & 0.8025 & 0.2874 & 0.5622 & \textbf{0.9060} & \textbf{0.7740} & 0.6009 & 0.8572 & 0.5305 & 0.6385 \\% & 140.6 \\
% \alg ($K=4$) & \textbf{0.9122} & \textbf{0.6873} & \textbf{0.1135} & \textbf{0.8037} & \textbf{0.2882} & \textbf{0.5661} & 0.8955 & 0.7734 & \textbf{0.6046} & 0.8625 & \textbf{0.5321} & \textbf{0.6399} \\% & 175.7 \\
% \midrule
% $\Delta$ & \textcolor{red}{+3.17\%} & \textcolor{red}{+6.30\%} & \textcolor{red}{+15.54\%} & \textcolor{red}{+6.20\%} & \textcolor{red}{+21.58\%} & \textcolor{red}{+16.89\%} & \textcolor{red}{+8.81\%} & \textcolor{red}{+14.65\%} & \textcolor{red}{+11.19\%} & \textcolor{red}{+5.47\%} & \textcolor{red}{+5.11\%} & \textcolor{red}{8.75\%} \\%& \\
% \bottomrule
% \end{tabular}}
% \end{table*}
\begin{table*}[t]
\centering
\caption{The Micro F1 score of unconstrained decoding (Vanilla), trie-based constrained decoding, and \alg on Entity Disambiguation across six datasets and two models. \textbf{Avg.} is the average score across datasets. The best score is bold, and $\Delta$ (in percentage) is the score improvement from Vanilla to \alg ($K=4$).}
\label{tab:entity}
\resizebox{\textwidth}{!}{\begin{tabular}{ccccccccccccccc}
\toprule
\textbf{Dataset} & \multicolumn{2}{c}{\textbf{ACE2004}} & \multicolumn{2}{c}{\textbf{AIDA}} & \multicolumn{2}{c}{\textbf{AQUAINT}} & \multicolumn{2}{c}{\textbf{CWEB}} & \multicolumn{2}{c}{\textbf{MSNBC}} & \multicolumn{2}{c}{\textbf{WIKI}} & \multicolumn{2}{c}{\textbf{Avg.}} \\% & \textbf{Time} \\
\midrule
\textbf{Models} & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 & BLINK & AY2 \\% & \\
\midrule
Vanilla & 0.794 & 0.802 & 0.830 & 0.770 & 0.791 & 0.792 & 0.630 & 0.615 & 0.761 & 0.742 & 0.778 & 0.787 & 0.764 & 0.751 \\% & 39.4 \\
\midrule
Trie & 0.809 & 0.848 & 0.830 & 0.796 & 0.792 & 0.816 & 0.629 & 0.635 & 0.770 & 0.761 & 0.779 & 0.803 & 0.768 & 0.776 \\% & 99.5 \\
\midrule
\alg ($K=1$) & 0.809 & 0.841 & 0.830 & 0.797 & 0.792 & 0.814 & 0.629 & 0.635 & 0.770 & 0.765 & 0.779 & 0.803 & 0.768 & 0.776 \\% & 55.4 \\
\alg ($K=2$) & 0.844 & 0.848 & 0.852 & 0.803 & 0.813 & 0.821 & 0.664 & 0.647 & 0.803 & 0.781 & 0.805 & 0.811 & 0.797 & 0.785 \\% & 89.2 \\
\alg ($K=3$) & 0.844 & 0.848 & 0.852 & 0.808 & \textbf{0.813} & \textbf{0.831} & 0.665 & \textbf{0.651} & 0.796 & \textbf{0.782} & 0.808 & 0.810 & 0.796 & 0.788 \\% & 118.5 \\
\alg ($K=4$) & \textbf{0.844} & \textbf{0.849} & \textbf{0.860} & \textbf{0.810} & 0.812 & 0.827 & \textbf{0.667} & 0.651 & \textbf{0.810} & 0.781 & \textbf{0.810} & \textbf{0.816} & \textbf{0.800} & \textbf{0.789} \\% & 146.6 \\
\midrule
$\Delta$ (\%) & \textcolor{red}{+6.37} & \textcolor{red}{+5.89} & \textcolor{red}{+3.63} & \textcolor{red}{+5.27} & \textcolor{red}{+2.61} & \textcolor{red}{+4.34} & \textcolor{red}{+5.95} & \textcolor{red}{+5.80} & \textcolor{red}{+6.41} & \textcolor{red}{+5.13} & \textcolor{red}{+4.13} & \textcolor{red}{+3.73} & \textcolor{red}{+4.79} & \textcolor{red}{+5.00} \\%& \\
\bottomrule
\end{tabular}}
\vspace{-5pt}
\end{table*}
\begin{table*}[t]
\centering
% \small
\caption{The R-Precision of all methods on Document Retrieval across eleven datasets. Similar to \Cref{tab:entity}.}
\label{tab:document}
\resizebox{\textwidth}{!}{
\begin{tabular}{ccccccccccccc}
\toprule
\textbf{Dataset} & \textbf{AY2} & \textbf{CWEB} & \textbf{ELI5} & \textbf{FEV} & \textbf{HoPo} & \textbf{NQ} & \textbf{zsRE} & \textbf{TREx} & \textbf{TQA} & \textbf{WNED} & \textbf{WoW} & \textbf{Avg.} \\% & \textbf{Time} \\
\midrule
Vanilla & 0.884 & 0.647 & 0.0982 & 0.757 & 0.237 & 0.484 & 0.823 & 0.675 & 0.544 & 0.818 & 0.506 & 0.588 \\% & 54.2 \\
\midrule
Trie-based & 0.894 & 0.670 & 0.104 & 0.786 & 0.265 & 0.538 & 0.860 & 0.754 & 0.556 & 0.852 & 0.502 & 0.616 \\% & 116.3 \\
\midrule
\alg ($K=1$) & 0.900 & 0.671 & 0.104 & 0.790 & 0.266 & 0.522 & 0.862 & 0.749 & 0.561 & 0.847 & 0.508 & 0.616 \\% & 61.4 \\
\alg ($K=2$) & 0.905 & 0.683 & 0.116 & 0.803 & 0.284 & 0.563 & 0.895 & 0.773 & 0.591 & \textbf{0.863} & 0.528 & 0.637 \\% & 104.3 \\
\alg ($K=3$) & 0.907 & 0.683 & 0.113 & 0.803 & 0.287 & 0.562 & \textbf{0.906} & \textbf{0.774} & 0.601 & 0.857 & 0.531 & 0.639 \\% & 140.6 \\
\alg ($K=4$) & \textbf{0.912} & \textbf{0.687} & \textbf{0.114} & \textbf{0.804} & \textbf{0.288} & \textbf{0.566} & 0.895 & 0.773 & \textbf{0.605} & 0.863 & \textbf{0.532} & \textbf{0.640} \\% & 175.7 \\
\midrule
$\Delta$ (\%) & \textcolor{red}{+3.17} & \textcolor{red}{+6.30} & \textcolor{red}{+15.5} & \textcolor{red}{+6.20} & \textcolor{red}{+21.6} & \textcolor{red}{+16.9} & \textcolor{red}{+8.81} & \textcolor{red}{+14.7} & \textcolor{red}{+11.2} & \textcolor{red}{+5.47} & \textcolor{red}{+5.11} & \textcolor{red}{+8.75} \\%& \\
\bottomrule
\end{tabular}}
\vspace{-5pt}
\end{table*}
\section{RELATED WORKS }
\paragraph{Constrained decoding.} The objective of constrained decoding is to ensure that the generated sequence complies with the specified set of constraints. These constraints could be in the form of lexical constraints \citep{Hokamp2017, Anderson2017, Post2018} or regular expressions (REGEX) or grammatical structures defined by context-free grammars (CFGs) \citep{Geng2023, Kellner2024} or just a list of possible outputs \citep{Cao2021}.
Several approaches have been proposed to impose these constraints during the sequence generation process. For example, \citet{Cao2021} uses a trie structure, and \citet{Miao2019} utilizes Metropolis-Hastings sampling. While these approaches address the need to generate valid outputs that satisfy all the specified constraints, they often struggle with computational efficiency when the constraint set is large \citep{Stahlberg2019}. Therefore, generating valid outputs on-the-fly without exhaustive enumeration remains a challenge.
\paragraph{Generative retrieval.}
% Extreme classification or information retrieval in general, aims at finding documents relevant to a given query from a very large set of documents (typically $O(B)$). Traditionally this is done by learning query and document embeddings and either doing an approximate nearest neighbor search or by parsing a learned tree/hierarchy over the document set [\cite{menon22},\cite{bhatia15}].
Recently, new approaches for retrieval have emerged whereby LLMs directly generate relevant documents for a given query without explicitly computing embeddings or doing nearest neighbor search \citep{rajput24}. There are primarily two ways of doing this: (1) the entire set of documents is put in the context of LLM, and then LLM generates relevant documents from this context \citep{lee2024can}, and (2) documents are not put in the context but to make sure that the generated documents are from the retrieval set, the decoding process is constrained \citep{Cao2021}. In applications where the retrieval set is enormous, the first approach would require the LLM to have $\mathcal O(N)$ context length where $N$ is the size of documents. As of today, this is too computationally expensive. In this paper, we focus on the second set of approaches.
% \paragraph{Discuss other tasks that we can potentially apply our method to.}
% \documentclass[twoside]{article}
% \usepackage{aistats2025}
% If your paper is accepted, change the options for the package
% aistats2025 as follows:
%
%\usepackage[accepted]{aistats2025}
%
% This option will print headings for the title of your paper and
% headings for the authors names, plus a copyright note at the end of
% the first column of the first page.
% If you set papersize explicitly, activate the following three lines:
%\special{papersize = 8.5in, 11in}
%\setlength{\pdfpageheight}{11in}
%\setlength{\pdfpagewidth}{8.5in}
% If you use natbib package, activate the following three lines:
%\usepackage[round]{natbib}
%\renewcommand{\bibname}{References}
%\renewcommand{\bibsection}{\subsubsection*{\bibname}}
% If you use BibTeX in apalike style, activate the following line:
%\bibliographystyle{apalike}
%\begin{document}
% If your paper is accepted and the title of your paper is very long,
% the style will print as headings an error message. Use the following
% command to supply a shorter title of your paper so that it can be
% used as headings.
%
%\runningtitle{I use this title instead because the last one was very long}
% If your paper is accepted and the number of authors is large, the
% style will print as headings an error message. Use the following
% command to supply a shorter version of the authors names so that
% they can be used as headings (for example, use only the surnames)
%
%\runningauthor{Surname 1, Surname 2, Surname 3, ...., Surname n}
% Supplementary material: To improve readability, you must use a single-column format for the supplementary material.
\onecolumn
\aistatstitle{Efficient and Asymptotically Unbiased\\Constrained Decoding for Large Language Models\\(Supplementary Materials)}
% \vspace{-10em}
% something here.
% We present the proofs of theorems, the PyTorch code for our parallel prefix-verification (\texttt{PPV}) algorithm, and additional experimental settings and results in the supplementary materials.
\input{Appendices/proofs}
\input{Appendices/code}
\input{Appendices/addition_exp_settings}
\input{Appendices/add_results}
% \end{document}
\section{Proofs of Theorems}
\label{app:proof}
\subsection{Proof of \Cref{thm:alg}}
The logic of the proof can be decomposed into two parts. First, we compute the probability that upon $K$ rounds of samples, at least one sample is accepted. We demonstrate that the distribution will be unbiased in this case. Second, when $K$ new samples are re-drawn, we demonstrate that the importance sampling can upper bound the remaining KL divergence. Specifically, assume we use $Q$ to denote the sequence distribution of \Cref{code:IS}, i.e., the distribution that $K$ samples $\{\va_k\}_{k=1}^K$ are drawn using constrained decoding and further resample based on scores $\{x_k\}_{k=1}^K$. Then we upper bound $KL(P_\mathcal S \| Q)$ and use it to derive the eventual bound.
We first derive a few quantities useful for the proof.
According to \Cref{alg:disc}, $x$ is the importance score depending solely on sequence $\va$, and we therefore rewrite it as $$x(\va) = \prod_{i \in \{1, \cdots, |\va|\}} P_L( \mathcal A_{\va_{<i}, \mathcal S} | \va_{<i}).$$Intuitively, $x(\va)$ denotes the probability of alternatives in $\mathcal S$ along the token paths, which is smaller than the probability of all alternatives (which should be $1$ at each decoding step).
This is because at each decoding step, due to the application of constrained decoding, only tokens in $\mathcal A_{\va_{<i}, \mathcal S}$ are considered, leading to a total candidate probability smaller than 1.
Readers are encouraged to verify that $\sum_{\va \in \mathcal S } P_L(\va) / x(\va) = 1$.
With this, the distribution for one sampled sequence using constrained (denote as $\hat P$) is
\begin{align*}
\hat P(\va) = \frac {P_L(\va) / x(\va)} {\sum_{\va \in \mathcal S } P_L(\va) / x(\va)} = \frac {P_L(\va) }{ x(\va)}.
\end{align*}
\paragraph{Probability of acceptance.}
Consider one particular draw $\va_k$ among the $K$ rounds. The acceptance of $\va_k$ is $x(\va_k)$, and in expectation, each draw is accepted with probability
$$
\sum_{\va \in \mathcal S} \hat P(\va) x(\va) = P_L(\mathcal S) = 1 - p_b.
$$
Therefore, at least one sample is accepted during $K$ rounds of sample are $1 - p_b^K$.
% Therefore, we only need to handle the remaining situation where importance sampling is involved, and use the fact that $KL(P\|tP+(1-t)Q)$ can be bound by $(1-t) KL(P\|Q)$.
\paragraph{Unbiased distribution for the accepted sample.} To show that when a sample is accepted during $K$ rounds, we only need to show that for each round, the distribution is unbiased. Conditioned on acceptance, the probability that $\va \in \mathcal S$ is returned is $\hat P(\va) x(\va)/ (1-p_b) = P_L(\va) / (1-p_b) = P_\mathcal S (\va) $.
\paragraph{KL Divergence between $P_\mathcal S$ and $Q$.} With probability $p_b^K$, $K$ rounds of samples are all rejected, and \alg has to resample $K$ samples and conduct importance sampling.
We aim to analyze the KL divergence between the target distribution $P_\mathcal{S}$ and the distribution $Q$ resulting from importance sampling.
Since the buffer size is $K$, the probability $Q(\va)$ is $K$ times the probability that sequence $\va$ is in the first position (which is $\hat P(\va)$) and subsequently wins during the importance sampling step (due to symmetry among the $K$ sequences).
The challenge lies in computing the winning probability $W(\va)$ of sequence $\va$.
Specifically, we have
\begin{align*} W(\va) = \sum_{\va_2, \ldots, \va_K \in \mathcal{S}} \prod_{j=2}^K \hat P(\va_j) \cdot \frac{x(\va)}{\sum_{j=1}^K x(\va_j)},
\end{align*}
where $x(\va)$ is the importance score defined previously:
\begin{align*}
x(\va) = \prod_{i=1}^{|\va|} P_L(\mathcal{A}_{\va{<i}, \mathcal{S}} | \va_{<i}).
\end{align*}
The KL divergence between $P_\mathcal{S}$ and $Q$ can then be expressed as (for simplicity, we set $A = 1 - p_b$)
\begin{align*}
\mathrm{KL}(P_\mathcal{S} \| Q)
% &=\sum_{\va \in \mathcal{S}} \hat P(\va) \ln \frac{\hat P(\va) }{Q(\va)} \\
&= \sum_{\va \in \mathcal{S}} \frac{P_L(\va)}{A} \ln \frac{P_L(\va) / A}{Q(\va)} \\
&= \sum_{\va \in \mathcal{S}} -\frac{P_L(\va)}{A} \ln \frac{A K P_L(\va) W(\va)}{P_L(\va) x(\va)} \\
&= \sum_{\va_1 \in \mathcal{S}} -\frac{P_L(\va_1)}{A} \ln \left( \sum_{\va_2, \ldots, \va_K \in \mathcal{S}} \prod_{j=2}^K \hat P(\va_j) \cdot \frac{A K}{\sum_{j=1}^K x(\va_j)} \right) \\
&= \sum_{\va_1 \in \mathcal{S}} -\frac{\hat P(\va_1) x(\va_1)}{A} \ln \left( \sum_{\va_2, \ldots, \va_K \in \mathcal{S}} \prod_{j=2}^K \hat P(\va_j) \cdot \frac{A K}{\sum_{j=1}^K x(\va_j)} \right) .
\end{align*}
Applying Jensen's inequality (noting that $\ln \mathbb{E}[f] \geq \mathbb{E}[\ln f]$ for a convex function $\ln$), and considering the negative sign, we obtain
\begin{align*}
\mathrm{KL}(P_\mathcal{S} \| Q) &\leq \sum_{\va_1 \in \mathcal{S}} -\frac{\hat P(\va_1) x(\va_1)}{A} \sum_{\va_2, \ldots, \va_K \in \mathcal{S}} \prod_{j=2}^K \hat P(\va_j) \ln \left( \frac{A K}{\sum_{j=1}^K x(\va_j)} \right) \\
&= \sum_{\va_1, \ldots, \va_K \in \mathcal{S}} \prod_{j=1}^K \hat P(\va_j) \left( -\frac{x(\va_1)}{A} \right) \ln \left( \frac{A K}{\sum_{j=1}^K x(\va_j)} \right) \\
&= \sum_{\va_1, \ldots, \va_K \in \mathcal{S}} \prod_{j=1}^K \hat P(\va_j) \cdot \frac{x(\va_1)}{A} \ln \left( \frac{\sum_{j=1}^K x(\va_j)}{A K} \right).
\end{align*}
Due to the symmetry among the indices $\{1,\cdots, K\}$ in the summation (each $\va_j$ is drawn independently and identically), we can replace $x(\va_1)$ with any $x(\va_j)$ without affecting the overall value. Therefore, we can generalize the expression by summing over all positions $j$ and averaging:
\begin{align*}
\mathrm{KL}(P_\mathcal{S} \| Q)
&\leq \frac{1}{K} \sum_{r=1}^K \sum_{\va_1, \ldots, \va_K \in \mathcal{S}} \prod_{j=1}^K \hat P(\va_j) \cdot \frac {x(\va_r)}{A} \ln \left( \frac{\sum_{j=1}^K x(\va_j)}{A K} \right) \\
&=\sum_{\va_1, \ldots, \va_K \in \mathcal{S}} \prod_{j=1}^K \hat P(\va_j) \cdot \frac {\sum _{j=1}^K x(\va_j)}{AK} \ln \left( \frac{\sum_{j=1}^K x(\va_j)}{A K} \right) .
\end{align*}
Imagine $ \tilde X = (\va_1,\cdots, \va_K)$ is a random variable equaling to $ \frac{\sum_{j=1}^K x(\va_j)}{AK}$ with probability $\tilde P(\tilde X) = \prod_{j=1}^K \hat P(\va_j)$. Then, we have
\begin{align*}
KL(P_\mathcal S \| Q) &\leq \sum \tilde P(\tilde X) f(\tilde X) = \mathbb E_{\tilde X} f(\tilde X),
\end{align*}
where $f(x) \triangleq x \ln x \leq |x-1|+ \frac{(x-1)^2}{2} $. This implies that
\begin{align}
KL(P_\mathcal S \| Q) &= \mathbb E_{\tilde X} [|\tilde X -1|] + \mathbb E_{\tilde X}[\frac{(\tilde X - 1)^2}{2}] \\
&\leq \sqrt{\mathbb E_{\tilde X}[(\tilde X -1)^2]} + \frac 1 2 \mathbb E_{\tilde X}[(\tilde X -1)^2], \label{eq:KL_upper}
\end{align}
where the inequality is by Jensen's inequality. The problem is therefore converted to upper bounding the variance of $\tilde X$ (whose expectation is $1$ because the expectation of $x(\va)$ is $A$). Due to the independence of each $\va_j$, we have
\begin{align*}
\text{Var}(\tilde X) = \frac 1 K \text{Var}(\frac{\va}{A}) = \frac 1 K \text{Var}(\frac{\va }{A}) \leq \frac{A(1-A)}{KA^2},
\end{align*}
where the inequality is because all r.v. with a value ranging from $[0,1] $ and expectation $\xi$ has a variance upper bound by $\xi (1-\xi)$. Plugging this into the \Cref{eq:KL_upper}, we have
\begin{align*}
KL(P_\mathcal S \| Q) \leq \sqrt{\frac{p_b}{K(1-p_b)}} + \frac {p_b}{2K(1-p_b)}.
\end{align*}
% Now, notice that for all positive $x, f(x) = x \ln x \leq |x-1|+ \frac{(x-1)^2}{2}.$
\paragraph{KL Divergence between $P_\mathcal S$ and $P_{IS}^K$.} The final step is to bound the divergence of \alg using all the above results. According to the acceptance probability and the unbiasedness of accepted samples, we have
\begin{align*}
P^K_{IS} = p_b^K Q + (1-p_b^K) P_\mathcal S.
\end{align*}
The final bound is obtained by using \Cref{lemma}.
\begin{lemma}\label{lemma}
For any two distribution $P,Q$, any scalar $t \in [0,1]$, we have
\begin{align*}
KL(P\|tP + (1-t)Q) \leq (1-t) KL(P\|Q).
\end{align*}
\end{lemma}
% Feel free to start with a proof scratch in the main paper that you believe is necessary to present.
% }
% \haotian{Generated by GPT based on my prompts.}
% \begin{proof}
% We will prove the two statements in the theorem:
% 1. **Expected Number of Sampling Steps:**
% \[
% \mathbb{E}[k] = \frac{1 - P_b^K}{1 - P_b} \leq \frac{1}{1 - P_b}.
% \]
% 2. **Upper Bound on KL Divergence:**
% \[
% \text{KL}\left(P_{\mathcal{S}} \,\|\, P_{\text{IS}}^{K}\right) \leq \mathcal{O}\left( \frac{P_b^K}{1 - P_b} \left( \frac{P_b}{K} + \sqrt{\frac{2\ln K}{K}} \right) \right),
% \]
% where \( P_b = 1 - P_L(\mathcal{S}) \).
% ---
% **Part 1: Expected Number of Sampling Steps**
% At each sampling attempt, the probability that the algorithm samples a sequence in \( \mathcal{S} \) is \( P_{\mathcal{S}} = 1 - P_b \). Conversely, the probability of sampling a sequence not in \( \mathcal{S} \) is \( P_b \).
% The algorithm continues sampling until either:
% - It accepts a sample from \( \mathcal{S} \) (with probability \( 1 - P_b^K \)), or
% - It reaches the maximum number of sampling steps \( K \).
% The probability that the algorithm stops at step \( k \) (for \( k = 1, 2, \dots, K-1 \)) is:
% \[
% P(k) = P_b^{k-1} \cdot (1 - P_b).
% \]
% At step \( K \), the algorithm stops regardless, so:
% \[
% P(K) = P_b^{K-1}.
% \]
% **Calculating the Expected Value:**
% \begin{align*}
% \mathbb{E}[k] &= \sum_{k=1}^{K} k \cdot P(k) \\
% &= \sum_{k=1}^{K-1} k \cdot P_b^{k-1} (1 - P_b) + K \cdot P_b^{K-1}.
% \end{align*}
% **Computing the Sum:**
% \begin{align*}
% \sum_{k=1}^{K-1} k \cdot P_b^{k-1} &= (1 - P_b) \sum_{k=1}^{K-1} k \cdot P_b^{k-1} \\
% &= (1 - P_b) \cdot \frac{1 - P_b^{K-1}(1 + (K - 1)(1 - P_b))}{(1 - P_b)^2} \\
% &= \frac{1 - P_b^{K-1} - (K - 1) P_b^{K} + (K - 1) P_b^{K-1}}{(1 - P_b)}.
% \end{align*}
% Simplifying and combining terms, we get:
% \begin{align*}
% \mathbb{E}[k] &= \frac{1 - P_b^{K}}{1 - P_b}.
% \end{align*}
% Since \( P_b^{K} \geq 0 \), it follows that:
% \[
% \mathbb{E}[k] = \frac{1 - P_b^{K}}{1 - P_b} \leq \frac{1}{1 - P_b}.
% \]
% ---
% **Part 2: Upper Bound on KL Divergence**
% The algorithm's output distribution \( P_{\text{IS}}^{K} \) can be viewed as a mixture:
% \[
% P_{\text{IS}}^{K} = (1 - P_b^{K}) \cdot P_{\mathcal{S}} + P_b^{K} \cdot Q,
% \]
% where \( Q \) is the distribution when the algorithm uses importance sampling over \( K \) samples (because no sample from \( \mathcal{S} \) was obtained in \( K \) attempts).
% **Decomposing the KL Divergence:**
% \begin{align*}
% \text{KL}(P_{\mathcal{S}} \,\|\, P_{\text{IS}}^{K}) &= \sum_{a \in \mathcal{S}} P_{\mathcal{S}}(a) \ln \frac{P_{\mathcal{S}}(a)}{(1 - P_b^{K}) P_{\mathcal{S}}(a) + P_b^{K} Q(a)} \\
% &= \sum_{a \in \mathcal{S}} P_{\mathcal{S}}(a) \ln \frac{1}{(1 - P_b^{K}) + P_b^{K} \frac{Q(a)}{P_{\mathcal{S}}(a)}}.
% \end{align*}
% **Using a Convexity Argument:**
% Since \( \ln x \) is concave for \( x > 0 \), we can use Jensen's inequality to bound the KL divergence:
% \[
% \text{KL}(P_{\mathcal{S}} \,\|\, P_{\text{IS}}^{K}) \leq P_b^{K} \cdot \text{KL}(P_{\mathcal{S}} \,\|\, Q).
% \]
% **Bounding \( \text{KL}(P_{\mathcal{S}} \,\|\, Q) \):**
% The main challenge is to bound \( \text{KL}(P_{\mathcal{S}} \,\|\, Q) \). To proceed, we consider the variance introduced by the importance sampling estimator.
% **Modeling the Importance Sampling:**
% Let \( \{x_i\} \) be the importance weights corresponding to the sequences in \( \mathcal{S} \), where \( x_i \in [0, 1] \) and \( \sum_{i} p_i x_i = A \) with \( A = 1 - P_b \).
% Define the average importance weight over \( K \) samples:
% \[
% \tilde{X} = \frac{1}{K} \sum_{j=1}^{K} x_{i_j},
% \]
% where \( x_{i_j} \) are independent and identically distributed random variables.
% **Applying Hoeffding's Inequality:**
% Hoeffding's inequality states that for \( \epsilon > 0 \):
% \[
% P\left( |\tilde{X} - A| \geq A \epsilon \right) \leq 2 \exp\left( -2 K A^2 \epsilon^2 \right).
% \]
% **Bounding the KL Divergence:**
% We can express the KL divergence in terms of the expectation of a function \( f(x) = x \ln x \):
% \[
% L = \mathbb{E}\left[ f\left( \frac{\tilde{X}}{A} \right) \right].
% \]
% Using the bound on the deviation of \( \tilde{X} \) from \( A \), we can write:
% \[
% L \leq 2 \exp\left( -2 K A^2 \epsilon^2 \right) f\left( \frac{1}{A} \right) + f(1 + \epsilon).
% \]
% Choosing \( \epsilon = \sqrt{ \frac{\ln K}{2 K A^2} } \), we get:
% \[
% L \leq \frac{2}{K} f\left( \frac{1}{A} \right) + f\left( 1 + \sqrt{ \frac{\ln K}{2 K A^2} } \right).
% \]
% Since \( f \) is continuous and \( A = 1 - P_b \), which is close to 1, the terms involving \( f \) can be bounded by constants.
% **Final Bound on the KL Divergence:**
% Combining the above, we obtain:
% \[
% \text{KL}(P_{\mathcal{S}} \,\|\, P_{\text{IS}}^{K}) \leq P_b^{K} \left( \frac{P_b}{K} + \sqrt{ \frac{2 \ln K}{K} } \right) \cdot C,
% \]
% where \( C \) is a constant depending on \( f \) and \( A \).
% Therefore,
% \[
% \text{KL}(P_{\mathcal{S}} \,\|\, P_{\text{IS}}^{K}) \leq \mathcal{O}\left( \frac{P_b^{K}}{1 - P_b} \left( \frac{P_b}{K} + \sqrt{ \frac{2 \ln K}{K} } \right) \right).
% \]
% ---
% **Conclusion**
% We have shown that:
% 1. The expected number of sampling steps is:
% \[
% \mathbb{E}[k] = \frac{1 - P_b^{K}}{1 - P_b} \leq \frac{1}{1 - P_b}.
% \]
% 2. The KL divergence between \( P_{\mathcal{S}} \) and \( P_{\text{IS}}^{K} \) is bounded by:
% \[
% \text{KL}\left(P_{\mathcal{S}} \,\|\, P_{\text{IS}}^{K}\right) \leq \mathcal{O}\left( \frac{P_b^{K}}{1 - P_b} \left( \frac{P_b}{K} + \sqrt{ \frac{2 \ln K}{K} } \right) \right).
% \]
% \end{proof}
\subsection{Proof of \Cref{thm:lower_bound}}
To rigorously prove the theorem, we only need to construct an LLM output distribution, determine the candidate set $\mathcal S$, and derive the divergence. In fact, assume there are only two tokens, i.e., $\mathcal V = \{v_1, v_2\}$, and all sequences are with length $2$ (so four possible sequences in total, with an ending $\texttt{EOK}$ token after each sequence being draw with probability $1$. We assume that
$$
P_L(v_1v_1) = P_L(v_1v_2) = \frac{1-p_b}2, P_L(v_2v_1 )=p_b\epsilon, P_L(v_2v_2) = p_b(1-\epsilon).
$$
Assume $\mathcal S = \{v_1v_1,v_1v_2,v_2v_1\}$. In this case, the ground-truth distribution $P_{\mathcal S} $ is
$$
P_L(v_1v_1) = P_L(v_1v_2) = \frac 1 2 \cdot \frac{1 - p_b}{1 - p_b + p_b\epsilon} , P_L(v_2v_1 )=\frac{p_b \epsilon}{1 - p_b + p_b\epsilon} .
$$
On the contrary, if we apply constrained decoding, then for the first token to be decoded, both $v_1, v_2$ are valid. Therefore, the probability of choosing $v_1, v_2$ is $1-p_b$ and $p_b$, separately. When $v_1$ is chosen, as $v_1v_1$ and $v_1v_2$ belong to $\mathcal S$, the probability for each of them conditioned on $v_1$ being chosen is $\frac 1 2$. However, when $v_2$ is chosen, since $v_2v_2$ does not satisfy the constraint, the model can only choose $v_2v_1$. This gives the distribution of $P_\mathcal S^{CD}$ as
$$
P_\mathcal S^{CD} (v_1 v_1) = P_\mathcal S^{CD}(v_1v_2) = \frac {1-p_b} 2, P_\mathcal S^{CD}(v_2v_1) = p_b.
$$
Consequently, the KL divergence between the two distributions is (denote $\frac{1 - p_b}{1 - p_b + p_b\epsilon} $ as $\kappa$
\begin{align*}
KL(P_\mathcal S \| P_\mathcal S ^{CD}) &= 2\times \frac 1 2 \kappa \ln \frac{\kappa / 2}{(1-p_b) / 2} + \frac{p_b \epsilon}{1 - p_b + p_b \epsilon}\ln (\frac{p_b \epsilon}{p_b(1 - p_b + p_b \epsilon ) } ) \\
&= \kappa \ln \frac{\kappa}{1 - p_b} + \frac{p_b \epsilon}{1 - p_b + p_b \epsilon}\ln (\frac{ \epsilon}{(1 - p_b + p_b \epsilon ) } ).
\end{align*}
The final result is obtained by considering $\epsilon \to 0$.
\section{\texttt{Torch} code for \ppv}
\lstset{style=mypython}
\begin{lstlisting}[caption={Parallel Prefix-Verification (PPV) algorithm}]
def ppv(P, a, X, M, vocab_size):
"""
Args:
P (torch.Tensor): Token distribution tensor of shape (vocab_size,).
a (torch.Tensor): Partial sequence tensor of shape (len_a,).
X (torch.Tensor): Constraint matrix of shape (N, l_max), sorted lexicographically.
M (int): Number of top tokens to consider.
vocab_size (int): Size of the vocabulary.
Returns:
torch.Tensor: Mask of valid tokens of shape (vocab_size,).
"""
# Get top M token indices
_, vt = torch.topk(P, M) # vt: (M,)
# Build candidate sequences V by appending vt to a
V = torch.cat([a.unsqueeze(0).repeat(M, 1), vt.unsqueeze(1)], dim=1).to(X.device) # V: (M, len_a + 1)
# Initialize binary search bounds
N = X.shape[0]
low = torch.zeros(M, dtype=torch.long, device=X.device)
high = torch.full_like(low, N)
# Perform binary search
while (low < high).any():
mid = torch.div(low + high, 2, rounding_mode="floor")
mid_seqs = X[mid][:, :V.shape[1]] # Shape: (M, len_a + 1)
# Lexicographical comparison
cmp = (V > mid_seqs).float() - (V < mid_seqs).float()
cmp_result = cmp.cumsum(dim=1)[:, -1] # Final comparison result per sequence
# Update bounds based on comparison results
greater = cmp_result > 0
low = torch.where(greater, mid + 1, low)
high = torch.where(~greater, mid, high)
# Check for matches after binary search
valid = (low < N)
V_valid = V[valid]
matching_seqs = X[low[valid]][:, :V.shape[1]]
# Verify exact matches
matches = (V_valid == matching_seqs).all(dim=1)
valid_indices = vt[valid][matches]
# Create mask of valid tokens
mask = torch.zeros(vocab_size, dtype=torch.bool, device=X.device)
mask[valid_indices] = True
return mask
\end{lstlisting}
\newpage
\section{Additional Experimental Settings}
\paragraph{Entity Disambiguation.} We use two BART checkpoints provided by~\citet{Cao2021} to retrieve entities from $\sim$6M Wikipedia pages. Two of the BART checkpoints are used to evaluate on Entity Disambiguation:
The first BART checkpoint, denoted as \texttt{BART-BLINK}, is pre-trained on the BLINK data~\citep{wu2019scalable}, i.e., 9M unique document-entity pairs from Wikipedia. The other BART checkpoint, \texttt{BART-AIDA} is further fine-tuned on AIDA-CoNLL dataset~\citep{hoffart2011robust}. We test Entity Disambiguation with 6 datasets: AIDA-CoNLL, MSNBC, AQUAINT, ACE2004, WNED-CWEB (CWEB) and WNED-WIKI (WIKI)~\citep{guo2018robust}.
\paragraph{Document Retrieval.} We use the BART checkpoint for Document Retrieval provided by~\citet{Cao2021}, which is trained on all the BLINK and KILT data~\citep{petroni2020kilt}. Document Retrieval is tested with KILT benchmark~\citep{petroni2020kilt}, consisting fact checking with FEVER~\citep{thorne2018fever}; open-domain question answering using Natural Questions~\citep{adelani2021masakhaner}, HotpotQA~\citep{yang2018hotpotqa}, TriviaQA~\citep{joshi2017triviaqa}, ELI5~\citep{fan2019eli5}; slot filling with T-REx~\citep{elsahar2018t}, Zero Shot RE~\citep{levy2017zero}; entity disambiguation on AIDA CoNLL-YAGO, WNED-WIKI and WNED-CEWB; dialog with Wizard of Wikipedia~\citep{dinan2018wizard}.
\paragraph{Entity Retrieval \& Fact checking} Without fine-tuning, the language model is provided with a few demonstrations from the target corpus and leverages its in-context learning ability to generate entities accordingly in these two tasks. For each task, three query-doc pairs from the validation set are provided as in-context examples to demonstrate the task intentions and output formats. We use the datasets from zero-shot information retrieval benchmark BEIR~\citep{thakur2021beir}, i.e., DBPedia-Entity~\citep{hasibi2017dbpedia}, Climate-FEVER~\citep{diggelmann2020climate} for Entity Retrieval and Fact Checking respectively. The prompts are given as follows:
\begin{itemize}
\item \begin{verbatim}
Retrieve entity from DBpedia for my query.
Query: <what is foreign policy for angola> Entity: <Foreign relations of Angola>
Query: <animalia> Entity: <Animalia (book)>
Query: <what is asphalt and what is it used for> Entity: <Asphalt>
Query: <{$newQuery}> Entity:
\end{verbatim}
\item \begin{verbatim}
Retrieve entity from Wikipedia for my query.
Query: <Global warming is driving polar bears toward extinction>
Entity: <Habitat destruction>
Query: <The sun has gone into ‘lockdown’ which could cause freezing weather, earthquakes
and famine, say scientists>
Entity: <Famine>
Query: <the Great Barrier Reef is in fine fettle>
Entity: <Great Barrier Reef>
Query: <{$newQuery}>
Entity:
\end{verbatim}
\end{itemize}
\section{Additional Results}
\begin{table}[h]
\centering
\caption{The Micro F1 score and evaluation time for \texttt{DISC} with different $K$ (allowed sampling steps) and $M$ (top candidate tokens). The task is Entity Disambiguation across six datasets with BART AY2 checkpoint. \textbf{Avg.} is the average score across datasets. }
\resizebox{\textwidth}{!}{
\begin{tabular}{ccccccccccccc!{\vrule width \lightrulewidth}cc}
\toprule
\multirow{2}{*}{} & \multicolumn{2}{c}{\textbf{ACE2004}} & \multicolumn{2}{c}{\textbf{AIDA}} & \multicolumn{2}{c}{\textbf{AQUAINT}} & \multicolumn{2}{c}{\textbf{CWEB}} & \multicolumn{2}{c}{\textbf{MSNBC}} & \multicolumn{2}{c!{\vrule width \lightrulewidth}}{\textbf{WIKI}} & \multicolumn{2}{c}{\textbf{Avg.}} \\
& F1 & Time & F1 & Time & F1 & Time & F1 & Time & F1 & Time & F1 & Time & F1 & Time \\
\midrule
\multicolumn{15}{c}{$K=1$} \\
\midrule
$M=1000$ & 0.8405 & 3.377 & 0.8379 & 55.436 & 0.7799 & 9.089 & 0.6484 & 142.65 & 0.7927 & 8.946 & 0.8009 & 89.52 & 0.7834 & 51.50 \\
$M=500$ & 0.8483 & 2.961 & 0.8395 & 52.747 & 0.7951 & 8.636 & 0.6504 & 135.50 & 0.7668 & 8.350 & 0.8003 & 86.47 & 0.7834 & 49.11 \\
$M=100$ & 0.8483 & 2.730 & 0.8459 & 49.429 & 0.7937 & 8.104 & 0.6518 & 127.27 & 0.7774 & 8.553 & 0.7948 & 81.23 & 0.7853 & 46.22 \\
$M=50$ & 0.8327 & 2.736 & 0.8421 & 48.054 & 0.7992 & 7.912 & 0.6517 & 126.86 & 0.7835 & 8.121 & 0.7965 & 80.15 & 0.7843 & 45.64 \\
\midrule
\multicolumn{15}{c}{$K=2$} \\
\midrule
$M=1000$ & 0.8521 & 4.751 & 0.8535 & 86.152 & 0.8157 & 13.944 & 0.6616 & 226.15 & 0.7973 & 13.358 & 0.8038 & 147.77 & 0.7973 & 82.02 \\
$M=500$ & 0.8483 & 4.558 & 0.8531 & 83.668 & 0.8088 & 14.043 & 0.6641 & 212.14 & 0.7851 & 12.798 & 0.8047 & 139.36 & 0.7940 & 77.76 \\
$M=100$ & 0.8405 & 4.641 & 0.8546 & 80.060 & 0.8198 & 12.915 & 0.6634 & 204.21 & 0.7896 & 12.761 & 0.8046 & 133.91 & 0.7954 & 74.75 \\
$M=50$ & 0.8560 & 4.224 & 0.8493 & 74.577 & 0.8061 & 12.575 & 0.6618 & 195.74 & 0.8003 & 11.783 & 0.8043 & 126.26 & 0.7963 & 70.86 \\
\midrule
\multicolumn{15}{c}{$K=3$} \\
\midrule
$M=1000$ & 0.8249 & 7.056 & 0.8544 & 123.357 & 0.8143 & 22.207 & 0.6652 & 337.48 & 0.8049 & 18.632 & 0.8102 & 217.88 & 0.7956 & 121.10 \\
$M=500$ & 0.8327 & 6.166 & 0.8562 & 109.944 & 0.8226 & 18.060 & 0.6623 & 281.26 & 0.7988 & 16.643 & 0.8102 & 187.58 & 0.7971 & 103.27 \\
$M=100$ & 0.8405 & 5.447 & 0.8493 & 98.833 & 0.8171 & 16.148 & 0.6642 & 256.41 & 0.7973 & 15.214 & 0.8097 & 168.73 & 0.7963 & 93.46 \\
$M=50$ & 0.8405 & 5.878 & 0.8542 & 102.032 & 0.8157 & 17.204 & 0.6646 & 263.40 & 0.8018 & 15.205 & 0.8059 & 171.65 & 0.7971 & 95.89 \\
\midrule
\multicolumn{15}{c}{$K=4$} \\
\midrule
$M=1000$ & 0.8405 & 7.806 & 0.8549 & 137.484 & 0.8322 & 22.944 & 0.6668 & 357.17 & 0.8018 & 20.691 & 0.8081 & 241.67 & 0.8007 & 131.29 \\
$M=500$ & 0.8366 & 7.246 & 0.8549 & 134.291 & 0.8184 & 21.457 & 0.6635 & 344.06 & 0.7896 & 19.285 & 0.8094 & 223.37 & 0.7954 & 124.95 \\
$M=100$ & 0.8444 & 7.065 & 0.8555 & 130.161 & 0.8171 & 21.470 & 0.6655 & 331.44 & 0.8018 & 19.462 & 0.8080 & 221.01 & 0.7987 & 121.77 \\
$M=50$ & 0.8327 & 6.680 & 0.8598 & 121.303 & 0.8088 & 20.674 & 0.6667 & 318.97 & 0.8079 & 18.179 & 0.8069 & 211.05 & 0.7971 & 116.14 \\
\bottomrule
\end{tabular}}
\label{tab:M}
\end{table}
\begin{table}[h]
\centering
\caption{The Micro F1 score for \texttt{DISC} with different beam search sizes with different $K$. The task is Entity Disambiguation across six datasets with BART AY2 checkpoint. \textbf{Avg.} is the average score across datasets. We show that \texttt{DISC} can be combined with beam search decoding strategy to improve performance.}
\resizebox{\textwidth}{!}{
\begin{tabular}{ccccccccccccc|cc}
\toprule
\textbf{Dataset} & \multicolumn{2}{c}{\textbf{ACE2004}} & \multicolumn{2}{c}{\textbf{AIDA}} & \multicolumn{2}{c}{\textbf{AQUAINT}} & \multicolumn{2}{c}{\textbf{CWEB}} & \multicolumn{2}{c}{\textbf{MSNBC}} & \multicolumn{2}{c|}{\textbf{WIKI}} & \multicolumn{2}{c}{\textbf{Avg.}} \\
\midrule
Beam size & 1 & 2 & 1 & 2 & 1 & 2 & 1 & 2 & 1 & 2 & 1 & 2 & 1 & 2 \\
\midrule
\texttt{DISC} ($K=1$) & 0.841 & 0.852 & 0.797 & 0.869 & 0.814 & 0.824 & 0.635 & 0.673 & 0.765 & 0.803 & 0.803 & 0.818 & 0.776 & 0.807 \\
\texttt{DISC} ($K=2$) & 0.848 & 0.856 & 0.803 & 0.871 & 0.821 & 0.827 & 0.647 & 0.676 & 0.781 & 0.806 & 0.811 & 0.819 & 0.785 & 0.809 \\
\texttt{DISC} ($K=3$) & 0.848 & 0.860 & 0.808 & 0.872 & 0.831 & 0.825 & 0.651 & 0.677 & 0.782 & 0.797 & 0.810 & 0.820 & 0.788 & 0.809 \\
\texttt{DISC} ($K=4$) & 0.849 & 0.860 & 0.810 & 0.873 & 0.827 & 0.825 & 0.651 & 0.678 & 0.781 & 0.805 & 0.816 & 0.820 & 0.789 & 0.810 \\
\bottomrule
\end{tabular}}
\label{tab:beam}
\end{table}
In this section, we present additional experimental results to show that \texttt{DISC} doesn't require huge $M$ (the number of top candidate tokens) and is compatible with other advanced decoding strategies such as beam search. In~\cref{tab:M}, we evaluate \texttt{DISC} on Entity Disambiguation task with $M$ ranges from $\{50, 100, 500, 1000\}$. From the results in \Cref{tab:M}, we observe that reducing the value of $M$ from $1000$ to $50$ has a minimal impact on the Micro F1 scores across all datasets. For instance, with $K=1$, decreasing $M$ from $1000$ to $50$ only results in a slight fluctuation in F1 scores—often within a margin of $0.01$. This indicates that even a small set of top candidate tokens ($M=50$) is sufficient for \texttt{DISC} to achieve near-optimal performance.
Moreover, a smaller $M$ significantly reduces the evaluation time. For example, with $K=1$ on the CWEB dataset, the evaluation time drops from $142.65$ seconds at $M=1000$ to $126.86$ seconds at $M=50$. This demonstrates that choosing a smaller $M$ not only maintains performance but also enhances computational efficiency.
Increasing the allowed sampling steps $K$ generally improves the F1 scores across datasets, as seen when comparing results from $K=1$ to $K=4$. However, this comes at the cost of increased evaluation time. For example, with $M=1000$, the average time increases from $51.50$ seconds at $K=1$ to $131.29$ seconds at $K=4$. Therefore, there’s a trade-off between performance gains and computational overhead when adjusting $K$.
In \Cref{tab:beam}, we evaluate the compatibility of \texttt{DISC} with beam search decoding strategies. Incorporating beam search with a beam size of $2$ consistently improves the Micro F1 scores across all datasets and values of $K$. For instance, at $K=1$, the average F1 score increases from $0.776$ with greedy decoding to $0.807$ with beam search. This enhancement suggests that beam search helps explore more candidate sequences, leading to better disambiguation outcomes.
Combining beam search with higher values of $K$ continues to yield marginal improvements. However, the gains taper off as $K$ increases, indicating that most of the performance benefits are captured with smaller $K$ and beam sizes. This synergy between \texttt{DISC} and beam search showcases the method’s flexibility and its ability to integrate with advanced decoding strategies to further boost performance.
In summary, these additional experiments demonstrate that \texttt{DISC} is both efficient and adaptable. It does not require a large number of top candidate tokens to perform effectively, and it can seamlessly incorporate beam search to enhance its disambiguation capabilities. This makes \texttt{DISC} a practical choice for entity disambiguation tasks where computational resources and time are critical considerations.
\begin{thebibliography}{}
\bibitem[Achiam et~al., 2023]{achiam2023gpt}
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F.~L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et~al. (2023).
\newblock Gpt-4 technical report.
\newblock {\em arXiv preprint arXiv:2303.08774}.
\bibitem[Adelani et~al., 2021]{adelani2021masakhaner}
Adelani, D.~I., Abbott, J., Neubig, G., D’souza, D., Kreutzer, J., Lignos, C., Palen-Michel, C., Buzaaba, H., Rijhwani, S., Ruder, S., et~al. (2021).
\newblock Masakhaner: Named entity recognition for african languages.
\newblock {\em Transactions of the Association for Computational Linguistics}, 9:1116--1131.
\bibitem[Anderson et~al., 2017]{Anderson2017}
Anderson, P., Fernando, B., Johnson, M., and Gould, S. (2017).
\newblock Guided open vocabulary image captioning with constrained beam search.
\newblock In {\em Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing}.
\bibitem[Bayer and Unterauer, 1977]{bayer1977prefix}
Bayer, R. and Unterauer, K. (1977).
\newblock Prefix b-trees.
\newblock {\em ACM Transactions on Database Systems (TODS)}, 2(1):11--26.
\bibitem[Beurer-Kellner et~al., 2024]{Kellner2024}
Beurer-Kellner, L., Fischer, M., and Vechev, M. (2024).
\newblock Guiding {LLM}s the right way: Fast, non-invasive constrained generation.
\newblock In {\em Proceedings of the 41st International Conference on Machine Learning}.
\bibitem[Bevilacqua et~al., 2022]{bevilacqua2022autoregressive}
Bevilacqua, M., Ottaviano, G., Lewis, P., Yih, S., Riedel, S., and Petroni, F. (2022).
\newblock Autoregressive search engines: Generating substrings as document identifiers.
\newblock {\em Advances in Neural Information Processing Systems}, 35:31668--31683.
\bibitem[Black et~al., 2022]{black2022gpt}
Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L., Golding, L., He, H., Leahy, C., McDonell, K., Phang, J., et~al. (2022).
\newblock Gpt-neox-20b: An open-source autoregressive language model.
\newblock In {\em Proceedings of BigScience Episode\# 5--Workshop on Challenges \& Perspectives in Creating Large Language Models}, pages 95--136.
\bibitem[Brown et~al., 2020]{brown2020language}
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J.~D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et~al. (2020).
\newblock Language models are few-shot learners.
\newblock {\em Advances in neural information processing systems}, 33:1877--1901.
\bibitem[De~Cao1 et~al., 2021]{Cao2021}
De~Cao1, N., Izacard, G., Riedel, S., and Petroni, F. (2021).
\newblock Autoregressive entity retrieval.
\newblock In {\em Proceedings of the Ninth International Conference on Learning Representations}.
\bibitem[Diggelmann et~al., 2020]{diggelmann2020climate}
Diggelmann, T., Boyd-Graber, J., Bulian, J., Ciaramita, M., and Leippold, M. (2020).
\newblock Climate-fever: A dataset for verification of real-world climate claims.
\newblock {\em arXiv preprint arXiv:2012.00614}.
\bibitem[Dinan et~al., 2018]{dinan2018wizard}
Dinan, E., Roller, S., Shuster, K., Fan, A., Auli, M., and Weston, J. (2018).
\newblock Wizard of wikipedia: Knowledge-powered conversational agents.
\newblock {\em arXiv preprint arXiv:1811.01241}.
\bibitem[Dubey et~al., 2024]{dubey2024llama}
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et~al. (2024).
\newblock The llama 3 herd of models.
\newblock {\em arXiv preprint arXiv:2407.21783}.
\bibitem[Elsahar et~al., 2018]{elsahar2018t}
Elsahar, H., Vougiouklis, P., Remaci, A., Gravier, C., Hare, J., Laforest, F., and Simperl, E. (2018).
\newblock T-rex: A large scale alignment of natural language with knowledge base triples.
\newblock In {\em Proceedings of the Eleventh International Conference on Language Resources and Evaluation (LREC 2018)}.
\bibitem[Fan et~al., 2019]{fan2019eli5}
Fan, A., Jernite, Y., Perez, E., Grangier, D., Weston, J., and Auli, M. (2019).
\newblock Eli5: Long form question answering.
\newblock {\em arXiv preprint arXiv:1907.09190}.
\bibitem[Fischer-Hbner and Berthold, 2017]{fischer2017privacy}
Fischer-Hbner, S. and Berthold, S. (2017).
\newblock Privacy-enhancing technologies.
\newblock In {\em Computer and information security Handbook}, pages 759--778. Elsevier.
\bibitem[Geng et~al., 2023]{Geng2023}
Geng, S., Josifoski, M., Peyrard, M., and West, R. (2023).
\newblock Grammar-constrained decoding for structured nlp tasks without finetuning.
\newblock In {\em Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing}.
\bibitem[Gilbert et~al., 2023]{gilbert2023large}
Gilbert, S., Harvey, H., Melvin, T., Vollebregt, E., and Wicks, P. (2023).
\newblock Large language model ai chatbots require approval as medical devices.
\newblock {\em Nature Medicine}, pages 1--3.
\bibitem[Guo and Barbosa, 2018]{guo2018robust}
Guo, Z. and Barbosa, D. (2018).
\newblock Robust named entity disambiguation with random walks.
\newblock {\em Semantic Web}, 9(4):459--479.
\bibitem[Hasibi et~al., 2017]{hasibi2017dbpedia}
Hasibi, F., Nikolaev, F., Xiong, C., Balog, K., Bratsberg, S.~E., Kotov, A., and Callan, J. (2017).
\newblock Dbpedia-entity v2: a test collection for entity search.
\newblock In {\em Proceedings of the 40th International ACM SIGIR Conference on Research and Development in Information Retrieval}, pages 1265--1268.
\bibitem[Hoffart et~al., 2011]{hoffart2011robust}
Hoffart, J., Yosef, M.~A., Bordino, I., F{\"u}rstenau, H., Pinkal, M., Spaniol, M., Taneva, B., Thater, S., and Weikum, G. (2011).
\newblock Robust disambiguation of named entities in text.
\newblock In {\em Proceedings of the 2011 conference on empirical methods in natural language processing}, pages 782--792.
\bibitem[Hokamp and Liu, 2017]{Hokamp2017}
Hokamp, C. and Liu, Q. (2017).
\newblock Lexically constrained decoding for sequence generation using grid beam search.
\newblock In {\em Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics}.
\bibitem[Hwang and Chang, 2023]{hwang2023review}
Hwang, G.-J. and Chang, C.-Y. (2023).
\newblock A review of opportunities and challenges of chatbots in education.
\newblock {\em Interactive Learning Environments}, 31(7):4099--4112.
\bibitem[Joshi et~al., 2017]{joshi2017triviaqa}
Joshi, M., Choi, E., Weld, D.~S., and Zettlemoyer, L. (2017).
\newblock Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension.
\newblock {\em arXiv preprint arXiv:1705.03551}.
\bibitem[Jouppi et~al., 2017]{jouppi2017datacenter}
Jouppi, N.~P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., et~al. (2017).
\newblock In-datacenter performance analysis of a tensor processing unit.
\newblock In {\em Proceedings of the 44th annual international symposium on computer architecture}, pages 1--12.
\bibitem[Le and Titov, 2018]{le2018improving}
Le, P. and Titov, I. (2018).
\newblock Improving entity linking by modeling latent relations between mentions.
\newblock {\em arXiv preprint arXiv:1804.10637}.
\bibitem[Lee et~al., 2024]{lee2024can}
Lee, J., Chen, A., Dai, Z., Dua, D., Sachan, D.~S., Boratko, M., Luan, Y., Arnold, S.~M., Perot, V., Dalmia, S., et~al. (2024).
\newblock Can long-context language models subsume retrieval, rag, sql, and more?
\newblock {\em arXiv preprint arXiv:2406.13121}.
\bibitem[Lee et~al., 2023]{lee2023benefits}
Lee, P., Bubeck, S., and Petro, J. (2023).
\newblock Benefits, limits, and risks of gpt-4 as an ai chatbot for medicine.
\newblock {\em New England Journal of Medicine}, 388(13):1233--1239.
\bibitem[Levy et~al., 2017]{levy2017zero}
Levy, O., Seo, M., Choi, E., and Zettlemoyer, L. (2017).
\newblock Zero-shot relation extraction via reading comprehension.
\newblock {\em arXiv preprint arXiv:1706.04115}.
\bibitem[Liu et~al., 2023]{liu2023visual}
Liu, H., Li, C., Wu, Q., and Lee, Y.~J. (2023).
\newblock Visual instruction tuning.
\newblock {\em arXiv preprint arXiv:2304.08485}.
\bibitem[Miao et~al., 2019]{Miao2019}
Miao, Y., Zhou, H., and Mou, L. (2019).
\newblock Cgmh: Constrained sentence generation by metropolis-hastings sampling.
\newblock In {\em Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing}.
\bibitem[Owens et~al., 2008]{owens2008gpu}
Owens, J.~D., Houston, M., Luebke, D., Green, S., Stone, J.~E., and Phillips, J.~C. (2008).
\newblock Gpu computing.
\newblock {\em Proceedings of the IEEE}, 96(5):879--899.
\bibitem[Petroni et~al., 2020]{petroni2020kilt}
Petroni, F., Piktus, A., Fan, A., Lewis, P., Yazdani, M., De~Cao, N., Thorne, J., Jernite, Y., Karpukhin, V., Maillard, J., et~al. (2020).
\newblock Kilt: a benchmark for knowledge intensive language tasks.
\newblock {\em arXiv preprint arXiv:2009.02252}.
\bibitem[Pezoa et~al., 2016]{pezoa2016foundations}
Pezoa, F., Reutter, J.~L., Suarez, F., Ugarte, M., and Vrgo{\v{c}}, D. (2016).
\newblock Foundations of json schema.
\newblock In {\em Proceedings of the 25th international conference on World Wide Web}, pages 263--273.
\bibitem[Post and Vilar, 2018]{Post2018}
Post, M. and Vilar, D. (2018).
\newblock Fast lexically constrained decoding with dynamic beam allocation for neural machine translation.
\newblock In {\em Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics}.
\bibitem[Rajput, 2024]{rajput24}
Rajput, Shashank, e.~a. (2024).
\newblock Recommender systems with generative retrieval.
\newblock In {\em Proceedings of the Thirty-Seventh Annual Conference on Advances in Neural Information Processing Systems}.
\bibitem[Rajput et~al., 2024]{rajput2024recommender}
Rajput, S., Mehta, N., Singh, A., Hulikal~Keshavan, R., Vu, T., Heldt, L., Hong, L., Tay, Y., Tran, V., Samost, J., et~al. (2024).
\newblock Recommender systems with generative retrieval.
\newblock {\em Advances in Neural Information Processing Systems}, 36.
\bibitem[Skjuve et~al., 2021]{skjuve2021my}
Skjuve, M., F{\o}lstad, A., Fostervold, K.~I., and Brandtzaeg, P.~B. (2021).
\newblock My chatbot companion-a study of human-chatbot relationships.
\newblock {\em International Journal of Human-Computer Studies}, 149:102601.
\bibitem[Stahlberg and Kumar, 2019]{Stahlberg2019}
Stahlberg, F. and Kumar, S. (2019).
\newblock Neural grammatical error correction with finite state transducers.
\newblock In {\em Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics}.
\bibitem[Sun et~al., 2024]{sun2024learning}
Sun, W., Yan, L., Chen, Z., Wang, S., Zhu, H., Ren, P., Chen, Z., Yin, D., Rijke, M., and Ren, Z. (2024).
\newblock Learning to tokenize for generative retrieval.
\newblock {\em Advances in Neural Information Processing Systems}, 36.
\bibitem[Thakur et~al., 2021]{thakur2021beir}
Thakur, N., Reimers, N., R{\"u}ckl{\'e}, A., Srivastava, A., and Gurevych, I. (2021).
\newblock Beir: A heterogenous benchmark for zero-shot evaluation of information retrieval models.
\newblock {\em arXiv preprint arXiv:2104.08663}.
\bibitem[Thorne et~al., 2018]{thorne2018fever}
Thorne, J., Vlachos, A., Christodoulopoulos, C., and Mittal, A. (2018).
\newblock Fever: a large-scale dataset for fact extraction and verification.
\newblock {\em arXiv preprint arXiv:1803.05355}.
\bibitem[Trinh et~al., 2024]{trinh2024solving}
Trinh, T.~H., Wu, Y., Le, Q.~V., He, H., and Luong, T. (2024).
\newblock Solving olympiad geometry without human demonstrations.
\newblock {\em Nature}, 625(7995):476--482.
\bibitem[Van~Gysel and de~Rijke, 2018]{van2018pytrec_eval}
Van~Gysel, C. and de~Rijke, M. (2018).
\newblock Pytrec\_eval: An extremely fast python interface to trec\_eval.
\newblock In {\em The 41st International ACM SIGIR Conference on Research \& Development in Information Retrieval}, pages 873--876.
\bibitem[Wu et~al., 2019]{wu2019scalable}
Wu, L., Petroni, F., Josifoski, M., Riedel, S., and Zettlemoyer, L. (2019).
\newblock Scalable zero-shot entity linking with dense entity retrieval.
\newblock {\em arXiv preprint arXiv:1911.03814}.
\bibitem[Yang et~al., 2018]{yang2018hotpotqa}
Yang, Z., Qi, P., Zhang, S., Bengio, Y., Cohen, W.~W., Salakhutdinov, R., and Manning, C.~D. (2018).
\newblock Hotpotqa: A dataset for diverse, explainable multi-hop question answering.
\newblock {\em arXiv preprint arXiv:1809.09600}.
\bibitem[Ye et~al., 2023]{ye2023mplugowl}
Ye, Q., Xu, H., Xu, G., Ye, J., Yan, M., Zhou, Y., Wang, J., Hu, A., Shi, P., Shi, Y., Li, C., Xu, Y., Chen, H., Tian, J., Qi, Q., Zhang, J., and Huang, F. (2023).
\newblock mplug-owl: Modularization empowers large language models with multimodality.
\newblock {\em arXiv preprint arXiv:2304.14178}.
\bibitem[Zhang et~al., 2023]{zhang2023safetybench}
Zhang, Z., Lei, L., Wu, L., Sun, R., Huang, Y., Long, C., Liu, X., Lei, X., Tang, J., and Huang, M. (2023).
\newblock Safetybench: Evaluating the safety of large language models with multiple choice questions.
\newblock {\em arXiv preprint arXiv:2309.07045}.
\bibitem[Zhu et~al., 2023]{zhu2023minigpt4}
Zhu, D., Chen, J., Shen, X., Li, X., and Elhoseiny, M. (2023).
\newblock Minigpt-4: Enhancing vision-language understanding with advanced large language models.
\newblock {\em arXiv preprint arXiv:2304.10592}.
\end{thebibliography}