# KIX: A Metacognitive Generalization Framework

Arun Kumar
Dept. of Computer Science
University of Minnesota, Twin Cities
Minneapolis, MN, USA
<EMAIL_ADDRESS>
&Paul Schrater
Depts. of Computer Science and Psychology
University of Minnesota, Twin Cities
Minneapolis, MN, USA
<EMAIL_ADDRESS>

###### Abstract

Humans and other animals aptly exhibit general intelligence behaviors in
solving a variety of tasks with flexibility and ability to adapt to novel
situations by reusing and applying high level knowledge acquired over time.
But artificial agents are more of a specialist, lacking such generalist
behaviors. Artificial agents will require understanding and exploiting
critical structured knowledge representations. We present a metacognitive
generalization framework, Knowledge-Interaction-eXecution (KIX), and argue
that interactions with objects leveraging type space facilitate the learning
of transferable interaction concepts and generalization. It is a natural way
of integrating knowledge into reinforcement learning and promising to act as
an enabler for autonomous and generalist behaviors in artificial intelligence
systems.

##  1 Introduction

Figure 1:  Learning for generalist behaviors necessitates building a natural
yet abstract understanding of the environment, as well as meta-reasoning and
interaction abilities. The metacognitive knowledge-guided approach captures
the inherent abstractions and relations between objects to generate
interaction goals, enabling the agent to interact with objects and generalize
to solve complex tasks.

Flexibility, ability to actively think and evaluate options –learning and
reasoning to enable adaptation– are hallmarks of human intelligence. Core to
this intelligence is how easily people can find natural decomposition in world
environments and tasks, enabling flexible, generalist behaviors. Replicating
such autonomy and general intelligence abilities in artificial agents
continues to pose significant challenges, which we argue requires a knowledge-
centric metacognitive approach.

Meta-reasoning about next actions, action values, and possible next states
facilitates flexible and transferable decision making for novel situations,
problem solving, concept learning, and knowledge acquisition. This reasoning
involves more than just memory lookup. Meta-reasoning is more broadly sense-
making in a knowledge system – enabling people to intuitively direct their
learning to make sense of relationships between entities, interact with the
same entities in different ways, and ground their interactions for both short
and long-term goals. Metacognitive reasoning appears to be central to
generalization. In contrast, meta-reasoning systems in artificial intelligence
are more focused on resource allocation problems. Although resource allocation
underlies efficient functioning, we need to ideate metacognitive
generalization, figure 1, to enable generalist behaviors in artificial agents.
Generalist behaviors have key characteristics which both constitute a basis of
devising a metacognitive framework and compartmentalize the behaviors in a
useful way.

Our proposal builds on three main principles for solving complex tasks in a
flexible manner. Natural decomposition and abstraction: World environments are
naturally composed of entities and categorized by their types. It is similar
to how a natural hierarchy of objects and classes is formed. Agents interact
with entities in the instance space, which inherit properties and
relationships from a higher-level type space. Interactions: World environments
are essentially knowledge graphs that represent the underlying objects and
their relationships. In order to complete its tasks and achieve its goals, an
agent’s metacognitive process employs these knowledge graphs to create a plan
that specifies which entities it must interact with and what interactions it
should have. Composition of interactions: Agents’ flexible behaviors stem from
their ability to learn interactions and leverage interaction pathways for
graph traversal at the type-space and instance levels as needed for quick
adaptation. Integrating these principles into metacognitive learning,
reasoning, and control provides levers and connecting points for instilling
generalist behaviors in agents.

Our proposed meta-cognitive framework captures the essence of generalist
behaviors and may have far-reaching implications across domains. In this
article, we provide a description of the key components of the proposed
framework, which exemplify generalist characteristics followed by experimental
evaluation, as well as how we envision it enabling flexibility through meta-
level knowledge guided interactions.

##  2 Background

In this section, we briefly describe closely related work as well as the
general setting of meta-reasoning, graph neural networks, and policy gradient
methods.

Related work Compositionality and abstractions are central for generalization
learning. Abstraction requires a method for grounding representations to the
actual objects. In previous work, Veerapaneni et al. (2020) presented an
object-centric model in which object representations were constructed from raw
observations and local simpler functions with uniformity, and then used for
prediction and planning. Similarly, Vijay et al. (2019) used a symbolic
object-centric representation to study the effect of adding prior knowledge
through relations in an object graph. These works only consider object space,
which inhibits the abstraction of objects and relational structures that occur
naturally in the real world and are used for making decisions. Some of the
main questions are what an abstraction should entail and how it can be used to
compose interactions that can be used as building blocks for leveraging
natural relational structures. We attempt to address these questions, which
have been largely overlooked in previous works. Moreover, what does natural
abstraction mean? It implies preserving relational higher-order structures and
abstractions that occur naturally in the real world. Our proposal preserves
abstract knowledge in the form of knowledge graphs, where abstractions are
grounded using type space. Learning a compositional structure is accomplished
by learning over knowledge graphs, with neighborhoods aggregated until
stability is reached. Different from previous works, that’s what our proposal
entails.

Other methods emphasize aspects which are different from ours, such as
encouraging exploration through intrinsic reward or information gain Ten et
al. (2022); Oudeyer et al. (2007); Pathak et al. (2017), or motivation
Moerland et al. (2023); Zhu et al. (2023); exploring for a variety of
experiences by developing internal goals for empowerment Klyubin et al. (2005)
; using a curriculum to progress from easier to more difficult tasks Soviany
et al. (2022); Campero et al. (2020) or inducing a learning hierarchy either
based on teacher signals Portelas et al. (2020) Pateria et al. (2021);
Hutsebaut-Buysse et al. (2022) or rewards Furelos-Blanco et al. (2023). These
methods focus on the execution space, influencing reward signals for learning
policies in one way or another but largely ignore the higher order structures
and relations that objects in the environment have. Our work distinguishes
itself by incorporating natural relational structures in the form of knowledge
graphs and allowing such meta-level knowledge structures to guide interactions
that govern low-level execution.

Meta-reasoning Meta-reasoning is often defined as the set of processes that
monitor our thought processes and regulate resource allocation Ackerman and
Thompson (2017). Its basic structure Nelson (1990) consists of two levels:
meta and object. The meta level’s functions include control and monitoring.
The object’s state changes as control passes from the meta to the object
level. The monitoring process allows the meta controller to track the status
of object-level processes and actions, as well as the consequences of those
actions. However, meta-cognitive reasoning in its current form as a resource
manager is insufficient to exhibit characteristics of generalist behaviors,
necessitating generalist meta-cognition approaches such as the one proposed
here.

Graph neural networks A graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ is
defined by a set of nodes $\mathcal{V}$ with node attributes $\mathtt{X}_{v}$
for $v\in\mathcal{V}$ and a set of edges $\mathcal{E}$ between nodes. Edge
from node $u\in\mathcal{V}$ to node $v\in\mathcal{V}$ is given by
$(u,v)\in\mathcal{E}$ with attributes $\mathtt{e}_{u,v}$. One typical task is
to predict the label or value of the entire graph. Graph neural networks (GNN)
(Gilmer et al. (2017); Hamilton et al. (2017) and others) use node
connectivity as well as node and edge features to learn a permutation-
invariant embedding vector for each node $h_{v},v\in\mathcal{V}$ and the
entire graph $h_{\mathcal{G}}$. Typically, the representation is learned
through message passing or neighborhood feature aggregation, in which node
embedding vectors are updated in each layer based on message aggregation from
neighbors:

$h_{v}^{(k)}=\mathcal{F}^{\prime}\left(h_{v}^{(k-1)},\bigoplus_{u\in\mathcal{N}(v)}\mathcal{F}^{\prime\prime}\left(h_{v}^{(k-1)},h_{u}^{(k-1)},\mathtt{e}_{u,v}\right)\right)$

where $h_{v}^{(k)}$ is the updated representation of node $v$ embedding in
layer $k$, $\mathcal{N}(v)$ is the set of neighboring nodes of node $v$,
$h_{u}^{(k-1)}$ is the representation of neighboring node $u$ in layer $k-1$,
$\bigoplus$ is a differentiable, permutation invariant aggregation function,
$\mathcal{F}^{\prime}$ and $\mathcal{F}^{\prime\prime}$ are differentiable
functions usually neural networks. Node features are transformed prior to
aggregation, and the methods used to transform and aggregate features differ
between graph network variants. Motivated by the attention mechanism Vaswani
et al. (2017) in the transformer model, Veličković et al. (2017) employs self-
attention over node features and aggregating over neighborhood nodes in graph
attention networks (GAT). We use the GATv2 operator Brody et al. (2021), which
appears to improve static attention problems in graph attention layers.

Policy gradient methods We consider a standard reinforcement learning setting
Sutton and Barto (2018) in which an agent interacts with the environment over
a number of steps. A Markov decision process (MDP) is defined as a tuple
$(\mathcal{S},\mathcal{A},\mathcal{P},r,\gamma)$ in which $\mathcal{S}$ is a
finite state space, $\mathcal{A}$ is a finite discrete action space,
$\mathcal{P}:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow\mathbb{R}$
is the transition model, $r:\mathcal{S}\times\mathcal{S}\rightarrow\mathbb{R}$
is the reward function and $\gamma\in(0,1)$ is the discount factor. At each
step t, the agent takes an action $a_{t}\in\mathcal{A}$ in state
$s_{t}\in\mathcal{S}$ according to its policy $\pi$, which is a mapping of
states to actions. As a result, the agent reaches state
$s_{t+1}\in\mathcal{S}$, receives a scalar reward $r_{t}\in\mathbb{R}$, and
accumulates returns; $R_{t}=\sum_{k=0}^{\inf}\gamma^{k}r_{t+k+1}$ represents
the accumulated return from step t. The agent’s goal is to maximize the
expected return for each state. The state-action value function
$Q^{\pi}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ is defined as
$Q^{\pi}(s,a)=\mathbb{E}[R_{t}|s_{t}=s,a]$, which is the expected return from
selecting action $a$ in state $s_{t}$ following policy $\pi$. The value of
state $s$ under policy $\pi$ is defined as
$V_{\pi}(s)=\mathbb{E}[R_{t}|s_{t}=s]$, which is the expected return from the
state $s$. Policy gradient methods, like REINFORCE Williams (1992) , directly
learn a parameterized policy $\pi(s;\theta)$ and update its parameters
$\theta$ using approximate gradient ascent of a performance measure,
$\nabla_{\theta}J(\theta)=\mathbb{E}_{\pi}[G_{t}\nabla_{\theta}\ln\pi_{\theta}(A_{t}|S_{t})]$.
To reduce the variance in the expected return estimate, a baseline is
subtracted from it; typically, a learned value function is used as the
baseline. The resulting quantity $A(s,a)=Q(s,a)-V(s)$ can be viewed as an
estimate of the advantage of action $a$ in state $s$. This method is called
advantage actor critic (A2C) Sutton and Barto (2018); Degris et al. (2012),
and the agent acting according to the policy $\pi$ is the actor, while the
value baseline is the critic. Several improved methods have been developed to
address sample efficiency, stability, and training difficulty. Proximal policy
optimization (PPO) Schulman et al. (2017) is an on-policy policy optimization
technique that addresses concerns in actor critic methods using clipped
objectives. We adapt PPO for learning policies.

