# Denoising Self-Attentive Sequential Recommendation

Huiyuan Chen<EMAIL_ADDRESS>Visa ResearchUSA ,  Yusan Lin<EMAIL_ADDRESS>Visa ResearchUSA ,  Menghai Pan<EMAIL_ADDRESS>Visa ResearchUSA ,  Lan Wang
<EMAIL_ADDRESS>Visa ResearchUSA ,  Chin-Chia Michael Yeh<EMAIL_ADDRESS>Visa
ResearchUSA ,  Xiaoting Li<EMAIL_ADDRESS>Visa ResearchUSA ,  Yan Zheng
<EMAIL_ADDRESS>Visa ResearchUSA ,  Fei Wang<EMAIL_ADDRESS>Visa
ResearchUSA and  Hao Yang<EMAIL_ADDRESS>Visa ResearchUSA

(2022)

###### Abstract.

Transformer-based sequential recommenders are very powerful for capturing both
short-term and long-term sequential item dependencies. This is mainly
attributed to their unique self-attention networks to exploit pairwise item-
item interactions within the sequence. However, real-world item sequences are
often noisy, which is particularly true for implicit feedback. For example, a
large portion of clicks do not align well with user preferences, and many
products end up with negative reviews or being returned. As such, the current
user action only depends on a subset of items, not on the entire sequences.
Many existing Transformer-based models use full attention distributions, which
inevitably assign certain credits to irrelevant items. This may lead to sub-
optimal performance if Transformers are not regularized properly.

Here we propose the Rec-denoiser model for better training of self-attentive
recommender systems. In Rec-denoiser, we aim to adaptively prune noisy items
that are unrelated to the next item prediction. To achieve this, we simply
attach each self-attention layer with a trainable binary mask to prune noisy
attentions, resulting in sparse and clean attention distributions. This
largely purifies item-item dependencies and provides better model
interpretability. In addition, the self-attention network is typically not
Lipschitz continuous and is vulnerable to small perturbations. Jacobian
regularization is further applied to the Transformer blocks to improve the
robustness of Transformers for noisy sequences. Our Rec-denoiser is a general
plugin that is compatible to many Transformers. Quantitative results on real-
world datasets show that our Rec-denoiser outperforms the state-of-the-art
baselines.

Sequential Recommendation, Sparse Transformer, Noise Analysis, Differentiable
Mask

††journalyear: 2022††copyright: acmcopyright††conference: Sixteenth ACM
Conference on Recommender Systems; September 18–23, 2022; Seattle, WA,
USA††booktitle: Sixteenth ACM Conference on Recommender Systems (RecSys ’22),
September 18–23, 2022, Seattle, WA, USA††price: 15.00††doi:
10.1145/3523227.3546788††isbn: 978-1-4503-9278-5/22/09††ccs: Information
systems Recommender systems

##  1\. Introduction

Sequential recommendation aims to recommend the next item based on a user’s
historical actions (Rendle et al., 2010; Hidasi et al., 2016; Ma et al., 2020;
Wang et al., 2021d; Wang et al., 2021a), e.g., to recommend a bluetooth
headphone after a user purchases a smart phone. Learning sequential user
behaviors is, however, challenging since a user’s choices on items generally
depend on both long-term and short-term preferences. Early Markov Chain models
(Rendle et al., 2010; He et al., 2017) have been proposed to capture short-
term item transitions by assuming that a user’s next decision is derived from
a few preceding actions, while neglecting long-term preferences. To alleviate
this limitation, many deep neural networks have been proposed to model the
entire users’ sequences and achieve great success, including recurrent neural
networks (Hidasi et al., 2016; Xu et al., 2019) and convolutional neural
networks (Yan et al., 2019; Yuan et al., 2019; Tang and Wang, 2018).

Recently, Transformers have shown promising results in various tasks, such as
machine translation (Vaswani et al., 2017). One key component of Transformers
is the self-attention network, which is capable of learning long-range
dependencies by computing attention weights between each pair of objects in a
sequence. Inspired by the success of Transformers, several self-attentive
sequential recommenders have been proposed and achieve the state-of-the-art
performance (Kang and McAuley, 2018; Sun et al., 2019; Wu et al., 2020b; Wu et
al., 2020a). For example, SASRec (Kang and McAuley, 2018) is the pioneering
framework to adopt self-attention network to learn the importance of items at
different positions. BERT4Rec (Sun et al., 2019) further models the
correlations of items from both left-to-right and right-to-left directions.
SSE-PT (Wu et al., 2020b) is a personalized Transformer model that provides
better interpretability of engagement patterns by introducing user embeddings.
LSAN (Li et al., 2021) adopts a novel twin-attention sequential framework,
which can capture both long-term and short-term user preference signals.
Recently, Transformers4Rec (de Souza Pereira Moreira et al., 2021) performs an
empirical analysis with broad experiments of various Transformer architectures
for the task of sequential recommendation.

Although encouraging performance has been achieved, the robustness of
sequential recommenders is far less studied in the literature. Many real-world
item sequences are naturally noisy, containing both true-positive and false-
positive interactions (Wang et al., 2021b, c; Chen et al., 2021b). For
example, a large portion of clicks do not align well with user preferences,
and many products end up with negative reviews or being returned. In addition,
there is no any prior knowledge about how a user’s historical actions should
be generated in online systems. Therefore, developing robust algorithms to
defend noise is of great significance for sequential recommendation.

Clearly, not every item in a sequence is aligned well with user preferences,
especially for implicit feedbacks (e.g., clicks, views, etc.) (Chen et al.,
2019). Unfortunately, the vanilla self-attention network is not Lipschitz
continuous111Roughly speaking, a function is Lipschitz continuous if changing
its input by a certain amount will not significantly change its output., and
is vulnerable to the quality of input sequences (Kim et al., 2021). Recently,
in the tasks of language modeling, people found that a large amount of BERT’s
attentions focus on less meaningful tokens, like ”[SEP]” and ”.”, which leads
to a misleading explanation (Clark et al., 2019). It is thus likely to obtain
sub-optimal performance if self-attention networks are not well regularized
for noisy sequences. We use the following example to further explain above
concerns.

Figure 1 illustrates an example of left-to-right sequential recommendation
where a user’s sequence contains some noisy or irrelevant items. For example,
a father may interchangeably purchase (phone, headphone, laptop) for his son,
and (bag, pant) for his daughter, resulting in a sequence: (phone, bag,
headphone, pant, laptop). In the setting of sequential recommendation, we
intend to infer the next item, e.g., laptop, based on the user’s previous
actions, e.g., (phone, bag, headphone, pant). However, the correlations among
items are unclear, and intuitively pant and laptop are neither complementary
nor compatible to each other, which makes the prediction untrustworthy. A
trustworthy model should be able to only capture correlated items while
ignoring these irrelevant items within sequences. Existing self-attentive
sequential models (e.g., SASRec (Kang and McAuley, 2018) and BERT4Rec (Sun et
al., 2019)) are insufficient to address noisy items within sequences. The
reason is that their full attention distributions are dense and would assign
certain credits to all items, including irrelevant items. This causes a lack
of focus and makes models less interpretable (Child et al., 2019; Zaheer et
al., 2020).

Figure 1. An illustrative example of sequential recommendation where a
sequence contains noisy or irrelevant items in left-to-right self-attention
networks.

To address the above issues, one straightforward strategy is to design sparse
Transformer architectures that sparsify the connections in the attention
layers, which have been actively investigated in language modeling tasks
(Child et al., 2019; Zaheer et al., 2020). Several representative models are
Star Transformer (Guo et al., 2019), Sparse Transformer (Child et al., 2019),
Longformer (Beltagy et al., 2020), and BigBird (Zaheer et al., 2020). These
sparse attention patterns could mitigate noisy issues and avoid allocating
credits to unrelated contents for the query of interest. However, these models
largely rely on pre-defined attention schemas, which lacks flexibility and
adaptability in practice. Unlike end-to-end training approaches, whether these
sparse patterns could generalize well to sequential recommendation remains
unknown and is still an open research question.

Contributions. In this work, we propose to design a denoising strategy, Rec-
Denoiser, for better training of self-attentive sequential recommenders. Our
idea stems from the recent findings that not all attentions are necessary and
simply pruning redundant attentions could further improve the performance
(Child et al., 2019; Zaheer et al., 2020; Correia et al., 2019; Sukhbaatar et
al., 2019; Yeh et al., 2022). Rather than randomly dropping out attentions, we
introduce differentiable masks to drop task-irrelevant attentions in the self-
attention layers, which can yield exactly zero attention scores for noisy
items.

The introduced sparsity in the self-attention layers has several benefits: 1)
Irrelevant attentions with parameterized masks can be learned to be dropped in
a data-driven way. Taking Figure 1 as an example, our Rec-denoiser would prune
the sequence (phone, bag, headphone) for pant, and (phone, bag, headphone,
pant) for laptop in the attention maps. Namely, we seek next item prediction
explicitly based on a subset of more informative items. 2) Our Rec-Denoiser
still takes full advantage of Transformers as it does not change their
architectures, but only the attention distributions. As such, Rec-Denoiser is
easy to implement and is compatible to any Transformers, making them less
complicated as well as improving their interpretability.

In our proposed Rec-Denoiser, there are two major challenges. First, the
discreteness of binary masks (i.e., $0$ is dropped while $1$ is kept) is,
however, intractable in the back-propagation. To remedy this issue, we relax
the discrete variables with a continuous approximation through probabilistic
reparameterization (Jang et al., 2017). As such, our differentiable masks can
be trained jointly with original Transformers in an end-to-end fashion. In
addition, the scaled dot-product attention is not Lipschitz continuous and is
thus vulnerable to input perturbations (Kim et al., 2021). In this work,
Jacobian regularization (Jakubovitz and Giryes, 2018; Hoffman et al., 2019) is
further applied to the entire Transformer blocks, to improve the robustness of
Transformers for noisy sequences. Experimental results on real-world benchmark
datasets demonstrate the effectiveness and robustness of the proposed Rec-
Denoiser. In summary, our contributions are:

  * •

We introduce the idea of denoising item sequences for better of training self-
attentive sequential recommenders, which greatly reduces the negative impacts
of noisy items.

  * •

We present a general Rec-Denoiser framework with differentiable masks that can
achieve sparse attentions by dynamically pruning irrelevant information,
leading to better model performance.

  * •

We propose an unbiased gradient estimator to optimize the binary masks, and
apply Jacobian regularization on the gradients of Transformer blocks to
further improve its robustness.

  * •

The experimental results demonstrate significant improvements that Rec-
Denoiser brings to self-attentive recommenders ($5.05\%\sim 19.55\%$
performance gains), as well as its robustness against input perturbations.

##  2\. Related Work

In this section, we briefly review the related work on sequential
recommendation and sparse Transformers. We also highlight the differences
between the existing efforts and ours.

###  2.1. Sequential Recommendation

Leveraging sequences of user-item interactions is crucial for sequential
recommendation. User dynamics can be caught by Markov Chains for inferring the
conditional probability of an item based on the previous items (Rendle et al.,
2010; He et al., 2017). More recently, growing efforts have been dedicated to
deploying deep neural networks for sequential recommendation such as recurrent
neural networks (Hidasi et al., 2016; Xu et al., 2019), convolutional neural
networks (Yan et al., 2019; Yuan et al., 2019; Tang and Wang, 2018), memory
networks (Chen et al., 2018; Huang et al., 2018), and graph neural networks
(Wu et al., 2019; Chang et al., 2021; Chen et al., 2022). For example, GRU4Rec
(Hidasi et al., 2016) employs a gated recurrent unit to study temporal
behaviors of users. Caser (Tang and Wang, 2018) learns sequential patterns by
using convolutional filters on local sequences. MANN (Chen et al., 2018)
adopts memory-augmented neural networks to model user historical records. SR-
GNN (Wu et al., 2019) converts session sequences into graphs and uses graph
neural networks to capture complex item-item transitions.

Transformer-based models have shown promising potential in sequential
recommendation (Kang and McAuley, 2018; Lian et al., 2020; Sun et al., 2019;
Chen et al., 2021a; Li et al., 2020; Wu et al., 2020b; Wu et al., 2020a; Liu
et al., 2021), due to their ability of modeling arbitrary dependencies in a
sequence. For example, SASRec (Kang and McAuley, 2018) first adopts self-
attention network to learn the importance of items at different positions. In
the follow-up studies, several Transformer variants have been designed for
different scenarios by adding bidirectional attentions (Sun et al., 2019),
time intervals (Li et al., 2020), personalization (Wu et al., 2020b),
importance sampling (Lian et al., 2020), and sequence augmentation (Liu et
al., 2021). However, very few studies pay attention to the robustness of self-
attentive recommender models. Typically, users’ sequences contain lots of
irrelevant items since they may subsequently purchase a series of products
with different purposes (Wang et al., 2021b). As such, the current user action
only depends on a subset of items, not on the entire sequences. However, the
self-attention module is known to be sensitive to noisy sequences (Kim et al.,
2021), which may lead to sub-optimal generalization performance. In this
paper, we aim to reap the benefits of Transformers while denoising the noisy
item sequences by using learnable masks in an end-to-end fashion.

