[tab] 6-appendix.tex (latex)
\appendix
\section{Appendix}

\subsection{RLHF for conditional language modeling as MDP}
\label{app:lang_model_rl}
We can formulate conditional language modeling as a sequence to sequence problem. The goal is to map a source sequence $\vx = \left(x_1, x_2, ...x_L
ight)$ into a target sequence $\vy = \left(y_1, y_2, ....y_T
ight)$, that is to learn a mapping from $\vx$ to $\vy$. Machine translation is a classic example of a sequence to sequence problem \citep{sutskever2014sequence}.

The model can be described as a Markov Decision Process (MDP) \citep{bellman1957}. An MDP, $\gM \defeq (\gS, \gA, 	ransitions, r, d)$, consists of finite sets of states $\gS$ and actions $\gA$, a transition distribution $	ransitions(s'|s,a), s,s'\in\gS, a\in\gA$, a reward function $r:\gS 	imes\gA 
ightarrow \mathbb{R}$, and an initial state distribution $d: \gS 
ightarrow [0,1]$. In the offline setting the agent learns from a dataset $\data$ which contains sequences $\left(s_n, a_n, r_{n+1}
ight)$. The dataset $\data$ is often assumed to be generated by an unknown {\it behaviour policy} $\mu$ characterised by a distribution over actions conditioned on the state: $\mu(a|s)$. The elements of MDP have the following meaning in conditional language modeling:
\begin{itemize}
    \item 	extbf{States} ($s$): The state $s_n$ consists of the input sequence and the partially generated output up to the $n$-th token: $s_n = [\vx, \vy_{1:n-1}]$.
    \item 	extbf{Actions} ($a$): The actions are the tokens $a_n = y_n$ to be generated by policy ($\pi$).
    \item 	extbf{Rewards} (r): Rewards can be given or learned. In our experiments we train a deep neural network on human preferences to assigns a score to the full generated sequence. Reward is produced after end-of-sequence (EOS) token, and it is zero at all other steps.
   \item 	extbf{Transition operator} (T): The transition operator defines the dynamics of an environment: $	ransitions(s^\prime|a,s) := 	ransitions \left(y_n | s_n = [\vx, \vy_{1:n-1}]
ight)$. In our setup, it is a deterministic operator that concatenates newly produced token (action) to the current sequence (state).
\end{itemize}

This RLHF formulation of conditional language modeling can be also seen as a contextual bandit problem with a very large action space. %


\subsection{Negative results with offline RL}
\label{sec:negative_offline_rl}
In addition to the experiments from Section 
ef{sec:experiments}, we also run experiments with Q-learning with BVE type of one-step RL approaches and reward conditioned approaches, such as decision transformers \citep{chen2021decision}. However, we could not obtain notable improvements with them over the supervised baseline. Q-function-based methods for learning a policy and value function performed worse than supervised learning when used as a policy even after training from initialization from a supervised checkpoint. This matches previous observations by \cite{mathieu2021starcraft} who showed that offline Q-learning has difficulty in large action spaces  (vocabulary size of 
um{32000} tokens) and state-value based methods are preferable. Furthermore, as  \cite{brandfonbrener2022does} showed, return-conditioned methods tend to learn sub-optimal policies on tasks with sparse continuous rewards.

\subsection{Data and Model details}
\label{app:data_model}
In all our experiments, the base policy architecture is a minor modification of a standard Transformer architecture \citep{vaswani2017attention}. We use a vocabulary of 
um{32000} tokens and a max sequence length of 
um{128} at decoding time. We use the 	exttt{sentencepiece} tool \citep{kudo2018subword} for building the tokenizers.


\paragraph{\grow{} step}
During the \grow{} step, we sampled from the latest checkpoint of the policy with tempered softmax using temperature $0.8$ following the procedure proposed by \cite{li2022competition} to generate the dataset. Morevero, in our analysis, we found that temperature $0.8$ often covers a broad range of rewards in the dataset.

\paragraph{Thresholds in \improve{} step}
If we set the threshold $	au$ to a value such that $	au  \geq  V^{\pi_0}$, and learn a policy $\pi^{1}$ on data filtered with this threshold, the updated average value $V^{\pi_1}$ of this policy satisfies $V^{\pi_1} > V^{\pi_0}$. Then, iterative threshold ensures policy improvement in the \improve{} step. In our experiments, we do several improvement steps until we reach the maximum threshold. At each step, we trained the model for 
um{500000} SGD steps and pick the checkpoint with the best reward model score.

\paragraph{IWSLT 2014 De-En}
We use train, validation and test sets from IWSLT 2014 De-En dataset \citep{cettolo2014report} which includes source sentences in German with human translations (references) in English. In each \grow{} step, we generate 
um{100} candidate translation for each source sentence in the training set, effectively yielding us $|\data_g|=$
um{16000000}. For this dataset, we use a tiny version of the Transformer encoder-decoder \citep{vaswani2017attention} architecture with the feedforward MLP layers of size 
um{512}, feedforward dimension of 
um{1024}, 
um{4} attention heads and 
um{6} encoder and decoder layers.

\paragraph{WMT 2020 Zh-En}
We use the source-reference pairs in Chinese and English from the work of \cite{koehn2020findings} for our training, validation and test sets. Exact details on the datasets and preprocessing can be found in \cite{yu2020wmt}. In each \grow{} step, we generate 
um{25} candidates for each source sentence in the training set, effectively yielding us $|\data_g|=$ 
um{890000000}. We choose to generate 
um{25} candidates per source as the WMT dataset is significantly ($\sim$ 100 times) larger than the IWSLT dataset.  We use an architecture that mimics the Transformer-base encoder-decoder \citep{vaswani2017attention} architecture with model dimension 
um{512}, feedforward dimension of 
um{2048}, 
um{8} attention heads and 
um{6} encoder and decoder layers.

\paragraph{Web Domain En-Zh}
Finally, we use the in-house dataset for English to Chinese translation with custom training, fine-tuning and test sets. This training corpus is our biggest dataset \footnote{For computational reasons, we run   on the fine-tuning corpus with a model warm-started from a supervised model trained on the entire massive training corpus.}, so we use a modified version of Transformer-big encoder-decoder \citep{vaswani2017attention} architecture with model dimension 
um{1024}, feedforward dimension of 
um{8192}, 
um{16} attention heads and 
um{6} encoder and decoder layers. 


In Table 
ef{tbl:data_sources}, we list all the datasets with their sizes. In all the experiments, unless stated otherwise, we report the average reward scores on the validation set.
 \begin{table}[ht!]
 \caption{Details of the datasets and their sizes used in   experiments.}
 \begin{tabular}{@{}lrrrl@{}}
 	oprule
 Dataset    & $|\data_g|$ & \# Eval samples & \# Candidates\\& & & per source \\ \midrule
 IWSLT 2014 de-en & 
um{16000000}       &       
um{7466}          & 
um{100}                      \\
 WMT 2020 zh-en  & 
um{890000000}      &       
um{2000}          & 
um{25}                       \\
 Web Domain Finetuning en-zh  & 
um{3000000}        & 
um{6000}                 & 
um{1000}                     \\ \bottomrule
 \end{tabular}
 \label{tbl:data_sources}
 \end{table}

\subsection{Reward model}
\label{app:reward}
We used learned reward models that assign a score to the whole translation. We considered two types of reward models for translation \citep{freitag-etal-2022-results}: i) 	extit{Reference-free} reward models estimate of how good a translation is based on the source and candidate sentences only, ii) 	extit{Reference-based} reward models additionally use the reference human translation to decide how good a candidate translation is. The reference-free reward models are more flexible since they do not require reference translations. Moreover, by their design, reference-based models assign higher scores for sentences similar to the reference, and this could happen even when the reference is erroneous or incomplete. Finally, reference-free models open possibilities to evaluate and discover candidate translations that supersede the quality of references. Thus, in our experiments, we chose to work with reference-free reward models. On the other hand, since the reference-free reward models are not grounded in a human reference translation, it is more vulnerable to the distribution shifts in the candidates generated by the model and reward hacking. In practice, we pre-computed and stored $R(\vx, \vy)$ for the generated $\data_{g}$.

