# Retrieval-Generation Alignment for End-to-End
Task-Oriented Dialogue System

Weizhou Shen1, Yingqi Gao1, Canbin Huang1, Fanqi Wan1, Xiaojun Quan1, Wei
Bi2$*$
1School of Computer Science and Engineering, Sun Yat-sen University, China
2Tencent AI Lab
{shenwzh3, gaoyq28, huangcb3<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS>Corresponding authors.

###### Abstract

Developing an efficient retriever to retrieve knowledge from a large-scale
knowledge base (KB) is critical for task-oriented dialogue systems to
effectively handle localized and specialized tasks. However, widely used
generative models such as T5 and ChatGPT often struggle to differentiate
subtle differences among the retrieved KB records when generating responses,
resulting in suboptimal quality of generated responses. In this paper, we
propose the application of maximal marginal likelihood to train a perceptive
retriever by utilizing signals from response generation for supervision. In
addition, our approach goes beyond considering solely retrieved entities and
incorporates various meta knowledge to guide the generator, thus improving the
utilization of knowledge. We evaluate our approach on three task-oriented
dialogue datasets using T5 and ChatGPT as the backbone models. The results
demonstrate that when combined with meta knowledge, the response generator can
effectively leverage high-quality knowledge records from the retriever and
enhance the quality of generated responses. The codes and models of this paper
are available at https://github.com/shenwzh3/MK-TOD.

##  1 Introduction

Task-oriented dialogue systems (TOD) assist users to accomplish daily tasks
such as restaurants, scheduling appointments, and navigating traffic by
leveraging external knowledge bases. Among them, pipeline systems (Henderson
et al., 2014; Hosseini-Asl et al., 2020) involve several intermediate stages
such as dialog state tracking and system policy learning for retrieving
knowledge and generating responses. In contrast, end-to-end task-oriented
dialog systems (E2E-TOD) (Wu et al., 2022; Tian et al., 2022) have gained
increasing concentration for their ability to directly generate responses
based on the knowledge base without intermediate annotations. Although the
end-to-end paradigm appears to be more compatible with practical scenarios and
large-scale language models, it imposes challenges in acquiring and utilizing
external knowledge as no belief state is provided for knowledge retrieval.

Figure 1: A demonstrative case in E2E-TOD. The table displays retrieved
entities sorted by retrieval order. The correct entity is highlighted in blue.
However, the response generator mistakenly selects the false entity,
highlighted in red, leading to an erroneous response.

Retrieval-augmented generation (Lewis et al., 2020; Ren et al., 2021; Singh et
al., 2021) has demonstrated success in various knowledge-intensive tasks by
employing a held-out dense retriever to retrieve knowledge and then taking the
knowledge to generate results. Q-TOD (Tian et al., 2022) applies this approach
to E2E-TOD and significantly outperforms previous methods that combine
knowledge retrieval and response generation into a single model (Madotto et
al., 2018; Qin et al., 2020; Raghu et al., 2021). However, our preliminary
study in Section 5.3 shows that under this framework the correlation between
the performance of knowledge retriever and that of response generator is
relatively weak, meaning that simply improving the retriever may not lead to a
better generator. We characterize this phenomenon as the _misalignment_
between the retrieval and generation processes in E2E-TOD systems. This
misalignment poses a bottleneck for current dialogue systems, as improvements
in the retriever component do not necessarily translate to enhanced generation
quality.

Through qualitative analysis, we hypothesize that the misalignment between
retrieval and generation is attributed to the homogeneity of retrieved
knowledge entities. As illustrated in Figure 1, the retrieved entities exhibit
a high degree of similarity, with only minor variations in their values.
Consequently, since the response generator is trained on reference responses
that predominantly consist of language tokens rather than knowledge-related
tokens, it struggles to differentiate between similar entities and may
inadvertently select inappropriate entities for response generation.

In this paper, we introduce Meta Knowledge for end-to-end Task-Oriented
Dialogue system (MK-TOD) as a solution to address the retrieval-generation
misalignment. MK-TOD aims to correlate the performance of the knowledge
retriever and response generator for improved system performance. To enhance
the knowledge retriever, we propose the application of maximum marginal
likelihood (Singh et al., 2021) for progressive retriever updating during the
training of the response generator. In order to enable the response generator
to distinguish between entities, we explore several methods for utilizing
retrieval-related _meta knowledge_. Here, meta knowledge refers to various
information about the retrieved entities, such as retrieval order, retrieval
confidence, and co-occurrence rate. We propose three approaches for
incorporating the meta knowledge: adding special prefix tokens, using prompts,
and applying contrastive learning. Additionally, we investigate the
introduction of negative knowledge during the generator’s training to enhance
its discriminative ability.

We apply MK-TOD to several backbone models, including T5 (Raffel et al., 2020)
and the large language model ChatGPT (OpenAI, 2022). We compare MK-TOD with
other E2E-TOD systems on three benchmark datasets, namely SMD, CamRest, and
WoZ (Eric et al., 2017; Wen et al., 2017; Eric et al., 2020). The empirical
results demonstrate the superiority of our proposed system over the current
state-of-the-art systems with similar model scales. Additionally, our system
effectively enhances the performance of ChatGPT in E2E-TOD with in-context
learning. Furthermore, through comprehensive analysis, we uncover that our
meta-knowledge approach successfully alleviates the misalignment between the
retriever and generator. This approach empowers the generator to better
differentiate between similar entities during response generation.

Figure 2: The MK-TOD framework comprises a knowledge retriever and a response
generator. Given the dialogue context, the retriever retrieves entities from
the knowledge base. Each entity is concatenated with its corresponding meta
knowledge and subsequently input into the generator to generate the response.
The optimization process involves two likelihoods: the normal text generation
likelihood and the marginal likelihood.

##  2 Related Works

###  2.1 End-to-End Task-Oriented Dialogue

The existing work on the usage of external knowledge in end-to-end task-
oriented dialogue systems can be divided into three categories. The first
category takes the whole knowledge base as the model input, and conducts
knowledge selection and response generation in one single model. For instance,
Mem2seq (Madotto et al., 2018), KB-Retriever (Qin et al., 2019), GLMP (Wu et
al., 2019) and CDNET (Raghu et al., 2021) employ memory networks for querying
knowledge. UnifiedSKG (Xie et al., 2022) directly concatenates entities as the
input of Transformers. The second category directly encodes knowledge into
model parameters. GPT-KE (Madotto et al., 2020) pre-trains their model on
augmented dialog data to embed the knowledge base, while ECO (Huang et al.,
2022) applies tri-constraints on top of GPT-KE to ensure entity consistency.
The third category is to use an individual retriever to retrieve knowledge.
For example, Q-TOD (Tian et al., 2022) decouples the dialogue system into a
retriever and a generator and uses the generator to generate a query for
knowledge retrieval. DialogKG (Rony et al., 2022) inputs the flattened records
to a graph neural network to select entities. And MAKER (Wan et al., 2023)
introduces a multi-grained retrival with both entity and attribute selection.
As mentioned earlier, although the retrieve-then-generate framework has been
one of the most successful paradigms to date, it can lead to misalignment
between the retriever and the generator in end-to-end task-oriented dialogue
systems.

###  2.2 Retrieval-Augmented Generation

With the success of the dual-encoder neural retriever (Karpukhin et al.,
2020), the retrieval-augmented generation framework is widely applied to
various knowledge-intensive tasks. This framework uses a retriever to retrieve
knowledge from a knowledge base and inputs the retrieval results into a
generator to generate the answer. Among them, RAG (Lewis et al., 2020)
separately encodes each retrieved knowledge record with the query and
marginalizes the probabilities of the answer based on each entity. FiD
(Izacard and Grave, 2021) encodes each retrieved knowledge like RAG and fuses
their hidden states in the decoder. FiD-KD (Izacard and Grave, 2022) and EMDR2
(Singh et al., 2021) are both based on the FiD framework but with different
retriever training methods: FID-KD uses knowledge distillation while EMDR2
uses marginal likelihood. REPLUG (Shi et al., 2023) applies the method of RAG
to large language models but only updates the retriever during training.

##  3 Methodology

The framework of our proposed MK-TOD is depicted in Figure 2. MK-TOD consists
of a retriever and a response generator. In each dialogue turn, the retriever
retrieves a set of relevant entities, which are then combined with retrieval-
related meta knowledge and the dialogue context. The generator utilizes this
information to generate a response for the current turn. In the following
section, we will introduce the notations and provide an overview of our
method. Then, we will delve into the detailed explanations of two crucial
components: maximum marginal likelihood and meta knowledge.

###  3.1 Notations

Given a dialogue $\mathcal{D}=\\{u_{1},r_{1},...,u_{T},r_{T}\\}$ of $T$ turns,
where $u_{t}$ and $r_{t}$ are the $t$-th turn user utterance and system
response, respectively. We use $c_{t}$ to represent the dialog context of the
$t$-th turn, where $c_{t}=\\{u_{1},r_{1},...,u_{t-1},r_{t-1},u_{t}\\}$. An
external knowledge base (KB) is provided in the form of a set of entities,
i.e., $\mathcal{K}=\\{e_{1},e_{2},...,e_{B}\\}$, where each entity $e_{i}$
consists of several attribute-value pairs and $B$ is the size of knowledge
base. End-to-end task-oriented dialog systems take dialogue context $c_{t}$
and knowledge base $\mathcal{K}$ as input and generate an informative natural
language response $r_{t}$.

###  3.2 System Overview

The retriever module comprises a context encoder and an entity encoder. The
context encoder transforms the current dialogue context $c_{t}$ into a vector
representation $h_{c_{t}}$. On the other hand, the entity encoder concatenates
the attribute-value pairs of each entity as plain text and encodes it into a
vector representation $h_{e_{i}}$. The matching score $s_{t,i}$ is computed by
taking the dot product between $h_{c_{t}}$ and $h_{e_{i}}$. Consequently, the
top-$K$ entities with the highest scores are selected as candidate entities
$\mathcal{E}_{t}$ for the current dialogue turn. Furthermore, if meta
knowledge is utilized, each entity in $\mathcal{E}_{t}$ is augmented with its
corresponding meta knowledge.

The generator takes the retrieved entities $\mathcal{E}_{t}$ and dialogue
context $c_{t}$ as inputs to generate the final system response $r_{t}$. The
probability of generating the response $r_{t}$ given the entities
$\mathcal{E}_{t}$ and dialogue context $c_{t}$ can be calculated as follows:

$p(r_{t}|c_{t},\mathcal{E}_{t};\theta)=\prod_{j=1}^{|r_{t}|}p(r_{t,j}|r_{t,<j},c_{t},\mathcal{E}_{t};\theta),$
(1)

where $\theta$ denotes the parameters of the generator.

Similar to most text generation tasks, we incorporate the negative log-
likelihood (NLL) loss as a training objective to train the generator:

$\mathcal{L}_{NLL}=-\text{log}p(r_{t}|c_{t},\mathcal{E}_{t};\theta).$ (2)

###  3.3 Maximum Marginal Likelihood

Due to the absence of retrieval labels for training the retriever, we depend
on supervision signals from the generator. However, since it is not possible
to backpropagate gradients through the NLL loss in Eq. (2) to the retriever,
we propose updating the retriever’s parameters by maximizing the marginal
likelihood (MML) of the response $r_{t}$. The marginal likelihood offers a
Bayesian perspective to compute $p(r_{t}|c_{t},\mathcal{K})$ by integrating
the likelihood over all the entities in the knowledge base:

$p(r_{t}|c_{t};\phi,\theta)=\sum_{e_{i}\in\mathcal{K}}q(e_{i}|c_{t};\phi)p(r_{t}|c_{t},e_{i};\theta),$
(3)

where $\phi$ denotes the parameters of the retriever and $q(e_{i}|c_{t};\phi)$
is the retrieval probability of entity $e_{i}$. Note that computing
$q(e_{i}|c_{t};\phi)$ for all entities in the entire knowledge base would
incur an unaffordable computational cost for Eq. (3). Therefore, following the
approach of EMDR2 Singh et al. (2021), we compute $q(e_{i}|c_{t};\phi)$ over
the retrieved entities $\mathcal{E}_{t}$ instead of the entire knowledge base
$\mathcal{K}$:

$p(r_{t}|c_{t};\phi,\theta)=\sum_{e_{t,i}\in\mathcal{E}_{t}}q(e_{t,i}|c_{t};\phi)p(r_{t}|c_{t},e_{t,i};\theta),$
(4)

where $q(e_{t,i}|c_{t};\phi)$ is implemented as follows:

$q(e_{t,i}|c_{t};\phi)=\frac{\text{exp}({s_{t,i}})}{\sum_{e_{t,j}\in\mathcal{E}_{t}}\text{exp}({s_{t,j}})}.$
(5)

The loss function for the marginal likelihood is defined as follows:

$\mathcal{L}_{MML}=-\text{log}\sum_{e_{t,i}\in\mathcal{E}_{t}}q(e_{t,i}|c_{t};\phi)p(r_{t}|c_{t},e_{t,i};\theta).$
(6)

By incorporating $q(e_{t,j}|c_{t};\phi)$, we can propagate gradients back to
the retriever and update its parameters. The ultimate training loss function
for MK-TOD is defined as follows:

$\mathcal{L}=\alpha\mathcal{L}_{NLL}+\beta\mathcal{L}_{MML},$ (7)

where $\alpha$ and $\beta$ are hyperparameters.

###  3.4 Meta Knowledge

We introduce the concept of retrieval-related meta knowledge, which
encompasses various information about the retrieved entities to guide the
generator and enhance the alignment between retrieval and generation. Three
key factors are considered in the meta knowledge: retrieval order, retrieval
confidence, and co-occurrence.

Retrieval order: The retriever evaluates entities based on their matching with
the dialogue context, prioritizing those that exhibit a higher degree of
matching. We incorporate the retrieval order of each entity as a part of the
meta knowledge.

Retrieval confidence: To provide more retrieval information, we categorize
retrieved entities into low-confidence, middle-confidence, and high-confidence
based on retrieval scores. The thresholds for categorizing entities are hyper-
parameters111We established the following ranges: (-infinity, 0.4] indicates
low confidence, (0.4, 0.75] indicates medium confidence, and (0.75, +infinity)
indicates high confidence.. Retrieval confidence, in conjunction with
retrieval order, enables the generator to disregard entities with low
confidence but high retrieval order.

Co-occurrence: Entities that have already appeared in the dialogue context are
more likely to be relevant for future responses. Thus, we inform the generator
about the occurrence of entities in the dialogue context through meta
knowledge.

To implement the above meta knowledge in our system, we design three
approaches: prefix, prompt, and contrastive learning.

####  3.4.1 Prefix

In this approach, we create a mapping function that assigns special tokens
representing meta knowledge to each entity. For instance, an entity ranked
second in retrieval order, with middle retrieval confidence, and not yet
mentioned in the context would be mapped to the set of <2nd-entity>, <mid-
confidence>, <new-entity>222The complete mapping is discussed in Appendix B.1.
These prefix tokens are then concatenated with the corresponding entity and
input to the generator during both training and inference stages.

####  3.4.2 Prompt

To fully leverage the generator’s language modeling capability, we explore
using prompts to incorporate meta knowledge. Here, we design a mapping
function that assigns each entity a set of prompts, which are natural language
sentences representing the meta knowledge. For example, a prompt can be “This
is the top-1 recalled entity with low confidence”.333The complete mapping is
discussed in Appendix B.2 Similar to the prefix approach, these prompts are
concatenated with the corresponding entities and fed into the generator.

####  3.4.3 Contrastive Learning

We can also train the generator to distinguish between entities by employing
contrastive learning. In this approach, we select a subset of entities from
the retrieved entities $\mathcal{E}_{t}$ based on their retrieval order,
forming a positive entity set $\mathcal{E}_{t}^{*}$.444Retrieval confidence
and co-occurrence are not considered to avoid sparsity in the positive entity
set. For each entity $e_{t,i}$ in $\mathcal{E}_{t}^{*}$, we compute its
length-normalized log-likelihood of generating the response:

$d_{t,i}=\frac{\text{log}p(r_{t}|c_{t},e_{t,i};\theta)}{|r_{t}|},$ (8)

where $|r_{t}|$ is the length of $r_{t}$. Additionally, we calculate the log-
likelihood of generating the response without any entity as the baseline
likelihood:

$d_{t}^{-}=\frac{\text{log}p(r_{t}|c_{t};\theta)}{|r_{t}|}.$ (9)

We employ a pairwise marginal ranking loss that ensures the likelihood of
positive entities greater than the baseline likelihood by a certain margin:

$\mathcal{L}_{CTR}=\sum_{e_{t,i}\in\mathcal{E}_{t}^{*}}\text{max}(0,d_{t}^{-}-d_{t,i}+\lambda),$
(10)

where $\lambda$ is a hyperparameter. We then add this loss term to the loss
function of MK-TOD:

$\mathcal{L}=\alpha\mathcal{L}_{NLL}+\beta\mathcal{L}_{MML}+\gamma\mathcal{L}_{CTR}.$
(11)

###  3.5 Negative Entity

Inspired by negative sampling in information retrieval (Karpukhin et al.,
2020), we also consider incorporating negative entities into the generator.
The negative entity, denoted as $e_{t}^{-}\notin\mathcal{E}_{t}$, is chosen as
the entity with the lowest retrieval score from $\mathcal{K}$. Special meta
knowledge is designed for the negative entity as well.555Details are provided
in Appendix B.1 and B.2 Note that the negative entity is different from the
baseline likelihood in the above contrastive learning (Section 3.4.3).

###  3.6 Model Inference

During inference, we first retrieve entities using the retriever. Then, we
prepend each entity with its corresponding meta knowledge. Finally, we
concatenate the entities with the dialogue context and input the resulting
sequence to the generator to generate the final response. Notably, we do not
include negative entities during inference.

###  3.7 Discussion

The conceptual of meta-knowledge is to empower the generator with the
generator’s capability to leverage retrieved information to enhance the
generation process. Noteworthy is the fact that, while we deliberately furnish
the generator with meta-knowledge such as retrieval order and confidence, this
does not obligate the generator to strictly use the higher-ranking entities.
Instead, our approach encourages the generators to autonomously differentiate
between these entities. Specifically, for T5, we integrate meta-knowledge as a
part of the model input, facilitating its assimilation during the training
process. With respect to ChatGPT, our method capitalizes on the model’s pre-
existing knowledge to comprehend the essence of meta-knowledge. The results
presented in Sections 5.5 exemplify the behaviors acquired by the generator
and how these behaviors contribute to enhanced performance.

##  4 Experimental Settings

###  4.1 Dataset

We evaluate our MK-TOD on three task-oriented dialogue datasets: MultiWOZ 2.1
(MWOZ) (Eric et al., 2020), CamRest (Wen et al., 2017), and Stanford Multi-
Domain (SMD) (Eric et al., 2017). We compare the methods with two different
settings about the knowledge bases: First, each dialogue has a small session-
level knowledge base associated with the user goal, which is the typical
setting of previous work. Second, all conversations shares a dataset-level
large-scale knowledge base by injecting all the session-level knowledge bases.
There are 223 and 112 entities in the large-scale knowledge base for MWOZ and
CamRest, respectively. Other detailed statistics of these datasets are shown
in Appendix A.

For all three datasets, we employ BLEU (Papineni et al., 2002) and Entity F1
(Eric et al., 2017) as the metrics to evaluate the quality of generated
responses. Entity F1 assesses the presence of accurate knowledge in the
responses by calculating the micro-averaged precision and recall scores of
attribute values. Additionally, for experiments conducted on large-scale
knowledge bases, we introduce Recall@K as a performance metric for the
retriever. Recall@K measures the percentage of gold entities appearing in the
retrieved entities.

Model | MWOZ | CamRest
---|---|---
BLEU | Entity F1 | Recall@7 | BLEU | Entity F1 | Recall@7
DF-Net (Qin et al., 2020) | 6.45 | 27.31 | - | - | - | -
EER (He et al., 2020b) | 11.60 | 31.86 | - | 20.61 | 57.59 | -
FG2Seq (He et al., 2020a) | 10.74 | 33.68 | - | 19.20 | 59.35 | -
CDNET (Raghu et al., 2021) | 10.90 | 31.40 | - | 16.50 | 63.60 | -
Q-TOD (T5-Large) (Tian et al., 2022) | 15.52 | 46.74 | 92.97 | 21.44 | 63.88 | 95.52
MAKER (T5-Base) (Wan et al., 2023) | 16.25 | 50.87 | - | 26.19 | 72.09 | -
MAKER (T5-Large) (Wan et al., 2023) | 18.23 | 52.12 | - | 25.34 | 72.43 | -
Oursprefix (Base) | 16.39 | 50.35 | 92.51 | 25.23 | 71.15 | 94.35
Oursprompt (Base) | 17.56 | 50.69 | 92.74 | 26.69 | 71.67 | 92.24
Oursctr (Base) | 15.96 | 51.35 | 92.74 | 26.85 | 73.51 | 93.88
Oursprefix (Large) | 16.69 | 53.59 | 92.93 | 27.32 | 72.77 | 91.08
Oursprompt (Large) | 17.15 | 52.99 | 94.42 | 26.88 | 72.92 | 95.41
Oursctr (Large) | 17.40 | 53.26 | 95.22 | 27.82 | 71.98 | 95.38
ChatGPT (OpenAI, 2022) | 6.79 | 30.31 | 92.74∗ | 14.76 | 52.92 | 94.35∗
Oursprefix (LLM) | 7.01 | 30.69 | 92.74∗ | 14.51 | 52.38 | 94.35∗
Oursprompt (LLM) | 7.31 | 32.04 | 92.74∗ | 14.91 | 53.58 | 94.35∗

Table 1: Overall results of E2E-TOD systems with large-scale knowledge bases
on MWOZ and CamRest, where “∗” means that we directly use the retriever co-
trained with T5-Base using MML.

###  4.2 Implementation Details

We utilize BERT (Devlin et al., 2019) as the context encoder and entity
encoder for the retriever. As for the generator, we employ T5 (Raffel et al.,
2020) and ChatGPT (OpenAI, 2022). Note that ChatGPT is not fine-tuned but
instead undergoes in-context learning using our datasets.666Further details
are provided in Appendix C The retriever for ChatGPT is directly copied from
the retriever trained with T5 using MML. All experiments are performed on a
single 24G NVIDIA RTX 3090 GPU, and the best checkpoints are selected based on
Entity F1 scores on the validation set. Hyperparameter settings are listed in
Appendix E.

Consistent with previous studies (Qin et al., 2019), we initialize the
retriever through pre-training with distant supervision to prevent collapsed
representations. Additional details on the pre-training process can be found
in Appendix D.

###  4.3 Baseline Methods

We include several strong baselines for comparison.

Implicit retrieval: These methods combine knowledge retrieval and response
generation in a single model, including DF-Net (Qin et al., 2020), EER (He et
al., 2020b), FG2Seq (He et al., 2020a), CDNET (Raghu et al., 2021) and GPT-KE
(Madotto et al., 2020).

Explicit retrieval: These approaches decouple the TOD system into a knowledge
retriever and a response generator, including Q-TOD (Tian et al., 2022),
DialoKG (Rony et al., 2022) and MAKER (Wan et al., 2023).

Large language models: Large language models (LLMs), such as ChatGPT (OpenAI,
2022), have demonstrated remarkable capabilities in engaging in dialogues with
humans. We establish a baseline LLM utilizing ChatGPT as the response
generator by leveraging the gpt-3.5-turbo API. To enhance its performance in
TOD tasks, we integrate our knowledge retriever with the system.

##  5 Results and Analysis

In this section, we present the overall results obtained using both large-
scale knowledge bases and condensed knowledge bases. Besides, we demonstrate
the phenomenon of retrieval-generation misalignment and conduct ablation
studies.

Model | MWOZ | CamRest | SMD
---|---|---|---
BLEU | Entity F1 | BLEU | Entity F1 | BLEU | Entity F1
DF-Net (Qin et al., 2020) | 9.40 | 35.10 | - | - | 14.40 | 62.70
GPT-2+KE (Madotto et al., 2020) | 15.05 | 39.58 | 18.00 | 54.85 | 17.35 | 59.78
EER (He et al., 2020b) | 13.60 | 35.60 | 19.20 | 65.70 | 17.20 | 59.00
FG2Seq (He et al., 2020a) | 14.60 | 36.50 | 20.20 | 66.40 | 16.80 | 61.10
CDNET (Raghu et al., 2021) | 11.90 | 38.70 | 17.80 | 62.90 | 21.80 | 68.60
DialoKG (Rony et al., 2022) | 12.60 | 43.50 | 23.40 | 75.60 | 20.00 | 65.90
Q-TOD (T5-Base) (Tian et al., 2022) | - | - | - | - | 20.14 | 68.22
Q-TOD (T5-Large) (Tian et al., 2022) | 17.62 | 50.61 | 23.75 | 74.22 | 21.33 | 71.11
Oursprefix (Base) | 16.97 | 51.99 | 26.39 | 72.43 | 23.96 | 68.60
Oursprompt (Base) | 17.05 | 52.42 | 25.00 | 72.09 | 23.54 | 68.28
Oursctr (Base) | 17.33 | 51.86 | 26.76 | 73.60 | 24.77 | 67.86
Oursprefix (Large) | 18.02 | 53.13 | 25.68 | 71.98 | 24.97 | 72.87
Oursprompt (Large) | 16.66 | 52.96 | 26.40 | 72.80 | 25.21 | 73.04
Oursctr (Large) | 17.55 | 52.97 | 26.20 | 71.72 | 25.43 | 73.31
ChatGPT (OpenAI, 2022) | 7.47 | 32.87 | 15.29 | 54.71 | 14.60 | 58.11
Oursprefix (ChatGPT) | 7.22 | 32.78 | 15.56 | 54.96 | 15.07 | 58.41
Oursprompt (ChatGPT) | 7.58 | 35.84 | 16.07 | 56.83 | 15.24 | 59.72

Table 2: Overall results of E2E-TOD systems with condensed knowledge bases on
MWOZ, SMD, and CamRest. The best results are highlighted in bold, and the
second-best results are underlined.

###  5.1 Overall Results with Large-Scale KBs

Comparing our method with others in the setting of retrieving knowledge from a
large-scale knowledge base aligns better with real-world TOD scenarios.
Therefore, we begin by comparing our proposed MK-TOD approach with baselines
in the context of large-scale knowledge bases. The results of this comparison
are displayed in Table 1. The upper part of Table 1 shows the results of
methods employing a fine-tuned response generator. “Oursprefix”, “Oursprompt”,
and “Oursctr” denote our method implementing meta knowledge using prefix,
prompt, and contrastive learning techniques, respectively. “Base” and “Large”
following the method names indicate the use of T5-Base or T5-Large as the
response generator.

Overall, MK-TOD outperforms all previous methods with the same scale of
generator model, indicating the effect of our proposed meta-knowledge. Further
more, Q-TOD’s retriever can achieve comparable and even higher performance
than ours due to their utilization of an additional query generator. However,
even when employing T5-Base and a relatively weaker retriever, our method
still surpasses Q-TOD in terms of BLEU and Entity F1 by a significant margin.
This indicates that our proposed method effectively utilizes the retrieved
knowledge better than Q-TOD. It is also noteworthy that the Entity F1 score of
T5-Large in CamRest is poorer than that of T5-Base, we attribute this to the
small size of CamRest’s training data, which includes merely 406 dialogues,
leading to overfitting particularly for larger models.

The bottom part of Table 2 presents the results of methods employing ChatGPT.
Since ChatGPT is not fine-tunable, we did not apply contrastive learning for
meta knowledge. According to the results, we found out that relying solely on
in-context learning does not enable ChatGPT to perform as well as the fine-
tuned methods in the context of E2E-TOD. However, our proposed approach
outperforms the baseline. Additionally, our proposed prefix method for
implementing meta knowledge yields only marginal improvement or even performs
worse than ChatGPT. This is attributed to ChatGPT’s limited ability to learn
the special prefix tokens representing meta knowledge from a limited number of
in-context demonstrations and concise explanatory text. In contrast, our
proposed prompt method significantly enhances its performance.

###  5.2 Overall Results with Condensed KBs

To make a comprehensive comparison with the previous methods, we also follow
the previous works’ setting the conduct evaluations on the condensed knowledge
base. The results in Table 2 indicate that our proposed method surpasses the
baselines on MWOZ and SMD with the same model scale, validating the efficacy
of our approach. However, among the three meta knowledge implementations, it
is challenging to determine a clear preference as the fine-tuned generator
tends to learn all of them.

For the evaluation with ChatGPT on the condensed knowledge base, we can still
observe the performance gain of ChatGPT when enhanced with our proposed meta-
knowledge. Biside, the performance gain is more significant than that of the
large-scale knowledge bases, suggesting that ChatGPT has a higher demand for
retrieval quality.

###  5.3 Retrieval-Generation Misalignment

To investigate the influence of retrieval performance on the E2E-TOD
generator, we select six retrievers on MWOZ with a large-scale knowledge base.
The details of the retrievers can be found in Appendix G. We then use
different generators to generate responses based on the retrieval results. As
generators, we choose Q-TOD, FiD (Izacard and Grave, 2021), and ChatGPT. The
Entity F1 scores of these generators are depicted in Figures 3(a) and 3(b) as
the retrieval performance varies with different retrievers.

The solid lines in the figures show that the performance of generators does
not consistently align with that of retrieval performance. Furthermore, the
performances of Q-TOD and FiD with oracle entities are even worse than those
with a weak retriever. We refer to this phenomenon as retrieval-generation
misalignment. In contrast, the dashed lines, which depict the results of the
generators with our proposed meta knowledge, exhibit greater consistency
between the retrieval performance and the generators. This indicates that our
proposed method mitigates the misalignment issue. The correlation coefficients
shown in parentheses next to method names further confirm this observation.

(a) T5-Based model

(b) ChatGPT

Figure 3: Entity F1 scores of generators ((a) FiD&Q-TOD and (b) ChatGPT) as
the retrieval performance varies with different retrievers. Bracketed numbers
following model names refer to the correlation coefficients between retrieval
performance and Entity F1.

###  5.4 Ablation Study

We assess the impact of maximum marginal likelihood, various types of meta
knowledge, and the inclusion of negative samples. Unless otherwise specified,
the ablation study is performed on the MWOZ dataset using T5-Base as the
generator, considering the computational resource constraints.

####  5.4.1 Maximum Marginal Likelihood

Table 3 presents the impact of the maximum marginal likelihood (MML) loss. The
methods labeled as “w/o MML” utilize the warmed-up retriever described in
Section 4.2, without the joint training with the response generator. The
results demonstrate that the inclusion of maximum marginal likelihood enables
further enhancement of the retriever during training. Consequently, the
improved retrievers lead to enhanced final generated responses.

Model | Large-scalse
---|---
BLEU | Entity F1 | Recall@7
Oursprefix | 16.39 | 50.35 | 92.51
w/o MML | 16.07 | 49.56 | 91.39
Oursprompt | 17.56 | 50.69 | 92.74
w/o MML | 16.67 | 50.41 | 91.39
Oursctr | 15.96 | 51.35 | 92.74
w/o MML | 14.78 | 50.54 | 91.39

Table 3: Ablation study of the MML loss.

####  5.4.2 Types of Meta Knowledge

Method | Type | Condensed | Large-scale
---|---|---|---
BLEU | Entity F1 | BLEU | Entity F1 | Recall@7
Oursprefix | all | 16.97 | 51.99 | 16.39 | 50.35 | 92.51
order | 16.97 | 51.64 | 16.20 | 49.88 | 92.27
conf | 16.15 | 51.70 | 13.96 | 47.35 | 89.11
cooccur | 16.70 | 51.14 | 15.61 | 49.66 | 91.39
Oursprompt | all | 17.05 | 52.42 | 17.56 | 50.69 | 92.74
order | 16.15 | 49.88 | 15.60 | 49.47 | 91.15
conf | 17.02 | 51.66 | 16.84 | 50.16 | 91.93
cooccur | 16.99 | 51.78 | 16.20 | 50.38 | 92.35

Table 4: Ablation study of different types of meta knowledge on MWOZ with
condensed and large-scale knowledge bases. “order”, “conf”, “cooccur” and
“all” mean using only retrival order, retrieval confidence, co-occurrence, or
all types of meta knowledge, respectively.

We compare different types of meta knowledge, and the results are presented in
Table 4. The findings indicate that using a single type of meta knowledge
yields inferior performance compared to combining all three types.
Furthermore, an interesting observation emerges when using the prefix: the
retrieval order outperforms other types of meta knowledge. In contrast, when
using the prompt, the results are reversed. We attribute this phenomenon to
the design of the prefix and prompt. Representing meta knowledge with a prefix
introduces a higher diversity in ranking order since a distinct prefix is
assigned to each ranking order. This increased diversity enables the generator
to better distinguish the recalled entities. On the other hand, the
distinction between retrieval confidence and co-occurrence in the prefix
setting is less obvious. In contrast, when representing meta knowledge with a
prompt, the retrieval order becomes less diverse, since only the numbers
representing the retrieval order are varied.

####  5.4.3 Negative Samples

Type | Condensed | Large-scalse
---|---|---
BLEU | Entity F1 | BLEU | Entity F1 | Recall@7
Oursprefix (Base) | 16.97 | 51.99 | 16.39 | 50.35 | 92.51
w/o neg | 16.68 | 50.45 | 15.94 | 49.46 | 90.24
Oursprompt (Base) | 17.05 | 52.42 | 17.56 | 50.69 | 92.74
w/o neg | 16.98 | 51.35 | 15.99 | 49.85 | 91.15
Oursctr (Base) | 17.33 | 51.86 | 15.96 | 51.35 | 92.74
w/o neg | 17.11 | 50.34 | 15.79 | 48.32 | 92.29
Oursprefix (LLM) | 7.32 | 32.38 | 6.83 | 30.47 | 92.74
w/o neg | 7.22 | 32.78 | 7.01 | 30.69 | 92.74
Oursprompt (LLM) | 7.29 | 35.98 | 6.97 | 31.88 | 92.74
w/o neg | 7.58 | 36.18 | 7.31 | 32.04 | 92.74

Table 5: Ablation study of negative entities.

We conduct an investigation into the impact of negative entities on the
performance of dialogue systems. The results presented in Table 5 demonstrate
that the inclusion of negative entities significantly improves the performance
of dialogue systems when applied to T5-Base. This performance enhancement can
be attributed to two main factors. Firstly, the presence of negative entities
facilitates easier entity distinction for the generator, enabling it to learn
more effectively. Secondly, the introduction of negative entities aids in
training the retriever through the MML loss in Equation (6). This concept is
somewhat analogous to the motivation behind incorporating negative samples in
knowledge retrieval tasks (Karpukhin et al., 2020).

However, when applied to ChatGPT, negative entities do not contribute to model
performance. The reason is that ChatGPT cannot be fine-tuned, meaning that
solely adding negative entities to the in-context demonstrations does not
effectively teach ChatGPT to differentiate between entities. Consequently, we
opt not to include negative entities when employing our method with ChatGPT.

###  5.5 Behavior of Generator

(a) Retrieval order

(b) Retrieval confidence

Figure 4: The percentage of samples utilizing the entities to generate
responses with respect to (a) retrieval order and (b) retrieval preference.

We examine how the generator utilizes the retrieved knowledge with the
assistance of meta knowledge on the MWOZ test set. For our model and the
baseline, which is T5-Large, we gather all their responses that contain
entities, and analyze the percentage of retrieved entities that appear in the
responses according to retrieval order and confidence. As illustrated in
Figure 4, our generator exhibits a higher propensity than the baseline to
utilize entities with both a high retrieval order and high confidence. Further
more, We have assessed the retrieval results on the MWOZ test set. The
findings demonstrate that our retriever adeptly recalls 80.69% of the gold
entities with a top-1 retrieval order, which directly correlates with the
system-wide performance enhancement. This observation suggests that our
proposed meta knowledge aids the generator in developing an inductive bias to
prioritize entities that are highlighted by the retriever.

##  6 Conclusion

This paper aims to address the retrieval-generation misalignment in end-to-end
task-oriented dialogue systems by introducing maximal marginal likelihood to
train a perceptive retriever that leverages signals from response generation.
To enable the response generator to better distinguish between entities, we
explore several methods for incorporating retrieval-related meta knowledge. We
also propose to incorporate negative entities to enhance the discriminative
capability. Experimental results demonstrate that when combined with meta
knowledge, the response generator effectively leverages high-quality retrieval
knowledge, leading to enhanced quality in the generated responses. Through
analysis, we observe that previous retrieval-augmented generator models suffer
from severe retrieval-generation misalignment, while our method mitigates this
misalignment.

## Limitations

There are three potential limitations of the paper that warrant consideration.
Firstly, the employment of the marginal likelihood method necessitates
computing the likelihood for each retrieved entity, resulting in higher
computational resource requirements compared to solely using negative log-
likelihood (NLL). Secondly, despite conducting various comparisons and
ablation studies in this paper, there are certain aspects missing in our
proposed meta knowledge, such as investigating the combined utilization of
prompt and contrastive learning, as well as exploring the utilization of
retrieval order alongside co-occurrence. Lastly, the theoretical rationale
behind the contribution of our proposed meta knowledge to task-oriented
dialogue (TOD) is not thoroughly discussed.

## Acknowledgements

This work was supported by the National Natural Science Foundation of China
(No. 62176270), the Guangdong Basic and Applied Basic Research Foundation (No.
2023A1515012832), and the Tencent AI Lab Rhino-Bird Focused Research Program.
We thank Ke Yang for his efforts in the preliminary experiments.

## References

  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4171–4186, Minneapolis, Minnesota. Association for Computational Linguistics.
  * Eric et al. (2020) Mihail Eric, Rahul Goel, Shachi Paul, Abhishek Sethi, Sanchit Agarwal, Shuyang Gao, Adarsh Kumar, Anuj Goyal, Peter Ku, and Dilek Hakkani-Tur. 2020.  MultiWOZ 2.1: A consolidated multi-domain dialogue dataset with state corrections and state tracking baselines.  In _Proceedings of the Twelfth Language Resources and Evaluation Conference_ , pages 422–428, Marseille, France. European Language Resources Association.
  * Eric et al. (2017) Mihail Eric, Lakshmi Krishnan, Francois Charette, and Christopher D. Manning. 2017\.  Key-value retrieval networks for task-oriented dialogue.  In _Proceedings of the 18th Annual SIGdial Meeting on Discourse and Dialogue_ , pages 37–49, Saarbrücken, Germany. Association for Computational Linguistics.
  * He et al. (2020a) Zhenhao He, Yuhong He, Qingyao Wu, and Jian Chen. 2020a.  Fg2seq: Effectively encoding knowledge for end-to-end task-oriented dialog.  In _ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 8029–8033. IEEE.
  * He et al. (2020b) Zhenhao He, Jiachun Wang, and Jian Chen. 2020b.  Task-oriented dialog generation with enhanced entity representation. In _INTERSPEECH_ , pages 3905–3909.
  * Henderson et al. (2014) Matthew Henderson, Blaise Thomson, and Steve Young. 2014.  Word-based dialog state tracking with recurrent neural networks.  In _Proceedings of the 15th Annual Meeting of the Special Interest Group on Discourse and Dialogue (SIGDIAL)_ , pages 292–299.
  * Hosseini-Asl et al. (2020) Ehsan Hosseini-Asl, Bryan McCann, Chien-Sheng Wu, Semih Yavuz, and Richard Socher. 2020.  A simple language model for task-oriented dialogue.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 20179–20191. Curran Associates, Inc.
  * Huang et al. (2022) Guanhuan Huang, Xiaojun Quan, and Qifan Wang. 2022.  Autoregressive entity generation for end-to-end task-oriented dialog.  In _Proceedings of the 29th International Conference on Computational Linguistics_ , pages 323–332.
  * Izacard and Grave (2021) Gautier Izacard and Edouard Grave. 2021.  Leveraging passage retrieval with generative models for open domain question answering.  In _Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume_ , pages 874–880, Online. Association for Computational Linguistics.
  * Izacard and Grave (2022) Gautier Izacard and Edouard Grave. 2022.  Distilling knowledge from reader to retriever for question answering.  _arXiv preprint arXiv:2012.04584_.
  * Karpukhin et al. (2020) Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020.  Dense passage retrieval for open-domain question answering.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 6769–6781.
  * 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, Sebastian Riedel, and Douwe Kiela. 2020.  Retrieval-augmented generation for knowledge-intensive nlp tasks.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 9459–9474. Curran Associates, Inc.
  * Madotto et al. (2020) Andrea Madotto, Samuel Cahyawijaya, Genta Indra Winata, Yan Xu, Zihan Liu, Zhaojiang Lin, and Pascale Fung. 2020.  Learning knowledge bases with parameters for task-oriented dialogue systems.  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , pages 2372–2394.
  * Madotto et al. (2018) Andrea Madotto, Chien-Sheng Wu, and Pascale Fung. 2018.  Mem2Seq: Effectively incorporating knowledge bases into end-to-end task-oriented dialog systems.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1468–1478, Melbourne, Australia. Association for Computational Linguistics.
  * OpenAI (2022) OpenAI. 2022.  Chatgpt: Conversational ai language model.  Website.  https://openai.com/chatgpt.
  * Papineni et al. (2002) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002.  Bleu: a method for automatic evaluation of machine translation.  In _Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics_ , pages 311–318.
  * Qin et al. (2019) Libo Qin, Yijia Liu, Wanxiang Che, Haoyang Wen, Yangming Li, and Ting Liu. 2019\.  Entity-consistent end-to-end task-oriented dialogue system with KB retriever.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 133–142, Hong Kong, China. Association for Computational Linguistics.
  * Qin et al. (2020) Libo Qin, Xiao Xu, Wanxiang Che, Yue Zhang, and Ting Liu. 2020.  Dynamic fusion network for multi-domain end-to-end task-oriented dialog.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 6344–6354, Online. Association for Computational Linguistics.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _Journal of Machine Learning Research_ , 21:1–67.
  * Raghu et al. (2021) Dinesh Raghu, Atishya Jain, Sachindra Joshi, et al. 2021.  Constraint based knowledge base distillation in end-to-end task oriented dialogs.  In _Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021_ , pages 5051–5061.
  * Ren et al. (2021) Ruiyang Ren, Yingqi Qu, Jing Liu, Wayne Xin Zhao, QiaoQiao She, Hua Wu, Haifeng Wang, and Ji-Rong Wen. 2021.  RocketQAv2: A joint training method for dense passage retrieval and passage re-ranking.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 2825–2835, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Rony et al. (2022) Md Rashad Al Hasan Rony, Ricardo Usbeck, and Jens Lehmann. 2022.  Dialokg: Knowledge-structure aware task-oriented dialogue generation.  _arXiv preprint arXiv:2204.09149_.
  * Shi et al. (2023) Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettlemoyer, and Wen tau Yih. 2023.  Replug: Retrieval-augmented black-box language models.  _arXiv preprint arXiv:2301.12652_.
  * Singh et al. (2021) Devendra Singh, Siva Reddy, Will Hamilton, Chris Dyer, and Dani Yogatama. 2021.  End-to-end training of multi-document reader and retriever for open-domain question answering.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 25968–25981. Curran Associates, Inc.
  * Tian et al. (2022) Xin Tian, Yingzhan Lin, Mengfei Song, Siqi Bao, Fan Wang, Huang He, Shuqi Sun, and Hua Wu. 2022.  Q-tod: A query-driven task-oriented dialogue system.  _arXiv preprint arXiv:2210.07564_.
  * Wan et al. (2023) Fanqi Wan, Weizhou Shen, Ke Yang, Xiaojun Quan, and Wei Bi. 2023.  Multi-grained knowledge retrieval for end-to-end task-oriented dialog.  In _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 11196–11210, Toronto, Canada. Association for Computational Linguistics.
  * Wen et al. (2017) Tsung-Hsien Wen, David Vandyke, Nikola Mrkšić, Milica Gašić, Lina M. Rojas-Barahona, Pei-Hao Su, Stefan Ultes, and Steve Young. 2017.  A network-based end-to-end trainable task-oriented dialogue system.  In _Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 1, Long Papers_ , pages 438–449, Valencia, Spain. Association for Computational Linguistics.
  * Wu et al. (2019) Chien-Sheng Wu, Richard Socher, and Caiming Xiong. 2019.  Global-to-local memory pointer networks for task-oriented dialogue.  In _International Conference on Learning Representations_.
  * Wu et al. (2022) Jie Wu, Ian G Harris, and Hongzhi Zhao. 2022.  Graphmemdialog: Optimizing end-to-end task-oriented dialog systems using graph memory networks.  In _the AAAI Conference on Artificial Intelligence_ , pages 11504–11512.
  * Xie et al. (2022) Tianbao Xie, Chen Henry Wu, Peng Shi, Ruiqi Zhong, Torsten Scholak, Michihiro Yasunaga, Chien-Sheng Wu, Ming Zhong, Pengcheng Yin, Sida I Wang, et al. 2022\.  Unifiedskg: Unifying and multi-tasking structured knowledge grounding with text-to-text language models.  _arXiv preprint arXiv:2201.05966_.

##  Appendix A Dataset Statistics

We follow the previous work (Raghu et al., 2021) to split the datasets. The
statistics of the three datasets are shown in Table 6.

Dataset | # Dialogues | # Turns
---|---|---
Train/Val/Test | Train/Val/Test
MWOZ (Eric et al., 2020) | 1839/117/141 | 9943/576/711
SMD (Eric et al., 2017) | 2425/302/304 | 6291/777/808
CamRest (Wen et al., 2017) | 406/135/135 | 2095/675/643

Table 6: Statistics of the three datasets.

##  Appendix B Mapping Rules of Meta Knowledge

###  B.1 Prefix of Meta Knowledge

The mapping rules from different forms of meta knowledge to the prefix tokens
are shown in Table 10. We use the same set of prefix for T5 and ChatGPT.
Particularly, for ChatGPT, we design a prompt paraphrase to explain the prefix
to ChatGPT. This explanation prompt is shown below:

‘‘Each record of knowledge base is accompanied by three tags. The first tag
indicates whether this entity appeared before. <new-entity> means this is a
new entity, and <old-entity> means this entity appeared before. The second tag
indicates the authenticity of the third tag. There are three types <low-
confidence>, <mid-confidence> and <high-confidence> indicating low, middle,
high retrieval confidence respectively. A higher retrieval confidence means
the entity is potentially more related to the user goal. The third tag
indicates its importance to the dialogue. <nth-entity> means it is the nth
important entity in the knowledge base, for example, <1th-entity> is the top-1
important and <other-entity> means it is not important.’’

###  B.2 Prompt of Meta Knowledge

The mapping rules from meta knowledge to the prompt are illustrated in Table
11. In the case of ChatGPT, more intricate prompts are devised, and these are
presented in Table 12.

###  B.3 Discussion

How generator learn to distinguish entities with meta knowledge?

Our method aims to enable the models to comprehend meta knowledge, such as
retrieval order and confidence scores, and then compare the retrieved
entities. To achieve this, we employ distinct strategies based on whether the
models are fine-tunable or not. For fine-tunable models like T5, we include
the meta knowledge as a part of the model input, _enabling its assimilation
during the training process_. In the case of non-fine-tunable models like
ChatGPT, our approach _utilizes the model’s existing knowledge to grasp the
essence of meta knowledge_. This is implemented by providing an explanation of
meta knowledge within the system prompt, as elaborated in Appendix B.1 and
Appendix B.2.

How are the prompts of meta knowledge designed?

When it comes to prompts for T5 models, given their fine-tuning capability,
our emphasis lies in the effectiveness of implementation with the pytorch
code, and we leave the comprehension of prompts to the model itself during the
training phrase. To achieve this, we craft the prompts in a manner that
permits them to be tokenized into sequences of uniform length with the
tokenizer, which facilitate the efficient mapping to entities through the
torch.gather operation.

For prompts for ChatGPT, the final prompts are yield thought several trials
and evaluations on the valid set of MWOZ.

##  Appendix C In-context Learning Demonstration

The inputs for ChatGPT to include meta knowledge implemented by the prefix and
prompt approaches are shown in Figure 5 and Figure 6, respectively.

##  Appendix D Pre-training for Retriever

To pre-train the retriever, we employ a distant supervision method. This
involves labeling the entity that exhibits the highest occurrence of attribute
values in both the dialogue context and system response as the pseudo positive
entity. Subsequently, we conduct pre-training of the retriever using in-batch
contrastive learning, considering the positive entities from other examples
within the same mini-batch as negative entities.

##  Appendix E Hyperparameter Settings

The hyperparameters of our system with both condensed and large-scale
knowledge bases are shown in Table 13. We also retrieve different numbers of
entities for different datasets in our experiments, as the details shown in
Table 7.

| Condensed KB | Large-scale KB
---|---|---
MWOZ | CamRest | SMD | MWOZ | CamRest
KB size | 7 | 7 | 8 | 223 | 112
# Retrieved entities for T5-Base | 6 | 6 | 8 | 7 | 7
# Retrieved entities for T5-Large | 5 | 5 | - | 5 | 5

Table 7: The number of retrieved entities under different settings.

##  Appendix F Result of MK-TOD with T5-Large and Large-scale Knowledge Bases

The results of our method built upon T5-Large are shown in Table 8.

Model | MWOZ | CamRest
---|---|---
BLEU | Entity F1 | Recall@5 | BLEU | Entity F1 | Recall@5
Oursprefix (Large) | 16.69 | 53.59 | 84.58 | 27\. 32 | 72.77 | 87.05
Oursprompt (Large) | 17.15 | 52.99 | 88.66 | 26.88 | 72.92 | 92.94
Oursctr (Large) | 17.40 | 53.26 | 93.19 | 27.82 | 71.98 | 92.70

Table 8: Overall results of E2E-TOD systems with large-scale knowledge bases
on MWOZ and CamRest.

##  Appendix G Different Retrievers for Section 5.3

In Section 5.3, we investigated the retrieval-generation misalignment by
introducing several retrievers with different performances. The details of
these retrievers are introduced as follows.

BM25: The BM25 retriever computes the BM25 score between the dialogue context
and each entity.

Frequency: This is a rule-based method. For each entity in the knowledge base,
we compute the number of its attribute values that appear in the context. We
then take the entities with the most attribute values appearing in the
dialogue context as the recalled entities.

Pre-train: This retriever is the pre-trained retriever introduced in Appendix
D.

Ours: This is the retriever introduced in our method (Oursprompt(Base)).

Q-TOD: This is the retriever of Q-TOD.

Oracle: This method uses the condensed knowledge base as the retrieved entity
set. The gold entity, which must appear in the condensed knowledge base, is
marked as the top-1 retrieved entity with high retrieval confidence, while
other entities are marked with low retrieval confidence.

We show the performance (Recall@7) of these retrievers in Table 9.

BM25 | Frequency | Pre-train | Ours | Q-TOD | Oracle
---|---|---|---|---|---
75.51 | 88.66 | 91.39 | 92.74 | 92.97 | 100

Table 9: The performance (Recall@7) of different retrievers for the retrieval-
generation misalignment study.

Meta Knowledge | Prefix
---|---
Retrieval Order
Firstly recalled entity | <1th-entity>
Secondly recalled entity | <2th-entity>
Thirdly recalled entity | <3th-entity>
Fourthly recalled entity | <4th-entity>
Fifthly recalled entity | <5th-entity>
The entities recalled behind the 5th entity and the easy negative entity | <other-entity>
Retrieval Confidence
Entity with retrieval score >= 0.75 | <high-confidence>
Entity with retrieval score < 0.75 and >= 0.25 | <mid-confidence>
Entity with retrieval score < 0.25 and the easy negative entity | <low-confidence>
Co-occurrence Relation
Entity existed in the dialogue context | <old-entity>
Entity not existed in the dialoglue context and the easy negative entity | <new-entity>

Table 10: The mapping rules from different types of meta knowledge to the
prefix token.

Meta Knowledge | Prompt
---|---
Retrieval Order
Firstly recalled entity | ‘‘The top-1 recalled:’’
Secondly recalled entity | ‘‘The top-2 recalled:’’
Thirdly recalled entity | ‘‘The top-3 recalled:’’
Fourthly recalled entity | ‘‘The top-4 recalled:’’
Fifthly recalled entity | ‘‘The top-5 recalled:’’
The entities recalled behind the 5th entity and the easy negative entity | ‘‘The negative entity recalled:’’
Retrieval Confidence
Entity with retrieval score >= 0.75 | ‘‘with high confidence:’’
Entity with retrieval score < 0.75 and >= 0.25 | ‘‘with middle confidence:’’
Entity with retrieval score < 0.25 and the easy negative entity | ‘‘with low confidence:’’
Co-occurrence Relation
Entity existed in the dialogue context | ‘‘existed in history:’’
Entity not existed in the dialoglue context and the easy negative entity | ‘‘newly recalled:’’

Table 11: The mapping rules from different types of meta knowledge to the
prompt for T5.

Meta Knowledge | Prompt
---|---
Retrieval Order
Firstly recalled entity | ‘‘this entity is top-1 important.’’
Secondly recalled entity | ‘‘this entity is top-2 important.’’
Thirdly recalled entity | ‘‘this entity is top-3 important.’’
Fourthly recalled entity | ‘‘this entity is top-4 important.’’
Fifthly recalled entity | ‘‘this entity is top-5 important.’’
The entities recalled behind the 5th entity and the easy negative entity | ‘‘this entity is not important.’’
Retrieval Confidence
Entity with retrieval score >= 0.75 | ‘‘It has high possibility that’’
Entity with retrieval score < 0.75 and >= 0.25 | ‘‘It has medium possibility that’’
Entity with retrieval score < 0.25 and the easy negative entity | ‘‘It has low possibility that’’
Co-occurrence Relation
Entity existed in the dialogue context | ‘‘This entity has appeared before.’’
Entity not existed in the dialoglue context and the easy negative entity | ‘‘This is a new entity.’’

Table 12: The mapping rules from different types of meta knowledge to the
prompt for ChatGPT.

Hyperparameters | Condensed KB | Large-scale KB
---|---|---
T5-Base | T5-Large | T5-Base | T5-Large
Batch size | 2 | 1 | 2 | 1
Gradient accumulation steps | 32 | 64 | 32 | 64
Training gradient steps | 1500
Learning rate schedule | Linear
Retriever learning rate | 1e-4
Response generator learning rate | 1e-4
Gradient weight decay | 0.01
Gradient clipping | 0.01
Retriever max input length | 128
Generator max input length for context | 200
Generator max input length for entity | 100
Max output length | 64
Loss weight $\alpha$ for Eq. 7 | 1
Loss weight $\beta$ for Eq. 7 | 1
Loss weight $\gamma$ for Eq. 11 | 1
Margin $\lambda$ for contrastive learning in Eq. 10 | 0.01

Table 13: Hyperparameter settings of our system. Figure 5: The input prompt
and demonstration for ChatGPT with meta knowledge as the prefix. Figure 6: The
input prompt and demonstration for ChatGPT with meta knowledge as prompt.


#  Resource Optimization for Tail-Based Control in Wireless Networked Control
Systems

Rasika Vijithasena12, Rafaela Scaciota12, Mehdi Bennis1, ,
and Sumudu Samarakoon12  1Centre for Wireless Communication, University of
Oulu, Finland
2Infotech Oulu, University of Oulu, Finland
Email: {rasika.vijithasena, rafaela.scaciotatimoesdasilva, mehdi.bennis,
<EMAIL_ADDRESS>
###### Abstract

Achieving control stability is one of the key design challenges of scalable
Wireless Networked Control Systems (WNCS) under limited communication and
computing resources. This paper explores the use of an alternative control
concept defined as _tail-based control_ , which extends the classical Linear
Quadratic Regulator (LQR) cost function for multiple dynamic control systems
over a shared wireless network. We cast the control of multiple control
systems as a network-wide optimization problem and decouple it in terms of
sensor scheduling, plant state prediction, and control policies. Toward this,
we propose a solution consisting of a scheduling algorithm based on Lyapunov
optimization for sensing, a mechanism based on Gaussian Process Regression
(GPR) for state prediction and uncertainty estimation, and a control policy
based on Reinforcement Learning (RL) to ensure tail-based control stability. A
set of discrete time-invariant mountain car control systems is used to
evaluate the proposed solution and is compared against four variants that use
state-of-the-art scheduling, prediction, and control methods. The experimental
results indicate that the proposed method yields $22\%$ reduction in overall
cost in terms of communication and control resource utilization compared to
state-of-the-art methods.

###### Index Term:

Tail-based control, GPR, WNCS, Lyapunov optimization, RL policy

##  I Introduction

00footnotetext:  This work is funded by the European Union Projects 6G-INTENSE
(GA 101139266), VERGE (GA 101096034) and the project Infotech R2D2. Views and
opinions expressed are however those of the author(s) only and do not
necessarily reflect those of the European Union. Neither the European Union
nor the granting authority can be held responsible for them.

Advances in industrial control systems have spurred extensive research and
innovation in wireless control systems in Industry 4.0 [1]. In particular,
Wireless Networked Control Systems (WNCS), defined as closed-loop control
systems consisting of sensors, actuators, and controllers over wireless
communication networks, are gaining prominence due to their cost-effectiveness
in deployment and maintenance, adaptability in installations, and the
potential to enhance safety. One of the significant challenges in the design
of WNCS is the need for time-criticality that mandates seamless low-latency
and reliability in wireless connectivity, which in turns impacts control
stability [1]. In this view, resource allocation and scheduling of these
control systems are crucial in attaining both latency and reliability targets.

The existing literature on WNCS focuses mainly on meeting latency and
reliability through efficient resource allocation and scheduling [2, 3, 4]. In
[2], the authors utilize a Round Robin (RR) scheduling, which is an
established static scheduling technique. In this approach, every sensor or
controller transmits information consistently. This method cannot stabilize a
large number of control systems due to the absence of scheduling decisions
based on channel state information (CSI). To address this, a dynamic CSI-aware
round-robin scheduling approach has been proposed in [3, 4] so that energy
consumption within the system is minimized while ensuring communication
reliability. Nonetheless, this method does not ensure stability for a
substantial number of control systems, as the waiting time for these systems
to refresh their state/action information increases with the number of served
control systems. As a remedy, the predictive control approach has been
introduced in [5] to ensure stability under scalability with limited
communication resources. Therefore, the freshness of the state information
referred to as Age-of-Information (AoI) is considered for scheduling control
systems, while Gaussian Process Regression (GPR) is used to predict the states
(and actions) for the non-scheduled sensors. To address the increased
computation cost incurs in GPR, the work of [6] opportunistically switches
between high and low complexity prediction models to utilize resources while
ensuring controllability at scale. However, in both [5] and [6], no effort has
been made to find alternative stability criteria that potentially reduce the
total energy consumption.

The main contribution of this paper is to propose a novel solution that
improves the use of communication and control resources in WNCS. Toward this
end, we define a _tail-based_ stability objective that relaxes the strict
stability conditions defined in the classical Linear Quadratic Regulator (LQR)
formulation. Then the joint control and communication in the WNCS is carried
out with three sub-tasks: i) scheduling a control system using the Lyapunov
optimization framework, ii) estimating the states of the unscheduled systems
using GPR, and iii) deriving the best control policy, referred to as _tail-
based control_ , hereinafter, over the proposed stability criteria using
Reinforcement Learning (RL). In practice, the proposed solution is showcased
through the path following of multiple autonomous robots. The robots navigate
within bounded path areas rather than strictly adhering to a desired path.
Action commands are only issued when the robots deviate from these boundaries.
This approach improves cost efficiency by allowing for more relaxed controls
within the specified boundaries. The paper is organized as follows. Sec. II
presents the system model architecture. The joint communication and control
problem is formalized in Sec. III and the proposed solution is derived in Sec.
IV. In Sec. V, numerical evaluations and results are discussed. Finally,
conclusions are drawn in Section VI.