###  2.2. Sparse Transformer

Recently, many lightweight Transformers seek to achieve sparse attention maps
since not all attentions carry important information in the self-attention
layers (Guo et al., 2019; Child et al., 2019; Beltagy et al., 2020; Zaheer et
al., 2020; Kitaev et al., 2019). For instance, Reformer (Kitaev et al., 2019)
computes attentions based on locality-sensitive hashing, leading to lower
memory consumption. Star Transformer (Guo et al., 2019) replaces the fully-
connected structure of self-attention with a star-shape topology. Sparse
Transformer (Child et al., 2019) and Longformer (Beltagy et al., 2020) achieve
sparsity by using various sparse patterns, such as diagonal sliding windows,
dilated sliding windows, local and global sliding windows. BigBird (Zaheer et
al., 2020) uses random and several fixed patterns to build sparse blocks. It
has been shown that these sparse attentions can obtain the state-of-the-art
performance and greatly reduce computational complexity. However, many of them
rely on fixed attention schemas that lack flexibility and require tremendous
engineering efforts.

Another line of work is to use learnable attention distributions (Correia et
al., 2019; Sukhbaatar et al., 2019; Peters et al., 2019; Malaviya et al.,
2018). Mostly, they calculate attention weights with variants of sparsemax
that replaces the softmax normalization in the self-attention networks. This
allows to produce both sparse and bounded attentions, yielding a compact and
interpretable set of alignments. Our Rec-denoiser is related to this line of
work. Instead of using sparsemax, we design a trainable binary mask for the
self-attention network. As a result, our proposed Rec-denoiser can
automatically determine which self-attention connections should be deleted or
kept in a data-driven way.

##  3\. Problem and Background

In this section, we first formulate the problem of sequential recommendation,
and then revisit several self-attentive models. We further discuss the
limitations of the existing work.

###  3.1. Problem Setup

In sequential recommendation, let $\mathcal{U}$ be a set of users,
$\mathcal{I}$ a set of items, and
$\mathcal{S}=\\{\mathcal{S}^{1},\mathcal{S}^{2},\ldots,\mathcal{S}^{|\mathcal{U}|}\\}$
a set of users’ actions. We use
$\mathcal{S}^{u}=(S_{1}^{u},S_{2}^{u},\ldots,S_{|S^{u}|}^{u})$ to denote a
sequence of items for user $u\in\mathcal{U}$ in a chronological order, where
$S_{t}^{u}\in\mathcal{I}$ is the item that user $u$ has interacted with at
time $t$, and $|S^{u}|$ is the length of sequence.

Given the interaction history $\mathcal{S}^{u}$, sequential recommendation
seeks to predict the next item $S_{|S^{u}+1|}^{u}$ at time step $|S^{u}+1|$.
During the training process (Kang and McAuley, 2018; Sun et al., 2019), it
will be convenient to regard the model’s input as
$(S_{1}^{u},S_{2}^{u},\ldots,S_{|S^{u}-1|}^{u})$ and its expected output is a
shifted version of the input sequence:
$(S_{2}^{u},S_{3}^{u},\ldots,S_{|S^{u}|}^{u})$.

###  3.2. Self-attentive Recommenders

Owing to the ability of learning long sequences, Transformer architectures
(Vaswani et al., 2017) have been widely used in sequential recommendation,
like SASRec (Kang and McAuley, 2018), BERT4Rec (Sun et al., 2019), and
TiSASRec (Li et al., 2020). Here we briefly introduce the design of SASRec and
discuss its limitations.

####  3.2.1. Embedding Layer

Transformer-based recommenders maintain an item embedding table
$\mathbf{T}\in\mathbb{R}^{|\mathcal{I}|\times d}$, where $d$ is the size of
the embedding. For each sequence
$(S_{1}^{u},S_{2}^{u},\ldots,S_{|S^{u}-1|}^{u})$, it can be converted into a
fixed-length sequence $(s_{1},s_{2},\ldots,s_{n})$, where $n$ is the maximum
length (e.g., keeping the most recent $n$ items by truncating or padding
items). The embedding for $(s_{1},s_{2},\ldots,s_{n})$ is denoted as
$\mathbf{E}\in\mathbb{R}^{n\times d}$, which can be retrieved from the table
$\mathbf{T}$. To capture the impacts of different positions, one can inject a
learnable positional embedding $\mathbf{P}\in\mathbb{R}^{n\times d}$ into the
input embedding as:

(1) $\hat{\mathbf{E}}=\mathbf{E}+\mathbf{P},$

where $\hat{\mathbf{E}}\in\mathbb{R}^{n\times d}$ is an order-aware embedding,
which can be directly fed to any Transformer-based models.

####  3.2.2. Transformer Block

A Transformer block consists of a self-attention layer and a point-wise feed-
forward layer.

Self-attention Layer: The key component of Transformer block is the self-
attention layer that is highly efficient to uncover sequential dependencies in
a sequence (Vaswani et al., 2017). The scaled dot-product attention is a
popular attention kernel:

(2)
$\textsf{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\textsf{softmax}\left(\frac{\mathbf{QK}^{T}}{\sqrt{d}}\right)\mathbf{V},$

where
$\textsf{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})\in\mathbb{R}^{n\times
d}$ is the output item representations;
$\mathbf{Q}=\hat{\mathbf{E}}\mathbf{W}^{Q}$,
$\mathbf{K}=\hat{\mathbf{E}}\mathbf{W}^{K}$, and
$\mathbf{V}=\hat{\mathbf{E}}\mathbf{W}^{V}$ are the queries, keys and values,
respectively;
$\\{\mathbf{W}^{Q},\mathbf{W}^{K},\mathbf{W}^{V}\\}\in\mathbb{R}^{d\times d}$
are three projection matrices; and $\sqrt{d}$ is the scale factor to produce a
softer attention distribution. In sequential recommendation, one can utilize
either left-to-right unidirectional attentions222This can be achieved by
forbidding all links between $\mathbf{Q}_{i}$ and $\mathbf{K}_{j}$ for all
$j>i$. (e.g., SASRec (Kang and McAuley, 2018) and TiSASRec (Li et al., 2020))
or bidirectional attentions (e.g., BERT4Rec (Sun et al., 2019)) to predict the
next item. Moreover, one can apply $H$ attention functions in parallel to
enhance expressiveness:
$\mathbf{H}\leftarrow\textsf{MultiHead}(\hat{\mathbf{E}})$ (Vaswani et al.,
2017).

Point-wise Feed-forward Layer: As the self attention layer is built on linear
projections, we can endow the non-linearity by introducing a point-wise feed-
forward layers:

(3)
$\displaystyle\mathbf{F_{i}}=\textsf{FFN}(\mathbf{H}_{i})=\textsf{ReLU}(\mathbf{H}_{i}\mathbf{W}^{(1)}+\mathbf{b}^{(1)})\mathbf{W}^{(2)}+\mathbf{b}^{(2)},$

where $\mathbf{W}^{(*)}\in\mathbb{R}^{d\times d}$,
$\mathbf{b}^{(*)}\in\mathbb{R}^{d}$ are the learnable weights and bias.

In practice, it is usually beneficial to learn hierarchical item dependencies
by stacking more Transformer blocks. Also, one can adopt the tricks of
residual connection, dropout, and layer normalization for stabilizing and
accelerating the training. Here we simply summarize the output of $L$
Transformer blocks as:
$\mathbf{F}^{(L)}\leftarrow\text{Transformer}(\hat{\mathbf{E}})$.

####  3.2.3. Learning Objective

After stacked $L$ Transformer blocks, one can predict the next item (given the
first $t$ items) based on $\mathbf{F}_{t}^{(L)}$. In this work, we use inner
product to predict the relevance of item $i$ as:

$r_{i,t}=\langle\mathbf{F}_{t}^{(L)},\mathbf{T}_{i}\rangle,$

where $\mathbf{T}_{i}\in\mathbb{R}^{d}$ is the embedding of item $i$. Recall
that the model inputs a sequence $\mathbf{s}=(s_{1},s_{2},\ldots,s_{n})$ and
its desired output is a shifted version of the same sequence
$\mathbf{o}=(o_{1},o_{2},\ldots,o_{n})$, we can adopt the binary cross-entropy
loss as:

(4)
$\displaystyle\mathcal{L}_{BCE}=-\sum_{\mathcal{S}^{u}\in\mathcal{S}}\sum_{t=1}^{n}\left[\log(\sigma(r_{o_{t},t}))+\log(1-\sigma(r_{o^{\prime}_{t},t}))\right]+\alpha\cdot\|\mathbf{\Theta}\|_{F}^{2},$

where $\mathbf{\Theta}$ is the model parameters, $\alpha$ is the regularizer
to prevent over-fitting, $o^{\prime}_{t}\not\in\mathcal{S}^{u}$ is a negative
sample corresponding to $o_{t}$, and $\sigma(\cdot)$ is the sigmoid function.
More details can be found in SASRec (Kang and McAuley, 2018) and BERT4Rec (Sun
et al., 2019).

###  3.3. The Noisy Attentions Problem

Despite the success of SASRec and its variants, we argue that they are
insufficient to address the noisy items in sequences. The reason is that the
full attention distributions (e.g., Eq. (2)) are dense and would assign
certain credits to irrelevant items. This complicates the item-item
dependencies, increases the training difficulty, and even degrades the model
performance. To address this issue, several attempts have been proposed to
manually define sparse attention schemas in language modeling tasks (Guo et
al., 2019; Child et al., 2019; Beltagy et al., 2020; Zaheer et al., 2020).
However, these fixed sparse attentions cannot adapt to the input data (Correia
et al., 2019), leading to sub-optimal performance.

On the other hand, several dropout techniques are specifically designed for
Transformers to keep only a small portion of attentions, including LayerDrop
(Fan et al., 2020), DropHead (Zhou et al., 2020), and UniDrop (Wu et al.,
2021). Nevertheless, these randomly dropout approaches are susceptible to
bias: the fact that attentions can be dropped randomly does not mean that the
model allows them to be dropped, which may lead to over-aggressive pruning
issues. In contrast, we propose a simple yet effective data-driven method to
mask out irrelevant attentions by using differentiable masks.

##  4\. Rec-Denoiser

In this section, we present our Rec-Denoiser that consists of two parts:
differentiable masks for self-attention layers and Jacobian regularization for
Transformer blocks.

###  4.1. Differentiable Masks

The self-attention layer is the cornerstone of Transformers to capture long-
range dependencies. As shown in Eq. (2), the softmax operator assigns a non-
zero weight to every item. However, full attention distributions may not
always be advantageous since they may cause irrelevant dependencies,
unnecessary computation, and unexpected explanation. We next put forward
differentiable masks to address this concern.

####  4.1.1. Learnable Sparse Attentions

Not every item in a sequence is aligned well with user preferences in the same
sense that not all attentions are strictly needed in self-attention layers.
Therefore, we attach each self-attention layer with a trainable binary mask to
prune noisy or task-irrelevant attentions. Formally, for the $l$-th self-
attention layer in Eq. (2), we introduce a binary matrix
$\mathbf{Z}^{(l)}\in\\{0,1\\}^{n\times n}$, where $\mathbf{Z}^{(l)}_{u,v}$
denotes whether the connection between query $u$ and key $v$ is present. As
such, the $l$-th self-attention layer becomes:

(5) $\begin{gathered}\mathbf{A}^{(l)}=\textsf{
softmax}\left(\frac{\mathbf{Q}^{(l)}{\mathbf{K}^{(l)}}^{T}}{\sqrt{d}}\right),\\\
\mathbf{M}^{(l)}=\mathbf{A}^{(l)}\odot\mathbf{Z}^{(l)},\\\
\textsf{Attention}(\mathbf{Q}^{(l)},\mathbf{K}^{(l)},\mathbf{V}^{(l)})=\mathbf{M}^{(l)}\mathbf{V}^{(l)},\end{gathered}$