During the development of the model, we relied on the methodology of ``unit tests'' for reward models where we tested them in various hand-crafted settings, such as on various permutations and repetitions of sentences. The unit-tests ensured the high quality of rewards generated for  . Even though we used the most robust available reference-free reward model, we found that it was still not perfect, and sometimes showed signs of delusions. For example, the reward of a translation occasionally increased when we repeated the sentence, independent of the initial quality of the translation.

\subsection{Alignment between human evaluation scores and the reward model scores}
\label{app:human_evals}

In Figure 
ef{fig:iwslt_rest_violin_human_evals} and 
ef{fig:wmt_rest_violin_human_evals}, we showed the distribution of human preferences and reward model scores for   with \BC{} and $	exttt{G=1 I=4}$. Here, we found out that our reward model has a very high variance on samples with low human preference score from   (\BC{}, $	exttt{G=1 I=4}$) and our supervised learning baseline (\BC{}, $	exttt{G=0 I=0}$). Our hypothesis for this is that the training dataset for the reward model was dominated by translations of high quality \footnote{The data came from WMT 2021 metrics shared task dataset \citep{freitag2021results}, which reviews the human evaluation of good translation models}. Thus the model overfitted to the samples with high scores, and it has high variance in its prediction over the samples with lesser quality. This gap can be addressed with the incremental retraining of the reward model.