_Notation:_ Scalars are denoted by lowercase symbols and sets are denoted by
calligraphic letters. Boldface lowercase letters denote column vectors, while
boldface uppercase letters denote matrices. The $D\times D$ identity matrix is
denoted by $\mathbf{I}_{D}$ and the spaces of $D\times D$ positive semi-
definite, definite and positive definite matrices are denoted by
$\mathbb{S}^{D}$, $\mathbb{S}^{D}_{+}$ and $\mathbb{S}^{D}_{++}$,
respectively. A multivariate Gaussian distribution with mean $\mathbf{m}$ and
covariance $\mathbf{R}$ is denoted by $\mathcal{N}(\mathbf{m},\mathbf{R})$
while $\operatorname{\mathbb{E}}[\cdot]$ is used for statistical expectation.
The operators $||\cdot||_{p}$, $(\cdot)^{T}$, and $\text{Tr}[.]$ stand for
$p$-norm, transpose, and trace, respectively. The indicator function is
defined as $\mathbb{I}_{\\{a>b\\}}=1$, if $a>b$ is held, and
$\mathbb{I}_{\\{a>b\\}}=0$ otherwise.

##  II System Model

###  II-A Wireless Networked Control Systems

Figure 1: An illustration of $M$ number of WNCS where sensor-controller
communication takes place over a shared wireless network.

