# A Probabilistic Semi-Supervised Approach with Triplet Markov Chains

###### Abstract

Triplet Markov chains are general generative models for sequential data which
take into account three kinds of random variables: (noisy) observations, their
associated discrete labels and latent variables which aim at strengthening the
distribution of the observations and their associated labels. However, in
practice, we do not have at our disposal all the labels associated to the
observations to estimate the parameters of such models. In this paper, we
propose a general framework based on a variational Bayesian inference to train
parameterized triplet Markov chain models in a semi-supervised context. The
generality of our approach enables us to derive semi-supervised algorithms for
a variety of generative models for sequential Bayesian classification.

Index Terms—  Generative Models; Variational Inference; Semi-Supervised
Learning; Triplet Markov Chains.

##  1 Introduction

This paper focuses on semi-supervised learning for sequential Bayesian
classification in general generative models. Let us start by recalling the
principle of Bayesian classification and semi-supervised estimation.

###  1.1 Sequential Bayesian classification

We denote as ${\mathbf{x}}_{T}=(x_{0},\dots,x_{T})$ a sequence of observed
random variables (r.v.) and ${\mathbf{z}}_{T}=(z_{0},\dots,z_{T})$ a sequence
of latent r.v. We also introduce a sequence of labels
${\mathbf{y}}_{T}=(y_{0},\dots,\;y_{T})$ associated to the previous sequence
${\mathbf{x}}_{T}$. We will assume that $x_{t}\in\mathbb{R}^{d_{x}}$,
$z_{t}\in\mathbb{R}^{d_{z}}$ while the label $y_{t}$ is discrete, so
$y_{t}\in\Omega=\\{\omega_{1},\dots,\omega_{C}\\}$. As far as notations are
concerned, we do not distinguish r.v. and their realizations. For example,
${\mathbf{x}}_{T}$ can represent a noisy grayscale image while
${\mathbf{y}}_{T}$ represents the original black and white image. For this
application, the latent variable $z_{t}$ can be used to govern the conditional
distribution of the noise given the original label.

When the labels associated to ${\mathbf{x}}_{T}$ are not observed, the
objective associated to Bayesian classification consists in computing, for all
$t$, the posterior distributions

$p(y_{t}|{\mathbf{x}}_{T})=\frac{\sum_{{\mathbf{y}}_{0:t-1},{\mathbf{y}}_{t+1:T}}\int
p({\mathbf{y}}_{T},{\mathbf{x}}_{T},{\mathbf{z}}_{T}){\rm
d}{\mathbf{z}}_{T}}{\sum_{{\mathbf{y}}_{T}}\int
p({\mathbf{y}}_{T},{\mathbf{x}}_{T},{\mathbf{z}}_{T}){\rm
d}{\mathbf{z}}_{T}}\text{.}$ (1)

Consequently, we first need to define a parameterized model
$p_{\theta}({\mathbf{x}}_{T},{\mathbf{y}}_{T},{\mathbf{z}}_{T})$ which aims at
describing the r.v. involved in the problem and from which it is possible to
estimate $\theta$ and next to compute (1) in a reasonable computational cost.

The estimation of $\theta$ (i.e. the learning step) can be realized from
sequences where we have at our disposal $({\mathbf{x}}_{T},{\mathbf{y}}_{T})$
(supervised learning) or only ${\mathbf{x}}_{T}$ (unsupervised learning). This
general problem is commonly used in many fields, such as speech recognition
[1], natural language processing [2], and activity recognition [3].

###  1.2 Semi-Supervised learning

The problem we consider in this paper is a little bit different. In many real-
world applications, it is expensive or impossible to obtain labels for the
entire sequence due to various reasons such as the high cost of labeling, the
lack of expertise, or the lack of time. So from now on we assume that we have
at our disposal a sequence of observations ${\mathbf{x}}_{T}$ with partially
observed labels and that i) we want to train relevant generative models; and
ii) we look for estimating the missing labels associated to each sequence. In
other words, decomposing a sequence of labels ${\mathbf{y}}_{T}$ as

${\mathbf{y}}_{T}=({\mathbf{y}}_{T}^{\mathcal{L}},{\mathbf{y}}_{T}^{\mathcal{U}})\text{,}$

where ${\mathbf{y}}_{T}^{\mathcal{L}}=\\{y_{t}\\}_{t\in{\mathcal{L}}}$ (resp.
${\mathbf{y}}_{T}^{\mathcal{U}}=\\{y_{t}\\}_{t\in\mathcal{U}}$) denotes the
observed (resp. the unobserved) labels (so ${\mathcal{L}}$ (resp.
$\mathcal{U}$) denotes the time index of observed (resp. unobserved) labels)
we now look for estimating $\theta$ from
$({\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})$, and next computing, for
all $t\in\mathcal{U}$,

$p(y_{t}|{\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})=\frac{\sum_{y_{s}\text{,
}s\in\mathcal{U}\backslash\\{t\\}}\int
p({\mathbf{y}}_{T},{\mathbf{x}}_{T},{\mathbf{z}}_{T}){\rm
d}{\mathbf{z}}_{T}}{\sum_{y_{s}\text{, }s\in\mathcal{U}}\int
p({\mathbf{y}}_{T},{\mathbf{x}}_{T},{\mathbf{z}}_{T}){\rm
d}{\mathbf{z}}_{T}}\text{.}$

###  1.3 Scope of the paper

In this paper, we show that it is possible to propose a very general framework
for semi-supervised learning in sequential generative models. In particular,
we show that it is sufficient to consider a generative model in which the
triplet process $\\{z_{t},x_{t},y_{t}\\}_{t\geq 0}$ is Markovian. As we will
see later, such a general generative model encompasses well known generative
models such that the Variational Sequential Labeler (VSL) [4] or the Semi-
supervised Variational Recurrent Neural Network (SVRNN) [5]. Due to this
general interpretation, we are able to propose new generative models that
outperform the previous ones for some semi-supervised learning tasks. The
estimation of these general models is based on an adaptation of the
variational Bayesian framework [6] for sequential and partially observed data.

The paper is organised as follows. In section 2, we recall the general triplet
Markov chain (TMC) model [7, 8] and the principle of variational inference. In
section 3, we show that our previous general models can be estimated with
modified variational inference techniques in the context of semi-supervised
learning. We next show that this general inference technique encompasses
popular semi-supervised learning algorithms and we propose our Deep TMC model.
Finally, in section 4, we compare different approaches for the image
segmentation problem.

##  2 Background

This section introduces the TMC model and the principle of variational
Bayesian inference techniques.

###  2.1 Triplet Markov Chains

Let us first start with the pair process $\\{x_{t},y_{t}\\}_{t\geq 0}$. A
popular model for describing this process is the hidden Markov chain model in
which the sequences of label is assumed to be Markovian and the (noisy)
observations are independent given the labels. In addition, the observation at
a current time only depends on the label at the same time. In other words,

$p({\mathbf{x}}_{T},{\mathbf{y}}_{T})\overset{\rm
HMC}{=}p(y_{0})\prod_{t=1}^{T}p(y_{t}|y_{t-1})\prod_{t=0}^{T}p(x_{t}|y_{t})\text{.}$

However, this model may be poor in practice due to the Markovian assumption on
the label. A simple way to relax it is to consider the pairwise Markov chain
(PMC) model in which the pair $\\{x_{t},y_{t}\\}_{t\geq 0}$ is assumed to be
Markovian,

$p({\mathbf{x}}_{T},{\mathbf{y}}_{T})\overset{\rm
PMC}{=}p(y_{0},x_{0})\prod_{t=1}^{T}p(y_{t},x_{t}|y_{t-1},x_{t-1})\text{.}$

Even if this model is a direct generalization of the previous HMC model, it
may be also unsatisfying in practice. The reason why is that it relies on the
particular choice of the transition distribution
$p(y_{t},x_{t}|y_{t-1},x_{t-1})$ which may be difficult in practice. A simple
way to address this problem is to introduce a new latent process
$\\{z_{t}\\}_{t\geq 0}$ which aims at making the previous distribution more
robust. By denoting $v_{t}=(z_{t},x_{t},y_{t})\text{,}$ the general TMC model
satisfies

$p({\mathbf{z}}_{T},{\mathbf{x}}_{T},{\mathbf{y}}_{T})=p({\mathbf{v}}_{T})\overset{\rm
TMC}{=}p(v_{0})\prod_{t=1}^{T}p(v_{t}|v_{t-1})$ (2)

and so, relies on the transition distribution

$p(v_{t}|v_{t-1})=p(z_{t},x_{t},y_{t}|z_{t-1},x_{t-1},y_{t-1})\text{.}$

It encompasses the previous PMC and HMC models and also leads to a general
class of generative models. For example, it is possible to first describe the
distribution of the latent variable, then the one of the label given by the
latent variable, and finally the one of the noisy observation given the latent
variable and the label. In this case, the transition distribution reads as

$p(v_{t}|v_{t-1})=p(z_{t}|v_{t-1})p(y_{t}|z_{t},v_{t-1})p(x_{t}|y_{t},z_{t},v_{t-1})\text{.}$

Even if the nature of the previous distributions is standard, the distribution
of interest

$p({\mathbf{x}}_{T},{\mathbf{y}}_{T})=\int
p({\mathbf{z}}_{T},{\mathbf{x}}_{T},{\mathbf{y}}_{T}){\rm d}{\mathbf{z}}_{T}$

can be complex. In order to be used in practice, we have to first choose a
parameterized transition distribution $p_{\theta}(v_{t}|v_{t-1})$ and then
estimate the parameter $\theta$ from partially labelled observations. This
estimation step relies on the variational inference framework that we now
recall.

###  2.2 Variational Bayesian Inference

In this section, we give up the temporal aspect of the problem and we only
consider an observation $x$ and a latent r.v. $z$; the distribution
$p_{\theta}(x,z)=p(z)p(x|z)$ is assumed to be known and we want to estimate
$\theta$ from a realization $x$. A popular estimator is the Maximum-Likelihood
(ML) estimate $\hat{\theta}=\operatorname*{arg\,max}_{\theta}p_{\theta}(x)$
due to its statistical properties [9, 10]. However, a direct maximization of
$p_{\theta}(x)$ is not always possible, particularly in models with latent
variables where the likelihood $p_{\theta}(x)=\int p_{\theta}(x,z){\rm d}z$ is
not computable. In the variational inference framework, a variational lower
bound called evidence lower bound (ELBO) on the log-likelihood is optimized in
order to estimate the parameters $\theta$ [6]. This variational lower bound
relies on the introduction of a parameterized variational distribution
$q_{\phi}(z|x)$ which aims at mimicking the true posterior $p_{\theta}(z|x)$
and which is parameterized by a set of parameters $\phi$; the ELBO reads