\begin{figure}[htbp!]
    \centering
    \includegraphics[width=0.8\linewidth]{assets/plots/iwslt_baseline_rest_twostep_violin_plots.png}
    \caption{	extbf{[IWSLT 2014 De-En]:} distribution of human preference and reward model scores for   (BC, $	exttt{I=4, G=1}$) in side by side evaluation with supervised model. The human preference scores less than or equal to 
um{3} have significantly higher variance in rewards compared to the human preference scores above 
um{3}. The plots on the left-hand side are for the samples generated from a supervised baseline and the right-hand side are for the samples generated with  .}
    \label{fig:iwslt_rest_violin_human_evals}
\end{figure}

\begin{figure}[htbp!]
    \centering
    \includegraphics[width=0.8\linewidth]{assets/plots/wmt_human_Eval_violin_rest.png}
    \caption{	extbf{[WMT 2020 Zh-En]:} distribution of human preference and reward model scores for   (BC, $	exttt{I=4, G=1}$) in side by side evaluation with supervised model. The human preference scores lower than or equal to 
um{3} has a variance in terms of reward model scores. The reward model has less certainty on the scores of the candidates having lower scores.}
    \label{fig:wmt_rest_violin_human_evals}
\end{figure}

\subsection{Results on test sets}
\label{app:test set results}
Figures 
ef{fig:improve_steps_rest_test} and 
ef{fig:wmt_rest_test} present the results on the test sets in IWSLT 2014 De-En and WMT 2020 Zh-En datasets. Similar to the validations set, $
est$ outperforms other value-based offline RL baselines. Also, increasing the number of both \grow{} and \improve{} steps helps to further increase the rewards. 

\begin{figure}[ht!]
    \centering
    \begin{floatrow}
    \includegraphics[scale=.24]{assets/plots/iwslt_multiple_improve_test.png}
    \hspace{.3cm}
    \includegraphics[scale=.24]{assets/plots/iwslt_second_grow_step_test.png}
    \end{floatrow}
    \caption{	extbf{[IWSLT 2014 De-En]: ReST results on test set.} All variations of   (purple) outperform supervised learning baseline (grey). As we increase the number of \grow{} and \improve{} steps, the reward of the model on the test set continues to increase. 
    }
    \label{fig:improve_steps_rest_test}
\end{figure}

\begin{figure}[ht!]
    \centering
    \begin{floatrow}
    \includegraphics[scale=.28]{assets/plots/wmt_multiple_improve_test.png}
    \hspace{.3cm}
    \includegraphics[scale=.28]{assets/plots/wmt_offline_algos_test.png}
    \end{floatrow}
    \caption{	extbf{[WMT 2020 Zh-En]: ReST results on test set.} We see that $
est$ clearly outperforms the baselines (right) and the reward of the model on the test set increases with the number of \improve{} steps.  
    }
    \label{fig:wmt_rest_test}
\end{figure}


\subsection{Additional experiments}
This section reports the results of several ablations in   which explain some of the design choices that we made. 

\paragraph{GOLD loss compared to BC loss in \improve{} step} 
In Figure 
ef{fig:wmt_gold_vs_bc}, we compare GOLD against BC losses with the increasing number of   \improve{} steps on WMT 2020 Zh-En dataset. The performance of both BC and GOLD improves with the number of \improve{} steps, but BC performs better than GOLD in every setting.