##  3 Model

Our proposal for meta-cognitive generalization centers around the ideas of
natural decomposition of tasks that an agent can accomplish through
interactions with objects. Meta-policy recommends localized agent-object
interactions. The agent executes the interactions following interaction
policies and constructs a compositional solution. Here, we discuss the main
ideas of the framework (see figure 2) and conjecture that it aptly exhibits
task generalization characteristics.

Figure 2: The meta-cognitive generalization framework consists of three
levels: meta-level knowledge (K), interaction (I), and execution (X), which
are based on levels of abstraction and computation to provide generalization
capabilities. The knowledge or type space contains type space knowledge graphs
representing higher order abstractions. The instance space contains object-
level instance graphs and functions as an intermediary. The execution space is
the lowest level abstraction in which the agent acts on meta-level
recommendations. Computations proceed through recognition, recommendation,
execution, and interactions to gather experiences for updates to interaction
and meta policies.

Knowledge An agent can perform low-level actions $u=\\{u_{1},u_{2},...\\}$ in
the environment $\Xi$. The environment consists of objects, which the agent’s
perception systems associate with the entity types. Agent maintains its
perceived knowledge in the form of knowledge graphs: a graph of objects and
their relationships at the instance level, and a graph of entity types and
their abstract relationships at the type level. A type space knowledge graph
is an ordered pair $G_{K}=(E,\Lambda)$ that consists of a set of entity type
nodes, $E=\\{E_{1},E_{2},...,E_{m}\\}$, and a set of edges or type level
relations, $\Lambda=\\{\Lambda_{1},\Lambda_{2},...,\Lambda_{\kappa}\\}$. The
same type of relation between objects and abstract entity types has the same
meaning; for example, adjacency relations indicate that two objects are
adjacent to each other, whether in instance or type space. Using the knowledge
graph and meta-level action of the meta-policy,
$\lambda=\\{\lambda_{1},\lambda_{2},...,\lambda_{p}\\}$, the meta-level of the
agent determines how it will interact with an underlying object. To interact
with an object, the agent activates an interaction type, $\lambda$, following
the meta-policy guidance. For an interaction, the agent will perform a series
of low-level actions. For example, if there is a door, the agent can activate
‘open the door’ at the abstract level, even though it requires many low-level
actions such as moving and toggling. The agent uses an activation relation to
keep track of which objects it wants to interact with and what type of
interaction. Thus, the agent’s knowledge space
$\mathcal{K}=(A,E,\Lambda,\lambda,G_{K})$ comprises of a type space knowledge
graph, a set of entity types, abstract relations, and interactions.

Interaction The environment consists of objects,
$e=\\{e_{1},e_{2},...,e_{n}\\}$, their types, $\text{type}(e_{i})\in
E,\text{for}\ e_{i}\in e$, and is partially observable. Agent keeps track of
the visible objects in its instance space and builds an instance graph, which
is an ordered pair of nodes (observed objects) and edges (observed object
relations), $G_{I}=(e,\Lambda_{I})$. Thus, the agent’s instance space,
$\mathcal{I}=(A,e,\Lambda_{I},\lambda,G_{I})$, is made up of its instance
space knowledge graph, relationships between objects, and available
interactions. After constructing the knowledge graphs, the agent’s goal is to
identify the objects for interaction which is done by the meta-level. Meta-
level recommends an interaction goal for the agent. For the recommended
interaction goals, the agent learns interaction policies for interacting with
objects, as well as the relationships between interaction types and entity
types. With the ability to interact with objects, the agent accumulates
interaction knowledge to update the meta-level policy. Therefore, when viewing
tasks as meta-policy guided interactions with objects, problems are naturally
decomposed into interaction goals.

eXecution By pursuing interaction goals as directed by the meta-level, the
agent learns to interact with objects and act at the execution level. The
execution space for the agent is its partially observed view $X$. During an
interaction, the agent takes actions $u$ on low-level states $s$. After
learning both the interaction and meta policies, the agent acts in the
environment, aided by interaction with objects. Therefore, unlike the models
that are tightly coupled with low-level representations through end-to-end
policies, agent actions are aided by meta-level knowledge-guided object
interactions, providing a mechanism for generalist behaviors.

Mapping between execution space and instance space Assume that the agent has
an object recognition system that can access its partially observed view $X$.
The agent generates an instance-level graph using the observed objects. The
function $\Phi:X\to G_{I}$ maps the agent’s object recognition to its instance
graph by capturing observed objects and their relationships. As the agent
interacts with objects, the environment changes, and its recognition system
captures a new partially observed view. The mapping function then generates
the corresponding instance graph.

Mapping between instance space and type space The function $\Psi:G_{I}\to
G_{K}$ provides a mapping between instance and type graphs. If the agent
encounters a novel object for which it does not have a corresponding typed
entity, it must either create a new type in its type space or assign the
object to a type that is closely related. In the framework, $\Psi$ is
learnable and represents a belief of mapping objects to their types. However,
in this article, we treat the mapping as a fixed mapping from object nodes in
$G_{I}$ to their corresponding type nodes in $G_{K}$.

Proposal for activation object and interaction Using the meta-policy, the
agent’s recommendation function $\Upsilon:G_{K},G_{I},\pi\to
e_{i},\lambda_{p}$ recommends a candidate object $e_{i}$ for interaction and
interaction goal $\lambda_{p}$. The recommender iterates through all visible
objects in the agent’s instance graph $G_{I}$, then activates the object by
adding an activation relationship between the agent and the object to the
instance graph, as well as a corresponding activation relation between the
agent and the object’s type node in the type space graph $G_{K}$. Activated
$G_{K}$ represents the activated candidate object and serves as the input
state for meta-policy, which returns both the state value and meta-action
probabilities. The state value represents the value of activating the
respective object. Thus, the recommender proposes an object for interaction
$e_{i}$ by sampling from the values of activating objects, and the proposed
interaction $\lambda_{p}$ is sampled from the meta-action probabilities of the
proposed object. Acting on these recommendations, the agent collects
interaction experiences $G_{K},\lambda$, that is, interaction trajectories
$G_{K}^{0}\curvearrowright G_{K}^{1}...\curvearrowright G_{K}^{t}$, and learns
meta-level policy.

Interactions At each step, the agent activates the object $e_{i}$ for
interaction $\lambda_{p}$, which is its interaction goal,
$g^{i}_{p}=(e_{i},\lambda_{p})$. Following this interaction goal, the agent
interacts with the environment. After each round of interaction, the status of
the interaction is passed to the meta level. Using the meta-policy, the meta-
level recommender suggests next target object and interaction. The process
continues until the task is completed or the environment is terminated.
Throughout the process, the agent collects experience trajectories at both the
interaction and meta levels. The policies are then updated based on the
collected experiences.

Meta-policy learning Meta-level representations of type space graphs with
activation relations are learned using a message-passing graph neural network.
Nodes have abstract level attributes such as type encoding, while edges have
edge type encoding for different types of relationships. We use GAT
convolutional layers Brody et al. (2021); Fey and Lenssen (2019) to learn
graph representations given as,

$h_{v}^{\prime}=\alpha_{v,v}\theta
h_{v}+\sum_{u\in\mathcal{N}(v)}\alpha_{u,v}\theta h_{u}$

where $h_{v}^{\prime}$ is the updated embedded representation of node $v$,
$\mathcal{N}(v)$ is the set of neighbors of node $v$, and $h_{u}$ is the
representation of the neighboring node. $\theta$ are learnable weights shared
by source and target nodes in the neighborhood, and the attention coefficients
$\alpha_{u,v}$ are computed using weighted node features.

The hidden representation after the final graph attention layer is pooled for
graph rollout, which is then passed to the fully connected actor and critic
layers of the model. The meta policy $\pi_{\theta}(.|G_{K})$, parameterized by
$\theta$, returns a state value and an interaction or meta-action, $\lambda$,
distribution. The recommender selects a target object and an interaction
$\lambda_{p}$, and the agent works through low-level interactions based on
guidance. At the end of the interaction, the status of the interaction is
returned to the meta level, and the agent collects the trajectory of meta-
level experiences, $\mathcal{D}_{\tau_{m}}$. An estimate of advantages is
computed as
$A_{m}=r_{m}+\gamma_{m}V_{\pi_{\theta}}(G_{K}^{\prime})-V_{\pi_{\theta}}(G_{K})$.

Meta level PPO optimizes a surrogate loss:

$\mathcal{L}_{m}^{\text{actor}}(\theta)=-\mathbb{E}_{\mathcal{D}_{\tau_{m}}}\left[\min\left(r_{\theta}A_{m},\text{clip}(r_{\theta},1-\epsilon,1+\epsilon)A_{m}\right)\right]$

with $r_{\theta}=\frac{\pi_{\theta}(.|G_{K})}{\pi_{\theta}^{r}(.|G_{K})}$
being the ratio of probabilities of actions at meta level roll-out. The critic
loss is the mean squared error between the estimated values and the returns:

$\mathcal{L}_{m}^{\text{critic}}(\theta)=\mathbb{E}_{\mathcal{D}_{\tau_{m}}}\left[\frac{1}{2}(V_{\pi_{\theta}}(G_{K})-R_{m})^{2}\right]$

Considering the mean entropy $\mathcal{H}_{\theta}$, the overall minimization
objective at the meta-level becomes:

$\mathcal{L}_{m}(\theta)=\mathcal{L}_{m}^{\text{actor}}(\theta)-\zeta_{m}^{h}\mathcal{H}_{\theta}+\zeta_{m}^{c}\mathcal{L}_{m}^{\text{critic}}(\theta)$

where $\zeta_{m}^{h}$, $\zeta_{m}^{c}$ are the entropy and value coefficients,
respectively. The model parameters $\theta$ are optimized as
$\theta\leftarrow\theta+\alpha_{m}\nabla_{\theta}L_{m}(\theta)$, and the
updated meta-policy is used for the next episode of learning.

Interaction policy learning Interaction policies are a repository of policies
for different interactions which can be recommended by the meta-level,
$\mu_{\phi}=\\{\mu_{\phi_{1}},\mu_{\phi_{2}},...,\mu_{\phi_{p}}\\}$ and each
policy is parameterized separately. When an interaction goal with an object
$e_{i}$ and interaction $\lambda_{p}$ is recommended,
$g^{i}_{p}=(e_{i},\lambda_{p})$, the agent attempts to achieve it by executing
a series of low-level actions from $u$. The agent observes a low-level state
$s$ from a convolutional neural network embedding, samples an action
$u\sim\mu_{\phi_{p}}(.|s)$ using the respective interaction policy, acts in
the environment, and collects the trajectory of experiences
$\mathcal{D}_{\tau_{p}}$. The experiences collected are used to learn
interaction policies. An estimate of advantages is computed as
$A_{p}=r_{p}+\gamma_{p}V_{\mu_{\phi_{p}}}(s^{\prime})-V_{\mu_{\phi_{p}}}(s)$.

Each interaction level PPO optimizes a surrogate loss:

$\mathcal{L}_{p}^{\text{actor}}(\phi_{p})=-\mathbb{E}_{\mathcal{D}_{\tau_{p}}}\left[\min\left(r_{\phi_{p}}A_{p},\text{clip}(r_{\phi_{p}},1-\epsilon,1+\epsilon)A_{p}\right)\right]$

with $r_{\phi_{p}}=\frac{\mu_{\phi_{p}}(.|s)}{\mu_{\phi_{p}}^{r}(.|s)}$ is the
ratio of probabilities of actions of policy and roll-out policy. The critic
loss is the mean squared error between the estimated values and the returns:

$\mathcal{L}_{p}^{\text{critic}}(\phi_{p})=\mathbb{E}_{\mathcal{D}_{\tau_{p}}}\left[\frac{1}{2}(V_{\mu_{\phi_{p}}}(s)-R_{p})^{2}\right]$

Considering the mean entropy $\mathcal{H}_{\phi_{p}}$, the overall
minimization objective becomes:

$\mathcal{L}_{p}(\phi_{p})=\mathcal{L}_{p}^{\text{actor}}(\phi_{p})-\zeta_{p}^{h}\mathcal{H}_{\phi_{p}}+\zeta_{p}^{c}\mathcal{L}_{p}^{\text{critic}}(\phi_{p})$

for each interaction $p$, where $\zeta_{p}^{h}$, $\zeta_{p}^{c}$ are the
entropy and value coefficients, respectively. For each interaction policy, the
model parameters $\phi_{p}$ are optimized as
$\phi_{p}\leftarrow\phi_{p}+\alpha\nabla_{\phi_{p}}L_{p}(\phi_{p})$, and the
updated interaction policies are used for the next episode of learning.

As a result, the agent follows a trajectory of activated interactions with
objects and aids the formation of a task solution through meta-guided
interactions. The proposed framework is generic to be applied to a wide range
of domains.

##  4 Experiments

Experimental set up To assess the capabilities of the proposed agents, we
require a challenging environment with ingredients for higher order reasoning
to solve tasks. Our goal here is to investigate whether the introduction of
meta-knowledge guided interactions benefits agents. MiniGrid world
environments Chevalier-Boisvert et al. (2018) are suitable for our purpose.
These are procedurally generated environments in which the agent interacts
with objects such as doors, keys, balls, and boxes.

Figure 3: MiniGrid ObstructedMaze Full environment is an appropriate testbed
because it is partially observable, requires complex reasoning, and new tasks
can be created.

The environments are challenging due to partially observability, sparsity of
rewards and random goal locations. The MiniGrid-ObstructedMaze-Full
environment, figure 3, in particular, is a difficult environment with only a
few claims for solving it. Due to the large state space size, stochasticity
and partial observability, agents cannot memorize but must learn.

In the environment, the agent is initially placed in the middle room, with a
7x7 view. At the low level, the agent can move forward, left, right, pick,
drop, and toggle, and its goal is to pick up the blue ball, which is randomly
placed in one of the corner rooms. Some of the doors are locked and require a
key of the same color to open. Keys are hidden inside boxes, but the agent
cannot see inside the boxes without toggling. The original task is referred as
Task 0. Furthermore, we identify the following additional task scenarios that
were not seen during training but were presented to the agent during testing:

  * •

No context in training (Task 1) - a relation that the agent did not come
across during training. It is simulated by placing the goal inside an empty
box. During training, the blue ball was never in a box, although the agent
encountered boxes and balls during training, but boxes contained keys.

  * •

Different context in testing (Task 2) - using a relation that the agent
encountered during training but in a new context during testing. It is
simulated by locking doors in the central room and placing keys directly in
the room. During training, doors in the central room were never locked, nor
was there direct access to keys.

  * •

Dynamic changes in the environment (Task 3) - when some event occurs, the
environment changes. During training, the goal location remained unchanged. To
simulate, as soon as the agent enters the room with the goal, relocate the
goal to one of the neighboring rooms.

Agent training

Initialize model weights $\theta,\phi$ , parameters, $T$

for  _ $t=1$ to $T$_ do

// collect exp for meta exp size K

for  _$k=1$ to $K$_ do

make instance graph $G_{I}\leftarrow\Phi(X)$

map $G_{I}$ to $G_{K}$, $G_{K}\leftarrow\Psi(G_{I})$

proposal $(e_{i},\lambda_{p})\leftarrow\Upsilon(G_{K},G_{I},\pi)$

// 1st variant

interaction exp following $\mu_{\lambda_{p}}$

// 2nd variant

reaching exp following $\mu_{\text{reach}}$

if  _success_ then

interaction exp following $\mu_{\lambda_{p}}$

end if

collect exp $\mathcal{D}_{\tau_{p}}$ and compute $A_{p}$

end for

collect meta exp $\mathcal{D}_{\tau_{m}}$ and compute $A_{m}$

// Update policies, both variants

update $\pi$ by optimizing $\mathcal{L}_{m}(\theta)$

for  _ $\lambda_{p}$_ do

update $\mu_{\lambda_{p}}$ by optimizing $\mathcal{L}_{p}(\phi_{p})$

end for

// 2nd variant only

update $\mu_{\text{reach}}$ by optimizing
$\mathcal{L}_{\text{reach}}(\phi_{\text{reach}})$

end for

Algorithm 1 KIX

We devise two variants, KIX1 and KIX2. Both variants learn meta-policy and
interaction policies. However the second variant additionally learns a
reachability policy, so it interacts with an object once it has reached it.
All meta, interactions, and reachability policies are implemented as actor-
critic PPO policies. Algorithm 1 provides a high-level overview of the
implementation. The input state embeddings for interaction and reachability
actor-critic networks are constructed by a convolutional network with three
Conv2D layers with 16, 32, and 64 output channels, a max pooling layer added
after the first convolution layer, and exponential linear units used as
activation functions. The meta-level recommends various objects for
interaction, so an activation indicator for the object of interest is added to
a new input channel, creating a four-channel input. Actors and critics are
fully connected networks, each with one hidden layer. Meta-level policies are
learned over type space knowledge graphs. In a type space graph, node object
type information is encoded as an attribute vector. Edges in type space graph
represent type-level relationships between nodes, such as "visible,"
"adjacent," "carrying," and "activated." A one-hot encoded edge attribute
vector is used to encode these relationships. Embeddings for meta-level actor-
critic are provided by two graph attention convolution layers, each with
16-dimensional embeddings and four attention heads. During message passing,
these layers share weights between source and target. The final graph
attention convolution layer is used for graph readout with a global add pool.
Meta-level actors and critics are fully connected networks, with the critic
learning a single value and the actor learning the probability of taking a
meta-level action, such as ’pickup’, ’drop’, ’reveal’, ’open’, or ’open with
key’. The meta-level actions serve as interaction goals. The environment
rewards the agent with $r=1-(c/c_{T})$ where $c$ is the number of steps taken
and $c_{T}$ is the max possible steps. For interaction policy learning, the
agent receives a reward of 1 when it has a successful interaction. At the
meta-level it is rewarded with 0.1 for a successful proposed interaction and
$1+r$ for achieving the environment goal. A base agent is also implemented,
with a similar network and a PPO policy, and it receives environmental
rewards. Grid search is used to select hyperparameters. All the agents are
trained for the same number of environment steps using 32 parallel workers.

Figure 4: The return profiles of agents performing various evaluation tasks
show that the proposed agents have higher returns on all tasks and benefit
from meta-guided interactions.

Meta-guided interactions aid agents Our proposal is motivated by the idea
that, rather than acting directly at the execution level, meta-level guided
interactions with objects aid agents in generalizing while performing tasks.
Because the agent’s environmental return is determined by the number of steps
taken to complete the task, completing tasks in fewer steps yields a higher
return. Therefore, the distribution of environmental returns is a reliable
indicator of whether meta-level interaction recommendations benefit agents. We
evaluate the distribution of environmental returns for each agent across all
four tasks. For each task, the trained agents are run for 12K episodes by
parallel workers, and their environmental rewards are recorded. In order to
understand the effects of meta-guided interactions and to ensure a fair
comparison between agents where differences in success rates are not a factor,
we examine the top $k$ returns with $k=100$. Figure 4 illustrates the return
profiles of agents. It is clear from the distribution of returns that the
proposed agents have received a higher return on all tasks than the base
agent, furthermore, this is even more so for dynamic tasks like task 3. It
demonstrates that meta-guided interactions not only help agents complete their
task goals in fewer steps, but also aid them to perform a wider range of tasks
in static or dynamic environments. Furthermore, the higher returns of the
second variant demonstrate that agents benefit from the introduction of
reachability as a separate interaction. In the absence of a separate
reachability policy, reaching is implicit in each interaction policy; however,
separating allows for a shared learning of reachability followed by
interactions which could be beneficial particularly in dynamically changing
environments.