$\displaystyle\tilde{Q}(\theta,\phi)=-\int\log\left(\frac{q_{\phi}(z|x)}{p_{\theta}(x,z)}\right)q_{\phi}(z|x){\rm
d}z$ (3)

and satisfies, for all $(\theta,\phi)$

$\log(p_{\theta}(x))\geq\tilde{Q}(\theta,\phi)\text{.}$

Equality holds if $q_{\phi}(z|x)=p_{\theta}(z|x)$. In this particular case,
the alternating maximization w.r.t. $\theta$ and $q_{\phi}$ of the ELBO,
$\tilde{Q}(\theta,\phi)$, coincides with the EM algorithm [11].

Variational inference consists in maximizing $Q(\theta,\phi)$ with respect to
$(\theta,\phi)$ for a given class of distributions $q_{\phi}$. The choice of
the variational distribution $q_{\phi}(z|x)$ is critical; $q_{\phi}(z|x)$
should be close to $p_{\theta}(z|x)$ but should also be chosen in a such way
that the associated ELBO can be exactly computed or easily approximated while
remaining differentiable w.r.t. $(\theta,\phi)$. A simple way to approximate
$\tilde{Q}(\theta,\phi)$ with a Monte Carlo method is to use the
reparametrization trick [12] which consists in choosing a parametric
distribution $q_{\phi}(z|x)$ such that a sample $z^{(i)}\sim q(z|x)$ can be
written as a differentiable function of $\phi$.

##  3 Semi-supervised Variational Inference for TMCs

Let us now turn back to the TMC model (2) described by a parameterized
transition distribution $p_{\theta}(v_{t}|v_{t-1})$.

###  3.1 Variational inference

In order to estimate $\theta$ from partially labelled observations
$({\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})$, we now want to maximize
the likelihood

$p_{\theta}({\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})=\sum_{y_{s}\text{,
}s\in\mathcal{U}}\int
p_{\theta}({\mathbf{z}}_{T},{\mathbf{x}}_{T},{\mathbf{y}}_{T}){\rm
d}{\mathbf{z}}_{T}$

which is not computable in the general case. We thus adapt the variational
Bayesian framework of section 2.2, where
$x\leftarrow({\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})$ and
$z\leftarrow({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}})$.

The ELBO (3) now reads

$\displaystyle Q(\theta,\phi)$ $\displaystyle=-\sum_{y_{s}\text{,
}s\in\mathcal{U}}\int
q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|{\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})\times$
$\displaystyle\log\left(\frac{q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|{\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})}{p_{\theta}({\mathbf{z}}_{T},{\mathbf{x}}_{T},{\mathbf{y}}_{T})}\right){\rm
d}{\mathbf{z}}_{T}\text{.}$ (4)

Let us now discuss on the computation of (3.1). First, it is worthwhile to
remark that it does not depend on the choice of the generative model: any
parameterized TMC model can be used since
$p_{\theta}({\mathbf{z}}_{T},{\mathbf{x}}_{T},{\mathbf{y}}_{T})$ is known. So
its computation only depends on the choice of the variational distribution
$q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|{\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})$.
It can be factorized in two ways. For sake of clarity, we will omit the
initial distribution of the variables at time $t=0$. The first factorization
coincides with

$\displaystyle
q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|{\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})=$
$\displaystyle\prod_{t=1}^{T}q_{\phi}(z_{t}|{\mathbf{z}}_{t-1},{\mathbf{y}}_{t-1},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})\times$
$\displaystyle\prod_{t\in\mathcal{U}}^{T}q_{\phi}(y_{t}|{\mathbf{y}}_{t-1},{\mathbf{z}}_{t},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})\text{,}$
(5)

(remember that
${\mathbf{y}}_{t-1}=({\mathbf{y}}_{t-1}^{\mathcal{U}},{\mathbf{y}}_{t-1}^{{\mathcal{L}}})$)
while the second one coincides with

$\displaystyle
q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|{\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})=$
$\displaystyle\prod_{t=1}^{T}q_{\phi}(z_{t}|{\mathbf{z}}_{t-1},{\mathbf{y}}_{t},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})\times$
$\displaystyle\\!\\!\prod_{t\in\mathcal{U}}^{T}q_{\phi}(y_{t}|{\mathbf{y}}_{t-1},{\mathbf{z}}_{t-1},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})\text{.}$
(6)

As we explained in the previous section, the choice of the variational
distribution depends on the generative model. In particular here, it has to
take into account the factorization of the parameterized transition
distribution $p_{\theta}(v_{t}|v_{t-1})$.

It remains to compute the ELBO (3.1) which is nothing more than expectation
according to
$q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|{\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})$.
We thus propose to use a Monte-Carlo approximation based on the
reparametrization trick in order to obtain a differentiable approximation
$\hat{Q}(\theta,\phi)$ of $Q(\theta,\phi)$. More precisely, we use the
classical reparameterization trick to sample sequentially according to the
continuous distribution
$q(z_{t}|{\mathbf{z}}_{t-1},{\mathbf{y}}_{t},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})$
(or
$q(z_{t}|{\mathbf{z}}_{t-1},{\mathbf{y}}_{t},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})$,
while we use the Gumbel-Softmax (G-S) trick [13, 14] to sample according to
$q(y_{t}|$
${\mathbf{y}}_{t-1},{\mathbf{z}}_{t},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})$
(or $q(y_{t}|$
${\mathbf{y}}_{t-1},{\mathbf{z}}_{t-1},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})$)
since the labels are discrete.

###  3.2 Particular semi-supervised algorithms for TMCs

As we have seen, semi-supervised algorithms for TMC depend on two key
ingredients: the generative model described by the transition distribution
$p_{\theta}(v_{t}|v_{t-1})$ which has an impact on the performance of the
model for a specific task (classification, prediction, detection, or
generation), and the choice of the variational distribution
$q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|{\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})$.
Actually, our general framework encompasses two existing solutions and we
propose a new one.

####  3.2.1 Variational Sequential Labeler (VSL)

The VSL [4] is a semi-supervised learning model for sequential data and it has
originally been proposed for the sequence labeling tasks in natural language
processing. It can be seen as a particular case of the general framework we
derived in section 3. This particular setting coincides with

$\displaystyle p_{\theta}(v_{t}|v_{t-1})\overset{\rm
VSL}{=}p_{\theta}(y_{t}|z_{t})p_{\theta}(z_{t}|x_{t-1},z_{t-1})p_{\theta}(x_{t}|z_{t})\text{,}$
(7)

while the associated variational distribution satisfies factorization (3.1)
with

$\displaystyle
q_{\phi}(z_{t}|{\mathbf{z}}_{t-1},{\mathbf{y}}_{t-1},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})=q_{\phi}(z_{t}|{\mathbf{x}}_{T})\text{,}$
(8) $\displaystyle
q_{\phi}(y_{t}|{\mathbf{y}}_{t-1},{\mathbf{z}}_{t},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})=p_{\theta}(y_{t}|z_{t})\text{.}$
(9)

In this case, the ELBO (3.1) reduces to

$\displaystyle Q(\theta,\phi)\overset{\rm VSL}{=}$
$\displaystyle\sum_{t\in{\mathcal{L}}}\int
q_{\phi}(z_{t}|{\mathbf{x}}_{T})\log p_{\theta}(y_{t}|z_{t}){\rm d}z_{t}+$
$\displaystyle\sum_{t=0}^{T}\int q_{\phi}(z_{t}|{\mathbf{x}}_{T})\Bigg{[}\log
p_{\theta}(x_{t}|z_{t})-$
$\displaystyle\quad\quad\log\left(\frac{q_{\phi}(z_{t}|{\mathbf{x}}_{T})}{p(z_{t}|x_{t-1},z_{t-1})}\right)\Bigg{]}{\rm
d}z_{t}\text{.}$ (10)

It can be observed that it consists of two terms and that the previous
assumptions enable to interpret it as an expectation according to
$q_{\phi}({\mathbf{z}}_{T}|{\mathbf{x}}_{T})$. Thus, it is not necessary to
sample discrete variables according to the G-S trick. Moreover, a
regularization term $\beta$ can be introduced in the second part of the ELBO
in order to encourage good performance on labeled data while leveraging the
context of the noisy image during reconstruction. While this model simplifies
the inference, it should be noted that in the generative process, the
observation $x_{t}$ is conditionally independent of its associated label and
may not be adapted to some applications.

####  3.2.2  Semi-supervised Variational Recurrent Neural Network (SVRNN)

The generative model used in the SVRNN can also be seen a particular version
of the TMC model where the latent variable $z_{t}$ consists of the pair
$z_{t}=(z^{\prime}_{t},h_{t})$. The associated transition distribution reads:

$\displaystyle\\!\\!\\!\\!p_{\theta}(v_{t}|v_{t-1}\\!)\\!=\\!p_{\theta}(y_{t}|v_{t-1}\\!)p_{\theta}(z_{t}|y_{t},\\!v_{t-1}\\!)p_{\theta}(x_{t}|y_{t},z_{t},\\!v_{t-1}\\!)\text{,}$
(11)

where

$\displaystyle p_{\theta}(y_{t}|v_{t-1}\\!)$ $\displaystyle=$ $\displaystyle
p_{\theta}(y_{t}|h_{t-1})\text{,}$ $\displaystyle
p_{\theta}(z_{t}|y_{t},v_{t-1})$ $\displaystyle=$
$\displaystyle\\!\delta_{f_{\theta}(z^{\prime}_{t},y_{t},x_{t},h_{t-1})}(h_{t})\\!\\!\times\\!p_{\theta}(z^{\prime}_{t}|y_{t},h_{t-1})\text{,}$
$\displaystyle p_{\theta}(x_{t}|y_{t},z_{t},\\!v_{t-1}\\!)$ $\displaystyle=$
$\displaystyle p_{\theta}(x_{t}|y_{t},z^{\prime}_{t},h_{t-1})\text{,}$

and where $f_{\theta}$ is a deterministic function parameterized by a
Recurrent Neural Network (RNN), for example.