where $\mathbf{A}^{(l)}$ is the original full attentions, $\mathbf{M}^{(l)}$
denotes the sparse attentions, and $\odot$ is the element-wise product.
Intuitively, the mask $\mathbf{Z}^{(l)}$ (e.g., $1$ is kept and $0$ is
dropped) requires minimal changes to the original self-attention layer. More
importantly, they are capable of yielding exactly zero attention scores for
irrelevant dependencies, resulting in better interpretability. The idea of
differentiable masks is not new. In the language modeling, differentiable
masks have been shown to be very powerful to extract short yet sufficient
sentences, which achieves better performance (Bastings et al., 2019; De Cao et
al., 2020).

One way to encourage sparsity of $\mathbf{M}^{(l)}$ is to explicitly penalize
the number of non-zero entries of $\mathbf{Z}^{(l)}$, for $1\leq l\leq L$, by
minimizing:

(6)
$\mathcal{R}_{M}=\sum_{l=1}^{L}\left\|\mathbf{Z}^{(l)}\right\|_{0}=\sum_{l=1}^{L}\sum_{u=1}^{n}\sum_{v=1}^{n}\mathbb{I}\left[\mathbf{Z}^{(l)}_{u,v}\neq
0\right],$

where $\mathbb{I}[c]$ is an indicator that is equal to $1$ if the condition
$c$ holds and $0$ otherwise; and $\|\cdot\|_{0}$ denotes the $L_{0}$ norm that
is able to drive irrelevant attentions to be exact zeros.

However, there are two challenges for optimizing $\mathbf{Z}^{(l)}$: non-
differentiability and large variance. $L_{0}$ is discontinuous and has zero
derivatives almost everywhere. Additionally, there are $2^{n^{2}}$ possible
states for the binary mask $\mathbf{Z}^{(l)}$ with large variance. Next, we
propose an efficient estimator to solve this stochastic binary optimization
problem.

####  4.1.2. Efficient Gradient Computation

Since $\mathbf{Z}^{(l)}$ is jointly optimized with the original Transformer-
based models, we combine Eq. (4) and Eq. (6) into one unified objective:

(7)
$\mathcal{L}(\mathbf{Z},\mathbf{\Theta})=\mathcal{L}_{BCE}(\\{\mathbf{A}^{(l)}\odot\mathbf{Z}^{(l)}\\},\mathbf{\Theta})+\beta\cdot\sum_{l=1}^{L}\sum_{u=1}^{n}\sum_{v=1}^{n}\mathbb{I}\left[\mathbf{Z}^{(l)}_{u,v}\neq
0\right],$

where $\beta$ controls the sparsity of masks and we denote $\mathbf{Z}$ as
$\mathbf{Z}:=\\{\mathbf{Z}^{(1)},\cdots,\mathbf{Z}^{(L)}\\}$. We further
consider each $\mathbf{Z}^{(l)}_{u,v}$ is drawn from a Bernoulli distribution
parameterized by $\mathbf{\Pi}^{(l)}_{u,v}$ such that
$\mathbf{Z}^{(l)}_{u,v}\sim\text{Bern}(\mathbf{\Pi}^{(l)}_{u,v})$ (Louizos et
al., 2019). As the parameter $\mathbf{\Pi}^{(l)}_{u,v}$ is jointly trained
with the downstream tasks, a small value of $\mathbf{\Pi}^{(l)}_{u,v}$
suggests that the attention $\mathbf{A}^{(l)}_{u,v}$ is more likely to be
irrelevant, and could be removed without side effects. By doing this, Eq. (7)
becomes:

(8) $\displaystyle{\mathcal{L}}(\mathbf{Z},\mathbf{\Theta})$
$\displaystyle=\mathbb{E}_{\mathbf{Z}\sim\prod^{L}_{l=1}\text{Bern}(\mathbf{Z}^{(l)};\mathbf{\Pi}^{(l)})}\left[\mathcal{L}_{BCE}(\mathbf{Z},\mathbf{\Theta})\right]+\beta\cdot\sum_{l=1}^{L}\sum_{u=1}^{n}\sum_{v=1}^{n}\mathbf{\Pi}^{(l)}_{u,v},$

where $\mathbb{E}(\cdot)$ is the expectation. The regularization term is now
continuous, but the first term $\mathcal{L}_{BCE}(\mathbf{Z},\mathbf{\Theta})$
still involves the discrete variables $\mathbf{Z}^{(l)}$. One can address this
issue by using existing gradient estimators, such as REINFORCE (Williams,
1992) and Straight Through Estimator (Bengio et al., 2013), etc. These
approaches, however, suffer from either biased gradients or high variance.
Alternatively, we directly optimize discrete variables using the recently
proposed augment-REINFORCE-merge (ARM) (Dimitriev and Zhou, 2021; Yin and
Zhou, 2019; Dong et al., 2020), which is unbiased and has low variance.

In particular, we adopt the reparameterization trick (Jang et al., 2017),
which reparameterizes $\mathbf{\Pi}^{(l)}_{u,v}\in[0,1]$ to a deterministic
function $g(\cdot)$ with parameter $\mathbf{\Phi}^{(l)}_{u,v}$, such that:

(9) $\mathbf{\Pi}^{(l)}_{u,v}=g(\mathbf{\Phi}^{(l)}_{u,v}),$

since the deterministic function $g(\cdot)$ should be bounded within $[0,1]$,
we simply choose the standard sigmoid function as our deterministic function:
$g(x)={1}/({1+e^{-x}})$. As such, the second term in Eq. (8) becomes
differentiable with the continuous function $g(\cdot)$. We next present the
ARM estimator for computing the gradients of binary variables in the first
term of Eq. (8) (Dimitriev and Zhou, 2021; Yin and Zhou, 2019; Dong et al.,
2020).

According to Theorem 1 in ARM (Yin and Zhou, 2019), we can compute the
gradients for Eq. (8) as:

(10)
$\displaystyle\nabla^{ARM}_{\mathbf{\Phi}}{\mathcal{L}}(\mathbf{\Phi},\mathbf{\Theta})$
$\displaystyle=\mathbb{E}_{\mathbf{U}\sim\prod_{l=1}^{L}\operatorname{Uni}(\mathbf{U}^{(l)};0,1)}[(\mathcal{L}_{BCE}(\mathbb{I}[\mathbf{U}>g(-\mathbf{\Phi})],\mathbf{\Theta})-\mathcal{L}_{BCE}(\mathbb{I}[\mathbf{U}<g(\mathbf{\Phi})],\mathbf{\Theta}))\cdot(\mathbf{U}-\frac{1}{2})]+\beta\cdot\nabla_{\mathbf{\Phi}}g(\mathbf{\Phi}),$

where $\text{Uni}(0,1)$ denotes the Uniform distribution within $[0,1]$, and
$\mathcal{L}_{BCE}(\mathbb{I}[\mathbf{U}>g(-\mathbf{\Phi})],\mathbf{\Theta})$
is the cross-entropy loss obtained by setting the binary masks
$\mathbf{Z}^{(l)}$ to $1$ if $\mathbf{U}^{(l)}>g(-\mathbf{\Phi}^{(l)})$ in the
forward pass, and $0$ otherwise. The same strategy is applied to
$\mathcal{L}_{BCE}(\mathbb{I}[\mathbf{U}<g(\mathbf{\Phi})],\mathbf{\Theta})$.
Moreover, ARM is an unbiased estimator due to the linearity of expectations.

Note that we need to evaluate $\mathcal{L}_{BCE}(\cdot)$ twice to compute
gradients in Eq. (10). Given the fact that $u\sim\text{Uni}(0,1)$ implies
$(1-u)\sim\text{Uni}(0,1)$, we can replace $\mathbf{U}$ with $(1-\mathbf{U})$
in the indicator function inside
$\mathcal{L}_{BCE}(\mathbb{I}[\mathbf{U}>g(-\mathbf{\Phi})],\mathbf{\Theta})$:

$\mathbb{I}[\mathbf{U}>g(-\mathbf{\Phi})]\Leftrightarrow\mathbb{I}[1-\mathbf{U}>g(-\mathbf{\Phi})]\Leftrightarrow\mathbb{I}[\mathbf{U}<g(\mathbf{\Phi})].$

To this end, we can further reduce the complexity by considering the variants
of ARM – Augment-Reinforce (AR) (Yin and Zhou, 2019):

(11)
$\displaystyle\nabla^{AR}_{\mathbf{\Phi}}{\mathcal{L}}(\mathbf{\Phi},\mathbf{\Theta})=\mathbb{E}_{\mathbf{U}\sim\prod_{l=1}^{L}\operatorname{Uni}(\mathbf{U}^{(l)};0,1)}\left[\mathcal{L}_{BCE}(\mathbb{I}[\mathbf{U}<g(\mathbf{\Phi})],\mathbf{\Theta})\cdot(1-2\mathbf{U})\right]+\beta\cdot\nabla_{\mathbf{\Phi}}g(\mathbf{\Phi}),$

where only requires one-forward pass. The gradient estimator
$\nabla^{AR}_{\mathbf{\Phi}}{\mathcal{L}}(\mathbf{\Phi},\mathbf{\Theta})$ is
still unbiased but may pose higher variance, comparing to
$\nabla^{ARM}_{\mathbf{\Phi}}{\mathcal{L}}(\mathbf{\Phi},\mathbf{\Theta})$. In
the experiments, we can trade off the variance of the estimator with
complexity.

In the training stage, we update
$\nabla_{\mathbf{\Phi}}{\mathcal{L}}(\mathbf{\Phi},\mathbf{\Theta})$ (either
Eq. (10) or Eq. (11)) and
$\nabla_{\mathbf{\Theta}}{\mathcal{L}}(\mathbf{\Phi},\mathbf{\Theta})$333$\nabla_{\mathbf{\Theta}}{\mathcal{L}}(\mathbf{\Phi},\mathbf{\Theta})$
is the original optimization for Transformers, and one can learn
$\mathbf{\Theta}$ via standard Stochastic Gradient Descent. during the back
propagation. In the inference stage, we can use the expectation of
$\mathbf{Z}^{(l)}_{u,v}\sim\text{Bern}(\mathbf{\Pi}^{(l)}_{u,v})$ as the mask
in Eq. (5), i.e.,
$\mathbb{E}(\mathbf{Z}^{(l)}_{u,v})=\mathbf{\Pi}^{(l)}_{u,v}=g(\boldsymbol{\Phi}^{(l)}_{u,v})$.
Nevertheless, this will not yield a sparse attention $\mathbf{M}^{(l)}$ since
the sigmoid function is smooth unless the hard sigmoid function is used in Eq.
(9). Here we simply clip those values $g(\boldsymbol{\Phi}^{(l)}_{u,v})\leq
0.5$ to zeros such that a sparse attention matrix is guaranteed and the
corresponding noisy attentions are eventually eliminated.

###  4.2. Jacobian Regularization

As recently proved by (Kim et al., 2021), the standard dot-product self-
attention is not Lipschitz continuous and is vulnerable to the quality of
input sequences. Let $f^{(l)}$ be the $l$-th Transformer block (Sec 3.2.2)
that contains both a self-attention layer and a point-wise feed-forward layer,
and $\mathbf{x}$ be the input. We can measure the robustness of the
Transformer block using the residual error:
$f^{(l)}(\mathbf{x}+\boldsymbol{\epsilon})-f^{(l)}(\mathbf{x})$, where
$\boldsymbol{\epsilon}$ is a small perturbation vector and the norm of
$\boldsymbol{\epsilon}$ is bounded by a small scalar $\delta$, i.e.,
$\|\boldsymbol{\epsilon}\|_{2}\leq\delta$. Following the Taylor expansion, we
have:

$f^{(l)}_{i}(\mathbf{x}+\boldsymbol{\epsilon})-f^{(l)}_{i}(\mathbf{x})\approx\left[\frac{\partial
f^{(l)}_{i}(\mathbf{x})}{\partial\mathbf{x}}\right]^{\top}\boldsymbol{\epsilon}.$

Let $\mathbf{J}^{(l)}(\mathbf{x})$ represent the Jacobian matrix at
$\mathbf{x}$ where $\mathbf{J}^{(l)}_{ij}=\frac{\partial
f^{(l)}_{i}(\mathbf{x})}{\partial x_{j}}$. Then, we set
$\mathbf{J}^{(l)}_{i}(\mathbf{x})=\frac{\partial
f^{(l)}_{i}(\mathbf{x})}{\partial\mathbf{x}}$ to denote the $i$-th row of
$\mathbf{J}^{(l)}(\mathbf{x}).$ According to Hölder’s
inequality444https://en.wikipedia.org/wiki/Hölder’s_inequality, we have:

$\|f^{(l)}_{i}(\mathbf{x}+\boldsymbol{\epsilon})-f^{(l)}_{i}(\mathbf{x})\|_{2}\approx\|\mathbf{J}^{(l)}_{i}(\mathbf{x})^{\top}\boldsymbol{\epsilon}\|_{2}\leq\|\mathbf{J}^{(l)}_{i}(\mathbf{x})^{\top}\|_{2}\|\boldsymbol{\epsilon}\|_{2}.$

Above inequality indicates that regularizing the $L_{2}$ norm on the Jacobians
enforces a Lipschitz constraint at least locally, and the residual error is
strictly bounded. Thus, we propose to regularize Jacobians with Frobenius norm
for each Transformer block as:

(12) $\mathcal{R}_{J}=\sum_{l=1}^{L}\|\mathbf{J}^{(l)}\|_{F}^{2}.$

Importantly, $\|\mathbf{J}^{(l)}\|_{F}^{2}$ can be approximated via various
Monte-Carlo estimators (Hutchinson, 1989; Meyer et al., 2021). In this work,
we adopt the classical Hutchinson estimator (Hutchinson, 1989). For each
Jocobian matrix $\mathbf{J}^{(l)}\in\mathbb{R}^{n\times n}$, we have:

$\|\mathbf{J}^{(l)}\|_{F}^{2}=\operatorname{Tr}\left(\mathbf{J}^{(l)}{\mathbf{J}^{(l)}}^{\top}\right)=\mathbb{E}_{\boldsymbol{\eta}\in\mathcal{N}\left(0,\mathbf{I}_{n}\right)}\left[\left\|\boldsymbol{\eta}^{\top}\mathbf{J}^{(l)}\right\|_{F}^{2}\right],$

where $\boldsymbol{\eta}\in\mathcal{N}(0,\mathbf{I}_{n})$ is the normal
distribution vector. We further make use of random projections to compute the
norm of Jacobians $\mathcal{R}_{J}$ and its gradient
$\nabla_{\mathbf{\Theta}}\mathcal{R}_{J}(\mathbf{\Theta})$ (Hoffman et al.,
2019), which significantly reduces the running time in practice.

Input: The training sequences $\mathcal{S}$, the number of attention layers
$L$, the number of heads $H$, and the regularization coefficients $\alpha$,
$\beta$, and $\gamma$.

1

2Initialize model parameters $\mathbf{\Theta}$, mask parameters
$\mathbf{\Phi}$;

3

4for  _each mini-batch_ do

5 for  _$l\leftarrow 1$ to $L$_ do

6 Compute the full attention $\mathbf{A}^{(l)}$ in Eq. (5)

7 Draw $\mathbf{U}^{(l)}\sim\text{Uniform}(0,1)$;

8 Compute the mask
$\mathbf{Z}^{(l)}=\mathbb{I}[\mathbf{U}^{(l)}<g(\mathbf{\Phi}^{(l)})]$;

9 Compute the sparse attention
$\mathbf{M}^{(l)}=\mathbf{A}^{(l)}\odot\mathbf{Z}^{(l)}$;

10 Feed $\mathbf{M}^{(l)}$ to the rest modules of Transformer;

11

12 end for

13 Compute the loss $\mathcal{L}_{Rec-Denoiser}$ in Eq. (13);

14 Update the model parameters $\mathbf{\Theta}$ and mask parameters
$\mathbf{\Phi}$ (Eq. (11)) via Stochastic Gradient Descent;

15

16 end for

Output: A well-trained Transformer.

Algorithm 1 Rec-denoiser with AR estimator

###  4.3. Optimization

####  4.3.1. Joint Training

Putting together loss in Eq. (4), Eq. (6), and Eq. (12), the overall objective
of Rec-Denoiser is:

(13) $\mathcal{L}_{Rec-
Denoiser}=\mathcal{L}_{BCE}+\beta\cdot\mathcal{R}_{M}+\gamma\cdot\mathcal{R}_{J},$

where $\beta$ and $\gamma$ are regularizers to control the sparsity and
robustness of self-attention networks, respectively. Algorithm 1 summarizes
the overall training of Rec-Denoiser with the AR estimator.

Lastly, it is worth mentioning that our Rec-Denoiser is compatible to many
Transformer-based sequential recommender models since our differentiable masks
and gradient regularizations will not change their main architectures. If we
simply set all masks $\mathbf{Z}^{(l)}$ to be all-ones matrix and
$\beta=\gamma=0$, our model boils down to their original designs. If we
randomly set subset of masks $\mathbf{Z}^{(l)}$ to be zeros, it is equivalent
to structured Dropout like LayerDrop (Fan et al., 2020), DropHead (Zhou et
al., 2020). In addition, our Rec-Denoiser can work together with linearized
self-attention networks (Katharopoulos et al., 2020; Zhou et al., 2021) to
further reduce the complexity of attentions. We leave this extension in the
future.

####  4.3.2. Model Complexity

The complexity of Rec-Denoiser comes from three parts: a basic Transformer,
differentiable masks, and Jacobian regularization. The complexity of basic
Transformer keeps the same as SASRec (Kang and McAuley, 2018) or BERT4Rec (Sun
et al., 2019). The complexity of differentiable masks requires either one-
forward pass (e.g., AR with high variance) or two-forward pass (e.g., ARM with
low variance) of the model. In sequential recommenders, the number of
Transformer blocks is often very small (e.g., $L=2$ in SASRec (Kang and
McAuley, 2018) and BERT4Rec (Sun et al., 2019) ). It is thus reasonable to use
the ARM estimator without heavy computations. Besides, we compare the
performance of AR and ARM estimators in Sec 5.3.

Moreover, the random project techniques are surprisingly efficient to compute
the norms of Jacobians (Hoffman et al., 2019). As a result, the overall
computational complexity remains the same order as the original Transformers
during the training. However, during the inference, our attention maps are
very sparse, which enables much faster feed-forward computations.

##  5\. Experiments

Here we present our empirical results. Our experiments are designed to answer
the following research questions:

  * •

RQ1: How effective is the proposed Rec-Denoiser compared to the state-of-the-
art sequential recommenders?

  * •

RQ2: How can Rec-Denoiser reduce the negative impacts of noisy items in a
sequence?

  * •

RQ3: How do different components (e.g., differentiable masks and Jacobian
regularization) affect the overall performance of Rec-Denoiser?

###  5.1. Experimental Setting

####  5.1.1. Dataset

We evaluate our models on five benchmark datasets:
Movielens555https://grouplens.org/datasets/movielens/20m/,
Amazon666https://jmcauley.ucsd.edu/data/amazon/ (we choose the three commonly
used categories: Beauty, Games, and Movies&TV), and
Steam777https://cseweb.ucsd.edu/~jmcauley/datasets.html (Li et al., 2020).
Their statistics are shown in Table 1. Among different datasets, MovieLens is
the most dense one while Beauty has the fewest actions per user. We use the
same procedure as (Kang and McAuley, 2018; Rendle et al., 2010; Li et al.,
2020) to perform preprocessing and split data into train/valid/test sets,
i.e., the last item of each user’s sequence for testing, the second-to-last
for validation, and the remaining items for training.

Table 1. The statistics of five benchmark datasets. Dataset | #Users | #Items | Avg actions/user | #Actions
---|---|---|---|---
MovieLens | 6,040 | 3,416 | 163.5 | 0.987M
Beauty | 51,369 | 19,369 | 4.39 | 0.225M
Games | 30,935 | 12,111 | 6.46 | 0.2M
Movies&TV | 40,928 | 37,564 | 25.55 | 1.05M
Steam | 114,796 | 8,648 | 7.58 | 0.87M

####  5.1.2.  Baselines

Here we include two groups of baselines. The first group includes general
sequential methods (Sec 5.2): 1) FPMC (Rendle et al., 2010): a mixture of
matrix factorization and first-order Markov chains model; 2) GRU4Rec (Hidasi
et al., 2016): a RNN-based method that models user action sequences; 3) Caser
(Tang and Wang, 2018): a CNN-based framework that captures high-order
relationships via convolutional operations; 4) SASRec (Kang and McAuley,
2018): a Transformer-based method that uses left-to-right self-attention
layers; 5) BERT4Rec (Sun et al., 2019): an architecture that is similar to
SASRec, but using bidirectional self-attention layers; 6) TiSASRec (Li et al.,
2020): a time-aware self-attention model that further considers the relative
time intervals between any two items; 7) SSE-PT (Wu et al., 2020b): a
framework that introduces personalization into self-attention layers; 8) Rec-
Denoiser: our proposed Rec-Denoiser that can choose any self-attentive models
as its backbone.

The second group contains sparse Transformers (Sec 5.3): 1) Sparse Transformer
(Child et al., 2019): it uses a fixed attention pattern, where only specific
cells summarize previous locations in the attention layers; 2) $\alpha$-entmax
sparse attention (Correia et al., 2019): it simply replaces softmax with
$\alpha$-entmax to achieve sparsity. Note that we do not compare with other
popular sparse Transformers like Star Transformer (Guo et al., 2019),
Longformer (Beltagy et al., 2020), and BigBird (Zaheer et al., 2020). These
Transformers are specifically designed for thousands of tokens or longer in
the language modeling tasks. We leave their explorations for recommendations
in the future. We also do not compare with LayerDrop (Fan et al., 2020) and
DropHead (Zhou et al., 2020) since the number of Transformer blocks and heads
are often very small (e.g., $L=2$ in SARRec) in sequential recommendation.
Other sequential architectures like memory networks (Chen et al., 2018; Huang
et al., 2018) and graph neural networks (Wu et al., 2019; Chang et al., 2021)
have been outperformed by the above baselines, we simply omit these baselines
and focus on Transformer-based models. The goal of experiments is to see
whether the proposed differentiable mask techniques can reduce the negative
impacts of noisy items in the self-attention layers.

####  5.1.3. Evaluation metrics

For easy comparison, we adopt two common Top-N metrics, Hit$@N$ and NDCG$@N$
(with default value $N=10$), to evaluate the performance of sequential models
(Kang and McAuley, 2018; Li et al., 2020; Sun et al., 2019). Typically,
Hit$@N$ counts the rates of the ground-truth items among top-$N$ items, while
NDCG$@N$ considers the position and assigns higher weights to higher
positions. Following the work (Kang and McAuley, 2018; Li et al., 2020), for
each user, we randomly sample $100$ negative items, and rank these items with
the ground-truth item. We calculate Hit$@10$ and NDCG$@10$ based on the
rankings of these $101$ items.

####  5.1.4. Parameter settings

For all baselines, we initialize the hyper-parameters as the ones suggested by
their original work. They are then well tuned on the validation set to achieve
optimal performance. The final results are conducted on the test set. We
search the dimension size of items within $\\{10,20,30,40,50\\}$. As our Rec-
Denoiser is a general plugin, we use the same hyper-parameters as the basic
Transformers, e.g., number of Transformer blocks, batch size, learning rate in
Adam optimizer, etc. According to Table 1, we set the maximum length of item
sequence $n=50$ for dense datasets MovieLens and Movies&TV, and $n=25$ for
sparse datasets Beauty, Games, and Steam. In addition, we set the number of
Transformer blocks $L=2$, and the number of heads $H=2$ for self-attentive
models. For Rec-Denoiser, two extra regularizers $\beta$ and $\gamma$ are both
searched within $\\{10^{-1},10^{-2},\ldots,10^{-5}\\}$. We choose ARM
estimator due to the shallow structures of self-attentive recommenders.

Table 2. Overall Performance of different models. ”RI” denotes the relative
improvement of Rec-Denoisers over their backbones. The best performing results
are boldfaced, and the second best ones are underlined.