We consider a WNCS setting which consists of a set $\mathcal{M}$ of $M$ non-
linear control systems, in which, each system comprises a plant, sensor,
controller, and an actuator along with a shared base station (BS) as presented
in Fig. 1. At each time $k$, in the system $m\in\mathcal{M}$, the sensor
observes the $D$-dimensional plant state
$\mathbf{x}_{m,k}=[{x}_{m,k}^{d}]_{d\in\\{1,\dots,D\\}}\in\mathcal{X}_{m}$
where the state space of a system is $\mathcal{X}_{m}$. Then state is
transmitted to the controller located at the BS over a shared wireless
channel. Based on the received plant state, the controller computes the
$N$-dimensional control action $\mathbf{{u}}_{m,k}\in\mathcal{U}$ where
$\mathcal{U}$ is the action space and transmits to the actuator through a
wired channel to stabilize the control system.

The discrete-time dynamics of the system $m$ is

$\mathbf{{x}}_{m,k+1}=\mathbf{{A}}_{m}\mathbf{{x}}_{m,k}+\mathbf{B}_{m}\mathbf{u}_{m,k}+\mathbf{{w}}_{k},$
(1)

where $\mathbf{{A}}_{m}\in\mathbb{R}^{D\times D}$ and
$\mathbf{{B}}_{m}\in\mathbb{R}^{D\times N}$ are constant matrices that
characterize the dynamics and $\mathbf{{w}}_{k}\in\mathbb{R}^{D}$ is the plant
noise that is sampled from an independent and identically distributed (i.i.d)
Gaussian distribution with zero mean and the covariance of $\mathbf{W}$. We
assume that $\mathbf{{w}}_{k}$ contains the errors that occur due to the
linearization of the non-linear control system. Furthermore,
$\mathbf{{A}}_{m}$ is assumed to be unstable if any eigenvalue has a magnitude
greater than one [7], i.e. $|\lambda_{d}|>1,$ where $d\in\\{1,...,{D}\\}$, and
$\lambda_{1},...,\lambda_{D}$ correspond to the eigenvalues of
$\mathbf{{A}}_{m}$. This suggests that, unless intervened, the plant’s state
will continuously increase (become unstable) over time.