Figure 5: Wasserstein distances between traversal distributions, when agents
complete tasks 1-3, and the trained task 0, provide comparative insights.

Agents profile While performing the tasks, agents interact with objects in the
environment according to the recommended interaction policies; as a result,
agents navigate and visit rooms. Because interactions are recommended by the
meta-level, which is only trained on the training task, tracking agents’ room
visits across all tasks provides an insight about their learning. Let $p_{0}$
be the distribution of agent room visits for trained task 0, and
$p_{i},i\in\\{1,2,3\\}$ be the distributions for tasks 1–3. For both variants
of the agent, we then compute the Wasserstein distance, a metric from optimal
transport theory, between the visit distribution for tasks 1-3 and the trained
task 0, as shown in figure 5. The distribution of visits to each agent is
expected to be different from the training task, due to changes in the
evaluation tasks. However, our interest is in how differently the two variants
behave in evaluation tasks, as indicated by the computed metric. In task 1,
the goal room remained unchanged, but a new relation was added to the
evaluation. For both variants, the computed metric value for task 1 is
similar, indicating that they transfer their learning to this task in a
similar manner. However, there is a noticeable difference for tasks 2 and 3.
In task 2, doors in the central room are locked with keys kept in the room - a
static task, whereas in task 3, the goal is dynamically relocated based on
observance - a dynamic task. The computed metric for the dynamic task for the
second variant is smaller than the first variant, implying that its evaluation
visit distribution is closer to that of the training task. It suggests that
the learning in the second variant can better capture dynamic changes. When
combined with the return profiles, it reinforces that having an additional
reachability policy benefits the agent in a wider range of tasks, particularly
dynamic tasks, and improves generalization.

##  5 Conclusion

In this work, we proposed a metacognitive framework that promotes generalist
behaviors through three main contributions. First, generalist and flexible
behaviors necessitate abstraction and composition, but current models lack
natural abstractions and compositions that can preserve both object
abstractions and relational knowledge structures. In our model, natural
abstractions and compositions are represented and grounded using type space
knowledge graphs which also preserve relational knowledge between objects.
Second, rather than directly acting in the raw environment, interactions with
objects play an important role in how an agent navigates the environment to
complete tasks. To accomplish this, we propose an intermediary interaction
level. The meta-level of the agent learns from type space knowledge graphs and
recommends what objects to interact with, as well as what interaction should
be. In this sense, our model also functions as an interaction goals generating
machine. Third, interaction policies learned through interaction goals serve
as a repository of policies for the agent to use in the execution space,
providing for the necessary loose coupling for generalization via interactions
as intermediaries. The agent’s meta-level recommends interactions based on
knowledge structures and learned meta-policy, and the execution level applies
interaction policies to complete tasks. We demonstrated that metacognitive
knowledge guided interactions aid agents in generalization. Our framework is
modular and adaptable, with the potential to have implications across domains
through the induction or integration of knowledge structures, whether
extrinsic, intrinsic, or derived, as well as control variables at the meta or
interaction level.