Dataset | MovieLens | Beauty | Games | Movies&TV | Steam
---|---|---|---|---|---
Metrics | Hit@10 | NDCG@10 | Hit@10 | NDCG@10 | Hit@10 | NDCG@10 | Hit@10 | NDCG@10 | Hit@10 | NDCG@10
FPMC (Rendle et al., 2010) | 0.7478 | 0.4889 | 0.2810 | 0.1792 | 0.5231 | 0.3410 | 0.4806 | 0.3174 | 0.6012 | 0.4084
GRU4Rec (Hidasi et al., 2016) | 0.5582 | 0.3383 | 0.2123 | 0.1205 | 0.2943 | 0.1939 | 0.4210 | 0.2343 | 0.4184 | 0.2687
Caser (Tang and Wang, 2018) | 0.7213 | 0.4672 | 0.2670 | 0.1531 | 0.4315 | 0.2652 | 0.4987 | 0.3120 | 0.7137 | 0.4810
SASRec (Kang and McAuley, 2018) | 0.7434 | 0.5012 | 0.4345 | 0.2765 | 0.6748 | 0.4622 | 0.6521 | 0.4093 | 0.7723 | 0.5514
SASRec+Denoiser | 0.7980 | 0.5610 | 0.4783 | 0.3025 | 0.7391 | 0.5439 | 0.7056 | 0.4718 | 0.8345 | 0.5946
+RI (%) | 7.34% | 11.93% | 10.08% | 9.40% | 9.53% | 17.68% | 8.20% | 15.27% | 5.05% | 7.83%
BERT4Rec (Sun et al., 2019) | 0.7549 | 0.5245 | 0.4528 | 0.3013 | 0.6812 | 0.4815 | 0.6701 | 0.4216 | 0.7901 | 0.5641
BERT4Rec+Denoiser | 0.8045 | 0.5814 | 0.4883 | 0.3348 | 0.7415 | 0.5310 | 0.7212 | 0.4875 | 0.8410 | 0.6223
+RI (%) | 6.57% | 10.85% | 7.84% | 11.12% | 8.85% | 10.28% | 7.63% | 15.63% | 6.45% | 10.32%
TiSASRec (Li et al., 2020) | 0.7365 | 0.5164 | 0.4532 | 0.2911 | 0.6613 | 0.4517 | 0.6412 | 0.4034 | 0.7704 | 0.5517
TiSASRec+Denoiser | 0.7954 | 0.5582 | 0.4962 | 0.3312 | 0.7331 | 0.4984 | 0.6914 | 0.4671 | 0.8414 | 0.6320
+RI (%) | 7.80% | 8.10% | 9.49% | 13.78% | 10.86% | 10.34% | 7.93% | 15.79% | 9.22% | 14.56%
SSE-PT (Wu et al., 2020b) | 0.7413 | 0.5041 | 0.4326 | 0.2731 | 0.6810 | 0.4713 | 0.6378 | 0.4127 | 0.7641 | 0.5703
SSE-PT+Denoiser | 0.8010 | 0.5712 | 0.4952 | 0.3265 | 0.7396 | 0.5152 | 0.6972 | 0.4571 | 0.8310 | 0.6133
+RI (%) | 8.01% | 13.31% | 14.47% | 19.55% | 8.61% | 11.68% | 9.31% | 10.76% | 8.76% | 13.51%

###  5.2. Overall Performance (RQ1)

Table 2 presents the overall recommendation performance of all methods on the
five datasets. Our proposed Rec-denoisers consistently obtain the best
performance for all datasets. Additionally, we have the following
observations:

  * •

The self-attentive sequential models (e.g., SASRec, BERT4Rec, TiSASRec, and
SSE-PT) generally outperform FPMC, GRU4Rec, and Caser with a large margin,
verifying that the self-attention networks have good ability of capture long-
range item dependencies for the task of sequential recommendation.

  * •

Comparing the original SASRec and its variants BERT4Rec, TiSASRec and SSE-PT,
we find that the self-attentive models can gets benefit from incorporating
additional information such as bi-directional attentions, time intervals, and
user personalization. Such auxiliary information is important to interpret the
dynamic behaviors of users.

  * •

The relative improvements of Rec-denoisers over their backbones are
significant for all cases. For example, SASRec+Denoiser has on average
$8.04\%$ improvement with respect to Hit$@10$ and over $12.42\%$ improvements
with respect to NDCG$@10$. Analogously, BERT4Rec+Denoiser outperforms the
vanilla BERT4Rec by average $7.47\%$ in Hit$@10$ and $11.64\%$ in NDCG$@10$.
We also conduct the significant test between Rec-denoisers and their
backbones, where all $p$-values$<1e^{-6}$, showing that the improvements of
Rec-denoisers are statistically significant in all cases.

Figure 2.  Overall performance on the training data are corrupted by synthetic
noises.

These improvements of our proposed models are mainly attributed to the
following reasons: 1) Rec-denoisers inherit full advantages of the self-
attention networks as in SASRec, BERT4Rec, TiSASRec, and SSE-PT; 2) Through
differentiable masks, irrelevant item-item dependencies are removed, which
could largely reduce the negative impacts of noisy data; 3) Jacobian
regularization enforces the smoothness of gradients, limiting quick changes of
the output against input perturbations. In general, smoothness improves the
generalization of sequential recommendation. Overall, the experimental results
demonstrate the superiority of our Rec-Denoisers.

###  5.3. Robustness to Noises (RQ2)

As discussed before, the observed item sequences often contain some noisy
items that are uncorrelated to each other. Generally, the performance of self-
attention networks is sensitive to noisy input. Here we analyze how robust our
training strategy is for noisy sequences. To achieve this, we follow the
strategy (Ma et al., 2020) that corrupts the training data by randomly
replacing a portion of the observed items in the training set with uniformly
sampled items that are not in the validation or test set. We range the ratio
of the corrupted training data from $0\%$ to $25\%$. We only report the
results of SASRec and SASRec-Denoiser in terms of Hit$@10$. The performance of
other self-attentive models is the same and omitted here due to page
limitations. In addition, we compare with two recent sparse Transformers:
Sparse Transformer (Child et al., 2019) and $\alpha$-entmax sparse attention
(Correia et al., 2019).

All the simulated experiments are repeated five times and the average results
are shown in Figure 2. Clearly, the performance of all models degrades with
the increasing noise ratio. We observe that our Rec-denoiser (use either ARM
or AR estimators) consistently outperforms $\alpha$-entmax and Sparse
Transformer under different ratios of noise on all datasets. $\alpha$-entmax
heavily relies on one trainable parameter $\alpha$ to filter out the noise,
which may be over tuned during the training, while Sparse Transformer adopts a
fixed attention pattern, which may lead to uncertain results, especially for
short item sequences like Beauty and Games.

In contrast, our differentaible masks have much more flexibility to adapt to
noisy sequences. The Jacobian regularization further encourages the smoothness
of our gradients, leading to better generalization. From the results, the AR
estimator performs better than $\alpha$-entmax but worse than ARM. This result
is expected since ARM has much low variance. In summary, both ARM and AR
estimators are able to reduce the negative impacts of noisy sequences, which
could improve the robustness of self-attentive models.

###  5.4. Study of Rec-Denoiser (RQ3)

We further investigate the parameter sensitivity of Rec-Denoiser. For the
number of blocks $L$ and the number of heads $H$, we find that self-attentive
models typically benefit from small values (e.g., $H,L\leq 4$), which is
similar to (Li et al., 2021; Sun et al., 2019). In this section, we mainly
study the following hyper-parameters: 1) the maximum length $n$, 2) the
regularizers $\beta$ and $\gamma$ to control the sparsity and smoothness. Here
we only study the SASRec and SASRec-Denoiser due to page limitations.

####  5.4.1.  Maximum length $n$

Figure 3 shows the Hit$@10$ for maximum length $n$ from $20$ to $80$ while
keeping other optimal hyper-parameters unchanged. We only test on the densest
and sparsest datasets: MovieLeans and Beauty. Intuitively, the larger sequence
we have, the larger probability that the sequence contains noisy items. We
observed that our SASRec-Denoiser improves the performance dramatically with
longer sequences. This demonstrates that our design is more suitable for
longer inputs, without worrying about the quality of sequences.

Figure 3. Effect of maximum length $n$ on ranking performance (Hit$@10$).

####  5.4.2.  The regularizers $\beta$ and $\gamma$

There are two major regularization parameters $\beta$ and $\gamma$ for
sparsity and gradient smoothness, respectively. Figure 4 shows the performance
by changing one parameter while fixing the other as $0.01$. As can be seen,
our performance is relatively stable with respect to different settings. In
the experiments, the best performance can be achieved at $\beta=0.01$ and
$\gamma=0.001$ for the MovieLens dataset.

Figure 4. Effect of regularizers $\beta$ and $\gamma$ on ranking performance
(Hit$@10$).

##  6\. CONCLUSION AND FUTURE WORK

In this work, we propose Rec-Denoiser to adaptively eliminate the negative
impacts of the noisy items for self-attentive recommender systems. The
proposed Rec-Denoiser employs differentiable masks for the self-attention
layers, which can dynamically prune irrelevant information. To further tackle
the vulnerability of self-attention networks to small perturbations, Jacobian
regularization is applied to the Transformer blocks to improve the robustness.
Our experimental results on multiple real-world sequential recommendation
tasks illustrate the effectiveness of our design.

Our proposed Rec-Denoiser framework (e.g., differentiable masks and Jacobian
regularization) can be easily applied to any Transformer-based models in many
tasks besides sequential recommendation. In the future, we will continue to
demonstrate the contributions of our design in many real-world applications.