###  II-B Wireless communication model

The plant’s state is transmitted to the controller located at the BS from the
sensor over a wireless communication channel, referred to as the uplink (UL),
which is presumed to be a Rayleigh block fading channel. This channel changes
independently over time and remains static over time. In UL, the received
signal $\mathbf{y}_{m,k}$ in BS for control system $m$ and time $k$ is

$\mathbf{y}_{m,k}=\sqrt{p_{m,k}}\mathbf{\Omega}_{m}\mathbf{H}_{m,k}\mathbf{x}_{m,k}+\mathbf{n}_{m,k},$
(2)

where $p_{m,k}\in[0,P_{\text{max}}]$ is the transmission power of each control
system with a maximum power $P_{\text{max}}$,
$\mathbf{H}_{m,k}\in\mathbb{R}^{D\times D}$ is the wireless channel of the UL
communication of each control system,
$\mathbf{\Omega}_{m}\in\mathbb{R}^{D\times D}$ is the observation matrix at
the UL state communication that characterizes the state-observation, and
$\mathbf{n}_{m,k}$ is the additive white Gaussian noise with zero-mean and
variance of $N_{0}$. Thus, for a given bandwidth $\omega$, the UL signal-to-
noise ratio (SNR) is given by

$\gamma_{m,k}={p_{m,k}\|\mathbf{H}_{m,k}\|^{2}}/{N_{0}\omega}.$ (3)