On the other hand, the variational distribution
$q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|$
${\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})$ satisfies the factorization
(3.1) with

$\displaystyle
q(z^{\prime}_{t}|{\mathbf{z}}_{t-1},{\mathbf{y}}_{t},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})=q_{\phi}(z^{\prime}_{t}|x_{t},y_{t},h_{t-1})\text{,}$
$\displaystyle
q(y_{t}|{\mathbf{y}}_{t-1},{\mathbf{z}}_{t-1},{\mathbf{x}}_{T},{\mathbf{y}}_{t+1:T}^{{\mathcal{L}}})=q_{\phi}(y_{t}|x_{t},h_{t-1})\text{,}$

but their final ELBO does not coincide with (3.1). The reason why is that they
derive it from the static case and add a penalization term that encourages
$p_{\theta}(y_{t}|v_{t-1})$ and $q_{\phi}(y_{t}|x_{t},h_{t-1})$ to be close to
the empirical distribution of the data.

####  3.2.3 Deep TMCs

We finally present a very general TMC model from which one can apply any of
the previous techniques. The set of parameters $(\theta,\phi)$ can be
described by any differentiable flexible function $\psi(\cdot)$. In
particular, we consider the case where the parameters are produced by a (deep)
neural network.

Due to the different factorizations of the generating (resp. variational)
distributions, we consider a general notation $p_{\theta}(x_{t}|\cdot)$,
$p_{\theta}(z_{t}|\cdot)$ and $p_{\theta}(y_{t}|\cdot)$ (resp.
$q_{\phi}(y_{t}|\cdot)$, $q_{\phi}(z_{t}|\cdot)$) in order to avoid presenting
a specific dependence between variables. These dependencies are specified for
each model and are presented in the previous sections.

Let $\zeta(y_{t};\cdot)$ and $\varsigma(y_{t};\cdot)$ (resp.
$\lambda(x_{t};\cdot)$; and $\eta(z_{t};\cdot)$, $\tau(z_{t};\cdot)$) be two
probability distributions on $\Omega$ (resp. probability density functions on
$\mathbb{R}^{d_{x}}$; and $\mathbb{R}^{d_{z}}$ ). The general model is
described by:

$\displaystyle p_{\theta}(v_{t}|v_{t-1})$
$\displaystyle=p_{\theta}(x_{t}|\cdot)p_{\theta}(z_{t}|\cdot)p_{\theta}(y_{t}|\cdot)\text{,}$
$\displaystyle p_{\theta}(x_{t}|\cdot)$
$\displaystyle=\lambda(x_{t};\psi_{px}(\cdot))\text{,}$ (12) $\displaystyle
p_{\theta}(z_{t}|\cdot)$ $\displaystyle=\eta(z_{t};\psi_{pz}(\cdot))\text{,}$
(13) $\displaystyle p_{\theta}(y_{t}|\cdot)$
$\displaystyle=\zeta(y_{t};\psi_{py}(\cdot))\text{,}$ (14)

(remember that $\psi_{px}(\cdot)$ denotes the parameters of the distribution
$p_{\theta}(x_{t}|\cdot)$ and can depend on $v_{t-1}$ or
$(v_{t-1},z_{t},y_{t})$ or $(v_{t-1},y_{t})$, etc… according to the original
factorization of $p_{\theta}(v_{t}|v_{t-1})$). Finally, the variational
distribution is given by

$\displaystyle q_{\phi}(z_{t}|\cdot)$
$\displaystyle=\tau(z_{t};\psi_{qz}(\cdot))\text{,}$ (15) $\displaystyle
q_{\phi}(y_{t}|\cdot)$
$\displaystyle=\varsigma(y_{t};\psi_{qy}(\cdot))\text{.}$ (16)

The parameters $\theta$ (resp. $\phi$) are derived from neural networks
$(\psi_{px},\psi_{pz},\psi_{py})$ (resp. $(\psi_{qz},\psi_{qy})$). Note that
in the VLS model, $\psi_{qy}$ is not longer needed since the assumption
$q_{\phi}(y_{t}|z_{t})=p_{\theta}(y_{t}|z_{t})$ is made.

##  4 Simulations

In this section, we present the results of the proposed models on semi-
supervised binary image segmentation. Our goal is to recover the segmentation
of a binary image ($\Omega=\\{\omega_{1},\omega_{2}\\}$) from the noisy
observations ${\mathbf{x}}_{T}$ when a partially segmentation
${\mathbf{y}}_{T}^{\mathcal{L}}$ is available.

In particular, $\zeta(y_{t};\cdot)$ (resp. $\varsigma(y_{t};\cdot)$) is set as
a Bernoulli distribution with parameters $\rho_{py,t}$ (resp. $\rho_{qy,t}$).
As for the distribution $\lambda(x_{t};\cdot)$ (resp. $\eta(z_{t};\cdot)$ and
$\tau(z_{t};\cdot)$), we set it as a Gaussian distribution with parameters
$[\mu_{px,t},{\rm diag}(\sigma_{px,t})]$ (resp. $[\mu_{pz,t},{\rm
diag}(\sigma_{pz,t})]$ and $[\mu_{qz,t},{\rm diag}(\sigma_{qz,t})]$). where
${\rm diag(.)}$ denotes the diagonal matrix deduced from the values of
$\sigma_{\cdot,t}$.

###  4.1 Deep mTMC

In our simulations, we consider three particular cases of the deep TMC model.
We start with the deep minimal TMC (d-mTMC) [15], where the choice of
parameters describes the transition:

$\displaystyle p_{\theta}(v_{t}|v_{t-1})\\!\overset{\rm
mTMC}{=}\\!p_{\theta}(y_{t}|y_{t-1})p_{\theta}(z_{t}|z_{t-1})p_{\theta}(x_{t}|y_{t},z_{t})\text{.}$
(17)

So this model assumes a Markovian distribution for the labels and the latent
variables aim at learning the distribution of the noise given the label and
the latent variable. In order to capture temporal dependencies in the input
data and to have an efficient computation of the variational distribution for
the d-mTMC model, we use a deterministic function to generate $\tilde{h}_{t}$
which takes as input $(x_{t},y_{t},z_{t},\tilde{h}_{t-1})$. Then the
variational distribution
$q_{\phi}({\mathbf{z}}_{T},{\mathbf{y}}_{T}^{\mathcal{U}}|$
${\mathbf{x}}_{T},{\mathbf{y}}_{T}^{\mathcal{L}})$ satisfies the factorization
(3.1) with $q_{\phi}(z_{t}|x_{t},y_{t},\tilde{h}_{t-1})$ and
$q_{\phi}(y_{t}|x_{t},\tilde{h}_{t-1})$.

In this case, the parameters are given by:

$\displaystyle[\mu_{px,t},\sigma_{px,t}]=\psi_{px}(y_{t},z_{t}),$
$\displaystyle[\mu_{pz,t},\sigma_{pz,t}]=\psi_{pz}(z_{t-1}),$
$\displaystyle\rho_{py,t}=\psi_{px}(y_{t-1}),$
$\displaystyle[\mu_{qz,t},\sigma_{qz,t}]=\psi_{qz}(x_{t},y_{t},\tilde{h}_{t-1})\text{,
}$ $\displaystyle\rho_{qy,t}=\psi_{qy}(x_{t},\tilde{h}_{t-1})\text{.}$

###  4.2 Experiments settings

We used the Binary Shape Database [16] and focused on both cattle-type and
camel-type images. To transform these images into a $1$-D signal (
${\mathbf{x}}_{T}$ ), we used a Hilbert-Peano filling curve [17]. To evaluate
the models presented in Section 3.2, we introduced non-linear blurring to
highlight their ability to learn and correct for signal corruption. More
precisely, we generated an artificial noise for the cattle-type by generating
$x_{t}$ according to

$x_{t}|y_{t},x_{t-1}\sim\mathcal{N}\Big{(}\sin(a_{y_{t}}+x_{t-1});\sigma^{2}\Big{)},$
(18)

where $a_{\omega_{1}}=0$ , $a_{\omega_{2}}=0.4$ and $\sigma^{2}=0.25$. We now
consider the camel-type image which is corrupted with a stationary
multiplicative noise (non-elementary noise) given by

$x_{t}|y_{t},z_{t}\sim\mathcal{N}\left(a_{y_{t}};b_{t_{t}}^{2}\right)*z_{t},$
(19)

where $z_{t}\sim\mathcal{N}(0,1)$, $a_{\omega_{1}}=0\text{,
}a_{\omega_{2}}=0.5$ and $b_{\omega_{1}}=b_{\omega_{2}}=0.2$.

The generated images are presented in Fig.1(a) and Fig.2(a), respectively.
More details about the image generation process are available in [15].
Additionally, we randomly selected pixels
$y_{t}\in{\mathbf{y}}_{T}^{\mathcal{L}}$, with a percentage of the pixels
being labeled and the rest considered unobserved (e.g. Fig.1(c) and Fig.2(c)
).

Each model was trained using stochastic gradient descent to optimize the
negative associated ELBO, with the Adam optimizer [18]. The neural networks
$\psi_{(\cdot)}$’s were designed with two hidden layers using rectified linear
units and appropriate outputs, such as linear, softplus, and sigmoid. To
ensure a fair comparison, we matched the total number of parameters of all
models to be approximately equal. As a result, the number of hidden units for
each hidden layer differs for each model. In fact, the SVRNN, mTMC, and VLS
models have 22, 25, and 41 hidden units, respectively. We used an RNN cell to
generate $\tilde{h}_{t}$ (resp. $h_{t}$) for the d-mTMC (resp. SVRNN) model.
In the VLS model, we used the parameterization approach for
$q_{\phi}(z_{t}|{\mathbf{x}}_{T})$ presented in [4], which involves using a
bi-directional RNN cell. We also set the regularization term to $0,1$.

###  4.3 Results

The performance of the models is evaluated in terms of the error rate (ER) of
the reconstruction of the unobserved pixels. Table 1 presents the error rates
obtained for reconstructing unobserved pixels on different images. The
notation image $\%$ is used to indicate the specific image and the percentage
of unobserved labels in the image. As shown in the table, the d-mTMC
consistently outperforms the VSL and the SVRNN, achieving a lower error rate
for each case.

Additionally, we observe that when dealing with elementary noise, the
performance of the VLS model is superior to that of SVRNN. However, this
capability is lost as we increase the percentage of unobserved labels, even
with elementary noise.