\begin{figure}[!ht]
    \RawFloats
    \centering
    \begin{minipage}{0.45\linewidth}
    \centering
    \includegraphics[width=0.96\linewidth,height=0.25	extheight]{assets/plots/wmt_GOLD_vs_BC_rest_ablations.png}
    \caption{	extbf{[WMT 2020 Zh-En]:} GOLD loss vs BC loss in multi-step improvement. The performance of all methods improves with more steps, but BC loss always performs better than other offline RL losses that we tried.}
    \label{fig:wmt_gold_vs_bc}
    \end{minipage}%
    \hspace{4mm}
    \begin{minipage}{0.45\linewidth}
    \centering
    \includegraphics[width=0.96\linewidth,height=0.27	extheight]{assets/plots/iwslt_baseline_comparisons.png}
    \caption{	extbf{[IWSLT 2014 De-En]:} BVMPO, BC, GOLD losses used in  . BC loss with three \improve{} steps yields the best results.}
    \label{fig:iwslt_baseline_comparisons}
    \end{minipage}
\end{figure}

\paragraph{Loss comparison on IWSLT 2014 De-En}
In Figure 
ef{fig:iwslt_baseline_comparisons}, we compare BVMPO, BC and GOLD losses in \improve{} step on IWSLT 2014. The results are consistent with WMT dataset: in short,   with BC loss and multiple \improve{} steps outperforms other approaches.


\paragraph{Selecting threshold based on percentiles of reward model scores}
Using a single threshold for all the source-candidate pairs may lead to a scenario with no training data for certain (harder) source sentences or contexts. Alternatively, we can filter the candidate sentences based on the percentile of the scores given their source. This way of filtering based on percentiles effectively results in changing the thresholds for each source sentence. The results are presented in Figure 
ef{fig:iwslt_percentile_comparisons}. As the filtering percentile increases, the performance increases, but it starts saturating, and the gains become smaller after $p=90$. The results are generally similar to those with global threshold-based filtering.

\begin{figure}[!t]
    \centering
    \begin{floatrow}
    \includegraphics[width=0.4\linewidth]{assets/plots/iwslt_percentile_comparisons.png}
    \hspace{.3cm}
    \includegraphics[width=0.4\linewidth]{assets/plots/iwslt_rest_percentile_improvements_comparisons.png}
    \end{floatrow}
    \caption{	extbf{[IWSLT 2014 De-En] percentiles}: Filtering data based on percentiles of reward scores for each source candidate.   performs better with increasing percentile and the performance saturates after $p=90$.}
    \label{fig:iwslt_percentile_comparisons}
\end{figure}

\paragraph{Selecting threshold based on linear interpolation between the mean and max scores of candidates}

An alternative way of specifying filters based on a source sentence is to compute the filter based on a max reward score ($V^{\max}(source)$) and the mean candidate score ($V^{	ext{mean}}(source)$) for a give source sentence. Then, we select the threshold for each source sentence as $	au_{\gamma}(source) = \gamma V^{\max}(source) + (1-\gamma) V^{	ext{mean}}(source)$. At each \improve{} step, we increase the parameter $\gamma$. We report the results with respect to different $\gamma$ values in Figure 
ef{fig:iwslt_interpolation}. In a nutshell, computing the threshold by interpolating the max and mean scores for a given candidate gives results similar to the percentile-based way of computing the thresholds per source. Also we can see that the schedule of thresholds and the number of \improve{} steps have a large impact on the final performance of the model.

\begin{figure}[t!]
    \centering
    \begin{floatrow}
    \includegraphics[width=0.35\linewidth]{assets/plots/iwslt_interpolation_comparisons.png}
    \hspace{.3cm}
    \includegraphics[width=0.52\linewidth]{assets/plots/iwslt_rest_interpolation_comparisons.png}
    \end{floatrow}
    \caption{	extbf{[IWSLT 2014 De-En] interpolation experiments:} Threshold in \improve{} step is chosen for each candidate sentence as $\gamma V^{\max}(source) + (1-\gamma) V^{	ext{mean}}(source)$. On the left, experiments with $	exttt{I=1}$ provide improvements and increasing $\gamma$, which also helps substantially. On the right, we present the results with $\gamma$ starting with $0.5$ instead of $0.5$. As it can be seen from this figure, the number of \improve{} steps and how the thresholds are selected greatly influence the final performance of the model. }
    \label{fig:iwslt_interpolation}
\end{figure}








\subsection{Baseline Algorithms}
\label{app:baseline_algos}

This section provides details on the baseline losses used in Figure~
ef{fig:baseline_comps}.