For a successful UL communication, $\gamma_{m,k}\geq\gamma_{0}$ needs to be
satisfied with a given threshold $\gamma_{0}$.

###  II-C Scheduling UL Communication

Since all controller units are co-located at the BS and share the UL
communication channel, a scheduler deployed at the BS is utilized to enable
interference-free UL communication. Towards this, the BS uses a scheduling
indicator $a_{m,k}$ where $a_{m,k}=1$ represents the channel is assigned for
the UL of system $m$ at time $k$ and $a_{m,k}=0$ otherwise. When $a_{m,k}=1$,
the minimum mean square error (MMSE) estimator is used to decode the received
signal as

$\tilde{\mathbf{x}}_{m,k}=\mathbb{E}[\mathbf{x}_{m,k}|\mathbf{y}_{m,k}]=\mathbf{G}_{m,k}\mathbf{y}_{m,k}=\mathbf{x}_{m,k}+\mathbf{v}_{m,k},$
(4)

where
$\mathbf{G}_{m,k}=(\sqrt{p_{m,k}}\bm{S}_{m,k}\mathbf{H}^{T}_{m,k}\mathbf{\Omega}_{m}^{T})/(p_{m,k}\mathbf{\Omega}_{m}\mathbf{H}_{m,k}\bm{S}_{m,k}\\\
\mathbf{H}_{m,k}^{T}\mathbf{\Omega}_{m}^{T}+N_{0}\mathbf{I}_{D})$ is the
linear MMSE matrix of control system $m$ and
$\bm{S}_{m,k}=\mathbb{E}[\mathbf{x}_{k}\mathbf{x}_{k}^{T}]$ is the covariance
matrix. Here, $\mathbf{v}_{m,k}$ is the MMSE estimation error that is assumed
to be a Gaussian random vector with zero mean and $\mathbf{V}_{m,k}$
covariance matrix given by