|  Cattle $40\%$ |  Camel $40\%$ |  Camel $60\%$
---|---|---|---
VSL | 15,64% | 41,84% | 41,80%
SVRNN | 16,55% | 12,12% | 21,38%
d-mTMM |  1,93% |  2,60% |  3,62%
Table 1: Error rates of the reconstruction of the unobserved pixels on
different images with different percentages of unobserved pixels.

Moreover, our algorithm achieves superior performance with a more complex
noise (the camel-type image). Fig. 2 displays the performance of our proposed
algorithm compared to the VSL and the SVRNN on the camel-type image with
$60\%$ of unobserved labels.

(a) ${\mathbf{x}}_{T}$

(b) True labels

(c) $\mathcal{U}$ \- 40%

(d) VLS

(e) SVRNN

(f) d-mTMC

Fig. 1: Semi-supervised image segmentation with d-TMC models with $40\%$ of
unlabeled pixels.

(a) ${\mathbf{x}}_{T}$

(b) True labels

(c) $\mathcal{U}$ \- 60%

(d) VLS

(e) SVRNN

(f) d-mTMC

Fig. 2: Semi-supervised image segmentation with d-TMC models with $60\%$ of
unlabeled pixels.

##  5 Conclusion

In this paper, we have presented a general framework for semi-supervised
estimation in generative models. In particular, by considering the TMC model,
we have shown that it is possible to obtain a wide variety of generative
models and to estimate them in the common framework of variational inference
in the case where only a part of the observations are labelled. Our
experiments demonstrate the effectiveness of the proposed approach in
achieving state-of-the-art performance on the task of image segmentation.

## References

  * [1] L. R. Rabiner,  “A tutorial on Hidden Markov Models and selected applications in speech recognition,”  Proceedings of the IEEE, vol. 77, no. 2, pp. 257–286, February 1989\.
  * [2] S. Morwal, N. Jahan, and D. Chopra,  “Named entity recognition using hidden markov model (hmm),”  International Journal on Natural Language Computing (IJNLC) Vol, vol. 1, 2012.
  * [3] J.-L. Reyes-Ortiz, L. Oneto, A. Samà, X. Parra, and D. Anguita,  “Transition-aware human activity recognition using smartphones,”  Neurocomputing, vol. 171, pp. 754–767, 2016.
  * [4] M. Chen, Q. Tang, K. Livescu, and K. Gimpel,  “Variational sequential labelers for semi-supervised learning,”  arXiv preprint arXiv:1906.09535, 2019.
  * [5] J. Butepage, H. Kjellstrom, and D. Kragic,  “Predicting the what and how-a probabilistic semi-supervised approach to multi-task human activity modeling,”  in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, 2019, pp. 0–0.
  * [6] I. J. Michael, G. Zoubin, T. S. Jaakola, and L. K. Saul,  “An introduction to variational methods for graphical models,”  in MACHINE LEARNING. 1999, pp. 183–233, MIT Press.
  * [7] W. Pieczynski,  “Chaines de Markov triplet,”  Comptes Rendus de l’Academie des Sciences - Mathematiques, vol. 335, pp. 275–278, 2002,  in French.
  * [8] W. Pieczynski and F. Desbouvries,  “On triplet Markov chains,”  in International Symposium on Applied Stochastic Models and Data Analysis, (ASMDA), 2005.
  * [9] H. White,  “Maximum likelihood estimation of misspecified models,”  Econometrica, vol. 50, no. 1, pp. 1–25, January 1982.
  * [10] R. Douc and E. Moulines,  “Asymptotic properties of the maximum likelihood estimation in misspecified hidden Markov models,”  Annals of Statistics, vol. 40, no. 5, pp. 2697–2732, 2012.
  * [11] D. G. Tzikas, A. C. Likas, and N. P. Galatsanos,  “The variational approximation for Bayesian inference,”  IEEE Signal Processing Magazine, vol. 25, no. 6, pp. 131–146, 2008\.
  * [12] D. P. Kingma and M. Welling,  “Auto-encoding variational Bayes,”  in 2nd International Conference on Learning Representations, ICLR, 2014.
  * [13] C J. Maddison, A. Mnih, and Y. W. Teh,  “The concrete distribution: A continuous relaxation of discrete random variables,”  arXiv preprint arXiv:1611.00712, 2016.
  * [14] E. Jang, S. Gu, and B. Poole,  “Categorical reparameterization with gumbel-softmax,”  arXiv preprint arXiv:1611.01144, 2016.
  * [15] H. Gangloff, K. Morales, and Y. Petetin,  “Deep parameterizations of pairwise and triplet markov models for unsupervised classification of sequential data,”  Computational Statistics & Data Analysis, vol. 180, pp. 107663, 2023.
  * [16] LEMS-Computer Vision Group,  “Binary shape,”  https://vision.lems.brown.edu/content/available-software-and-databases.
  * [17] H. Sagan,  Space-filling curves,  Springer, 2012.
  * [18] D. Kingma and J. Ba,  “Adam: A method for stochastic optimization,”  International conference on learning representations, 12 2014.


# Progressive Query Expansion for Retrieval Over Cost-constrained Data Sources

Muhammad Shihab Rashid  University of California, RiversideRiversideCA
<EMAIL_ADDRESS>,  Jannat Ara Meem  University of California,
RiversideRiversideCA<EMAIL_ADDRESS>,  Yue Dong  University of California,
RiversideRiversideCA<EMAIL_ADDRESS>and  Vagelis Hristidis  University of
California, RiversideRiversideCA<EMAIL_ADDRESS>
(2018)

###### Abstract.

Query expansion has been employed for a long time to improve the accuracy of
query retrievers. Earlier works relied on pseudo-relevance feedback (PRF)
techniques, which augment a query with terms extracted from documents
retrieved in a first stage. However, the documents may be noisy hindering the
effectiveness of the ranking. To avoid this, recent studies have instead used
Large Language Models (LLMs) to generate additional content to expand a query.
These techniques are prone to hallucination and also focus on the LLM usage
cost. However, the cost may be dominated by the retrieval in several important
practical scenarios, where the corpus is only available via APIs which charge
a fee per retrieved document. We propose combining classic PRF techniques with
LLMs and create a progressive query expansion algorithm $\mathsf{ProQE}$ that
iteratively expands the query as it retrieves more documents. $\mathsf{ProQE}$
is compatible with both sparse or dense retrieval systems. Our experimental
results on four retrieval datasets show that $\mathsf{ProQE}$ outperforms
state-of-the-art baselines by 37% and is the most cost-effective.

retriever, cost-constrained, pseudo-relevance-feedback, large language models

††copyright: acmcopyright††journalyear: 2018††doi:
XXXXXXX.XXXXXXX††conference: Make sure to enter the correct conference title
from your rights confirmation emai; June 03–05, 2018; Woodstock, NY††price:
15.00††isbn: 978-1-4503-XXXX-X/18/06††ccs: Computer systems organization
Embedded systems††ccs: Computer systems organization Redundancy††ccs: Computer
systems organization Robotics††ccs: Networks Network reliability

##  1\. Introduction

Figure 1. An overview of ProQE for sparse retrieval

Many Information Retrieval and AI tasks depend on the availability of an
effective Retriever module (Lewis et al., 2020). A Retriever extracts $k$
relevant documents or passages, given a query. It serves as a standalone task
as a core component in modern search engines, or as an intermediate step for
retrieval-augmented question-answering or other downstream tasks. There are
two main paradigms for retrievers: 1) sparse or lexical-based retrievers such
as BM25 (Robertson and Walker, 1994), and 2) dense or embedding-based
retrievers like DPR (Karpukhin et al., 2020) and Contriever (Izacard et al.,
2021). Dense retrievers have been shown to perform better when large amounts
of labeled data are available, whereas BM25 remains competitive on out-of-
domain datasets (Thakur et al., 2021).

On the other side, query expansion is a popular approach to improve the
accuracy of both types of retrievers (Carpineto and Romano, 2012). A popular
method to expand the query has been the use of pseudo-relevance-feedback (PRF)
(Rocchio Jr, 1971; Zhai and Lafferty, 2001), which addresses the query-to-
document vocabulary mismatch problem in Sparse retrieval. Key terms are
extracted from the top-$k$ relevant documents in the first pass retrieval and
appended to the original query to perform the final retrieval. However, the
documents returned from the first stage retrieval may not be relevant, and key
terms added from these documents may introduce noises, thus hindering the
effectiveness of PRF. In contrast, recent LLM-based approaches like query2doc
(Wang et al., 2023b), CoT (Jagerman et al., 2023) and GRF (Mackie et al.,
2023b) skip the first-stage retrieval and use LLMs to generate additional
content to append to the original query. These approaches use pre-trained LLMs
as black boxes and have shown improved results. However, they are prone to
hallucination (Zhang et al., 2023) and thus can generate highly irrelevant
content.

A salient assumption of these LLM query expansion works is that the cost of
retrieving documents is low, compared to the cost of accessing the LLM. For
example, the document collection may be stored in Elastic Search, which has a
very low per-query cost. We argue that this assumption does not always hold in
practical settings, where the dominant cost is the retrieval of result
documents. This is the case when the document corpus is not available or
indexed locally, but is accessed via APIs. For example, legal document
retrieval systems like PACER (pac, 2024), Westlaw (wes, 2024), and LexisNexis
(lex, 2024) charge a fee for retrieving each document. These fees can be as
high as 0.1 USD per page of a document (pac, 2024).

Our key contribution for improving retrieval accuracy is to combine classic
pseudo-relevance feedback expansion techniques with modern LLM-based query
expansion techniques. To mitigate the drawbacks of introducing noises from
pseudo relevance feedback, we employ an LLM as a relevance judge for each
returned result. Specifically, we propose $\mathsf{ProQE}$, shown in Figure 1,
which is a progressive query expansion algorithm that iteratively expands the
query as it retrieves more documents.

Designing an effective progressive query expansion algorithm entails the
following key choices and principles: 1) decide how the terms in each
retrieved document, whose relevance is uncertain, should be used to
potentially adapt the query. Such relevancy determination is similar to the
exploration vs. exploitation trade-off: if we retrieve more documents using
the original query, we may obtain more diverse results, leading to diverse
potential expansion terms, whereas aggressively refining the query using the
early retrieved documents may improve the focus and accuracy of the retrieval
but limited expansion terms. 2) such an expansion method should perform well
for most of the popular ranking algorithms, both based on sparse and dense
retrieval. This will allow our method to be applicable to a wide range of
black-box (e.g. API-based) ranking systems which may input a list of keywords
or weights and return a ranked list of results.