\subsubsection{BVMPO}
The BVMPO approach is similar to DIME proposed by \cite{abdolmaleki2021multi}. The main difference is that we use a state-value function instead of Q funtion with v-trace \citep{espeholt2018impala}, similarly to V-MPO \citep{song2019v}. We train separate neural networks for policy and value function. The policy is pre-trained with BC and the value network is pre-trained with BVE to estimate the behavior value $V^{\beta}$. BVMPO uses the behavioral cloning policy $\pi^{\beta}(a|s)$ as a behaviour prior:
\begin{equation}
    \gL_{BVMPO} = \gL_{VMPO} + \lambda 	extup{KL}(\pi^\beta | \pi).
\end{equation}
As opposed to DIME, BVMPO starts with a fixed coefficient $\lambda=1$ and anneals it to $1e-5$ during the training. We found that starting with a large $\lambda$ and annealing it stabilizes the training without any noticeable drop in performance. Similarly, annealing the behavior regularization co-efficient was also shown to be effective in offline Q-learning \citep{agarwal2022beyond}.

\subsubsection{GOLD}
The GOLD loss was introduced by \cite{pang2021}. When using it in \improve{} step, we start with a BC policy. In all GOLD experiments, we used $k=5$ steps to unroll the MLE model into the future for reward computations. Increasing $k$ did not improve the performance.

\subsubsection{Offline Actor Critic (OAC)}
Offline actor critic is an actor critic approach proposed by \cite{mathieu2021starcraft} where the state-value function is trained to estimate the behavior value in the dataset \citep{gulcehre2021regularized}. Unlike \cite{mathieu2021starcraft} we did not use the V-trace as it did not improve over the vanilla advantage function in our initial experiments.


\subsubsection{Value Functions}

We trained state-value functions on IWSLT 2014 De-En and WMT 2020 Zh-En datasets after the initial grow step. The value functions estimate the behavior value of the policy on the training dataset. The value functions predict the sum of discounted rewards for each word until the end of the sentence. We used a transformer-based encoder-decoder architecture for the value functions. The encoder gets the source sentence as the input, and the decoder predicts Monte Carlo discounted returns conditioned on the source and history of generated sequences of tokens in the target. 



On the training sets of IWSLT and WMT, our value functions achieved Kendall-tau correlation of $92.4\%$ and Spearman correlation above $99\%$ with respect to the reward model scores in the dataset. As we do RL only offline, the value functions do not need to generalize beyond the training dataset.







\subsection{ : Population interpretation}
\label{app:Restpopulation}
Here, we provide a population interpretation of the   algorithm. In this setup, the initial BC training step can be thought of as identifying $	heta$ by minimizing
\begin{align*}
    	extup{KL}(p(\vx,\vy) ||\pi_{	heta}(\vx,\vy))&=\sum_{\vx,\vy} p(\vx,\vy) \log \frac{p(\vx,\vy)}{\pi_{	heta}(\vx,\vy)}\\
    &=\sum_{\vx,\vy} p(\vx,\vy) \log \frac{p(\vy|\vx)}{\pi_{	heta}(\vy|\vx)}\\
    &=\mathbb{E}_{\vx \sim p(\vx)}[ 	extup{KL}(p(\vy|\vx) ||\pi_{	heta}(\vy |\vx))],
\end{align*}
where $	extup{KL}$ is the Kullback--Leibler discrepancy, $p(\vx,\vy)=p(\vx)p(\vy|\vx)$ is the (training) data distribution and $\pi_{	heta}(\vx,\vy)=p(\vx)\pi_{	heta}(\vy|\vx)$. Let us call the minimizer $	heta_0$. 

Assume now that one has the parameter $	heta_k$ at the end of \improve{} steps of the $k${th} \grow{} step. Then at the next $(k+1)${th} \grow{} step, we consider \improve{} steps using a sequence of increasing thresholds $(	au_i)_{i=1}^{I}$ which are used to filter the available data distribution (mixture of original data distribution and synthetic data distribution) to obtain a new distribution 
\begin{equation*}
    \pi_{	heta_k}^{	au_i}(\vx,\vy) \propto \{(1-\lambda) p(\vx,\vy)+\lambda \pi_{	heta_k}(\vx,\vy)\}\mathbb{I}(R(\vx,\vy)> 	au_i).
\end{equation*}
We then obtain the next parameter $	heta$ by minimizing the KL 
\begin{equation*}
      	extup{KL}(\pi^{	au_i}_{	heta_k}(\vx,\vy)||\pi_	heta(\vx,\vy))
\end{equation*}
and keep increasing the threshold until the reward $\mathbb{E}_{\pi_	heta(\vx,\vy)}[R(\vx,\vy)]$ stops increasing. This yields $	heta_{k+1}$.