$\mathbf{V}_{m,k}=\mathbb{E}[\mathbf{v}_{m,k}\mathbf{v}_{m,k}^{T}]=\bm{S}_{m,k}-\mathbf{G}_{m,k}\sqrt{p_{m,k}}\mathbf{\Omega}_{m}\mathbf{H}_{m,k}\bm{S}_{m,k}.$
(5)

On other hand, if $a_{m,k}=0$, the controller $m$ relies on an estimate
$\hat{\mathbf{x}}_{m,k}$ of the state to compute the control action. While
opportunistically scheduling a control system, it is highly likely that
control decisions derived from estimated and outdated states may become
inaccurate, in which, ensuring stability could be challenging. In this view,
the AoI can be used to measure the freshness of the states considering the
time interval between two consecutive state information updates. AoI linearly
increases with time and can be represented as

$\beta_{m,k}=\begin{cases}1,&\text{if\leavevmode\nobreak\
}a_{m,k}=1\text{\leavevmode\nobreak\ and\leavevmode\nobreak\
}\gamma_{m,k}\geq\gamma_{0},\\\ 1+\beta_{m,k-1},&\text{otherwise}.\end{cases}$
(6)

The controller computes the control action based on the estimated state or
predicted state considering AoI and transmits it to the actuator. We consider
downlink (DL) to be ideal, leading to the assumption that the actuator
receives control decisions without any errors. Thereafter, the actuator
applies these error-free control decisions on the plant to maintain stability.

##  III Joint Communication and Control Problem

Our objective of this research is to optimize the resources utilized for
communication and control while ensuring the control stability and
communication reliability.

###  III-A Communication Cost

The increase in AoI reflects fewer communications, which is beneficial in
terms of energy savings, but has an adverse effect on the estimated state with
increased uncertainty. Since it can cause the system to be unstable, it is
essential to strike a balance between AoI measures over all systems. At the
same time, transmission power impacts the reliability of communication and the
amount of energy consumed, and thus, power allocation for communication with
all systems should be treated fairly. In this view, the total communication
cost ${C}$ can be expressed as follows:

${C}\left([\bar{\beta}_{m},\bar{\hat{p}}_{m}]_{m\in\mathcal{M}}\right)=\psi_{\beta}\sum_{m=1}^{M}\mathcal{G}(\bar{\beta}_{m})\\\
+\psi_{p}\sum_{m=1}^{M}\mathcal{G}(\bar{\hat{p}}_{m}),$ (7)

where $\hat{p}_{m,k}=a_{m,k}{p}_{m,k}$ and $\mathcal{G}(.)=\log(1+.)$ is a
non-decreasing concave function that captures the fairness of AoI and
transmission power for each control system [8]. Hereinafter, the notation of
$\bar{z}=\frac{1}{K}\sum_{k=1}^{K}z_{k}$ is used for the time average value of
any quantity $z_{k}$. The positive coefficients $\psi_{\beta}$ and $\psi_{p}$
can be defined to reflect the relative importance of the costs related to AoI
and the transmission power, respectively.

###  III-B Control Cost

Traditional LQR-based control designs are rooted on the stability defined over
a singular point of operation and thus, the optimal control always ensures
that the plant state is always maintained around the desired state. This
design is inefficient for systems that consider extreme state conditions (tail
conditions) as undesirable. Towards this, the objective of tail-based control
is to define a state range for the desired stability, in contrast to the
classical stability conditions.

In this view, for any system $m$, suppose that the desired range of stability
$\|\mathbf{x}_{m,k}\|_{p}>\eta$ at time $k$ is defined by the threshold
$\eta(>0)$. Hence, the deviations from stability and the effort spent on
control is used to define the control cost $J_{m,k}$ for a given control
system $m$ at time $k$ as follows:

${J}_{m,k}=\underbrace{\mathbf{x}^{T}_{m,k}\mathbf{Q}\mathbf{x}_{m,k}f_{m,k}}_{\text{Cost
related to
stability}}+\underbrace{\mathbf{u}^{T}_{m,k}\mathbf{Y}\mathbf{u}_{m,k}}_{\text{Cost
related to controlling}},$ (8)

where the indicator
$f_{m,k}=\mathbb{I}_{\\{\|\mathbf{x}_{m,k}\|_{p}>\eta\\}}$. Here,
$\mathbf{Q}\in\mathbb{S}^{D}$ and $\mathbf{Y}\in\mathbb{R}$, which correspond
to control system design parameters.

###  III-C Control-Constrained Optimization Problem

We use the quadratic Lyapunov function to evaluate control stability, which
measures the performance of each control system as a function of the state as
follows:

$\mathcal{L}(\mathbf{x}_{m,k})=\mathbf{x}^{T}_{m,k}\,\mathbf{Z}\,\mathbf{x}_{m,k}f_{m,k},$
(9)

where $\mathbf{Z}\in\mathbb{S}^{D}_{++}$ is a unique solution to the discrete
Lyapunov equation
$\mathbf{{A}}^{c^{T}}_{m}\mathbf{Z}+\mathbf{Z}\mathbf{A}^{c}_{m}=-\mathbf{I}_{D}$,
and $\mathbf{{A}}^{c}_{m}$ is a closed-loop state transition matrix defined as
$\mathbf{A}^{c}_{m}=\mathbf{A}_{m}-{\mathbf{B}}_{m}\mathbf{\Phi}_{m}$. Here,
$\mathbf{\Phi}_{m}$ is the feedback gain matrix of the control system $m$,
which is given by,

$\mathbf{\Phi}_{m}={\arg\min}_{\mathbf{\Phi}^{\prime}_{m}}\frac{1}{K}\sum_{k=1}^{K}(\mathbf{u}_{m,k}-\mathbf{\Phi}^{\prime}_{m}\mathbf{x}_{m,k}).$
(10)

However, the centralized scheduler only has access to the predicted state
$\hat{\mathbf{x}}_{m,k}$. Therefore, the expected Lyapunov value is calculated
as [5],