A key feature of $\mathsf{ProQE}$ is its plug-and-play capability, allowing it
to integrate seamlessly with any sparse or dense retrieval methods. The
process operates as follows. For Sparse retrievals, we first retrieve the top
one new document using the original query. Through two LLM calls, it extracts
potential expansion terms from this document and further scores the relevance.
Our scoring function takes the relevance score and all previously retrieved
terms as input, and updates the weights. The terms are appended to the query
based on their updated weights. We repeat this process $n$ times. Retrieving
only 1 new document at each iteration helps by saving unnecessary retrieval
costs. Further, evaluating each document ensures that terms from more relevant
documents receive higher weights, allowing for progressive updates to the
query terms based on LLM feedback. Finally, after $n$ iterations, the final
query is formulated by prompting the LLM using chain-of-thought (Wei et al.,
2022) to retrieve additional context and appending it to the intermediate
query. For dense retrieval models, separate embeddings of the original query,
expansion terms, and CoT output are created and then combined using a weighted
average to form the final query embedding.

We extensively compare our method to state-of-the-art pseudo relevance
feedback and generative query expansion approaches, for multiple types of
sparse and dense retrieval models on four popular datasets: Natural Questions
(NQ) (Kwiatkowski et al., 2019), Web Questions (WQ) (Berant et al., 2013),
TREC (Craswell et al., 2020) DL19, and DL20. $\mathsf{ProQE}$ achieves an
average gain of 37% on MRR and R@1 ranking accuracy compared to the baselines.
We also show that $\mathsf{ProQE}$ is the cheapest among all other baselines.
Our contributions are summarized as follows:

  * •

We introduce the problem of LLM-assisted retrieval over cost-constrained
black-box data sources.

  * •

Our progressive query expansion algorithms serve as a plug-and-play module
that works for both sparse and dense retrieval systems.

  * •

Our comprehensive experimental results demonstrate that our approach
outperforms various baseline expansion methods, over various retrieval models
(dense and sparse) on four datasets. We make our code available to the
research community.111https://anonymous.4open.science/r/ProQE

##  2\. ProQE: Progressive Query Expansion

### Problem Definition.

Given a collection of documents $\mathcal{D}$, a retriever aims to rank the
top-$k$ documents based on relevancy to a query $q$, either indexed with
sparse or dense document representations. The query expansion task aims to
generate an expanded query $q^{\prime}$ from the original query $q$ by adding
additional terms.

In cost-constraint setup, we assume $\mathcal{D}$ is not indexed locally and
is only accessible via a retrieval API $\mathcal{A}$, which charges a fee
$\mathcal{C}$ for the retrieval of each new document given a query. We also
assume the cost typically does not depend on any other variables (i.e. the
size of the query), given how current non-local indexed retrieval systems
charge fees (pac, 2024; wes, 2024). By following how commercial systems such
as ScrapeOps web content retriever (scr, 2024)’s practice, we also assume that
the query interface only charges for retrieving new unique documents. That is,
if we expand the query and resubmit and the same document $p$ is returned,
there is no additional cost.

ProQE extracts key terms from each retrieved document and uses these terms at
each iteration to modify the query, before retrieving more documents.

### ProQE for Sparse Retrieval.

Our method first retrieves top-1 new document $p_{1}$ calling $\mathcal{A}$
using the original query $q$. The relevance of the passage $rel(p_{i})$ is
assessed by prompting an LLM $\mathcal{L}$ with a pointwise ranking
instruction (Liang et al., 2022), ”Is the following passage related to the
query?”. In parallel, $m$ potential expansion terms are extracted using
$\mathcal{L}$ with $q$ and $p_{1}$ given as input with the instruction ”Given
the query and passage, extract {m} keywords that may be useful to better
retrieve relevant passages.”. The weights $w(t_{i})$ of terms $t_{1}\cdots
t_{m}$ are updated using the following equation, and kept in a global
dictionary with total terms $M$.

(1) $w(t_{i})=\begin{cases}w(t_{i})+\beta,&\text{if $rel(p_{1})=1$}\\\
w(t_{i})-\gamma,&\text{if $rel(p_{1})=0$}\end{cases}$

The terms with $w(t_{i})>0$ are considered as expansion terms and are repeated
$int(w(t_{i}))$ times and appended to the original query. The original query
is boosted $\alpha$ times to form the intermediate query $q^{+}$.

(2) $q^{+}=concat(\\{q\\}\times\alpha,\sum_{i}^{M}\\{t_{i}\\}\times w(t_{i}))$

This process is iterated $n$ times. At the beginning of each iteration, the
intermediate query $q^{+}$ is used to retrieve $p_{1}$ and a new $q^{+}$ is
generated at the end. By boosting and decreasing the weights of expansion
terms based on feedback from the LLM and the retrieved passage, only relevant
terms are appended to the query, thereby reducing noise. The iterative process
facilitates focused retrieval in each turn, leading to the generation of
effective query terms. Any irrelevant term added in one iteration is corrected
in subsequent iterations. We tune $m$, $n$, $\alpha$, $\beta$, and $\gamma$ on
dev sets and show that the number of iterations does not vary the final
performance much. We discuss parameter details further in Section 3.

Finally, after $n$-th iteration, we prompt $\mathcal{L}$ using chain-of-
thought instruction (Jagerman et al., 2023): ”Answer the following query, give
rationale before answering.” and receive the output $\theta_{c}$. We stop the
iteration at $n$ as further updates do not improve the performance and may
deteriorate. The final query $q^{\prime}$ is formulated as
$q^{\prime}=concat(q^{+},\theta_{c})$. We observed that appending $\theta_{c}$
with $q$ at the start of the iterations adversely impacts performance, as the
non-factual outputs from the LLM can misdirect the progressive update of
queries via relevant passages.

### ProQE for Dense Retrieval.

Query expansion with key terms typically works best for sparse retrievals as
expansion targets vocabulary mismatch and is uncommon for API-based retrieval
systems. Nonetheless, for completeness, we show that $\mathsf{ProQE}$ also
improves the dense retrieval system. Appending a term multiple times does not
boost its weight in a dense retrieval system as the whole semantic meaning is
captured in an embedding. We use an encoder from a dense retriever model to
create embeddings for the original query $\vec{\mathcal{E}_{q}}$. After each
iteration, intermediate query embedding $\vec{\mathcal{E}_{q^{+}}}$ is
computed as follows.

(3)
$\vec{\mathcal{E}_{q^{+}}}=\sigma\times\vec{\mathcal{E}_{q}}+\tau\times\frac{1}{M}\sum_{i}^{M}w(t_{i})\times\vec{\mathcal{E}_{ti}}$

where $\sigma$ is the query weight and $\tau$ is the term weight for dense
models. After $n$ iterations, similarly, we create the embedding for the CoT
output $\vec{\mathcal{E}_{\theta c}}$ and compute the final query embedding
$\vec{\mathcal{E}_{q^{\prime}}}=\sigma\times\vec{\mathcal{E}_{q^{+}}}+\delta\times\vec{\mathcal{E}_{\theta
c}}$, where $\delta$ is the CoT weight. We use this final query embedding to
search the corpus embeddings using similarity search to retrieve the
documents.

##  3\. Experimental Evaluation

### Datasets.

Following previous work on passage retrieval, we choose the popular benchmark
datasets Natural Questions (NQ) (Kwiatkowski et al., 2019), Web Questions (WQ)
(Berant et al., 2013), TREC (Craswell et al., 2020) DL19, and DL20. For TREC
datasets, there are multiple relevant passages per query contrary to NQ and
WQ. To ensure fairness among datasets, we consider the passages with a score
of 3 to be the relevant ones.

### Implementation.

For experiments, we indexed the document corpus with Pyserini. For LLM choice,
we compared with GPT-3.5, Flan T5-XL, Llama-2 on dev set and chose T5-XL as it
has the best cost-to-performance ratio, also supported by previous work
(Jagerman et al., 2023; Rashid et al., 2024b). We tuned our sparse weight
parameters $\alpha$, $\beta$, $\gamma$ with a range from 0 to 5 and step size
of 1, dense weight parameters $\sigma$, $\tau$, and $\delta$ with a range from
0 to 1 and step size of 0.1, iteration number $n$ (range from 2 to 15 with
step size of 1), and number of potential expansion terms $m$ (range from 3 to
7 with step size of 1) on the dev sets of our datasets and chose the values
$\alpha$ = 1, $\beta$ = 1, $\gamma$ = 0, $\sigma$ = 0.8, $\tau$ = 0.2,
$\delta$ = 0.2, $n$ = 5, and $m$ = 5. Note that, our choices of $\alpha$,
$\sigma$, $\tau$, $\delta$, and $m$ are also supported by previous work (Wang
et al., 2023b; Mackie et al., 2023b).

Method | NQ | WQ | TREC DL 19 | TREC DL 20
---|---|---|---|---
MRR | R@1 | MRR | R@1 | MRR | R@1 | MRR | R@1
Sparse Retrieval
BM25 | 29.84 | 20.77 | 28.16 | 19.00 | 33.59 | 20.93 | 12.85 | 10.00
+RM3 | 28.76 | 20.24 | 31.16 | 22.78 | 30.09 | 20.93 | 10.89 | 8.00
+Rocchio PRF | 25.42 | 17.61 | 26.28 | 18.65 | 28.33 | 18.60 | 10.40 | 8.00
+query2doc ZS | 32.65 | 24.73 | 38.05 | 30.41 | 26.44 | 13.95 | 11.02 | 8.50
+query2doc FS | 35.11 | 25.70 | 38.58 | 29.13 | 37.03 | 20.93 | 12.86 | 10.50
+CoT | 35.42 | 26.48 | 44.07 | 35.48 | 35.78 | 25.58 | 13.73 | 11.00
+GRF | 33.51 | 26.34 | 42.22 | 34.99 | 23.31 | 11.62 | 10.69 | 8.50
+$\mathsf{ProQE}$ | 39.48 | 33.01 | 47.70 | 40.89 | 34.12 | 27.90 | 14.71 | 12.50
docT5 | - | - | - | - | 44.87 | 32.55 | 13.96 | 10.50
+$\mathsf{ProQE}$ | - | - | - | - | 43.05 | 34.88 | 14.02 | 13.60
Dense Retrieval
DPR | 22.67 | 10.33 | 24.69 | 13.13 | - | - | - | -
+CoT | 23.13 | 10.55 | 25.29 | 12.99 | - | - | - | -
+query2doc | 23.38 | 11.02 | 25.23 | 13.04 | - | - | - | -
+$\mathsf{ProQE}$ | 24.73 | 12.32 | 26.42 | 14.18 | - | - | - | -
TCT-Colbert | - | - | - | - | 46.66 | 37.20 | 16.33 | 14.01
\+ $\mathsf{ProQE}$ | - | - | - | - | 47.17 | 39.53 | 16.45 | 14.01
Table 1. Results (MRR and R@1) on all datasets for 20 passages. Best
performing are marked bold.