## References

  * (1)
  * Bastings et al. (2019) Jasmijn Bastings, Wilker Aziz, and Ivan Titov. 2019.  Interpretable Neural Predictions with Differentiable Binary Variables. In _ACL_. 2963–2977.
  * Beltagy et al. (2020) Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020\.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ (2020).
  * Bengio et al. (2013) Yoshua Bengio, Nicholas Léonard, and Aaron Courville. 2013\.  Estimating or propagating gradients through stochastic neurons for conditional computation.  _arXiv preprint arXiv:1308.3432_ (2013).
  * Chang et al. (2021) Jianxin Chang, Chen Gao, Yu Zheng, Yiqun Hui, Yanan Niu, Yang Song, Depeng Jin, and Yong Li. 2021\.  Sequential Recommendation with Graph Neural Networks. In _SIGIR_. 378–387.
  * Chen et al. (2021a) Huiyuan Chen, Yusan Lin, Fei Wang, and Hao Yang. 2021a.  Tops, bottoms, and shoes: building capsule wardrobes via cross-attention tensor network. In _RecSys_. 453–462.
  * Chen et al. (2021b) Huiyuan Chen, Lan Wang, Yusan Lin, Chin-Chia Michael Yeh, Fei Wang, and Hao Yang. 2021b.  Structured graph convolutional networks with stochastic masks for recommender systems. In _SIGIR_. 614–623.
  * Chen et al. (2022) Huiyuan Chen, Chin-Chia Michael Yeh, Fei Wang, and Hao Yang. 2022.  Graph Neural Transport Networks with Non-local Attentions for Recommender Systems. In _WWW_. 1955–1964.
  * Chen et al. (2019) Qiwei Chen, Huan Zhao, Wei Li, Pipei Huang, and Wenwu Ou. 2019.  Behavior sequence transformer for e-commerce recommendation in alibaba. In _DLP-KDD_.
  * Chen et al. (2018) Xu Chen, Hongteng Xu, Yongfeng Zhang, Jiaxi Tang, Yixin Cao, Zheng Qin, and Hongyuan Zha. 2018.  Sequential recommendation with user memory networks. In _WSDM_. 108–116.
  * Child et al. (2019) Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019\.  Generating long sequences with sparse transformers.  _arXiv preprint arXiv:1904.10509_ (2019).
  * Clark et al. (2019) Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. 2019.  What Does BERT Look at? An Analysis of BERT’s Attention. In _ACL Workshop BlackboxNLP_. 276–286.
  * Correia et al. (2019) Gonçalo M Correia, Vlad Niculae, and André FT Martins. 2019.  Adaptively Sparse Transformers. In _EMNLP_. 2174–2184.
  * De Cao et al. (2020) Nicola De Cao, Michael Sejr Schlichtkrull, Wilker Aziz, and Ivan Titov. 2020.  How Do Decisions Emerge across Layers in Neural Models? Interpretation with Differentiable Masking. In _EMNLP_. 3243–3255.
  * de Souza Pereira Moreira et al. (2021) Gabriel de Souza Pereira Moreira, Sara Rabhi, Jeong Min Lee, Ronay Ak, and Even Oldridge. 2021\.  Transformers4Rec: Bridging the Gap between NLP and Sequential/Session-Based Recommendation. In _RecSys_.
  * Dimitriev and Zhou (2021) Aleksandar Dimitriev and Mingyuan Zhou. 2021.  ARMS: Antithetic-REINFORCE-Multi-Sample Gradient for Binary Variables. In _ICML_. 2717–2727.
  * Dong et al. (2020) Zhe Dong, Andriy Mnih, and George Tucker. 2020.  DisARM: An Antithetic Gradient Estimator for Binary Latent Variables. In _NeurIPS_.
  * Fan et al. (2020) Angela Fan, Edouard Grave, and Armand Joulin. 2020\.  Reducing Transformer Depth on Demand with Structured Dropout. In _ICLR_.
  * Guo et al. (2019) Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. 2019\.  Star-Transformer. In _NAACL-HLT_. 1315–1325.
  * He et al. (2017) Ruining He, Wang-Cheng Kang, and Julian McAuley. 2017\.  Translation-based recommendation. In _RecSys_. 161–169.
  * Hidasi et al. (2016) Balázs Hidasi, Alexandros Karatzoglou, Linas Baltrunas, and Domonkos Tikk. 2016.  Session-based Recommendations with Recurrent Neural Networks. In _ICLR_.
  * Hoffman et al. (2019) Judy Hoffman, Daniel A Roberts, and Sho Yaida. 2019\.  Robust learning with jacobian regularization.  _arXiv preprint arXiv:1908.02729_ (2019).
  * Huang et al. (2018) Jin Huang, Wayne Xin Zhao, Hongjian Dou, Ji-Rong Wen, and Edward Y Chang. 2018.  Improving sequential recommendation with knowledge-enhanced memory networks. In _SIGIR_. 505–514.
  * Hutchinson (1989) Michael F Hutchinson. 1989\.  A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines.  _Communications in Statistics-Simulation and Computation_ (1989).
  * Jakubovitz and Giryes (2018) Daniel Jakubovitz and Raja Giryes. 2018.  Improving dnn robustness to adversarial attacks using jacobian regularization. In _ECCV_. 514–529.
  * Jang et al. (2017) Eric Jang, Shixiang Gu, and Ben Poole. 2017.  Categorical Reparameterization with Gumbel-Softmax. In _ICLR_.
  * Kang and McAuley (2018) Wang-Cheng Kang and Julian McAuley. 2018.  Self-attentive sequential recommendation. In _ICDM_. 197–206.
  * Katharopoulos et al. (2020) Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. 2020.  Transformers are rnns: Fast autoregressive transformers with linear attention. In _ICML_. 5156–5165.
  * Kim et al. (2021) Hyunjik Kim, George Papamakarios, and Andriy Mnih. 2021\.  The lipschitz constant of self-attention. In _ICML_. 5562–5571.
  * Kitaev et al. (2019) Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2019\.  Reformer: The Efficient Transformer. In _ICLR_.
  * Li et al. (2020) Jiacheng Li, Yujie Wang, and Julian McAuley. 2020.  Time interval aware self-attention for sequential recommendation. In _WSDM_. 322–330.
  * Li et al. (2021) Yang Li, Tong Chen, Peng-Fei Zhang, and Hongzhi Yin. 2021\.  Lightweight Self-Attentive Sequential Recommendation. In _CIKM_. 967–977.
  * Lian et al. (2020) Defu Lian, Yongji Wu, Yong Ge, Xing Xie, and Enhong Chen. 2020.  Geography-aware sequential location recommendation. In _KDD_. 2009–2019.
  * Liu et al. (2021) Zhiwei Liu, Ziwei Fan, Yu Wang, and Philip S. Yu. 2021\.  Augmenting Sequential Recommendation with Pseudo-Prior Items via Reversely Pre-Training Transformer. In _SIGIR_. 1608–1612.
  * Louizos et al. (2019) Christos Louizos, Max Welling, and Diederik P Kingma. 2019\.  Learning Sparse Neural Networks through $L\\_0$ Regularization. In _ICLR_.
  * Ma et al. (2020) Jianxin Ma, Chang Zhou, Hongxia Yang, Peng Cui, Xin Wang, and Wenwu Zhu. 2020\.  Disentangled self-supervision in sequential recommenders. In _KDD_. 483–491.
  * Malaviya et al. (2018) Chaitanya Malaviya, Pedro Ferreira, and André FT Martins. 2018\.  Sparse and Constrained Attention for Neural Machine Translation. In _ACL_. 370–376.
  * Meyer et al. (2021) Raphael A Meyer, Cameron Musco, Christopher Musco, and David P Woodruff. 2021.  Hutch++: Optimal stochastic trace estimation. In _SOSA_. 142–155.
  * Peters et al. (2019) Ben Peters, Vlad Niculae, and André FT Martins. 2019\.  Sparse Sequence-to-Sequence Models. In _ACL_. 1504–1519.
  * Rendle et al. (2010) Steffen Rendle, Christoph Freudenthaler, and Lars Schmidt-Thieme. 2010\.  Factorizing personalized markov chains for next-basket recommendation. In _WWW_. 811–820.
  * Sukhbaatar et al. (2019) Sainbayar Sukhbaatar, Édouard Grave, Piotr Bojanowski, and Armand Joulin. 2019.  Adaptive Attention Span in Transformers. In _ACL_. 331–335.
  * Sun et al. (2019) Fei Sun, Jun Liu, Jian Wu, Changhua Pei, Xiao Lin, Wenwu Ou, and Peng Jiang. 2019.  BERT4Rec: Sequential recommendation with bidirectional encoder representations from transformer. In _CIKM_. 1441–1450.
  * Tang and Wang (2018) Jiaxi Tang and Ke Wang. 2018.  Personalized top-n sequential recommendation via convolutional sequence embedding. In _WSDM_. 565–573.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need. In _NeurIPS_. 5998–6008.
  * Wang et al. (2021a) Jianling Wang, Kaize Ding, and James Caverlee. 2021a.  Sequential Recommendation for Cold-start Users with Meta Transitional Learning. In _SIGIR_. 1783–1787.
  * Wang et al. (2021b) Wenjie Wang, Fuli Feng, Xiangnan He, Liqiang Nie, and Tat-Seng Chua. 2021b.  Denoising implicit feedback for recommendation. In _WSDM_. 373–381.
  * Wang et al. (2021c) Wenjie Wang, Fuli Feng, Xiangnan He, Hanwang Zhang, and Tat-Seng Chua. 2021c.  Clicks can be cheating: Counterfactual recommendation for mitigating clickbait issue. In _SIGIR_. 1288–1297.
  * Wang et al. (2021d) Zhenlei Wang, Jingsen Zhang, Hongteng Xu, Xu Chen, Yongfeng Zhang, Wayne Xin Zhao, and Ji-Rong Wen. 2021d.  Counterfactual data-augmented sequential recommendation. In _SIGIR_. 347–356.
  * Williams (1992) Ronald J Williams. 1992\.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine learning_ 8, 3-4 (1992), 229–256.
  * Wu et al. (2020a) Jibang Wu, Renqin Cai, and Hongning Wang. 2020a.  Déjà vu: A contextualized temporal attention mechanism for sequential recommendation. In _WWW_. 2199–2209.
  * Wu et al. (2020b) Liwei Wu, Shuqing Li, Cho-Jui Hsieh, and James Sharpnack. 2020b.  SSE-PT: Sequential recommendation via personalized transformer. In _RecSys_. 328–337.
  * Wu et al. (2019) Shu Wu, Yuyuan Tang, Yanqiao Zhu, Liang Wang, Xing Xie, and Tieniu Tan. 2019\.  Session-based recommendation with graph neural networks. In _AAAI_. 346–353.
  * Wu et al. (2021) Zhen Wu, Lijun Wu, Qi Meng, Yingce Xia, Shufang Xie, Tao Qin, Xinyu Dai, and Tie-Yan Liu. 2021\.  UniDrop: A Simple yet Effective Technique to Improve Transformer without Extra Cost. In _NAACL-HLT_. 3865–3878.
  * Xu et al. (2019) Chengfeng Xu, Pengpeng Zhao, Yanchi Liu, Jiajie Xu, Victor S Sheng S. Sheng, Zhiming Cui, Xiaofang Zhou, and Hui Xiong. 2019\.  Recurrent convolutional neural network for sequential recommendation. In _WWW_. 3398–3404.
  * Yan et al. (2019) An Yan, Shuo Cheng, Wang-Cheng Kang, Mengting Wan, and Julian McAuley. 2019.  CosRec: 2D convolutional neural networks for sequential recommendation. In _CIKM_. 2173–2176.
  * Yeh et al. (2022) Chin-Chia Michael Yeh, Mengting Gu, Yan Zheng, Huiyuan Chen, Javid Ebrahimi, Zhongfang Zhuang, Junpeng Wang, Liang Wang, and Wei Zhang. 2022\.  Embedding Compression with Hashing for Efficient Representation Learning in Graph. In _KDD_.
  * Yin and Zhou (2019) Mingzhang Yin and Mingyuan Zhou. 2019.  ARM: Augment-REINFORCE-merge gradient for stochastic binary networks. In _ICLR_.
  * Yuan et al. (2019) Fajie Yuan, Alexandros Karatzoglou, Ioannis Arapakis, Joemon M Jose, and Xiangnan He. 2019.  A simple convolutional generative network for next item recommendation. In _WSDM_. 582–590.
  * Zaheer et al. (2020) Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020\.  Big bird: Transformers for longer sequences. In _NeurIPS_. 17283–17297.
  * Zhou et al. (2021) Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. 2021.  Informer: Beyond efficient transformer for long sequence time-series forecasting. In _AAAI_. 11106–11115.
  * Zhou et al. (2020) Wangchunshu Zhou, Tao Ge, Furu Wei, Ming Zhou, and Ke Xu. 2020.  Scheduled DropHead: A Regularization Method for Transformer Models. In _EMNLP_. 1971–1980.


# RL with KL penalties is better viewed as Bayesian inference

T  Tomasz Korbak
University of Sussex
<EMAIL_ADDRESS>
&Ethan Perez
New York University
<EMAIL_ADDRESS>
&Christopher L Buckley
University of Sussex
<EMAIL_ADDRESS>

###### Abstract

Reinforcement learning (RL) is frequently employed in fine-tuning large
language models (LMs), such as GPT-3, to penalize them for undesirable
features of generated sequences, such as offensiveness, social bias,
harmfulness or falsehood. The RL formulation involves treating the LM as a
policy and updating it to maximise the expected value of a reward function
which captures human preferences, such as non-offensiveness. In this paper, we
analyze challenges associated with treating a language model as an RL policy
and show how avoiding those challenges requires moving beyond the RL paradigm.
We start by observing that the standard RL approach is flawed as an objective
for fine-tuning LMs because it leads to distribution collapse: turning the LM
into a degenerate distribution. Then, we analyze KL-regularised RL, a widely
used recipe for fine-tuning LMs, which additionally constrains the fine-tuned
LM to stay close to its original distribution in terms of Kullback-Leibler
(KL) divergence. We show that KL-regularised RL is equivalent to variational
inference: approximating a Bayesian posterior which specifies how to update a
prior LM to conform with evidence provided by the reward function. We argue
that this Bayesian inference view of KL-regularised RL is more insightful than
the typically employed RL perspective. The Bayesian inference view explains
how KL-regularised RL avoids the distribution collapse problem and offers a
first-principles derivation for its objective. While this objective happens to
be equivalent to RL (with a particular choice of parametric reward), there
exist other objectives for fine-tuning LMs which are no longer equivalent to
RL. That observation leads to a more general point: RL is not an adequate
formal framework for problems such as fine-tuning language models. These
problems are best viewed as Bayesian inference: approximating a pre-defined
target distribution.

Keywords:

language models, controllable language generation, distribution collapse, KL-
control, reinforcement learning from human feedback, Bayesian inference,
variational inference

### Acknowledgements

Tomasz Korbak was supported by the Leverhulme Doctoral Scholarship and
OpenPhilantropy. This paper greatly benefited from discussions with Hady
Elsahar, Marc Dymetman and Germán Kruszewski.

##  1 Introduction

Large language models (LMs), such as GPT-3 Brown et al. (2020), tend to
generate outputs that reflect undesirable features of their training data such
as offensiveness (Gehman et al., 2020), social bias (Bender et al., 2021),
harmfulness (Bai et al., 2022) or dishonesty (Lin et al., 2021). Addressing
these biases and constraining LMs to be honest, helpful and harmless is an
essential part of the problem of aligning LMs with human preferences (Askell
et al., 2021). One intuitive approach to aligning LMs is reinforcement
learning (RL): capturing human preferences as a reward function and fine-
tuning the LM to maximise the reward expected under LM distribution. A
practical recipe for implementing this idea is RL from human feedback (Ziegler
et al., 2019): first, a reward model is trained to predict which of two texts
a human prefers and then a pretrained LM is fine-tuned to maximise reward
given by the reward model while also being penalized for Kullback-Leibler (KL)
divergence from its initial distribution. However, despite immense popularity
of RL from human feedback (Stiennon et al., 2020; Perez et al., 2022; Bai et
al., 2022), the motivation for KL penalty is not widely understood.