$\mathbb{E}\left[\mathcal{L}(\mathbf{x}_{m,k})|\hat{\mathbf{x}}_{m,k}\right]=\hat{\mathbf{x}}^{T}_{m,k}\,\mathbf{Z}\,\hat{\mathbf{x}}_{m,k}+\text{Tr}\left[\mathbf{Z}\,\mathbf{\Psi}_{m,k}\right],$
(11)

where $\mathbf{\Psi}_{m,k}$ is the covariance matrix of the state prediction
error, which is discussed in IV-B. For control stability, it is required that
the expected future value of the Lyapunov function
$\mathcal{L}(\mathbf{x}_{m,k+1})$, should be less than the current value of
$\mathcal{L}(\mathbf{x}_{m,k})$ scaled by a certain rate $\zeta_{m}\in(0,1]$
as

$\mathbb{E}\left[\mathcal{L}(\mathbf{x}_{m,k+1})|\hat{\mathbf{x}}_{m,k},\mathbf{u}_{m,k},\mathbf{H}_{m,k},p_{m,k}\right]\\\
\leq\zeta_{m}\mathbb{E}\left[\mathcal{L}(\mathbf{x}_{m,k})|\hat{\mathbf{x}}_{m,k}\right],$
(12)

where the expectation in right hand side of (12) is with respect to the plant
noise $\mathbf{w}_{k}$ in (1) and the signal estimation error
$\mathbf{v}_{m,k}$ defined in (4).

Based on the aforementioned costs of communication and control along the
stability and UL channel sharing (one system is scheduled at a time), the
formulation of the control-constrained optimization problem is given by,

$\displaystyle\underset{(\mathbf{a}_{k},\mathbf{p}_{k},\mathbf{u}_{k}),\forall
k}{\operatorname*{minimize}}$
$\displaystyle\quad\frac{1}{M}\sum_{m=1}^{M}{C}\left(\bar{\beta}_{m},\bar{\hat{p}}_{m}\right)+\frac{1}{M}\sum_{m=1}^{M}\bar{J}_{m}$
(13a) s.t. $\displaystyle\quad\gamma_{m,k}\geq a_{m,k}\gamma_{0},\quad\forall
m\in\mathcal{M},\forall k,$ (13b)
$\displaystyle\quad\mathbf{a}_{k}\in\mathcal{A},\mathbf{p}_{k}\in\mathcal{P},\mathbf{u}_{k}\in\mathcal{U},\quad\forall
k,$ (13c) $\displaystyle\quad\eqref{Future_Lyapunov},\quad\forall
m\in\mathcal{M},\forall k,$ (13d)

where
$\mathcal{A}=\\{\mathbf{a}_{k}=[a_{m,k}]_{m\in\mathcal{M}}|\sum_{m}a_{m,k}\leq
1,a_{m,k}\in\\{0,1\\}\\}$ and
$\mathcal{P}=\\{\mathbf{p}_{k}=[p_{m,k}]_{m\in\mathcal{M}}|p_{m,k}\in[0,P_{\text{max}}\\}$.

##  IV Communication and Control Co-Design

Due to the complex nature of the joint control and communication problem posed
in (13), we propose a solution method consisting of three components: an UL
scheduling algorithm, a GPR-based prediction model for state prediction, and a
control policy to satisfy the tail-based stability.

###  IV-A Lyapunov Optimization for UL Scheduling

In the stability constraint (12), the decisions related to state prediction,
UL scheduling, and tail-based control are coupled, in which, solving (13)
becomes a daunting task. Hence, (12) can be recast in an alternative form as
follows [5, Sec. IIIA]:

${\lim\sup_{K\to\infty}}\frac{1}{K}\sum_{k=1}^{K}\xi_{m,k}\geq{\lim\sup_{K\to\infty}}\frac{1}{K}\sum_{k=1}^{K}\frac{\Gamma_{m,k}}{\Upsilon_{m,k}}$
(14)

where $\xi_{m,k}=a_{m,k}\mathbb{I}_{\\{\gamma_{m,k}>\gamma_{0}\\}}$, and

$\displaystyle\Upsilon_{m,k}=$
$\displaystyle\begin{multlined}\mathrm{Tr}[(\mathbf{B}_{m}\mathbf{\Phi}_{m})^{T}\mathbf{Z}(\mathbf{B}_{m}\mathbf{\Phi}_{m})\mathbf{\Psi}_{m,k}]\\\
-\mathrm{Tr}[(\mathbf{B}_{m}\mathbf{\Phi}_{m})^{T}\mathbf{Z}(\mathbf{B}_{m}\mathbf{\Phi}_{m})\mathbf{V}_{m,k}],\end{multlined}\mathrm{Tr}[(\mathbf{B}_{m}\mathbf{\Phi}_{m})^{T}\mathbf{Z}(\mathbf{B}_{m}\mathbf{\Phi}_{m})\mathbf{\Psi}_{m,k}]\\\
-\mathrm{Tr}[(\mathbf{B}_{m}\mathbf{\Phi}_{m})^{T}\mathbf{Z}(\mathbf{B}_{m}\mathbf{\Phi}_{m})\mathbf{V}_{m,k}],$
$\displaystyle\Gamma_{m,k}=$
$\displaystyle\begin{multlined}\left\|(\mathbf{A}_{m}^{c}-\zeta_{m}\mathbf{I}_{D})_{m,k}\right\|^{2}_{\mathbf{Z}^{\frac{1}{2}}}+\mathrm{Tr}[\mathbf{B}_{m}^{T}\mathbf{Z}\mathbf{B}_{m}\mathbf{\Psi}_{m,k}]\\\
+\mathrm{Tr}\left[(\mathbf{A}_{m}^{T}\mathbf{Z}\mathbf{A}_{m}-\zeta_{m}\mathbf{Z})\mathbf{\Psi}_{m,k}\right]+\mathrm{Tr}[\mathbf{Z}\mathbf{W}].\end{multlined}\left\|(\mathbf{A}_{m}^{c}-\zeta_{m}\mathbf{I}_{D})_{m,k}\right\|^{2}_{\mathbf{Z}^{\frac{1}{2}}}+\mathrm{Tr}[\mathbf{B}_{m}^{T}\mathbf{Z}\mathbf{B}_{m}\mathbf{\Psi}_{m,k}]\\\
+\mathrm{Tr}\left[(\mathbf{A}_{m}^{T}\mathbf{Z}\mathbf{A}_{m}-\zeta_{m}\mathbf{Z})\mathbf{\Psi}_{m,k}\right]+\mathrm{Tr}[\mathbf{Z}\mathbf{W}].$

Due to the decoupling of scheduling, prediction, and control variables as
(14), we can now minimize the communication cost $C$ independently given the
control policy as follows:

$\displaystyle\underset{(\mathbf{a}_{k},\mathbf{p}_{k}),\forall
k}{\operatorname*{minimize}}\quad$
$\displaystyle\frac{1}{M}\sum_{m=1}^{M}{C}\left(\bar{\beta}_{m},\bar{\hat{p}}_{m}\right)$
(15a) s.t. $\displaystyle\bar{a}_{m}\geq\bar{{G}_{l}}({c}_{m,k}),\quad\forall
m\in\mathcal{M},$ (15c) $\displaystyle\eqref{eqopt1e_1a},\leavevmode\nobreak\
\eqref{eqopt1e_1d},$

where ${c}_{m,k}=\Gamma_{m,k}/\Upsilon_{m,k}$ is the lower-bound stability of
$\xi_{m,k}$ of each control system. Here,
${G}_{l}(\cdot)=\max\left[\min(\cdot,1),0\right]]$ is used as the lower bound
of (15c) to ensure the feasibility of the scheduling constraints
$\sum_{m=1}^{M}a_{m,k}\leq 1$ and $a_{m,k}\in\\{0,1\\}$ for all $k$. To solve
the optimization over a time horizon, we resort to the stochastic Lyapunov
optimization framework [9]. Therein, the notion of virtual queues is used to
track the dynamics of the time average variables, allowing us to decompose
(15) into a series of sub-problems that can be solved at each time step $k$.
The details of the solution are similar to the work of [5] and left out due to
the space limitations. We direct the interested readers to [5, Sec.IV] for the
detailed derivations.

###  IV-B Prediction of Missing States using GPR

In the absence of UL, either due to poor SNR or not being scheduled,
controllers are required to predict the state of their corresponding systems
using previous observations. For this, we employ a GPR-based state prediction
technique in each controller. Since the system state dynamics results in a
time series, the dynamics can be modeled as sampling from a Gaussian process.

Let
$\mathbf{X}_{m,\bm{k^{\prime}}}^{d}=[\tilde{x}_{m,k^{\prime}}^{d}]_{k^{\prime}\in\mathcal{K}}$
be past observations where
$\mathcal{K}=\\{\kappa|\kappa<k,a_{m,\kappa}=1,\gamma_{m.\kappa}\geq\gamma_{0}\\}$
with the corresponding time steps
$\bm{k}^{\prime}=[k^{\prime}]_{k^{\prime}\in\mathcal{K}}$. Then the latent
function of the regression model over the dimension $d$ can be learned as
$\tilde{x}_{m,k^{\prime}}^{d}=g_{d}(k^{\prime})+\tau$ with a noise
$\tau\sim\mathcal{N}\left({0},\sigma^{2}_{\kappa}\right)$. For estimating
$g_{d}$, first, it is necessary to rely on a correlation between the state
observations over different time steps. In this view, we consider a periodic
kernel to model the correlation between the outputs according to their
temporal behaviors as [10],

${R}(k,k^{\prime})=\textstyle
h^{2}\exp\left(-\frac{2}{l^{2}}\sin^{2}\left(\frac{\pi\left(k-k^{\prime}\right)}{s}\right)\right),$
(16)

where ${R}(k,k^{\prime})$ represents the covariance between the outputs
according to times $k$ and $k^{\prime}$. Here, $h$, $l$, $s$ are the hyper-
parameters that represent output scaling, time scaling, and frequency,
respectively, which are denoted as $\mathbf{\Theta}=(h,l,s)$, hereinafter.
Then, the joint distribution of past observations and the output of $g_{d}(k)$
follows:

$\left[\begin{array}[]{c}\mathbf{X}^{d}_{m,\bm{k}^{\prime}}\\\
g_{d}(k)\end{array}\right]\sim\mathcal{N}\left(\mathbf{0},\left[\begin{array}[]{cc}\mathbf{R}(\bm{k}^{\prime})&\mathbf{r}(\bm{k}^{\prime},k)\\\
\mathbf{r}(k,\bm{k}^{\prime})&\mathbf{r}(k,k)\end{array}\right]\right),$ (17)

where $\mathbf{R}(\bm{k}^{\prime})=[R(i,j)]_{i,j\in\mathcal{K}}$,
$\mathbf{r}(k,\bm{k}^{\prime})=[R(k,j)]_{j\in\mathcal{K}}$,
$\mathbf{r}(\bm{k}^{\prime},k)=\mathbf{r}(k,\bm{k}^{\prime})^{T}$ and
$K=|\mathcal{K}|$. Then, the posterior distribution of $g_{d}(k)$ based on the
past observations $\mathbf{X}^{d}_{m,\bm{k^{\prime}}}$ can be analytically
determined as,

$\text{Pr}\left(g_{d}(k)|\mathbf{X}^{d}_{m,\bm{k^{\prime}}},k,{\mathbf{\Theta}}\right)\sim\mathcal{N}\left(\hat{x}^{d}_{m,k},(\sigma^{d}_{m,k})^{2}\right).$
(18)

Here,
$\sigma_{m,k}^{d}=\sqrt{({R}(k,k)-\mathbf{r}(k,\bm{k}^{\prime})\mathbf{R}(\bm{k}^{\prime})^{-1}\mathbf{r}(k,\bm{k}^{\prime})^{T})}$
is the standard deviation and the predicted state over the dimension $d$ is
given by

$\hat{x}_{m,k}^{d}=\mathbf{r}(k,\bm{k}^{\prime})^{T}\mathbf{R}(\bm{k}^{\prime})^{-1}\mathbf{X}^{d}_{m,\bm{k}^{\prime}}.$
(19)

Finally, the prediction error covariance matrix is calculated as