### Baselines.

We sampled from each retrieval category, sparse and dense with unsupervised
and supervised variants to show the effectiveness of $\mathsf{ProQE}$. For
sparse retrieval, we compare with BM25 and docT5 (Nogueira et al., 2019) as
retrievers. docT5 uses a trained T5-large model to generate a query given a
document and the generated query is appended at the end of the document. We
use Pyserini’s prebuilt index msmarco-v1-passage-d2q-t5. For dense retrieval,
we compare with DPR (Karpukhin et al., 2020) and TCT-Colbert (Lin et al.,
2020). We use DPR’s question encoder fine-tuned on NQ and multiset and the
prebuilt index of Pyserini. TCT-Colbert fine-tunes a student encoder with
distillation from a teacher ColBERT (Khattab and Zaharia, 2020) model. We use
castorini /tct_colbert-msmarco as the query encoder. We chose both state-of-
the-art pseudo-relevance feedback and generative models to compare against
$\mathsf{ProQE}$ as comparing methods. Specifically, we choose the following:
RM3 (Abdul-Jaleel et al., 2004): A PRF approach that expands the query from
top-k retrieved documents. We use fb_terms = 10, fb_doc = 10, and query-
weight = 0.5 following Pyserini’s instructions (pys, 2024).
Rocchio PRF (Rocchio Jr, 1971): Classic Rocchio formula with fb_terms = 5 and
fb_docs = 3.
query2doc (Wang et al., 2023b): Additional passage generated from queries
using LLMs. We show both the zero-shot (ZS) and few-shot (FS) variants.
CoT (Jagerman et al., 2023): Chain-of-Thought prompting output from LLMs.
Generative relevance feedback (GRF) (Mackie et al., 2023b): Multiple types of
additional content such as news articles, essays, keywords, queries, and
entities generated from LLMs given the original query and combined together.

### Main Results: Sparse

We show our main evaluation results in Table 1 for $k$ = 20 passages. We
choose the popular Mean Reciprocal Rank (MRR) and Recall@k as our evaluation
metrics following previous work. We see that in both types of retrievals,
$\mathsf{ProQE}$ improves the baselines by up to an average of 37% from the
variant without expansion. PRF methods like RM3, and Rocchio do not perform
well due to the top retrieved passages not being relevant. Among LLM-based
generative methods, both query2doc and CoT perform really well and
significantly improve the performance. However, to get the best results,
feedback from both retrieved passages and LLMs are needed. The few-shot
variant of query2doc has better MRR in DL 19 dataset than ours but worse R@1.
Generating larger passages with LLM in some scenarios may retrieve better
results at $k$ ¿¿ 1 position. However, in RAG tasks, only the top few passages
are considered hence R@1 metric is more important. docT5 uses a fine-tuned
model hence the performance is better than BM25. We see that ProQE still
improves a trained model’s performance.

### Main Results: Dense

We see that ProQE improves native dense retrievals by an average of 8%,
although the margin is lower than sparse. Unsupervised dense retrieval is not
typically suited for query expansion and is uncommon for cost-constrained data
sources. Regardless, our method is applicable to any such system if released
in the future.

### Analysis: Cost.

We see that ProQE has the highest MRR and R@1 scores, meaning it will need to
retrieve fewer documents to get the relevant ones. Hence ours is the most
cost-effective solution. Even if cost is not a consideration, ours still
performs the best.

Figure 2. Impact of iterations for NQ and WQ dev sets.

### Analysis: Impact of iterations.

We show the impact of iteration number in Figure 2. We see diminishing results
after 5 iterations. However, the performance does not vary much after 5. This
shows that, regardless of what iteration number is chosen, it will improve the
performance of native retrieval systems and other strong baselines.

##  4\. Related Work

Query Expansion. To resolve the lexical mismatch between query and relevant
documents, relevance feedback from documents (Lv and Zhai, 2009; Rocchio Jr,
1971) or knowledge sources (Dalton et al., 2014; Xiong and Callan, 2015; Meij
et al., 2010) are used to expand the query. In cases where the gold labels are
not available, the top retrieved documents are used as pseudo-relevant
documents like KL (Zhai and Lafferty, 2001), RM3 (Lavrenko and Croft, 2017)
etc. PRF methods are primarily used in Sparse retrievals and may introduce
noise in expanded terms, affecting its reliability. Recently, there have been
systems for learned sparse retrievals like SPLADE (Formal et al., 2021), which
is a neural retrieval model that uses BERT and sparse regularization to learn
query and document sparse expansions. PRF methods have also been adapted by
embedding-based dense retrieval models (Karpukhin et al., 2020) like ANCE-PRF
(Yu et al., 2021), ColBERT-PRF (Wang et al., 2023a) which extracts relevant
embeddings from retrieved documents to incorporate to the query embedding.
Both learned sparse and dense retrieval models require training data with gold
relevance labels which becomes exponentially difficult to collect if the
corpus is not available locally. Further, our algorithm can work with both
sparse, learned and unsupervised dense retrieval models.

LLM Augmentation. The use of LLMs (Brown et al., 2020) have spread to
different augmentation techniques such as query rewriting (Jeronymo et al.,
2023; Wu et al., 2021), query-specific reasoning (Ferraretto et al., 2023),
document augmenting (doc2query) (Nogueira et al., 2019) etc. Some very recent
LLM based query expansion works include query2doc (Wang et al., 2023b), where
an LLM generated document is augmented; GRF (Mackie et al., 2023b, a), where
additional context such as keywords, news, facts generated using LLM are
appended, and CoT (Jagerman et al., 2023), where a chain of thought answer is
appended to the query. These works exclusively use LLMs as additional context
which have been shown to hallucinate. Other works include HyDE (Gao et al.,
2022), and GAR (Arora et al., 2023) which require trained models to compute
the embeddings of the generated documents.

Cost-aware Methods. To the best of our knowledge, we are the first to consider
the cost of retrieval as a constraint. Other cost-aware methods like FrugalGPT
(Chen et al., 2023), EcoRank (Rashid et al., 2024b), etc. consider the costs
of LLM APIs and are used for either direct question-answering (Rashid et al.,
2021, 2024a), reasoning, or text re-ranking tasks. We show that, in practical
scenarios, retrieval API costs can dominate the total cost of retrieval
augmented generation.

##  5\. Conclusion

We introduce the problem of retrieval over cost-constrained data sources and
propose a novel progressive query expansion algorithm with a weighted scoring
function that iteratively expands the query as it retrieves more documents and
uses LLMs to navigate the relevant expansion-terms space. Our method can work
with any type of retriever. Our experimental results show that ProQE achieves
an average gain of 37% over other baselines on MRR and R@1 and is also the
most cost-effective.