In this paper, we discuss an underappreciated perspective on KL-regularised RL
— the objective employed by RL from human feedback for fine-tuning LMs — which
explains its empirical success. We start with describing a problem that arises
from naively applying the standard RL objective: distribution collapse. The
optimal policy under the RL objective would be a minimal-entropy LM generating
a small set of sequences that obtain the highest reward. Then, we discuss how
KL-regularised RL avoids distribution collapse due to its KL penalty. This
constraint, we argue, transforms the problem from RL to Bayesian inference:
updating a prior to conform with evidence provided by the reward. The Bayesian
perspective moves KL-regularised RL closer to other divergence-minimisation-
based approaches to fine-tuning LMs (Khalifa et al., 2021) and, more broadly,
to other divergence-minimisation-based accounts of agency (Friston et al.,
2010; Levine, 2018; Hafner et al., 2020). These divergence minimisation
approaches naturally avoid the distribution collapse problem because they
formalize the agent as a generative model. In contrast, RL avoids distribution
collapse only with reward functions that make it equivalent to divergence
minimisation. Therefore, we conclude, RL is not an adequate formal framework
for problems such as fine-tuning LMs.

##  2 Fine-tuning language models via standard RL

Let $\mathcal{X}$ be the set of sequences of tokens from some vocabulary. An
LM $\pi$ can be seen as a probability distribution over $\mathcal{X}$. While
most modern LMs are autoregressive, for simplicity we will only talk about
full sequences, e.g. $\pi(x)$ denotes the probability of a sequence
$x\in\mathcal{X}$. Similarly, a reward function $r$ assigns sequences
$x\in\mathcal{X}$ with scalar rewards. In practice, $r(x)$ could represent
human preferences we would like $\pi$ to be aligned with, e.g. a non-
offensiveness reward would assign low values to sequences that are offensive.

If $\pi_{\theta}$ is our parametric LM (with parameters $\theta$), the RL
objective for fine-tuning it with our reward function $r$ is just the reward
expected under LM distribution:

$J_{\text{RL}}(\theta)=\mathbb{E}_{x\sim\pi_{\theta}}r(x).$ (1)

Intuitively, maximising $J_{\text{RL}}(\theta)$ means sampling a number of
sequences from the LM and rewarding the LM for good sequences and penalising
for bad ones (e.g. offensive sentences). This approach to LM alignment is
appealing in several ways, especially when compared with the standard self-
supervised language modelling objective of predicting the next token in a
static dataset. Because the samples come from the LM itself (as opposed to a
static dataset), the sampling distribution naturally follows what the LM has
already learned and the reward is only evaluated on LM’s current best guesses
about the correct behaviour. For instance, if the reward is non-offensiveness
which involves, but is not limited to avoiding curse words, the LM could
quickly learn to avoid curses and then focus on avoiding more elaborate forms
of toxicity, wasting no time on samples containing curse words.

The problem with the RL objective is that it treats the LM as a policy, not as
a generative model. While a generative model is supposed to capture a diverse
distribution of samples, a policy is supposed to chose the optimal action.
Since we don’t have a notion of state for LMs, the RL objective reduces to
searching for $x^{*}$, the sequence with highest reward. If there is one, the
optimal policy $\pi^{*}$ is a degenerate, deterministic generative model that
puts entire probability mass on that single sequence:

$\pi^{*}=\text{argmax}_{\theta}J_{\text{RL}}(\theta)=\delta_{x^{*}},$ (2)

where $\delta_{x^{*}}$ is a Dirac delta distribution centred on $x^{*}$. If
there are multiple optimal sequences $x^{*}$, probability mass would be put
only on them.

This failure mode is not purely theoretical. Empirically, distribution
collapse induced by maximising reward manifests as decreased fluency and
diversity of samples from the LM, which can be measured in terms of
perplexity, entropy and the frequency of repetitions. Degeneration of this
kind was observed in multiple language generation tasks ranging from
translation (Choshen et al., 2019), summarisation (Paulus et al., 2018), story
generation (Tambwekar et al., 2019), video captioning (Pasunuru & Bansal,
2017), dialogue (Jaques et al., 2019), to code generation (Korbak et al.,
2021) and LM debiasing (Khalifa et al., 2021).

While the distribution collapse problem is exacerbated by RL failure modes
such as insufficient exploration or reward hacking, it is distinct from
exploration-exploitation trade-off or reward misspecification. Even with
perfect exploration (if we sampled sequences uniformly from $\mathcal{X}$ as
opposed to sampling from $\pi_{\theta}$), the optimal policy will still put
all probability mass on $x^{*}$. Similarly, even if $r$ perfectly captures
human preferences across the whole space of possible sequences $\mathcal{X}$
and if $x^{*}$ is truly the best thing, we still wouldn’t want the LM to
generate only $x^{*}$. Essentially, the distribution collapse problem arises
from the fact that the RL objective for LM alignment is flawed: it doesn’t
care about preserving distributional properties of an LM and will always
penalise the LM for putting any probability mass on non-optimal sequences
until the LM collapses into a degenerate distribution.

##  3 Fine-tuning language models via KL-regularised RL

There is an obvious solution to the distribution collapse problem: including
preserving distributional properties of an LM as part of the reward function.
The notion of preserving distributional properties of an LM $\pi_{\theta}$ can
be formalised as penalising for Kullback-Leibler (KL) divergence between
$\pi_{\theta}$ and some other, pretrained LM $\pi_{0}$. Typically,
$\pi_{\theta}$ is initialised to $\pi_{0}$ and then fine-tuned to maximise the
following objective:

$J_{\text{KL-
RL}}(\theta)=\mathbb{E}_{x\sim\pi_{\theta}}[r(x)]-\beta\text{KL}(\pi_{\theta},\pi_{0}).$
(3)

The first term in (3) is equivalent to $J_{\text{RL}}(\theta)$ in (1) while
the second additionally constrains $\pi_{\theta}$ to stay close (in terms of
KL) to $\pi_{0}$. Almost always some reward needs to be sacrificed for that;
the coefficient $\beta$ determines the trade-off of how much reward is needed
to justify departing from $\pi_{0}$ by a certain distance. This objective is
commonly used as part of a popular recipe for fine-tuning LMs termed “RL from
Human Feedback” (RLHF) and works surprisingly well in practice (Ziegler et
al., 2019; Stiennon et al., 2020; Perez et al., 2022; Bai et al., 2022).
Earlier approaches to fine-tuning LMs employing this objective used the called
it “conservative fine-tuning” (Jaques et al., 2017) or KL-control Jaques et
al. (2019). Here, we focus only on the policy optimisation part of this setup,
which we term “KL-regularised RL”.

The KL-regularised RL objective (3) can easily be reformulated as just
expected reward as in (1). We only have to define a new reward function
$r^{\prime}_{\theta}(x)$ which incorporates the original reward $r$ and the KL
penalty, using the definition of KL divergence:

$\displaystyle J_{\text{KL-RL}}(\theta)$
$\displaystyle=\mathbb{E}_{x\sim\pi_{\theta}}[r^{\prime}_{\theta}(x)],\text{where}$
(4) $\displaystyle r^{\prime}_{\theta}(x)$
$\displaystyle=r(x)+\beta(\log\pi_{0}(x)-\log\pi_{\theta}(x)).$

This new reward function additionally rewards sequences likely under $\pi_{0}$
(therefore fluent) and unlikely under $\pi_{\theta}$ itself (an entropy
bonus). But even in (4), $J_{\text{KL-RL}}(\theta)$ is not a standard RL
objective: now the reward depends on policy parameters $\theta$, which makes
it non-stationary and coupled with $\pi_{\theta}$. But is framing the
maximisation of (4) as RL really necessary? In the next section, we will
develop an alternative view of this objective – as an approximate solution to
a Bayesian inference problem – and argue that it is more appealing than the RL
framing.

##  4 KL-regularised RL as variational inference

Fine-tuning a pretrained LM $\pi_{0}$ to align with preferences encoded by a
reward function $r$ is essentially a Bayesian inference problem. Intuitively,
Bayesian inference is the problem updating a distribution to conform with new
evidence. Given the prior probability $p(h)$ of a hypothesis $h$ and
likelihood $p(e|h)$ of evidence $e$ assuming $h$, the posterior probability of
$h$ is given by the Bayes theorem: $p(h|e)\propto p(e|h)p(h)$. In our setting,
we’re updating $\pi_{\theta}$, which initially equal to a prior $\pi_{0}$ to
conform with evidence provided by the assumption that $\pi_{\theta}$ is
optimal in terms of $r$. A reward function can be represented as a
distribution over $\mathcal{X}$ that makes high-reward sequences more likely
that low-reward sequences. A simple way of doing that is exponentiating reward
$r$ and renormalizing it. Then, the posterior is given by:

$\pi^{*}_{\text{KLRL}}(x)=\frac{1}{Z}\pi_{0}(x)\exp(r(x)/\beta),$ (5)

where $\pi_{0}$ is the prior, $\exp(r(x)/\beta)$ is the evidence provided by
the reward function (scaled by temperature $\beta$) and $Z$ is a constant
ensuring that $\pi^{*}_{\text{KLRL}}$ is a normalised probability
distribution. $\pi^{*}_{\text{KLRL}}$ represents a version $\pi_{0}$ updated
to account for the reward $r$. It also happens to coincide with the optimal
policy for $J_{\text{KLRL}}$:

$\pi^{*}_{\text{KLRL}}=\text{argmax}_{\theta}J_{\text{KLRL}}(\theta)$ (6)

Moreover, the KL-regularised RL objective can be cast as minimising the KL
divergence between the LM $\pi_{\theta}$ and this target distribution
$\pi^{*}_{\text{KLRL}}$:

$J_{\text{KLRL}}(\theta)\propto-\text{KL}(\pi_{\theta},\pi^{*}_{\text{KLRL}})$
(7)

This KL divergence is different KL that the penalty term
$\text{KL}(\pi_{\theta},\pi_{0})$ in (3). Minimising this new divergence
coincides with a variational inference (Blei et al., 2017), a well-known
approach to approximating Bayesian inference. More formally,
$J_{\text{KLRL}}(\theta)$ is the evidence lower bound (ELBO) on the log
likelihood of $\pi_{\theta}$ being optimal under $r$, assuming a prior
$\pi_{0}$. Minimising this bound makes $\pi_{\theta}$ approximate the true
posterior $\pi^{*}_{\text{KLRL}}$. A derivation of these equalities can be
found in the appendix below.

Why is this picture insightful? It explains where the KL penalty term
$\beta\text{KL}(\pi_{\theta},\pi_{0})$ in KL-regularised RL’s original
objective comes from. It is necessary to transform the problem from RL to
minimising a divergence from a target distribution $\pi^{*}_{\text{RLHP}}$.
This in turn makes the distributional character of an LM a first-class citizen
which explains why KL-regularised RL is able to maintain the fluency and
diversity of the original LM $\pi_{0}$.

There is a family of other divergence minimisation approaches to fine-tuning
LMs which are not equivalent to RL. Take Generative Distributional Control
(GDC) (Khalifa et al., 2021), an approach to fine-tuning LMs that obtains
results comparable with KL-regularised RL but minimises a slightly different
divergence:

$J_{\text{GDC}}(\theta)=-\text{KL}(\pi^{*}_{\text{GDC}},\pi_{\theta}),$ (8)

where $\pi^{*}_{\text{GDC}}$ is an exponential family distribution similar to
$\pi^{*}_{\text{KLRL}}$. The difference between $J_{\text{GDC}}(\theta)$ and
$J_{\text{KLRL}}(\theta)$ is in the order of arguments (forward vs reverse
KL). However, $J_{\text{GDC}}(\theta)$ is no longer equivalent to RL (Korbak
et al., 2022) because the expectation in forward KL divergence is with respect
to a $\pi^{*}_{\text{KLRL}}$, not $\pi_{\theta}$. Similarly, standard
supervised training objective can be seen as minimising
$\text{KL}(\pi^{*}_{\text{MLE}},\pi_{\theta})$, a divergence from the
empirical distribution $\pi^{*}_{\text{MLE}}$ provided by the training set.

One can therefore mount a double dissociation argument in favour of the
divergence minimisation perspective on KL-regularised RL: RL without
distribution matching fails, divergence minimisation without RL works.
Therefore, it’s the divergence minimisation aspect of KL-regularised RL that
accounts for its success. In consequence, calling it RL is just a
redescription of it that happens to be correct under a particular choice of
reward function $r^{\prime}_{\theta}$, but does not provide motivation for
this choice of $r^{\prime}_{\theta}$ and does not hold for alternative
divergence minimisation approaches to fine-tuning LMs such as GDC.

##  5 Conclusion