$\mathbf{\Psi}_{m,k}=\operatorname*{diag}\left(\sigma_{m,k}^{2}(1),\dots,\sigma^{2}_{m,k}(D)\right).$
(20)

###  IV-C Tail-based Control

The control cost minimization can be represented as a series of sub-problems,
in which, for system $m$ it is as follows:

$\displaystyle\operatorname*{minimize}_{(\mathbf{u}_{m,k}\in\mathcal{U}),\forall
k}\quad$ $\displaystyle\frac{1}{M}\sum_{m=1}^{M}\bar{J}_{m}.$ (21)

Due to the inherent complexities and potential non-linearities of tail-based
control approach, it is a challenge to find a close form expression for the
control policy by using either convex optimization, dynamic programming, or
numerical methods. In contrast, it is efficient at optimizing a non-linear,
complex cost function by refining the policy in an iterative manner through
trial and error. Hence, the derivation of the optimal control policy
$\bm{\pi}^{*}$, the probabilistic mapping between plant states and control
actions, has been cast as a solving Markov Decision Process (MDP) using RL.
The MDP problem is formulated as a tuple
$(\mathcal{X}_{m},\mathcal{U},F,{R_{e}},\nu)$ where $F$ defines the dynamics
as per (1) and $\nu$ is the discount factor. Considering the tail-based
control cost $J_{m}$ in (8), The reward
${R_{e}}(\mathbf{x}_{m,k},\mathbf{u}_{m,k})$ is generated from the tail-based
control cost $J_{m,k}$ in (8) as follows:

${R_{e}}(\mathbf{x}_{m,k},\mathbf{u}_{m,k})=-J_{m,k}.$ (22)

In the above MDP, the optimal policy $\bm{\pi}^{*}$, which maximizes the
discounted cumulative rewards can be represented as

$\bm{\pi}^{*}=\max_{\bm{\pi}}\left[\sum_{k=1}^{\infty}\nu^{k}{R_{e}}(\mathbf{x}_{m,k},\mathbf{u}_{m,k}|\bm{\pi})\right].$
(23)

Afterwards, the optimal control decision is derived by sampling the control
policy given the state, i.e.,
$\mathbf{u}_{m,k}={B}_{u}(\mathbf{x}_{m,k};\bm{\pi})$ where $B_{u}$ represents
the function in our RL model. It maps the current state of the system,
$\mathbf{x}_{m,k}$, to an optimal control action, $\mathbf{u}_{m,k}$, using
the policy $\bm{\pi}$ that has been learned through the RL process.

###  IV-D Proposed Algorithm

Due to the decoupling of UL scheduling, state prediction, and control policy.
we implement the overall solution in two-step. In the first step, the RL-based
policy (23) is trained over the reward function (22) for a generic system.
During the second step, the feedback gain matrix is computed using (10) and
used for UL scheduling followed by state prediction. Finally, the control
decisions are computed using either received or estimated state. These steps
are presented in Algorithm 1.

1: Step 1: Offline training of the tail-based control policy

2: Initialization: $\mathcal{X},\mathcal{U},\nu$, $\mathbf{x}_{0}$,
$\mathbb{\eta}$, and $\bm{\pi}$

3: for each epoch do

4: for each step $k$ do

5: Sample action $\mathbf{u}_{k}={B}(\mathbf{x}_{k};\bm{\pi})$

6: Calculate reward ${R_{e}}$ using (22)

7: Update state $\mathbf{x}_{k+1}$ using (1)

8: Update policy $\bm{\pi}$ using (23)

9: end for

110: end for

10: Step 2: Online decision making at the controller

11: Initialization: $V,\psi_{\beta},\psi_{p}$, and $\mathbf{x}_{m,0}$

12: for $k$ =1 to $K$ do

13: for ${m}$ = 1 to $M$ do

14: Compute the feedback gain matrix $\mathbf{\Phi}$ by (10)

15: Calculate $\hat{\mathbf{x}}_{m,k}$ with (19) and $\mathbf{\Psi}_{m,k}$
using (20)

16: Find $a_{m,k}$ by solving (15)

17: if $a_{m,k}=1$ then

18: Schedule system $m$ and decode $\tilde{\mathbf{x}}_{m,k}$ using (4)

19: else

20: Predict $\hat{\mathbf{x}}_{m,k}$ using (19)

21: end if

22: Apply control action $\mathbf{u}_{m,k}={B_{u}}(\mathbf{x}_{m,k};\bm{\pi})$

23: Update $\mathbf{x}_{m,k+1}$ based on $\mathbf{u}_{m,k}$ using (1)

24: end for

25: end for

Algorithm 1 Communication and Control Co-Design

##  V Numerical Results

In this section, we present numerical results to validate our theoretical
analysis. We employ the Moore’s mountain car problem [11] to evaluate the
performance of the proposed algorithm. The discretized state space dynamics
that use scalar control actions with a sampling period $T=0.01$ s and the
equilibrium point $\epsilon=[\frac{\pi}{2b},0]$ results in

$\mathbf{A}_{m}=\begin{bmatrix}(1+\alpha{b})&1\\\
\alpha{b}&1\end{bmatrix},\quad\mathbf{B}_{m}=\begin{bmatrix}1\\\ 1\\\
\end{bmatrix},\quad\text{and}\quad\mathbf{Y}=1,$ (24)

with $\alpha$ as the gravity. We use the following as the default simulation
parameters $\mathbf{Q}=\mathbf{I}$, $\mathbf{x}_{m,0}=[-1.5,0]$, $b=3$,
$\alpha=0.0025$ Nkg-1, $\mathbf{\Omega}=\mathbf{I}$, $P_{max}=30$ dBm,
$N_{0}=1$, $V=1000$, $\zeta_{i}=0.1$, $\psi_{\beta}=1$, $\psi_{p}=1$,
$\gamma_{0}=20$ dB, $\sigma_{m}^{2}=0.02$, and $p_{m,k}=28.2$ dBm unless
specified otherwise. For the GPR, $10$ past state observations have been used.
For each simulation, $1000$ time steps have been carried out. Furthermore, all
analyses are based on averaging data collected over $100$ simulations.

The RL model is developed using the Python MushroomRL library [12], which
provides common interfaces to develop RL models. A Gaussian policy with a
learnable standard deviation has been considered to train the policy. The RL
model is trained for $100$ epochs, each with $200$ episodes with
$\eta=[0.1,0.1]$, $\mathcal{X}=[-\infty,+\infty],\mathcal{U}=[-10,10]$, and
$\nu=0.9$.

Since the proposed solution consists of three parts: a novel Control-Aware
Scheduler (CA), the adoption of GPR and a RL-based control policy, the
corresponding concepts of the existing solutions, the RR scheduler, without
predictions or Autoregressive Integrated Moving Average (ARIMA) [13], and the
classical LQR control policy, respectively, are used for performance
comparison. In this regard, we compare the proposed solution Full with four
other variants summarized in Table I.

Table I: Methods compared in simulations Method | Scheduling | State prediction | Control |
---|---|---|---|---
V1 | RR | No | Tail | $\left.\begin{array}[]{l}\\\ \\\ \\\ \\\ \end{array}\right\\}\text{Variants}$
V2 | RR | ARIMA | Tail
V3 | RR | ARIMA | Classic
V4 | CA | GPR | Classic
Full | CA | GPR | Tail |
Figure 2: RL convergence of the tail-based control and LQR problems.

First, we analyze the convergence of the trained RL policy via simulations for
both the proposed tail-based control and LQR formulations as illustrated in
Fig. 2. Though there exists a closed form solution for LQR formulation, we
have used RL for deriving both policies for a fair comparison. It can be noted
from Fig. 2 that both policies achieve convergence with respective to their
own objectives after about $70$ epochs.

In Fig. 3, we compare the total cost that is the combination of communication
cost (7) and control cost (8), the cost related to controlling, and the cost
related to stability for all methods with settings having different number of
control systems. Specifically, in Fig. 3(a), we can observe that V1, exhibits
lower costs compared to other three variants while the proposed solution Full
yields the lowest cost. In summary, the reductions of total cost with Full
compared to V1 are $15\%$ at $M=6$ while $22\%$ at $M=21$.

Fig. 3(b) shows that for $M=6$, V1 and V2 have costs related to controlling
$4.3$ and $2.8$ times lower than the Full method, respectively. However, as
the number of systems scales to $M=21$, V2 incurs higher costs, while V1
exhibits slightly higher costs than Full. This is due to the fact that the
ARIMA method tends to produce inaccurate state estimates leading to system
instability as well as to apply control decisions with higher magnitude. On
the contrary, V1, which lacks a prediction method, avoids applying control
decisions when systems are unscheduled, thus, avoiding the undesirable state
fluctuations evident in V2. Meanwhile, the classic control policy employed in
V3 and V4 is focused on imposing strict stability, which is impractical when
control decisions are applied inconsistently, thus, yielding increased costs.
Furthermore, we can observe in Fig. 3(c), the costs related to stability are
comparatively much higher (about $10^{6}$ times) than the costs related to
controlling, which are dominant in the overall cost computation. As a result,
the total costs in Fig. 3(a) and the costs related to stability in Fig. 3(c)
show similar trends. Finally, we would like to point out that certain
instances showcase fluctuations in cost that disagree with trends, which can
be rectified by augmenting the number of simulations.

(a) Objective value as per (13a).

(b) Cost related to controlling.

(c) Cost related to stability.

Figure 3: Performance of total cost, cost related to controlling, and cost
related to stability per control system for different $M$.

##  VI Conclusion

This paper proposes a novel approach to improve the utilization of
communication and control resources in WNCS. To this end, a new notion of
stability, referred to as tail-based stability, deals with extreme conditions
relaxing the classical stability conditions. Then, joint control and
communication are carried out in three stages: (i) scheduling control systems
using the Lyapunov optimization framework, (ii) predicting missing plant
states with GPR, and (iii) deriving the control policy using RL. The results
show that by combining the above three mechanisms, the proposed method can
outperform designs that rely on different state-of-the-art scheduling,
prediction, and controlling methods. Simultaneously controlling heterogeneous
systems and extending the analysis for different propagation environments and
multi-antenna systems along with the ablation studies via varying system
parameters are left as future directions.

## References

  * [1] A. Varghese _et al._ , “Wireless requirements and challenges in Industry 4.0,” in _2014 International Conference on Contemporary Computing and Informatics (IC3I)_ , 2014, pp. 634–638.
  * [2] L. Schenato _et al._ , “Foundations of control and estimation over lossy networks,” _Proceedings of the IEEE_ , vol. 95, no. 1, pp. 163–187, 2007.
  * [3] Y. Xu _et al._ , “Stability analysis of networked control systems with round-robin scheduling and packet dropouts,” _Journal of the Franklin Institute_ , vol. 350, no. 8, 2013.
  * [4] X. Liu _et al._ , “A framework for opportunistic scheduling in wireless networks,” _Computer networks_ , vol. 41, no. 4, pp. 451–474, 2003.
  * [5] A. M. Girgis _et al._ , “Predictive control and communication co-design via two-way gaussian process regression and aoi-aware scheduling,” _IEEE Transactions on Communications_ , vol. 69, no. 10, pp. 7077–7093, 2021.
  * [6] K. Ranasinghe, “Predictive over-the-air sensing and controlling under limited communication and computational resources,” Master’s thesis, K. Ranasinghe, 2022.
  * [7] K. Ogata, _Discrete-time control systems_. Prentice-Hall, Inc., 1995.
  * [8] L. Li _et al._ , “Proportional fairness in multi-rate wireless lans,” in _The 27th Conference on Computer Communications_. IEEE, 2008, pp. 1004–1012.
  * [9] M. J. Neely, “Queue stability and probability 1 convergence via lyapunov optimization,” _arXiv preprint arXiv:1008.3519_ , 2010.
  * [10] C. E. Rasmussen _et al._ , “Gaussian processes for machine learning,” 2006.
  * [11] A. W. Moore, “Efficient memory-based learning for robot control,” University of Cambridge, Computer Laboratory, Tech. Rep., 1990.
  * [12] T. M. Team, “Mushroomrl documentation,” 2023. [Online]. Available: https://mushroomrl.readthedocs.io/en/latest/index.html
  * [13] P. J. Brockwell _et al._ , _Introduction to time series and forecasting_. Springer, 2002.