## References

  * (1)
  * lex (2024) 2024\.  _LexisNexis: Legal and Professional Solutions with AI_.  Retrieved March 24, 2024 from https://www.lexisnexis.com/en-us/home.page
  * pac (2024) 2024\.  _PACER: Public Access to Court Electronic Records_.  Retrieved March 24, 2024 from https://pacer.uscourts.gov/pacer-pricing-how-fees-work
  * pys (2024) 2024\.  _Pyserini: Interactive Search_.  Retrieved March 28, 2024 from https://github.com/castorini/pyserini/blob/master/docs/usage-interactive-search.md
  * wes (2024) 2024\.  _Thomson Reuters: Westlaw Dockets_.  Retrieved March 24, 2024 from https://legal.thomsonreuters.com/en/products/westlaw/dockets-coverage
  * scr (2024) 2024\.  _Your complete toolbox for web scraping_.  Retrieved March 27, 2024 from https://scrapeops.io/docs/intro/
  * Abdul-Jaleel et al. (2004) Nasreen Abdul-Jaleel, James Allan, W Bruce Croft, Fernando Diaz, Leah Larkey, Xiaoyan Li, Mark D Smucker, and Courtney Wade. 2004.  UMass at TREC 2004: Novelty and HARD.  _Computer Science Department Faculty Publication Series_ (2004), 189.
  * Arora et al. (2023) Daman Arora, Anush Kini, Sayak Ray Chowdhury, Nagarajan Natarajan, Gaurav Sinha, and Amit Sharma. 2023.  GAR-meets-RAG Paradigm for Zero-Shot Information Retrieval.  _arXiv preprint arXiv:2310.20158_ (2023).
  * Berant et al. (2013) Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. 2013.  Semantic parsing on freebase from question-answer pairs. In _Proceedings of the 2013 conference on empirical methods in natural language processing_. 1533–1544.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020\.  Language models are few-shot learners.  _Advances in neural information processing systems_ 33 (2020), 1877–1901.
  * Carpineto and Romano (2012) Claudio Carpineto and Giovanni Romano. 2012.  A survey of automatic query expansion in information retrieval.  _Acm Computing Surveys (CSUR)_ 44, 1 (2012), 1–50.
  * Chen et al. (2023) Lingjiao Chen, Matei Zaharia, and James Zou. 2023.  Frugalgpt: How to use large language models while reducing cost and improving performance.  _arXiv preprint arXiv:2305.05176_ (2023).
  * Craswell et al. (2020) Nick Craswell, Bhaskar Mitra, Emine Yilmaz, Daniel Campos, and Ellen M Voorhees. 2020.  Overview of the TREC 2019 deep learning track.  _arXiv preprint arXiv:2003.07820_ (2020).
  * Dalton et al. (2014) Jeffrey Dalton, Laura Dietz, and James Allan. 2014.  Entity query feature expansion using knowledge base links. In _Proceedings of the 37th international ACM SIGIR conference on Research & development in information retrieval_. 365–374.
  * Ferraretto et al. (2023) Fernando Ferraretto, Thiago Laitz, Roberto Lotufo, and Rodrigo Nogueira. 2023.  Exaranker: Explanation-augmented neural ranker.  _arXiv preprint arXiv:2301.10521_ (2023).
  * Formal et al. (2021) Thibault Formal, Benjamin Piwowarski, and Stéphane Clinchant. 2021.  SPLADE: Sparse lexical and expansion model for first stage ranking. In _Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval_. 2288–2292.
  * Gao et al. (2022) Luyu Gao, Xueguang Ma, Jimmy Lin, and Jamie Callan. 2022.  Precise zero-shot dense retrieval without relevance labels.  _arXiv preprint arXiv:2212.10496_ (2022).
  * Izacard et al. (2021) Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. 2021.  Unsupervised dense information retrieval with contrastive learning.  _arXiv preprint arXiv:2112.09118_ (2021).
  * Jagerman et al. (2023) Rolf Jagerman, Honglei Zhuang, Zhen Qin, Xuanhui Wang, and Michael Bendersky. 2023.  Query expansion by prompting large language models.  _arXiv preprint arXiv:2305.03653_ (2023).
  * Jeronymo et al. (2023) Vitor Jeronymo, Luiz Bonifacio, Hugo Abonizio, Marzieh Fadaee, Roberto Lotufo, Jakub Zavrel, and Rodrigo Nogueira. 2023.  Inpars-v2: Large language models as efficient dataset generators for information retrieval.  _arXiv preprint arXiv:2301.01820_ (2023).
  * Karpukhin et al. (2020) Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020.  Dense passage retrieval for open-domain question answering.  _arXiv preprint arXiv:2004.04906_ (2020).
  * Khattab and Zaharia (2020) Omar Khattab and Matei Zaharia. 2020.  Colbert: Efficient and effective passage search via contextualized late interaction over bert. In _Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval_. 39–48.
  * Kwiatkowski et al. (2019) Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. 2019\.  Natural questions: a benchmark for question answering research.  _Transactions of the Association for Computational Linguistics_ 7 (2019), 453–466.
  * Lavrenko and Croft (2017) Victor Lavrenko and W Bruce Croft. 2017.  Relevance-based language models. In _ACM SIGIR Forum_ , Vol. 51. ACM New York, NY, USA, 260–267.
  * Lewis et al. (2020) Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. 2020\.  Retrieval-augmented generation for knowledge-intensive nlp tasks.  _Advances in Neural Information Processing Systems_ 33 (2020), 9459–9474.
  * Liang et al. (2022) Percy Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, et al. 2022\.  Holistic evaluation of language models.  _arXiv preprint arXiv:2211.09110_ (2022).
  * Lin et al. (2020) Sheng-Chieh Lin, Jheng-Hong Yang, and Jimmy Lin. 2020.  Distilling dense representations for ranking using tightly-coupled teachers.  _arXiv preprint arXiv:2010.11386_ (2020).
  * Lv and Zhai (2009) Yuanhua Lv and ChengXiang Zhai. 2009.  A comparative study of methods for estimating query language models with pseudo feedback. In _Proceedings of the 18th ACM conference on Information and knowledge management_. 1895–1898.
  * Mackie et al. (2023a) Iain Mackie, Shubham Chatterjee, and Jeffrey Dalton. 2023a.  Generative and pseudo-relevant feedback for sparse, dense and learned sparse retrieval.  _arXiv preprint arXiv:2305.07477_ (2023).
  * Mackie et al. (2023b) Iain Mackie, Shubham Chatterjee, and Jeffrey Dalton. 2023b.  Generative relevance feedback with large language models. In _Proceedings of the 46th International ACM SIGIR Conference on Research and Development in Information Retrieval_. 2026–2031.
  * Meij et al. (2010) Edgar Meij, Dolf Trieschnigg, Maarten De Rijke, and Wessel Kraaij. 2010.  Conceptual language models for domain-specific retrieval.  _Information Processing & Management_ 46, 4 (2010), 448–469.
  * Nogueira et al. (2019) Rodrigo Nogueira, Wei Yang, Jimmy Lin, and Kyunghyun Cho. 2019.  Document expansion by query prediction.  _arXiv preprint arXiv:1904.08375_ (2019).
  * Rashid et al. (2021) Muhammad Shihab Rashid, Fuad Jamour, and Vagelis Hristidis. 2021.  QuAX: Mining the Web for High-utility FAQ. In _Proceedings of the 30th ACM International Conference on Information & Knowledge Management_. 1518–1527.
  * Rashid et al. (2024b) Muhammad Shihab Rashid, Jannat Ara Meem, Yue Dong, and Vagelis Hristidis. 2024b.  EcoRank: Budget-Constrained Text Re-ranking Using Large Language Models.  _arXiv preprint arXiv:2402.10866_ (2024).
  * Rashid et al. (2024a) Muhammad Shihab Rashid, Jannat Ara Meem, and Vagelis Hristidis. 2024a.  NORMY: Non-Uniform History Modeling for Open Retrieval Conversational Question Answering. In _2024 IEEE 18th International Conference on Semantic Computing (ICSC)_. 101–109.  https://doi.org/10.1109/ICSC59802.2024.00022
  * Robertson and Walker (1994) Stephen E Robertson and Steve Walker. 1994.  Some simple effective approximations to the 2-poisson model for probabilistic weighted retrieval. In _SIGIR’94: Proceedings of the Seventeenth Annual International ACM-SIGIR Conference on Research and Development in Information Retrieval, organised by Dublin City University_. Springer, 232–241.
  * Rocchio Jr (1971) Joseph John Rocchio Jr. 1971.  Relevance feedback in information retrieval.  _The SMART retrieval system: experiments in automatic document processing_ (1971).
  * Thakur et al. (2021) Nandan Thakur, Nils Reimers, Andreas Rücklé, Abhishek Srivastava, and Iryna Gurevych. 2021.  Beir: A heterogenous benchmark for zero-shot evaluation of information retrieval models.  _arXiv preprint arXiv:2104.08663_ (2021).
  * Wang et al. (2023b) Liang Wang, Nan Yang, and Furu Wei. 2023b.  Query2doc: Query expansion with large language models.  _arXiv preprint arXiv:2303.07678_ (2023).
  * Wang et al. (2023a) Xiao Wang, Craig Macdonald, Nicola Tonellotto, and Iadh Ounis. 2023a.  ColBERT-PRF: Semantic pseudo-relevance feedback for dense passage and document retrieval.  _ACM Transactions on the Web_ 17, 1 (2023), 1–39.
  * Wei et al. (2022) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022\.  Chain-of-thought prompting elicits reasoning in large language models.  _Advances in neural information processing systems_ 35 (2022), 24824–24837.
  * Wu et al. (2021) Zeqiu Wu, Yi Luan, Hannah Rashkin, David Reitter, Hannaneh Hajishirzi, Mari Ostendorf, and Gaurav Singh Tomar. 2021.  Conqrr: Conversational query rewriting for retrieval with reinforcement learning.  _arXiv preprint arXiv:2112.08558_ (2021).
  * Xiong and Callan (2015) Chenyan Xiong and Jamie Callan. 2015.  Query expansion with freebase. In _Proceedings of the 2015 international conference on the theory of information retrieval_. 111–120.
  * Yu et al. (2021) HongChien Yu, Chenyan Xiong, and Jamie Callan. 2021.  Improving query representations for dense retrieval with pseudo relevance feedback. In _Proceedings of the 30th ACM International Conference on Information & Knowledge Management_. 3592–3596.
  * Zhai and Lafferty (2001) Chengxiang Zhai and John Lafferty. 2001.  Model-based feedback in the language modeling approach to information retrieval. In _Proceedings of the tenth international conference on Information and knowledge management_. 403–410.
  * Zhang et al. (2023) Yue Zhang, Yafu Li, Leyang Cui, Deng Cai, Lemao Liu, Tingchen Fu, Xinting Huang, Enbo Zhao, Yu Zhang, Yulong Chen, et al. 2023\.  Siren’s song in the AI ocean: a survey on hallucination in large language models.  _arXiv preprint arXiv:2309.01219_ (2023).


proof that derives the fact $\operatorname{pruned}(G,\pi_{0},[\nu]_{i})$ for
some $i\leq|\nu|$.

######  Proof A.19.

Let $\nu^{*}$ be the canonical leaf of $\mathcal{T}(G,\pi_{0})$, and let
$|\nu^{*}|=l$. Let $\nu$ be a non-canonical leaf, and let $|\nu|=k$. Since
$\nu$ is not canonical, one of the following cases must be true:

  * •

$\bar{\phi}(G,\pi_{0},\nu^{*})>\bar{\phi}(G,\pi_{0},\nu)$: since the node
invariants are compared lexicographically, we have following two subcases:

    * –

there is some $m$, such that $m\leq k$ and $m\leq l$, and such that
$\bar{\phi}(G,\pi_{0},[\nu^{*}]_{m})>\bar{\phi}(G,\pi_{0},[\nu]_{m})$. Assume
that $m$ is the smallest number with that property. Let $\pi_{1}$ and
$\pi_{2}$ be the colorings assigned to the nodes $[\nu^{*}]_{m}$ and
$[\nu]_{m}$. Now we have that $\bar{\phi}(G,\pi_{0},[\nu^{*}]_{m-1})$ and
$\bar{\phi}(G,\pi_{0},[\nu]_{m-1})$ are equal, and
$\operatorname{hash}(G,\pi_{1})>\operatorname{hash}(G,\pi_{2})$. Since there
are proofs that derive the fact
$\bar{\phi}(G,\pi_{0},[\nu^{*}]_{m-1})=\bar{\phi}(G,\pi_{0},[\nu]_{m-1})$
(Lemma 18), and the facts $\bar{R}(G,\pi_{0},[\nu^{*}]_{m})=\pi_{1}$ and
$\bar{R}(G,\pi_{0},[\nu]_{m})=\pi_{2}$ (Lemma 16), we can apply the rule
$\mathtt{PruneInvariant}$ to derive the fact
$\operatorname{pruned}(G,\pi_{0},[\nu]_{m})$.

    * –