The divergence minimisation perspective on KL-regularised RL we presented
stems from a general framework known as control as inference Levine (2018).
Control as inference provides a formalisation of intelligent decision making
as inference on a probabilistic graphical model representing the agent, its
preferences and environmental dynamics. While control as inference is
typically considered with graphical models parameterised to make it equivalent
to RL, it does not have to be. Moreover, there are frameworks such as active
inference (Friston et al., 2010; Buckley et al., 2017) or action and
perception as divergence minimisation (Hafner et al., 2020) that further
generalise control as inference to a principle of minimising the KL divergence
from a probability distribution representing desired behaviour of the agent.
In contrast with RL, they conceptualise the agent as a generative model, not
as a decision rule represented as a probability distribution out of
convenience. Therefore, they naturally avoid the distribution collapse problem
and preserve the distributional properties of the agent. In consequence,
frameworks based on Bayesian inference or divergence minimisation are more
adequate than RL for problems such as fine-tuning language models.

## References

  * Askell et al. (2021) Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, Nelson Elhage, Zac Hatfield-Dodds, Danny Hernandez, Jackson Kernion, Kamal Ndousse, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, and Jared Kaplan.  A general language assistant as a laboratory for alignment, 2021.  URL https://arxiv.org/abs/2112.00861.
  * Bai et al. (2022) Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, Nicholas Joseph, Saurav Kadavath, Jackson Kernion, Tom Conerly, Sheer El-Showk, Nelson Elhage, Zac Hatfield-Dodds, Danny Hernandez, Tristan Hume, Scott Johnston, Shauna Kravec, Liane Lovitt, Neel Nanda, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, Ben Mann, and Jared Kaplan.  Training a helpful and harmless assistant with reinforcement learning from human feedback, 2022.
  * Bender et al. (2021) Emily M. Bender, Timnit Gebru, Angelina McMillan-Major, and Shmargaret Shmitchell.  On the dangers of stochastic parrots: Can language models be too big?  In _Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency_ , FAccT ’21, pp. 610–623, New York, NY, USA, 2021. Association for Computing Machinery.  ISBN 9781450383097.  doi: 10.1145/3442188.3445922.  URL https://doi.org/10.1145/3442188.3445922.
  * Blei et al. (2017) David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe.  Variational inference: A review for statisticians.  _Journal of the American Statistical Association_ , 112(518):859–877, apr 2017.  doi: 10.1080/01621459.2017.1285773.  URL https://doi.org/10.1080%2F01621459.2017.1285773.
  * Brown et al. (2020) Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  _CoRR_ , abs/2005.14165, 2020.  URL https://arxiv.org/abs/2005.14165.
  * Buckley et al. (2017) Christopher L. Buckley, Chang Sub Kim, Simon McGregor, and Anil K. Seth.  The free energy principle for action and perception: A mathematical review.  _Journal of Mathematical Psychology_ , 81:55–79, 2017.  ISSN 0022-2496.  doi: https://doi.org/10.1016/j.jmp.2017.09.004.  URL https://www.sciencedirect.com/science/article/pii/S0022249617300962.
  * Choshen et al. (2019) Leshem Choshen, Lior Fox, Zohar Aizenbud, and Omri Abend.  On the weaknesses of reinforcement learning for neural machine translation.  _arXiv preprint arXiv:1907.01752_ , 2019.
  * Friston et al. (2010) Karl J Friston, Jean Daunizeau, James Kilner, and Stefan J Kiebel.  Action and behavior: a free-energy formulation.  _Biological cybernetics_ , 102(3):227–260, 2010\.
  * Gehman et al. (2020) Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A Smith.  Realtoxicityprompts: Evaluating neural toxic degeneration in language models.  _arXiv preprint arXiv:2009.11462_ , 2020.
  * Hafner et al. (2020) Danijar Hafner, Pedro A. Ortega, Jimmy Ba, Thomas Parr, Karl Friston, and Nicolas Heess.  Action and perception as divergence minimization, 2020.
  * Jaques et al. (2017) Natasha Jaques, Shixiang Gu, Dzmitry Bahdanau, José Miguel Hernández-Lobato, Richard E. Turner, and Douglas Eck.  Sequence tutor: Conservative fine-tuning of sequence generation models with kl-control.  In Doina Precup and Yee Whye Teh (eds.), _Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 1645–1654. PMLR, 2017.  URL http://proceedings.mlr.press/v70/jaques17a.html.
  * Jaques et al. (2019) Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Àgata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind W. Picard.  Way off-policy batch deep reinforcement learning of implicit human preferences in dialog.  _CoRR_ , abs/1907.00456, 2019.  URL http://arxiv.org/abs/1907.00456.
  * Khalifa et al. (2021) Muhammad Khalifa, Hady Elsahar, and Marc Dymetman.  A distributional approach to controlled text generation.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=jWkw45-9AbL.
  * Korbak et al. (2021) Tomasz Korbak, Hady Elsahar, Marc Dymetman, and Germán Kruszewski.  Energy-based models for code generation under compilability constraints, 2021.
  * Korbak et al. (2022) Tomasz Korbak, Hady Elsahar, Germán Kruszewski, and Marc Dymetman.  On reward maximization and distribution matching for fine-tuning language models, 2022.  URL https://openreview.net/forum?id=8f95ajHrIFc.
  * Levine (2018) Sergey Levine.  Reinforcement learning and control as probabilistic inference: Tutorial and review, 2018.  URL https://arxiv.org/abs/1805.00909.
  * Lin et al. (2021) Stephanie Lin, Jacob Hilton, and Owain Evans.  Truthfulqa: Measuring how models mimic human falsehoods.  _arXiv preprint arXiv:2109.07958_ , 2021.
  * Pasunuru & Bansal (2017) Ramakanth Pasunuru and Mohit Bansal.  Reinforced video captioning with entailment rewards.  In Martha Palmer, Rebecca Hwa, and Sebastian Riedel (eds.), _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, EMNLP 2017, Copenhagen, Denmark, September 9-11, 2017_ , pp. 979–985. Association for Computational Linguistics, 2017.  doi: 10.18653/v1/d17-1103.  URL https://doi.org/10.18653/v1/d17-1103.
  * Paulus et al. (2018) Romain Paulus, Caiming Xiong, and Richard Socher.  A deep reinforced model for abstractive summarization.  In _6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings_. OpenReview.net, 2018.  URL https://openreview.net/forum?id=HkAClQgA-.
  * Perez et al. (2022) Ethan Perez, Saffron Huang, Francis Song, Trevor Cai, Roman Ring, John Aslanides, Amelia Glaese, Nat McAleese, and Geoffrey Irving.  Red teaming language models with language models, 2022.  URL https://arxiv.org/abs/2202.03286.
  * Stiennon et al. (2020) Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano.  Learning to summarize from human feedback.  In _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , NIPS’20, Red Hook, NY, USA, 2020. Curran Associates Inc.  ISBN 9781713829546.
  * Tambwekar et al. (2019) Pradyumna Tambwekar, Murtaza Dhuliawala, Lara J. Martin, Animesh Mehta, Brent Harrison, and Mark O. Riedl.  Controllable neural story plot generation via reward shaping.  In Sarit Kraus (ed.), _Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI 2019, Macao, China, August 10-16, 2019_ , pp. 5982–5988. ijcai.org, 2019.  doi: 10.24963/ijcai.2019/829.  URL https://doi.org/10.24963/ijcai.2019/829.
  * Ziegler et al. (2019) Daniel M. Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B. Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving.  Fine-tuning language models from human preferences.  _CoRR_ , abs/1909.08593, 2019.  URL http://arxiv.org/abs/1909.08593.

## Appendix

Let’s assume we have a prior distribution over sequences of tokens
$\pi_{0}(x)$ and a reward function $r$ which is (for technical reasons) always
negative (from $-\infty$ to 0). We can also represent $r$ as a binary random
variable $\mathcal{O}$ (the optimality variable). $\mathcal{O}=1$ if a certain
LM $\pi$ is optimal. We can define $\mathcal{O}$ in terms of $r$ as

$p(\mathcal{O}=1|x)=\exp(r(x))$ (9)

which is normalised because $r(x)$ is always negative. For instance, if $r(x)$
is a log probability that a sequence $x$ is non-offensive,
$p(\mathcal{O}=1|x)$ is a probability that $x$ is non-offensive and the
marginal $p(\mathcal{O}=1)$ is the average offensiveness score of $\pi$ (or a
probability that a random sample from $\pi$ is non-offensive). The problem of
aligning LMs can be seen as inferring $p(x|\mathcal{O}=1)$, a distribution
over sequences of tokens conditioned on being non-offensive. This can be
computed by applying Bayes rule as

$p(x|\mathcal{O}=1)=\frac{p(\mathcal{O}=1|x)p(x)}{p(\mathcal{O}=1)}=\frac{1}{Z}a(x)\exp(r(x)/\beta)$
(10)

where we chose the prior $p(x)=\pi_{0}(x)$, redefined the marginal
$p(\mathcal{O}=1)$ as the normalising constant $Z$, used the definition of
$p(\mathcal{O}=1|x)$ and chose $\beta=1$. $p(x|\mathcal{O}=1)$ here is
equivalent to $\pi^{*}_{\text{KL-RL}}$, the optimal policy under objective in
(3) (up to the choice of $\beta$ which can be absorbed into $r$).

$p(x|\mathcal{O}=1)$ is a non-parametric distribution: it doesn’t have to lie
in the family of distributions representable by a parametric model. In
general, we’d like to find a parametric model $\pi_{\theta}$ closest to
$\pi^{*}_{\text{KL-RL}}$. This can be formalised as finding $\pi_{\theta}$
minimising $\text{KL}(\pi_{\theta},\pi^{*}_{\text{KL-RL}})$. Here, however, we
will derive this objective from a yet more general perspective: inferring a
random latent variable $x$ that best explains the assumption that certain LM
$\pi$ is optimal given a prior $\pi_{0}(x)$. This can be seen as maximising
the log-likelihood of $\mathcal{O}=1$ via variational inference:

$\displaystyle\log p(\mathcal{O}=1)$
$\displaystyle=\log\sum_{x}p(\mathcal{O}=1,x)$ (11)
$\displaystyle=\log\Big{[}\sum_{x}p(\mathcal{O}=1|x)\pi_{0}(x)\Big{]}$ (12)
$\displaystyle=\log\Big{[}\sum_{a}\pi_{\theta}(x)p(\mathcal{O}=1|x)\frac{\pi_{0}(x)}{\pi_{\theta}(x)}\Big{]}$
(13)
$\displaystyle\geq\sum_{x}\pi_{\theta}(x)\log\Big{[}p(\mathcal{O}=1|x)\frac{\pi_{0}(x)}{\pi_{\theta}(x)}\Big{]}$
(14)
$\displaystyle=\mathbb{E}_{x\sim\pi_{\theta}}\log\Big{[}\exp(r(x))\frac{\pi_{0}(x)}{\pi_{\theta}(x)}\Big{]}$
(15)

In this derivation, we first introduce a latent variable $x$ using the sum
rule of probability (11), factorize a joint distribution (12), introduce a
variational distribution $\pi_{\theta}$ over that latent variable (13), use
Jensen’s inequality to obtain a bound (ELBo) (14) and, finally in (15), use
the definition of $p(\mathcal{O}=1|x)$. This new bound can be alternatively
expressed in two different ways:

$\mathbb{E}_{x\sim\pi_{\theta}}[r(x)]-\text{KL}(\pi_{\theta},a)$ (16)
$-\mathbb{E}_{x\sim\pi_{\theta}}\log\frac{\pi_{\theta}(x)}{a(x)\exp(r(x))}$
(17)

(16) is just KL-regularised RL objective $J_{\text{KL-RL}}(\theta)$ with
$\beta=1$. (17) is proportional (up to a constant $-\log Z$) to negative
$\text{KL}(\pi_{\theta},\pi^{*}_{\text{KL-RL}})$, where $\pi^{*}_{\text{KL-
RL}}=\frac{1}{Z}a(x)\exp(r(x))$ is the target distribution (or optimal policy
for $J_{\text{KL-RL}}(\theta)$). Their equivalence proves that KL-regularised
reward maximisation is equivalent to minimising divergence from
$\pi^{*}_{\text{KL-RL}}$.

More broadly, the derivation above shows that $J_{\text{KL-RL}}(\theta)$ can
be derived from first principles under a framework called (Levine, 2018). The
central idea here is to start from a Bayesian inference problem: inferring a
distribution over $x$ (an LM) that reconciles the assumption that this LM is
optimal ($p(\mathcal{O}=1)=1$, which plays we role of evidence) with a prior
$\pi_{0}(x)$. KL-regularised RL arises as we solve this inference problem
approximately via variational inference, i.e. by introducing a variational
distribution $\pi_{\theta}$ and optimising it to maximise a lower bound on
$\log p(\mathcal{O}=1)$.