## References

  * Ackerman and Thompson [2017] Rakefet Ackerman and Valerie A Thompson.  Meta-reasoning: Monitoring and control of thinking and reasoning.  Trends in cognitive sciences, 21(8):607–617, 2017.
  * Brody et al. [2021] Shaked Brody, Uri Alon, and Eran Yahav.  How attentive are graph attention networks?  arXiv preprint arXiv:2105.14491, 2021.
  * Campero et al. [2020] Andres Campero, Roberta Raileanu, Heinrich Küttler, Joshua B Tenenbaum, Tim Rocktäschel, and Edward Grefenstette.  Learning with amigo: Adversarially motivated intrinsic goals.  arXiv preprint arXiv:2006.12122, 2020.
  * Chevalier-Boisvert et al. [2018] Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal.  Minimalistic gridworld environment for openai gym.  GitHub repository, 2018.
  * Degris et al. [2012] Thomas Degris, Patrick M Pilarski, and Richard S Sutton.  Model-free reinforcement learning with continuous action in practice.  In 2012 American Control Conference (ACC), pages 2177–2182. IEEE, 2012.
  * Fey and Lenssen [2019] Matthias Fey and Jan Eric Lenssen.  Fast graph representation learning with pytorch geometric.  arXiv preprint arXiv:1903.02428, 2019.
  * Furelos-Blanco et al. [2023] Daniel Furelos-Blanco, Mark Law, Anders Jonsson, Krysia Broda, and Alessandra Russo.  Hierarchies of reward machines.  In International Conference on Machine Learning, pages 10494–10541. PMLR, 2023.
  * Gilmer et al. [2017] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl.  Neural message passing for quantum chemistry.  In International conference on machine learning, pages 1263–1272. PMLR, 2017.
  * Hamilton et al. [2017] Will Hamilton, Zhitao Ying, and Jure Leskovec.  Inductive representation learning on large graphs.  Advances in neural information processing systems, 30, 2017.
  * Hutsebaut-Buysse et al. [2022] Matthias Hutsebaut-Buysse, Kevin Mets, and Steven Latré.  Hierarchical reinforcement learning: A survey and open research challenges.  Machine Learning and Knowledge Extraction, 4(1):172–221, 2022.
  * Klyubin et al. [2005] Alexander S Klyubin, Daniel Polani, and Chrystopher L Nehaniv.  Empowerment: A universal agent-centric measure of control.  In 2005 ieee congress on evolutionary computation, volume 1, pages 128–135. IEEE, 2005.
  * Moerland et al. [2023] Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al.  Model-based reinforcement learning: A survey.  Foundations and Trends® in Machine Learning, 16(1):1–118, 2023.
  * Nelson [1990] Thomas O Nelson.  Metamemory: A theoretical framework and new findings.  In Psychology of learning and motivation, volume 26, pages 125–173. Elsevier, 1990.
  * Oudeyer et al. [2007] Pierre-Yves Oudeyer, Frdric Kaplan, and Verena V Hafner.  Intrinsic motivation systems for autonomous mental development.  IEEE transactions on evolutionary computation, 11(2):265–286, 2007\.
  * Pateria et al. [2021] Shubham Pateria, Budhitama Subagdja, Ah-hwee Tan, and Chai Quek.  Hierarchical reinforcement learning: A comprehensive survey.  ACM Computing Surveys (CSUR), 54(5):1–35, 2021.
  * Pathak et al. [2017] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  In International conference on machine learning, pages 2778–2787. PMLR, 2017.
  * Portelas et al. [2020] Rémy Portelas, Cédric Colas, Katja Hofmann, and Pierre-Yves Oudeyer.  Teacher algorithms for curriculum learning of deep rl in continuously parameterized environments.  In Conference on Robot Learning, pages 835–853. PMLR, 2020.
  * Schulman et al. [2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  arXiv preprint arXiv:1707.06347, 2017.
  * Soviany et al. [2022] Petru Soviany, Radu Tudor Ionescu, Paolo Rota, and Nicu Sebe.  Curriculum learning: A survey.  International Journal of Computer Vision, 130(6):1526–1565, 2022\.
  * Sutton and Barto [2018] Richard S Sutton and Andrew G Barto.  Reinforcement learning: An introduction.  MIT press, 2018.
  * Ten et al. [2022] Alexandr Ten, Pierre-Yves Oudeyer, and Clément Moulin-Frier.  Curiosity-driven exploration.  The Drive for Knowledge: The Science of Human Information Seeking, page 53, 2022.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * Veerapaneni et al. [2020] Rishi Veerapaneni, John D Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine.  Entity abstraction in visual model-based reinforcement learning.  In Conference on Robot Learning, pages 1439–1456. PMLR, 2020.
  * Veličković et al. [2017] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio.  Graph attention networks.  arXiv preprint arXiv:1710.10903, 2017.
  * Vijay et al. [2019] Varun Kumar Vijay, Abhinav Ganesh, Hanlin Tang, and Arjun Bansal.  Generalization to novel objects using prior relational knowledge.  arXiv preprint arXiv:1906.11315, 2019.
  * Williams [1992] Ronald J Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  Machine learning, 8:229–256, 1992.
  * Zhu et al. [2023] Zhuangdi Zhu, Kaixiang Lin, Anil K Jain, and Jiayu Zhou.  Transfer learning in deep reinforcement learning: A survey.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023\.


# ConTextual Masked Auto-Encoder for Retrieval-based Dialogue Systems

Zhenpeng Su, Xing Wu, Wei Zhou, Guangyuan Ma, Songlin Hu  Zhenpeng
Su1,2111These authors contributed equally to this work., Xing Wu1,2,3111These
authors contributed equally to this work., Wei Zhou1,2222Corresponding
authors., Guangyuan Ma1,2, Songlin Hu1,2222Corresponding authors.
1Institute of Information Engineering, Chinese Academy of Sciences, Beijing,
China
2School of Cyber Security, University of Chinese Academy of Sciences, Beijing,
China
3Kuaishou Technology, Beijing, China
<EMAIL_ADDRESS>

###### Abstract

Dialogue response selection aims to select an appropriate response from
several candidates based on a given user and system utterance history. Recent
studies have been improving the accuracy of dialogue response selection
through post-training, mostly relying on naive masked language modeling
methods. However, the recently developed generative methods have shown
promising text representation capabilities in IR community, which could
potentially lead to better dialogue semantics modeling. Thus, in this paper,
we propose Dial-MAE (Dialogue Contextual Masking Auto-encoder), a
straightforward yet effective post-training technique tailored for dialogue
response selection. Dial-MAE uses an asymmetric encoder-decoder architecture
that learns to better compress the semantics of the dialogue into dialogue-
dense vectors. The process of Dial-MAE involves a deep encoder creating a
dialogue embedding with the masked dialogue context, followed by a shallow
decoder that uses this embedding along with the highly masked response to
restore the original response. Our experiments have demonstrated that Dial-MAE
is highly effective, achieving state-of-the-art performance on two commonly
evaluated benchmarks.

Dialogue systems based on retrieval are a popular research topic. Pre-trained
language models (PLMs) have been widely used in dialogue response selection,
such as deep bidirectional Transformer encoder Vaswani et al. (2017) LMs like
BERT Devlin et al. (2019), which is used in many response selection models and
achieves superior performance. Recent works Han et al. (2021); Zhang et al.
(2022); Whang et al. (2021) on dialogue response selection models are more
based on cross-encoder, feeding the compared dialogue context and response
directly into LM, and using attention over all tokens to output a relevance
score. Although effective in this way, these cross-encoder models have to
compute the matching for every possible combination of context-response pairs,
which is time-consuming Lan et al. (2021). In practice, cross encoders are
often used as re-ranking. Another common use of deep LM is bi-encoder, also
known as dense retrieval, which encodes dialogue context and response into
context vectors and response vectors, respectively. Bi-encoder has a faster
computational speed but often performs worse than the cross-encoder model. For
the bi-encoder, the correlations between context and responses are computed
using cosine similarity or dot product functions in vector space.

Since PLMs are pre-trained on common corpora such as Books Corpus and
Wikipedia, post-training with data from the target domain can further improve
fine-tuning. Previous studies have paid more attention to cross-encoders,
Whang et al. (2020); Gu et al. (2020) propose domain-adaptive post-training
objective are the same as BERT, including masked language model (MLM)
objective and next sentence prediction (NSP) objective. Furthermore, Han et
al. (2021) propose a simple yet powerful fine-grained post-training approach,
which builds on MLM objective and achieves very good results by classifying
the relationship between a given context and responses into finer-grained
labels. Considering the bi-encoder, Lan et al. (2021) propose a dialogue-dense
retrieval model, using language model encoding to encode context and responses
separately into dense vectors. They conducted post-training via MLM objective
and then used a contrastive learning method to fine-tune, making the bi-
encoder model achieve competitive results, surpassing most current cross-
encoder models.

Although these works have proven effective for post-training using MLM
objective before fine-tuning, Gao and Callan (2021) find train bi-encoder has
a difficulty that is not seen in the cross-encoder. Their experiments showed
that in standard PLMs like BERT, pre-training with MLM objective, were not
trained to aggregate sophisticated information into a single vector
representation like $[CLS]$ token. It is manifested in the dialogue retrieval
model that it is difficult for LMs to get a dialogue vector with sufficient
semantics only post-training with MLM objective. How to fully aggregate
dialogue contextual semantics through post-training is not discussed. In this
paper, inspired by the IR communityLu et al. (2021); Gao and Callan (2021); Wu
et al. (2022), we propose a post-training method tailored to enhance the
ability of dialogue-dense vector representation.

Recently, several methods Lu et al. (2021); Gao and Callan (2021); Wu et al.
(2022) have emerged to improve the representation ability of dense vectors
through pre-training tailored for dense retrieval. They propose to force the
encoder to provide a better text representation through an auxiliary MLM task
of the decoder. The role of the auxiliary task is to help the encoder fully
aggregate semantic information. Inspired by this, in this paper, we propose a
simple and effective generative post-training method to compress dialogue
semantic information to enhance the representation of dialogue-dense vectors.
Previous domain adaptive post-training methods focus on single-text and do not
consider dialogue contextual information. Differently, during the modeling
process, we consider both the semantics of the dialogue context and the
semantic relevance of the response. With the help of the dialogue context
embedding output by the encoder, the auxiliary task uses a weak decoder to
reconstruct the masked response text. As shown in Figure 1, firstly, we
perform random masking operations on the context and response, respectively.
We use an asymmetric masking strategy where the response has a higher masking
rate than the context. Next, the masked dialogue context is fed into the
encoder, and the original text is reconstructed using MLM objective. Then, we
consider the MLM objective for the decoder. The input of the decoder includes
the dialogue embedding output by the encoder and the masked response text, and
the original response text is reconstructed from both inputs via MLM. It is
worth noting that we use a shallow decoder and aggressively increase the
masking ratio on the decoder side. Therefore, the reconstruction on the
decoder side is difficult to accomplish only by relying on masked the response
and must rely more on the dialogue embedding output by the encoder. This
forces the encoder to sufficiently aggregate the semantics of the dialogue
context to aid the decoder in its MLM task.

Figure 1: The model design for Dial-MAE. The input of the encoder is the
dialogue context, and its next response and dialogue context embedding output
by the encoder are used as the input to the decoder. The encoder is deep and
has enough parameters to learn a good dialogue embedding, shallow decoder help
encoders achieve this goal through MLM.

To verify the effectiveness of our proposed Dial-MAE, we conduct experiments
on three benchmarks(Ubuntu Lowe et al. (2015), Douban Wu et al. (2017), and
E-commerce Zhang et al. (2018)). Our method achieves state-of-the-art
performances on two of these benchmarks: Ubuntu and E-commerce. In addition,
we conduct ablation experiments to verify the effectiveness of our method. Our
contributions are as follows:

  1. 1.

We propose Dial-MAE, a novel generative post-training method tailored for
dialogue retrieval, which considers both context and its response.

  2. 2.

We design a novel asymmetric encoder-decoder architecture to enhance the
representational power of dialogue embedding.

  3. 3.

Experimental results show that in dialogue response selection, our method
achieves the state-of-the-art on two benchmarks with faster response speed.

##  1 Related Work

###  1.1 Dialogue Response Selection

Dialogue response selection aims to select the most appropriate response from
a range of candidates. Earlier studies Kadlec et al. (2015); Lowe et al.
(2015) focused on single-turn response selection. Recent studies pay more
attention to multi-turn dialogue response selection. Lowe et al. (2015)
propose to calculate the matching degree of dialogue and response based on
RNN, and contributed a benchmark dataset called Ubuntu V1. Kadlec et al.
(2015) propose to use CNN and LSTM as encoders to encode the context and
response. Zhou et al. (2016) propose a multi-view model that encodes both
word-level and utterance-level representations. Wu et al. (2017) propose to
use word embeddings and their sequential representations encoded by GRU to
build a matching matrix between dialogue context and response, and contributed
a dataset called Douban. Zhang et al. (2018) propose to form a fine-grained
context representation through deep discourse aggregation, and contributed a
dataset called E-commerce. Zhou et al. (2018) propose a deep attention
matching network that applies the attention mechanism to the response
selection dialogue system. Tao et al. (2019) propose to perform context and
response matching by stacking multiple interaction blocks. Yuan et al. (2019)
propose a multi-hop selector network to select the relevant utterances in the
context of response matching.

After pre-trained language models showed impressive improvements in various
downstream NLP tasks, some researchers began to apply them to response
selection. BERT-VFT Whang et al. (2020) first applies the pre-trained language
model BERT to dialogue response selection, and achieved state-of-the-art
results. SA-BERT Gu et al. (2020) adds speaker embedding to the model, in
order to make the model aware of the speaker change information. Multi-Task
Learning is also an effective way, UMSBERT+ Whang et al. (2021) proposes a set
of strategies, which aids the response selection model towards maintaining
dialogue coherence. Alternatively, Xu et al. (2021a) propose learning a
context-response matching model with multiple auxiliary self-supervised tasks.
BERT-FP Han et al. (2021) proposes to classify the relationship between a
given utterance and a target utterance into more fine-grained labels, which
makes the model learns the semantic relevance and coherence between the
utterances. DR-BERT Lan et al. (2021) proposes a dense retrieval model to
encode context and response. Zhang et al. (2022) propose two-level supervised
contrastive learning so that the learned dialogue representations can be
further separated in the embedding space. Although many previous studies have
used domain-adaptive post-training, there is no research on tailoring post-
training tasks to enhance the representation ability of dialogue-dense
vectors.

###  1.2 Pre-training Methods tailored for Dense Retrieval

Many pre-training methods tailored for dense retrieval have recently been
proposed and proven effective in tasks such as paragraph retrieval and
document retrieval. Some previous studies Lu et al. (2021); Gao and Callan
(2021); Wu et al. (2022); Liu and Shao (2022) encourage the encoder to learn
better text representations by auxiliary self-supervised auto-encoding tasks.
Specifically, Lu et al. (2021) propose to pre-train the auto-encoder by using
a weak auto-regressive decoder, with restricted capacity and attention
flexibility to push the encoder to provide better text representations. Gao
and Callan (2021) utilize MLM objective as an auxiliary self-supervised
encoding task. Liu and Shao (2022) propose to apply asymmetric mask ratios to
an asymmetric encoder-decoder structure. Wu et al. (2022) propose context-
supervised auto-encoding, in which a masked version of a given text and its
neighboring text’s embeddings. Precisely, self-supervised masking auto-
encoding reconstructs the masked input text span, and context-supervised pre-
training reconstructs the other text span with the help of its neighboring
span embedding. In this paper, we propose a post-training method tailored for
dialogue response selection, which considers dialogue contextual information
to enhance the representation of dialogue-dense vectors.

##  2 Methodology

This section first introduces masked language model pre-training as
preliminary knowledge. Then we introduce detailed post-training, including the
construction of data and the auxiliary task. Finally, we introduce the details
of fine-tuning.

###  2.1 Masked Language Model Pre-training

The MLM task is used for the pre-training of many transformer architecture
models. MLM is an unsupervised method that masks parts of the input tokens and
requires the model to predict them based on the unmasked tokens. Formally,
given an input sentence $X=[x_{1},x_{2},...,x_{n}]$. We select a certain
percentage of tokens from $X$ and replace them with a special token $[MASK]$
to get polluted $\widetilde{X}$. We denote the tokens replaced by $[MASK]$ as
$m(X)$. Then, LM is used to transform the polluted input as the hidden states:

$\displaystyle=LM([CLS],\tilde{X})$ (1)

$[CLS]$ is a special token that is prepended at the beginning of the text. $l$
represents the number of layers of the transformer. For masked $x_{i}$ its
corresponding hidden feature $h_{i}$ is used to predict the actual $x_{i}$. We
formulate this process as:

$\displaystyle L_{mlm}=-\sum_{x_{i}\in
m(X)}\log{p(x_{i}|LM([CLS],\tilde{X}))}$ (2)

###  2.2 Dial-MAE: Dialogue Contextual Masking Auto-encoder

Dial-MAE learns dialogue context information, which jointly models the
semantics of the tokens inside a dialogue context and its response. We first
describe how to build training data from all utterances of the dialogue
session and then introduce the Dial-MAE post-training method. We randomly
sample multiple consecutive utterances as context and the next utterance as
its response. Multiple utterances of the context are connected using $[SEG]$.
For each dialogue scene, we sample multiple sets of such context and response
pairs. The sampled context and response will serve as input to the encoder and
decoder, respectively.

Then, we introduce the post-training design for Dial-MAE, as shown in Figure
1, We use an asymmetric encoder-decoder: a deep encoder to generate dialogue
context embedding, a shallow transformer-based decoder (eg., 1 layer or 2
layers) for response reconstruction. We apply a BERT-like encoder $Enc(.)$
with 12 layers whose input is the masked dialogue context. The deep encoder
has enough parameters to learn good dialogue representations, following the
common practice, we select the $[CLS]$ token’s final hidden state as the
dialogue context embedding. Considering the decoder side, its role is to
assist the encoder to learn a better semantic representation of the dialogue.
The input of the decoder $Dec(.)$ is the masked response as well as the
dialogue context embedding, and it reconstructs the masked response tokens
with the help of the context embedding. It is worth noting that we employ an
asymmetric masking operation(eg., 30% for encoder, 75% for decoder). On the
decoder side, an aggressive mask rate and fewer model parameters will force
its MLM task to rely more on the encoder’s context embedding, which helps the
encoder side learn better representations. Formally, we denote the dialogue
context as $c$ and the response as $r$. We apply random mask operation to
context to get $\tilde{c}$, denoting tokens replaced by $[MASK]$ in context as
$m_{eco}(c)$. Similarly, we apply a random masking operation with a higher
masking ratio for response to get $\tilde{r}$, denoting tokens replaced by
$[MASK]$ in response as $m_{dec}(r)$. The encoding process can be expressed
as:

$\displaystyle=Enc([CLS],\tilde{c})$ (3)
$\displaystyle=Dec(h_{cls}^{c},\tilde{r})$ (4)

On the encoder side, the original context is learned to be reconstructed by
optimizing the cross-entropy loss:

$\displaystyle L_{enc}=-\sum_{c_{i}\in
m_{eco}(c)}\log{p(c_{i}|Enc([CLS],\tilde{c}))}$ (5)

Differently, on the decoder side, the decoder reconstructs the original
response with the help of the context embedding $h_{cls}^{c}$. We formulate
this process as:

$\displaystyle L_{dec}=-\sum_{r_{i}\in
m_{dec}(r)}\log{p(r_{i}|Dec(h_{cls}^{c},\tilde{r}))}$ (6)

Then, we add the encoder and decoder losses to obtain a summed loss:

$\displaystyle L=L_{enc}+L_{dec}$ (7)

###  2.3 Fine-tuning for dialogue response selection

At the end of Dial-MAE post-training, fine-tuning is conducted on the
downstream dialogue response selection to verify the effectiveness of post-
training. As shown in Figure 2, in the fine-tuning stage, we only keep the
encoder and discard the decoder. The encoder weights are used to initialize a
dialogue context encoder $f_{c}$ and a response encoder $f_{r}$, respectively.

The dialogue consists of a context $c$ that includes multiple utterances and a
response $r$ with one utterance. After the dialogue context and response pass
through the encoder, the context vector and response vector are respectively
output. We train a dialogue response selection model using a contrastive
learning loss function.

$\displaystyle
L_{ft}=\frac{exp(d(c,r^{+}))}{exp(d(c,r^{+}))+\sum_{j}exp(d(c,r_{j}^{-}))}$
(8)

$r^{+}$ is the correct response corresponding to the dialogue context $c$.
$r^{-}$ represents negative samples within a mini-batch. At inference time, we
use the dot product $d(c,r)$ to measure the similarity between the context
vector and the response vector:

$\displaystyle d(c,r)=f_{c}(c)\cdot f_{r}(r)$ (9)

Figure 2: We discard the decoder, initialize the context encoder and response
encoder using the encoder part of Dial-MAE, and fine-tune using contrastive
learning. At inference time, We use a dot product to measure similarity.

##  3 Experiment

In this section, we first introduce our experimental details, including
datasets, evaluation metrics, baselines, post-training, and fine-tuning. Then
we introduce the experimental results.

###  3.1 Datasets

We tested our model on widely used benchmarks that include Ubuntu Corpus,
E-commerce Corpus, and the Douban Corpus. The statistics for the three
datasets are presented in Table 1.

  1. 1.

Ubuntu Corpus. Ubuntu IRC Corpus V1 Lowe et al. (2015) is a publicly available
domain-specific dialogue dataset. Each set of conversations has two
participants discussing how to troubleshoot Ubuntu systems.

  2. 2.

E-commerce Corpus. E-commerce Corpus Zhang et al. (2018) comprises genuine
conversations in Chinese between customers and customer service personnel,
collected from Taobao, a Chinese e-commerce platform.

  3. 3.

Douban Corpus. Douban Corpus Wu et al. (2017) is a multi-turn Chinese open-
domain dialogue dataset, it is collected from the Douban platform, a popular
Chinese social media.

Dataset | Ubuntu | E-commerce | Douban
---|---|---|---
| train | val | test | train | val | test | train | val | test
context-response pairs | 1M | 500k | 500k | 1M | 10k | 10k | 1M | 50k | 10k
pos : neg | 1:1 | 1:9 | 1:9 | 1:1 | 1:1 | 1:9 | 1:1 | 1:1 | 1.2:8.8
avg turns | 10.13 | 10.11 | 10.11 | 5.11 | 5.48 | 5.64 | 6.69 | 6.75 | 6.45

Table 1: Statistics related to data for the Ubuntu, E-commerce Corpus, and
Douban.

###  3.2 Evaluation Metric

We evaluated our model using R10@k, following previous studies Han et al.
(2021); Zhang et al. (2022), we evaluate our model using R10@k. The notation
R10@k represents recall, indicating that among ten possible responses, the
correct answer is included within the top k options. Different from the other
two datasets, the Douban Corpus includes multiple correct candidates for
context in the test set. For the Douban dataset, in addition to recall, we
also considered mean reciprocal rank (MRR), precision-at-one (P@1), and mean
average precision(MAP).

###  3.3 Baselines

We compared our generative post-trained model, Dial-MAE, with the following
previous models.

Single-turn matching models Lowe et al. (2015) proposes basic models with RNN,
Kadlec et al. (2015) proposes to use CNN and LSTM as encoders.

SMN Wu et al. (2017) breaks down the context-response pair into multiple pairs
of utterance-response. After matching each individual utterance and response,
the resulting matching vectors are then aggregated to obtain the final
matching score.

DUA Zhang et al. (2018) proposes to form a fine-grained context representation
through deep discourse aggregation.

DAM Zhou et al. (2018) proposes to calculate the matching score between the
context and response through transformer encoder-based model.

IoI Tao et al. (2019) leverages multiple interaction block chains to get more
interactions between utterances and responses.

ESIM Chen and Wang (2019) proposes a sequential matching model based only on
chain sequence for multi-turn response selection.

MSN Yuan et al. (2019) proposes a multi-hop selector network to select the
relevant utterances in the context for response matching.

RoBERTa-SS-DA Lu et al. (2020) proposes the speaker segmentation approach.

BERT-VFT Whang et al. (2020) proposes to use BERT’s MLM and NSP to post-train
the model and then fine-tuned to the response selection task.

SA-BERT Gu et al. (2020) adds speaker embedding to the model, in order to make
the model aware of the speaker change information.

UMSBERT+ Whang et al. (2021) proposes a multi-task learning framework
consisting of three tasks to maintain dialogue coherence.

BERT-SL Xu et al. (2021a) proposes four self-supervised auxiliary tasks in a
multi-task manner.

DR-BERT Lan et al. (2021) proposes a dense retrieval model to encode context
and response.

BERT-FP Han et al. (2021) proposes to classify the relationship between a
given utterance and a target utterance into more fine-grained labels.

BERT-TL Zhang et al. (2022) proposes two-level supervised contrastive learning
to enhance the discriminative ability.

###  3.4 Implementation Details

We first introduce the experimental setup for post-training, followed by the
experimental setup for contrastive learning.

#### Post-training.

Dial-MAE’s encoder is initialized with a pre-trained 12-layer BERT-base model,
while the decoder is initialized from scratch. Specifically, following the
previous works, for Chinese datasets, we employ bert-base-
chinese111https://huggingface.co/bert-base-chinese. For the English dataset,
we utilize the bert-base-uncased222https://huggingface.co/bert-base-uncased.
We pre-train the model using the AdamW optimizer for a maximum of 15k steps, a
global batch size of 1k, and a linear schedule with a warmup ratio of 0.1 on
all three datasets. We set the input sequence lengths to 256 and 64 for the
encoder and decoder, respectively. There are differences in some settings
between the English dataset and the Chinese datasets. In fact, for the Chinese
datasets E-commerce and Douban, we followed Cot-MAEWu et al. (2022)’s
parameter settings. The masking ratio of the encoder is 30%, the masking rate
of the decoder is 45%, the learning ratio is 1e-4, and the decoder is two
layers. Differently, for the English dataset Ubuntu, the masking ratio of the
encoder is 30%, the masking ratio of the decoder is 75%, and the decoder is
one layer. We also adjust the learning rate to 3e-4 to ensure the loss
function converges. We set a widely used random seed as 42 for
reproducibility. After post-training, we discard the decoder, only leaving the
encoder for fine-tuning.

#### Fine-tuning.

We fine-tune using contrastive learning on each dataset. During training, we
follow Lan et al. (2021) regarding every utterance in the dialogue sense as a
response and its previous utterances as a context. Our model is optimized by
AdamW optimizer, and the linear learning ratio scheduler is used. We tuned the
hypermeters of individual tasks on their development sets. There are
differences in some settings between the English dataset and the Chinese
datasets. For Ubuntu, we fine-tune for 5 epochs, the learning rate is set to
5e-5, and the batch size is set to 64. For Douban and E-commerce, we fine-tune
for 2 epochs, the learning rate is set to 1e-4, and the batch size is set to
128. We set a widely used random seed as 42 for reproducibility.

Models | Ubuntu | Douban | E-commerce
---|---|---|---
| R10@1 | R10@2 | R10@5 | MRR | MAP | P@1 | R10@1 | R10@2 | R10@5 | R10@1 | R10@2 | R10@5
TF-IDF Lowe et al. (2015) | 0.410 | 0.545 | 0.708 | 0.359 | 0.331 | 0.359 | 0.096 | 0.172 | 0.405 | 0.159 | 0.256 | 0.477
RNN Lowe et al. (2015) | 0.403 | 0.547 | 0.819 | 0.422 | 0.390 | 0.208 | 0.118 | 0.422 | 0.208 | 0.118 | 0.223 | 0.589
CNN Kadlec et al. (2015) | 0.549 | 0.684 | 0.896 | 0.440 | 0.417 | 0.226 | 0.121 | 0.252 | 0.647 | 0.328 | 0.515 | 0.792
LSTM Kadlec et al. (2015) | 0.638 | 0.784 | 0.949 | 0.537 | 0.485 | 0.320 | 0.187 | 0.343 | 0.720 | 0.365 | 0.536 | 0.828
SMN Wu et al. (2017) | 0.726 | 0.847 | 0.961 | 0.569 | 0.529 | 0.397 | 0.233 | 0.396 | 0.724 | 0.453 | 0.654 | 0.886
DUA Zhang et al. (2018) | 0.752 | 0.868 | 0.962 | 0.599 | 0.551 | 0.421 | 0.243 | 0.421 | 0.780 | 0.501 | 0.700 | 0.921
DAM Zhou et al. (2018) | 0.767 | 0.874 | 0.969 | 0.601 | 0.550 | 0.427 | 0.254 | 0.410 | 0.757 | 0.526 | 0.727 | 0.933
IOI Tao et al. (2019) | 0.796 | 0.894 | 0.974 | 0.621 | 0.573 | 0.444 | 0.269 | 0.451 | 0.786 | 0.563 | 0.768 | 0.950
ESIM Chen and Wang (2019) | 0.796 | 0.894 | 0.975 | - | - | - | - | - | - | 0.570 | 0.767 | 0.948
MSN Yuan et al. (2019) | 0.800 | 0.899 | 0.978 | 0.632 | 0.587 | 0.470 | 0.295 | 0.452 | 0.788 | 0.606 | 0.770 | 0.937
RoBERTa-SS-DA Lu et al. (2020) | 0.826 | 0.909 | 0.978 | 0.646 | 0.602 | 0.460 | 0.280 | 0.495 | 0.847 | 0.627 | 0.835 | 0.980
BERT-VFT Whang et al. (2020) | 0.855 | 0.928 | 0.985 | - | - | - | - | - | - | - | - | -
SA-BERT Gu et al. (2020) | 0.855 | 0.928 | 0.983 | 0.659 | 0.619 | 0.496 | 0.313 | 0.481 | 0.847 | 0.704 | 0.879 | 0.985
UMSBERT+ Whang et al. (2021) | 0.875 | 0.942 | 0.988 | 0.664 | 0.625 | 0.499 | 0.318 | 0.482 | 0.858 | 0.764 | 0.905 | 0.986
BERT-SL Xu et al. (2021b) | 0.884 | 0.946 | 0.990 | - | - | - | - | - | - | 0.776 | 0.919 | 0.991
DR-BERT Lan et al. (2021) $\clubsuit$ | 0.910 | 0.962 | 0.993 | - | - | - | - | - | - | - | - | -
BERT-FP Han et al. (2021) | 0.911 | 0.962 | 0.994 | 0.680 | 0.644 | 0.512 | 0.324 | 0.542 | 0.870 | 0.870 | 0.956 | 0.993
BERT-TL Zhang et al. (2022) | 0.910 | 0.962 | 0.993 | 0.718 | 0.675 | 0.564 | 0.367 | 0.571 | 0.874 | 0.927 | 0.974 | 0.997
BERT+ | 0.910 | 0.960 | 0.992 | 0.655 | 0.622 | 0.477 | 0.302 | 0.511 | 0.848 | 0.914 | 0.927 | 0.990
Dial-MAE | 0.918 | 0.964 | 0.993 | 0.670 | 0.637 | 0.500 | 0.320 | 0.533 | 0.857 | 0.930 | 0.977 | 0.997

Table 2: Main experiment results on E-commerce Corpus, Douban Corpus, and
Ubuntu Corpus. BERT+ means that only use BERT’s MLM task to do domain-adaptive
post-training, and then use contrastive learning to fine-tune. The best score
on a given dataset is marked in bold. $\clubsuit:$ According to the published
code, for Douban and E-commerce, they adjusted the hyperparameters on the test
set without cross-validation, we think the results are misleading, and this
part has been removed.

###  3.5 Results and Discussions

We show the main results in Table 2, which shows that Dial-MAE achieves new
state-of-the-art on the Ubuntu dataset and E-commerce dataset, and third-best
results on the Douban dataset. Compared to BERT-FP, our model achieved an
absolute improvement of 0.7%p in R10@1 on the Ubuntu Corpus and 6%p in R10@1
on the E-commerce. Compared to BERT-TL, our model achieves an absolute
improvement of 0.8%p in R10@1 on the Ubuntu Corpus and a slight improvement of
0.3%p in E-commerce. However, Dial-MAE performed worse than BERT-FP and BERT-
TL in Douban, and only achieved third place. We think this may be because
retrieval in the open domain data is difficult for the bi-encoder models.
BERT+ means that only use BERT’s MLM task to do domain-adaptive post-training,
and then use contrastive learning to fine-tune. In comparison to BERT+, Dial-
MAE achieved an absolute improvement in R10@1 by 0.8%p, 1.8%p, 1.6%p on Ubuntu
Corpus, Douban Corpus, and E-commerce Corpus, respectively. This suggests that
our custom generative post-training approach for dialogue retrieval models is
effective. We believe the improvement comes from two aspects. On the one hand,
the post-training method considers both the semantics of the tokens inside the
context and its response. On the other hand, the asymmetric encoder-decoder
structure with an asymmetric masking strategy facilitates post-training, which
forces the encoder to learn better dialogue embeddings.

###  3.6 Ablation Study

In this section, we analyze the experimental results to demonstrate the
effectiveness of the proposed Dial-MAE method. In the following experimental
analysis, due to high computing budgets, most experiments use Ubuntu Corpus.

Models | Ubuntu | Douban | E-commerce
---|---|---|---
| R10@1 | R10@2 | R10@5 | R10@1 | R10@2 | R10@5 | R10@1 | R10@2 | R10@5
BERT+ | 0.910 | 0.960 | 0.992 | 0.302 | 0.511 | 0.848 | 0.914 | 0.927 | 0.990
w/o Contrastive loss | 0.620 | 0.739 | 0.887 | 0.190 | 0.346 | 0.648 | 0.351 | 0.521 | 0.804
Dial-MAE | 0.918 | 0.964 | 0.993 | 0.320 | 0.533 | 0.857 | 0.930 | 0.977 | 0.997
w/o Contrastive loss | 0.783 | 0.867 | 0.950 | 0.264 | 0.408 | 0.704 | 0.483 | 0.639 | 0.853

Table 3: Ablation results on the test sets of the three benchmarks.

#### The Impact of Auxilary Network.

We remove the contrastive loss in BERT+ and Dial-MAE, then evaluate their
performance changes. As shown in Table 3, Dial-MAE achieved an absolute
improvement in R10@1 by 16.3%p, 7.4%p, 13.2%p on Ubuntu Corpus, Douban Corpus,
and E-commerce Corpus, respectively. This shows that our proposed generative
post-training can enhance the representation ability of dialogue embeddings
compared to only using MLM as a post-training task. We believe the gain comes
from our auxiliary network helping the encoder aggregate dialogue contextual
information. We then use contrastive learning to fine-tune the post-training
models, and the performance of the models can be further improved. We also
give the fine-tuning schedule on Ubuntu Corpus as shown in Figure 3, with the
accuracy steadily improving as the training time increases, and Dial-MAE
consistently outperforms BERT+. This result shows that both the contrastive
loss and the auxiliary MLM loss are crucial in our method.

Figure 3: Fine-tuning schedules on the dev set of Ubuntu Corpus. A longer
fine-tuning schedule gives a noticeable improvement. The performance of Dial-
MAE is always better than BERT+.

#### Impact of Mask Rate.

Enc | Dec |  |  | R10@1 | R10@2 | R10@5
---|---|---|---|---|---|---
0.15 | 0 |  |  | 91.0 | 96.0 | 99.2
0.15 | 0.45 |  |  | 91.5 | 96.2 | 99.3
0.15 | 0.75 |  |  | 91.5 | 96.3 | 99.3
0.30 | 0.45 |  |  | 91.7 | 96.5 | 99.3
0.30 | 0.75 |  |  | 91.9 | 96.5 | 99.3
0.30 | 0.90 |  |  | 91.6 | 96.4 | 99.3
0.45 | 0.75 |  |  | 91.8 | 96.4 | 99.4

Table 4: Impact of mask rate on the dev set of Ubuntu Corpus. "Enc" denotes
encoder, "Dec" denotes decoder.

Wu et al. (2022) find that using a larger mask rate in both the encoder and
decoder can enhance the performance of the contextual masking auto-encoder. As
shown in Table 4, in our experiments, we find that an aggressive mask rate
helps the learning of Dial-MAE. when the encoder mask rate equals 30%, and the
decoder mask rate equals 75%, Dial-MAE achieves the best performance. When the
encoder mask rate stays below 30%, the performance of Dial-MAE improves as the
decoder mask rate increases. When the encoder mask rate rises to 45%, Dial-
MAE’s performance declines slightly. We believe this is due to the encoder
doesn’t provide enough dialogue context semantic information when its mask
rate is too high. In addition, from the experimental results, no matter what
set of mask rates, Dial-MAE obviously exceeds the result of post-training for
MLM tasks alone, which proves the robustness of Dial-MAE.

#### Impact of Decoder Layer Number.

As shown in Figure 4, we further explore the impact of different decoder layer
numbers on Dial-MAE performance. we find that using only one layer of the
decoder yields the best results. Fewer decoder parameters can force the
auxiliary MLM task to rely more on dialogue context embeddings output by the
encoder. We believe that the more layers of the decoder, the stronger the
decoding ability, and the decoder’s dependence on context embedding will
decrease, leading to insufficient constraints on encoder training. In general,
no matter what set of layers, R@1 obviously exceeds the result of post-
training for MLM tasks alone as shown in Table 4, which proves the robustness
of Dial-MAE.

Figure 4: Impact of layer number on Ubuntu Corpus.

##  4 Conclusion

In this paper, we propose a generative post-training method tailored for
dialogue response, considering the semantics of dialogue context and its
corresponding responses. Precisely, we leverage a shallow decoder to force the
encoder output dialogue embeddings to be more expressive. Experimental results
show that our post-training method leads to considerable improvements,
achieving state-of-the-art on two benchmark datasets, we also demonstrate the
effectiveness of Dial-MAE through ablation experiments. We believe that
improving the representation ability of dialogue-dense vectors through
auxiliary networks is a promising research direction.

## References

  * Chen and Wang (2019) Qian Chen and Wen Wang. 2019.  Sequential attention-based network for noetic end-to-end response selection.  _arXiv preprint arXiv:1901.02609_.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pages 4171–4186. Association for Computational Linguistics.
  * Gao and Callan (2021) Luyu Gao and Jamie Callan. 2021.  Condenser: a pre-training architecture for dense retrieval.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021_ , pages 981–993. Association for Computational Linguistics.
  * Gu et al. (2020) Jia-Chen Gu, Tianda Li, Quan Liu, Zhen-Hua Ling, Zhiming Su, Si Wei, and Xiaodan Zhu. 2020.  Speaker-aware BERT for multi-turn response selection in retrieval-based chatbots.  In _CIKM ’20: The 29th ACM International Conference on Information and Knowledge Management, Virtual Event, Ireland, October 19-23, 2020_ , pages 2041–2044.
  * Han et al. (2021) Janghoon Han, Taesuk Hong, Byoungjae Kim, Youngjoong Ko, and Jungyun Seo. 2021.  Fine-grained post-training for improving retrieval-based dialogue systems.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021, Online, June 6-11, 2021_ , pages 1549–1558.
  * Kadlec et al. (2015) Rudolf Kadlec, Martin Schmid, and Jan Kleindienst. 2015.  Improved deep learning baselines for ubuntu corpus dialogs.  _CoRR_ , abs/1510.03753.
  * Lan et al. (2021) Tian Lan, Deng Cai, Yan Wang, Yixuan Su, Xian-Ling Mao, and Heyan Huang. 2021\.  Exploring dense retrieval for dialogue response selection.  _CoRR_ , abs/2110.06612.
  * Liu and Shao (2022) Zheng Liu and Yingxia Shao. 2022.  Retromae: Pre-training retrieval-oriented transformers via masked auto-encoder.  _arXiv preprint arXiv:2205.12035_.
  * Lowe et al. (2015) Ryan Lowe, Nissan Pow, Iulian Serban, and Joelle Pineau. 2015.  The ubuntu dialogue corpus: A large dataset for research in unstructured multi-turn dialogue systems.  In _Proceedings of the SIGDIAL 2015 Conference, The 16th Annual Meeting of the Special Interest Group on Discourse and Dialogue, 2-4 September 2015, Prague, Czech Republic_ , pages 285–294.
  * Lu et al. (2020) Junyu Lu, Xiancong Ren, Yazhou Ren, Ao Liu, and Zenglin Xu. 2020.  Improving contextual language models for response retrieval in multi-turn conversation.  In _Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval_ , pages 1805–1808.
  * Lu et al. (2021) Shuqi Lu, Di He, Chenyan Xiong, Guolin Ke, Waleed Malik, Zhicheng Dou, Paul Bennett, Tie-Yan Liu, and Arnold Overwijk. 2021.  Less is more: Pretrain a strong siamese encoder for dense text retrieval using a weak decoder.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021_ , pages 2780–2791. Association for Computational Linguistics.
  * Tao et al. (2019) Chongyang Tao, Wei Wu, Can Xu, Wenpeng Hu, Dongyan Zhao, and Rui Yan. 2019.  One time of interaction may not be enough: Go deep with an interaction-over-interaction network for response selection in dialogues.  In _Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers_ , pages 1–11. Association for Computational Linguistics.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pages 5998–6008.
  * Whang et al. (2020) Taesun Whang, Dongyub Lee, Chanhee Lee, Kisu Yang, Dongsuk Oh, and Heuiseok Lim. 2020.  An effective domain adaptive post-training method for BERT in response selection.  In _Interspeech 2020, 21st Annual Conference of the International Speech Communication Association, Virtual Event, Shanghai, China, 25-29 October 2020_ , pages 1585–1589.
  * Whang et al. (2021) Taesun Whang, Dongyub Lee, Dongsuk Oh, Chanhee Lee, Kijong Han, Dong-hun Lee, and Saebyeok Lee. 2021.  Do response selection models really know what’s next? utterance manipulation strategies for multi-turn response selection.  In _Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021_ , pages 14041–14049.
  * Wu et al. (2022) Xing Wu, Guangyuan Ma, Meng Lin, Zijia Lin, Zhongyuan Wang, and Songlin Hu. 2022\.  Contextual mask auto-encoder for dense passage retrieval.  _arXiv preprint arXiv:2208.07670_.
  * Wu et al. (2017) Yu Wu, Wei Wu, Chen Xing, Ming Zhou, and Zhoujun Li. 2017.  Sequential matching network: A new architecture for multi-turn response selection in retrieval-based chatbots.  In _Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics, ACL 2017, Vancouver, Canada, July 30 - August 4, Volume 1: Long Papers_ , pages 496–505.
  * Xu et al. (2021a) Ruijian Xu, Chongyang Tao, Daxin Jiang, Xueliang Zhao, Dongyan Zhao, and Rui Yan. 2021a.  Learning an effective context-response matching model with self-supervised tasks for retrieval-based dialogues.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 14158–14166.
  * Xu et al. (2021b) Ruijian Xu, Chongyang Tao, Daxin Jiang, Xueliang Zhao, Dongyan Zhao, and Rui Yan. 2021b.  Learning an effective context-response matching model with self-supervised tasks for retrieval-based dialogues.  In _Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021_ , pages 14158–14166.
  * Yuan et al. (2019) Chunyuan Yuan, Wei Zhou, Mingming Li, Shangwen Lv, Fuqing Zhu, Jizhong Han, and Songlin Hu. 2019.  Multi-hop selector network for multi-turn response selection in retrieval-based chatbots.  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 2019, Hong Kong, China, November 3-7, 2019_ , pages 111–120. Association for Computational Linguistics.
  * Zhang et al. (2022) Wentao Zhang, Shuang Xu, and Haoran Huang. 2022.  Two-level supervised contrastive learning for response selection in multi-turn dialogue.  _CoRR_ , abs/2203.00793.
  * Zhang et al. (2018) Zhuosheng Zhang, Jiangtong Li, Pengfei Zhu, Hai Zhao, and Gongshen Liu. 2018.  Modeling multi-turn conversation with deep utterance aggregation.  In _Proceedings of the 27th International Conference on Computational Linguistics, COLING 2018, Santa Fe, New Mexico, USA, August 20-26, 2018_ , pages 3740–3752. Association for Computational Linguistics.
  * Zhou et al. (2016) Xiangyang Zhou, Daxiang Dong, Hua Wu, Shiqi Zhao, Dianhai Yu, Hao Tian, Xuan Liu, and Rui Yan. 2016.  Multi-view response selection for human-computer conversation.  In _Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016_ , pages 372–381.
  * Zhou et al. (2018) Xiangyang Zhou, Lu Li, Daxiang Dong, Yi Liu, Ying Chen, Wayne Xin Zhao, Dianhai Yu, and Hua Wu. 2018.  Multi-turn response selection for chatbots with deep attention matching network.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July 15-20, 2018, Volume 1: Long Papers_ , pages 1118–1127. Association for Computational Linguistics.

##  Appendix A Qualitative Analysis

To qualitatively analyze our post-training method, as shown in Table 5, we
provide two examples. Dial-MAE can sort out the most appropriate response more
accurately than BERT+. The response sorted by BERT+ have some token overlap
with the dialogue context but are not semantically related. Compared with
BERT+, Dial-MAE can better understand dialogue semantics due to the joint
modeling of context and response through post-training. This further
demonstrates the effectiveness of our method.

Relevant | Model |  Rank 1st response
---|---|---
|  |  Context: A:I already have these disks in my system just want to migrate my current homefolder to the new drive. B:Mount your new disk to a temporarely mount point move or copy your home folder to the new disc after that delete you old homedir than mount the new disk to _path_.
✗ | BERT+ |  Response: Do you use the ubuntu desktop or server i am using the desktop on my laptop i have only started with the server a little before _number_ came out maybe before april.
✓ | Dial-MAE |  Response: I presume i use gparted to get the mountpoints correct or am i wrong
|  |  Context: A:I ve had firefox crash six times today usually it s once every other day gnome-panel has crashed four times. B:Anyone notice how the latest updates are making stuff in ubuntu crash.
✗ | BERT+ |  Response: Ok then i will try that can the whole cd be copied onto a usb i have a _number_ gb flash drive
✓ | Dial-MAE |  Response: So is there a way to figure out what s making something crash didn’t happen to me but my install is pretty basic.

Table 5: Examples of rank 1st response recalled by different models on the dev
set of the Ubuntu Corpus.