the invariant $\bar{\phi}(G,\pi_{0},\nu)$ is a prefix of
$\bar{\phi}(G,\pi_{0},\nu^{*})$. This means that $k<l$, and the invariants
$\bar{\phi}(G,\pi_{0},[\nu^{*}]_{k})$ and $\bar{\phi}(G,\pi_{0},\nu)$ are
equal. Let $\pi_{1}$ and $\pi_{2}$ be the colorings assigned to the nodes
$[\nu^{*}]_{k}$ and $\nu$, respectively. Since $\nu$ is a leaf, the coloring
$\pi_{2}$ is discrete. On the other hand, the node $[\nu^{*}]_{k}$ is not a
leaf, so $\pi_{1}$ is not discrete. By Lemma 18, there is a proof that derives
the fact $\bar{\phi}(G,\pi_{0},[\nu^{*}]_{k})=\bar{\phi}(G,\pi_{0},\nu)$, and
by Lemma 16 there are proofs that derive the facts
$\bar{R}(G,\pi_{0},[\nu^{*}]_{k})=\pi_{1}$ and
$\bar{R}(G,\pi_{0},\nu)=\pi_{2}$. Thus, we can apply the rule
$\mathtt{PruneLeaf}$ to derive the fact
$\operatorname{pruned}(G,\pi_{0},\nu)$.

  * •

$\bar{\phi}(G,\pi_{0},\nu^{*})=\bar{\phi}(G,\pi_{0},\nu)$, but
$G_{\nu^{*}}>G_{\nu}$: since the node invariants of $\nu^{*}$ and $\nu$ are
equal, it must hold $l=k$, because node invariants are compared
lexicographically. By Lemma 18, there is a proof that derives the fact
$\bar{\phi}(G,\pi_{0},\nu^{*})=\bar{\phi}(G,\pi_{0},\nu)$. Let $\pi^{*}$ and
$\pi$ be the colorings assigned to $\nu^{*}$ and $\nu$, respectively. By Lemma
16, there are proofs deriving the facts $\bar{R}(G,\pi_{0},\nu^{*})=\pi^{*}$
and $\bar{R}(G,\pi_{0},\nu)=\pi$. Since $\nu^{*}$ and $\nu$ are leaves, both
$\pi^{*}$ and $\pi$ are discrete. From $G_{\nu^{*}}=G^{\pi^{*}}$ and
$G_{\nu}=G^{\pi}$ we have $G^{\pi^{*}}>G^{\pi}$. Now we can apply the rule
$\mathtt{PruneLeaf}$ to derive the fact
$\operatorname{pruned}(G,\pi_{0},\nu)$.

  * •

$\bar{\phi}(G,\pi_{0},\nu^{*})=\bar{\phi}(G,\pi_{0},\nu)$ and
$G_{\nu^{*}}=G_{\nu}$, but $\nu^{*}<_{lex}\nu$: again, since the invariants of
$\nu^{*}$ and $\nu$ are equal, it must hold $l=k$. Assume that
$\pi^{*}=\bar{R}(G,\pi_{0},\nu^{*})$ and $\pi=\bar{R}(G,\pi_{0},\nu)$. Notice
that $G^{\pi^{*}}=G^{\pi}$ implies
$(G^{\pi^{*}},\pi_{0}^{\pi^{*}})=(G^{\pi},\pi_{0}^{\pi})$, since both
$\pi^{*}$ and $\pi$ are discrete colorings finer than $\pi_{0}$. Now it holds
$(G,\pi_{0})^{\sigma}=(G,\pi_{0})$, for $\sigma=\pi^{*}\pi^{-1}$, so
$\sigma\in Aut(G,\pi_{0})$. Since
${\pi^{*}}^{\sigma}=\sigma^{-1}\pi^{*}=(\pi^{*}\pi^{-1})^{-1}\pi^{*}=\pi{\pi^{*}}^{-1}\pi^{*}=\pi$,
we have
$\bar{R}(G,\pi_{0},\nu)=\pi={\pi^{*}}^{\sigma}=\bar{R}(G,\pi_{0},\nu^{*})^{\sigma}=\bar{R}(G^{\sigma},\pi_{0}^{\sigma},{\nu^{*}}^{\sigma})=\bar{R}(G,\pi_{0},{\nu^{*}}^{\sigma})$,
so it must hold ${\nu^{*}}^{\sigma}=\nu$, since no two distinct leaves
correspond to the same coloring. Now we can apply the rule
$\mathtt{PruneAutomorphism}$ to derive the fact
$\operatorname{pruned}(G,\pi_{0},\nu)$.111111In fact, if $i$ is the smallest
number such that $[\nu^{*}]_{i}<_{lex}[\nu]_{i}$, we can derive the fact
$\operatorname{pruned}(G,\pi_{0},[\nu]_{i})$.

Thus, in all cases, we can derive the fact
$\operatorname{pruned}(G,\pi_{0},[\nu]_{i})$ for some $i\leq|\nu|$.

######  Lemma 20.

If a subtree $\mathcal{T}(G,\pi_{0},\nu)$ does not contain the canonical leaf,
then there is a proof that derives the fact
$\operatorname{pruned}(G,\pi_{0},[\nu]_{i})$ for some $i\leq|\nu|$.

######  Proof A.20.

Let define the _height_ $h(\nu)$ of a node $\nu\in\mathcal{T}(G,\pi_{0})$ as a
maximal distance to some of its descendant leaves. We prove the lemma by
induction on the node height. If $h(\nu)=0$, then $\nu$ is a leaf, and the
statement follows from Lemma 19. Let us assume that the statement holds for
all nodes whose height is smaller than some $l>0$, and let $\nu$ be an
arbitrary node such that $h(\nu)=l$. Assume that $W$ is a target cell assigned
to the node $\nu$. Then, by Lemma 17, there is a proof of the fact
$\bar{T}(G,\pi_{0},\nu)=W$. For each $w\in W$, the height of the node
$[\nu,w]$ is smaller than $l$, and the subtree
$\mathcal{T}(G,\pi_{0},[\nu,w])$ also does not contain the canonical leaf, so
by induction hypothesis, there is a proof that derives either the fact
$\operatorname{pruned}(G,\pi_{0},[\nu,w])$, or a fact
$\operatorname{pruned}(G,\pi_{0},[\nu]_{i})$ for some $i\leq|\nu|$. If the
second case is true for at least one $w\in W$, then the statement also holds
for the node $\nu$. Otherwise, for each $w\in W$ there is a proof of the fact
$\operatorname{pruned}(G,\pi_{0},[\nu,w])$, so we can apply the rule
$\mathtt{PruneParent}$ to derive the fact
$\operatorname{pruned}(G,\pi_{0},\nu)$. By proving the induction step, we
proved the lemma.

######  Lemma 21.

If a node $\nu$ is an ancestor of the canonical leaf of the tree
$\mathcal{T}(G,\pi_{0})$, then there is a proof that derives the fact
$\operatorname{on\\_path}(G,\pi_{0},\nu)$.

######  Proof A.21.

The lemma will be proved by induction on $k=|\nu|$. For $k=0$, the node $\nu$
is the root node $[\ ]$. We can apply the rule $\mathtt{PathAxiom}$ to derive
the fact $\operatorname{on\\_path}(G,\pi_{0},[\ ])$. Assume that the statement
holds for some $k$, and let $\nu=[\nu^{\prime},v]$ be an ancestor of the
canonical node such that $|\nu^{\prime}|=k$. Since $\nu^{\prime}$ is also an
ancestor of the canonical leaf, by induction hypothesis there is a proof that
derives the fact $\operatorname{on\\_path}(G,\pi_{0},\nu^{\prime})$. Let $W$
be the target cell assigned to the node $\nu^{\prime}$. By Lemma 17, there is
a proof that derives the fact $\bar{T}(G,\pi_{0},\nu^{\prime})=W$. Notice that
$v\in W$ and for each $w\in W\setminus\\{v\\}$ the node $[\nu^{\prime},w]$ is
not an ancestor of the canonical leaf, that is, the subtree
$\mathcal{T}(G,\pi_{0},[\nu^{\prime},w])$ does not contain the canonical leaf.
By Lemma 20, either there is a proof that derives the fact
$\operatorname{pruned}(G,\pi_{0},[\nu^{\prime},w])$, or a proof deriving the
fact $\operatorname{pruned}(G,\pi_{0},[\nu^{\prime}]_{i})$ for some $i\leq k$.
But in the second case the pruned node $[\nu^{\prime}]_{i}$ would be an
ancestor of the canonical leaf, which is not possible, by Lemma 12. Therefore,
there exists a proof that derives the fact
$\operatorname{pruned}(G,\pi_{0},[\nu^{\prime},w])$ for each $w\in
W\setminus\\{v\\}$. Thus, we can apply the rule $\mathtt{ExtendPath}$ to
derive the fact $\operatorname{on\\_path}(G,\pi_{0},[\nu^{\prime},v])$. We
proved the induction step, so the lemma holds.

######  Lemma 22.

If $\pi^{*}$ is the coloring that corresponds to the canonical leaf $\nu^{*}$
of the tree $\mathcal{T}(G,\pi_{0})$, then there is a proof that derives the
fact $C(G,\pi_{0})=(G^{\pi^{*}},\pi_{0}^{\pi^{*}})$.

######  Proof A.22.

Since the canonical leaf $\nu^{*}$ is an ancestor of itself, by Lemma 21,
there is a proof that derives the fact
$\operatorname{on\\_path}(G,\pi_{0},\nu^{*})$. Also, by Lemma 16, there is a
proof that derives the fact $\bar{R}(G,\pi_{0},\nu^{*})=\pi^{*}$. Since
$\nu^{*}$ is a leaf, $\pi^{*}$ is discrete, so we can apply the rule
$\mathtt{CanonicalLeaf}$ to derive the fact
$C(G,\pi_{0})=(G^{\pi^{*}},\pi_{0}^{\pi^{*}})$.

######  Theorem 23.

Let $(G,\pi_{0})$ be a colored graph, and let $(G^{\prime},\pi^{\prime})$ be
its canonical form. Then there is a proof that corresponds to $(G,\pi_{0})$
deriving the fact $C(G,\pi_{0})=(G^{\prime},\pi^{\prime})$.

######  Proof A.23.

By Lemma 22, we can derive the fact
$C(G,\pi_{0})=(G^{\pi^{*}},\pi_{0}^{\pi^{*}})$, where $\pi^{*}$ is the
coloring corresponding to the canonical leaf $\nu^{*}$. Since the colored
graph $(G^{\pi^{*}},\pi_{0}^{\pi^{*}})$ is the canonical graph of
$(G,\pi_{0})$ by definition, our theorem holds.